skypilot-nightly 1.0.0.dev20250623__py3-none-any.whl → 1.0.0.dev20250625__py3-none-any.whl

sky/utils/status_lib.py

@@ -54,3 +54,13 @@ class StorageStatus(enum.Enum):
 
     # Finished uploading, in terminal state
     READY = 'READY'
+
+
+class VolumeStatus(enum.Enum):
+    """Volume status as recorded in table 'volumes'."""
+
+    # Volume is ready to be used
+    READY = 'READY'
+
+    # Volume is being used
+    IN_USE = 'IN_USE'

sky/utils/validator.py

@@ -14,9 +14,19 @@ def case_insensitive_enum(validator, enums, instance, schema):
             f'{instance!r} is not one of {enums!r}')
 
 
+def case_sensitive_enum(validator, enums, instance, schema):
+    del validator, schema  # Unused.
+    if instance not in enums:
+        yield jsonschema.ValidationError(
+            f'{instance!r} is not one of {enums!r}')
+
+
 # Move this to a function to delay initialization
 def get_schema_validator():
     """Get the schema validator class, initializing it only when needed."""
     return jsonschema.validators.extend(
         jsonschema.Draft7Validator,
-        validators={'case_insensitive_enum': case_insensitive_enum})
+        validators={
+            'case_insensitive_enum': case_insensitive_enum,
+            'case_sensitive_enum': case_sensitive_enum
+        })

sky/volumes/client/sdk.py

@@ -0,0 +1,64 @@
+"""SDK functions for managed jobs."""
+import json
+import typing
+from typing import List
+
+from sky import sky_logging
+from sky.adaptors import common as adaptors_common
+from sky.server import common as server_common
+from sky.server.requests import payloads
+from sky.usage import usage_lib
+from sky.utils import annotations
+from sky.utils import context
+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:
+    """Creates or registers a volume.
+    """
+    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)
+    response = requests.post(f'{server_common.get_server_url()}/volumes/apply',
+                             json=json.loads(body.model_dump_json()),
+                             cookies=server_common.get_api_cookie_jar())
+    return server_common.get_request_id(response)
+
+
+@context.contextual
+@usage_lib.entrypoint
+@server_common.check_server_healthy_or_start
+@annotations.client_api
+def ls() -> server_common.RequestId:
+    """Lists all volumes."""
+    response = requests.get(f'{server_common.get_server_url()}/volumes',
+                            cookies=server_common.get_api_cookie_jar())
+    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:
+    """Deletes a volume."""
+    body = payloads.VolumeDeleteBody(names=names)
+    response = requests.post(f'{server_common.get_server_url()}/volumes/delete',
+                             json=json.loads(body.model_dump_json()),
+                             cookies=server_common.get_api_cookie_jar())
+    return server_common.get_request_id(response)

sky/volumes/server/core.py

