konduktor-nightly 0.1.0.dev20251128104812__py3-none-any.whl

konduktor/__init__.py

@@ -0,0 +1,49 @@
+"""The Konduktor package."""
+
+import os
+import subprocess
+
+from konduktor.execution import launch
+from konduktor.resource import Resources
+from konduktor.serving import Serving
+from konduktor.task import Task
+
+__all__ = ['launch', 'Resources', 'Task', 'Serving']
+
+# Replaced with the current commit when building the wheels.
+_KONDUKTOR_COMMIT_SHA = '2c7cb302ac6c8f558befdf1ad323e30c47f1ba71'
+os.makedirs(os.path.expanduser('~/.konduktor'), exist_ok=True)
+
+
+def _get_git_commit():
+    if 'KONDUKTOR_COMMIT_SHA' not in _KONDUKTOR_COMMIT_SHA:
+        # This is a release build, so we don't need to get the commit hash from
+        # git, as it's already been set.
+        return _KONDUKTOR_COMMIT_SHA
+
+    # This is a development build (pip install -e .), so we need to get the
+    # commit hash from git.
+    try:
+        cwd = os.path.dirname(__file__)
+        commit_hash = subprocess.check_output(
+            ['git', 'rev-parse', 'HEAD'],
+            cwd=cwd,
+            universal_newlines=True,
+            stderr=subprocess.DEVNULL,
+        ).strip()
+        changes = subprocess.check_output(
+            ['git', 'status', '--porcelain'],
+            cwd=cwd,
+            universal_newlines=True,
+            stderr=subprocess.DEVNULL,
+        ).strip()
+        if changes:
+            commit_hash += '-dirty'
+        return commit_hash
+    except Exception:  # pylint: disable=broad-except
+        return _KONDUKTOR_COMMIT_SHA
+
+
+__commit__ = _get_git_commit()
+__version__ = '1.0.0.dev0.1.0.dev20251128104812'
+__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/adaptors/common.py

@@ -0,0 +1,118 @@
+# 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.
+
+"""Lazy import for modules to avoid import error when not used."""
+
+import functools
+import importlib
+import os
+import threading
+from typing import Any, Callable, Optional, Tuple
+
+import filelock
+
+
+class LazyImport:
+    """Lazy importer for heavy modules or cloud modules only when enabled.
+
+    We use this for pandas and networkx, as they can be time-consuming to import
+    (0.1-0.2 seconds). With this class, we can avoid the unnecessary import time
+    when the module is not used.
+
+    We also use this for cloud adaptors, because we do not want to import the
+    cloud dependencies when it is not enabled.
+    """
+
+    def __init__(
+        self,
+        module_name: str,
+        import_error_message: Optional[str] = None,
+        set_loggers: Optional[Callable] = None,
+    ):
+        self._module_name = module_name
+        self._module = None
+        self._import_error_message = import_error_message
+        self._set_loggers = set_loggers
+        self._lock = threading.RLock()
+
+    def load_module(self):
+        # Avoid extra imports when multiple threads try to import the same
+        # module. The overhead is minor since import can only run in serial
+        # due to GIL even in multi-threaded environments.
+        with self._lock:
+            if self._module is None:
+                try:
+                    self._module = importlib.import_module(self._module_name)
+                    if self._set_loggers is not None:
+                        self._set_loggers()
+                except ImportError as e:
+                    if self._import_error_message is not None:
+                        raise ImportError(self._import_error_message) from e
+                    raise
+        return self._module
+
+    def __getattr__(self, name: str) -> Any:
+        # Attempt to access the attribute, if it fails, assume it's a submodule
+        # and lazily import it
+        try:
+            if name in self.__dict__:
+                return self.__dict__[name]
+            return getattr(self.load_module(), name)
+        except AttributeError:
+            # Dynamically create a new LazyImport instance for the submodule
+            submodule_name = f'{self._module_name}.{name}'
+            lazy_submodule = LazyImport(submodule_name, self._import_error_message)
+            setattr(self, name, lazy_submodule)
+            return lazy_submodule
+
+
+def load_lazy_modules(modules: Tuple[LazyImport, ...]):
+    """Load lazy modules before entering a function to error out quickly."""
+
+    def decorator(func):
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            for m in modules:
+                m.load_module()
+            return func(*args, **kwargs)
+
+        return wrapper
+
+    return decorator
+
+
+class LockedClientProxy:
+    """Proxy for GCP client that locks access to the client."""
+
+    def __init__(
+        self,
+        client,
+        lock_path=os.path.expanduser('~/.konduktor/gcs_storage.lock'),
+        timeout=10,
+    ):
+        self._client = client
+        self._lock = filelock.FileLock(lock_path, timeout=timeout)
+
+    def __getattr__(self, attr):
+        target = getattr(self._client, attr)
+
+        if callable(target):
+
+            @functools.wraps(target)
+            def locked_method(*args, **kwargs):
+                with self._lock:
+                    return target(*args, **kwargs)
+
+            return locked_method
+        else:
+            # Attribute (not method) access just passes through
+            return target

