Merge conflicts

Author: agrimk

ads/aqua/app.py

@@ -64,6 +64,7 @@ def __init__(self) -> None:
             set_auth("resource_principal")
         self._auth = default_signer({"service_endpoint": OCI_ODSC_SERVICE_ENDPOINT})
         self.ds_client = oc.OCIClientFactory(**self._auth).data_science
+        self.compute_client = oc.OCIClientFactory(**default_signer()).compute
         self.logging_client = oc.OCIClientFactory(**default_signer()).logging_management
         self.identity_client = oc.OCIClientFactory(**default_signer()).identity
         self.region = extract_region(self._auth)

ads/aqua/common/entities.py

@@ -157,6 +157,8 @@ class AquaMultiModelRef(Serializable):
         Optional environment variables to override during deployment.
     artifact_location : Optional[str]
         Artifact path of model in the multimodel group.
+    fine_tune_weights_location : Optional[str]
+        For fine tuned models, the artifact path of the modified model weights
     """
 
     model_id: str = Field(..., description="The model OCID to deploy.")
@@ -171,6 +173,9 @@ class AquaMultiModelRef(Serializable):
     artifact_location: Optional[str] = Field(
         None, description="Artifact path of model in the multimodel group."
     )
+    fine_tune_weights_location: Optional[str] = Field(
+        None, description="For fine tuned models, the artifact path of the modified model weights"
+    )
 
     class Config:
         extra = "ignore"

ads/aqua/common/utils.py

@@ -1158,9 +1158,11 @@ def validate_cmd_var(cmd_var: List[str], overrides: List[str]) -> List[str]:
 
 
 def build_pydantic_error_message(ex: ValidationError):
-    """Added to handle error messages from pydantic model validator.
+    """
+    Added to handle error messages from pydantic model validator.
     Combine both loc and msg for errors where loc (field) is present in error details, else only build error
-    message using msg field."""
+    message using msg field.
+    """
 
     return {
         ".".join(map(str, e["loc"])): e["msg"]
@@ -1185,67 +1187,71 @@ def is_pydantic_model(obj: object) -> bool:
 
 @cached(cache=TTLCache(maxsize=1, ttl=timedelta(minutes=5), timer=datetime.now))
 def load_gpu_shapes_index(
-    auth: Optional[Dict] = None,
+    auth: Optional[Dict[str, Any]] = None,
 ) -> GPUShapesIndex:
     """
-    Loads the GPU shapes index from Object Storage or a local resource folder.
+    Load the GPU shapes index, preferring the OS bucket copy over the local one.
 
-    The function first attempts to load the file from an Object Storage bucket using fsspec.
-    If the loading fails (due to connection issues, missing file, etc.), it falls back to
-    loading the index from a local file.
+    Attempts to read `gpu_shapes_index.json` from OCI Object Storage first;
+    if that succeeds, those entries will override the local defaults.
 
     Parameters
     ----------
-    auth: (Dict, optional). Defaults to None.
-        The default authentication is set using `ads.set_auth` API. If you need to override the
-        default, use the `ads.common.auth.api_keys` or `ads.common.auth.resource_principal` to create appropriate
-        authentication signer and kwargs required to instantiate IdentityClient object.
+    auth
+        Optional auth dict (as returned by `ads.common.auth.default_signer()`)
+        to pass through to `fsspec.open()`.
 
     Returns
     -------
-    GPUShapesIndex: The parsed GPU shapes index.
+    GPUShapesIndex
+        Merged index where any shape present remotely supersedes the local entry.
 
     Raises
     ------
