konduktor-nightly 0.1.0.dev20250209104336__py3-none-any.whl → 0.1.0.dev20250313070642__py3-none-any.whl

konduktor/__init__.py

@@ -3,12 +3,22 @@
 import os
 import subprocess
 
+from konduktor.execution import launch
+from konduktor.resource import Resources
+from konduktor.task import Task
+
+__all__ = [
+    'launch',
+    'Resources',
+    'Task',
+]
+
 # Replaced with the current commit when building the wheels.
-_KONDUKTOR_COMMIT_SHA = "10848f624606ce6072ec6c646c4b3a6cb347585a"
+_KONDUKTOR_COMMIT_SHA = '9d16b4c393e6b4dd288eca4ea6baf66591089b91'
 
 
 def _get_git_commit():
-    if "KONDUKTOR_COMMIT_SHA" not in _KONDUKTOR_COMMIT_SHA:
+    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
@@ -18,24 +28,24 @@ def _get_git_commit():
     try:
         cwd = os.path.dirname(__file__)
         commit_hash = subprocess.check_output(
-            ["git", "rev-parse", "HEAD"],
+            ['git', 'rev-parse', 'HEAD'],
             cwd=cwd,
             universal_newlines=True,
             stderr=subprocess.DEVNULL,
         ).strip()
         changes = subprocess.check_output(
-            ["git", "status", "--porcelain"],
+            ['git', 'status', '--porcelain'],
             cwd=cwd,
             universal_newlines=True,
             stderr=subprocess.DEVNULL,
         ).strip()
         if changes:
-            commit_hash += "-dirty"
+            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"
+__version__ = '1.0.0.dev0.1.0.dev20250313070642'
 __root_dir__ = os.path.dirname(os.path.abspath(__file__))

konduktor/adaptors/common.py

