Add config validation and telemetry (#872)

Author: VipulMascarenhas

ads/aqua/constants.py

@@ -32,6 +32,9 @@
 AQUA_GA_LIST = ["id19sfcrra6z"]
 AQUA_MODEL_TYPE_SERVICE = "service"
 AQUA_MODEL_TYPE_CUSTOM = "custom"
+AQUA_MODEL_ARTIFACT_CONFIG = "config.json"
+AQUA_MODEL_ARTIFACT_CONFIG_MODEL_NAME = "_name_or_path"
+AQUA_MODEL_ARTIFACT_CONFIG_MODEL_TYPE = "model_type"
 
 TRAINING_METRICS_FINAL = "training_metrics_final"
 VALIDATION_METRICS_FINAL = "validation_metrics_final"

ads/aqua/model/model.py

@@ -19,6 +19,7 @@
     get_artifact_path,
     read_file,
     copy_model_config,
+    load_config,
 )
 from ads.aqua.constants import (
     LICENSE_TXT,
@@ -32,12 +33,16 @@
     UNKNOWN,
     VALIDATION_METRICS,
     VALIDATION_METRICS_FINAL,
+    AQUA_MODEL_ARTIFACT_CONFIG,
+    AQUA_MODEL_ARTIFACT_CONFIG_MODEL_NAME,
+    AQUA_MODEL_ARTIFACT_CONFIG_MODEL_TYPE,
+    AQUA_MODEL_TYPE_CUSTOM,
 )
 from ads.aqua.model.constants import *
 from ads.aqua.model.entities import *
 from ads.common.auth import default_signer
 from ads.common.oci_resource import SEARCH_TYPE, OCIResource
