skypilot-nightly 1.0.0.dev20250502__py3-none-any.whl → 1.0.0.dev20251203__py3-none-any.whl

sky/utils/volume.py

@@ -0,0 +1,165 @@
+"""Volume utilities."""
+from dataclasses import dataclass
+import enum
+import time
+from typing import Any, Dict, Optional
+
+from sky import exceptions
+from sky import global_user_state
+from sky import models
+from sky.utils import common_utils
+from sky.utils import resources_utils
+from sky.utils import schemas
+from sky.utils import status_lib
+
+MIN_RUNPOD_NETWORK_VOLUME_SIZE_GB = 10
+
+
+class VolumeAccessMode(enum.Enum):
+    """Volume access mode."""
+    READ_WRITE_ONCE = 'ReadWriteOnce'
+    READ_WRITE_ONCE_POD = 'ReadWriteOncePod'
+    READ_WRITE_MANY = 'ReadWriteMany'
+    READ_ONLY_MANY = 'ReadOnlyMany'
+
+
+class VolumeType(enum.Enum):
+    """Volume type."""
+    PVC = 'k8s-pvc'
+    RUNPOD_NETWORK_VOLUME = 'runpod-network-volume'
+
+    @classmethod
+    def supported_types(cls) -> list:
+        """Return list of supported volume type values."""
+        return [vt.value for vt in cls]
+
+
+EPHEMERAL_VOLUME_TYPES = [VolumeType.PVC.value]
+
+
+@dataclass
+class VolumeInfo:
+    """Represents volume info."""
+    name: str
+    path: str
+    volume_name_on_cloud: Optional[str] = None
+    volume_id_on_cloud: Optional[str] = None
+
+
+class VolumeMount:
+    """Volume mount specification."""
+
+    def __init__(self,
+                 path: str,
+                 volume_name: str,
+                 volume_config: models.VolumeConfig,
+                 is_ephemeral: bool = False):
+        self.path: str = path
+        self.volume_name: str = volume_name
+        self.volume_config: models.VolumeConfig = volume_config
+        self.is_ephemeral: bool = is_ephemeral
+
+    def pre_mount(self) -> None:
+        """Update the volume status before actual mounting."""
+        # Skip pre_mount for ephemeral volumes as they don't exist yet
+        if self.is_ephemeral:
+            return
+        # TODO(aylei): for ReadWriteOnce volume, we also need to queue the
+        # mount request if the target volume is already mounted to another
+        # cluster. For now, we only support ReadWriteMany volume.
+        global_user_state.update_volume(self.volume_name,
+                                        last_attached_at=int(time.time()),
+                                        status=status_lib.VolumeStatus.IN_USE)
+
+    @classmethod
+    def resolve(cls, path: str, volume_name: str) -> 'VolumeMount':
+        """Resolve the volume mount by populating metadata of volume."""
+        record = global_user_state.get_volume_by_name(volume_name)
+        if record is None:
+            raise exceptions.VolumeNotFoundError(
+                f'Volume {volume_name} not found.')
+        assert 'handle' in record, 'Volume handle is None.'
+        volume_config: models.VolumeConfig = record['handle']
+        return cls(path, volume_name, volume_config)
+
+    @classmethod
+    def from_yaml_config(cls, config: Dict[str, Any]) -> 'VolumeMount':
+        common_utils.validate_schema(config, schemas.get_volume_mount_schema(),
+                                     'Invalid volume mount config: ')
+
+        path = config.pop('path', None)
+        volume_name = config.pop('volume_name', None)
+        is_ephemeral = config.pop('is_ephemeral', False)
+        volume_config: models.VolumeConfig = models.VolumeConfig.model_validate(
+            config.pop('volume_config', None))
+        return cls(path, volume_name, volume_config, is_ephemeral)
+
+    @classmethod
+    def resolve_ephemeral_config(cls, path: str,
+                                 config: Dict[str, Any]) -> 'VolumeMount':
+        """Create an ephemeral volume mount from inline config.
+
+        Args:
+            path: The mount path for the volume.
+            config: The volume configuration dict with size, and
+                optional type, labels, and config fields, etc.
+
+        Returns:
+            A VolumeMount instance for the ephemeral volume.
+        """
+        volume_type = config.get('type')
+        if volume_type and volume_type.lower() not in EPHEMERAL_VOLUME_TYPES:
+            raise ValueError(f'Unsupported ephemeral volume type: '
+                             f'{volume_type}. Supported types: '
+                             f'{", ".join(EPHEMERAL_VOLUME_TYPES)}')
+        size_config = config.get('size')
+        if size_config is None:
+            raise ValueError('Volume size must be specified for ephemeral '
+                             'volumes.')
+        try:
+            size = resources_utils.parse_memory_resource(size_config,
+                                                         'size',
+                                                         allow_rounding=True)
+            if size == '0':
+                raise ValueError('Size must be no less than 1Gi')
+        except ValueError as e:
+            raise ValueError(
+                f'Invalid size {size_config} for ephemeral volume: {e}') from e
+
+        # Create VolumeConfig for ephemeral volume
+        # Note: the empty fields will be populated during provisioning
+        volume_config = models.VolumeConfig(
+            name='',
+            type=config.get('type', ''),
+            # Default to kubernetes cloud here for backward compatibility,
+            # but this will be reset to the correct cloud during provisioning.
+            cloud='kubernetes',
+            region=None,
+            zone=None,
+            name_on_cloud='',
+            size=size,
+            config=config.get('config', {}),
+            labels=config.get('labels'),
+        )
+
+        return cls(path, '', volume_config, is_ephemeral=True)
+
+    def to_yaml_config(self) -> Dict[str, Any]:
+        return {
+            'path': self.path,
+            'volume_name': self.volume_name,
+            'volume_config': self.volume_config.model_dump(),
+            'is_ephemeral': self.is_ephemeral,
+        }
+
+    @property
+    def name(self) -> str:
+        """Return the volume name for use in provisioning."""
+        return self.volume_name
+
+    def __repr__(self):
+        return (f'VolumeMount('
+                f'\n\tpath={self.path},'
+                f'\n\tvolume_name={self.volume_name},'
+                f'\n\tis_ephemeral={self.is_ephemeral},'
+                f'\n\tvolume_config={self.volume_config})')

