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

sky/volumes/server/core.py

@@ -0,0 +1,270 @@
+"""Volume management core."""
+
+import contextlib
+import os
+from typing import Any, Dict, Generator, List, Optional
+import uuid
+
+import filelock
+
+from sky import global_user_state
+from sky import models
+from sky import provision
+from sky import sky_logging
+from sky.schemas.api import responses
+from sky.utils import common_utils
+from sky.utils import registry
+from sky.utils import rich_utils
+from sky.utils import status_lib
+from sky.utils import ux_utils
+
+logger = sky_logging.init_logger(__name__)
+
+# Filelocks for the storage management.
+VOLUME_LOCK_PATH = os.path.expanduser('~/.sky/.{volume_name}.lock')
+VOLUME_LOCK_TIMEOUT_SECONDS = 20
+
+
+def volume_refresh():
+    """Refreshes the volume status."""
+    volumes = volume_list(is_ephemeral=False)
+    for volume in volumes:
+        volume_name = volume.name
+        usedby_pods = volume.usedby_pods
+        with _volume_lock(volume_name):
+            latest_volume = global_user_state.get_volume_by_name(volume_name)
+            if latest_volume is None:
+                logger.warning(f'Volume {volume_name} not found.')
+                continue
+            status = latest_volume.get('status')
+            if not usedby_pods:
+                if status != status_lib.VolumeStatus.READY:
+                    logger.info(f'Update volume {volume_name} '
+                                f'status to READY')
+                    global_user_state.update_volume_status(
+                        volume_name, status=status_lib.VolumeStatus.READY)
+            else:
+                if status != status_lib.VolumeStatus.IN_USE:
+                    logger.info(f'Update volume {volume_name} '
+                                f'status to IN_USE, usedby: {usedby_pods}')
+                    global_user_state.update_volume_status(
+                        volume_name, status=status_lib.VolumeStatus.IN_USE)
+
+
+def volume_list(
+        is_ephemeral: Optional[bool] = None) -> List[responses.VolumeRecord]:
+    """Gets the volumes.
+
+    Returns:
+        [
+            {
+                'name': str,
+                'type': str,
+                'launched_at': int timestamp of creation,
+                'cloud': str,
+                'region': str,
+                'zone': str,
+                'size': str,
+                'config': Dict[str, Any],
+                'name_on_cloud': str,
+                'user_hash': str,
+                'workspace': str,
+                'last_attached_at': int timestamp of last attachment,
+                'last_use': last command,
+                'status': sky.VolumeStatus,
+                'usedby_pods': List[str],
+                'usedby_clusters': List[str],
+                'is_ephemeral': bool,
+            }
+        ]
+    """
+    with rich_utils.safe_status(ux_utils.spinner_message('Listing volumes')):
+        volumes = global_user_state.get_volumes(is_ephemeral=is_ephemeral)
+        cloud_to_configs: Dict[str, List[models.VolumeConfig]] = {}
+        for volume in volumes:
+            config = volume.get('handle')
+            if config is None:
+                volume_name = volume.get('name')
+                logger.warning(f'Volume {volume_name} has no handle.')
+                continue
+            cloud = config.cloud
+            if cloud not in cloud_to_configs:
+                cloud_to_configs[cloud] = []
+            cloud_to_configs[cloud].append(config)
+
+        cloud_to_used_by_pods, cloud_to_used_by_clusters = {}, {}
+        for cloud, configs in cloud_to_configs.items():
+            used_by_pods, used_by_clusters = provision.get_all_volumes_usedby(
+                cloud, configs)
+            cloud_to_used_by_pods[cloud] = used_by_pods
+            cloud_to_used_by_clusters[cloud] = used_by_clusters
+
+        all_users = global_user_state.get_all_users()
+        user_map = {user.id: user.name for user in all_users}
+        records = []
+        for volume in volumes:
+            volume_name = volume.get('name')
+            record = {
+                'name': volume_name,
+                'launched_at': volume.get('launched_at'),
+                'user_hash': volume.get('user_hash'),
+                'user_name': user_map.get(volume.get('user_hash'), ''),
+                'workspace': volume.get('workspace'),
+                'last_attached_at': volume.get('last_attached_at'),
+                'last_use': volume.get('last_use'),
+                'usedby_pods': [],
+                'usedby_clusters': [],
+                'is_ephemeral': volume.get('is_ephemeral', False),
+            }
+            status = volume.get('status')
+            if status is not None:
+                record['status'] = status.value
+            else:
+                record['status'] = ''
+            config = volume.get('handle')
+            if config is None:
+                logger.warning(f'Volume {volume_name} has no handle.')
+                continue
+            cloud = config.cloud
+            usedby_pods, usedby_clusters = provision.map_all_volumes_usedby(
+                cloud,
+                cloud_to_used_by_pods[cloud],
+                cloud_to_used_by_clusters[cloud],
+                config,
+            )
+            record['type'] = config.type
+            record['cloud'] = config.cloud
+            record['region'] = config.region
+            record['zone'] = config.zone
+            record['size'] = config.size
+            record['config'] = config.config
+            record['name_on_cloud'] = config.name_on_cloud
+            record['usedby_pods'] = usedby_pods
+            record['usedby_clusters'] = usedby_clusters
+            records.append(responses.VolumeRecord(**record))
+        return records
+
+
+def volume_delete(names: List[str], ignore_not_found: bool = False) -> None:
+    """Deletes volumes.
+
+    Args:
+        names: List of volume names to delete.
+        ignore_not_found: If True, ignore volumes that are not found.
+
+    Raises:
+        ValueError: If the volume does not exist
+          or is in use or has no handle.
+    """
+    with rich_utils.safe_status(ux_utils.spinner_message('Deleting volumes')):
+        for name in names:
+            volume = global_user_state.get_volume_by_name(name)
+            if volume is None:
+                if ignore_not_found:
+                    continue
+                raise ValueError(f'Volume {name} not found.')
+            config = volume.get('handle')
+            if config is None:
+                raise ValueError(f'Volume {name} has no handle.')
+            cloud = config.cloud
+            usedby_pods, usedby_clusters = provision.get_volume_usedby(
+                cloud, config)
+            if usedby_clusters:
+                usedby_clusters_str = ', '.join(usedby_clusters)
+                cluster_str = 'clusters' if len(
+                    usedby_clusters) > 1 else 'cluster'
+                raise ValueError(f'Volume {name} is used by {cluster_str}'
+                                 f' {usedby_clusters_str}.')
+            if usedby_pods:
+                usedby_pods_str = ', '.join(usedby_pods)
+                pod_str = 'pods' if len(usedby_pods) > 1 else 'pod'
+                raise ValueError(
+                    f'Volume {name} is used by {pod_str} {usedby_pods_str}.')
+            logger.debug(f'Deleting volume {name} with config {config}')
+            with _volume_lock(name):
+                provision.delete_volume(cloud, config)
+                global_user_state.delete_volume(name)
+        logger.info(f'Deleted volumes: {names}')
+
+
+def volume_apply(
+    name: str,
+    volume_type: str,
+    cloud: str,
+    region: Optional[str],
+    zone: Optional[str],
+    size: Optional[str],
+    config: Dict[str, Any],
+    labels: Optional[Dict[str, str]] = None,
+    use_existing: Optional[bool] = None,
+    is_ephemeral: bool = False,
+) -> None:
+    """Creates or registers a volume.
+
+    Args:
+        name: The name of the volume.
+        volume_type: The type of the volume.
+        cloud: The cloud of the volume.
+        region: The region of the volume.
+        zone: The zone of the volume.
+        size: The size of the volume.
+        config: The configuration of the volume.
+        labels: The labels of the volume.
+        use_existing: Whether to use an existing volume.
+        is_ephemeral: Whether the volume is ephemeral.
+    """
+    with rich_utils.safe_status(ux_utils.spinner_message('Creating volume')):
+        # Reuse the method for cluster name on cloud to
+        # generate the storage name on cloud.
+        cloud_obj = registry.CLOUD_REGISTRY.from_str(cloud)
+        assert cloud_obj is not None
+        region, zone = cloud_obj.validate_region_zone(region, zone)
+        if use_existing:
+            name_on_cloud = name
+        else:
+            name_uuid = str(uuid.uuid4())[:6]
+            name_on_cloud = common_utils.make_cluster_name_on_cloud(
+                name, max_length=cloud_obj.max_cluster_name_length())
+            name_on_cloud += '-' + name_uuid
+        config = models.VolumeConfig(
+            name=name,
+            type=volume_type,
+            cloud=str(cloud_obj),
+            region=region,
+            zone=zone,
+            size=size,
+            config=config,
+            name_on_cloud=name_on_cloud,
+            labels=labels,
+        )
+        logger.debug(
+            f'Creating volume {name} on cloud {cloud} with config {config}')
+        with _volume_lock(name):
+            current_volume = global_user_state.get_volume_by_name(name)
+            if current_volume is not None:
+                logger.info(f'Volume {name} already exists.')
+                return
+            config = provision.apply_volume(cloud, config)
+            global_user_state.add_volume(
+                name,
+                config,
+                status_lib.VolumeStatus.READY,
+                is_ephemeral,
+            )
+        logger.info(f'Created volume {name} on cloud {cloud}')
+
+
+@contextlib.contextmanager
+def _volume_lock(volume_name: str) -> Generator[None, None, None]:
+    """Context manager for volume lock."""
+    try:
+        with filelock.FileLock(VOLUME_LOCK_PATH.format(volume_name=volume_name),
+                               VOLUME_LOCK_TIMEOUT_SECONDS):
+            yield
+    except filelock.Timeout as e:
+        raise RuntimeError(
+            f'Failed to update user due to a timeout '
+            f'when trying to acquire the lock at '
+            f'{VOLUME_LOCK_PATH.format(volume_name=volume_name)}. '
+            'Please try again or manually remove the lock '
+            f'file if you believe it is stale.') from e

