support input tags for model creation

VipulMascarenhas · VipulMascarenhas · commit 1583829b652e · 2024-12-09T14:54:18.000-08:00
diff --git a/ads/aqua/app.py b/ads/aqua/app.py
@@ -2,6 +2,7 @@
 # Copyright (c) 2024 Oracle and/or its affiliates.
 # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
 
+import json
 import os
 from dataclasses import fields
 from typing import Dict, Union
@@ -135,6 +136,8 @@ def create_model_version_set(
         description: str = None,
         compartment_id: str = None,
         project_id: str = None,
+        freeform_tags: dict = None,
+        defined_tags: dict = None,
         **kwargs,
     ) -> tuple:
         """Creates ModelVersionSet from given ID or Name.
@@ -153,7 +156,10 @@ def create_model_version_set(
             Project OCID.
         tag: (str, optional)
             calling tag, can be Tags.AQUA_FINE_TUNING or Tags.AQUA_EVALUATION
-
+        freeform_tags: (dict, optional)
+            Freeform tags for the model version set
+        defined_tags: (dict, optional)
+            Defined tags for the model version set
         Returns
         -------
         tuple: (model_version_set_id, model_version_set_name)
@@ -182,13 +188,15 @@ def create_model_version_set(
                 mvs_freeform_tags = {
                     tag: tag,
                 }
+                mvs_freeform_tags = {**mvs_freeform_tags, **(freeform_tags or {})}
                 model_version_set = (
                     ModelVersionSet()
                     .with_compartment_id(compartment_id)
                     .with_project_id(project_id)
                     .with_name(model_version_set_name)
                     .with_description(description)
                     .with_freeform_tags(**mvs_freeform_tags)
+                    .with_defined_tags(**(defined_tags or {}))
                     # TODO: decide what parameters will be needed
                     # when refactor eval to use this method, we need to pass tag here.
                     .create(**kwargs)
@@ -340,7 +348,9 @@ def build_cli(self) -> str:
         """
         cmd = f"ads aqua {self._command}"
         params = [
-            f"--{field.name} {getattr(self,field.name)}"
+            f"--{field.name} {json.dumps(getattr(self, field.name))}"
+            if isinstance(getattr(self, field.name), dict)
+            else f"--{field.name} {getattr(self, field.name)}"
             for field in fields(self.__class__)
             if getattr(self, field.name) is not None
         ]
diff --git a/ads/aqua/model/entities.py b/ads/aqua/model/entities.py
@@ -291,6 +291,8 @@ class ImportModelDetails(CLIBuilderMixin):
     inference_container_uri: Optional[str] = None
     allow_patterns: Optional[List[str]] = None
     ignore_patterns: Optional[List[str]] = None
+    freeform_tags: Optional[dict] = None
+    defined_tags: Optional[dict] = None
 
     def __post_init__(self):
         self._command = "model register"
