[Ready For Review][AQUA][Multi-Model] Enhance AQUA CLI to Accept Fine-Tuned Weights Under Base Model in Multi-Model Deployment (#1209)

Author: mrDzurb

ads/aqua/app.py

@@ -5,6 +5,7 @@
 import json
 import os
 import traceback
+from concurrent.futures import ThreadPoolExecutor
 from dataclasses import fields
 from datetime import datetime, timedelta
 from itertools import chain
@@ -58,6 +59,8 @@
 class AquaApp:
     """Base Aqua App to contain common components."""
 
+    MAX_WORKERS = 10  # Number of workers for asynchronous resource loading
+
     @telemetry(name="aqua")
     def __init__(self) -> None:
         if OCI_RESOURCE_PRINCIPAL_VERSION:
@@ -128,20 +131,69 @@ def update_model_provenance(
             update_model_provenance_details=update_model_provenance_details,
         )
 
-    # TODO: refactor model evaluation implementation to use it.
     @staticmethod
     def get_source(source_id: str) -> Union[ModelDeployment, DataScienceModel]:
-        if is_valid_ocid(source_id):
-            if "datasciencemodeldeployment" in source_id:
-                return ModelDeployment.from_id(source_id)
-            elif "datasciencemodel" in source_id:
-                return DataScienceModel.from_id(source_id)
+        """
+        Fetches a model or model deployment based on the provided OCID.
+
+        Parameters
+        ----------
+        source_id : str
+            OCID of the Data Science model or model deployment.
+
+        Returns
+        -------
+        Union[ModelDeployment, DataScienceModel]
+            The corresponding resource object.
+
+        Raises
+        ------
+        AquaValueError
+            If the OCID is invalid or unsupported.
+        """
+        logger.debug(f"Resolving source for ID: {source_id}")
+        if not is_valid_ocid(source_id):
+            logger.error(f"Invalid OCID format: {source_id}")
+            raise AquaValueError(
+                f"Invalid source ID: {source_id}. Please provide a valid model or model deployment OCID."
+            )
+
+        if "datasciencemodeldeployment" in source_id:
+            logger.debug(f"Identified as ModelDeployment OCID: {source_id}")
+            return ModelDeployment.from_id(source_id)
 
+        if "datasciencemodel" in source_id:
+            logger.debug(f"Identified as DataScienceModel OCID: {source_id}")
+            return DataScienceModel.from_id(source_id)
+
+        logger.error(f"Unrecognized OCID type: {source_id}")
         raise AquaValueError(
-            f"Invalid source {source_id}. "
-            "Specify either a model or model deployment id."
+            f"Unsupported source ID type: {source_id}. Must be a model or model deployment OCID."
         )
 
+    def get_multi_source(
+        self,
+        ids: List[str],
+    ) -> Dict[str, Union[ModelDeployment, DataScienceModel]]:
+        """
+        Retrieves multiple DataScience resources concurrently.
+
+        Parameters
+        ----------
+        ids : List[str]
+            A list of DataScience OCIDs.
+
+        Returns
+        -------
+        Dict[str, Union[ModelDeployment, DataScienceModel]]
+            A mapping from OCID to the corresponding resolved resource object.
+        """
+        logger.debug(f"Fetching {ids} sources in parallel.")
+        with ThreadPoolExecutor(max_workers=self.MAX_WORKERS) as executor:
+            results = list(executor.map(self.get_source, ids))
+
+        return dict(zip(ids, results))
+
     # TODO: refactor model evaluation implementation to use it.
     @staticmethod
     def create_model_version_set(

ads/aqua/common/entities.py

@@ -6,7 +6,7 @@
 from typing import Any, Dict, List, Optional
 
 from oci.data_science.models import Model
-from pydantic import BaseModel, Field, model_validator
+from pydantic import BaseModel, ConfigDict, Field, model_validator
 
 from ads.aqua import logger
 from ads.aqua.config.utils.serializer import Serializable
@@ -80,24 +80,29 @@ class GPUShapesIndex(Serializable):
 
 class ComputeShapeSummary(Serializable):
     """
-    Represents the specifications of a compute instance's shape.
+    Represents the specifications of a compute instance shape,
+    including CPU, memory, and optional GPU characteristics.
     """
 
     core_count: Optional[int] = Field(
-        default=None, description="The number of CPU cores available."
+        default=None,
+        description="Total number of CPU cores available for the compute shape.",
     )
     memory_in_gbs: Optional[int] = Field(
-        default=None, description="The amount of memory (in GB) available."
+        default=None,
+        description="Amount of memory (in GB) available for the compute shape.",
     )
     name: Optional[str] = Field(
-        default=None, description="The name identifier of the compute shape."
+        default=None,
+        description="Full name of the compute shape, e.g., 'VM.GPU.A10.2'.",
     )
     shape_series: Optional[str] = Field(
-        default=None, description="The series or category of the compute shape."
+        default=None,
+        description="Shape family or series, e.g., 'GPU', 'Standard', etc.",
     )
     gpu_specs: Optional[GPUSpecs] = Field(
         default=None,
-        description="The GPU specifications associated with the compute shape.",
+        description="Optional GPU specifications associated with the shape.",
     )
 
     @model_validator(mode="after")
@@ -136,27 +141,46 @@ def set_gpu_specs(cls, model: "ComputeShapeSummary") -> "ComputeShapeSummary":
         return model
 
 
-class LoraModuleSpec(Serializable):
+class LoraModuleSpec(BaseModel):
     """
-    Lightweight descriptor for LoRA Modules used in fine-tuning models.
+    Descriptor for a LoRA (Low-Rank Adaptation) module used in fine-tuning base models.
+
+    This class is used to define a single fine-tuned module that can be loaded during
+    multi-model deployment alongside a base model.
 
     Attributes
     ----------
     model_id : str
-        The unique identifier of the fine tuned model.
-    model_name : str
-        The name of the fine-tuned model.
-    model_path : str
-        The model-by-reference path to the LoRA Module within the model artifact
+        The OCID of the fine-tuned model registered in the OCI Model Catalog.
+    model_name : Optional[str]
+        The unique name used to route inference requests to this model variant.
+    model_path : Optional[str]
+        The relative path within the artifact pointing to the LoRA adapter weights.
     """
 
-    model_id: Optional[str] = Field(None, description="The fine tuned model OCID to deploy.")
-    model_name: Optional[str] = Field(None, description="The name of the fine-tuned model.")
+    model_config = ConfigDict(protected_namespaces=(), extra="allow")
+
+    model_id: str = Field(
+        ...,
+        description="OCID of the fine-tuned model (must be registered in the Model Catalog).",
+    )
+    model_name: Optional[str] = Field(
+        default=None,
+        description="Name assigned to the fine-tuned model for serving (used as inference route).",
+    )
     model_path: Optional[str] = Field(
-        None,
-        description="The model-by-reference path to the LoRA Module within the model artifact.",
+        default=None,
+        description="Relative path to the LoRA weights inside the model artifact.",
     )
 
+    @model_validator(mode="before")
+    @classmethod
+    def validate_lora_module(cls, data: dict) -> dict:
+        """Validates that required structure exists for a LoRA module."""
+        if "model_id" not in data or not data["model_id"]:
+            raise ValueError("Missing required field: 'model_id' for fine-tuned model.")
+        return data
+
 
 class AquaMultiModelRef(Serializable):
     """
@@ -203,6 +227,22 @@ class AquaMultiModelRef(Serializable):
         description="For fine tuned models, the artifact path of the modified model weights",
     )
 
+    def all_model_ids(self) -> List[str]:
+        """
+        Returns all associated model OCIDs, including the base model and any fine-tuned models.
+
+        Returns
+        -------
+        List[str]
+            A list of all model OCIDs associated with this multi-model reference.
+        """
+        ids = {self.model_id}
+        if self.fine_tune_weights:
+            ids.update(
+                module.model_id for module in self.fine_tune_weights if module.model_id
+            )
+        return list(ids)
+
     class Config:
         extra = "ignore"
         protected_namespaces = ()