konduktor-nightly 0.1.0.dev20250710105129__tar.gz → 0.1.0.dev20251103104940__tar.gz

{konduktor_nightly-0.1.0.dev20250710105129 → konduktor_nightly-0.1.0.dev20251103104940}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: konduktor-nightly
-Version: 0.1.0.dev20250710105129
+Version: 0.1.0.dev20251103104940
 Summary: GPU Cluster Health Management
 Author: Andrew Aikawa
 Author-email: asai@berkeley.edu

{konduktor_nightly-0.1.0.dev20250710105129 → konduktor_nightly-0.1.0.dev20251103104940}/konduktor/__init__.py

@@ -5,16 +5,13 @@ 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',
-]
+__all__ = ['launch', 'Resources', 'Task', 'Serving']
 
 # Replaced with the current commit when building the wheels.
-_KONDUKTOR_COMMIT_SHA = '5760f2c18ed487270e9244d22c7209eee12821c1'
+_KONDUKTOR_COMMIT_SHA = 'd5fddf4e144c4887227e1c6943c70bcd72d364d5'
 os.makedirs(os.path.expanduser('~/.konduktor'), exist_ok=True)
 
 
@@ -48,5 +45,5 @@ def _get_git_commit():
 
 
 __commit__ = _get_git_commit()
-__version__ = '1.0.0.dev0.1.0.dev20250710105129'
+__version__ = '1.0.0.dev0.1.0.dev20251103104940'
 __root_dir__ = os.path.dirname(os.path.abspath(__file__))

konduktor_nightly-0.1.0.dev20251103104940/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']

konduktor_nightly-0.1.0.dev20251103104940/konduktor/backends/constants.py

@@ -0,0 +1,21 @@
+KONDUKTOR_SSH_PORT = 2222
+
+# Common labels used across JobSets and Deployments
+JOB_NAME_LABEL = 'trainy.ai/job-name'
+DEPLOYMENT_NAME_LABEL = 'trainy.ai/deployment-name'
+AIBRIX_NAME_LABEL = 'model.aibrix.ai/name'
+USERID_LABEL = 'trainy.ai/user-id'
+USER_LABEL = 'trainy.ai/username'
+ACCELERATOR_LABEL = 'trainy.ai/accelerator'
+NUM_ACCELERATORS_LABEL = 'trainy.ai/num-accelerators'
+MAX_EXECUTION_TIME_LABEL = 'kueue.x-k8s.io/max-exec-time-seconds'
+
+# Start/stop/status related labels
+STOP_USERID_LABEL = 'trainy.ai/stop-userid'
+STOP_USERNAME_LABEL = 'trainy.ai/stop-username'
+
+# Secret labels
+SECRET_BASENAME_LABEL = 'trainy.ai/secret-basename'
+SECRET_KIND_LABEL = 'trainy.ai/secret-kind'
+SECRET_OWNER_LABEL = 'trainy.ai/secret-owner'
+ROOT_NAME = 'trainy.ai/root-name'

konduktor_nightly-0.1.0.dev20251103104940/konduktor/backends/deployment.py