@@ -0,0 +1,199 @@
+"""Volume management core."""
+
+import contextlib
+import os
+from typing import Any, Dict, Generator, List, Optional
+import uuid
+
+import filelock
+
+import sky
+from sky import global_user_state
+from sky import models
+from sky import provision
+from sky import sky_logging
+from sky.utils import common_utils
+from sky.utils import status_lib
+
+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 = global_user_state.get_volumes()
+    for volume in volumes:
+        volume_name = volume.get('name')
+        config = volume.get('handle')
+        if config is None:
+            logger.warning(f'Volume {volume_name} has no handle.'
+                           'Skipping status refresh...')
+            continue
+        cloud = config.cloud
+        usedby = provision.get_volume_usedby(cloud, config)
+        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:
+                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}')
+                    global_user_state.update_volume_status(
+                        volume_name, status=status_lib.VolumeStatus.IN_USE)
+
+
+def volume_list() -> List[Dict[str, Any]]:
+    """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,
+            }
+        ]
+    """
+    volumes = global_user_state.get_volumes()
+    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'),
+        }
+        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
+        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
+        records.append(record)
+    return records
+
+
+def volume_delete(names: List[str]) -> None:
+    """Deletes volumes.
+
+    Args:
+        names: List of volume names to delete.
+
+    Raises:
+        ValueError: If the volume does not exist
+          or is in use or has no handle.
+    """
+    for name in names:
+        volume = global_user_state.get_volume_by_name(name)
+        if volume is None:
+            raise ValueError(f'Volume {name} not found.')
+        if volume.get('status') == status_lib.VolumeStatus.IN_USE:
+            raise ValueError(f'Volume {name} is in use.')
+        config = volume.get('handle')
+        if config is None:
+            raise ValueError(f'Volume {name} has no handle.')
+        logger.debug(f'Deleting volume {name} with config {config}')
+        cloud = config.cloud
+        with _volume_lock(name):
+            provision.delete_volume(cloud, config)
+            global_user_state.delete_volume(name)
+
+
+def volume_apply(name: str, volume_type: str, cloud: str, region: Optional[str],
+                 zone: Optional[str], size: Optional[str],
+                 config: Dict[str, Any]) -> 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.
+
+    """
+    # Reuse the method for cluster name on cloud to
+    # generate the storage name on cloud.
+    cloud_obj = sky.CLOUD_REGISTRY.from_str(cloud)
+    assert cloud_obj is not None
+    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,
+    )
+    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)
+
+
+@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,85 @@
+"""REST API for storage management."""
+
+import fastapi
+
+import sky
+from sky import clouds
+from sky import sky_logging
+from sky.server.requests import executor
+from sky.server.requests import payloads
+from sky.server.requests import requests as requests_lib
+from sky.volumes import volume
+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."""
+    executor.schedule_request(
+        request_id=request.state.request_id,
+        request_name='volume_list',
+        request_body=payloads.RequestBody(),
+        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."""
+    executor.schedule_request(
+        request_id=request.state.request_id,
+        request_name='volume_delete',
+        request_body=volume_delete_body,
+        func=core.volume_delete,
+        schedule_type=requests_lib.ScheduleType.LONG,
+    )
+
+
+@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
+
+    supported_volume_types = [
+        volume_type.value for volume_type in volume.VolumeType
+    ]
+    if volume_type not in supported_volume_types:
+        raise fastapi.HTTPException(
+            status_code=400, detail=f'Invalid volume type: {volume_type}')
+    cloud = sky.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.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.VolumeAccessMode
+        ]
+        if volume_config is None:
+            volume_config = {}
+        access_mode = volume_config.get('access_mode')
+        if access_mode is None:
+            volume_config[
+                'access_mode'] = volume.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}')
+    executor.schedule_request(
+        request_id=request.state.request_id,
+        request_name='volume_apply',
+        request_body=volume_apply_body,
+        func=core.volume_apply,
+        schedule_type=requests_lib.ScheduleType.LONG,
+    )

sky/volumes/utils.py

@@ -0,0 +1,158 @@
+"""Volume utils."""
+import abc
+from datetime import datetime
+from typing import Any, Dict, List, Optional
+
+import prettytable
+
+from sky import sky_logging
+from sky.utils import log_utils
+from sky.volumes import volume
+
+logger = sky_logging.init_logger(__name__)
+
+
+def _get_infra_str(cloud: Optional[str], region: Optional[str],
+                   zone: Optional[str]) -> str:
+    """Get the infrastructure string for the volume."""
+    infra = ''
+    if cloud:
+        infra += cloud
+    if region:
+        infra += f'/{region}'
+    if zone:
+        infra += f'/{zone}'
+    return infra
+
+
+class VolumeTable(abc.ABC):
+    """The volume table."""
+
+    @abc.abstractmethod
+    def format(self) -> str:
+        """Format the volume table for display."""
+        pass
+
+
+class PVCVolumeTable(VolumeTable):
+    """The PVC volume table."""
+
+    def __init__(self, volumes: List[Dict[str, Any]], show_all: bool = False):
+        super().__init__()
+        self.table = self._create_table(show_all)
+        self._add_rows(volumes, show_all)
+
+    def _create_table(self, show_all: bool = False) -> prettytable.PrettyTable:
+        """Create the PVC volume table."""
+        #  If show_all is True, show the table with the columns:
+        #   NAME, TYPE, INFRA, SIZE, USER, WORKSPACE,
+        #   AGE, LAST_USE, STATUS
+        #  If show_all is False, show the table with the columns:
+        #   NAME, TYPE, INFRA, SIZE, USER, WORKSPACE,
+        #   AGE, LAST_USE, STATUS, NAME_ON_CLOUD,
+        #   STORAGE_CLASS, ACCESS_MODE
+
+        if show_all:
+            columns = [
+                'NAME',
+                'TYPE',
+                'INFRA',
+                'SIZE',
+                'USER',
+                'WORKSPACE',
+                'AGE',
+                'STATUS',
+                'LAST_USE',
+                'NAME_ON_CLOUD',
+                'STORAGE_CLASS',
+                'ACCESS_MODE',
+            ]
+        else:
+            columns = [
+                'NAME',
+                'TYPE',
+                'INFRA',
+                'SIZE',
+                'USER',
+                'WORKSPACE',
+                'AGE',
+                'STATUS',
+                'LAST_USE',
+            ]
+
+        table = log_utils.create_table(columns)
+        return table
+
+    def _add_rows(self,
+                  volumes: List[Dict[str, Any]],
+                  show_all: bool = False) -> None:
+        """Add rows to the PVC volume table."""
+        for row in volumes:
+            # Convert last_attached_at timestamp to human readable string
+            last_attached_at = row.get('last_attached_at')
+            if last_attached_at is not None:
+                last_attached_at_str = datetime.fromtimestamp(
+                    last_attached_at).strftime('%Y-%m-%d %H:%M:%S')
+            else:
+                last_attached_at_str = '-'
+            size = row.get('size', '')
+            if size:
+                size = f'{size}Gi'
+            infra = _get_infra_str(row.get('cloud'), row.get('region'),
+                                   row.get('zone'))
+            table_row = [
+                row.get('name', ''),
+                row.get('type', ''),
+                infra,
+                size,
+                row.get('user_name', '-'),
+                row.get('workspace', '-'),
+                log_utils.readable_time_duration(row.get('launched_at', 0)),
+                row.get('status', ''),
+                last_attached_at_str,
+            ]
+            if show_all:
+                table_row.append(row.get('name_on_cloud', ''))
+                table_row.append(
+                    row.get('config', {}).get('storage_class_name', '-'))
+                table_row.append(row.get('config', {}).get('access_mode', ''))
+
+            self.table.add_row(table_row)
+
+    def format(self) -> str:
+        """Format the PVC volume table for display."""
+        return str(self.table)
+
+
+def format_volume_table(volumes: List[Dict[str, Any]],
+                        show_all: bool = False) -> str:
+    """Format the volume table for display.
+
+    Args:
+        volume_table (dict): The volume table.
+
+    Returns:
+        str: The formatted volume table.
+    """
+    volumes_per_type: Dict[str, List[Dict[str, Any]]] = {}
+    supported_volume_types = [
+        volume_type.value for volume_type in volume.VolumeType
+    ]
+    for row in volumes:
+        volume_type = row.get('type', '')
+        if volume_type in supported_volume_types:
+            if volume_type not in volumes_per_type:
+                volumes_per_type[volume_type] = []
+            volumes_per_type[volume_type].append(row)
+        else:
+            logger.warning(f'Unknown volume type: {volume_type}')
+            continue
+    table_str = ''
+    for volume_type, volume_list in volumes_per_type.items():
+        if volume_type == volume.VolumeType.PVC.value:
+            table = PVCVolumeTable(volume_list, show_all)
+            table_str += table.format()
+    if table_str:
+        return table_str
+    else:
+        return 'No existing volumes.'