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

sky/utils/kubernetes/ssh_utils.py

@@ -0,0 +1,221 @@
+"""Utility functions for managing SSH node pools."""
+import os
+import re
+import subprocess
+from typing import Any, Callable, Dict, List, Optional
+import uuid
+
+import yaml
+
+from sky.utils import ux_utils
+
+DEFAULT_SSH_NODE_POOLS_PATH = os.path.expanduser('~/.sky/ssh_node_pools.yaml')
+RED = '\033[0;31m'
+NC = '\033[0m'  # No color
+
+
+def check_host_in_ssh_config(hostname: str) -> bool:
+    """Return True iff *hostname* matches at least one `Host`/`Match` stanza
+    in the user's OpenSSH client configuration (including anything pulled in
+    via Include).
+
+    It calls:  ssh -vvG <hostname> -o ConnectTimeout=0
+    which:
+      • -G  expands the effective config without connecting
+      • -vv prints debug lines that show which stanzas are applied
+      • ConnectTimeout=0 avoids a DNS lookup if <hostname> is a FQDN/IP
+
+    No config files are opened or parsed manually.
+
+    Parameters
+    ----------
+    hostname : str
+        The alias/IP/FQDN you want to test.
+
+    Returns
+    -------
+    bool
+        True  – a specific stanza matched the host
+        False – nothing but the global defaults (`Host *`) applied
+    """
+    # We direct stderr→stdout because debug output goes to stderr.
+    proc = subprocess.run(
+        ['ssh', '-vvG', hostname, '-o', 'ConnectTimeout=0'],
+        text=True,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT,
+        check=False,  # we only want the text, not to raise
+    )
+
+    # Look for lines like:
+    #   debug1: ~/.ssh/config line 42: Applying options for <hostname>
+    # Anything other than "*"
+    pattern = re.compile(r'^debug\d+: .*Applying options for ([^*].*)$',
+                         re.MULTILINE)
+
+    return bool(pattern.search(proc.stdout))
+
+
+class UniqueKeySafeLoader(yaml.SafeLoader):
+    """Custom YAML loader that raises an error if there are duplicate keys."""
+
+    def construct_mapping(self, node, deep=False):
+        mapping = set()
+        for key_node, _ in node.value:
+            key = self.construct_object(key_node, deep=deep)
+            if key in mapping:
+                raise yaml.constructor.ConstructorError(
+                    note=(f'Duplicate key found: {key!r}.\n'
+                          'Please remove one of them from the YAML file.'))
+            mapping.add(key)
+        return super().construct_mapping(node, deep)
+
+
+def load_ssh_targets(file_path: str) -> Dict[str, Any]:
+    """Load SSH targets from YAML file."""
+    if not os.path.exists(file_path):
+        with ux_utils.print_exception_no_traceback():
+            raise ValueError(f'SSH Node Pools file not found: {file_path}')
+
+    try:
+        with open(file_path, 'r', encoding='utf-8') as f:
+            targets = yaml.load(f, Loader=UniqueKeySafeLoader)
+        return targets
+    except yaml.constructor.ConstructorError as e:
+        with ux_utils.print_exception_no_traceback():
+            raise ValueError(e.note) from e
+    except (yaml.YAMLError, IOError, OSError) as e:
+        with ux_utils.print_exception_no_traceback():
+            raise ValueError(f'Error loading SSH Node Pools file: {e}') from e
+
+
+def get_cluster_config(
+        targets: Dict[str, Any],
+        cluster_name: Optional[str] = None,
+        file_path: str = DEFAULT_SSH_NODE_POOLS_PATH) -> Dict[str, Any]:
+    """Get configuration for specific clusters or all clusters."""
+    if not targets:
+        with ux_utils.print_exception_no_traceback():
+            raise ValueError(
+                f'No clusters defined in SSH Node Pools file {file_path}')
+
+    if cluster_name:
+        if cluster_name not in targets:
+            with ux_utils.print_exception_no_traceback():
+                raise ValueError(f'Cluster {cluster_name!r} not found in '
+                                 f'SSH Node Pools file {file_path}')
+        return {cluster_name: targets[cluster_name]}
+
+    # Return all clusters if no specific cluster is specified
+    return targets
+
+
+def prepare_hosts_info(
+    cluster_name: str,
+    cluster_config: Dict[str, Any],
+    upload_ssh_key_func: Optional[Callable[[str, str], str]] = None
+) -> List[Dict[str, str]]:
+    """Prepare list of hosts with resolved user, identity_file, and password.
+
+    Args:
+        cluster_name: The name of the cluster.
+        cluster_config: The configuration for the cluster.
+        upload_ssh_key_func: A function to upload the SSH key to the remote
+            server and wait for the key to be uploaded. This function will take
+            the key name and the local key file path as input, and return the
+            path for the remote SSH key file on the API server. This function
+            will only be set in `sky ssh up -f` mode, and if this function is
+            set, any ssh config will not be allowed as we don't support
+            uploading any ssh config to the API server.
+
+    Returns:
+        A list of hosts with resolved user, identity_file, and password.
+    """
+    if 'hosts' not in cluster_config or not cluster_config['hosts']:
+        with ux_utils.print_exception_no_traceback():
+            raise ValueError(
+                f'No hosts defined in cluster {cluster_name} configuration')
+
+    # Get cluster-level defaults
+    cluster_user = cluster_config.get('user', '')
+    cluster_identity_file = os.path.expanduser(
+        cluster_config.get('identity_file', ''))
+    cluster_password = cluster_config.get('password', '')
+
+    # Check if cluster identity file exists
+    if cluster_identity_file and not os.path.isfile(cluster_identity_file):
+        with ux_utils.print_exception_no_traceback():
+            raise ValueError(
+                f'SSH Identity File Missing: {cluster_identity_file}')
+
+    use_cluster_config_msg = (f'Cluster {cluster_name} uses SSH config '
+                              'for hostname {host}, which is not '
+                              'supported by the -f flag. Please use a '
+                              'dict with `ip` field instead.')
+
+    def _maybe_hardcode_identity_file(i: int, identity_file: str) -> str:
+        if upload_ssh_key_func is None:
+            return identity_file
+        if not os.path.exists(os.path.expanduser(identity_file)):
+            with ux_utils.print_exception_no_traceback():
+                raise ValueError(
+                    f'Identity file {identity_file} does not exist.')
+        key_name = f'{cluster_name}-{i}-{str(uuid.uuid4())[:4]}'
+        key_file_on_api_server = upload_ssh_key_func(key_name, identity_file)
+        return key_file_on_api_server
+
+    hosts_info = []
+    for i, host in enumerate(cluster_config['hosts']):
+        # Host can be a string (IP or SSH config hostname) or a dict
+        if isinstance(host, str):
+            # Check if this is an SSH config hostname
+            is_ssh_config_host = check_host_in_ssh_config(host)
+            if upload_ssh_key_func is not None and is_ssh_config_host:
+                with ux_utils.print_exception_no_traceback():
+                    raise ValueError(use_cluster_config_msg.format(host=host))
+
+            hosts_info.append({
+                'ip': host,
+                'user': '' if is_ssh_config_host else cluster_user,
+                'identity_file': '' if is_ssh_config_host else
+                                 _maybe_hardcode_identity_file(
+                                     i, cluster_identity_file),
+                'password': cluster_password,
+                'use_ssh_config': is_ssh_config_host
+            })
+        else:
+            # It's a dict with potential overrides
+            if 'ip' not in host:
+                print(f'{RED}Warning: Host missing \'ip\' field, '
+                      f'skipping: {host}{NC}')
+                continue
+
+            # Check if this is an SSH config hostname
+            is_ssh_config_host = check_host_in_ssh_config(host['ip'])
+            if upload_ssh_key_func is not None and is_ssh_config_host:
+                with ux_utils.print_exception_no_traceback():
+                    raise ValueError(use_cluster_config_msg.format(host=host))
+
+            # Use host-specific values or fall back to cluster defaults
+            host_user = '' if is_ssh_config_host else host.get(
+                'user', cluster_user)
+            host_identity_file = '' if is_ssh_config_host else (
+                _maybe_hardcode_identity_file(
+                    i, host.get('identity_file', cluster_identity_file)))
+            host_identity_file = os.path.expanduser(host_identity_file)
+            host_password = host.get('password', cluster_password)
+
+            if host_identity_file and not os.path.isfile(host_identity_file):
+                with ux_utils.print_exception_no_traceback():
+                    raise ValueError(
+                        f'SSH Identity File Missing: {host_identity_file}')
+
+            hosts_info.append({
+                'ip': host['ip'],
+                'user': host_user,
+                'identity_file': host_identity_file,
+                'password': host_password,
+                'use_ssh_config': is_ssh_config_host
+            })
+
+    return hosts_info

