From fac365f5885fb6f697bb70a8649d400238cee68c Mon Sep 17 00:00:00 2001
From: Dmitrii Cherkasov <dmitrii.cherkasov@oracle.com>
Date: Mon, 10 Jun 2024 16:14:17 -0700
Subject: [PATCH] Fixes an import issue.

---
 ads/aqua/common/utils.py         | 11 +++--------
 ads/common/decorator/threaded.py |  4 ++--
 2 files changed, 5 insertions(+), 10 deletions(-)

diff --git a/ads/aqua/common/utils.py b/ads/aqua/common/utils.py
index 933b5a9a5..369b62aae 100644
--- a/ads/aqua/common/utils.py
+++ b/ads/aqua/common/utils.py
@@ -28,17 +28,12 @@
 from ads.aqua.constants import *
 from ads.aqua.data import AquaResourceIdentifier
 from ads.common.auth import default_signer
-from ads.common.decorator import threaded
+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 copy_file, get_console_link, upload_to_os
-from ads.config import (
-    AQUA_SERVICE_MODELS_BUCKET,
-    CONDA_BUCKET_NS,
-    TENANCY_OCID,
-    THREADED_DEFAULT_TIMEOUT,
-)
+from ads.config import AQUA_SERVICE_MODELS_BUCKET, CONDA_BUCKET_NS, TENANCY_OCID
 from ads.model import DataScienceModel, ModelVersionSet
 
 logger = logging.getLogger("ads.aqua")
@@ -201,7 +196,7 @@ def read_file(file_path: str, **kwargs) -> str:
         return UNKNOWN
 
 
-@threaded(timeout=THREADED_DEFAULT_TIMEOUT)
+@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://"):
diff --git a/ads/common/decorator/threaded.py b/ads/common/decorator/threaded.py
index 2b5c9d0de..7c77ffcda 100644
--- a/ads/common/decorator/threaded.py
+++ b/ads/common/decorator/threaded.py
@@ -46,7 +46,7 @@ def __init__(
         )
 
 
-def threaded(timeout=None):
+def threaded(timeout: Optional[int] = THREADED_DEFAULT_TIMEOUT):
     """
     Decorator to run a function in a separate thread using a global thread pool.
 
@@ -83,7 +83,7 @@ def wrapper(*args, **kwargs):
             """
             future = thread_pool.submit(func, *args, **kwargs)
             try:
-                return future.result(timeout=timeout or THREADED_DEFAULT_TIMEOUT)
+                return future.result(timeout=timeout)
             except concurrent.futures.TimeoutError as ex:
                 logger.debug(
                     f"The function '{func.__name__}' "