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

sky/utils/command_runner.py

@@ -4,6 +4,7 @@ import hashlib
 import os
 import pathlib
 import shlex
+import sys
 import time
 from typing import Any, Callable, Iterable, List, Optional, Tuple, Type, Union
 
@@ -231,9 +232,9 @@ class CommandRunner:
             self,
             source: str,
             target: str,
-            node_destination: str,
+            node_destination: Optional[str],
             up: bool,
-            rsh_option: str,
+            rsh_option: Optional[str],
             # Advanced options.
             log_path: str = os.devnull,
             stream_logs: bool = True,
@@ -283,28 +284,43 @@ class CommandRunner:
                         RSYNC_EXCLUDE_OPTION.format(
                             shlex.quote(str(resolved_source / GIT_EXCLUDE))))
 
-        rsync_command.append(f'-e {shlex.quote(rsh_option)}')
+        if rsh_option is not None:
+            rsync_command.append(f'-e {shlex.quote(rsh_option)}')
+        maybe_dest_prefix = ('' if node_destination is None else
+                             f'{node_destination}:')
 
         if up:
             resolved_target = target
-            if target.startswith('~'):
-                remote_home_dir = _get_remote_home_dir_with_retry()
-                resolved_target = target.replace('~', remote_home_dir)
+            if node_destination is None:
+                # Is a local rsync. Directly resolve the target.
+                resolved_target = str(
+                    pathlib.Path(target).expanduser().resolve())
+            else:
+                if target.startswith('~'):
+                    remote_home_dir = _get_remote_home_dir_with_retry()
+                    resolved_target = target.replace('~', remote_home_dir)
             full_source_str = str(resolved_source)
             if resolved_source.is_dir():
                 full_source_str = os.path.join(full_source_str, '')
             rsync_command.extend([
                 f'{full_source_str!r}',
-                f'{node_destination}:{resolved_target!r}',
+                f'{maybe_dest_prefix}{resolved_target!r}',
             ])
         else:
             resolved_source = source
-            if source.startswith('~'):
-                remote_home_dir = _get_remote_home_dir_with_retry()
-                resolved_source = source.replace('~', remote_home_dir)
+            if node_destination is None:
+                resolved_target = str(
+                    pathlib.Path(target).expanduser().resolve())
+                resolved_source = str(
+                    pathlib.Path(source).expanduser().resolve())
+            else:
+                resolved_target = os.path.expanduser(target)
+                if source.startswith('~'):
+                    remote_home_dir = _get_remote_home_dir_with_retry()
+                    resolved_source = source.replace('~', remote_home_dir)
             rsync_command.extend([
-                f'{node_destination}:{resolved_source!r}',
-                f'{os.path.expanduser(target)!r}',
+                f'{maybe_dest_prefix}{resolved_source!r}',
+                f'{resolved_target!r}',
             ])
         command = ' '.join(rsync_command)
         logger.debug(f'Running rsync command: {command}')
@@ -964,3 +980,93 @@ class KubernetesCommandRunner(CommandRunner):
             # /~/xx, so we need to replace ~ with the remote home directory. We
             # only need to do this when ~ is at the beginning of the path.
             get_remote_home_dir=get_remote_home_dir)
