Addressing review comments

Author: kumar-shivam-ranjan

ads/aqua/extension/deployment_handler.py

@@ -55,29 +55,28 @@ def get(self, id=""):
             raise HTTPError(400, f"The request {self.request.path} is invalid.")
 
     @handle_exceptions
-    def delete(self,model_deployment_id):
+    def delete(self, model_deployment_id):
         return self.finish(AquaDeploymentApp().delete(model_deployment_id))
 
     @handle_exceptions
-    def put(self,*args,**kwargs):
-        try:
-            input_data = self.get_json_body()
-        except Exception as ex:
-            raise HTTPError(400, Errors.INVALID_INPUT_DATA_FORMAT) from ex
-        if not input_data:
-            raise HTTPError(400, Errors.NO_INPUT_DATA)
-
-        # required input parameters
-        model_deployment_id = input_data.get("model_deployment_id")
-        if not model_deployment_id:
-            raise HTTPError(
-                400, Errors.MISSING_REQUIRED_PARAMETER.format("model_deployment_id")
-            )
+    def put(self, *args, **kwargs):
+        """
+        Handles put request for the activating and deactivating OCI datascience model deployments
+        Raises
+        ------
+        HTTPError
+            Raises HTTPError if inputs are missing or are invalid
+        """
         url_parse = urlparse(self.request.path)
-        paths = url_parse.path.strip("/")
-        if paths.startswith("aqua/deployments/activate"):
+        paths = url_parse.path.strip("/").split("/")
+        if len(paths) != 4 or paths[0] != "aqua" or paths[1] != "deployments":
+            raise HTTPError(400, f"The request {self.request.path} is invalid.")
+
+        model_deployment_id = paths[2]
+        action = paths[3]
+        if action == "activate":
             return self.finish(AquaDeploymentApp().activate(model_deployment_id))
-        elif paths.startswith("aqua/deployments/deactivate"):
+        elif action == "deactivate":
             return self.finish(AquaDeploymentApp().deactivate(model_deployment_id))
         else:
             raise HTTPError(400, f"The request {self.request.path} is invalid.")

ads/aqua/extension/model_handler.py

@@ -8,14 +8,17 @@
 from tornado.web import HTTPError
 
 from ads.aqua.common.decorator import handle_exceptions
-from ads.aqua.common.enums import InferenceContainerTypeFamily
 from ads.aqua.common.errors import AquaRuntimeError, AquaValueError
-from ads.aqua.common.utils import get_hf_model_info, list_hf_models
+from ads.aqua.common.utils import (
+    get_container_config,
+    get_hf_model_info,
+    list_hf_models,
+)
 from ads.aqua.extension.base_handler import AquaAPIhandler
 from ads.aqua.extension.errors import Errors
 from ads.aqua.model import AquaModelApp
 from ads.aqua.model.entities import AquaModelSummary, HFModelSummary
-from ads.aqua.ui import ModelFormat
+from ads.aqua.ui import AquaContainerConfig, ModelFormat
 
 
 class AquaModelHandler(AquaAPIhandler):
@@ -75,7 +78,7 @@ def delete(self, id=""):
         if paths.startswith("aqua/model/cache"):
             return self.finish(AquaModelApp().clear_model_list_cache())
         elif id:
-            return self.finish(AquaModelApp().delete_registered_model(id))
+            return self.finish(AquaModelApp().delete_model(id))
         else:
             raise HTTPError(400, f"The request {self.request.path} is invalid.")
 
@@ -141,7 +144,7 @@ def post(self, *args, **kwargs):
         )
 
     @handle_exceptions
-    def put(self,id):
+    def put(self, id):
         try:
             input_data = self.get_json_body()
         except Exception as ex:
@@ -150,17 +153,26 @@ def put(self,id):
         if not input_data:
             raise HTTPError(400, Errors.NO_INPUT_DATA)
 