-    FileNotFoundError: If the GPU shapes index cannot be found in either Object Storage or locally.
-    json.JSONDecodeError: If the JSON is malformed.
+    json.JSONDecodeError
+        If any of the JSON is malformed.
     """
     file_name = "gpu_shapes_index.json"
-    data: Dict[str, Any] = {}
 
-    # Check if the CONDA_BUCKET_NS environment variable is set.
+    # Try remote load
+    remote_data: Dict[str, Any] = {}
     if CONDA_BUCKET_NS:
         try:
             auth = auth or authutil.default_signer()
-            # Construct the object storage path. Adjust bucket name and path as needed.
             storage_path = (
                 f"oci://{CONDA_BUCKET_NAME}@{CONDA_BUCKET_NS}/service_pack/{file_name}"
             )
-            logger.debug("Loading GPU shapes index from Object Storage")
-            with fsspec.open(storage_path, mode="r", **auth) as file_obj:
-                data = json.load(file_obj)
-            logger.debug("Successfully loaded GPU shapes index.")
-        except Exception as ex:
             logger.debug(
-                f"Failed to load GPU shapes index from Object Storage. Details: {ex}"
+                "Loading GPU shapes index from Object Storage: %s", storage_path
             )
-
-    # If loading from Object Storage failed, load from the local resource folder.
-    if not data:
-        try:
-            local_path = os.path.join(
-                os.path.dirname(__file__), "../resources", file_name
-            )
-            logger.debug(f"Loading GPU shapes index from {local_path}.")
-            with open(local_path) as file_obj:
-                data = json.load(file_obj)
-            logger.debug("Successfully loaded GPU shapes index.")
-        except Exception as e:
+            with fsspec.open(storage_path, mode="r", **auth) as f:
+                remote_data = json.load(f)
             logger.debug(
-                f"Failed to load GPU shapes index from {local_path}. Details: {e}"
+                "Loaded %d shapes from Object Storage",
+                len(remote_data.get("shapes", {})),
             )
+        except Exception as ex:
+            logger.debug("Remote load failed (%s); falling back to local", ex)
+
+    # Load local copy
+    local_data: Dict[str, Any] = {}
+    local_path = os.path.join(os.path.dirname(__file__), "../resources", file_name)
+    try:
+        logger.debug("Loading GPU shapes index from local file: %s", local_path)
+        with open(local_path) as f:
+            local_data = json.load(f)
+        logger.debug(
+            "Loaded %d shapes from local file", len(local_data.get("shapes", {}))
+        )
+    except Exception as ex:
+        logger.debug("Local load GPU shapes index failed (%s)", ex)
+
+    # Merge: remote shapes override local
+    local_shapes = local_data.get("shapes", {})
+    remote_shapes = remote_data.get("shapes", {})
+    merged_shapes = {**local_shapes, **remote_shapes}
 
-    return GPUShapesIndex(**data)
+    return GPUShapesIndex(shapes=merged_shapes)
 
 
 def get_preferred_compatible_family(selected_families: set[str]) -> str:

ads/aqua/extension/ui_handler.py

@@ -15,7 +15,7 @@
 from ads.aqua.extension.utils import validate_function_parameters
 from ads.aqua.model.entities import ImportModelDetails
 from ads.aqua.ui import AquaUIApp
-from ads.config import COMPARTMENT_OCID
+from ads.config import COMPARTMENT_OCID, IS_BYOR_ENABLED
 
 
 @dataclass
@@ -82,6 +82,10 @@ def get(self, id=""):
             return self.is_bucket_versioned()
         elif paths.startswith("aqua/containers"):
             return self.list_containers()
+        elif paths.startswith("aqua/capacityreservations/enabled"):
+            return self.is_capacity_reservations_enabled()
+        elif paths.startswith("aqua/capacityreservations"):
+            return self.list_capacity_reservations()
         else:
             raise HTTPError(400, f"The request {self.request.path} is invalid.")
 
@@ -103,6 +107,19 @@ def list_log_groups(self, **kwargs):
             AquaUIApp().list_log_groups(compartment_id=compartment_id, **kwargs)
         )
 
+    def is_capacity_reservations_enabled(self):
+        """Checks if the tenant is whitelisted for BYOR (Bring your own reservation) feature."""
+        return self.finish({"status": str(IS_BYOR_ENABLED).strip().lower() == "true"})
+
+    def list_capacity_reservations(self, **kwargs):
+        """Lists users compute reservations in a specified compartment."""
+        compartment_id = self.get_argument("compartment_id", default=COMPARTMENT_OCID)
+        return self.finish(
+            AquaUIApp().list_capacity_reservations(
+                compartment_id=compartment_id, **kwargs
+            )
+        )
+
     def list_logs(self, log_group_id: str, **kwargs):
         """Lists the specified log group's log objects."""
         return self.finish(AquaUIApp().list_logs(log_group_id=log_group_id, **kwargs))
@@ -279,4 +296,5 @@ def post(self, *args, **kwargs):
     ("bucket/versioning/?([^/]*)", AquaUIHandler),
     ("containers/?([^/]*)", AquaUIHandler),
     ("cli/?([^/]*)", AquaCLIHandler),
+    ("capacityreservations/?([^/]*)", AquaUIHandler),
 ]

ads/aqua/model/model.py

@@ -83,6 +83,10 @@
     ModelValidationResult,
 )
 from ads.aqua.model.enums import MultiModelSupportedTaskType