diff --git a/ads/aqua/model/model.py b/ads/aqua/model/model.py
@@ -127,7 +127,13 @@ class AquaModelApp(AquaApp):
 
     @telemetry(entry_point="plugin=model&action=create", name="aqua")
     def create(
-        self, model_id: str, project_id: str, compartment_id: str = None, **kwargs
+        self,
+        model_id: str,
+        project_id: str,
+        compartment_id: str = None,
+        freeform_tags: Optional[dict] = None,
+        defined_tags: Optional[dict] = None,
+        **kwargs,
     ) -> DataScienceModel:
         """Creates custom aqua model from service model.
 
@@ -140,7 +146,10 @@ def create(
         compartment_id: str
             The compartment id for custom model. Defaults to None.
             If not provided, compartment id will be fetched from environment variables.
-
+        freeform_tags: dict
+            Freeform tags for the model
+        defined_tags: dict
+            Defined tags for the model
         Returns
         -------
         DataScienceModel:
@@ -157,15 +166,25 @@ def create(
             )
             return service_model
 
+        # combine tags
+        combined_freeform_tags = {
+            **(service_model.freeform_tags or {}),
+            **(freeform_tags or {}),
+        }
+        combined_defined_tags = {
+            **(service_model.defined_tags or {}),
+            **(defined_tags or {}),
+        }
+
         custom_model = (
             DataScienceModel()
             .with_compartment_id(target_compartment)
             .with_project_id(target_project)
             .with_model_file_description(json_dict=service_model.model_file_description)
             .with_display_name(service_model.display_name)
             .with_description(service_model.description)
-            .with_freeform_tags(**(service_model.freeform_tags or {}))
-            .with_defined_tags(**(service_model.defined_tags or {}))
+            .with_freeform_tags(**combined_freeform_tags)
+            .with_defined_tags(**combined_defined_tags)
             .with_custom_metadata_list(service_model.custom_metadata_list)
             .with_defined_metadata_list(service_model.defined_metadata_list)
             .with_provenance_metadata(service_model.provenance_metadata)
@@ -414,7 +433,7 @@ def edit_registered_model(self, id, inference_container, enable_finetuning, task
                         except Exception as ex:
                             raise AquaRuntimeError(
                                 f"The given model already doesn't support finetuning: {ex}"
-                            )
+                            ) from ex
 
                 custom_metadata_list.remove("modelDescription")
                 if task:
@@ -766,6 +785,8 @@ def _create_model_catalog_entry(
         compartment_id: Optional[str],
         project_id: Optional[str],
         inference_container_uri: Optional[str],
+        freeform_tags: Optional[dict] = None,
+        defined_tags: Optional[dict] = None,
     ) -> DataScienceModel:
         """Create model by reference from the object storage path
 
