Add source model tag to imported models (#853)

Author: VipulMascarenhas

ads/aqua/extension/model_handler.py

@@ -179,7 +179,10 @@ def post(self, *args, **kwargs):
             )
 
         # Check pipeline_tag, it should be `text-generation`
-        if hf_model_info.pipeline_tag.lower() != ModelTask.TEXT_GENERATION:
+        if (
+            not hf_model_info.pipeline_tag
+            or hf_model_info.pipeline_tag.lower() != ModelTask.TEXT_GENERATION
+        ):
             raise AquaRuntimeError(
                 f"Unsupported pipeline tag for the chosen model: '{hf_model_info.pipeline_tag}'. "
                 f"AQUA currently supports the following tasks only: {', '.join(ModelTask.values())}. "

ads/aqua/model/model.py

@@ -592,10 +592,12 @@ def _create_model_catalog_entry(
         except Exception:
             logger.exception(f"Could not fetch model information for {model_name}")
         tags = (
-            {**shadow_model.freeform_tags, Tags.BASE_MODEL_CUSTOM: "true"}
+            {**shadow_model.freeform_tags, Tags.AQUA_SERVICE_MODEL_TAG: shadow_model.id}
             if shadow_model
             else {Tags.AQUA_TAG: "active", Tags.BASE_MODEL_CUSTOM: "true"}
         )
+        tags.update({Tags.BASE_MODEL_CUSTOM: "true"})
+
         # Remove `ready_to_import` tag that might get copied from service model.
         tags.pop(Tags.READY_TO_IMPORT, None)
         metadata = None
@@ -691,7 +693,7 @@ def _create_model_catalog_entry(
 
     def register(
         self, import_model_details: ImportModelDetails = None, **kwargs
-    ) -> str:
+    ) -> DataScienceModel:
         """Loads the model from huggingface and registers as Model in Data Science Model catalog
         Note: For the models that require user token, use `huggingface-cli login` to setup the token
         The inference container and finetuning container could be of type Service Manged Container(SMC) or custom. If it is custom, full container URI is expected. If it of type SMC, only the container family name is expected.

tests/unitary/with_extras/aqua/test_model.py

@@ -249,7 +249,8 @@ def test_create_model(self, mock_from_id, mock_validate, mock_create):
     @patch("ads.aqua.model.model.read_file")
     @patch.object(DataScienceModel, "from_id")
     @patch(
-        "ads.aqua.model.model.get_artifact_path", return_value="oci://bucket@namespace/prefix"
+        "ads.aqua.model.model.get_artifact_path",
+        return_value="oci://bucket@namespace/prefix",
     )
     def test_get_foundation_models(
         self,
@@ -356,7 +357,8 @@ def test_get_foundation_models(
     @patch("ads.aqua.model.model.read_file")
     @patch.object(DataScienceModel, "from_id")
     @patch(
-        "ads.aqua.model.model.get_artifact_path", return_value="oci://bucket@namespace/prefix"
+        "ads.aqua.model.model.get_artifact_path",
+        return_value="oci://bucket@namespace/prefix",
     )
     def test_get_model_fine_tuned(
         self, mock_get_artifact_path, mock_from_id, mock_read_file, mock_query_resource
@@ -563,6 +565,7 @@ def test_import_shadow_model(
         )
         ds_model.with_custom_metadata_list(custom_metadata_list)
         ds_model.set_spec(ds_model.CONST_MODEL_FILE_DESCRIPTION, {})
+        ds_model.dsc_model = MagicMock(id="test_model_id")
         DataScienceModel.from_id = MagicMock(return_value=ds_model)
         reload(ads.aqua.model.model)
         app = AquaModelApp()
@@ -587,6 +590,7 @@ def test_import_shadow_model(
             )  # The imported model should not have this tag
             assert model.freeform_tags == {
                 "aqua_custom_base_model": "true",
+                "aqua_service_model": "test_model_id",
                 **ds_freeform_tags,
             }
             expected_metadata = [