[AQUA] Added support for deploy stack model. (#1223)

Author: lu-ohai

ads/aqua/model/model.py

@@ -141,19 +141,19 @@ class AquaModelApp(AquaApp):
     @telemetry(entry_point="plugin=model&action=create", name="aqua")
     def create(
         self,
-        model_id: Union[str, AquaMultiModelRef],
+        model: Union[str, AquaMultiModelRef],
         project_id: Optional[str] = None,
         compartment_id: Optional[str] = None,
         freeform_tags: Optional[Dict] = None,
         defined_tags: Optional[Dict] = None,
         **kwargs,
-    ) -> DataScienceModel:
+    ) -> Union[DataScienceModel, DataScienceModelGroup]:
         """
-        Creates a custom Aqua model from a service model.
+        Creates a custom Aqua model or model group from a service model.
 
         Parameters
         ----------
-        model_id : Union[str, AquaMultiModelRef]
+        model : Union[str, AquaMultiModelRef]
             The model ID as a string or a AquaMultiModelRef instance to be deployed.
         project_id : Optional[str]
             The project ID for the custom model.
@@ -167,28 +167,18 @@ def create(
 
         Returns
         -------
-        DataScienceModel
-            The instance of DataScienceModel.
+        Union[DataScienceModel, DataScienceModelGroup]
+            The instance of DataScienceModel or DataScienceModelGroup.
         """
-        model_id = (
-            model_id.model_id if isinstance(model_id, AquaMultiModelRef) else model_id
-        )
-        service_model = DataScienceModel.from_id(model_id)
+        fine_tune_weights = []
+        if isinstance(model, AquaMultiModelRef):
+            fine_tune_weights = model.fine_tune_weights
+            model = model.model_id
+
+        service_model = DataScienceModel.from_id(model)
         target_project = project_id or PROJECT_OCID
         target_compartment = compartment_id or COMPARTMENT_OCID
 