konduktor/adaptors/gcp.py

@@ -0,0 +1,126 @@
+# 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.
+
+"""GCP cloud adaptors"""
+
+import json
+import os
+
+from konduktor.adaptors import common
+
+_IMPORT_ERROR_MESSAGE = (
+    'Failed to import dependencies for GCP. ' 'Try pip install "konduktor[gcp]"'
+)
+googleapiclient = common.LazyImport(
+    'googleapiclient', import_error_message=_IMPORT_ERROR_MESSAGE
+)
+google = common.LazyImport('google', import_error_message=_IMPORT_ERROR_MESSAGE)
+_LAZY_MODULES = (google, googleapiclient)
+
+_LOCK_PATH = '~/.konduktor/gcs_storage.lock'
+
+
+@common.load_lazy_modules(_LAZY_MODULES)
+def build(service_name: str, version: str, *args, **kwargs):
+    """Build a GCP service.
+
+    Args:
+        service_name: GCP service name (e.g., 'compute', 'storagetransfer').
+        version: Service version (e.g., 'v1').
+    """
+
+    return googleapiclient.discovery.build(service_name, version, *args, **kwargs)
+
+
+@common.load_lazy_modules(_LAZY_MODULES)
+def storage_client():
+    """Helper that connects to GCS Storage Client for GCS Bucket"""
+    from google.cloud import storage
+
+    return common.LockedClientProxy(
+        storage.Client(), lock_path=os.path.expanduser(_LOCK_PATH)
+    )
+
+
+@common.load_lazy_modules(_LAZY_MODULES)
+def anonymous_storage_client():
+    """Helper that connects to GCS Storage Client for Public GCS Buckets"""
+    from google.cloud import storage
+
+    return common.LockedClientProxy(
+        storage.Client(), lock_path=os.path.expanduser(_LOCK_PATH)
+    )
+
+
+@common.load_lazy_modules(_LAZY_MODULES)
+def not_found_exception():
+    """NotFound exception."""
+    from google.api_core import exceptions as gcs_exceptions
+
+    return gcs_exceptions.NotFound
+
+
+@common.load_lazy_modules(_LAZY_MODULES)
+def forbidden_exception():
+    """Forbidden exception."""
+    from google.api_core import exceptions as gcs_exceptions
+
+    return gcs_exceptions.Forbidden
+
+
+@common.load_lazy_modules(_LAZY_MODULES)
+def conflict_exception():
+    """Conflict exception."""
+    from google.api_core import exceptions as gcs_exceptions
+
+    return gcs_exceptions.Conflict
+
+
+@common.load_lazy_modules(_LAZY_MODULES)
+def http_error_exception():
+    """HttpError exception."""
+    from googleapiclient import errors
+
+    return errors.HttpError
+
+
+@common.load_lazy_modules(_LAZY_MODULES)
+def credential_error_exception():
+    """CredentialError exception."""
+    from google.auth import exceptions
+
+    return exceptions.DefaultCredentialsError
+
+
+@common.load_lazy_modules(_LAZY_MODULES)
+def get_credentials(cred_type: str, credentials_field: str):
+    """Get GCP credentials."""
+    from google.oauth2 import service_account
+    from google.oauth2.credentials import Credentials as OAuthCredentials
+
+    if cred_type == 'service_account':
+        # If parsing the gcp_credentials failed, then the user likely made a
+        # mistake in copying the credentials into the config yaml.
+        try:
+            service_account_info = json.loads(credentials_field)
+        except json.decoder.JSONDecodeError as e:
+            raise RuntimeError(
+                'gcp_credentials found in cluster yaml file but '
+                'formatted improperly.'
+            ) from e
+        credentials = service_account.Credentials.from_service_account_info(
+            service_account_info
+        )
+    elif cred_type == 'credentials_token':
+        # Otherwise the credentials type must be credentials_token.
+        credentials = OAuthCredentials(credentials_field)
+    return credentials

konduktor/authentication.py

