Adding edit registered model api

Author: kumar-shivam-ranjan

ads/aqua/extension/errors.py

@@ -8,3 +8,4 @@ class Errors(str):
     NO_INPUT_DATA = "No input data provided."
     MISSING_REQUIRED_PARAMETER = "Missing required parameter: '{}'"
     MISSING_ONEOF_REQUIRED_PARAMETER = "Either '{}' or '{}' is required."
+    INVALID_VALUE_OF_PARAMETER = "Invalid value of parameter: '{}'"

ads/aqua/extension/model_handler.py

@@ -8,6 +8,7 @@
 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.extension.base_handler import AquaAPIhandler
@@ -139,6 +140,28 @@ def post(self, *args, **kwargs):
             )
         )
 
+    @handle_exceptions
+    def put(self,id):
+        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)
+
+        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))
+
+
 
 class AquaModelLicenseHandler(AquaAPIhandler):
     """Handler for Aqua Model license REST APIs."""

ads/aqua/model/model.py

@@ -10,11 +10,11 @@
 import oci
 from cachetools import TTLCache
 from huggingface_hub import snapshot_download
-from oci.data_science.models import JobRun, Model
+from oci.data_science.models import JobRun, Model, UpdateModelDetails, Metadata
 
 from ads.aqua import ODSC_MODEL_COMPARTMENT_OCID, logger
 from ads.aqua.app import AquaApp
-from ads.aqua.common.enums import InferenceContainerTypeFamily, Tags
+from ads.aqua.common.enums import InferenceContainerTypeFamily, Tags, FineTuningContainerTypeFamily
 from ads.aqua.common.errors import AquaRuntimeError, AquaValueError
 from ads.aqua.common.utils import (
     LifecycleStatus,
@@ -75,7 +75,7 @@
     TENANCY_OCID,
 )
 from ads.model import DataScienceModel
-from ads.model.model_metadata import ModelCustomMetadata, ModelCustomMetadataItem
+from ads.model.model_metadata import ModelCustomMetadata, ModelCustomMetadataItem, MetadataCustomCategory
 from ads.telemetry import telemetry
 
 
@@ -332,6 +332,72 @@ def delete_registered_model(self,model_id):
         else:
             raise AquaRuntimeError(f"Failed to delete model:{model_id}. Only registered models can be deleted.")
 
+    @telemetry(entry_point="plugin=model&action=delete", name="aqua")
+    def edit_registered_model(self,id,inference_container,enable_finetuning,task):
+        """Edits the default config of unverified registered model.
+
+        Parameters
+        ----------
+        id: str
+            The model OCID.
+        inference_container: str.
+            The inference container family name
+        enable_finetuning: str
+            Flag to enable or disable finetuning over the model. Defaults to None
+
+        Returns
+        -------
+        Model:
+            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.")
+            else:
+                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
+                                             )
+                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":
+                        try:
+                            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}")
+
+                custom_metadata_list.remove("modelDescription")
+                if 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
+                )
+                return self.ds_client.update_model(id,update_model_details).data
+        else:
+            raise AquaRuntimeError(f"Failed to edit model:{id}. Only registered unverified models can be deleted.")
+
+
     def _fetch_metric_from_metadata(
         self,
         custom_metadata_list: ModelCustomMetadata,