skypilot-nightly 1.0.0.dev20250624__py3-none-any.whl → 1.0.0.dev20250626__py3-none-any.whl

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.'

sky/volumes/volume.py

@@ -0,0 +1,198 @@
+"""Volume types and access modes."""
+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 infra_utils
+from sky.utils import resources_utils
+from sky.utils import schemas
+from sky.utils import status_lib
+
+
+class VolumeType(enum.Enum):
+    """Volume type."""
+    PVC = 'k8s-pvc'
+
+
+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 VolumeMount:
+    """Volume mount specification."""
+
+    def __init__(self, path: str, volume_name: str,
+                 volume_config: models.VolumeConfig):
+        self.path: str = path
+        self.volume_name: str = volume_name
+        self.volume_config: models.VolumeConfig = volume_config
+
+    def pre_mount(self) -> None:
+        """Update the volume status before actual mounting."""
+        # 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)
+        volume_config: models.VolumeConfig = models.VolumeConfig.model_validate(
+            config.pop('volume_config', None))
+        return cls(path, volume_name, volume_config)
+
+    def to_yaml_config(self) -> Dict[str, Any]:
+        return {
+            'path': self.path,
+            'volume_name': self.volume_name,
+            'volume_config': self.volume_config.model_dump(),
+        }
+
+    def __repr__(self):
+        return (f'VolumeMount('
+                f'\n\tpath={self.path},'
+                f'\n\tvolume_name={self.volume_name},'
+                f'\n\tvolume_config={self.volume_config})')
+
+
+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,
+            resource_name: Optional[str] = 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
+            config: Additional configuration
+        """
+        self.name = name
+        self.type = type
+        self.infra = infra
+        self.size = size
+        self.resource_name = resource_name
+        self.config = config or {}
+
+        self.cloud: Optional[str] = None
+        self.region: Optional[str] = None
+        self.zone: Optional[str] = None
+
+    @classmethod
+    def from_dict(cls, config_dict: Dict[str, Any]) -> 'Volume':
+        """Create a Volume instance from a dictionary."""
+        return cls(name=config_dict.get('name'),
+                   type=config_dict.get('type'),
+                   infra=config_dict.get('infra'),
+                   size=config_dict.get('size'),
+                   resource_name=config_dict.get('resource_name'),
+                   config=config_dict.get('config', {}))
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert the Volume to a dictionary."""
+        return {
+            'name': self.name,
+            'type': self.type,
+            'infra': self.infra,
+            'size': self.size,
+            'resource_name': self.resource_name,
+            'config': self.config,
+            'cloud': self.cloud,
+            'region': self.region,
+            'zone': self.zone,
+        }
+
+    def normalize_config(
+            self,
+            name: Optional[str] = None,
+            infra: Optional[str] = None,
+            type: Optional[str] = None,  # pylint: disable=redefined-builtin
+            size: Optional[str] = None) -> None:
+        """Override the volume config with CLI options,
+           adjust and validate the config.
+
+        Args:
+            name: Volume name to override
+            infra: Infrastructure to override
+            type: Volume type to override
+            size: Volume size to override
+        """
+        if name is not None:
+            self.name = name
+        if infra is not None:
+            self.infra = infra
+        if type is not None:
+            self.type = type
+        if size is not None:
+            self.size = size
+
+        # Validate schema
+        common_utils.validate_schema(self.to_dict(),
+                                     schemas.get_volume_schema(),
+                                     'Invalid volumes config: ')
+
+        # Adjust the volume config (e.g., parse size)
+        self._adjust_config()
+
+        # Validate the volume config
+        self._validate_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
+
+    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_config(self) -> None:
+        """Validate the volume config."""
+        if not self.resource_name 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.')

{skypilot_nightly-1.0.0.dev20250624.dist-info → skypilot_nightly-1.0.0.dev20250626.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: skypilot-nightly
-Version: 1.0.0.dev20250624
+Version: 1.0.0.dev20250626
 Summary: SkyPilot: Run AI on Any Infra — Unified, Faster, Cheaper.
 Author: SkyPilot Team
 License: Apache 2.0
@@ -51,6 +51,7 @@ Requires-Dist: sqlalchemy
 Requires-Dist: psycopg2-binary
 Requires-Dist: casbin
 Requires-Dist: sqlalchemy_adapter
+Requires-Dist: prometheus_client>=0.8.0
 Requires-Dist: passlib
 Provides-Extra: aws
 Requires-Dist: awscli>=1.27.10; extra == "aws"