-        # Skip model copying if it is registered model or fine-tuned model
-        if (
-            service_model.freeform_tags.get(Tags.BASE_MODEL_CUSTOM, None) is not None
-            or service_model.freeform_tags.get(Tags.AQUA_FINE_TUNED_MODEL_TAG)
-            is not None
-        ):
-            logger.info(
-                f"Aqua Model {model_id} already exists in the user's compartment."
-                "Skipped copying."
-            )
-            return service_model
-
         # combine tags
         combined_freeform_tags = {
             **(service_model.freeform_tags or {}),
@@ -199,29 +189,112 @@ def create(
             **(defined_tags or {}),
         }
 
+        custom_model = None
+        if fine_tune_weights:
+            custom_model = self._create_model_group(
+                model_id=model,
+                compartment_id=target_compartment,
+                project_id=target_project,
+                freeform_tags=combined_freeform_tags,
+                defined_tags=combined_defined_tags,
+                fine_tune_weights=fine_tune_weights,
+                service_model=service_model,
+            )
+
+            logger.info(
+                f"Aqua Model Group {custom_model.id} created with the service model {model}."
+            )
+        else:
+            # Skip model copying if it is registered model or fine-tuned model
+            if (
+                Tags.BASE_MODEL_CUSTOM in service_model.freeform_tags
+                or Tags.AQUA_FINE_TUNED_MODEL_TAG in service_model.freeform_tags
+            ):
+                logger.info(
+                    f"Aqua Model {model} already exists in the user's compartment."
+                    "Skipped copying."
+                )
+                return service_model
+
+            custom_model = self._create_model(
+                compartment_id=target_compartment,
+                project_id=target_project,
+                freeform_tags=combined_freeform_tags,
+                defined_tags=combined_defined_tags,
+                service_model=service_model,
+                **kwargs,
+            )
+            logger.info(
+                f"Aqua Model {custom_model.id} created with the service model {model}."
+            )
+
+        # Track unique models that were created in the user's compartment
+        self.telemetry.record_event_async(
+            category="aqua/service/model",
+            action="create",
+            detail=service_model.display_name,
+        )
+
+        return custom_model
+
+    def _create_model(
+        self,
+        compartment_id: str,
+        project_id: str,
+        freeform_tags: Dict,
+        defined_tags: Dict,
+        service_model: DataScienceModel,
+        **kwargs,
+    ):
+        """Creates a data science model by reference."""
         custom_model = (
             DataScienceModel()
-            .with_compartment_id(target_compartment)
-            .with_project_id(target_project)
+            .with_compartment_id(compartment_id)
+            .with_project_id(project_id)
             .with_model_file_description(json_dict=service_model.model_file_description)
             .with_display_name(service_model.display_name)
             .with_description(service_model.description)
-            .with_freeform_tags(**combined_freeform_tags)
-            .with_defined_tags(**combined_defined_tags)
+            .with_freeform_tags(**freeform_tags)
+            .with_defined_tags(**defined_tags)
             .with_custom_metadata_list(service_model.custom_metadata_list)
             .with_defined_metadata_list(service_model.defined_metadata_list)
             .with_provenance_metadata(service_model.provenance_metadata)
             .create(model_by_reference=True, **kwargs)
         )
-        logger.info(
-            f"Aqua Model {custom_model.id} created with the service model {model_id}."
-        )
 
-        # Track unique models that were created in the user's compartment
-        self.telemetry.record_event_async(
-            category="aqua/service/model",
-            action="create",
-            detail=service_model.display_name,
+        return custom_model
+
+    def _create_model_group(
+        self,
+        model_id: str,
+        compartment_id: str,
+        project_id: str,
+        freeform_tags: Dict,
+        defined_tags: Dict,
+        fine_tune_weights: List,
+        service_model: DataScienceModel,
+    ):
+        """Creates a data science model group."""
+        custom_model = (
+            DataScienceModelGroup()
+            .with_compartment_id(compartment_id)
+            .with_project_id(project_id)
+            .with_display_name(service_model.display_name)
+            .with_description(service_model.description)
+            .with_freeform_tags(**freeform_tags)
+            .with_defined_tags(**defined_tags)
+            .with_custom_metadata_list(service_model.custom_metadata_list)
+            .with_base_model_id(model_id)
+            .with_member_models(
+                [
+                    {
+                        "inference_key": fine_tune_weight.model_name,
+                        "model_id": fine_tune_weight.model_id,
+                    }
+                    for fine_tune_weight in fine_tune_weights
+                ]
+            )
+            .create()
         )
 
         return custom_model
@@ -271,6 +344,16 @@ def create_multi(
         DataScienceModelGroup
             Instance of DataScienceModelGroup object.
         """
+        member_model_ids = [{"model_id": model.model_id} for model in models]
+        for model in models:
+            if model.fine_tune_weights:
+                member_model_ids.extend(
+                    [
+                        {"model_id": fine_tune_model.model_id}
+                        for fine_tune_model in model.fine_tune_weights
+                    ]
+                )
+
         custom_model_group = (
             DataScienceModelGroup()
             .with_compartment_id(compartment_id)
@@ -281,7 +364,7 @@ def create_multi(
             .with_defined_tags(**(defined_tags or {}))
             .with_custom_metadata_list(model_custom_metadata)
             # TODO: add member model inference key
-            .with_member_models([{"model_id": model.model_id for model in models}])
+            .with_member_models(member_model_ids)
         )
         custom_model_group.create()
 

ads/aqua/modeldeployment/constants.py

@@ -9,5 +9,12 @@
 This module contains constants used in Aqua Model Deployment.
 """
 
+from ads.common.extended_enum import ExtendedEnum
+
 DEFAULT_WAIT_TIME = 12000
 DEFAULT_POLL_INTERVAL = 10
+
+
+class DeploymentType(ExtendedEnum):
+    STACKED = "STACKED"
+    MULTI = "MULTI"

ads/aqua/modeldeployment/deployment.py

@@ -24,11 +24,13 @@
 from ads.aqua.common.errors import AquaRuntimeError, AquaValueError
 from ads.aqua.common.utils import (
     DEFINED_METADATA_TO_FILE_MAP,
+    build_params_string,
     build_pydantic_error_message,
     find_restricted_params,
     get_combined_params,
     get_container_params_type,
     get_ocid_substring,
+    get_params_dict,
     get_params_list,
     get_preferred_compatible_family,
     get_resource_name,
@@ -61,7 +63,11 @@
     ModelDeploymentConfigSummary,
     MultiModelDeploymentConfigLoader,
 )
-from ads.aqua.modeldeployment.constants import DEFAULT_POLL_INTERVAL, DEFAULT_WAIT_TIME
+from ads.aqua.modeldeployment.constants import (
+    DEFAULT_POLL_INTERVAL,
+    DEFAULT_WAIT_TIME,
+    DeploymentType,
+)
 from ads.aqua.modeldeployment.entities import (
     AquaDeployment,
     AquaDeploymentDetail,
@@ -76,6 +82,7 @@
     AQUA_DEPLOYMENT_CONTAINER_CMD_VAR_METADATA_NAME,
     AQUA_DEPLOYMENT_CONTAINER_METADATA_NAME,
     AQUA_DEPLOYMENT_CONTAINER_URI_METADATA_NAME,
+    AQUA_MODEL_DEPLOYMENT_FOLDER,
     AQUA_TELEMETRY_BUCKET,
     AQUA_TELEMETRY_BUCKET_NS,
     COMPARTMENT_OCID,
@@ -162,6 +169,7 @@ def create(
                 cmd_var (Optional[List[str]]): Command variables for the container runtime.
                 freeform_tags (Optional[Dict]): Freeform tags for model deployment.
                 defined_tags (Optional[Dict]): Defined tags for model deployment.
+                deployment_type (Optional[str]): The type of model deployment.
 
         Returns
         -------
@@ -206,13 +214,26 @@ def create(
 
         # Create an AquaModelApp instance once to perform the deployment creation.
         model_app = AquaModelApp()
-        if create_deployment_details.model_id:
+        if (
+            create_deployment_details.model_id
+            or create_deployment_details.deployment_type == DeploymentType.STACKED
+        ):
+            model = create_deployment_details.model_id
+            if not model:
+                if len(create_deployment_details.models) != 1:
+                    raise AquaValueError(
+                        "Invalid 'models' provided. Only one base model is required for model stack deployment."
+                    )
+                model = create_deployment_details.models[0]
+
+            service_model_id = model if isinstance(model, str) else model.model_id
             logger.debug(
-                f"Single model ({create_deployment_details.model_id}) provided. "
+                f"Single model ({service_model_id}) provided. "
                 "Delegating to single model creation method."
             )
+
             aqua_model = model_app.create(
-                model_id=create_deployment_details.model_id,
+                model=model,
                 compartment_id=compartment_id,
                 project_id=project_id,
                 freeform_tags=freeform_tags,
@@ -231,6 +252,7 @@ def create(
                 create_deployment_details=create_deployment_details,
                 container_config=container_config,
             )
+        # TODO: add multi model validation from deployment_type
         else:
             # Collect all unique model IDs (including fine-tuned models)
             source_model_ids = list(
@@ -685,7 +707,7 @@ def _build_model_group_config(
 
     def _create(
         self,
-        aqua_model: DataScienceModel,
+        aqua_model: Union[DataScienceModel, DataScienceModelGroup],
         create_deployment_details: CreateModelDeploymentDetails,
         container_config: Dict,
     ) -> AquaDeployment:
@@ -719,7 +741,10 @@ def _create(
         tags.update({Tags.TASK: aqua_model.freeform_tags.get(Tags.TASK, UNKNOWN)})
 
         # Set up info to get deployment config
-        config_source_id = create_deployment_details.model_id
+        config_source_id = (
+            create_deployment_details.model_id
+            or create_deployment_details.models[0].model_id
+        )
         model_name = aqua_model.display_name
 
         # set up env and cmd var
@@ -870,6 +895,20 @@ def _create(
         deployment_params = get_combined_params(config_params, user_params)
 
         params = f"{params} {deployment_params}".strip()
+
+        if isinstance(aqua_model, DataScienceModelGroup):
+            env_var.update({"VLLM_ALLOW_RUNTIME_LORA_UPDATING": "true"})
+            env_var.update(
+                {"MODEL": f"{AQUA_MODEL_DEPLOYMENT_FOLDER}{aqua_model.base_model_id}/"}
+            )
+
+            params_dict = get_params_dict(params)
+            # updates `--served-model-name` with service model id
+            params_dict.update({"--served-model-name": aqua_model.base_model_id})
+            # adds `--enable_lora` to parameters
+            params_dict.update({"--enable_lora": UNKNOWN})
+            params = build_params_string(params_dict)
+
         if params:
             env_var.update({"PARAMS": params})
         env_vars = container_spec.env_vars if container_spec else []

ads/aqua/modeldeployment/entities.py

@@ -325,6 +325,9 @@ class CreateModelDeploymentDetails(BaseModel):
     defined_tags: Optional[Dict] = Field(
         None, description="Defined tags for model deployment."
     )
+    deployment_type: Optional[str] = Field(
+        None, description="The type of model deployment."
+    )
 
     @model_validator(mode="before")
     @classmethod