@@ -0,0 +1,88 @@
+# 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 threading
+from typing import Any, Callable, Optional, Tuple
+
+
+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 (e.g., `networkx` should not be imported for
+    `sky status and `pandas` should not be imported for `sky exec`).
+
+    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

konduktor/adaptors/gcp.py

@@ -0,0 +1,112 @@
+# 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"""
+
+# pylint: disable=import-outside-toplevel
+import json
+
+from konduktor.adaptors import common
+
+_IMPORT_ERROR_MESSAGE = (
+    'Failed to import dependencies for GCP. ' 'Try pip install "skypilot[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)
+
+
+@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 storage.Client()
+
+
+@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 storage.Client.create_anonymous_client()
+
+
+@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 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/backends/__init__.py

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

konduktor/backends/backend.py

@@ -0,0 +1,86 @@
+# 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.
+
+"""Konduktor backend interface."""
+
+import typing
+from typing import Dict, Optional
+
+if typing.TYPE_CHECKING:
+    from konduktor.data import storage as storage_lib
+
+import konduktor
+from konduktor.utils import ux_utils
+
+Path = str
+
+
+class Backend:
+    """Backend interface: handles provisioning, setup, and scheduling."""
+
+    # NAME is used to identify the backend class from cli/yaml.
+    NAME = 'backend'
+
+    # --- APIs ---
+    def check_resources_fit_cluster(self, task: 'konduktor.Task') -> bool:
+        """Check whether resources of the task are satisfied by cluster."""
+        raise NotImplementedError
+
+    def sync_workdir(self, workdir: Path) -> None:
+        return self._sync_workdir(workdir)
+
+    def sync_file_mounts(
+        self,
+        all_file_mounts: Optional[Dict[Path, Path]],
+        storage_mounts: Optional[Dict[Path, 'storage_lib.Storage']],
+    ) -> None:
+        return self._sync_file_mounts(all_file_mounts, storage_mounts)
+
+    def add_storage_objects(self, task: 'konduktor.Task') -> None:
+        raise NotImplementedError
+
+    def execute(
+        self, task: 'konduktor.Task', detach_run: bool, dryrun: bool = False
+    ) -> Optional[str]:
+        """Execute the task on the cluster.
+
+        Returns:
+            Job id if the task is submitted to the cluster, None otherwise.
+        """
+        ux_utils.spinner_message('Submitting job')
+        return self._execute(task, detach_run, dryrun)
+
+    def post_execute(self) -> None:
+        """Post execute(): e.g., print helpful inspection messages."""
+        return self._post_execute()
+
+    def register_info(self, **kwargs) -> None:
+        """Register backend-specific information."""
+        pass
+
+    def _sync_workdir(self, workdir: Path) -> None:
+        raise NotImplementedError
+
+    def _sync_file_mounts(
+        self,
+        all_file_mounts: Optional[Dict[Path, Path]],
+        storage_mounts: Optional[Dict[Path, 'storage_lib.Storage']],
+    ) -> None:
+        raise NotImplementedError
+
+    def _execute(
+        self, task: 'konduktor.Task', detach_run: bool, dryrun: bool = False
+    ) -> Optional[str]:
+        raise NotImplementedError
+
+    def _post_execute(self) -> None:
+        raise NotImplementedError

konduktor/backends/jobset.py

@@ -0,0 +1,218 @@
+"""Batch job execution via k8s jobsets
+https://jobset.sigs.k8s.io/
+https://kueue.sigs.k8s.io/docs/tasks/run/jobsets/
+"""
+
+import threading
+import time
+import typing
+from typing import Dict, Optional, Tuple
+
+import colorama
+
+if typing.TYPE_CHECKING:
+    import konduktor
+    from konduktor.data import storage as storage_lib
+
+from konduktor import config, logging
+from konduktor.backends import backend, jobset_utils
+from konduktor.utils import kubernetes_utils, loki_utils, rich_utils, ux_utils
+
+Path = str
+logger = logging.get_logger(__file__)
+
+POLL_INTERVAL = 5
+DEFAULT_ATTACH_TIMEOUT = 300  # 5 minutes
+FLUSH_LOGS_TIMEOUT = 5
+
+
+class JobsetError(Exception):
+    pass
+
+
+def _raise_job_error(job):
+    """Checks a jobs conditions and statuses for error"""
+    for condition in job.status.conditions:
+        if 'ConfigIssue' in condition.message:
+            raise ValueError(
+                'Job failed with '
+                f'{colorama.Style.BRIGHT}{colorama.Fore.RED}'
+                f'ConfigIssue: ErrImagePull.{colorama.Style.RESET_ALL} '
+                f'Check that your '
+                f'{colorama.Style.BRIGHT}{colorama.Fore.YELLOW}'
+                f'`image_id`{colorama.Style.RESET_ALL} is correct and '
+                f'your container credentials are correct. Image specified '
+                f'in your task definition is '
+                f'{colorama.Style.BRIGHT}{colorama.Fore.RED}'
+                f'`{job.spec.template.spec.containers[0].image}`'
+                f'{colorama.Style.RESET_ALL}'
+            )
+        elif 'BackoffLimitExceeded' == condition.reason:
+            raise JobsetError('Job failed with non-zero exit code.')
+    logger.error(
+        'Job failed with unknown error. Check jobset status in k8s with '
+        f'{colorama.Style.BRIGHT}{colorama.Fore.YELLOW}'
+        f'`kubectl get job -o yaml {job.metadata.name}`'
+        f'{colorama.Style.RESET_ALL}'
+    )
+
+
+def _wait_for_jobset_start(namespace: str, job_name: str):
+    start = time.time()
+    timeout = config.get_nested(
+        ('kubernetes', 'pod_config'),
+        default_value=DEFAULT_ATTACH_TIMEOUT,
+        override_configs={},
+    )
+
+    while True:
+        jobsets = jobset_utils.get_jobset(namespace, job_name)
+        assert jobsets is not None, (
+            f'Jobset {job_name} ' f'not found in namespace {namespace}'
+        )
+        if jobsets['status']['replicatedJobsStatus'][0]['ready']:
+            logger.info(
+                f'task '
+                f'{colorama.Fore.CYAN}{colorama.Style.BRIGHT}{job_name}'
+                f'{colorama.Style.RESET_ALL} ready'
+            )
+            break
+        elif jobsets['status']['replicatedJobsStatus'][0]['succeeded']:
+            return
+        elif jobsets['status']['replicatedJobsStatus'][0]['failed']:
+            logger.info(
+                f'job '
+                f'{colorama.Fore.CYAN}{colorama.Style.BRIGHT}{job_name}'
+                f'{colorama.Style.RESET_ALL} '
+                f'{colorama.Fore.RED}{colorama.Style.BRIGHT}failed{colorama.Style.RESET_ALL}'
+            )
+            job = jobset_utils.get_job(namespace, job_name)
+            _raise_job_error(job)
+            return
+        if timeout != -1 and time.time() - start > timeout:
+            logger.error(
+                f'{colorama.Style.BRIGHT}'
+                f'{colorama.Fore.RED}Job timed out to schedule.'
+                f'{colorama.Style.RESET_ALL}. Deleting job'
+            )
+            jobset_utils.delete_jobset(namespace, job_name)
+            raise JobsetError(
+                'Job failed to start within '
+                f'timeout of {timeout} seconds. '
+                f'Increase or disable timeout '
+                f'{colorama.Style.BRIGHT}'
+                '`konduktor.provision_timeout: -1`'
+                f'{colorama.Style.RESET_ALL}'
+            )
+        time.sleep(POLL_INTERVAL)
+
+
+def _wait_for_jobset_completion(namespace: str, job_name: str) -> Tuple[bool, str]:
+    while True:
+        jobsets = jobset_utils.get_jobset(namespace, job_name)
+        assert jobsets is not None, (
+            f'Jobset {job_name} ' f'not found in namespace {namespace}'
+        )
+        if jobsets['status']['replicatedJobsStatus'][0]['succeeded']:
+            msg = (
+                f'task '
+                f'{colorama.Fore.CYAN}{colorama.Style.BRIGHT}{job_name}'
+                f'{colorama.Style.RESET_ALL} {colorama.Fore.GREEN}'
+                f'{colorama.Style.BRIGHT}finished{colorama.Style.RESET_ALL}'
+            )
+            return True, msg
+        elif jobsets['status']['replicatedJobsStatus'][0]['failed']:
+            msg = (
+                f'task '
+                f'{colorama.Fore.CYAN}{colorama.Style.BRIGHT}{job_name}'
+                f'{colorama.Style.RESET_ALL} {colorama.Fore.RED}'
+                f'{colorama.Style.BRIGHT}failed{colorama.Style.RESET_ALL}'
+            )
+            return False, msg
+        time.sleep(POLL_INTERVAL)
+
+
+class JobsetBackend(backend.Backend):
+    def _sync_file_mounts(
+        self,
+        all_file_mounts: Optional[Dict[Path, Path]],
+        storage_mounts: Optional[Dict[Path, 'storage_lib.Storage']],
+    ) -> None:
+        """Syncs files/directories to cloud storage before job launch.
+
+        This uploads any local files/dirs to cloud storage so they can be downloaded
+        by the pods when they start.
+        """
+        pass
+
+    def _sync_workdir(self, workdir: str) -> None:
+        """Syncs the working directory to cloud storage before job launch."""
+
+        pass
+
+    def _post_execute(self) -> None:
+        """
+        TODO(asaiacai): add some helpful messages/commands that a user can run
+        to inspect the status of their jobset.
+        """
+        pass
+
+    def _execute(
+        self, task: 'konduktor.Task', detach_run: bool = False, dryrun: bool = False
+    ) -> Optional[str]:
+        """Executes the task on the cluster. By creating a jobset
+
+        Returns:
+            Job id if the task is submitted to the cluster, None otherwise.
+        """
+
+        # we should consider just building an image with the cloud provider
+        # sdks baked in. These can initialize and pull files first before
+        # the working container starts.
+
+        # first define the pod spec then create the jobset definition
+        pod_spec = jobset_utils.create_pod_spec(task)
+        context = kubernetes_utils.get_current_kube_config_context_name()
+        namespace = kubernetes_utils.get_kube_config_context_namespace(context)
+        # TODO(asaiacai): need to set env variables in pod
+        jobset_utils.create_jobset(
+            namespace,
+            task,
+            pod_spec['kubernetes']['pod_config'],
+            dryrun=dryrun,
+        )
+
+        if not dryrun and not detach_run:
+            with ux_utils.print_exception_no_traceback():
+                with rich_utils.safe_status(
+                    ux_utils.spinner_message('waiting for job to start\n')
+                ):
+                    _wait_for_jobset_start(namespace, task.name)
+                try:
+                    log_thread = threading.Thread(
+                        target=loki_utils.tail_loki_logs_ws,
+                        args=(task.name,),
+                        daemon=True,
+                    )
+                    logger.info('streaming logs...')
+                    log_thread.start()
+                    is_success, msg = _wait_for_jobset_completion(namespace, task.name)
+                    log_thread.join(timeout=2.0)  # give the job sometime to flush logs
+                    if not is_success:
+                        logger.error(msg)
+                    else:
+                        logger.info(msg)
+                except KeyboardInterrupt:
+                    logger.info('detaching from log stream...')
+                except Exception as err:
+                    logger.error(
+                        f'Check if job resources are '
+                        f'active/queued with '
+                        f'{colorama.Style.BRIGHT}'
+                        f'`konduktor status`'
+                        f'{colorama.Style.RESET_ALL}'
+                    )
+                    raise JobsetError(f'error: {err}')
+        else:
+            logger.info('detaching from run.')
+        return task.name