sky/utils/resources_utils.py

@@ -8,6 +8,7 @@ import typing
 from typing import Dict, List, Optional, Set, Union
 
 from sky import skypilot_config
+from sky.skylet import constants
 from sky.utils import common_utils
 from sky.utils import registry
 from sky.utils import ux_utils
@@ -331,3 +332,68 @@ def make_launchables_for_valid_region_zones(
             # Batch the requests at the granularity of a single region.
             launchables.append(launchable_resources.copy(region=region.name))
     return launchables
+
+
+def parse_memory_resource(resource_qty_str: Union[str, int, float],
+                          field_name: str,
+                          ret_type: type = int,
+                          unit: str = 'gb',
+                          allow_plus: bool = False,
+                          allow_x: bool = False,
+                          allow_rounding: bool = False) -> str:
+    """Returns memory size in chosen units given a resource quantity string.
+
+    Args:
+        resource_qty_str: Resource quantity string
+        unit: Unit to convert to
+        allow_plus: Whether to allow '+' prefix
+        allow_x: Whether to allow 'x' suffix
+    """
+    assert unit in constants.MEMORY_SIZE_UNITS, f'Invalid unit: {unit}'
+
+    error_msg = (f'"{field_name}" field should be a '
+                 f'{constants.MEMORY_SIZE_PATTERN}+?,'
+                 f' got {resource_qty_str}')
+
+    resource_str = str(resource_qty_str)
+
+    # Handle plus and x suffixes, x is only used internally for jobs controller
+    plus = ''
+    if resource_str.endswith('+'):
+        if allow_plus:
+            resource_str = resource_str[:-1]
+            plus = '+'
+        else:
+            raise ValueError(error_msg)
+
+    x = ''
+    if resource_str.endswith('x'):
+        if allow_x:
+            resource_str = resource_str[:-1]
+            x = 'x'
+        else:
+            raise ValueError(error_msg)
+
+    try:
+        # We assume it is already in the wanted units to maintain backwards
+        # compatibility
+        ret_type(resource_str)
+        return f'{resource_str}{plus}{x}'
+    except ValueError:
+        pass
+
+    resource_str = resource_str.lower()
+    for mem_unit, multiplier in constants.MEMORY_SIZE_UNITS.items():
+        if resource_str.endswith(mem_unit):
+            try:
+                value = ret_type(resource_str[:-len(mem_unit)])
+                converted = (value * multiplier /
+                             constants.MEMORY_SIZE_UNITS[unit])
+                if not allow_rounding and ret_type(converted) != converted:
+                    raise ValueError(error_msg)
+                converted = ret_type(converted)
+                return f'{converted}{plus}{x}'
+            except ValueError:
+                continue
+
+    raise ValueError(error_msg)

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)