added inital code for supporting FT models in multi model

Author: elizjo

ads/aqua/finetuning/entities.py

@@ -3,14 +3,20 @@
 # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
 
 import json
-from typing import Any, Dict, List, Literal, Optional, Union
+from typing import Any, Dict, List, Literal, Optional, Tuple, Union
 
 from pydantic import Field, model_validator
 
 from ads.aqua.common.errors import AquaValueError
+from ads.aqua.common.utils import get_model_by_reference_paths
 from ads.aqua.config.utils.serializer import Serializable
 from ads.aqua.data import AquaResourceIdentifier
-from ads.aqua.finetuning.constants import FineTuningRestrictedParams
+from ads.aqua.finetuning.constants import (
+    FineTuneCustomMetadata,
+    FineTuningRestrictedParams,
+)
+from ads.common.object_storage_details import ObjectStorageDetails
+from ads.model.datascience_model import DataScienceModel
 
 
 class AquaFineTuningParams(Serializable):
@@ -179,3 +185,44 @@ class CreateFineTuningDetails(Serializable):
 
     class Config:
         extra = "ignore"
+
+
+@staticmethod
+def extract_base_model_ocid(aqua_model: DataScienceModel) -> Tuple[str, str]:
+    """Extracts the model_name, base model (config_source_id) OCID of the Fine Tuned Model
+    """
+    config_source_id = aqua_model.custom_metadata_list.get(
+        FineTuneCustomMetadata.FINE_TUNE_SOURCE
+    ).value
+    model_name = aqua_model.custom_metadata_list.get(
+        FineTuneCustomMetadata.FINE_TUNE_SOURCE_NAME
+    ).value
+
+    if not config_source_id or not model_name:
+        raise AquaValueError(
+        f"Either {FineTuneCustomMetadata.FINE_TUNE_SOURCE} or {FineTuneCustomMetadata.FINE_TUNE_SOURCE_NAME} is missing "
+        f"from custom metadata for the model {config_source_id}")
+
+    return config_source_id, model_name
+
+
+@staticmethod
+def set_fine_tune_env_var(aqua_model: DataScienceModel, env_var: Dict[str,str]) -> Dict[str,str]:
+    """Extracts the fine tuning source (fine_tune_output_path).
+    Sets the environment variable (env_var) of the fine tuned model to include FT_model (fine tuning source)"""
+
+    base_model_path, fine_tune_output_path = get_model_by_reference_paths(
+    aqua_model.model_file_description
+    )
+
+    if fine_tune_output_path and ObjectStorageDetails.is_oci_path(fine_tune_output_path):
+        os_path = ObjectStorageDetails.from_path(fine_tune_output_path)
+        fine_tune_output_path = os_path.filepath.rstrip("/")
+
+        env_var.update({"FT_MODEL": f"{fine_tune_output_path}"})
+
+        return env_var
+
+    raise AquaValueError(
+    "Fine tuned output path is not available in the model artifact."
+    )

ads/aqua/model/model.py

@@ -64,6 +64,7 @@
     VALIDATION_METRICS,
     VALIDATION_METRICS_FINAL,
 )
+from ads.aqua.finetuning.entities import extract_base_model_ocid, set_fine_tune_env_var
 from ads.aqua.model.constants import (
     AquaModelMetadataKeys,
     FineTuningCustomMetadata,
@@ -311,6 +312,12 @@ def create_multi(
             #         "Currently only service models are supported for multi model deployment."
             #     )
 
+            is_fine_tuned_model = Tags.AQUA_FINE_TUNED_MODEL_TAG in source_model.freeform_tags
+
+            if is_fine_tuned_model:
+                model.model_id, model.model_name = extract_base_model_ocid(source_model)
+                model.env_var = set_fine_tune_env_var(source_model, model.env_var)
+
             display_name_list.append(display_name)
 
             self._extract_model_task(model, source_model)

ads/aqua/modeldeployment/deployment.py

@@ -25,7 +25,6 @@
     build_pydantic_error_message,
     get_combined_params,
     get_container_params_type,
-    get_model_by_reference_paths,
     get_ocid_substring,
     get_params_dict,
     get_params_list,
@@ -46,7 +45,7 @@
     UNKNOWN_DICT,
 )
 from ads.aqua.data import AquaResourceIdentifier
-from ads.aqua.finetuning.finetuning import FineTuneCustomMetadata
+from ads.aqua.finetuning.entities import extract_base_model_ocid, set_fine_tune_env_var
 from ads.aqua.model import AquaModelApp
 from ads.aqua.model.constants import AquaModelMetadataKeys, ModelCustomMetadataFields
 from ads.aqua.modeldeployment.entities import (
@@ -210,7 +209,8 @@ def create(
                 container_config=container_config,
             )
         else:
-            model_ids = [model.model_id for model in create_deployment_details.models]
+            model_ids =[model.model_id for model in create_deployment_details.models]
+
             try:
                 model_config_summary = self.get_multimodel_deployment_config(
                     model_ids=model_ids, compartment_id=compartment_id
@@ -343,22 +343,6 @@ def _create(
         config_source_id = create_deployment_details.model_id
         model_name = aqua_model.display_name
 
-        is_fine_tuned_model = Tags.AQUA_FINE_TUNED_MODEL_TAG in aqua_model.freeform_tags
-
-        if is_fine_tuned_model:
-            try:
-                config_source_id = aqua_model.custom_metadata_list.get(
-                    FineTuneCustomMetadata.FINE_TUNE_SOURCE
-                ).value
-                model_name = aqua_model.custom_metadata_list.get(
-                    FineTuneCustomMetadata.FINE_TUNE_SOURCE_NAME
-                ).value
-            except ValueError as err:
-                raise AquaValueError(
-                    f"Either {FineTuneCustomMetadata.FINE_TUNE_SOURCE} or {FineTuneCustomMetadata.FINE_TUNE_SOURCE_NAME} is missing "
-                    f"from custom metadata for the model {config_source_id}"
-                ) from err
-
         # set up env and cmd var
         env_var = create_deployment_details.env_var or {}
         cmd_var = create_deployment_details.cmd_var or []
@@ -378,20 +362,11 @@ def _create(
 
         env_var.update({"BASE_MODEL": f"{model_path_prefix}"})
 
-        if is_fine_tuned_model:
-            _, fine_tune_output_path = get_model_by_reference_paths(
-                aqua_model.model_file_description
-            )
-
-            if not fine_tune_output_path:
-                raise AquaValueError(
-                    "Fine tuned output path is not available in the model artifact."
-                )
-
-            os_path = ObjectStorageDetails.from_path(fine_tune_output_path)
-            fine_tune_output_path = os_path.filepath.rstrip("/")
+        is_fine_tuned_model = Tags.AQUA_FINE_TUNED_MODEL_TAG in aqua_model.freeform_tags
 
-            env_var.update({"FT_MODEL": f"{fine_tune_output_path}"})
+        if is_fine_tuned_model:
+            config_source_id, model_name = extract_base_model_ocid(aqua_model)
+            env_var = set_fine_tune_env_var(aqua_model, env_var)
 
         container_type_key = self._get_container_type_key(
             model=aqua_model,
@@ -647,6 +622,12 @@ def _create_multi(
             config_data = {"params": params, "model_path": artifact_path_prefix}
             if model.model_task:
                 config_data["model_task"] = model.model_task
+
+            fine_tuned_model = model.env_var.get("FT_MODEL")
+
+            if fine_tuned_model:
+                config_data["FT_MODEL"] = fine_tuned_model
+
             model_config.append(config_data)
             model_name_list.append(model.model_name)