-from ads.common.utils import get_console_link, is_path_exists
+from ads.common.utils import get_console_link
 from ads.config import (
     AQUA_DEPLOYMENT_CONTAINER_METADATA_NAME,
     AQUA_EVALUATION_CONTAINER_METADATA_NAME,
@@ -688,9 +693,16 @@ def register(
         if not import_model_details:
             import_model_details = ImportModelDetails(**kwargs)
 
-        if not is_path_exists(
-            f"{import_model_details.os_path.rstrip('/')}/config.json"
-        ):
+        try:
+            model_config = load_config(
+                file_path=import_model_details.os_path,
+                config_file_name=AQUA_MODEL_ARTIFACT_CONFIG,
+            )
+        except Exception as ex:
+            logger.error(
+                f"Exception occurred while loading config file from {import_model_details.os_path}"
+                f"Exception message: {ex}"
+            )
             raise AquaRuntimeError(
                 f"The model path {import_model_details.os_path} does not contain the file config.json. "
                 f"Please check if the path is correct or the model artifacts are available at this location."
@@ -713,6 +725,30 @@ def register(
             )
         if model_service_id:
             verified_model_details = DataScienceModel.from_id(model_service_id)
+            try:
+                metadata_model_type = verified_model_details.custom_metadata_list.get(
+                    AQUA_MODEL_ARTIFACT_CONFIG_MODEL_TYPE
+                ).value
+                if metadata_model_type:
+                    if AQUA_MODEL_ARTIFACT_CONFIG_MODEL_TYPE in model_config:
+                        if (
+                            model_config[AQUA_MODEL_ARTIFACT_CONFIG_MODEL_TYPE]
+                            != metadata_model_type
+                        ):
+                            raise AquaRuntimeError(
+                                f"The {AQUA_MODEL_ARTIFACT_CONFIG_MODEL_TYPE} attribute in {AQUA_MODEL_ARTIFACT_CONFIG}"
+                                f" at {import_model_details.os_path} is invalid, expected {metadata_model_type} for "
+                                f"the model {import_model_details.model}. Please check if the path is correct or "
+                                f"the correct model artifacts are available at this location."
+                                f""
+                            )
+                    else:
+                        logger.debug(
+                            f"Could not find {AQUA_MODEL_ARTIFACT_CONFIG_MODEL_TYPE} attribute in "
+                            f"{AQUA_MODEL_ARTIFACT_CONFIG}. Proceeding with model registration."
+                        )
+            except:
+                pass
 
         # Copy the model name from the service model if `model` is ocid
         model_name = (
@@ -734,19 +770,14 @@ def register(
         # registered model will always have inference and evaluation container, but
         # fine-tuning container may be not set
         inference_container = ds_model.custom_metadata_list.get(
-            ModelCustomMetadataFields.DEPLOYMENT_CONTAINER,
-            ModelCustomMetadataItem(key=ModelCustomMetadataFields.DEPLOYMENT_CONTAINER),
+            ModelCustomMetadataFields.DEPLOYMENT_CONTAINER
         ).value
         evaluation_container = ds_model.custom_metadata_list.get(
             ModelCustomMetadataFields.EVALUATION_CONTAINER,
-            ModelCustomMetadataItem(key=ModelCustomMetadataFields.EVALUATION_CONTAINER),
         ).value
         try:
             finetuning_container = ds_model.custom_metadata_list.get(
                 ModelCustomMetadataFields.FINETUNE_CONTAINER,
-                ModelCustomMetadataItem(
-                    key=ModelCustomMetadataFields.FINETUNE_CONTAINER
-                ),
             ).value
         except:
             finetuning_container = None
@@ -756,18 +787,31 @@ def register(
             project_id=ds_model.project_id,
             model_card=str(
                 read_file(
-                    file_path=(
-                        f"{import_model_details.os_path.rstrip('/')}/config/{README}"
-                        if verified_model_details
-                        else f"{import_model_details.os_path.rstrip('/')}/{README}"
-                    ),
+                    file_path=f"{import_model_details.os_path.rstrip('/')}/{README}",
                     auth=default_signer(),
                 )
             ),
             inference_container=inference_container,
             finetuning_container=finetuning_container,
             evaluation_container=evaluation_container,
         )
+
+        if verified_model_details:
+            telemetry_model_name = model_name
+        else:
+            if AQUA_MODEL_ARTIFACT_CONFIG_MODEL_NAME in model_config:
+                telemetry_model_name = f"{AQUA_MODEL_TYPE_CUSTOM}_{model_config[AQUA_MODEL_ARTIFACT_CONFIG_MODEL_NAME]}"
+            elif AQUA_MODEL_ARTIFACT_CONFIG_MODEL_TYPE in model_config:
+                telemetry_model_name = f"{AQUA_MODEL_TYPE_CUSTOM}_{model_config[AQUA_MODEL_ARTIFACT_CONFIG_MODEL_TYPE]}"
+            else:
+                telemetry_model_name = AQUA_MODEL_TYPE_CUSTOM
+
+        self.telemetry.record_event_async(
+            category="aqua/model",
+            action="register",
+            detail=telemetry_model_name,
+        )
+
         return AquaModel(**aqua_model_attributes)
 
     def _if_show(self, model: DataScienceModel) -> bool:

tests/unitary/with_extras/aqua/test_model.py

@@ -26,7 +26,7 @@
     ModelProvenanceMetadata,
     ModelTaxonomyMetadata,
 )
-from ads.aqua.common.errors import AquaRuntimeError
+from ads.aqua.common.errors import AquaRuntimeError, AquaFileNotFoundError
 from ads.model.service.oci_datascience_model import OCIDataScienceModel
 
 
@@ -534,10 +534,10 @@ def test_get_model_fine_tuned(
     )
     @patch("ads.aqua.common.utils.copy_file")
     @patch("ads.common.object_storage_details.ObjectStorageDetails.list_objects")
-    @patch("ads.common.utils.is_path_exists", return_value=True)
+    @patch("ads.aqua.common.utils.load_config", return_value={})
     def test_import_verified_model(
         self,
-        mock_is_path_exists,
+        mock_load_config,
         mock_list_objects,
         mock_copy_file,
         artifact_location_set,
@@ -612,7 +612,7 @@ def test_import_verified_model(
             "aqua_service_model": "test_model_id",
             **ds_freeform_tags,
         }
-        mock_is_path_exists.assert_called()
+        mock_load_config.assert_called()
 
         assert model.inference_container == "odsc-tgi-serving"
         assert model.finetuning_container is None
@@ -621,8 +621,8 @@ def test_import_verified_model(
         assert model.ready_to_deploy is True
         assert model.ready_to_finetune is False
 
-    @patch("ads.common.utils.is_path_exists", return_value=True)
-    def test_import_any_model_no_containers_specified(self, mock_is_path_exists):
+    @patch("ads.aqua.common.utils.load_config", return_value={})
+    def test_import_any_model_no_containers_specified(self, mock_load_config):
         ObjectStorageDetails.is_bucket_versioned = MagicMock(return_value=True)
         ads.common.oci_datascience.OCIDataScienceMixin.init_client = MagicMock()
         DataScienceModel.upload_artifact = MagicMock()
@@ -655,8 +655,8 @@ def test_import_any_model_no_containers_specified(self, mock_is_path_exists):
                     os_path=os_path,
                 )
 
-    @patch("ads.common.utils.is_path_exists", return_value=True)
-    def test_import_model_with_project_compartment_override(self, mock_is_path_exists):
+    @patch("ads.aqua.common.utils.load_config", return_value={})
+    def test_import_model_with_project_compartment_override(self, mock_load_config):
         ObjectStorageDetails.is_bucket_versioned = MagicMock(return_value=True)
         ads.common.oci_datascience.OCIDataScienceMixin.init_client = MagicMock()
         DataScienceModel.upload_artifact = MagicMock()
@@ -704,8 +704,8 @@ def test_import_model_with_project_compartment_override(self, mock_is_path_exist
         assert model.compartment_id == compartment_override
         assert model.project_id == project_override
 
-    @patch("ads.common.utils.is_path_exists", return_value=False)
-    def test_import_model_with_missing_artifact(self, mock_is_path_exists):
+    @patch("ads.aqua.common.utils.load_config", side_effect=AquaFileNotFoundError)
+    def test_import_model_with_missing_config(self, mock_load_config):
         """Test for validating if error is returned when model artifacts are incomplete or not available."""
         os_path = "oci://aqua-bkt@aqua-ns/prefix/path"
         model_name = "oracle/aqua-1t-mega-model"
@@ -717,10 +717,10 @@ def test_import_model_with_missing_artifact(self, mock_is_path_exists):
                 os_path=os_path,
             )
 
-    @patch("ads.common.utils.is_path_exists", return_value=True)
+    @patch("ads.aqua.common.utils.load_config", return_value={})
     def test_import_any_model_smc_container(
         self,
-        mock_is_path_exists,
+        mock_load_config,
     ):
         my_model = "oracle/aqua-1t-mega-model"
         ObjectStorageDetails.is_bucket_versioned = MagicMock(return_value=True)