-        inference_container=input_data.get('inference_container')
-        if inference_container is not None and inference_container not in [
-            InferenceContainerTypeFamily.AQUA_TGI_CONTAINER_FAMILY,
-            InferenceContainerTypeFamily.AQUA_VLLM_CONTAINER_FAMILY,
-            InferenceContainerTypeFamily.AQUA_LLAMA_CPP_CONTAINER_FAMILY
-        ]:
-            raise HTTPError(400,Errors.INVALID_VALUE_OF_PARAMETER.format("inference_container"))
-        enable_finetuning=input_data.get('enable_finetuning')
-        task=input_data.get('task')
-        return self.finish(AquaModelApp().edit_registered_model(id,inference_container,enable_finetuning,task))
+        inference_container = input_data.get("inference_container")
+        containers = list(
+            AquaContainerConfig.from_container_index_json(
+                config=get_container_config(), enable_spec=True
+            ).inference.values()
+        )
+        family_values = [item.family for item in containers]
+
+        if inference_container is not None and inference_container not in family_values:
+            raise HTTPError(
+                400, Errors.INVALID_VALUE_OF_PARAMETER.format("inference_container")
+            )
 
+        enable_finetuning = input_data.get("enable_finetuning")
+        task = input_data.get("task")
+        return self.finish(
+            AquaModelApp().edit_registered_model(
+                id, inference_container, enable_finetuning, task
+            )
+        )
 
 
 class AquaModelLicenseHandler(AquaAPIhandler):

ads/aqua/model/model.py

@@ -10,11 +10,15 @@
 import oci
 from cachetools import TTLCache
 from huggingface_hub import snapshot_download
-from oci.data_science.models import JobRun, Model, UpdateModelDetails, Metadata
+from oci.data_science.models import JobRun, Metadata, Model, UpdateModelDetails
 
 from ads.aqua import ODSC_MODEL_COMPARTMENT_OCID, logger
 from ads.aqua.app import AquaApp
-from ads.aqua.common.enums import InferenceContainerTypeFamily, Tags, FineTuningContainerTypeFamily
+from ads.aqua.common.enums import (
+    FineTuningContainerTypeFamily,
+    InferenceContainerTypeFamily,
+    Tags,
+)
 from ads.aqua.common.errors import AquaRuntimeError, AquaValueError
 from ads.aqua.common.utils import (
     LifecycleStatus,
@@ -75,7 +79,11 @@
     TENANCY_OCID,
 )
 from ads.model import DataScienceModel
-from ads.model.model_metadata import ModelCustomMetadata, ModelCustomMetadataItem, MetadataCustomCategory
+from ads.model.model_metadata import (
+    MetadataCustomCategory,
+    ModelCustomMetadata,
+    ModelCustomMetadataItem,
+)
 from ads.telemetry import telemetry
 
 
@@ -324,16 +332,21 @@ def get(self, model_id: str, load_model_card: Optional[bool] = True) -> "AquaMod
         return model_details
 
     @telemetry(entry_point="plugin=model&action=delete", name="aqua")
-    def delete_registered_model(self,model_id):
-        ds_model=DataScienceModel.from_id(model_id)
-        is_registered_model=ds_model.freeform_tags.get(Tags.BASE_MODEL_CUSTOM,None)
-        if is_registered_model:
+    def delete_model(self, model_id):
+        ds_model = DataScienceModel.from_id(model_id)
+        is_registered_model = ds_model.freeform_tags.get(Tags.BASE_MODEL_CUSTOM, None)
+        is_fine_tuned_model = ds_model.freeform_tags.get(
+            Tags.AQUA_FINE_TUNED_MODEL_TAG, None
+        )
+        if is_registered_model or is_fine_tuned_model:
             return ds_model.delete()
         else:
-            raise AquaRuntimeError(f"Failed to delete model:{model_id}. Only registered models can be deleted.")
+            raise AquaRuntimeError(
+                f"Failed to delete model:{model_id}. Only registered models or finetuned model can be deleted."
+            )
 
     @telemetry(entry_point="plugin=model&action=delete", name="aqua")
-    def edit_registered_model(self,id,inference_container,enable_finetuning,task):
+    def edit_registered_model(self, id, inference_container, enable_finetuning, task):
         """Edits the default config of unverified registered model.
 
         Parameters
@@ -351,52 +364,61 @@ def edit_registered_model(self,id,inference_container,enable_finetuning,task):
             The instance of oci.data_science.models.Model.
 
         """
