skypilot-nightly 1.0.0.dev20250502__py3-none-any.whl → 1.0.0.dev20251203__py3-none-any.whl

sky/utils/serialize_utils.py

@@ -0,0 +1,16 @@
+"""Utilities for handling resource handles."""
+import copy
+import typing
+
+
+def prepare_handle_for_backwards_compatibility(
+        handle: typing.Any) -> typing.Any:
+    """Prepare a handle for backwards compatibility with older clients."""
+    # skylet_ssh_tunnel was causing backwards compatibility issues with older
+    # clients: AttributeError: Can't get attribute 'SSHTunnelInfo'
+    #
+    # But it is not needed on the client side, so we can just remove it.
+    if handle is not None and hasattr(handle, 'skylet_ssh_tunnel'):
+        handle = copy.deepcopy(handle)
+        handle.skylet_ssh_tunnel = None
+    return handle

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/subprocess_utils.py

@@ -6,18 +6,20 @@ import random
 import resource
 import shlex
 import subprocess
+import sys
 import threading
 import time
 import typing
-from typing import Any, Callable, Dict, List, Optional, Protocol, Tuple, Union
+from typing import (Any, Callable, Dict, List, Optional, Protocol, Set, Tuple,
+                    Union)
 
 import colorama
 
 from sky import exceptions
 from sky import sky_logging
 from sky.adaptors import common as adaptors_common
-from sky.skylet import constants
 from sky.skylet import log_lib
+from sky.skylet import subprocess_daemon
 from sky.utils import common_utils
 from sky.utils import timeline
 from sky.utils import ux_utils
@@ -107,7 +109,7 @@ def get_parallel_threads(cloud_str: Optional[str] = None) -> int:
 
 
 def run_in_parallel(func: Callable,
-                    args: List[Any],
+                    args: Union[List[Any], Set[Any]],
                     num_threads: Optional[int] = None) -> List[Any]:
     """Run a function in parallel on a list of arguments.
 
@@ -128,7 +130,7 @@ def run_in_parallel(func: Callable,
     if len(args) == 0:
         return []
     if len(args) == 1:
-        return [func(args[0])]
+        return [func(list(args)[0])]
 
     processes = (num_threads
                  if num_threads is not None else get_parallel_threads())
@@ -208,8 +210,11 @@ def kill_children_processes(parent_pids: Optional[Union[
             kill_process_with_grace_period(child, force=force)
 
 
-def kill_process_with_grace_period(proc: Union[multiprocessing.Process,
-                                               psutil.Process],
+GenericProcess = Union[multiprocessing.Process, psutil.Process,
+                       subprocess.Popen]
+
+
+def kill_process_with_grace_period(proc: GenericProcess,
                                    force: bool = False,
                                    grace_period: int = 10) -> None:
     """Kill a process with SIGTERM and wait for it to exit.
