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

konduktor/resource.py

@@ -0,0 +1,546 @@
+# 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.
+
+"""Resources: compute requirements of Tasks."""
+
+import functools
+from typing import Any, Dict, List, Optional, Union
+
+from konduktor import logging
+from konduktor.utils import (
+    accelerator_registry,
+    common_utils,
+    schemas,
+    ux_utils,
+    validator,
+)
+
+logger = logging.get_logger(__name__)
+
+_DEFAULT_DISK_SIZE_GB = 256
+
+
+class Resources:
+    """Resources: compute requirements of Tasks.
+
+    This class is immutable once created (to ensure some validations are done
+    whenever properties change). To update the property of an instance of
+    Resources, use `resources.copy(**new_properties)`.
+
+    Used:
+
+    * for representing resource requests for task
+
+    """
+
+    # If any fields changed, increment the version. For backward compatibility,
+    # modify the __setstate__ method to handle the old version.
+    _VERSION = 1
+
+    def __init__(
+        self,
+        cloud: Optional[Any] = None,
+        cpus: Union[None, int, float, str] = None,
+        memory: Union[None, int, float, str] = None,
+        accelerators: Optional[str] = None,
+        image_id: Union[str, None] = None,
+        disk_size: Optional[int] = None,
+        labels: Optional[Dict[str, str]] = None,
+        job_config: Optional[Dict[str, Union[int, str]]] = None,
+        # Internal use only.
+        # pylint: disable=invalid-name
+        _cluster_config_overrides: Optional[Dict[str, Any]] = None,
+        # used to prevent double validation of image (would happen from overrides)
+        _validate_image: bool = True,
+    ):
+        """Initialize a Resources object.
+
+        All fields are optional.  ``Resources.is_launchable`` decides whether
+        the Resources is fully specified to launch an instance.
+
+        Examples:
+          .. code-block:: python
+
+            # Specifying required resources; the system decides the
+            # cloud/instance type. The below are equivalent:
+            konduktor.Resources(accelerators='V100')
+            konduktor.Resources(accelerators='V100:1')
+            konduktor.Resources(accelerators={'V100': 1})
+            konduktor.Resources(cpus='2+', memory='16+', accelerators='V100')
+
+        Args:
+          cloud: the cloud to use. (deprecated) all jobs are submitted to k8s
+          instance_type: the instance type to use.
+          cpus: the number of CPUs required for the task.
+            If a str, must be a string of the form ``'2'`` or ``'2+'``, where
+            the ``+`` indicates that the task requires at least 2 CPUs.
+          memory: the amount of memory in GiB required. If a
+            str, must be a string of the form ``'16'`` or ``'16+'``, where
+            the ``+`` indicates that the task requires at least 16 GB of memory.
+          accelerators: the accelerators required. If a str, must be
+            a string of the form ``'V100'`` or ``'V100:2'``, where the ``:2``
+            indicates that the task requires 2 V100 GPUs. If a dict, must be a
+            dict of the form ``{'V100': 2}`` or ``{'tpu-v2-8': 1}``.
+
+          image_id: docker image to use
+
+          disk_size: the size of the OS disk in GiB.
+          labels: the labels to apply to the instance. These are useful for
+            assigning metadata that may be used by external tools.
+            Implementation depends on the chosen cloud - On AWS, labels map to
+            instance tags. On GCP, labels map to instance labels. On
+            Kubernetes, labels map to pod labels. On other clouds, labels are
+            not supported and will be ignored.
+          job_config: the configuration of the job spec
+        Raises:
+            ValueError: if some attributes are invalid.
+            exceptions.NoCloudAccessError: if no public cloud is enabled.
+        """
+        self._version = self._VERSION
+        if cloud is not None:
+            raise ValueError('cloud specified, but all jobs are submitted to k8s')
+        self._cloud = cloud
+
+        if disk_size is not None:
+            if round(disk_size) != disk_size:
+                with ux_utils.print_exception_no_traceback():
+                    raise ValueError(
+                        f'OS disk size must be an integer. Got: {disk_size}.'
+                    )
+            self._disk_size = int(disk_size)
+        else:
+            self._disk_size = _DEFAULT_DISK_SIZE_GB
+
+        # self._image_id is a dict of {region: image_id}.
+        # The key is None if the same image_id applies for all regions.
+        self._image_id = image_id
+        if isinstance(image_id, str):
+            self._image_id = image_id.strip()
+            # Validate Docker image format and existence
+            if _validate_image:
+                validator.validate_and_warn_image(self._image_id, 'task')
+
+        self._labels = labels
+        self._cluster_config_overrides = _cluster_config_overrides
+
+        self._set_cpus(cpus)
+        self._set_memory(memory)
+        self._set_accelerators(accelerators)
+        self.job_config = job_config or {}
+
+        # TODO: move these out of init to prevent repeated calls.
+        self._try_validate_cpus_mem()
+        self._try_validate_image_id()
+
+    def __repr__(self) -> str:
+        """Returns a string representation for display.
+
+        Examples:
+
+            >>> konduktor.Resources(accelerators='V100')
+            <Kubernetes>({'V100': 1})
+
+        """
+        accelerators = ''
+        if self.accelerators is not None:
+            accelerators = f', {self.accelerators}'
+
+        cpus = ''
+        if self._cpus is not None:
+            cpus = f', cpus={self._cpus}'
+
+        memory = ''
+        if self.memory is not None:
+            memory = f', mem={self.memory}'
+
+        image_id = ''
+        if self.image_id is not None:
+            image_id = f', image_id={self.image_id}'
+        else:
+            with ux_utils.print_exception_no_traceback():
+                raise ValueError(
+                    'no image id for the task was specified. You must '
+                    'specify an image id for this task (e.g. '
+                    '`nvcr.io/nvidia/pytorch:xx.xx-py3`'
+                )
+
+        disk_size = ''
+        if self.disk_size != _DEFAULT_DISK_SIZE_GB:
+            disk_size = f', disk_size={self.disk_size}'
+
+        # Do not show region/zone here as `konduktor status -a` would show them as
+        # separate columns. Also, Resources repr will be printed during
+        # failover, and the region may be dynamically determined.
+        hardware_str = f'{cpus}{memory}{accelerators}{image_id}' f'{disk_size}'
+        # It may have leading ',' (for example, instance_type not set) or empty
+        # spaces.  Remove them.
+        while hardware_str and hardware_str[0] in (',', ' '):
+            hardware_str = hardware_str[1:]
+
+        return f'({hardware_str})'
+
+    @property
+    def cloud(self):
+        return self._cloud
+
+    @property
+    @functools.lru_cache(maxsize=1)
+    def cpus(self) -> Optional[str]:
+        """Returns the number of vCPUs that each instance must have.
+
+        For example, cpus='4' means each instance must have exactly 4 vCPUs,
+        and cpus='4+' means each instance must have at least 4 vCPUs.
+
+        (Developer note: The cpus field is only used to select the instance type
+        at launch time. Thus, Resources in the backend's ResourceHandle will
+        always have the cpus field set to None.)
+        """
+        if self._cpus is not None:
+            return self._cpus
+        return None
+
+    @property
+    def memory(self) -> Optional[str]:
+        """Returns the memory that each instance must have in GB.
+
+        For example, memory='16' means each instance must have exactly 16GB
+        memory; memory='16+' means each instance must have at least 16GB
+        memory.
+
+        (Developer note: The memory field is only used to select the instance
+        type at launch time. Thus, Resources in the backend's ResourceHandle
+        will always have the memory field set to None.)
+        """
+        return self._memory
+
+    @property
+    @functools.lru_cache(maxsize=1)
+    def accelerators(self) -> Optional[Dict[str, int]]:
+        """Returns the accelerators field directly or by inferring.
+
+        For example, Resources(AWS, 'p3.2xlarge') has its accelerators field
+        set to None, but this function will infer {'V100': 1} from the instance
+        type.
+        """
+        if self._accelerators is not None:
+            return self._accelerators
+        return None
+
+    @property
+    def disk_size(self) -> int:
+        return self._disk_size
+
+    @property
+    def image_id(self) -> Optional[str]:
+        return self._image_id
+
+    @property
+    def labels(self) -> Optional[Dict[str, str]]:
+        return self._labels
+
+    @property
+    def cluster_config_overrides(self) -> Dict[str, Any]:
+        if self._cluster_config_overrides is None:
+            return {}
+        return self._cluster_config_overrides
+
+    def _set_cpus(
+        self,
+        cpus: Union[None, int, float, str],
+    ) -> None:
+        if cpus is None:
+            self._cpus = None
+            return
+
+        self._cpus = str(cpus)
+        if isinstance(cpus, str):
+            if cpus.endswith('+'):
+                num_cpus_str = cpus[:-1]
+            else:
+                num_cpus_str = cpus
+
+            try:
+                num_cpus = float(num_cpus_str)
+            except ValueError:
+                with ux_utils.print_exception_no_traceback():
+                    raise ValueError(
+                        f'The "cpus" field should be either a number or '
+                        f'a string "<number>+". Found: {cpus!r}'
+                    ) from None
+        else:
+            num_cpus = float(cpus)
+
+        if num_cpus <= 0:
+            with ux_utils.print_exception_no_traceback():
+                raise ValueError(
+                    f'The "cpus" field should be positive. Found: {cpus!r}'
+                )
+
+    def _set_memory(
+        self,
+        memory: Union[None, int, float, str],
+    ) -> None:
+        if memory is None:
+            self._memory = None
+            return
+
+        self._memory = str(memory)
+        if isinstance(memory, str):
+            if memory.endswith(('+', 'x')):
+                # 'x' is used internally for make sure our resources used by
+                # jobs controller (memory: 3x) to have enough memory based on
+                # the vCPUs.
+                num_memory_gb = memory[:-1]
+            else:
+                num_memory_gb = memory
+
+            try:
+                memory_gb = float(num_memory_gb)
+            except ValueError:
+                with ux_utils.print_exception_no_traceback():
+                    raise ValueError(
+                        f'The "memory" field should be either a number or '
+                        f'a string "<number>+". Found: {memory!r}'
+                    ) from None
+        else:
+            memory_gb = float(memory)
+
+        if memory_gb <= 0:
+            with ux_utils.print_exception_no_traceback():
+                raise ValueError(
+                    f'The "cpus" field should be positive. Found: {memory!r}'
+                )
+
+    def _set_accelerators(
+        self,
+        accelerators: Union[None, str, Dict[str, int]],
+        accelerator_args: Optional[Dict[str, str]] = None,
+    ) -> None:
+        """Sets accelerators.
+
+        Args:
+            accelerators: A string or a dict of accelerator types to counts.
+            accelerator_args: (deprecated) A dict of accelerator types to args.
+        """
+        if accelerators is not None:
+            if isinstance(accelerators, str):  # Convert to Dict[str, int].
+                if ':' not in accelerators:
+                    accelerators = {accelerators: 1}
+                else:
+                    splits = accelerators.split(':')
+                    parse_error = (
+                        'The "accelerators" field as a str '
+                        'should be <name> or <name>:<cnt>. '
+                        f'Found: {accelerators!r}'
+                    )
+                    if len(splits) != 2:
+                        with ux_utils.print_exception_no_traceback():
+                            raise ValueError(parse_error)
+                    try:
+                        num = float(splits[1])
+                        num = int(num)
+                        accelerators = {splits[0]: num}
+                    except ValueError:
+                        with ux_utils.print_exception_no_traceback():
+                            raise ValueError(parse_error) from None
+
+            # Canonicalize the accelerator names.
+            accelerators = {
+                accelerator_registry.canonicalize_accelerator_name(acc): acc_count
+                for acc, acc_count in accelerators.items()
+            }
+
+            acc, _ = list(accelerators.items())[0]
+
+        self._accelerators = accelerators
+
+    def _try_validate_cpus_mem(self) -> None:
+        """Try to validate the cpus and memory attributes.
+
+        Raises:
+            ValueError: if the attributes are invalid.
+        """
+        if self._cpus is None and self._memory is None:
+            return
+
+    def _try_validate_image_id(self) -> None:
+        """Try to validate the image_id attribute.
+
+        Raises:
+            ValueError: if the attribute is invalid.
+        """
+        if self._image_id is None:
+            with ux_utils.print_exception_no_traceback():
+                raise ValueError(
+                    'no image id for the task was specified. You must '
+                    'specify an image id for this task (e.g. '
+                    '`nvcr.io/nvidia/pytorch:xx.xx-py3`'
+                )
+
+    def get_accelerators_str(self) -> str:
+        accelerators = self.accelerators
+        accel_str = ''
+        if accelerators is None:
+            accel_str = '-'
+        elif isinstance(accelerators, dict) and len(accelerators) == 1:
+            accel_name, accel_count = list(accelerators.items())[0]
+            accel_str = f'{accel_name}:{accel_count}'
+        return accel_str
+
+    def get_completions(self) -> Optional[int]:
+        value = self.job_config.get('completions')
+        if value is not None:
+            value = int(value)
+            if value <= 0:
+                with ux_utils.print_exception_no_traceback():
+                    raise ValueError('completions must be a positive integer')
+            return value
+        return None
+
+    def get_max_restarts(self) -> Optional[int]:
+        value = self.job_config.get('max_restarts')
+        if value is not None:
+            value = int(value)
+            if value < 0:
+                with ux_utils.print_exception_no_traceback():
+                    raise ValueError('max_restarts must be a non-negative integer')
+            return value
+        return None
+
+    def get_accelerator_type(self) -> Optional[str]:
+        """Returns the first accelerator type from the accelerators dict.
+
+        Returns:
+            The accelerator type (e.g., 'V100', 'A100') or None if no accelerators
+        """
+        if self.accelerators is None or not self.accelerators:
+            return None
+        return next(iter(self.accelerators.keys()))  # type: ignore
+
+    def get_accelerator_count(self) -> Optional[int]:
+        """Returns the count of the first accelerator type from the accelerators dict.
+
+        Returns:
+            The accelerator count (e.g., 1, 2) or None if no accelerators
+        """
+        if self.accelerators is None or not self.accelerators:
+            return None
+        return next(iter(self.accelerators.values()))  # type: ignore
+
+    def copy(self, **override) -> 'Resources':
+        """Returns a copy of the given Resources."""
+        # used to prevent double validation of image (would happen from overrides)
+        new_image_id = override.pop('image_id', self.image_id)
+        resources = Resources(
+            cloud=override.pop('cloud', self.cloud),
+            cpus=override.pop('cpus', self._cpus),
+            memory=override.pop('memory', self.memory),
+            accelerators=override.pop('accelerators', self.accelerators),
+            disk_size=override.pop('disk_size', self.disk_size),
+            image_id=new_image_id,
+            labels=override.pop('labels', self.labels),
+            job_config=override.pop('job_config', self.job_config),
+            # used to prevent double validation of image (would happen from overrides)
+            _validate_image=(new_image_id != self.image_id),
+        )
+        assert len(override) == 0
+        return resources
+
+    @classmethod
+    def from_yaml_config(cls, config: Optional[Dict[str, Any]]) -> 'Resources':
+        if config is None:
+            return Resources()
+        common_utils.validate_schema(
+            config, schemas.get_resources_schema(), 'Invalid resources YAML: '
+        )
+
+        if config.get('job_config', None):
+            common_utils.validate_schema(
+                config['job_config'],
+                schemas.get_job_schema(),
+                'Invalid job config YAML: ',
+            )
+
+        def _override_resources(
+            base_resource_config: Dict[str, Any], override_configs: List[Dict[str, Any]]
+        ) -> List[Resources]:
+            resources_list = []
+            for override_config in override_configs:
+                new_resource_config = base_resource_config.copy()
+                # Labels are handled separately.
+                override_labels = override_config.pop('labels', None)
+                new_resource_config.update(override_config)
+
+                # Update the labels with the override labels.
+                labels = new_resource_config.get('labels', None)
+                if labels is not None and override_labels is not None:
+                    labels.update(override_labels)
+                elif override_labels is not None:
+                    labels = override_labels
+                new_resource_config['labels'] = labels
+
+                # Call from_yaml_config again instead of
+                # _from_yaml_config_single to handle the case, where both
+                # multiple accelerators and `any_of` is specified.
+                # This will not cause infinite recursion because we have made
+                # sure that `any_of` and `ordered` cannot be specified in the
+                # resource candidates in `any_of` or `ordered`, by the schema
+                # validation above.
+                resources_list.extend([Resources.from_yaml_config(new_resource_config)])
+
+            return resources_list
+
+        config = config.copy()
+
+        return Resources._from_yaml_config_single(config)
+
+    @classmethod
+    def _from_yaml_config_single(cls, config: Dict[str, str]) -> 'Resources':
+        resources_fields: Dict[str, Any] = {}
+        resources_fields['cpus'] = config.pop('cpus', None)
+        resources_fields['memory'] = config.pop('memory', None)
+        resources_fields['accelerators'] = config.pop('accelerators', None)
+        resources_fields['disk_size'] = config.pop('disk_size', None)
+        resources_fields['image_id'] = config.pop('image_id', None)
+        resources_fields['labels'] = config.pop('labels', None)
+        resources_fields['job_config'] = config.pop('job_config', None)
+
+        if resources_fields['cpus'] is not None:
+            resources_fields['cpus'] = str(resources_fields['cpus'])
+        if resources_fields['memory'] is not None:
+            resources_fields['memory'] = str(resources_fields['memory'])
+        # TODO(asaiacai): should we remove disk size
+        # since we aren't letting users set this at the host level?
+        if resources_fields['disk_size'] is not None:
+            resources_fields['disk_size'] = int(resources_fields['disk_size'])
+
+        assert not config, f'Invalid resource args: {config.keys()}'
+        return Resources(**resources_fields)
+
+    def to_yaml_config(self) -> Dict[str, Union[str, int]]:
+        """Returns a yaml-style dict of config for this resource bundle."""
+        config = {}
+
+        def add_if_not_none(key, value):
+            if value is not None and value != 'None':
+                config[key] = value
+
+        add_if_not_none('cloud', str(self.cloud))
+        add_if_not_none('cpus', self._cpus)
+        add_if_not_none('memory', self.memory)
+        add_if_not_none('accelerators', self.accelerators)
+
+        add_if_not_none('disk_size', self.disk_size)
+        add_if_not_none('image_id', self.image_id)
+        add_if_not_none('labels', self.labels)
+        add_if_not_none('job_config', self.job_config)
+        return config