-        ds_model=DataScienceModel.from_id(id)
-        if ds_model.freeform_tags.get(Tags.BASE_MODEL_CUSTOM,None):
-            if ds_model.freeform_tags.get(Tags.AQUA_SERVICE_MODEL_TAG,None):
-                raise AquaRuntimeError(f"Failed to edit model:{id}. Only registered unverified models can be edited.")
+        ds_model = DataScienceModel.from_id(id)
+        if ds_model.freeform_tags.get(Tags.BASE_MODEL_CUSTOM, None):
+            if ds_model.freeform_tags.get(Tags.AQUA_SERVICE_MODEL_TAG, None):
+                raise AquaRuntimeError(
+                    f"Failed to edit model:{id}. Only registered unverified models can be edited."
+                )
             else:
-                custom_metadata_list=ds_model.custom_metadata_list
-                freeform_tags=ds_model.freeform_tags
+                custom_metadata_list = ds_model.custom_metadata_list
+                freeform_tags = ds_model.freeform_tags
                 if inference_container:
-                    custom_metadata_list.add(key=ModelCustomMetadataFields.DEPLOYMENT_CONTAINER,
-                                             value=inference_container,
-                                             category=MetadataCustomCategory.OTHER,
-                                             description="Deployment container mapping for SMC",
-                                             replace=True
-                                             )
+                    custom_metadata_list.add(
+                        key=ModelCustomMetadataFields.DEPLOYMENT_CONTAINER,
+                        value=inference_container,
+                        category=MetadataCustomCategory.OTHER,
+                        description="Deployment container mapping for SMC",
+                        replace=True,
+                    )
                 if enable_finetuning is not None:
-                    if enable_finetuning.lower()=="true":
-                        custom_metadata_list.add(key=ModelCustomMetadataFields.FINETUNE_CONTAINER,
-                                                 value=FineTuningContainerTypeFamily.AQUA_FINETUNING_CONTAINER_FAMILY,
-                                                 category=MetadataCustomCategory.OTHER,
-                                                 description="Fine-tuning container mapping for SMC",
-                                                 replace=True
-                                                 )
-                        freeform_tags.update({Tags.READY_TO_FINE_TUNE:"true"})
-                    elif enable_finetuning.lower()=="false":
+                    if enable_finetuning.lower() == "true":
+                        custom_metadata_list.add(
+                            key=ModelCustomMetadataFields.FINETUNE_CONTAINER,
+                            value=FineTuningContainerTypeFamily.AQUA_FINETUNING_CONTAINER_FAMILY,
+                            category=MetadataCustomCategory.OTHER,
+                            description="Fine-tuning container mapping for SMC",
+                            replace=True,
+                        )
+                        freeform_tags.update({Tags.READY_TO_FINE_TUNE: "true"})
+                    elif enable_finetuning.lower() == "false":
                         try:
-                            custom_metadata_list.remove(ModelCustomMetadataFields.FINETUNE_CONTAINER)
+                            custom_metadata_list.remove(
+                                ModelCustomMetadataFields.FINETUNE_CONTAINER
+                            )
                             freeform_tags.pop(Tags.READY_TO_FINE_TUNE)
                         except Exception as ex:
-                            raise AquaRuntimeError(f"The given model already doesn't support finetuning: {ex}")
+                            raise AquaRuntimeError(
+                                f"The given model already doesn't support finetuning: {ex}"
+                            )
 
                 custom_metadata_list.remove("modelDescription")
                 if task:
-                    freeform_tags.update({"task":task})
+                    freeform_tags.update({"task": task})
 
                 updated_custom_metadata_list = [
                     Metadata(**metadata)
                     for metadata in custom_metadata_list.to_dict()["data"]
                 ]
                 update_model_details = UpdateModelDetails(
                     custom_metadata_list=updated_custom_metadata_list,
-                    freeform_tags=freeform_tags
+                    freeform_tags=freeform_tags,
                 )
-                return self.ds_client.update_model(id,update_model_details).data
+                return AquaApp().update_model(id, update_model_details).data
         else:
