[AQUA] Group Model Deployment & Stacked Model Deployment (#1217)

Co-authored-by: Lu Peng <bolu.peng@oracle.com>
Co-authored-by: Aryan Gosaliya <aryan.gosaliya@oracle.com>

Author: 3 people

ads/aqua/app.py

@@ -290,6 +290,7 @@ def create_model_catalog(
         model_taxonomy_metadata: Union[ModelTaxonomyMetadata, Dict],
         compartment_id: str,
         project_id: str,
+        freeform_tags: Dict = None,
         defined_tags: Dict = None,
         **kwargs,
     ) -> DataScienceModel:
@@ -303,6 +304,7 @@ def create_model_catalog(
             .with_custom_metadata_list(model_custom_metadata)
             .with_defined_metadata_list(model_taxonomy_metadata)
             .with_provenance_metadata(ModelProvenanceMetadata(training_id=UNKNOWN))
+            .with_freeform_tags(**(freeform_tags or {}))
             .with_defined_tags(
                 **(defined_tags or {})
             )  # Create defined tags when a model is created.

ads/aqua/common/enums.py

@@ -44,6 +44,8 @@ class Tags(ExtendedEnum):
     MODEL_FORMAT = "model_format"
     MODEL_ARTIFACT_FILE = "model_file"
     MULTIMODEL_TYPE_TAG = "aqua_multimodel"
+    STACKED_MODEL_TYPE_TAG = "aqua_stacked_model"
+    AQUA_FINE_TUNE_MODEL_VERSION = "fine_tune_model_version"
 
 
 class InferenceContainerType(ExtendedEnum):

ads/aqua/common/utils.py

@@ -643,14 +643,18 @@ def get_resource_name(ocid: str) -> str:
     return name
 
 
-def get_model_by_reference_paths(model_file_description: dict):
+def get_model_by_reference_paths(
+    model_file_description: dict, is_ft_model_v2: bool = False
+):
     """Reads the model file description json dict and returns the base model path and fine-tuned path for
         models created by reference.
 
     Parameters
     ----------
     model_file_description: dict
         json dict containing model paths and objects for models created by reference.
+    is_ft_model_v2: bool
+        Flag to indicate if it's fine tuned model v2. Defaults to False.
 
     Returns
     -------
@@ -666,8 +670,18 @@ def get_model_by_reference_paths(model_file_description: dict):
             "Please check if the model created by reference has the correct artifact."
         )
 
+    if is_ft_model_v2:
+        # model_file_description json for fine tuned model v2 contains only fine tuned model artifacts
+        # so first model is always the fine tuned model
+        ft_model_artifact = models[0]
+        fine_tune_output_path = f"oci://{ft_model_artifact['bucketName']}@{ft_model_artifact['namespace']}/{ft_model_artifact['prefix']}".rstrip(
+            "/"
+        )
+
+        return UNKNOWN, fine_tune_output_path
+
     if len(models) > 0:
-        # since the model_file_description json does not have a flag to identify the base model, we consider
+        # since the model_file_description json for legacy fine tuned model does not have a flag to identify the base model, we consider
         # the first instance to be the base model.
         base_model_artifact = models[0]
         base_model_path = f"oci://{base_model_artifact['bucketName']}@{base_model_artifact['namespace']}/{base_model_artifact['prefix']}".rstrip(

ads/aqua/constants.py

@@ -45,6 +45,8 @@
 AQUA_TROUBLESHOOTING_LINK = "https://github.com/oracle-samples/oci-data-science-ai-samples/blob/main/ai-quick-actions/troubleshooting-tips.md"
 MODEL_FILE_DESCRIPTION_VERSION = "1.0"
 MODEL_FILE_DESCRIPTION_TYPE = "modelOSSReferenceDescription"
+AQUA_FINE_TUNE_MODEL_VERSION = "v2"
+INCLUDE_BASE_MODEL = 1
 
 TRAINING_METRICS_FINAL = "training_metrics_final"
 VALIDATION_METRICS_FINAL = "validation_metrics_final"

ads/aqua/evaluation/evaluation.py

@@ -99,6 +99,7 @@
 from ads.jobs.builders.runtimes.base import Runtime
 from ads.jobs.builders.runtimes.container_runtime import ContainerRuntime
 from ads.model.datascience_model import DataScienceModel
+from ads.model.datascience_model_group import DataScienceModelGroup
 from ads.model.deployment import ModelDeploymentContainerRuntime
 from ads.model.deployment.model_deployment import ModelDeployment
 from ads.model.generic_model import ModelDeploymentRuntimeType
@@ -254,7 +255,11 @@ def create(
                         f"Make sure the {Tags.AQUA_MODEL_ID_TAG} tag is added to the deployment."
                     )
 
-                aqua_model = DataScienceModel.from_id(multi_model_id)
+                aqua_model = (
+                    DataScienceModelGroup.from_id(multi_model_id)
+                    if "datasciencemodelgroup" in multi_model_id
+                    else DataScienceModel.from_id(multi_model_id)
+                )
                 AquaEvaluationApp.validate_model_name(
                     aqua_model, create_aqua_evaluation_details
                 )
@@ -630,23 +635,23 @@ def create(
 
     @staticmethod
     def validate_model_name(
-        evaluation_source: DataScienceModel,
+        evaluation_source: Union[DataScienceModel, DataScienceModelGroup],
         create_aqua_evaluation_details: CreateAquaEvaluationDetails,
     ) -> None:
         """
         Validates the user input for the model name when creating an Aqua evaluation.
 
         This function verifies that:
         - The model group is not empty.
-        - The model multi metadata is present in the DataScienceModel metadata.
+        - The model multi metadata is present in the DataScienceModel or DataScienceModelGroup metadata.
         - The user provided a non-empty model name.
-        - The provided model name exists in the DataScienceModel metadata.
+        - The provided model name exists in the DataScienceModel or DataScienceModelGroup metadata.
         - The deployment configuration contains core metadata required for validation.
 
         Parameters
         ----------
-        evaluation_source : DataScienceModel
-            The DataScienceModel object containing metadata about each model in the deployment.
+        evaluation_source : Union[DataScienceModel, DataScienceModelGroup]
+            The DataScienceModel or DataScienceModelGroup object containing metadata about each model in the deployment.
         create_aqua_evaluation_details : CreateAquaEvaluationDetails
             Contains required and optional fields for creating the Aqua evaluation.
 
@@ -711,27 +716,30 @@ def validate_model_name(
             logger.debug(error_message)
             raise AquaRuntimeError(error_message)
 
-        try:
-            multi_model_metadata = json.loads(
-                evaluation_source.dsc_model.get_custom_metadata_artifact(
-                    metadata_key_name=ModelCustomMetadataFields.MULTIMODEL_METADATA
-                ).decode("utf-8")
-            )
-        except Exception as ex:
-            error_message = (
-                f"Error fetching {ModelCustomMetadataFields.MULTIMODEL_METADATA} "
-                f"from custom metadata for evaluation source ID '{evaluation_source.id}'. "
-                f"Details: {ex}"
-            )
-            logger.error(error_message)
-            raise AquaRuntimeError(error_message) from ex
+        if isinstance(evaluation_source, DataScienceModel):
+            try:
+                multi_model_metadata = json.loads(
+                    evaluation_source.dsc_model.get_custom_metadata_artifact(
+                        metadata_key_name=ModelCustomMetadataFields.MULTIMODEL_METADATA
+                    ).decode("utf-8")
+                )
+            except Exception as ex:
+                error_message = (
+                    f"Error fetching {ModelCustomMetadataFields.MULTIMODEL_METADATA} "
+                    f"from custom metadata for evaluation source ID '{evaluation_source.id}'. "
+                    f"Details: {ex}"
+                )
+                logger.error(error_message)
+                raise AquaRuntimeError(error_message) from ex
 
         # Build the list of valid model names from custom metadata.
         model_names = []
         for metadata in multi_model_metadata:
             model = AquaMultiModelRef(**metadata)
             model_names.append(model.model_name)
-            model_names.extend(ft.model_name for ft in (model.fine_tune_weights or []) if ft.model_name)
+            model_names.extend(
+                ft.model_name for ft in (model.fine_tune_weights or []) if ft.model_name
+            )
 
         # Check if the provided model name is among the valid names.
         if user_model_name not in model_names:

ads/aqua/extension/deployment_handler.py

@@ -119,7 +119,15 @@ def post(self, *args, **kwargs):  # noqa: ARG002
         if not input_data:
             raise HTTPError(400, Errors.NO_INPUT_DATA)
 
-        self.finish(AquaDeploymentApp().create(**input_data))
+        model_deployment_id = input_data.pop("model_deployment_id", None)
+        if model_deployment_id:
+            self.finish(
+                AquaDeploymentApp().update(
+                    model_deployment_id=model_deployment_id, **input_data
+                )
+            )
+        else:
+            self.finish(AquaDeploymentApp().create(**input_data))
 
     def read(self, id):
         """Read the information of an Aqua model deployment."""
@@ -436,7 +444,14 @@ def get(self, model_deployment_id):
             list_model_result = aqua_client.fetch_data()
             return self.finish(list_model_result)
         except Exception as ex:
-            raise HTTPError(500, str(ex))
+            error_type = type(ex).__name__
+            error_message = (
+                f"Error fetching data from endpoint '{endpoint}' [{error_type}]: {ex}"
+            )
+            logger.error(
+                error_message, exc_info=True
+            )  # Log with stack trace for diagnostics
+            raise HTTPError(500, error_message) from ex
 
 
 __handlers__ = [

ads/aqua/finetuning/constants.py

@@ -14,6 +14,7 @@ class FineTuneCustomMetadata(ExtendedEnum):
     SERVICE_MODEL_ARTIFACT_LOCATION = "artifact_location"
     SERVICE_MODEL_DEPLOYMENT_CONTAINER = "deployment-container"
     SERVICE_MODEL_FINE_TUNE_CONTAINER = "finetune-container"
+    FINE_TUNE_INCLUDE_BASE_MODEL_ARTIFACT = "include_base_model_artifact"
 
 
 class FineTuningRestrictedParams(ExtendedEnum):

ads/aqua/finetuning/finetuning.py

@@ -25,6 +25,7 @@
     upload_local_to_os,
 )
 from ads.aqua.constants import (
+    AQUA_FINE_TUNE_MODEL_VERSION,
     DEFAULT_FT_BATCH_SIZE,
     DEFAULT_FT_BLOCK_STORAGE_SIZE,
     DEFAULT_FT_REPLICA,
@@ -304,6 +305,11 @@ def create(
             "val_set_size": create_fine_tuning_details.validation_set_size,
             "training_data": ft_dataset_path,
         }
+        # needs to add 'fine_tune_model_version' tag when creating the ft model for the
+        # ft container to block merging base model artifact with ft model artifact.
+        ft_model_freeform_tags = {
+            Tags.AQUA_FINE_TUNE_MODEL_VERSION: AQUA_FINE_TUNE_MODEL_VERSION
+        }
 
         ft_model = self.create_model_catalog(
             display_name=create_fine_tuning_details.ft_name,
@@ -314,6 +320,7 @@ def create(
             compartment_id=target_compartment,
             project_id=target_project,
             model_by_reference=True,
+            freeform_tags=ft_model_freeform_tags,
             defined_tags=create_fine_tuning_details.defined_tags,
         )
         defined_metadata_dict = {}
@@ -446,6 +453,7 @@ def create(
 
         model_freeform_tags = {
             **model_freeform_tags,
+            **(ft_model.freeform_tags or {}),
             Tags.READY_TO_FINE_TUNE: "false",
             Tags.AQUA_TAG: UNKNOWN,
             Tags.AQUA_FINE_TUNED_MODEL_TAG: f"{source.id}#{source.display_name}",