Merge branch 'main' into track_md_logs_for_error_logging

Author: agrimk

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/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.

ads/aqua/client/openai_client.py

@@ -32,7 +32,7 @@ class ModelDeploymentBaseEndpoint(ExtendedEnum):
     """Supported base endpoints for model deployments."""
 
     PREDICT = "predict"
-    PREDICT_WITH_RESPONSE_STREAM = "predictwithresponsestream"
+    PREDICT_WITH_RESPONSE_STREAM = "predictWithResponseStream"
 
 
 class AquaOpenAIMixin:
@@ -51,9 +51,9 @@ def _patch_route(self, original_path: str) -> str:
         Returns:
             str: The normalized OpenAI-compatible route path (e.g., '/v1/chat/completions').
         """
-        normalized_path = original_path.lower().rstrip("/")
+        normalized_path = original_path.rstrip("/")
 
-        match = re.search(r"/predict(withresponsestream)?", normalized_path)
+        match = re.search(r"/predict(WithResponseStream)?", normalized_path)
         if not match:
             logger.debug("Route header cannot be resolved from path: %s", original_path)
             return ""
@@ -71,7 +71,7 @@ def _patch_route(self, original_path: str) -> str:
                 "Route suffix does not start with a version prefix (e.g., '/v1'). "
                 "This may lead to compatibility issues with OpenAI-style endpoints. "
                 "Consider updating the URL to include a version prefix, "
-                "such as '/predict/v1' or '/predictwithresponsestream/v1'."
+                "such as '/predict/v1' or '/predictWithResponseStream/v1'."
             )
             # route_suffix = f"v1/{route_suffix}"
 
@@ -124,13 +124,13 @@ def _patch_headers(self, request: httpx.Request) -> None:
 
     def _patch_url(self) -> httpx.URL:
         """
-        Strips any suffixes from the base URL to retain only the `/predict` or `/predictwithresponsestream` path.
+        Strips any suffixes from the base URL to retain only the `/predict` or `/predictWithResponseStream` path.
 
         Returns:
             httpx.URL: The normalized base URL with the correct model deployment path.
         """
-        base_path = f"{self.base_url.path.lower().rstrip('/')}/"
-        match = re.search(r"/predict(withresponsestream)?/", base_path)
+        base_path = f"{self.base_url.path.rstrip('/')}/"
+        match = re.search(r"/predict(WithResponseStream)?/", base_path)
         if match:
             trimmed = base_path[: match.end() - 1]
             return self.base_url.copy_with(path=trimmed)
@@ -144,7 +144,7 @@ def _prepare_request_common(self, request: httpx.Request) -> None:
 
         This includes:
         - Patching headers with streaming and routing info.
-        - Normalizing the URL path to include only `/predict` or `/predictwithresponsestream`.
+        - Normalizing the URL path to include only `/predict` or `/predictWithResponseStream`.
 
         Args:
             request (httpx.Request): The outgoing HTTPX request.
@@ -176,6 +176,7 @@ def __init__(
         http_client: Optional[httpx.Client] = None,
         http_client_kwargs: Optional[Dict[str, Any]] = None,
         _strict_response_validation: bool = False,
+        patch_headers: bool = False,
         **kwargs: Any,
     ) -> None:
         """
@@ -196,6 +197,7 @@ def __init__(
             http_client (httpx.Client, optional): Custom HTTP client; if not provided, one will be auto-created.
             http_client_kwargs (dict[str, Any], optional): Extra kwargs for auto-creating the HTTP client.
             _strict_response_validation (bool, optional): Enable strict response validation.
+            patch_headers (bool, optional): If True, redirects the requests by modifying the headers.
             **kwargs: Additional keyword arguments passed to the parent __init__.
         """
         if http_client is None:
@@ -207,6 +209,8 @@ def __init__(
             logger.debug("API key not provided; using default placeholder for OCI.")
             api_key = "OCI"
 
+        self.patch_headers = patch_headers
+
         super().__init__(
             api_key=api_key,
             organization=organization,
@@ -229,7 +233,8 @@ def _prepare_request(self, request: httpx.Request) -> None:
         Args:
             request (httpx.Request): The outgoing HTTP request.
         """
-        self._prepare_request_common(request)
+        if self.patch_headers:
+            self._prepare_request_common(request)
 
 
 class AsyncOpenAI(openai.AsyncOpenAI, AquaOpenAIMixin):
@@ -248,6 +253,7 @@ def __init__(
         http_client: Optional[httpx.Client] = None,
         http_client_kwargs: Optional[Dict[str, Any]] = None,
         _strict_response_validation: bool = False,
+        patch_headers: bool = False,
         **kwargs: Any,
     ) -> None:
         """
@@ -269,6 +275,7 @@ def __init__(
             http_client (httpx.AsyncClient, optional): Custom asynchronous HTTP client; if not provided, one will be auto-created.
             http_client_kwargs (dict[str, Any], optional): Extra kwargs for auto-creating the HTTP client.
             _strict_response_validation (bool, optional): Enable strict response validation.
+            patch_headers (bool, optional): If True, redirects the requests by modifying the headers.
             **kwargs: Additional keyword arguments passed to the parent __init__.
         """
         if http_client is None:
@@ -280,6 +287,8 @@ def __init__(
             logger.debug("API key not provided; using default placeholder for OCI.")
             api_key = "OCI"
 
+        self.patch_headers = patch_headers
+
         super().__init__(
             api_key=api_key,
             organization=organization,
@@ -302,4 +311,5 @@ async def _prepare_request(self, request: httpx.Request) -> None:
         Args:
             request (httpx.Request): The outgoing HTTP request.
         """
-        self._prepare_request_common(request)
+        if self.patch_headers:
+            self._prepare_request_common(request)