skypilot-nightly 1.0.0.dev20250509__py3-none-any.whl → 1.0.0.dev20251107__py3-none-any.whl

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,149 @@
+"""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,
+    )
+    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,
+        resource_name=volume.resource_name,
+        size=volume.size,
+        config=volume.config,
+        labels=volume.labels,
+    )
+    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)

sky/volumes/server/core.py

@@ -0,0 +1,258 @@
+"""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 = 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_pods, _ = 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_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() -> 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],
+            }
+        ]
+    """
+    with rich_utils.safe_status(ux_utils.spinner_message('Listing volumes')):
+        volumes = global_user_state.get_volumes()
+        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': [],
+            }
+            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]) -> 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.
+    """
+    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:
+                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)
+
+
+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,
+) -> 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.
+
+    """
+    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)
+        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)
+
+
+@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,122 @@
+"""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,
+            'resource_name': volume_validate_body.resource_name,
+        }
+        volume = volume_lib.Volume.from_yaml_config(volume_config)
+        volume.validate()
+    except Exception as 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
+
+    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
+        ]
+        if volume_config is None:
+            volume_config = {}
+        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,
+    )