+
+
+class LocalProcessCommandRunner(CommandRunner):
+    """Runner for local process commands."""
+
+    def __init__(self):
+        super().__init__('local')
+
+    @timeline.event
+    @context_utils.cancellation_guard
+    def run(
+            self,
+            cmd: Union[str, List[str]],
+            *,
+            require_outputs: bool = False,
+            port_forward: Optional[List[Tuple[int, int]]] = None,
+            # Advanced options.
+            log_path: str = os.devnull,
+            # If False, do not redirect stdout/stderr to optimize performance.
+            process_stream: bool = True,
+            stream_logs: bool = True,
+            ssh_mode: SshMode = SshMode.NON_INTERACTIVE,
+            separate_stderr: bool = False,
+            connect_timeout: Optional[int] = None,
+            source_bashrc: bool = False,
+            skip_num_lines: int = 0,
+            **kwargs) -> Union[int, Tuple[int, str, str]]:
+        """Use subprocess to run the command."""
+        del port_forward, ssh_mode, connect_timeout  # Unused.
+
+        command_str = self._get_command_to_run(cmd,
+                                               process_stream,
+                                               separate_stderr,
+                                               skip_num_lines=skip_num_lines,
+                                               source_bashrc=source_bashrc)
+
+        log_dir = os.path.expanduser(os.path.dirname(log_path))
+        os.makedirs(log_dir, exist_ok=True)
+
+        executable = None
+        command = [command_str]
+        if not process_stream:
+            if stream_logs:
+                command += [
+                    f'| tee {log_path}',
+                    # This also requires the executor to be '/bin/bash' instead
+                    # of the default '/bin/sh'.
+                    '; exit ${PIPESTATUS[0]}'
+                ]
+            else:
+                command += [f'> {log_path}']
+            executable = '/bin/bash'
+        command_str = ' '.join(command)
+        # For local process, the API server might not have this python path
+        # setup. But this command runner should only be triggered from the API
+        # server (in controller consolidation mode), so we can safely replace
+        # the python path with the executable of the API server.
+        command_str = command_str.replace(constants.SKY_PYTHON_CMD,
+                                          sys.executable)
+        logger.debug(f'Running command locally: {command_str}')
+        return log_lib.run_with_log(command_str,
+                                    log_path,
+                                    require_outputs=require_outputs,
+                                    stream_logs=stream_logs,
+                                    process_stream=process_stream,
+                                    shell=True,
+                                    executable=executable,
+                                    **kwargs)
+
+    @timeline.event
+    def rsync(
+        self,
+        source: str,
+        target: str,
+        *,
+        up: bool,
+        # Advanced options.
+        log_path: str = os.devnull,
+        stream_logs: bool = True,
+        max_retry: int = 1,
+    ) -> None:
+        """Use rsync to sync the source to the target."""
+        self._rsync(source,
+                    target,
+                    node_destination=None,
+                    up=up,
+                    rsh_option=None,
+                    log_path=log_path,
+                    stream_logs=stream_logs,
+                    max_retry=max_retry)

sky/utils/command_runner.pyi

@@ -271,3 +271,60 @@ class KubernetesCommandRunner(CommandRunner):
               stream_logs: bool = ...,
               max_retry: int = ...) -> None:
         ...
+
+
+class LocalProcessCommandRunner(CommandRunner):
+
+    def __init__(self) -> None:
+        ...
+
+    @typing.overload
+    def run(self,
+            cmd: Union[str, List[str]],
+            *,
+            port_forward: Optional[List[int]] = ...,
+            require_outputs: Literal[False] = ...,
+            log_path: str = ...,
+            process_stream: bool = ...,
+            stream_logs: bool = ...,
+            ssh_mode: SshMode = ...,
+            separate_stderr: bool = ...,
+            connect_timeout: Optional[int] = ...,
+            source_bashrc: bool = ...,
+            skip_lines: int = ...,
+            **kwargs) -> int:
+        ...
+
+    @typing.overload
+    def run(self,
+            cmd: Union[str, List[str]],
+            *,
+            port_forward: Optional[List[int]] = ...,
+            require_outputs: Literal[True],
+            log_path: str = ...,
+            process_stream: bool = ...,
+            stream_logs: bool = ...,
+            ssh_mode: SshMode = ...,
+            separate_stderr: bool = ...,
+            connect_timeout: Optional[int] = ...,
+            source_bashrc: bool = ...,
+            skip_lines: int = ...,
+            **kwargs) -> Tuple[int, str, str]:
+        ...
+
+    @typing.overload
+    def run(self,
+            cmd: Union[str, List[str]],
+            *,
+            port_forward: Optional[List[int]] = ...,
+            require_outputs: bool = ...,
+            log_path: str = ...,
+            process_stream: bool = ...,
+            stream_logs: bool = ...,
+            ssh_mode: SshMode = ...,
+            separate_stderr: bool = ...,
+            connect_timeout: Optional[int] = ...,
+            source_bashrc: bool = ...,
+            skip_lines: int = ...,
+            **kwargs) -> Union[Tuple[int, str, str], int]:
+        ...

sky/utils/common_utils.py

@@ -26,6 +26,7 @@ from sky.adaptors import common as adaptors_common
 from sky.skylet import constants
 from sky.usage import constants as usage_constants
 from sky.utils import annotations
+from sky.utils import common_utils
 from sky.utils import ux_utils
 from sky.utils import validator
 
@@ -298,6 +299,13 @@ def get_current_user() -> 'models.User':
     return models.User.get_current_user()
 
 