@@ -0,0 +1,124 @@
+# 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.
+
+"""
+The local machine's public key should not be uploaded to the remote VM, because
+it will cause private/public key pair mismatch when the user tries to launch new
+VM from that remote VM using SkyPilot, e.g., the node is used as a jobs
+controller. (Lambda cloud is an exception, due to the limitation of the cloud
+provider. See the comments in setup_lambda_authentication)
+"""
+
+import functools
+import os
+from typing import Tuple
+
+import filelock
+
+from konduktor import logging
+from konduktor.utils import common_utils
+
+logger = logging.get_logger(__name__)
+
+_SSH_KEY_PATH_PREFIX = '~/.konduktor/clients/{user_hash}/ssh'
+
+MAX_TRIALS = 64
+
+
+def get_ssh_key_and_lock_path() -> Tuple[str, str, str]:
+    user_hash = common_utils.get_user_hash()
+    user_ssh_key_prefix = _SSH_KEY_PATH_PREFIX.format(user_hash=user_hash)
+    os.makedirs(os.path.expanduser(user_ssh_key_prefix), exist_ok=True, mode=0o700)
+    private_key_path = os.path.join(user_ssh_key_prefix, 'konduktor-key')
+    public_key_path = os.path.join(user_ssh_key_prefix, 'konduktor-key.pub')
+    lock_path = os.path.join(user_ssh_key_prefix, '.__internal-konduktor-key.lock')
+    return private_key_path, public_key_path, lock_path
+
+
+def _generate_rsa_key_pair() -> Tuple[str, str]:
+    # Keep the import of the cryptography local to avoid expensive
+    # third-party imports when not needed.
+    # pylint: disable=import-outside-toplevel
+    from cryptography.hazmat.backends import default_backend
+    from cryptography.hazmat.primitives import serialization
+    from cryptography.hazmat.primitives.asymmetric import rsa
+
+    key = rsa.generate_private_key(
+        backend=default_backend(), public_exponent=65537, key_size=2048
+    )
+
+    private_key = (
+        key.private_bytes(
+            encoding=serialization.Encoding.PEM,
+            format=serialization.PrivateFormat.TraditionalOpenSSL,
+            encryption_algorithm=serialization.NoEncryption(),
+        )
+        .decode('utf-8')
+        .strip()
+    )
+
+    public_key = (
+        key.public_key()
+        .public_bytes(
+            serialization.Encoding.OpenSSH, serialization.PublicFormat.OpenSSH
+        )
+        .decode('utf-8')
+        .strip()
+    )
+
+    return public_key, private_key
+
+
+def _save_key_pair(
+    private_key_path: str, public_key_path: str, private_key: str, public_key: str
+) -> None:
+    key_dir = os.path.dirname(private_key_path)
+    os.makedirs(key_dir, exist_ok=True, mode=0o700)
+
+    with open(
+        private_key_path,
+        'w',
+        encoding='utf-8',
+        opener=functools.partial(os.open, mode=0o600),
+    ) as f:
+        f.write(private_key)
+
+    with open(
+        public_key_path,
+        'w',
+        encoding='utf-8',
+        opener=functools.partial(os.open, mode=0o644),
+    ) as f:
+        f.write(public_key)
+
+
+def get_or_generate_keys() -> Tuple[str, str]:
+    """Returns the aboslute private and public key paths."""
+    private_key_path, public_key_path, lock_path = get_ssh_key_and_lock_path()
+    private_key_path = os.path.expanduser(private_key_path)
+    public_key_path = os.path.expanduser(public_key_path)
+    lock_path = os.path.expanduser(lock_path)
+
+    lock_dir = os.path.dirname(lock_path)
+    # We should have the folder ~/.konduktor/generated/ssh to have 0o700 permission,
+    # as the ssh configs will be written to this folder as well in
+    # backend_utils.SSHConfigHelper
+    os.makedirs(lock_dir, exist_ok=True, mode=0o700)
+    with filelock.FileLock(lock_path, timeout=10):
+        if not os.path.exists(private_key_path):
+            public_key, private_key = _generate_rsa_key_pair()
+            _save_key_pair(private_key_path, public_key_path, private_key, public_key)
+    assert os.path.exists(public_key_path), (
+        'Private key found, but associated public key '
+        f'{public_key_path} does not exist.'
+    )
+    return private_key_path, public_key_path

konduktor/backends/__init__.py

@@ -0,0 +1,6 @@
+"""Batch job backends"""
+
+from konduktor.backends.deployment import DeploymentBackend
+from konduktor.backends.jobset import JobsetBackend
+
+__all__ = ['Backend', 'JobsetBackend', 'DeploymentBackend']