revert

sky/adaptors/aws.py

@@ -28,6 +28,7 @@
 
 # pylint: disable=import-outside-toplevel
 
+import functools
 import logging
 import threading
 import time
@@ -54,6 +55,29 @@
 _MAX_ATTEMPT_FOR_CREATION = 5
 
 
+class _ThreadLocalLRUCache(threading.local):
+
+    def __init__(self, maxsize=32):
+        super().__init__()
+        self.cache = functools.lru_cache(maxsize=maxsize)
+
+
+def _thread_local_lru_cache(maxsize=32):
+    # Create thread-local storage for the LRU cache
+    local_cache = _ThreadLocalLRUCache(maxsize)
+
+    def decorator(func):
+
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            # Use the thread-local LRU cache
+            return local_cache.cache(func)(*args, **kwargs)
+
+        return wrapper
+
+    return decorator
+
+
 def _assert_kwargs_builtin_type(kwargs):
     assert all(isinstance(v, (int, float, str)) for v in kwargs.values()), (
         f'kwargs should not contain none built-in types: {kwargs}')
@@ -95,7 +119,7 @@ def _create_aws_object(creation_fn_or_cls: Callable[[], Any],
 
 # The LRU cache needs to be thread-local to avoid multiple threads sharing the
 # same session object, which is not guaranteed to be thread-safe.
-@common.thread_local_lru_cache()
+@_thread_local_lru_cache()
 def session():
     """Create an AWS session."""
     return _create_aws_object(boto3.session.Session, 'session')

sky/adaptors/common.py

@@ -1,7 +1,6 @@
 """Lazy import for modules to avoid import error when not used."""
 import functools
 import importlib
-import threading
 from typing import Any, Optional, Tuple
 
 
@@ -64,26 +63,3 @@ def wrapper(*args, **kwargs):
         return wrapper
 
     return decorator
-
-
-class _ThreadLocalLRUCache(threading.local):
-
-    def __init__(self, maxsize=32):
-        super().__init__()
-        self.cache = functools.lru_cache(maxsize=maxsize)
-
-
-def thread_local_lru_cache(maxsize=32):
-    # Create thread-local storage for the LRU cache
-    local_cache = _ThreadLocalLRUCache(maxsize)
-
-    def decorator(func):
-
-        @functools.wraps(func)
-        def wrapper(*args, **kwargs):
-            # Use the thread-local LRU cache
-            return local_cache.cache(func)(*args, **kwargs)
-
-        return wrapper
-
-    return decorator

sky/adaptors/gcp.py

@@ -13,13 +13,7 @@
 _LAZY_MODULES = (google, googleapiclient)
 
 
-# The google-api-python-client library is built on top of the httplib2 library,
-# which is not thread-safe.
-# Reference: https://googleapis.github.io/google-api-python-client/docs/thread_safety.html # pylint: disable=line-too-long
-# We use a thread-local LRU cache to ensure that each thread has its own
-# httplib2.Http object.
 @common.load_lazy_modules(_LAZY_MODULES)
-@common.thread_local_lru_cache()
 def build(service_name: str, version: str, *args, **kwargs):
     """Build a GCP service.
 
@@ -34,6 +28,9 @@ def build(service_name: str, version: str, *args, **kwargs):
     if credentials is None:
         credentials, _ = google.auth.default()
 
+    # The google-api-python-client library is built on top of the httplib2
+    # library, which is not thread-safe.
+    # Reference: https://googleapis.github.io/google-api-python-client/docs/thread_safety.html # pylint: disable=line-too-long
     # Create a new Http() object for every request, to ensure that each thread
     # has its own httplib2.Http object for thread safety.
     def build_request(http, *args, **kwargs):