sky/utils/yaml_utils.py

@@ -0,0 +1,111 @@
+"""YAML utilities."""
+import io
+from typing import Any, Dict, List, Optional, TYPE_CHECKING, Union
+
+from sky.adaptors import common
+
+if TYPE_CHECKING:
+    import yaml
+else:
+    yaml = common.LazyImport('yaml')
+
+_c_extension_unavailable = False
+
+
+def safe_load(stream) -> Any:
+    global _c_extension_unavailable
+    if _c_extension_unavailable:
+        return yaml.load(stream, Loader=yaml.SafeLoader)
+
+    try:
+        return yaml.load(stream, Loader=yaml.CSafeLoader)
+    except AttributeError:
+        _c_extension_unavailable = True
+        return yaml.load(stream, Loader=yaml.SafeLoader)
+
+
+def safe_load_all(stream) -> Any:
+    global _c_extension_unavailable
+    if _c_extension_unavailable:
+        return yaml.load_all(stream, Loader=yaml.SafeLoader)
+
+    try:
+        return yaml.load_all(stream, Loader=yaml.CSafeLoader)
+    except AttributeError:
+        _c_extension_unavailable = True
+        return yaml.load_all(stream, Loader=yaml.SafeLoader)
+
+
+def read_yaml(path: Optional[str]) -> Dict[str, Any]:
+    if path is None:
+        raise ValueError('Attempted to read a None YAML.')
+    with open(path, 'r', encoding='utf-8') as f:
+        config = safe_load(f)
+    return config
+
+
+def read_yaml_str(yaml_str: str) -> Dict[str, Any]:
+    stream = io.StringIO(yaml_str)
+    parsed_yaml = safe_load(stream)
+    if not parsed_yaml:
+        # Empty dict
+        return {}
+    return parsed_yaml
+
+
+def read_yaml_all_str(yaml_str: str) -> List[Dict[str, Any]]:
+    stream = io.StringIO(yaml_str)
+    config = safe_load_all(stream)
+    configs = list(config)
+    if not configs:
+        # Empty YAML file.
+        return [{}]
+    return configs
+
+
+def read_yaml_all(path: str) -> List[Dict[str, Any]]:
+    with open(path, 'r', encoding='utf-8') as f:
+        return read_yaml_all_str(f.read())
+
+
+def dump_yaml(path: str,
+              config: Union[List[Dict[str, Any]], Dict[str, Any]],
+              blank: bool = False) -> None:
+    """Dumps a YAML file.
+
+    Args:
+        path: the path to the YAML file.
+        config: the configuration to dump.
+    """
+    with open(path, 'w', encoding='utf-8') as f:
+        contents = dump_yaml_str(config)
+        if blank and isinstance(config, dict) and len(config) == 0:
+            # when dumping to yaml, an empty dict will go in as {}.
+            contents = ''
+        f.write(contents)
+
+
+def dump_yaml_str(config: Union[List[Dict[str, Any]], Dict[str, Any]]) -> str:
+    """Dumps a YAML string.
+    Args:
+        config: the configuration to dump.
+    Returns:
+        The YAML string.
+    """
+
+    # https://github.com/yaml/pyyaml/issues/127
+    class LineBreakDumper(yaml.SafeDumper):
+
+        def write_line_break(self, data=None):
+            super().write_line_break(data)
+            if len(self.indents) == 1:
+                super().write_line_break()
+
+    if isinstance(config, list):
+        dump_func = yaml.dump_all  # type: ignore
+    else:
+        dump_func = yaml.dump  # type: ignore
+    return dump_func(config,
+                     Dumper=LineBreakDumper,
+                     sort_keys=False,
+                     default_flow_style=False)

