konduktor-nightly 0.1.0.dev20250409105017__py3-none-any.whl → 0.1.0.dev20250410104738__py3-none-any.whl

konduktor/__init__.py

@@ -14,7 +14,7 @@ __all__ = [
 ]
 
 # Replaced with the current commit when building the wheels.
-_KONDUKTOR_COMMIT_SHA = '0c99331021c3d7f5ff72801f64f7abe3675bde69'
+_KONDUKTOR_COMMIT_SHA = 'c7becbf4b64d992f56aad47b69770f3345a37297'
 os.makedirs(os.path.expanduser('~/.konduktor'), exist_ok=True)
 
 
@@ -48,5 +48,5 @@ def _get_git_commit():
 
 
 __commit__ = _get_git_commit()
-__version__ = '1.0.0.dev0.1.0.dev20250409105017'
+__version__ = '1.0.0.dev0.1.0.dev20250410104738'
 __root_dir__ = os.path.dirname(os.path.abspath(__file__))

konduktor/adaptors/aws.py

@@ -0,0 +1,221 @@
+# Proprietary Changes made for Trainy under the Trainy Software License
+# Original source: skypilot: https://github.com/skypilot-org/skypilot
+# which is Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""AWS cloud adaptors
+
+Thread safety notes:
+
+The results of session() is cached by each thread in a thread.local() storage.
+This means using their results is completely thread-safe.
+
+We do not cache the resource/client objects, because some credentials may be
+automatically rotated, but the cached resource/client object may not refresh the
+credential quick enough, which can cause unexpected NoCredentialsError. By
+creating the resource/client object from the thread-local session() object every
+time, the credentials will be explicitly refreshed.
+
+Calling session(), resource(), and client() is thread-safe, since they use a
+lock to protect each object's creation.
+
+
+This is informed by the following boto3 docs:
+- Unlike Resources and Sessions, clients are generally thread-safe.
+  https://boto3.amazonaws.com/v1/documentation/api/latest/guide/clients.html
+- Resource instances are not thread safe and should not be shared across
+  threads or processes
+  https://boto3.amazonaws.com/v1/documentation/api/latest/guide/resources.html
+- Similar to Resource objects, Session objects are not thread safe and
+  should not be shared across threads and processes.
+  https://boto3.amazonaws.com/v1/documentation/api/latest/guide/session.html
+"""
+
+# pylint: disable=import-outside-toplevel
+
+import functools
+import logging
+import threading
+import time
+from typing import Any, Callable
+
+from konduktor.adaptors import common
+from konduktor.utils import annotations, common_utils
+
+_IMPORT_ERROR_MESSAGE = (
+    'Failed to import dependencies for AWS. ' 'Try pip install konduktor-nightly[s3]'
+)
+boto3 = common.LazyImport('boto3', import_error_message=_IMPORT_ERROR_MESSAGE)
+botocore = common.LazyImport('botocore', import_error_message=_IMPORT_ERROR_MESSAGE)
+_LAZY_MODULES = (boto3, botocore)
+
+logger = logging.getLogger(__name__)
+_session_creation_lock = threading.RLock()
+
+version = 1
+
+# Retry 5 times by default for potential credential errors,
+_MAX_ATTEMPT_FOR_CREATION = 5
+
+
+class _ThreadLocalLRUCache(threading.local):
+    def __init__(self, maxsize=32):
+        super().__init__()
+        self.cache = annotations.lru_cache(scope='global', 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}'
+
+
+def _create_aws_object(creation_fn_or_cls: Callable[[], Any], object_name: str) -> Any:
+    """Create an AWS object.
+
+    Args:
+        creation_fn: The function to create the AWS object.
+
+    Returns:
+        The created AWS object.
+    """
+    attempt = 0
+    backoff = common_utils.Backoff()
+    while True:
+        try:
+            # Creating the boto3 objects are not thread-safe,
+            # so we add a reentrant lock to synchronize the session creation.
+            # Reference: https://github.com/boto/boto3/issues/1592
+
+            # NOTE: we need the lock here to avoid thread-safety issues when
+            # creating the resource, because Python module is a shared object,
+            # and we are not sure if the code inside 'session()' or
+            # 'session().xx()' is thread-safe.
+            with _session_creation_lock:
+                return creation_fn_or_cls()
+        except (
+            botocore_exceptions().CredentialRetrievalError,
+            botocore_exceptions().NoCredentialsError,
+        ) as e:
+            attempt += 1
+            if attempt >= _MAX_ATTEMPT_FOR_CREATION:
+                raise
+            time.sleep(backoff.current_backoff())
+            logger.info(
+                f'Retry creating AWS {object_name} due to '
+                f'{common_utils.format_exception(e)}.'
+            )
+
+
+# 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.
+@_thread_local_lru_cache()
+def session(check_credentials: bool = True):
+    """Create an AWS session."""
+    s = _create_aws_object(boto3.session.Session, 'session')
+    if check_credentials and s.get_credentials() is None:
+        # s.get_credentials() can be None if there are actually no credentials,
+        # or if we fail to get credentials from IMDS (e.g. due to throttling).
+        # Technically, it could be okay to have no credentials, as certain AWS
+        # APIs don't actually need them. But afaik everything we use AWS for
+        # needs credentials.
+        raise botocore_exceptions().NoCredentialsError()
+    return s
+
+
+# Avoid caching the resource/client objects. If we are using the assumed role,
+# the credentials will be automatically rotated, but the cached resource/client
+# object will only refresh the credentials with a fixed 15 minutes interval,
+# which can cause unexpected NoCredentialsError. By creating the resource/client
+# object every time, the credentials will be explicitly refreshed.
+# The creation of the resource/client is relatively fast (around 0.3s), so the
+# performance impact is negligible.
+# Reference: https://github.com/skypilot-org/skypilot/issues/2697
+def resource(service_name: str, **kwargs):
+    """Create an AWS resource of a certain service.
+
+    Args:
+        service_name: AWS resource name (e.g., 's3').
+        kwargs: Other options. We add max_attempts to the kwargs instead of
+            using botocore.config.Config() because the latter will generate
+            different keys even if the config is the same
+    """
+    _assert_kwargs_builtin_type(kwargs)
+
+    max_attempts = kwargs.pop('max_attempts', None)
+    if max_attempts is not None:
+        config = botocore_config().Config(retries={'max_attempts': max_attempts})
+        kwargs['config'] = config
+
+    check_credentials = kwargs.pop('check_credentials', True)
+
+    # Need to use the client retrieved from the per-thread session to avoid
+    # thread-safety issues (Directly creating the client with boto3.resource()
+    # is not thread-safe). Reference: https://stackoverflow.com/a/59635814
+    return _create_aws_object(
+        lambda: session(check_credentials=check_credentials).resource(
+            service_name, **kwargs
+        ),
+        'resource',
+    )
+
+
+def client(service_name: str, **kwargs):
+    """Create an AWS client of a certain service.
+
+    Args:
+        service_name: AWS service name (e.g., 's3', 'ec2').
+        kwargs: Other options.
+    """
+    _assert_kwargs_builtin_type(kwargs)
+
+    check_credentials = kwargs.pop('check_credentials', True)
+
+    # Need to use the client retrieved from the per-thread session to avoid
+    # thread-safety issues (Directly creating the client with boto3.client() is
+    # not thread-safe). Reference: https://stackoverflow.com/a/59635814
+
+    return _create_aws_object(
+        lambda: session(check_credentials=check_credentials).client(
+            service_name, **kwargs
+        ),
+        'client',
+    )
+
+
+@common.load_lazy_modules(modules=_LAZY_MODULES)
+def botocore_exceptions():
+    """AWS botocore exception."""
+    from botocore import exceptions
+
+    return exceptions
+
+
+@common.load_lazy_modules(modules=_LAZY_MODULES)
+def botocore_config():
+    """AWS botocore exception."""
+    from botocore import config
+
+    return config

konduktor/cli.py

@@ -773,10 +773,9 @@ def check(clouds: Tuple[str]):
 
     .. code-block:: bash
 
-      # Check credentials for all supported clouds.
-      konduktor check
-      # Check only specific clouds - GCP.
-      konduktor check gcp
+      # Check only specific clouds - gs, s3.
+      konduktor check gs
+      konduktor check s3
     """
     clouds_arg = clouds if len(clouds) > 0 else None
     konduktor_check.check(clouds=clouds_arg)

konduktor/data/aws/__init__.py

@@ -0,0 +1,15 @@
+"""Data sync between workstation <--> blob (s3, gcs, etc.) <--> worker pods"""
+
+from konduktor.data.aws.s3 import (
+    DEFAULT_AWS_CONFIG_PATH,
+    DEFAULT_AWS_CREDENTIAL_PATH,
+    S3CloudStorage,
+    S3Store,
+)
+
+__all__ = [
+    'S3Store',
+    'S3CloudStorage',
+    'DEFAULT_AWS_CREDENTIAL_PATH',
+    'DEFAULT_AWS_CONFIG_PATH',
+]