+def get_current_user_name() -> str:
+    """Returns the current user name."""
+    name = common_utils.get_current_user().name
+    assert name is not None
+    return name
+
+
 def set_current_user(user: 'models.User'):
     """Sets the current user."""
     global _current_user
@@ -754,7 +762,7 @@ def get_cleaned_username(username: str = '') -> str:
     Returns:
       A cleaned username.
     """
-    username = username or getpass.getuser()
+    username = username or common_utils.get_current_user_name()
     username = username.lower()
     username = re.sub(r'[^a-z0-9-_]', '', username)
     username = re.sub(r'^[0-9-]+', '', username)

sky/utils/context.py

@@ -254,7 +254,9 @@ class Popen(subprocess.Popen):
     def __init__(self, *args, **kwargs):
         env = kwargs.pop('env', None)
         if env is None:
-            env = os.environ
+            # Pass a copy of current context.environ to avoid race condition
+            # when the context is updated after the Popen is created.
+            env = os.environ.copy()
         super().__init__(*args, env=env, **kwargs)
 
 

sky/utils/controller_utils.py

@@ -2,7 +2,6 @@
 import copy
 import dataclasses
 import enum
-import getpass
 import os
 import tempfile
 import typing
@@ -498,7 +497,7 @@ def shared_controller_vars_to_fill(
     env_vars.update({
         # Should not use $USER here, as that env var can be empty when
         # running in a container.
-        constants.USER_ENV_VAR: getpass.getuser(),
+        constants.USER_ENV_VAR: common_utils.get_current_user_name(),
         constants.USER_ID_ENV_VAR: common_utils.get_user_hash(),
         # Skip cloud identity check to avoid the overhead.
         env_options.Options.SKIP_CLOUD_IDENTITY_CHECK.env_key: '1',

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.
@@ -672,18 +807,6 @@ def get_task_schema():
             'service': {
                 'type': 'object',
             },
-            'job': {
-                'type': 'object',
-                'required': [],
-                'additionalProperties': False,
-                'properties': {
-                    'priority': {
-                        'type': 'integer',
-                        'minimum': 0,
-                        'maximum': 1000,
-                    },
-                },
-            },
             'setup': {
                 'type': 'string',
             },
@@ -735,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(),
         }
     }
@@ -899,30 +1030,41 @@ def get_config_schema():
         if k != '$schema'
     }
     resources_schema['properties'].pop('ports')
-    controller_resources_schema = {
-        'type': 'object',
-        'required': [],
-        'additionalProperties': False,
-        'properties': {
-            'controller': {
-                'type': 'object',
-                'required': [],
-                'additionalProperties': False,
-                'properties': {
-                    'resources': resources_schema,
-                    'high_availability': {
-                        'type': 'boolean',
-                    },
-                    'autostop': _AUTOSTOP_SCHEMA,
-                }
+
+    def _get_controller_schema(add_consolidation_mode: bool = False):
+        controller_properties = {
+            'resources': resources_schema,
+            'high_availability': {
+                'type': 'boolean',
+                'default': False,
             },
-            'bucket': {
-                'type': 'string',
-                'pattern': '^(https|s3|gs|r2|cos)://.+',
-                'required': [],
+            'autostop': _AUTOSTOP_SCHEMA,
+        }
+        if add_consolidation_mode:
+            controller_properties['consolidation_mode'] = {
+                'type': 'boolean',
+                'default': False,
+            }
+
+        return {
+            'type': 'object',
+            'required': [],
+            'additionalProperties': False,
+            'properties': {
+                'controller': {
+                    'type': 'object',
+                    'required': [],
+                    'additionalProperties': False,
+                    'properties': controller_properties,
+                },
+                'bucket': {
+                    'type': 'string',
+                    'pattern': '^(https|s3|gs|r2|cos)://.+',
+                    'required': [],
+                }
             }
         }
-    }
+
     cloud_configs = {
         'aws': {
             'type': 'object',
@@ -1440,8 +1582,8 @@ def get_config_schema():
             'db': {
                 'type': 'string',
             },
-            'jobs': controller_resources_schema,
-            'serve': controller_resources_schema,
+            'jobs': _get_controller_schema(add_consolidation_mode=True),
+            'serve': _get_controller_schema(add_consolidation_mode=False),
             'allowed_clouds': allowed_clouds,
             'admin_policy': admin_policy_schema,
             'docker': docker_configs,