sky/volumes/server/server.py

@@ -0,0 +1,124 @@
+"""REST API for storage management."""
+
+import fastapi
+
+from sky import clouds
+from sky import exceptions
+from sky import sky_logging
+from sky.server.requests import executor
+from sky.server.requests import payloads
+from sky.server.requests import request_names
+from sky.server.requests import requests as requests_lib
+from sky.utils import registry
+from sky.utils import volume as volume_utils
+from sky.volumes.server import core
+
+logger = sky_logging.init_logger(__name__)
+
+router = fastapi.APIRouter()
+
+
+@router.get('')
+async def volume_list(request: fastapi.Request) -> None:
+    """Gets the volumes."""
+    auth_user = request.state.auth_user
+    auth_user_env_vars_kwargs = {
+        'env_vars': auth_user.to_env_vars()
+    } if auth_user else {}
+    request_body = payloads.RequestBody(**auth_user_env_vars_kwargs)
+    await executor.schedule_request_async(
+        request_id=request.state.request_id,
+        request_name=request_names.RequestName.VOLUME_LIST,
+        request_body=request_body,
+        func=core.volume_list,
+        schedule_type=requests_lib.ScheduleType.SHORT,
+    )
+
+
+@router.post('/delete')
+async def volume_delete(request: fastapi.Request,
+                        volume_delete_body: payloads.VolumeDeleteBody) -> None:
+    """Deletes a volume."""
+    await executor.schedule_request_async(
+        request_id=request.state.request_id,
+        request_name=request_names.RequestName.VOLUME_DELETE,
+        request_body=volume_delete_body,
+        func=core.volume_delete,
+        schedule_type=requests_lib.ScheduleType.LONG,
+    )
+
+
+@router.post('/validate')
+async def volume_validate(
+        _: fastapi.Request,
+        volume_validate_body: payloads.VolumeValidateBody) -> None:
+    """Validates a volume."""
+    # pylint: disable=import-outside-toplevel
+    from sky.volumes import volume as volume_lib
+
+    try:
+        volume_config = {
+            'name': volume_validate_body.name,
+            'type': volume_validate_body.volume_type,
+            'infra': volume_validate_body.infra,
+            'size': volume_validate_body.size,
+            'labels': volume_validate_body.labels,
+            'config': volume_validate_body.config,
+            'use_existing': volume_validate_body.use_existing,
+        }
+        volume = volume_lib.Volume.from_yaml_config(volume_config)
+        volume.validate()
+    except Exception as e:
+        requests_lib.set_exception_stacktrace(e)
+        raise fastapi.HTTPException(status_code=400,
+                                    detail=exceptions.serialize_exception(e))
+
+
+@router.post('/apply')
+async def volume_apply(request: fastapi.Request,
+                       volume_apply_body: payloads.VolumeApplyBody) -> None:
+    """Creates or registers a volume."""
+    volume_cloud = volume_apply_body.cloud
+    volume_type = volume_apply_body.volume_type
+    volume_config = volume_apply_body.config
+    if volume_config is None:
+        volume_config = {}
+    volume_config['use_existing'] = volume_apply_body.use_existing
+
+    supported_volume_types = [
+        volume_type.value for volume_type in volume_utils.VolumeType
+    ]
+    if volume_type not in supported_volume_types:
+        raise fastapi.HTTPException(
+            status_code=400, detail=f'Invalid volume type: {volume_type}')
+    cloud = registry.CLOUD_REGISTRY.from_str(volume_cloud)
+    if cloud is None:
+        raise fastapi.HTTPException(status_code=400,
+                                    detail=f'Invalid cloud: {volume_cloud}')
+    if volume_type == volume_utils.VolumeType.PVC.value:
+        if not cloud.is_same_cloud(clouds.Kubernetes()):
+            raise fastapi.HTTPException(
+                status_code=400,
+                detail='PVC storage is only supported on Kubernetes')
+        supported_access_modes = [
+            access_mode.value for access_mode in volume_utils.VolumeAccessMode
+        ]
+        access_mode = volume_config.get('access_mode')
+        if access_mode is None:
+            volume_config['access_mode'] = (
+                volume_utils.VolumeAccessMode.READ_WRITE_ONCE.value)
+        elif access_mode not in supported_access_modes:
+            raise fastapi.HTTPException(
+                status_code=400, detail=f'Invalid access mode: {access_mode}')
+    elif volume_type == volume_utils.VolumeType.RUNPOD_NETWORK_VOLUME.value:
+        if not cloud.is_same_cloud(clouds.RunPod()):
+            raise fastapi.HTTPException(
+                status_code=400,
+                detail='Runpod network volume is only supported on Runpod')
+    await executor.schedule_request_async(
+        request_id=request.state.request_id,
+        request_name=request_names.RequestName.VOLUME_APPLY,
+        request_body=volume_apply_body,
+        func=core.volume_apply,
+        schedule_type=requests_lib.ScheduleType.LONG,
+    )

