Merge branch 'main' of https://github.com/oracle/accelerated-data-science into feature/model_group

Author: lu-ohai

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
@@ -22,7 +23,7 @@
 from ads.aqua import logger
 from ads.aqua.common.entities import ModelConfigResult
 from ads.aqua.common.enums import ConfigFolder, Tags
-from ads.aqua.common.errors import AquaRuntimeError, AquaValueError
+from ads.aqua.common.errors import AquaValueError
 from ads.aqua.common.utils import (
     _is_valid_mvs,
     get_artifact_path,
@@ -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(
@@ -284,8 +336,11 @@ def if_artifact_exist(self, model_id: str, **kwargs) -> bool:
                 logger.info(f"Artifact not found in model {model_id}.")
                 return False
 
+    @cached(cache=TTLCache(maxsize=5, ttl=timedelta(minutes=1), timer=datetime.now))
     def get_config_from_metadata(
-        self, model_id: str, metadata_key: str
+        self,
+        model_id: str,
+        metadata_key: str,
     ) -> ModelConfigResult:
         """Gets the config for the given Aqua model from model catalog metadata content.
 
@@ -300,8 +355,9 @@ def get_config_from_metadata(
         ModelConfigResult
             A Pydantic model containing the model_details (extracted from OCI) and the config dictionary.
         """
-        config = {}
+        config: Dict[str, Any] = {}
         oci_model = self.ds_client.get_model(model_id).data
+
         try:
             config = self.ds_client.get_model_defined_metadatum_artifact_content(
                 model_id, metadata_key
@@ -321,7 +377,7 @@ def get_config_from_metadata(
             )
         return ModelConfigResult(config=config, model_details=oci_model)
 
-    @cached(cache=TTLCache(maxsize=1, ttl=timedelta(minutes=1), timer=datetime.now))
+    @cached(cache=TTLCache(maxsize=1, ttl=timedelta(minutes=5), timer=datetime.now))
     def get_config(
         self,
         model_id: str,
@@ -346,8 +402,10 @@ def get_config(
         ModelConfigResult
             A Pydantic model containing the model_details (extracted from OCI) and the config dictionary.
         """
-        config_folder = config_folder or ConfigFolder.CONFIG
+        config: Dict[str, Any] = {}
         oci_model = self.ds_client.get_model(model_id).data
+
+        config_folder = config_folder or ConfigFolder.CONFIG
         oci_aqua = (
             (
                 Tags.AQUA_TAG in oci_model.freeform_tags
@@ -357,9 +415,9 @@ def get_config(
             else False
         )
         if not oci_aqua:
-            raise AquaRuntimeError(f"Target model {oci_model.id} is not an Aqua model.")
+            logger.debug(f"Target model {oci_model.id} is not an Aqua model.")
+            return ModelConfigResult(config=config, model_details=oci_model)
 
-        config: Dict[str, Any] = {}
         artifact_path = get_artifact_path(oci_model.custom_metadata_list)
         if not artifact_path:
             logger.debug(

ads/aqua/cli.py

@@ -94,3 +94,18 @@ def _validate_value(flag, value):
                 "If you intend to chain a function call to the result, please separate the "
                 "flag and the subsequent function call with separator `-`."
             )
+    
+    @staticmethod
+    def install():
+        """Install ADS Aqua Extension from wheel file. Set enviroment variable `AQUA_EXTENSTION_PATH` to change the wheel file path.
+
+        Return 
+        ------
+        int:
+            Installatation status.
+        """
+        import subprocess
+
+        wheel_file_path = os.environ.get("AQUA_EXTENSTION_PATH", "/ads/extension/adsjupyterlab_aqua_extension*.whl")
+        status =  subprocess.run(f"pip install {wheel_file_path}",shell=True)
+        return status.check_returncode

ads/aqua/client/client.py

@@ -61,13 +61,20 @@ class HttpxOCIAuth(httpx.Auth):
 
     def __init__(self, signer: Optional[oci.signer.Signer] = None):
         """
-        Initialize the HttpxOCIAuth instance.
+        Initializes the authentication handler with the given or default OCI signer.
 
-        Args:
-            signer (oci.signer.Signer): The OCI signer to use for signing requests.
+        Parameters
+        ----------
+        signer : oci.signer.Signer, optional
+            The OCI signer instance to use. If None, a default signer will be retrieved.
         """
-
-        self.signer = signer or authutil.default_signer().get("signer")
+        try:
+            self.signer = signer or authutil.default_signer().get("signer")
+            if not self.signer:
+                raise ValueError("OCI signer could not be initialized.")
+        except Exception as e:
+            logger.error("Failed to initialize OCI signer: %s", e)
+            raise
 
     def auth_flow(self, request: httpx.Request) -> Iterator[httpx.Request]:
         """
@@ -80,21 +87,31 @@ def auth_flow(self, request: httpx.Request) -> Iterator[httpx.Request]:
             httpx.Request: The signed HTTPX request.
         """
         # Create a requests.Request object from the HTTPX request
-        req = requests.Request(
-            method=request.method,
-            url=str(request.url),
-            headers=dict(request.headers),
-            data=request.content,
-        )
-        prepared_request = req.prepare()
+        try:
+            req = requests.Request(
+                method=request.method,
+                url=str(request.url),
+                headers=dict(request.headers),
+                data=request.content,
+            )
+            prepared_request = req.prepare()
+            self.signer.do_request_sign(prepared_request)
+
+            # Replace headers on the original HTTPX request with signed headers
+            request.headers.update(prepared_request.headers)
+            logger.debug("Successfully signed request to %s", request.url)
 
-        # Sign the request using the OCI Signer
-        self.signer.do_request_sign(prepared_request)
+            # Fix for GET/DELETE requests that OCI Gateway expects with Content-Length
+            if (
+                request.method in ["GET", "DELETE"]
+                and "content-length" not in request.headers
+            ):
+                request.headers["content-length"] = "0"
 
-        # Update the original HTTPX request with the signed headers
-        request.headers.update(prepared_request.headers)
+        except Exception as e:
+            logger.error("Failed to sign request to %s: %s", request.url, e)
+            raise
 
-        # Proceed with the request
         yield request
 
 
@@ -330,8 +347,8 @@ def _prepare_headers(
             "Content-Type": "application/json",
             "Accept": "text/event-stream" if stream else "application/json",
         }
-        if stream:
-            default_headers["enable-streaming"] = "true"
+        # if stream:
+        #     default_headers["enable-streaming"] = "true"
         if headers:
             default_headers.update(headers)
 
@@ -495,7 +512,7 @@ def generate(
         prompt: str,
         payload: Optional[Dict[str, Any]] = None,
         headers: Optional[Dict[str, str]] = None,
-        stream: bool = True,
+        stream: bool = False,
     ) -> Union[Dict[str, Any], Iterator[Mapping[str, Any]]]:
         """
         Generate text completion for the given prompt.
@@ -521,7 +538,7 @@ def chat(
         messages: List[Dict[str, Any]],
         payload: Optional[Dict[str, Any]] = None,
         headers: Optional[Dict[str, str]] = None,
-        stream: bool = True,
+        stream: bool = False,
     ) -> Union[Dict[str, Any], Iterator[Mapping[str, Any]]]:
         """
         Perform a chat interaction with the model.