support input tags for finetuning

Author: VipulMascarenhas

ads/aqua/finetuning/entities.py

@@ -80,6 +80,10 @@ class CreateFineTuningDetails(DataClassSerializable):
         The log id for fine tuning job infrastructure.
     force_overwrite: (bool, optional). Defaults to `False`.
         Whether to force overwrite the existing file in object storage.
+    freeform_tags: (dict, optional)
+        Freeform tags for the fine-tuning model
+    defined_tags: (dict, optional)
+        Defined tags for the fine-tuning model
     """
 
     ft_source_id: str
@@ -101,3 +105,5 @@ class CreateFineTuningDetails(DataClassSerializable):
     log_id: Optional[str] = None
     log_group_id: Optional[str] = None
     force_overwrite: Optional[bool] = False
+    freeform_tags: Optional[dict] = None
+    defined_tags: Optional[dict] = None

ads/aqua/finetuning/finetuning.py

@@ -35,7 +35,11 @@
     ENV_AQUA_FINE_TUNING_CONTAINER,
     FineTuneCustomMetadata,
 )
-from ads.aqua.finetuning.entities import *
+from ads.aqua.finetuning.entities import (
+    AquaFineTuningParams,
+    AquaFineTuningSummary,
+    CreateFineTuningDetails,
+)
 from ads.common.auth import default_signer
 from ads.common.object_storage_details import ObjectStorageDetails
 from ads.common.utils import get_console_link
@@ -100,14 +104,14 @@ def create(
         if not create_fine_tuning_details:
             try:
                 create_fine_tuning_details = CreateFineTuningDetails(**kwargs)
-            except:
+            except Exception as ex:
                 allowed_create_fine_tuning_details = ", ".join(
                     field.name for field in fields(CreateFineTuningDetails)
                 ).rstrip()
                 raise AquaValueError(
                     "Invalid create fine tuning parameters. Allowable parameters are: "
                     f"{allowed_create_fine_tuning_details}."
-                )
+                ) from ex
 
         source = self.get_source(create_fine_tuning_details.ft_source_id)
 
@@ -148,28 +152,27 @@ def create(
                 "Specify the subnet id via API or environment variable AQUA_JOB_SUBNET_ID."
             )
 
-        if create_fine_tuning_details.replica > DEFAULT_FT_REPLICA:
-            if not (
-                create_fine_tuning_details.log_id
-                and create_fine_tuning_details.log_group_id
-            ):
-                raise AquaValueError(
-                    f"Logging is required for fine tuning if replica is larger than {DEFAULT_FT_REPLICA}."
-                )
+        if create_fine_tuning_details.replica > DEFAULT_FT_REPLICA and not (
+            create_fine_tuning_details.log_id
+            and create_fine_tuning_details.log_group_id
+        ):
+            raise AquaValueError(
+                f"Logging is required for fine tuning if replica is larger than {DEFAULT_FT_REPLICA}."
+            )
 
         ft_parameters = None
         try:
             ft_parameters = AquaFineTuningParams(
                 **create_fine_tuning_details.ft_parameters,
             )
-        except:
+        except Exception as ex:
             allowed_fine_tuning_parameters = ", ".join(
                 field.name for field in fields(AquaFineTuningParams)
             ).rstrip()
             raise AquaValueError(
                 "Invalid fine tuning parameters. Fine tuning parameters should "
                 f"be a dictionary with keys: {allowed_fine_tuning_parameters}."
-            )
+            ) from ex
 
         experiment_model_version_set_id = create_fine_tuning_details.experiment_id
         experiment_model_version_set_name = create_fine_tuning_details.experiment_name
@@ -197,11 +200,11 @@ def create(
                     auth=default_signer(),
                     force_overwrite=create_fine_tuning_details.force_overwrite,
                 )
-            except FileExistsError:
+            except FileExistsError as fe:
                 raise AquaFileExistsError(
                     f"Dataset {dataset_file} already exists in {create_fine_tuning_details.report_path}. "
                     "Please use a new dataset file name, report path or set `force_overwrite` as True."
-                )
+                ) from fe
             logger.debug(
                 f"Uploaded local file {ft_dataset_path} to object storage {dst_uri}."
             )
@@ -222,6 +225,8 @@ def create(
             description=create_fine_tuning_details.experiment_description,
             compartment_id=target_compartment,
             project_id=target_project,
+            freeform_tags=create_fine_tuning_details.freeform_tags,
+            defined_tags=create_fine_tuning_details.defined_tags,
         )
 
         ft_model_custom_metadata = ModelCustomMetadata()
@@ -273,6 +278,10 @@ def create(
             Tags.AQUA_TAG: UNKNOWN,
             Tags.AQUA_FINE_TUNED_MODEL_TAG: f"{source.id}#{source.display_name}",
         }
+        ft_job_freeform_tags = {
+            **ft_job_freeform_tags,
+            **(create_fine_tuning_details.freeform_tags or {}),
+        }
 
         ft_job = Job(name=ft_model.display_name).with_infrastructure(
             DataScienceJob()
@@ -286,6 +295,7 @@ def create(
                 or DEFAULT_FT_BLOCK_STORAGE_SIZE
             )
             .with_freeform_tag(**ft_job_freeform_tags)
+            .with_defined_tag(**(create_fine_tuning_details.defined_tags or {}))
         )
 
         if not subnet_id:
@@ -353,6 +363,7 @@ def create(
         ft_job_run = ft_job.run(
             name=ft_model.display_name,
             freeform_tags=ft_job_freeform_tags,
+            defined_tags=create_fine_tuning_details.defined_tags or {},
             wait=False,
         )
         logger.debug(
@@ -372,22 +383,25 @@ def create(
             for metadata in ft_model_custom_metadata.to_dict()["data"]
         ]
 
-        source_freeform_tags = source.freeform_tags or {}
-        source_freeform_tags.pop(Tags.LICENSE, None)
-        source_freeform_tags.update({Tags.READY_TO_FINE_TUNE: "false"})
-        source_freeform_tags.update({Tags.AQUA_TAG: UNKNOWN})
-        source_freeform_tags.pop(Tags.BASE_MODEL_CUSTOM, None)
+        model_freeform_tags = source.freeform_tags or {}
+        model_freeform_tags.pop(Tags.LICENSE, None)
+        model_freeform_tags.pop(Tags.BASE_MODEL_CUSTOM, None)
+
+        model_freeform_tags = {
+            **model_freeform_tags,
+            Tags.READY_TO_FINE_TUNE: "false",
+            Tags.AQUA_TAG: UNKNOWN,
+            Tags.AQUA_FINE_TUNED_MODEL_TAG: f"{source.id}#{source.display_name}",
+            **(create_fine_tuning_details.freeform_tags or {}),
+        }
+        model_defined_tags = create_fine_tuning_details.defined_tags or {}
 
         self.update_model(
             model_id=ft_model.id,
             update_model_details=UpdateModelDetails(
                 custom_metadata_list=updated_custom_metadata_list,
-                freeform_tags={
-                    Tags.AQUA_FINE_TUNED_MODEL_TAG: (
-                        f"{source.id}#{source.display_name}"
-                    ),
-                    **source_freeform_tags,
-                },
+                freeform_tags=model_freeform_tags,
+                defined_tags=model_defined_tags,
             ),
         )
 
@@ -462,12 +476,16 @@ def create(
                     region=self.region,
                 ),
             ),
-            tags=dict(
-                aqua_finetuning=Tags.AQUA_FINE_TUNING,
-                finetuning_job_id=ft_job.id,
-                finetuning_source=source.id,
-                finetuning_experiment_id=experiment_model_version_set_id,
-            ),
+            tags={
+                **{
+                    "aqua_finetuning": Tags.AQUA_FINE_TUNING,
+                    "finetuning_job_id": ft_job.id,
+                    "finetuning_source": source.id,
+                    "finetuning_experiment_id": experiment_model_version_set_id,
+                },
+                **model_freeform_tags,
+                **model_defined_tags,
+            },
             parameters={
                 key: value
                 for key, value in asdict(ft_parameters).items()
@@ -635,6 +653,6 @@ def validate_finetuning_params(self, params: Dict = None) -> Dict:
             raise AquaValueError(
                 f"Invalid fine tuning parameters. Allowable parameters are: "
                 f"{allowed_fine_tuning_parameters}."
-            )
+            ) from e
 
-        return dict(valid=True)
+        return {"valid": True}

tests/unitary/with_extras/aqua/test_finetuning.py

@@ -119,6 +119,9 @@ def test_create_fine_tuning(
         self.app.ds_client.update_model = MagicMock()
         self.app.ds_client.update_model_provenance = MagicMock()
 
+        ft_model_freeform_tags = {"ftag1": "fvalue1", "ftag2": "fvalue2"}
+        ft_model_defined_tags = {"dtag1": "dvalue1", "dtag2": "dvalue2"}
+
         create_aqua_ft_details = dict(
             ft_source_id="ocid1.datasciencemodel.oc1.iad.<OCID>",
             ft_name="test_ft_name",
@@ -134,6 +137,8 @@ def test_create_fine_tuning(
             validation_set_size=0.2,
             block_storage_size=1,
             experiment_name="test_experiment_name",
+            freeform_tags=ft_model_freeform_tags,
+            defined_tags=ft_model_defined_tags,
         )
 
         aqua_ft_summary = self.app.create(**create_aqua_ft_details)
@@ -167,10 +172,14 @@ def test_create_fine_tuning(
                 "url": f"https://cloud.oracle.com/data-science/models/{ft_source.id}?region={self.app.region}",
             },
             "tags": {
-                "aqua_finetuning": "aqua_finetuning",
-                "finetuning_experiment_id": f"{mock_mvs_create.return_value[0]}",
-                "finetuning_job_id": f"{mock_job_id.return_value}",
-                "finetuning_source": f"{ft_source.id}",
+                **{
+                    "aqua_finetuning": "aqua_finetuning",
+                    "finetuning_experiment_id": f"{mock_mvs_create.return_value[0]}",
+                    "finetuning_job_id": f"{mock_job_id.return_value}",
+                    "finetuning_source": f"{ft_source.id}",
+                },
+                **ft_model_freeform_tags,
+                **ft_model_defined_tags,
             },
             "time_created": f"{ft_model.time_created}",
         }