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

sky/utils/rich_utils.py

@@ -7,6 +7,7 @@ import threading
 import typing
 from typing import Callable, Iterator, Optional, Tuple, Union
 
+from sky import exceptions
 from sky.adaptors import common as adaptors_common
 from sky.utils import annotations
 from sky.utils import context
@@ -58,6 +59,7 @@ class Control(enum.Enum):
     EXIT = 'rich_exit'
     UPDATE = 'rich_update'
     HEARTBEAT = 'heartbeat'
+    RETRY = 'retry'
 
     def encode(self, msg: str) -> str:
         return f'<{self.value}>{msg}</{self.value}>'
@@ -365,6 +367,10 @@ def decode_rich_status(
                     yield line
                     continue
 
+                if control == Control.RETRY:
+                    raise exceptions.ServerTemporarilyUnavailableError(
+                        'The server is temporarily unavailable. Please try '
+                        'again.')
                 # control is not None, i.e. it is a rich status control message.
                 if threading.current_thread() is not threading.main_thread():
                     yield None

sky/utils/schemas.py

@@ -70,8 +70,36 @@ _AUTOSTOP_SCHEMA = {
 }
 
 
-def _get_single_resources_schema():
-    """Schema for a single resource in a resources list."""
+# Note: This is similar to _get_infra_pattern()
+# but without the wildcard patterns.
+def _get_volume_infra_pattern():
+    # Building the regex pattern for the infra field
+    # Format: cloud[/region[/zone]] or wildcards or kubernetes context
+    # Match any cloud name (case insensitive)
+    all_clouds = list(constants.ALL_CLOUDS)
+    all_clouds.remove('kubernetes')
+    cloud_pattern = f'(?i:({"|".join(all_clouds)}))'
+
+    # Optional /region followed by optional /zone
+    # /[^/]+ matches a slash followed by any characters except slash (region or
+    # zone name)
+    # The outer (?:...)? makes the entire region/zone part optional
+    region_zone_pattern = '(?:/[^/]+(?:/[^/]+)?)?'
+
+    # Kubernetes specific pattern - matches:
+    # 1. Just the word "kubernetes" or "k8s" by itself
+    # 2. "k8s/" or "kubernetes/" followed by any context name (which may contain
+    # slashes)
+    kubernetes_pattern = '(?i:kubernetes|k8s)(?:/.+)?'
+
+    # Combine all patterns with alternation (|)
+    # ^ marks start of string, $ marks end of string
+    infra_pattern = (f'^(?:{cloud_pattern}{region_zone_pattern}|'
+                     f'{kubernetes_pattern})$')
+    return infra_pattern
+
+
+def _get_infra_pattern():
     # Building the regex pattern for the infra field
     # Format: cloud[/region[/zone]] or wildcards or kubernetes context
     # Match any cloud name (case insensitive)
@@ -103,7 +131,11 @@ def _get_single_resources_schema():
     infra_pattern = (f'^(?:{cloud_pattern}{region_zone_pattern}|'
                      f'{wildcard_cloud}{wildcard_with_region}|'
                      f'{kubernetes_pattern})$')
+    return infra_pattern
 
+
+def _get_single_resources_schema():
+    """Schema for a single resource in a resources list."""
     return {
         '$schema': 'https://json-schema.org/draft/2020-12/schema',
         'type': 'object',
@@ -133,7 +165,7 @@ def _get_single_resources_schema():
                 # 3. Kubernetes patterns - e.g. "kubernetes/my-context",
                 #    "k8s/context-name",
                 #    "k8s/aws:eks:us-east-1:123456789012:cluster/my-cluster"
-                'pattern': infra_pattern,
+                'pattern': _get_infra_pattern(),
             },
             'cpus': {
                 'anyOf': [{
@@ -383,6 +415,66 @@ def get_resources_schema():
     }
 
 
+def get_volume_schema():
+    # pylint: disable=import-outside-toplevel
+    from sky.volumes import volume
+
+    return {
+        '$schema': 'https://json-schema.org/draft/2020-12/schema',
+        'type': 'object',
+        'required': ['name', 'type', 'infra'],
+        'additionalProperties': False,
+        'properties': {
+            'name': {
+                'type': 'string',
+            },
+            'type': {
+                'type': 'string',
+                'case_sensitive_enum': [
+                    type.value for type in volume.VolumeType
+                ],
+            },
+            'infra': {
+                'type': 'string',
+                'description': ('Infrastructure specification in format: '
+                                'cloud[/region[/zone]].'),
+                # Pattern validates:
+                # 1. cloud[/region[/zone]] - e.g. "aws", "aws/us-east-1",
+                #    "aws/us-east-1/us-east-1a"
+                # 2. Kubernetes patterns - e.g. "kubernetes/my-context",
+                #    "k8s/context-name",
+                #    "k8s/aws:eks:us-east-1:123456789012:cluster/my-cluster"
+                'pattern': _get_volume_infra_pattern(),
+            },
+            'size': {
+                'type': 'string',
+                'pattern': constants.MEMORY_SIZE_PATTERN,
+            },
+            'resource_name': {
+                'type': 'string',
+            },
+            'config': {
+                'type': 'object',
+                'required': [],
+                'properties': {
+                    'storage_class_name': {
+                        'type': 'string',
+                    },
+                    'access_mode': {
+                        'type': 'string',
+                        'case_sensitive_enum': [
+                            type.value for type in volume.VolumeAccessMode
+                        ],
+                    },
+                    'namespace': {
+                        'type': 'string',
+                    },
+                },
+            },
+        }
+    }
+
+
 def get_storage_schema():
     # pylint: disable=import-outside-toplevel
     from sky.data import storage
@@ -457,6 +549,49 @@ def get_storage_schema():
     }
 
 
+def get_volume_mount_schema():
+    """Schema for volume mount object in task config (internal use only)."""
+    return {
+        '$schema': 'https://json-schema.org/draft/2020-12/schema',
+        'type': 'object',
+        'required': [],
+        'additionalProperties': False,
+        'properties': {
+            'path': {
+                'type': 'string',
+            },
+            'volume_name': {
+                'type': 'string',
+            },
+            'volume_config': {
+                'type': 'object',
+                'required': [],
+                'additionalProperties': True,
+                'properties': {
+                    'cloud': {
+                        'type': 'string',
+                        'case_insensitive_enum': list(constants.ALL_CLOUDS)
+                    },
+                    'region': {
+                        'anyOf': [{
+                            'type': 'string'
+                        }, {
+                            'type': 'null'
+                        }]
+                    },
+                    'zone': {
+                        'anyOf': [{
+                            'type': 'string'
+                        }, {
+                            'type': 'null'
+                        }]
+                    },
+                },
+            }
+        }
+    }
+
+
 def get_service_schema():
     """Schema for top-level `service:` field (for SkyServe)."""
     # To avoid circular imports, only import when needed.
@@ -723,6 +858,14 @@ def get_task_schema():
             'config': _filter_schema(
                 get_config_schema(),
                 constants.OVERRIDEABLE_CONFIG_KEYS_IN_TASK),
+            # volumes config is validated separately using get_volume_schema
+            'volumes': {
+                'type': 'object',
+            },
+            'volume_mounts': {
+                'type': 'array',
+                'items': get_volume_mount_schema(),
+            },
             **_experimental_task_schema(),
         }
     }

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,
+    )