sky/volumes/volume.py

@@ -0,0 +1,215 @@
+"""Volume types and access modes."""
+from typing import Any, Dict, Optional
+
+from sky import clouds
+from sky.utils import common_utils
+from sky.utils import infra_utils
+from sky.utils import registry
+from sky.utils import resources_utils
+from sky.utils import schemas
+from sky.utils import volume as volume_lib
+
+VOLUME_TYPE_TO_CLOUD = {
+    volume_lib.VolumeType.PVC: clouds.Kubernetes(),
+    volume_lib.VolumeType.RUNPOD_NETWORK_VOLUME: clouds.RunPod(),
+}
+CLOUD_TO_VOLUME_TYPE = {
+    clouds.Kubernetes(): [volume_lib.VolumeType.PVC],
+    clouds.RunPod(): [volume_lib.VolumeType.RUNPOD_NETWORK_VOLUME],
+}
+
+
+class Volume:
+    """Volume specification."""
+
+    def __init__(
+            self,
+            name: Optional[str] = None,
+            type: Optional[str] = None,  # pylint: disable=redefined-builtin
+            infra: Optional[str] = None,
+            size: Optional[str] = None,
+            labels: Optional[Dict[str, str]] = None,
+            use_existing: Optional[bool] = None,
+            config: Optional[Dict[str, Any]] = None):
+        """Initialize a Volume instance.
+
+        Args:
+            name: Volume name
+            type: Volume type (e.g., 'k8s-pvc')
+            infra: Infrastructure specification
+            size: Volume size
+            labels: Volume labels
+            use_existing: Whether to use an existing volume
+            config: Additional configuration
+        """
+        self.name = name
+        self.type = type
+        self.infra = infra
+        self.size = size
+        self.labels = labels or {}
+        self.use_existing = use_existing
+        self.config = config or {}
+
+        self.cloud: Optional[str] = None
+        self.region: Optional[str] = None
+        self.zone: Optional[str] = None
+
+        self._normalize_config()
+
+    @classmethod
+    def from_yaml_config(cls, config: Dict[str, Any]) -> 'Volume':
+        """Create a Volume subclass instance from a dictionary via factory."""
+        vol_type_val = config.get('type')
+        try:
+            vt = (volume_lib.VolumeType(vol_type_val)
+                  if vol_type_val is not None else None)
+        except Exception:  # pylint: disable=broad-except
+            vt = None
+
+        if vt is None:
+            raise ValueError(f'Invalid volume type: {vol_type_val}')
+
+        if vt == volume_lib.VolumeType.PVC:
+            return PVCVolume(name=config.get('name'),
+                             type=vol_type_val,
+                             infra=config.get('infra'),
+                             size=config.get('size'),
+                             labels=config.get('labels'),
+                             use_existing=config.get('use_existing'),
+                             config=config.get('config', {}))
+        if vt == volume_lib.VolumeType.RUNPOD_NETWORK_VOLUME:
+            return RunpodNetworkVolume(name=config.get('name'),
+                                       type=vol_type_val,
+                                       infra=config.get('infra'),
+                                       size=config.get('size'),
+                                       labels=config.get('labels'),
+                                       use_existing=config.get('use_existing'),
+                                       config=config.get('config', {}))
+
+        raise ValueError(f'Invalid volume type: {vol_type_val}')
+
+    def to_yaml_config(self) -> Dict[str, Any]:
+        """Convert the Volume to a dictionary."""
+        return {
+            'name': self.name,
+            'type': self.type,
+            'infra': self.infra,
+            'size': self.size,
+            'labels': self.labels,
+            'use_existing': self.use_existing,
+            'config': self.config,
+            'cloud': self.cloud,
+            'region': self.region,
+            'zone': self.zone,
+        }
+
+    def _normalize_config(self) -> None:
+        """Normalize and validate the config."""
+        # Validate schema
+        common_utils.validate_schema(self.to_yaml_config(),
+                                     schemas.get_volume_schema(),
+                                     'Invalid volumes config: ')
+
+        # Adjust the volume config (e.g., parse size)
+        self._adjust_config()
+
+        # Resolve the infrastructure options to cloud, region, zone
+        infra_info = infra_utils.InfraInfo.from_str(self.infra)
+        self.cloud = infra_info.cloud
+        self.region = infra_info.region
+        self.zone = infra_info.zone
+
+        # Set cloud from volume type if not specified
+        cloud_obj_from_type = VOLUME_TYPE_TO_CLOUD.get(
+            volume_lib.VolumeType(self.type))
+        if self.cloud:
+            cloud_obj = registry.CLOUD_REGISTRY.from_str(self.cloud)
+            assert cloud_obj is not None
+            if not cloud_obj.is_same_cloud(cloud_obj_from_type):
+                raise ValueError(
+                    f'Invalid cloud {self.cloud} for volume type {self.type}')
+        else:
+            self.cloud = str(cloud_obj_from_type)
+
+    def _adjust_config(self) -> None:
+        """Adjust the volume config (e.g., parse size)."""
+        if self.size is None:
+            return
+        try:
+            size = resources_utils.parse_memory_resource(self.size,
+                                                         'size',
+                                                         allow_rounding=True)
+            if size == '0':
+                raise ValueError('Size must be no less than 1Gi')
+            self.size = size
+        except ValueError as e:
+            raise ValueError(f'Invalid size {self.size}: {e}') from e
+
+    def validate(self, skip_cloud_compatibility: bool = False) -> None:
+        """Validates the volume."""
+        self.validate_name()
+        self.validate_size()
+        if not skip_cloud_compatibility:
+            self.validate_cloud_compatibility()
+        # Extra, type-specific validations
+        self._validate_config_extra()
+
+    def validate_name(self) -> None:
+        """Validates if the volume name is set."""
+        assert self.name is not None, 'Volume name must be set'
+
+    def validate_size(self) -> None:
+        """Validates that size is specified for new volumes."""
+        if not self.use_existing and not self.size:
+            raise ValueError('Size is required for new volumes. '
+                             'Please specify the size in the YAML file or '
+                             'use the --size flag.')
+
+    def validate_cloud_compatibility(self) -> None:
+        """Validates region, zone, name, labels with the cloud."""
+        cloud_obj = registry.CLOUD_REGISTRY.from_str(self.cloud)
+        assert cloud_obj is not None
+
+        valid, err_msg = cloud_obj.is_volume_name_valid(self.name)
+        if not valid:
+            raise ValueError(f'Invalid volume name: {err_msg}')
+
+        if self.labels:
+            for key, value in self.labels.items():
+                valid, err_msg = cloud_obj.is_label_valid(key, value)
+                if not valid:
+                    raise ValueError(f'{err_msg}')
+
+    # Hook methods for subclasses
+    def _validate_config_extra(self) -> None:
+        """Additional type-specific validation.
+
+        Subclasses can override to enforce stricter rules.
+        """
+        return
+
+
+class PVCVolume(Volume):
+    """Kubernetes PVC-backed volume."""
+    pass
+
+
+class RunpodNetworkVolume(Volume):
+    """RunPod Network Volume."""
+
+    def _validate_config_extra(self) -> None:
+        if not self.use_existing and self.size is not None:
+            try:
+                size_int = int(self.size)
+                if size_int < volume_lib.MIN_RUNPOD_NETWORK_VOLUME_SIZE_GB:
+                    raise ValueError(
+                        f'RunPod network volume size must be at least '
+                        f'{volume_lib.MIN_RUNPOD_NETWORK_VOLUME_SIZE_GB}GB.')
+            except Exception as e:  # pylint: disable=broad-except
+                raise ValueError(f'Invalid volume size {self.size!r}: '
+                                 f'{e}') from e
+        if not self.zone:
+            raise ValueError('RunPod DataCenterId is required for network '
+                             'volumes. Set the zone in the infra field.')
+
+        return