@@ -778,6 +799,8 @@ def _create_model_catalog_entry(
             compartment_id (Optional[str]): Compartment Id of the compartment where the model has to be created
             project_id (Optional[str]): Project id of the project where the model has to be created
             inference_container_uri (Optional[str]): Inference container uri for BYOC
+            freeform_tags (dict): Freeform tags for the model
+            defined_tags (dict): Defined tags for the model
 
         Returns:
             DataScienceModel: Returns Datascience model instance.
@@ -918,13 +941,16 @@ def _create_model_catalog_entry(
             category="Other",
             replace=True,
         )
+        # override tags with freeform tags if set
+        tags = {**tags, **(freeform_tags or {})}
         model = (
             model.with_custom_metadata_list(metadata)
             .with_compartment_id(compartment_id or COMPARTMENT_OCID)
             .with_project_id(project_id or PROJECT_OCID)
             .with_artifact(os_path)
             .with_display_name(model_name)
             .with_freeform_tags(**tags)
+            .with_defined_tags(**(defined_tags or {}))
         ).create(model_by_reference=True)
         logger.debug(model)
         return model
@@ -1314,7 +1340,7 @@ def _download_model_from_hf(
             os_path=os_path,
             local_dir=local_dir,
             model_name=model_name,
-            exclude_pattern=f"{HF_METADATA_FOLDER}*"
+            exclude_pattern=f"{HF_METADATA_FOLDER}*",
         )
 
         return model_artifact_path
@@ -1402,6 +1428,8 @@ def register(
             compartment_id=import_model_details.compartment_id,
             project_id=import_model_details.project_id,
             inference_container_uri=import_model_details.inference_container_uri,
+            freeform_tags=import_model_details.freeform_tags,
+            defined_tags=import_model_details.defined_tags,
         )
         # registered model will always have inference and evaluation container, but
         # fine-tuning container may be not set
diff --git a/tests/unitary/with_extras/aqua/test_model.py b/tests/unitary/with_extras/aqua/test_model.py
@@ -748,13 +748,13 @@ def test_import_verified_model(
                     local_dir=str(tmpdir),
                     download_from_hf=True,
                     allow_patterns=["*.json"],
-                    ignore_patterns=["test.json"]
+                    ignore_patterns=["test.json"],
                 )
                 mock_snapshot_download.assert_called_with(
                     repo_id=model_name,
                     local_dir=f"{str(tmpdir)}/{model_name}",
                     allow_patterns=["*.json"],
-                    ignore_patterns=["test.json"]
+                    ignore_patterns=["test.json"],
                 )
                 mock_subprocess.assert_called_with(
                     shlex.split(
@@ -1119,6 +1119,61 @@ def test_import_tei_model_byoc(
         assert model.ready_to_deploy is True
         assert model.ready_to_finetune is False
 
+    @patch("ads.model.service.oci_datascience_model.OCIDataScienceModel.create")
+    @patch("ads.model.datascience_model.DataScienceModel.sync")
+    @patch("ads.model.datascience_model.DataScienceModel.upload_artifact")
+    @patch("ads.common.object_storage_details.ObjectStorageDetails.list_objects")
+    @patch.object(HfApi, "model_info")
+    @patch("ads.aqua.common.utils.load_config", return_value={})
+    def test_import_model_with_input_tags(
+        self,
+        mock_load_config,
+        mock_list_objects,
+        mock_upload_artifact,
+        mock_sync,
+        mock_ocidsc_create,
+        mock_get_hf_model_info,
+        mock_init_client,
+    ):
+        my_model = "oracle/aqua-1t-mega-model"
+        ObjectStorageDetails.is_bucket_versioned = MagicMock(return_value=True)
+
+        os_path = "oci://aqua-bkt@aqua-ns/prefix/path"
+        ds_freeform_tags = {
+            "OCI_AQUA": "active",
+        }
+        mock_list_objects.return_value = MagicMock(objects=[])
+
+        reload(ads.aqua.model.model)
+        app = AquaModelApp()
+        with patch.object(AquaModelApp, "list") as aqua_model_mock_list:
+            aqua_model_mock_list.return_value = [
+                AquaModelSummary(
+                    id="test_id1",
+                    name="organization1/name1",
+                    organization="organization1",
+                )
+            ]
+            model: AquaModel = app.register(
+                model=my_model,
+                os_path=os_path,
+                inference_container="odsc-vllm-or-tgi-container",
+                finetuning_container="odsc-llm-fine-tuning",
+                download_from_hf=False,
+                freeform_tags={"ftag1": "fvalue1", "ftag2": "fvalue2"},
+                defined_tags={"dtag1": "dvalue1", "dtag2": "dvalue2"},
+            )
+            assert model.tags == {
+                "aqua_custom_base_model": "true",
+                "model_format": "SAFETENSORS",
+                "ready_to_fine_tune": "true",
+                "dtag1": "dvalue1",
+                "dtag2": "dvalue2",
+                "ftag1": "fvalue1",
+                "ftag2": "fvalue2",
+                **ds_freeform_tags,
+            }
+
     @pytest.mark.parametrize(
         "data, expected_output",
         [
@@ -1163,6 +1218,18 @@ def test_import_tei_model_byoc(
                 },
                 "ads aqua model register --model oracle/oracle-1it --os_path oci://aqua-bkt@aqua-ns/path --download_from_hf True --inference_container odsc-tei-serving --inference_container_uri <region>.ocir.io/<your_tenancy>/<your_image>",
             ),
+            (
+                {
+                    "os_path": "oci://aqua-bkt@aqua-ns/path",
+                    "model": "oracle/oracle-1it",
+                    "inference_container": "odsc-vllm-serving",
+                    "freeform_tags": {"ftag1": "fvalue1", "ftag2": "fvalue2"},
+                    "defined_tags": {"dtag1": "dvalue1", "dtag2": "dvalue2"},
+                },
+                "ads aqua model register --model oracle/oracle-1it --os_path oci://aqua-bkt@aqua-ns/path "
+                "--download_from_hf True --inference_container odsc-vllm-serving --freeform_tags "
+                '{"ftag1": "fvalue1", "ftag2": "fvalue2"} --defined_tags {"dtag1": "dvalue1", "dtag2": "dvalue2"}',
+            ),
         ],
     )
     def test_import_cli(self, data, expected_output):
