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

konduktor/data/data_utils.py

@@ -16,7 +16,7 @@ from multiprocessing import pool
 from typing import Any, Callable, Dict, List, Optional, Tuple
 
 from konduktor import logging
-from konduktor.adaptors import gcp
+from konduktor.adaptors import aws, gcp
 from konduktor.utils import exceptions, log_utils, ux_utils
 
 Client = Any
@@ -36,6 +36,30 @@ def split_gcs_path(gcs_path: str) -> Tuple[str, str]:
     return bucket, key
 
 
+def split_s3_path(s3_path: str) -> Tuple[str, str]:
+    """Splits S3 Path into Bucket name and Relative Path to Bucket
+
+    Args:
+      s3_path: str; S3 Path, e.g. s3://imagenet/train/
+    """
+    path_parts = s3_path.replace('s3://', '').split('/')
+    bucket = path_parts.pop(0)
+    key = '/'.join(path_parts)
+    return bucket, key
+
+
+def create_s3_client(region: Optional[str] = None) -> Client:
+    """Helper method that connects to Boto3 client for S3 Bucket
+
+    Args:
+      region: str; Region name, e.g. us-west-1, us-east-2. If None, default
+        region us-east-1 is used.
+    """
+    if region is None:
+        region = 'us-east-1'
+    return aws.client('s3', region_name=region)
+
+
 def verify_gcs_bucket(name: str) -> bool:
     """Helper method that checks if the GCS bucket exists
 
@@ -49,6 +73,27 @@ def verify_gcs_bucket(name: str) -> bool:
         return False
 
 
+def verify_s3_bucket(name: str) -> bool:
+    """Helper method that checks if the S3 bucket exists
+
+    Args:
+      name: str; Name of the S3 Bucket (without s3:// prefix)
+    """
+    import boto3
+    from botocore.exceptions import ClientError
+
+    s3 = boto3.client('s3')
+    try:
+        s3.head_bucket(Bucket=name)
+        return True
+    except ClientError as e:
+        error_code = int(e.response['Error']['Code'])
+        if error_code == 404:
+            return False
+        # In case of permissions issues or other errors, you can log and assume False
+        return False
+
+
 def is_cloud_store_url(url):
     result = urllib.parse.urlsplit(url)
     # '' means non-cloud URLs.

konduktor/data/registry.py

@@ -1,11 +1,12 @@
+from konduktor.data.aws import S3CloudStorage
 from konduktor.data.gcp import GcsCloudStorage
 
 # Maps bucket's URIs prefix(scheme) to its corresponding storage class
 
 _REGISTRY = {
     'gs': GcsCloudStorage(),
+    's3': S3CloudStorage(),
     # TODO(asaiacai): Add other cloud stores here
-    # 's3': S3CloudStorage(),
     # 'r2': R2CloudStorage(),
     # 'cos': IBMCosCloudStorage(),
     # 'oci': OciCloudStorage(),

konduktor/data/storage.py

@@ -28,7 +28,7 @@ import urllib.parse
 from typing import Any, Dict, List, Literal, Optional, Tuple, Type, Union
 
 from konduktor import check, config, logging
-from konduktor.data import constants, data_utils, gcp, registry, storage_utils
+from konduktor.data import aws, constants, data_utils, gcp, registry, storage_utils
 from konduktor.utils import annotations, common_utils, exceptions, schemas, ux_utils
 
 logger = logging.get_logger(__file__)
@@ -79,12 +79,15 @@ class StoreType(enum.Enum):
     """Enum for the different types of stores."""
 
     GCS = 'GCS'
+    S3 = 'S3'
 
     @classmethod
     def from_cloud(cls, cloud: str) -> 'StoreType':
         # these need to match the cloud store classes in konduktor/cloud_stores.py
         if cloud.lower() == 'gs':
             return StoreType.GCS
+        elif cloud.lower() == 's3':
+            return StoreType.S3
         else:
             with ux_utils.print_exception_no_traceback():
                 raise ValueError(f'Unknown cloud: {cloud}')
@@ -93,6 +96,8 @@ class StoreType(enum.Enum):
     def from_store(cls, store: 'storage_utils.AbstractStore') -> 'StoreType':
         if store.__repr__() == 'GcsStore':
             return StoreType.GCS
+        elif store.__repr__() == 'S3Store':
+            return StoreType.S3
         else:
             with ux_utils.print_exception_no_traceback():
                 raise ValueError(f'Unknown store type: {store}')
@@ -100,6 +105,8 @@ class StoreType(enum.Enum):
     def store_prefix(self) -> str:
         if self == StoreType.GCS:
             return 'gs://'
+        elif self == StoreType.S3:
+            return 's3://'
         else:
             with ux_utils.print_exception_no_traceback():
                 raise ValueError(f'Unknown store type: {self}')
@@ -123,6 +130,8 @@ class StoreType(enum.Enum):
             if store_url.startswith(store_type.store_prefix()):
                 if store_type == StoreType.GCS:
                     bucket_name, sub_path = data_utils.split_gcs_path(store_url)
+                elif store_type == StoreType.S3:
+                    bucket_name, sub_path = data_utils.split_s3_path(store_url)
                 return store_type, bucket_name, sub_path, storage_account_name, region
         raise ValueError(f'Unknown store URL: {store_url}')
 
@@ -143,7 +152,7 @@ class StoreType(enum.Enum):
 
 
 # this should match the above StoreType enum
-STORE_TYPES = Literal[StoreType.GCS]
+STORE_TYPES = Literal[StoreType.GCS, StoreType.S3]
 
 
 class Storage(object):
@@ -310,6 +319,8 @@ class Storage(object):
             if isinstance(self.source, str):
                 if self.source.startswith('gs://'):
                     self.add_store(StoreType.GCS)
+                elif self.source.startswith('s3://'):
+                    self.add_store(StoreType.S3)
 
     @staticmethod
     def _validate_source(
@@ -533,12 +544,13 @@ class Storage(object):
             # When initializing from global_user_state, we override the
             # source from the YAML
             try:
-                # if s_type == StoreType.S3:
-                #     store = S3Store.from_metadata(
-                #         s_metadata,
-                #         source=self.source,
-                #         sync_on_reconstruction=self.sync_on_reconstruction)
-                if s_type == StoreType.GCS:
+                if s_type == StoreType.S3:
+                    store = aws.S3Store.from_metadata(
+                        s_metadata,
+                        source=self.source,
+                        sync_on_reconstruction=self.sync_on_reconstruction,
+                    )
+                elif s_type == StoreType.GCS:
                     store = gcp.GcsStore.from_metadata(
                         s_metadata,
                         source=self.source,
@@ -622,6 +634,8 @@ class Storage(object):
         store_cls: Type['storage_utils.AbstractStore']
         if store_type == StoreType.GCS:
             store_cls = gcp.GcsStore
+        elif store_type == StoreType.S3:
+            store_cls = aws.S3Store
         else:
             with ux_utils.print_exception_no_traceback():
                 raise exceptions.StorageSpecError(

konduktor/task.py

@@ -653,6 +653,19 @@ class Task:
             if storage.mode == storage_lib.StorageMode.COPY:
                 store_type = storage_plans[storage]
             # TODO(asaiacai): add back other stores here
+            elif store_type is storage_lib.StoreType.S3:
+                if isinstance(storage.source, str) and storage.source.startswith(
+                    's3://'
+                ):
+                    blob_path = storage.source
+                else:
+                    assert storage.name is not None, storage
+                    blob_path = 's3://' + storage.name
+                self.update_file_mounts(
+                    {
+                        mnt_path: blob_path,
+                    }
+                )
             elif store_type is storage_lib.StoreType.GCS:
                 if isinstance(storage.source, str) and storage.source.startswith(
                     'gs://'

konduktor/templates/pod.yaml.j2

@@ -79,7 +79,7 @@ kubernetes:
               {% if 'curl' in run_cmd %}
               PACKAGES="$PACKAGES curl";
               {% endif %}
-              {% if 'gs' in mount_secrets %}
+              {% if 'gs' in mount_secrets or 's3' in mount_secrets %}
               PACKAGES="$PACKAGES unzip wget";
               {% endif %}
               {% if 'git' in run_cmd %}
@@ -129,6 +129,10 @@ kubernetes:
               $(prefix_cmd) echo "Unpacking GCP secret"
               $(prefix_cmd) mkdir -p ~/.config
               $(prefix_cmd) unzip /run/konduktor/gs-secret/gcpcredentials -d ~/.config/gcloud
+              {% elif secret_type == "s3" %}
+              $(prefix_cmd) echo "Unpacking AWS secret"
+              $(prefix_cmd) mkdir -p ~/.aws
+              $(prefix_cmd) unzip /run/konduktor/s3-secret/awscredentials -d ~/.aws
               {% endif %}
               {% endfor %}
               end_epoch=$(date +%s);

konduktor/utils/common_utils.py

@@ -17,6 +17,7 @@ import getpass
 import hashlib
 import inspect
 import os
+import random
 import re
 import socket
 import sys
@@ -391,3 +392,31 @@ def format_exception(
     if use_bracket:
         return f'[{class_fullname(e.__class__)}] {e}'
     return f'{class_fullname(e.__class__)}: {e}'
+
+
+class Backoff:
+    """Exponential backoff with jittering."""
+
+    MULTIPLIER = 1.6
+    JITTER = 0.4
+
+    def __init__(self, initial_backoff: float = 5, max_backoff_factor: int = 5):
+        self._initial = True
+        self._backoff = 0.0
+        self._initial_backoff = initial_backoff
+        self._max_backoff = max_backoff_factor * self._initial_backoff
+
+    # https://github.com/grpc/grpc/blob/2d4f3c56001cd1e1f85734b2f7c5ce5f2797c38a/doc/connection-backoff.md
+    # https://github.com/grpc/grpc/blob/5fc3ff82032d0ebc4bf252a170ebe66aacf9ed9d/src/core/lib/backoff/backoff.cc
+
+    def current_backoff(self) -> float:
+        """Backs off once and returns the current backoff in seconds."""
+        if self._initial:
+            self._initial = False
+            self._backoff = min(self._initial_backoff, self._max_backoff)
+        else:
+            self._backoff = min(self._backoff * self.MULTIPLIER, self._max_backoff)
+        self._backoff += random.uniform(
+            -self.JITTER * self._backoff, self.JITTER * self._backoff
+        )
+        return self._backoff

{konduktor_nightly-0.1.0.dev20250409105017.dist-info → konduktor_nightly-0.1.0.dev20250410104738.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: konduktor-nightly
-Version: 0.1.0.dev20250409105017
+Version: 0.1.0.dev20250410104738
 Summary: GPU Cluster Health Management
 Author: Andrew Aikawa
 Author-email: asai@berkeley.edu
@@ -10,6 +10,10 @@ Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
+Provides-Extra: s3
+Requires-Dist: awscli (>=1.32.84,<2.0.0) ; extra == "s3"
+Requires-Dist: boto3 (>=1.34.84,<2.0.0) ; extra == "s3"
+Requires-Dist: botocore (>=1.34.84,<2.0.0) ; extra == "s3"
 Requires-Dist: click (>=8.1.7,<9.0.0)
 Requires-Dist: colorama (>=0.4.6,<0.5.0)
 Requires-Dist: filelock (>=3.18.0,<4.0.0)

{konduktor_nightly-0.1.0.dev20250409105017.dist-info → konduktor_nightly-0.1.0.dev20250410104738.dist-info}/RECORD

@@ -1,5 +1,6 @@
-konduktor/__init__.py,sha256=Ka1B4zUTmoP9pKaZEEQhTtqNHimV85pb9tDNBCwCQi0,1540
+konduktor/__init__.py,sha256=GTCKY5x_jINcGcz-nfNIu_QGaWttvSpxmiyltgOSLfs,1540
 konduktor/adaptors/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+konduktor/adaptors/aws.py,sha256=s47Ra-GaqCQibzVfmD0pmwEWHif1EGO5opMbwkLxTCU,8244
 konduktor/adaptors/common.py,sha256=mYb_6c3u5MghtiFfiW5OO-EH6t7cIR5npbkgUmz6FYE,3517
 konduktor/adaptors/gcp.py,sha256=pOQA2q8fFyr97Htn8EqvNM0XT-Ao8UwvExviiLaDats,4746
 konduktor/backends/__init__.py,sha256=1Q6sqqdeMYarpTX_U-QVywJYf7idiUTRsyP-E4BQSOw,129
@@ -7,7 +8,7 @@ konduktor/backends/backend.py,sha256=qh0bp94lzoTYZkzyQv2-CVrB5l91FkG2vclXg24UFC0
 konduktor/backends/jobset.py,sha256=lh_PihQgM0tmVryCpjSsZjWug8hBnJr7ua9lqk0qEAM,8251
 konduktor/backends/jobset_utils.py,sha256=NIwTvJdGhbDnXEceabiuUm9aHZ29LK3jVHfQzutB_ec,17297
 konduktor/check.py,sha256=JennyWoaqSKhdyfUldd266KwVXTPJpcYQa4EED4a_BA,7569
-konduktor/cli.py,sha256=90bnh3nIobfBkzqS_SXgw9Z8Zqh4ouwpLDj0kx_6kL8,23562
+konduktor/cli.py,sha256=Ii9-2mrc-1f2ksLasA-xRb-JnEi_9ZeCXZ3lJ1GG8H8,23515
 konduktor/config.py,sha256=J50JxC6MsXMnlrJPXdDUMr38C89xvOO7mR8KJ6fyils,15520
 konduktor/constants.py,sha256=T3AeXXxuQHINW_bAWyztvDeS8r4g8kXBGIwIq13cys0,1814
 konduktor/controller/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -48,14 +49,16 @@ konduktor/dashboard/frontend/postcss.config.mjs,sha256=rDHiqV72T-J860Ek4QFnUnMQe
 konduktor/dashboard/frontend/server.js,sha256=jcp6_Ww9YJD3uKY07jR3KMlAM6n1QZdxZnVY6Kh-J6A,1789
 konduktor/dashboard/frontend/tailwind.config.js,sha256=fCnc48wvioIDOe5ldQ_6RE7F76cP7aU7pDrxBPJx-Fk,366
 konduktor/data/__init__.py,sha256=KMR2i3E9YcIpiIuCxtRdS7BQ1w2vUAbbve7agziJrLo,213
+konduktor/data/aws/__init__.py,sha256=_6zWfNNAK1QGgyKqg_yPYWcXlnffchyvIMErYa6tw_U,331
+konduktor/data/aws/s3.py,sha256=E6JdT5K0nhy0ExQQN6M5vjemDSfBi6Cc3PCgjyCMJoQ,47777
 konduktor/data/constants.py,sha256=yXVEoTI2we1xOjVSU-bjRCQCLpVvpEvJ0GedXvSwEfw,127
-konduktor/data/data_utils.py,sha256=aIv3q2H1GSiN2w8WNjZgVaglm-hoiHSb4KR-MAiKKXs,8383
+konduktor/data/data_utils.py,sha256=yrnu8_cY63TXqfWfFG3yqY2w_tE9UQK9jIQAFQCDVg0,9668
 konduktor/data/gcp/__init__.py,sha256=rlQxACBC_Vu36mdgPyJgUy4mGc_6Nt_a96JAuaPz2pQ,489
 konduktor/data/gcp/constants.py,sha256=dMfOiFccM8O6rUi9kClJcbvw1K1VnS1JzzQk3apq8ho,1483
 konduktor/data/gcp/gcs.py,sha256=kDbUzf8ALYzsw_G3sBRn_enQ8fjI-UKV0jeWuFZiULA,42018
 konduktor/data/gcp/utils.py,sha256=FJQcMXZqtMIzjZ98b3lTTc0UbdPUKTDLsOsfJaaH5-s,214
-konduktor/data/registry.py,sha256=eLs8Wr5ugwOfXGPtg1utTvGIqdbVLsCf-a3PFS1NELc,652
-konduktor/data/storage.py,sha256=sm0ZfGZUZRiChza_jMRQY1xDIWtZpFQqwPuVOF8PM_Y,34742
+konduktor/data/registry.py,sha256=CUbMsN_Q17Pf4wRHkqZrycErEjTP7cLEdgcfwVGcEpc,696
+konduktor/data/storage.py,sha256=SDKRWDd7PCT9ytuz4cH0CejZj5QmWG_EZhUMVoTzWsc,35308
 konduktor/data/storage_utils.py,sha256=n4GivkN0KMqmyOTDznF0Z-hzsJvm7KCEh5i5HgFAT-4,20806
 konduktor/execution.py,sha256=UaHUdBmDaIYgiAXkRKJQOHniYPVIR4sr4yUbIqpgMrQ,18401
 konduktor/kube_client.py,sha256=aqwjDfNSneB5NOxV6CtqhkBeNl0UQNUt730R3ujG9Ow,6156
@@ -65,16 +68,16 @@ konduktor/manifests/dashboard_deployment.yaml,sha256=xJLd4FbPMAosI0fIv5_8y7dV9bw
 konduktor/manifests/dmesg_daemonset.yaml,sha256=pSWt7YOeTYjS0l0iki1fvHOs7MhY-sH-RQfVW6JJyno,1391
 konduktor/manifests/pod_cleanup_controller.yaml,sha256=hziL1Ka1kCAEL9R7Tjvpb80iw1vcq9_3gwHCu75Bi0A,3939
 konduktor/resource.py,sha256=68z8gC8Ivqktwv0R6ylMn9ZNocgkcRT0yIRGGKOdwcM,18491
-konduktor/task.py,sha256=edHgMLYECGux6WLCilqsNZNYr3dEcw_miWvu4FYpu5U,34713
+konduktor/task.py,sha256=Vu1TzYtLvSBz-HyHY2gsM2cMcUhMNQu44L3CWmYRXKE,35232
 konduktor/templates/jobset.yaml.j2,sha256=onYiHtXAgk-XBtji994hPu_g0hxnLzvmfxwjbdKdeZc,960
-konduktor/templates/pod.yaml.j2,sha256=24g0onMUl4BDhPO2ftTdZB5l_wvMjwCuDQ_fTdwBxD8,8474
+konduktor/templates/pod.yaml.j2,sha256=q_kFyuYP-RlHocQUV8Vsnh-hYurWO7Y_pwl0gytGMl4,8728
 konduktor/usage/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 konduktor/usage/constants.py,sha256=gCL8afIHZhO0dcxbJGpESE9sCC1cBSbeRnQ8GwNOY4M,612
 konduktor/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 konduktor/utils/accelerator_registry.py,sha256=lzUqkod1U3hDo8ETFU0CBXESXhPNZI8e-CJt4Oqf3Ko,550
 konduktor/utils/annotations.py,sha256=oy2-BLydkFt3KWkXDuaGY84d6b7iISuy4eAT9uXk0Fc,2225
 konduktor/utils/base64_utils.py,sha256=mF-Tw98mFRG70YE4w6s9feuQSCYZHOb8YatBZwMugyI,3130
-konduktor/utils/common_utils.py,sha256=1_j-nRikKmTnB8BFE0xQb7LquKVAOLaJnsy4LxZlNbI,13869
+konduktor/utils/common_utils.py,sha256=F5x7k4AdBB44u8PYRkaugORnZKnK3JLqGn1jHOKgUYo,14960
 konduktor/utils/constants.py,sha256=1DneiTR21lvKUcWdBGwC4I4fD4uPjbjLUilEnJS7rzA,216
 konduktor/utils/env_options.py,sha256=T41Slzf4Mzl-n45CGXXqdy2fCrYhPNZQ7RP5vmnN4xc,2258
 konduktor/utils/exceptions.py,sha256=GBOFIkk9nikqWGR0FXGXOWVVImoH7nWnMl_L3Oux3fo,6581
@@ -87,8 +90,8 @@ konduktor/utils/schemas.py,sha256=Gv7SEhFpv-eO5izqRz8d-eQ9z-lVmY05akm6HEXIIdc,17
 konduktor/utils/subprocess_utils.py,sha256=WoFkoFhGecPR8-rF8WJxbIe-YtV94LXz9UG64SDhCY4,9448
 konduktor/utils/ux_utils.py,sha256=NPNu3Igu2Z9Oq77ghJhy_fIxQZTXWr9BtKyxN3Wslzo,7164
 konduktor/utils/validator.py,sha256=tgBghVyedyzGx84-U2Qfoh_cJBE3oUk9gclMW90ORks,691
-konduktor_nightly-0.1.0.dev20250409105017.dist-info/LICENSE,sha256=MuuqTZbHvmqXR_aNKAXzggdV45ANd3wQ5YI7tnpZhm0,6586
-konduktor_nightly-0.1.0.dev20250409105017.dist-info/METADATA,sha256=TQwygSbHFhQFKOiGfeW3g3Eyc5YO9mrqQoamc9TtQro,4112
-konduktor_nightly-0.1.0.dev20250409105017.dist-info/WHEEL,sha256=FMvqSimYX_P7y0a7UY-_Mc83r5zkBZsCYPm7Lr0Bsq4,88
-konduktor_nightly-0.1.0.dev20250409105017.dist-info/entry_points.txt,sha256=k3nG5wDFIJhNqsZWrHk4d0irIB2Ns9s47cjRWYsTCT8,48
-konduktor_nightly-0.1.0.dev20250409105017.dist-info/RECORD,,
+konduktor_nightly-0.1.0.dev20250410104738.dist-info/LICENSE,sha256=MuuqTZbHvmqXR_aNKAXzggdV45ANd3wQ5YI7tnpZhm0,6586
+konduktor_nightly-0.1.0.dev20250410104738.dist-info/METADATA,sha256=_EvP828hXwzUqFcjDyh5t-9HVMnsMqtp8UPDdxN2kSI,4303
+konduktor_nightly-0.1.0.dev20250410104738.dist-info/WHEEL,sha256=FMvqSimYX_P7y0a7UY-_Mc83r5zkBZsCYPm7Lr0Bsq4,88
+konduktor_nightly-0.1.0.dev20250410104738.dist-info/entry_points.txt,sha256=k3nG5wDFIJhNqsZWrHk4d0irIB2Ns9s47cjRWYsTCT8,48
+konduktor_nightly-0.1.0.dev20250410104738.dist-info/RECORD,,