Adds timeout option for loading AQUA config. (#878)

Author: mrDzurb

ads/aqua/common/utils.py

@@ -28,10 +28,11 @@
 from ads.aqua.constants import *
 from ads.aqua.data import AquaResourceIdentifier
 from ads.common.auth import default_signer
+from ads.common.decorator.threaded import threaded
 from ads.common.extended_enum import ExtendedEnumMeta
 from ads.common.object_storage_details import ObjectStorageDetails
 from ads.common.oci_resource import SEARCH_TYPE, OCIResource
-from ads.common.utils import get_console_link, upload_to_os, copy_file
+from ads.common.utils import copy_file, get_console_link, upload_to_os
 from ads.config import AQUA_SERVICE_MODELS_BUCKET, CONDA_BUCKET_NS, TENANCY_OCID
 from ads.model import DataScienceModel, ModelVersionSet
 
@@ -195,6 +196,7 @@ def read_file(file_path: str, **kwargs) -> str:
         return UNKNOWN
 
 
+@threaded()
 def load_config(file_path: str, config_file_name: str, **kwargs) -> dict:
     artifact_path = f"{file_path.rstrip('/')}/{config_file_name}"
     if artifact_path.startswith("oci://"):
@@ -540,8 +542,10 @@ def get_container_image(
 
 
 def fetch_service_compartment() -> Union[str, None]:
-    """Loads the compartment mapping json from service bucket. This json file has a service-model-compartment key which
-    contains a dictionary of namespaces and the compartment OCID of the service models in that namespace.
+    """
+    Loads the compartment mapping json from service bucket.
+    This json file has a service-model-compartment key which contains a dictionary of namespaces
+    and the compartment OCID of the service models in that namespace.
     """
     config_file_name = (
         f"oci://{AQUA_SERVICE_MODELS_BUCKET}@{CONDA_BUCKET_NS}/service_models/config"
@@ -554,8 +558,8 @@ def fetch_service_compartment() -> Union[str, None]:
         )
     except Exception as e:
         logger.debug(
-            f"Config file {config_file_name}/{CONTAINER_INDEX} to fetch service compartment OCID could not be found. "
-            f"\n{str(e)}."
+            f"Config file {config_file_name}/{CONTAINER_INDEX} to fetch service compartment OCID "
+            f"could not be found. \n{str(e)}."
         )
         return
     compartment_mapping = config.get(COMPARTMENT_MAPPING_KEY)

ads/common/decorator/threaded.py

@@ -0,0 +1,97 @@
+#!/usr/bin/env python
+# -*- coding: utf-8; -*-
+
+# Copyright (c) 2021, 2024 Oracle and/or its affiliates.
+# Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
+
+
+import concurrent.futures
+import functools
+import logging
+from typing import Optional
+
+from git import Optional
+
+from ads.config import THREADED_DEFAULT_TIMEOUT
+
+logger = logging.getLogger(__name__)
+
+# Create a global thread pool with a maximum of 10 threads
+thread_pool = concurrent.futures.ThreadPoolExecutor(max_workers=10)
+
+
+class TimeoutError(Exception):
+    """
+    Custom exception to be raised when a function times out.
+
+    Attributes
+    ----------
+    message : str
+        The error message describing what went wrong.
+
+    Parameters
+    ----------
+    message : str
+        The error message.
+    """
+
+    def __init__(
+        self,
+        func_name: str,
+        timeout: int,
+        message: Optional[str] = "The operation could not be completed in time.",
+    ):
+        super().__init__(
+            f"{message} The function '{func_name}' exceeded the timeout of {timeout} seconds."
+        )
+
+
+def threaded(timeout: Optional[int] = THREADED_DEFAULT_TIMEOUT):
+    """
+    Decorator to run a function in a separate thread using a global thread pool.
+
+    Parameters
+    ----------
+    timeout (int, optional)
+        The maximum time in seconds to wait for the function to complete.
+        If the function does not complete within this time, "timeout" is returned.
+
+    Returns
+    -------
+    function: The wrapped function that will run in a separate thread with the specified timeout.
+    """
+
+    def decorator(func):
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            """
+            Wrapper function to submit the decorated function to the thread pool and handle timeout.
+
+            Parameters
+            ----------
+                *args: Positional arguments to pass to the decorated function.
+                **kwargs: Keyword arguments to pass to the decorated function.
+
+            Returns
+            -------
+            Any: The result of the decorated function if it completes within the timeout.
+
+            Raise
+            -----
+            TimeoutError
+                In case of the function exceeded the timeout.
+            """
+            future = thread_pool.submit(func, *args, **kwargs)
+            try:
+                return future.result(timeout=timeout)
+            except concurrent.futures.TimeoutError as ex:
+                logger.debug(
+                    f"The function '{func.__name__}' "
+                    f"exceeded the timeout of {timeout} seconds. "
+                    f"{ex}"
+                )
+                raise TimeoutError(func.__name__, timeout)
+
+        return wrapper
+
+    return decorator

ads/config.py

@@ -8,6 +8,7 @@
 import inspect
 import os
 from typing import Dict, Optional
+
 from ads.common.config import DEFAULT_CONFIG_PATH, DEFAULT_CONFIG_PROFILE, Config, Mode
 
 OCI_ODSC_SERVICE_ENDPOINT = os.environ.get("OCI_ODSC_SERVICE_ENDPOINT")
@@ -85,6 +86,8 @@
 AQUA_SERVICE_NAME = "aqua"
 DATA_SCIENCE_SERVICE_NAME = "data-science"
 
+THREADED_DEFAULT_TIMEOUT = os.environ.get("THREADED_DEFAULT_TIMEOUT", 5)
+
 
 def export(
     uri: Optional[str] = DEFAULT_CONFIG_PATH,