From fe0fc1789c17f3a6bd83ec6b43f12d0d012b7122 Mon Sep 17 00:00:00 2001 From: Sai Atmakuri Date: Sat, 2 Sep 2023 00:06:02 +0000 Subject: [PATCH] Pass labels to job config --- .../docker_image_batch_job_llm_fine_tuning_service.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/model-engine/model_engine_server/infra/services/docker_image_batch_job_llm_fine_tuning_service.py b/model-engine/model_engine_server/infra/services/docker_image_batch_job_llm_fine_tuning_service.py index 008d5dfc..9e25b8cf 100644 --- a/model-engine/model_engine_server/infra/services/docker_image_batch_job_llm_fine_tuning_service.py +++ b/model-engine/model_engine_server/infra/services/docker_image_batch_job_llm_fine_tuning_service.py @@ -73,10 +73,14 @@ async def create_fine_tune( if not di_batch_job_bundle.public and di_batch_job_bundle.owner != owner: raise LLMFineTuningMethodNotImplementedException("Fine-tuning method not accessible") + # TODO: Pass user-defined labels + labels = dict(team="egp", product="llm-fine-tune") + batch_job_id = await self.docker_image_batch_job_gateway.create_docker_image_batch_job( created_by=created_by, owner=owner, job_config=dict( + **labels, gateway_url=os.getenv("GATEWAY_URL"), user_id=owner, training_file=training_file, @@ -97,8 +101,7 @@ async def create_fine_tune( gpu_type=di_batch_job_bundle.gpu_type, storage=di_batch_job_bundle.storage, ), - # TODO: Pass user-defined labels - labels=dict(team="egp", product="llm-fine-tune"), + labels=labels, annotations=dict(fine_tuned_model=fine_tuned_model), mount_location=di_batch_job_bundle.mount_location, )