+from ads.aqua.model.utils import (
+    extract_base_model_from_ft,
+    extract_fine_tune_artifacts_path,
+)
 from ads.common.auth import default_signer
 from ads.common.oci_resource import SEARCH_TYPE, OCIResource
 from ads.common.utils import (
@@ -311,12 +315,21 @@ def create_multi(
             #         "Currently only service models are supported for multi model deployment."
             #     )
 
+            # check if model is a fine-tuned model and if so, add the fine tuned weights path to the fine_tune_weights_location pydantic field
+            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_from_ft(source_model)
+                model_artifact_path, model.fine_tune_weights_location = extract_fine_tune_artifacts_path(source_model)
+
+            else:
+                # Retrieve model artifact for base models
+                model_artifact_path = source_model.artifact
+
             display_name_list.append(display_name)
 
             self._extract_model_task(model, source_model)
 
-            # Retrieve model artifact
-            model_artifact_path = source_model.artifact
             if not model_artifact_path:
                 raise AquaValueError(
                     f"Model '{display_name}' (ID: {model.model_id}) has no artifacts. "

ads/aqua/model/utils.py

@@ -0,0 +1,52 @@
+#!/usr/bin/env python
+# Copyright (c) 2025 Oracle and/or its affiliates.
+# Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
+"""AQUA model utils"""
+
+from typing import Dict, Optional, Tuple
+
+from ads.aqua.common.entities import AquaMultiModelRef
+from ads.aqua.common.errors import AquaValueError
+from ads.aqua.common.utils import get_model_by_reference_paths
+from ads.aqua.finetuning.constants import FineTuneCustomMetadata
+from ads.common.object_storage_details import ObjectStorageDetails
+from ads.model.datascience_model import DataScienceModel
+
+
+def extract_base_model_from_ft(aqua_model: DataScienceModel) -> Tuple[str, str]:
+    """Extracts the model_name and base model OCID (config_source_id) OCID for a 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
+
+
+def extract_fine_tune_artifacts_path(aqua_model: DataScienceModel) -> Tuple[str, str]:
+    """Extracts the fine tuning source (fine_tune_output_path) and base model path from the DataScienceModel Object"""
+
+    base_model_path, fine_tune_output_path = get_model_by_reference_paths(
+        aqua_model.model_file_description
+    )
+
+    if not fine_tune_output_path or not ObjectStorageDetails.is_oci_path(
+        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("/")
+
+    return base_model_path, fine_tune_output_path

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,9 +45,12 @@
     UNKNOWN_DICT,
 )
 from ads.aqua.data import AquaResourceIdentifier
-from ads.aqua.finetuning.finetuning import FineTuneCustomMetadata
 from ads.aqua.model import AquaModelApp
-from ads.aqua.model.constants import AquaModelMetadataKeys, ModelCustomMetadataFields, ModelType
+from ads.aqua.model.constants import AquaModelMetadataKeys, ModelCustomMetadataFields
+from ads.aqua.model.utils import (
+    extract_base_model_from_ft,
+    extract_fine_tune_artifacts_path,
+)
 from ads.aqua.modeldeployment.entities import (
     AquaDeployment,
     AquaDeploymentConfig,
@@ -211,6 +213,7 @@ def create(
             )
         else:
             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 +346,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,19 +365,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
 
+        if is_fine_tuned_model:
+            config_source_id, model_name = extract_base_model_from_ft(aqua_model)
+            _, fine_tune_output_path = extract_fine_tune_artifacts_path(aqua_model)
             env_var.update({"FT_MODEL": f"{fine_tune_output_path}"})
 
         container_type_key = self._get_container_type_key(
@@ -647,6 +626,10 @@ def _create_multi(
             config_data = {"params": params, "model_path": artifact_path_prefix}
             if model.model_task:
                 config_data["model_task"] = model.model_task
+
+            if model.fine_tune_weights_location:
+                config_data["fine_tune_weights_location"] = model.fine_tune_weights_location
+
             model_config.append(config_data)
             model_name_list.append(model.model_name)
 

ads/aqua/ui.py

@@ -90,6 +90,26 @@ def list_logs(self, **kwargs) -> str:
         res = self.logging_client.list_logs(log_group_id=log_group_id, **kwargs).data
         return sanitize_response(oci_client=self.logging_client, response=res)
 
+    @telemetry(entry_point="plugin=ui&action=list_capacity_reservations", name="aqua")
+    def list_capacity_reservations(self, **kwargs) -> list:
+        """
+        Lists users compute reservations in a specified compartment
+
+        Returns
+        -------
+            json representation of `oci.core.models.ComputeCapacityReservationSummary`.
+
+        """
+        compartment_id = kwargs.pop("compartment_id", COMPARTMENT_OCID)
+        logger.info(f"Loading Capacity reservations from compartment: {compartment_id}")
+
+        reservations = self.compute_client.list_compute_capacity_reservations(
+            compartment_id=compartment_id, **kwargs
+        )
+        return sanitize_response(
+            oci_client=self.compute_client, response=reservations.data
+        )
+
     @telemetry(entry_point="plugin=ui&action=list_compartments", name="aqua")
     def list_compartments(self) -> str:
         """Lists the compartments in a tenancy specified by TENANCY_OCID env variable. This is a pass through the OCI list_compartments