sky/volumes/__init__.py

@@ -0,0 +1,13 @@
+"""Volumes."""
+
+from sky.volumes.client.sdk import apply
+from sky.volumes.client.sdk import delete
+from sky.volumes.client.sdk import ls
+from sky.volumes.volume import Volume
+
+__all__ = [
+    'apply',
+    'delete',
+    'ls',
+    'Volume',
+]

sky/volumes/client/sdk.py

@@ -0,0 +1,150 @@
+"""SDK functions for managed jobs."""
+import json
+import typing
+from typing import List
+
+from sky import exceptions
+from sky import sky_logging
+from sky.adaptors import common as adaptors_common
+from sky.schemas.api import responses
+from sky.server import common as server_common
+from sky.server import versions
+from sky.server.requests import payloads
+from sky.usage import usage_lib
+from sky.utils import annotations
+from sky.utils import context
+from sky.utils import ux_utils
+from sky.volumes import volume as volume_lib
+
+if typing.TYPE_CHECKING:
+    import requests
+else:
+    requests = adaptors_common.LazyImport('requests')
+
+logger = sky_logging.init_logger(__name__)
+
+
+@context.contextual
+@usage_lib.entrypoint
+@server_common.check_server_healthy_or_start
+@annotations.client_api
+def apply(volume: volume_lib.Volume) -> server_common.RequestId[None]:
+    """Creates or registers a volume.
+
+    Example:
+        .. code-block:: python
+
+            import sky.volumes
+            cfg = {
+                'name': 'pvc',
+                'type': 'k8s-pvc',
+                'size': '100GB',
+                'labels': {
+                    'key': 'value',
+                },
+            }
+            vol = sky.volumes.Volume.from_yaml_config(cfg)
+            request_id = sky.volumes.apply(vol)
+            sky.get(request_id)
+
+            or
+
+            import sky.volumes
+            vol = sky.volumes.Volume(
+                name='vol',
+                type='runpod-network-volume',
+                infra='runpod/ca/CA-MTL-1',
+                size='100GB',
+            )
+            request_id = sky.volumes.apply(vol)
+            sky.get(request_id)
+
+    Args:
+        volume: The volume to apply.
+
+    Returns:
+        The request ID of the apply request.
+    """
+    body = payloads.VolumeApplyBody(
+        name=volume.name,
+        volume_type=volume.type,
+        cloud=volume.cloud,
+        region=volume.region,
+        zone=volume.zone,
+        size=volume.size,
+        config=volume.config,
+        labels=volume.labels,
+        use_existing=volume.use_existing,
+    )
+    response = server_common.make_authenticated_request(
+        'POST', '/volumes/apply', json=json.loads(body.model_dump_json()))
+    return server_common.get_request_id(response)
+
+
+@context.contextual
+@usage_lib.entrypoint
+@server_common.check_server_healthy_or_start
+@annotations.client_api
+@versions.minimal_api_version(20)
+def validate(volume: volume_lib.Volume) -> None:
+    """Validates the volume.
+
+    All validation is done on the server side.
+
+    Args:
+        volume: The volume to validate.
+
+    Raises:
+        ValueError: If the volume is invalid.
+    """
+    body = payloads.VolumeValidateBody(
+        name=volume.name,
+        volume_type=volume.type,
+        infra=volume.infra,
+        size=volume.size,
+        config=volume.config,
+        labels=volume.labels,
+        use_existing=volume.use_existing,
+    )
+    response = server_common.make_authenticated_request(
+        'POST', '/volumes/validate', json=json.loads(body.model_dump_json()))
+    if response.status_code == 400:
+        with ux_utils.print_exception_no_traceback():
+            raise exceptions.deserialize_exception(
+                response.json().get('detail'))
+
+
+@context.contextual
+@usage_lib.entrypoint
+@server_common.check_server_healthy_or_start
+@annotations.client_api
+def ls() -> server_common.RequestId[List[responses.VolumeRecord]]:
+    """Lists all volumes.
+
+    Returns:
+        The request ID of the list request.
+    """
+    response = server_common.make_authenticated_request(
+        'GET',
+        '/volumes',
+    )
+    return server_common.get_request_id(response)
+
+
+@context.contextual
+@usage_lib.entrypoint
+@server_common.check_server_healthy_or_start
+@annotations.client_api
+def delete(names: List[str]) -> server_common.RequestId[None]:
+    """Deletes volumes.
+
+    Args:
+        names: List of volume names to delete.
+
+    Returns:
+        The request ID of the delete request.
+    """
+    body = payloads.VolumeDeleteBody(names=names)
+    response = server_common.make_authenticated_request(
+        'POST', '/volumes/delete', json=json.loads(body.model_dump_json()))
+    return server_common.get_request_id(response)