konduktor/serving.py

@@ -0,0 +1,153 @@
+# 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.
+
+"""Serving: configuration for long-running serving deployments."""
+
+from typing import Any, Dict, Optional, Union
+
+from konduktor import logging
+from konduktor.utils import common_utils, schemas, ux_utils
+
+logger = logging.get_logger(__name__)
+
+
+class Serving:
+    """Serving: configuration for deployments.
+
+    Immutable once created. Use `copy()` to create a modified copy.
+
+    Used:
+    * to represent serving config in tasks
+    """
+
+    _VERSION = 1
+
+    def __init__(
+        self,
+        min_replicas: Optional[int] = None,
+        max_replicas: Optional[int] = None,
+        ports: Optional[int] = 8000,
+        probe: Optional[str] = '/health',
+    ):
+        self._version = self._VERSION
+
+        if min_replicas is None and max_replicas is None:
+            with ux_utils.print_exception_no_traceback():
+                raise ValueError(
+                    'At least one of min_replicas or ' 'max_replicas must be specified.'
+                )
+
+        if min_replicas is None:
+            min_replicas = max_replicas
+        if max_replicas is None:
+            # Edge case: if min_replicas is 0, set max_replicas to 1
+            if min_replicas == 0:
+                max_replicas = 1
+            else:
+                max_replicas = min_replicas
+
+        if min_replicas is not None and min_replicas < 0:
+            with ux_utils.print_exception_no_traceback():
+                raise ValueError('min_replicas must be >= 0')
+
+        if (
+            max_replicas is not None
+            and min_replicas is not None
+            and max_replicas < min_replicas
+        ):
+            with ux_utils.print_exception_no_traceback():
+                raise ValueError(
+                    f'max_replicas ({max_replicas}) must '
+                    f'be >= min_replicas ({min_replicas})'
+                )
+
+        self._min_replicas = min_replicas
+        self._max_replicas = max_replicas
+        self._ports = ports
+        self._probe = probe
+
+    @property
+    def min_replicas(self) -> int:
+        assert self._min_replicas is not None
+        return self._min_replicas
+
+    @property
+    def max_replicas(self) -> int:
+        assert self._max_replicas is not None
+        return self._max_replicas
+
+    @property
+    def ports(self) -> int:
+        assert self._ports is not None
+        return self._ports
+
+    @property
+    def probe(self) -> Optional[str]:
+        return self._probe
+
+    def get(self, key: str, default=None):
+        return {
+            'min_replicas': self._min_replicas,
+            'max_replicas': self._max_replicas,
+            'ports': self._ports,
+            'probe': self._probe,
+        }.get(key, default)
+
+    def copy(self, **override) -> 'Serving':
+        """Returns a copy of this Serving with fields overridden."""
+        return Serving(
+            min_replicas=override.pop('min_replicas', self._min_replicas),
+            max_replicas=override.pop('max_replicas', self._max_replicas),
+            ports=override.pop('ports', self._ports),
+            probe=override.pop('probe', self._probe),
+        )
+
+    @classmethod
+    def from_yaml_config(
+        cls, config: Optional[Dict[str, Any]], task_run: Optional[str] = None
+    ) -> Optional['Serving']:
+        if config is None:
+            return None
+        common_utils.validate_schema(
+            config,
+            schemas.get_serving_schema(),
+            'Invalid serving config YAML: ',
+        )
+
+        if 'min_replicas' not in config and 'max_replicas' not in config:
+            raise ValueError(
+                'At least one of min_replicas or '
+                'max_replicas must be specified in serving'
+            )
+
+        # Determine default probe based on deployment type
+        default_probe = None  # No probing by default for general deployments
+        if task_run and 'vllm.entrypoints.openai.api_server' in task_run:
+            default_probe = '/health'  # Aibrix deployments get /health by default
+
+        return cls(
+            min_replicas=config.get('min_replicas', None),
+            max_replicas=config.get('max_replicas', None),
+            ports=config.get('ports', 8000),
+            probe=config.get('probe', default_probe),
+        )
+
+    def to_yaml_config(self) -> Dict[str, Union[int, str]]:
+        config: Dict[str, Union[int, str]] = {
+            'min_replicas': self._min_replicas if self._min_replicas is not None else 1,
+            'max_replicas': self._max_replicas if self._max_replicas is not None else 1,
+            'ports': self._ports if self._ports is not None else 8000,
+        }
+        # Only include probe if it's not None
+        if self._probe is not None:
+            config['probe'] = self._probe
+        return config