@@ -223,6 +228,9 @@ def kill_process_with_grace_period(proc: Union[multiprocessing.Process,
     if isinstance(proc, psutil.Process):
         alive = proc.is_running
         wait = proc.wait
+    elif isinstance(proc, subprocess.Popen):
+        alive = lambda: proc.poll() is None
+        wait = proc.wait
     else:
         alive = proc.is_alive
         wait = proc.join
@@ -240,11 +248,10 @@ def kill_process_with_grace_period(proc: Union[multiprocessing.Process,
         # The child process may have already been terminated.
         return
     except psutil.TimeoutExpired:
-        # Pass to finally to force kill the process.
-        pass
-    finally:
         logger.debug(f'Process {proc.pid} did not terminate after '
                      f'{grace_period} seconds')
+        # Continue to finally to force kill the process.
+    finally:
         # Attempt to force kill if the normal termination fails
         if not force:
             logger.debug(f'Force killing process {proc.pid}')
@@ -300,11 +307,17 @@ def run_with_retries(
     return returncode, stdout, stderr
 
 
-def kill_process_daemon(process_pid: int) -> None:
+def kill_process_daemon(process_pid: int, use_kill_pg: bool = False) -> None:
     """Start a daemon as a safety net to kill the process.
 
     Args:
         process_pid: The PID of the process to kill.
+        use_kill_pg: Whether to use kill process group to kill the process. If
+            True, the process will use os.killpg() to kill the target process
+            group on UNIX system, which is more efficient than using the daemon
+            to refresh the process tree in the daemon. Note that both
+            implementations have corner cases where subprocesses might not be
+            killed. Refer to subprocess_daemon.py for more details.
     """
     # Get initial children list
     try:
@@ -317,12 +330,8 @@ def kill_process_daemon(process_pid: int) -> None:
     daemon_script = os.path.join(
         os.path.dirname(os.path.abspath(log_lib.__file__)),
         'subprocess_daemon.py')
-    python_path = subprocess.check_output(constants.SKY_GET_PYTHON_PATH_CMD,
-                                          shell=True,
-                                          stderr=subprocess.DEVNULL,
-                                          encoding='utf-8').strip()
     daemon_cmd = [
-        python_path,
+        sys.executable,
         daemon_script,
         '--parent-pid',
         str(parent_pid),
@@ -335,6 +344,10 @@ def kill_process_daemon(process_pid: int) -> None:
         ','.join(map(str, initial_children)),
     ]
 
+    env = os.environ.copy()
+    if use_kill_pg:
+        env[subprocess_daemon.USE_KILL_PG_ENV_VAR] = '1'
+
     # We do not need to set `start_new_session=True` here, as the
     # daemon script will detach itself from the parent process with
     # fork to avoid being killed by parent process. See the reason we
@@ -346,6 +359,7 @@ def kill_process_daemon(process_pid: int) -> None:
         stderr=subprocess.DEVNULL,
         # Disable input
         stdin=subprocess.DEVNULL,
+        env=env,
     )
 
 

sky/utils/tempstore.py

@@ -0,0 +1,70 @@
+"""Temporary storage context manager."""
+
+import contextlib
+import contextvars
+import functools
+import os
+import tempfile
+import typing
+from typing import Any, Callable, Iterator, Optional, TypeVar
+
+_TEMP_DIR: contextvars.ContextVar[Optional[str]] = contextvars.ContextVar(
+    'temp_store_dir', default=None)
+
+
+@contextlib.contextmanager
+def tempdir() -> Iterator[str]:
+    """Context manager for temporary directory of current context.
+
+    This wraps tempfile.TemporaryDirectory and makes the temp dir available
+    throughout the context, eliminating the need to pass the temp dir to
+    the nested functions that need it.
+
+    This context manager is nestable - nested calls will create new temp dirs
+    and restore the previous temp dir when exiting.
+    """
+    with tempfile.TemporaryDirectory(prefix='sky-tmp') as temp_dir:
+        token = _TEMP_DIR.set(temp_dir)
+        try:
+            yield temp_dir
+        finally:
+            _TEMP_DIR.reset(token)
+
+
+# Keep the function signature same as tempfile.mkdtemp.
+# pylint: disable=redefined-builtin
+def mkdtemp(suffix: Optional[str] = None,
+            prefix: Optional[str] = None,
+            dir: Optional[str] = None) -> str:
+    """Create a temporary directory in the temp dir of current context.
+
+    The directory will be cleaned when the current context exits.
+    If there is no temp dir in current context, this function is equivalent to
+    tempfile.mkdtemp.
+    """
+    context_temp_dir = _TEMP_DIR.get()
+
+    if context_temp_dir is not None and dir is None:
+        dir = context_temp_dir
+    elif context_temp_dir is not None and dir is not None:
+        dir = os.path.join(context_temp_dir, dir)
+        os.makedirs(dir, exist_ok=True)
+
+    return tempfile.mkdtemp(suffix=suffix, prefix=prefix, dir=dir)
+
+
+F = TypeVar('F', bound=Callable[..., Any])
+
+
+def with_tempdir(func: F) -> F:
+    """Decorator that wraps a function call with tempdir() context manager.
+
+    Refer to `tempdir` for more details.
+    """
+
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        with tempdir():
+            return func(*args, **kwargs)
+
+    return typing.cast(F, wrapper)

sky/utils/thread_utils.py

@@ -0,0 +1,91 @@
+"""Utility functions for threads."""
+
+import threading
+from typing import Any, Dict, Generic, Optional, overload, TypeVar
+
+from sky.utils import common_utils
+
+
+class SafeThread(threading.Thread):
+    """A thread that can catch exceptions."""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._exc = None
+
+    def run(self):
+        try:
+            super().run()
+        except BaseException as e:  # pylint: disable=broad-except
+            self._exc = e
+
+    @property
+    def format_exc(self) -> Optional[str]:
+        if self._exc is None:
+            return None
+        return common_utils.format_exception(self._exc)
+
+
+# pylint: disable=invalid-name
+KeyType = TypeVar('KeyType')
+ValueType = TypeVar('ValueType')
+
+
+# Google style guide: Do not rely on the atomicity of built-in types.
+# Our launch and down process pool will be used by multiple threads,
+# therefore we need to use a thread-safe dict.
+# see https://google.github.io/styleguide/pyguide.html#218-threading
+class ThreadSafeDict(Generic[KeyType, ValueType]):
+    """A thread-safe dict."""
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        self._dict: Dict[KeyType, ValueType] = dict(*args, **kwargs)
+        self._lock = threading.Lock()
+
+    def __getitem__(self, key: KeyType) -> ValueType:
+        with self._lock:
+            return self._dict.__getitem__(key)
+
+    def __setitem__(self, key: KeyType, value: ValueType) -> None:
+        with self._lock:
+            return self._dict.__setitem__(key, value)
+
+    def __delitem__(self, key: KeyType) -> None:
+        with self._lock:
+            return self._dict.__delitem__(key)
+
+    def __len__(self) -> int:
+        with self._lock:
+            return self._dict.__len__()
+
+    def __contains__(self, key: KeyType) -> bool:
+        with self._lock:
+            return self._dict.__contains__(key)
+
+    def items(self):
+        with self._lock:
+            return self._dict.items()
+
+    def values(self):
+        with self._lock:
+            return self._dict.values()
+
+    @overload
+    def get(self, key: KeyType, default: ValueType) -> ValueType:
+        ...
+
+    @overload
+    def get(self,
+            key: KeyType,
+            default: Optional[ValueType] = None) -> Optional[ValueType]:
+        ...
+
+    def get(self,
+            key: KeyType,
+            default: Optional[ValueType] = None) -> Optional[ValueType]:
+        with self._lock:
+            return self._dict.get(key, default)
+
+    def pop(self, key: KeyType) -> Optional[ValueType]:
+        with self._lock:
+            return self._dict.pop(key, None)

sky/utils/timeline.py

@@ -4,7 +4,6 @@ The timeline follows the trace event format defined here:
 https://docs.google.com/document/d/1CvAClvFfyA5R-PhYUmn5OOQtYMH4h6I0nSsKchNAySU/preview
 """  # pylint: disable=line-too-long
 import atexit
-import functools
 import json
 import os
 import threading
@@ -12,13 +11,15 @@ import time
 import traceback
 from typing import Callable, Optional, Union
 
-import filelock
-
 from sky.utils import common_utils
 
 _events = []
 
 
+def _get_events_file_path():
+    return os.environ.get('SKYPILOT_TIMELINE_FILE_PATH')
+
+
 class Event:
     """Record an event.
 
@@ -28,6 +29,10 @@ class Event:
     """
 
     def __init__(self, name: str, message: Optional[str] = None):
+        self._skipped = False
+        if not _get_events_file_path():
+            self._skipped = True
+            return
         self._name = name
         self._message = message
         # See the module doc for the event format.
@@ -44,6 +49,8 @@ class Event:
             self._event['args'] = {'message': self._message}
 
     def begin(self):
+        if self._skipped:
+            return
         event_begin = self._event.copy()
         event_begin.update({
             'ph': 'B',
@@ -51,10 +58,13 @@ class Event:
         })
         event_begin['args'] = {'stack': '\n'.join(traceback.format_stack())}
         if self._message is not None:
-            event_begin['args']['message'] = self._message
+            event_begin['args'][
+                'message'] = self._message  # type: ignore[index]
         _events.append(event_begin)
 
     def end(self):
+        if self._skipped:
+            return
         event_end = self._event.copy()
         event_end.update({
             'ph': 'E',
@@ -76,63 +86,26 @@ def event(name_or_fn: Union[str, Callable], message: Optional[str] = None):
     return common_utils.make_decorator(Event, name_or_fn, message=message)
 
 
-class FileLockEvent:
-    """Serve both as a file lock and event for the lock."""
-
-    def __init__(self, lockfile: Union[str, os.PathLike], timeout: float = -1):
-        self._lockfile = lockfile
-        os.makedirs(os.path.dirname(os.path.abspath(self._lockfile)),
-                    exist_ok=True)
-        self._lock = filelock.FileLock(self._lockfile, timeout)
-        self._hold_lock_event = Event(f'[FileLock.hold]:{self._lockfile}')
-
-    def acquire(self):
-        was_locked = self._lock.is_locked
-        with Event(f'[FileLock.acquire]:{self._lockfile}'):
-            self._lock.acquire()
-        if not was_locked and self._lock.is_locked:
-            # start holding the lock after initial acquiring
-            self._hold_lock_event.begin()
-
-    def release(self):
-        was_locked = self._lock.is_locked
-        self._lock.release()
-        if was_locked and not self._lock.is_locked:
-            # stop holding the lock after initial releasing
-            self._hold_lock_event.end()
-
-    def __enter__(self):
-        self.acquire()
-        return self
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        self.release()
-
-    def __call__(self, f):
-        # Make this class callable as a decorator.
-        @functools.wraps(f)
-        def wrapper(*args, **kwargs):
-            with self:
-                return f(*args, **kwargs)
-
-        return wrapper
-
-
 def save_timeline():
-    file_path = os.environ.get('SKYPILOT_TIMELINE_FILE_PATH')
-    if not file_path:
+    events_file_path = _get_events_file_path()
+    if not events_file_path:
         return
+    global _events
+    events_to_write = _events
+    _events = []
     json_output = {
-        'traceEvents': _events,
+        'traceEvents': events_to_write,
         'displayTimeUnit': 'ms',
         'otherData': {
-            'log_dir': os.path.dirname(os.path.abspath(file_path)),
+            'log_dir': os.path.dirname(os.path.abspath(events_file_path)),
         }
     }
-    os.makedirs(os.path.dirname(os.path.abspath(file_path)), exist_ok=True)
-    with open(file_path, 'w', encoding='utf-8') as f:
+    os.makedirs(os.path.dirname(os.path.abspath(events_file_path)),
+                exist_ok=True)
+    with open(events_file_path, 'w', encoding='utf-8') as f:
         json.dump(json_output, f)
+    del events_to_write
 
 
-if os.environ.get('SKYPILOT_TIMELINE_FILE_PATH'):
+if _get_events_file_path():
     atexit.register(save_timeline)

sky/utils/ux_utils.py

@@ -1,17 +1,19 @@
 """Utility functions for UX."""
 import contextlib
 import enum
+import fnmatch
 import os
 import sys
 import traceback
 import typing
-from typing import Callable, Optional, Union
+from typing import Callable, Iterable, List, Optional, Union
 
 import colorama
 
 from sky import sky_logging
 from sky.skylet import constants
 from sky.utils import common_utils
+from sky.utils import env_options
 from sky.utils import rich_console_utils
 
 if typing.TYPE_CHECKING:
@@ -25,9 +27,16 @@ BOLD = '\033[1m'
 RESET_BOLD = '\033[0m'
 
 # Log path hint in the spinner during launching
+# (old, kept for backward compatibility)
 _LOG_PATH_HINT = (f'{colorama.Style.DIM}View logs: sky api logs -l '
                   '{log_path}'
                   f'{colorama.Style.RESET_ALL}')
+# Log hint: recommend sky logs --provision <cluster_name>
+_PROVISION_LOG_HINT = (
+    f'{colorama.Style.DIM}View logs: '
+    f'{BOLD}sky logs --provision {{cluster_name}}{RESET_BOLD}'
+    f'{colorama.Style.RESET_ALL}')
+# Legacy path hint retained for local-only cases where we don't have cluster
 _LOG_PATH_HINT_LOCAL = (f'{colorama.Style.DIM}View logs: '
                         '{log_path}'
                         f'{colorama.Style.RESET_ALL}')
@@ -57,10 +66,14 @@ def print_exception_no_traceback():
             if error():
                 raise ValueError('...')
     """
-    original_tracelimit = getattr(sys, 'tracebacklimit', 1000)
-    sys.tracebacklimit = 0
-    yield
-    sys.tracebacklimit = original_tracelimit
+    if env_options.Options.SHOW_DEBUG_INFO.get():
+        # When SKYPILOT_DEBUG is set, show the full traceback
+        yield
+    else:
+        original_tracelimit = getattr(sys, 'tracebacklimit', 1000)
+        sys.tracebacklimit = 0
+        yield
+        sys.tracebacklimit = original_tracelimit
 
 
 @contextlib.contextmanager
@@ -121,7 +134,10 @@ class RedirectOutputForProcess:
 
 def log_path_hint(log_path: Union[str, 'pathlib.Path'],
                   is_local: bool = False) -> str:
-    """Gets the log path hint for the given log path."""
+    """Gets the log path hint for the given log path.
+
+    Kept for backward compatibility when only paths are available.
+    """
     log_path = str(log_path)
     expanded_home = os.path.expanduser('~')
     if log_path.startswith(expanded_home):
@@ -134,6 +150,12 @@ def log_path_hint(log_path: Union[str, 'pathlib.Path'],
     return _LOG_PATH_HINT.format(log_path=log_path)
 
 
+def provision_hint(cluster_name: Optional[str]) -> Optional[str]:
+    if not cluster_name:
+        return None
+    return _PROVISION_LOG_HINT.format(cluster_name=cluster_name)
+
+
 def starting_message(message: str) -> str:
     """Gets the starting message for the given message."""
     # We have to reset the color before the message, because sometimes if a
@@ -145,7 +167,8 @@ def starting_message(message: str) -> str:
 def finishing_message(message: str,
                       log_path: Optional[Union[str, 'pathlib.Path']] = None,
                       is_local: bool = False,
-                      follow_up_message: Optional[str] = None) -> str:
+                      follow_up_message: Optional[str] = None,
+                      cluster_name: Optional[str] = None) -> str:
     """Gets the finishing message for the given message.
 
     Args:
@@ -161,7 +184,11 @@ def finishing_message(message: str,
     follow_up_message = follow_up_message if (follow_up_message
                                               is not None) else ''
     success_prefix = (f'{colorama.Style.RESET_ALL}{colorama.Fore.GREEN}✓ '
-                      f'{message}{colorama.Style.RESET_ALL}{follow_up_message}')
+                      f'{message}{colorama.Style.RESET_ALL}{follow_up_message}'
+                      f'{colorama.Style.RESET_ALL}')
+    hint = provision_hint(cluster_name)
+    if hint:
+        return f'{success_prefix}  {hint}'
     if log_path is None:
         return success_prefix
     path_hint = log_path_hint(log_path, is_local)
@@ -170,13 +197,17 @@ def finishing_message(message: str,
 
 def error_message(message: str,
                   log_path: Optional[Union[str, 'pathlib.Path']] = None,
-                  is_local: bool = False) -> str:
+                  is_local: bool = False,
+                  cluster_name: Optional[str] = None) -> str:
     """Gets the error message for the given message."""
     # We have to reset the color before the message, because sometimes if a
     # previous spinner with dimmed color overflows in a narrow terminal, the
     # color might be messed up.
     error_prefix = (f'{colorama.Style.RESET_ALL}{colorama.Fore.RED}⨯'
                     f'{colorama.Style.RESET_ALL} {message}')
+    hint = provision_hint(cluster_name)
+    if hint:
+        return f'{error_prefix}  {hint}'
     if log_path is None:
         return error_prefix
     path_hint = log_path_hint(log_path, is_local)
@@ -194,9 +225,16 @@ def retry_message(message: str) -> str:
 
 def spinner_message(message: str,
                     log_path: Optional[Union[str, 'pathlib.Path']] = None,
-                    is_local: bool = False) -> str:
-    """Gets the spinner message for the given message and log path."""
+                    is_local: bool = False,
+                    cluster_name: Optional[str] = None) -> str:
+    """Gets the spinner message for the given message and log path.
+
+    If cluster_name is provided, recommend `sky logs --provision <cluster>`.
+    """
     colored_spinner = f'[bold cyan]{message}[/]'
+    hint = provision_hint(cluster_name)
+    if hint:
+        return f'{colored_spinner}  {hint}'
     if log_path is None:
         return colored_spinner
     path_hint = log_path_hint(log_path, is_local)
@@ -247,9 +285,40 @@ def command_hint_messages(hint_type: CommandHintType,
                 f'{BOLD}sky jobs logs {job_id}{RESET_BOLD}'
                 f'\n{INDENT_SYMBOL}To stream controller logs:\t\t'
                 f'{BOLD}sky jobs logs --controller {job_id}{RESET_BOLD}'
-                f'\n{INDENT_SYMBOL}To view all managed jobs:\t\t'
-                f'{BOLD}sky jobs queue{RESET_BOLD}'
-                f'\n{INDENT_LAST_SYMBOL}To view managed job dashboard:\t\t'
-                f'{BOLD}sky jobs dashboard{RESET_BOLD}')
+                f'\n{INDENT_LAST_SYMBOL}To view all managed jobs:\t\t'
+                f'{BOLD}sky jobs queue{RESET_BOLD}')
     else:
         raise ValueError(f'Invalid hint type: {hint_type}')
+
+
+def is_glob_pattern(pattern: str) -> bool:
+    """Checks if a string contains common glob pattern wildcards."""
+    glob_chars = {'*', '?', '[', ']'}
+    # Also check for '**' as a specific globstar pattern
+    if '**' in pattern:
+        return True
+    for char in pattern:
+        if char in glob_chars:
+            return True
+    return False
+
+
+def get_non_matched_query(query_clusters: Iterable[str],
+                          cluster_names: Iterable[str]) -> List[str]:
+    """Gets the non-matched query clusters."""
+    glob_query_clusters = []
+    non_glob_query_clusters = []
+    for cluster_name in query_clusters:
+        if is_glob_pattern(cluster_name):
+            glob_query_clusters.append(cluster_name)
+        else:
+            non_glob_query_clusters.append(cluster_name)
+    not_found_clusters = [
+        query_cluster for query_cluster in non_glob_query_clusters
+        if query_cluster not in cluster_names
+    ]
+    not_found_clusters.extend([
+        query_cluster for query_cluster in glob_query_clusters
+        if not fnmatch.filter(cluster_names, query_cluster)
+    ])
+    return not_found_clusters

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