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

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/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/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