@@ -0,0 +1,204 @@
+import time
+import typing
+from typing import Dict, Optional
+
+import colorama
+
+if typing.TYPE_CHECKING:
+    import konduktor
+    from konduktor.data import storage as storage_lib
+
+from kubernetes.client.exceptions import ApiException
+
+from konduktor import config, kube_client, logging
+from konduktor.backends import backend, deployment_utils, pod_utils
+from konduktor.utils import kubernetes_utils, rich_utils, ux_utils
+
+Path = str
+logger = logging.get_logger(__file__)
+
+POLL_INTERVAL = 5
+DEFAULT_ATTACH_TIMEOUT = 300
+
+
+class DeploymentError(Exception):
+    pass
+
+
+def _wait_for_all_ready(namespace: str, name: str):
+    """Wait for Deployment, Service, and Autoscaler readiness."""
+    time.sleep(2)
+    start = time.time()
+    timeout = config.get_nested(
+        ('kubernetes', 'provision_timeout'),
+        default_value=DEFAULT_ATTACH_TIMEOUT,
+    )
+
+    while True:
+        context = kubernetes_utils.get_current_kube_config_context_name()
+
+        # Directly read objects instead of listing everything
+        try:
+            deployment = kube_client.apps_api(context).read_namespaced_deployment(
+                name=name, namespace=namespace
+            )
+            deployments_map = {name: deployment}
+        except ApiException:
+            deployments_map = {}
+
+        try:
+            service = kube_client.core_api(context).read_namespaced_service(
+                name=name, namespace=namespace
+            )
+            services_map = {name: service}
+        except ApiException:
+            services_map = {}
+
+        autoscalers_map = {}
+        try:
+            autoscaler_obj = deployment_utils.get_autoscaler(namespace, name)
+            if autoscaler_obj:
+                # detect aibrix vs general from deployment labels
+                labels = (deployment.metadata.labels or {}) if deployment else {}
+                is_aibrix = deployment_utils.AIBRIX_NAME_LABEL in labels
+                if is_aibrix:
+                    autoscalers_map[name] = {'kpa': autoscaler_obj}
+                else:
+                    autoscalers_map[name] = {'hpa': autoscaler_obj}
+        except ApiException:
+            pass
+
+        status = deployment_utils.get_model_status(
+            name, deployments_map, services_map, autoscalers_map
+        )
+
+        is_ready = (
+            status['deployment'] == 'ready'
+            and status['service'] == 'ready'
+            and (status['autoscaler'] == 'ready' or status['autoscaler'] is None)
+        )
+
+        states = {
+            'Deployment': status['deployment'],
+            'Service': status['service'],
+            'Autoscaler': status['autoscaler'],
+        }
+
+        # Figure out which components are missing
+        missing_parts = [name for name, state in states.items() if state == 'missing']
+
+        if missing_parts:
+            deployment_utils.delete_serving_specs(name, namespace)
+            missing_str = ', '.join(missing_parts)
+            raise DeploymentError(
+                f'Deployment failed. '
+                f'The following components are missing: {missing_str}.'
+            )
+
+        if is_ready:
+            logger.info(
+                f'task {colorama.Fore.CYAN}{colorama.Style.BRIGHT}'
+                f'{name}{colorama.Style.RESET_ALL} ready'
+            )
+            return
+
+        if timeout != -1 and time.time() - start > timeout:
+            logger.error(
+                f'{colorama.Style.BRIGHT}{colorama.Fore.RED}'
+                f'Model timed out waiting for readiness.'
+                f'{colorama.Style.RESET_ALL}'
+                f'Final status:\n{status}'
+            )
+            deployment_utils.delete_serving_specs(name, namespace)
+            raise DeploymentError(
+                f'Model failed to become ready within {timeout} seconds.\n'
+            )
+
+        time.sleep(POLL_INTERVAL)
+
+
+class DeploymentBackend(backend.Backend):
+    NAME = 'deployment'
+
+    def check_resources_fit_cluster(self, task: 'konduktor.Task') -> bool:
+        return True
+
+    def add_storage_objects(self, task: 'konduktor.Task') -> None:
+        pass
+
+    def register_info(self, **kwargs) -> None:
+        pass
+
+    def _sync_file_mounts(
+        self,
+        all_file_mounts: Optional[Dict[Path, Path]],
+        storage_mounts: Optional[Dict[Path, 'storage_lib.Storage']],
+    ) -> None:
+        pass
+
+    def _sync_workdir(self, workdir: str) -> None:
+        pass
+
+    def _post_execute(self) -> None:
+        pass
+
+    def _execute(
+        self,
+        task: 'konduktor.Task',
+        detach_run: bool = False,
+        dryrun: bool = False,
+    ) -> Optional[str]:
+        """Execute a task by launching a long-running Deployment."""
+
+        pod_spec = pod_utils.create_pod_spec(task)
+        context = kubernetes_utils.get_current_kube_config_context_name()
+        namespace = kubernetes_utils.get_kube_config_context_namespace(context)
+
+        if not dryrun and task.serving:
+            logger.debug(f'[DEBUG] Creating deployment for task: {task.name}')
+            deployment_utils.create_deployment(
+                namespace=namespace,
+                task=task,
+                pod_spec=pod_spec['kubernetes']['pod_config'],
+                dryrun=dryrun,
+            )
+
+            logger.debug(f'[DEBUG] Creating service for task: {task.name}')
+            deployment_utils.create_service(
+                namespace=namespace,
+                task=task,
+                dryrun=dryrun,
+            )
+
+            # Create podautoscaler for non-general deployments
+            logger.debug(f'[DEBUG] Creating podautoscaler for task: {task.name}')
+            deployment_utils.create_pod_autoscaler(
+                namespace=namespace,
+                task=task,
+                dryrun=dryrun,
+            )
+
+            # HTTP Add-on resources for general deployments
+            logger.debug(
+                f'[DEBUG] Creating HTTP Add-on resources for task: {task.name}'
+            )
+            deployment_utils.create_http_addon_resources(
+                namespace=namespace,
+                task=task,
+                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 resources to be ready.\n')
+                ):
+                    _wait_for_all_ready(namespace, task.name)
+            logger.info(
+                f"Model '{task.name}' is ready. "
+                f'Run `konduktor serve status` for details.'
+            )
+        else:
+            logger.info('detaching from run.')
+
+        return task.name