-            raise AquaRuntimeError(f"Failed to edit model:{id}. Only registered unverified models can be deleted.")
-
+            raise AquaRuntimeError(
+                f"Failed to edit model:{id}. Only registered unverified models can be deleted."
+            )
 
     def _fetch_metric_from_metadata(
         self,
@@ -1010,38 +1032,39 @@ def _validate_model(
         # gguf extension exist.
         if {ModelFormat.SAFETENSORS, ModelFormat.GGUF}.issubset(set(model_formats)):
             if (
-                import_model_details.inference_container.lower() == InferenceContainerTypeFamily.AQUA_LLAMA_CPP_CONTAINER_FAMILY
+                import_model_details.inference_container.lower()
+                == InferenceContainerTypeFamily.AQUA_LLAMA_CPP_CONTAINER_FAMILY
             ):
                 self._validate_gguf_format(
                     import_model_details=import_model_details,
                     verified_model=verified_model,
                     gguf_model_files=gguf_model_files,
                     validation_result=validation_result,
-                    model_name=model_name
+                    model_name=model_name,
                 )
             else:
                 self._validate_safetensor_format(
                     import_model_details=import_model_details,
                     verified_model=verified_model,
                     validation_result=validation_result,
                     hf_download_config_present=hf_download_config_present,
-                    model_name=model_name
+                    model_name=model_name,
                 )
         elif ModelFormat.SAFETENSORS in model_formats:
             self._validate_safetensor_format(
                 import_model_details=import_model_details,
                 verified_model=verified_model,
                 validation_result=validation_result,
                 hf_download_config_present=hf_download_config_present,
-                model_name=model_name
+                model_name=model_name,
             )
         elif ModelFormat.GGUF in model_formats:
             self._validate_gguf_format(
                 import_model_details=import_model_details,
                 verified_model=verified_model,
                 gguf_model_files=gguf_model_files,
                 validation_result=validation_result,
-                model_name=model_name
+                model_name=model_name,
             )
 
         return validation_result
@@ -1052,7 +1075,7 @@ def _validate_safetensor_format(
         verified_model: DataScienceModel = None,
         validation_result: ModelValidationResult = None,
         hf_download_config_present: bool = None,
-        model_name: str = None
+        model_name: str = None,
     ):
         if import_model_details.download_from_hf:
             # validates config.json exists for safetensors model from hugginface
@@ -1079,20 +1102,13 @@ def _validate_safetensor_format(
                 ) from ex
             else:
                 try:
-                    metadata_model_type = (
-                        verified_model.custom_metadata_list.get(
-                            AQUA_MODEL_ARTIFACT_CONFIG_MODEL_TYPE
-                        ).value
-                    )
+                    metadata_model_type = verified_model.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 AQUA_MODEL_ARTIFACT_CONFIG_MODEL_TYPE in model_config:
                             if (
-                                model_config[
-                                    AQUA_MODEL_ARTIFACT_CONFIG_MODEL_TYPE
-                                ]
+                                model_config[AQUA_MODEL_ARTIFACT_CONFIG_MODEL_TYPE]
                                 != metadata_model_type
                             ):
                                 raise AquaRuntimeError(
@@ -1110,9 +1126,7 @@ def _validate_safetensor_format(
                 except Exception:
                     pass
                 if verified_model:
-                    validation_result.telemetry_model_name = (
-                        verified_model.display_name
-                    )
+                    validation_result.telemetry_model_name = verified_model.display_name
                 elif (
                     model_config is not None
                     and AQUA_MODEL_ARTIFACT_CONFIG_MODEL_NAME in model_config
@@ -1124,9 +1138,7 @@ def _validate_safetensor_format(
                 ):
                     validation_result.telemetry_model_name = f"{AQUA_MODEL_TYPE_CUSTOM}_{model_config[AQUA_MODEL_ARTIFACT_CONFIG_MODEL_TYPE]}"
                 else:
-                    validation_result.telemetry_model_name = (
-                        AQUA_MODEL_TYPE_CUSTOM
-                    )
+                    validation_result.telemetry_model_name = AQUA_MODEL_TYPE_CUSTOM
 
     @staticmethod
     def _validate_gguf_format(