skypilot-nightly 1.0.0.dev2024053101__py3-none-any.whl → 1.0.0.dev2025022801__py3-none-any.whl

sky/utils/subprocess_utils.py

@@ -2,21 +2,26 @@
 from multiprocessing import pool
 import os
 import random
+import resource
+import shlex
 import subprocess
 import time
-from typing import Any, Callable, Iterable, List, Optional, Tuple, Union
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
 import colorama
 import psutil
 
 from sky import exceptions
 from sky import sky_logging
+from sky.skylet import constants
 from sky.skylet import log_lib
 from sky.utils import timeline
 from sky.utils import ux_utils
 
 logger = sky_logging.init_logger(__name__)
 
+_fd_limit_warning_shown = False
+
 
 @timeline.event
 def run(cmd, **kwargs):
@@ -42,27 +47,86 @@ def run_no_outputs(cmd, **kwargs):
                **kwargs)
 
 
-def get_parallel_threads() -> int:
-    """Returns the number of idle CPUs."""
+def _get_thread_multiplier(cloud_str: Optional[str] = None) -> int:
+    # If using Kubernetes, we use 4x the number of cores.
+    if cloud_str and cloud_str.lower() == 'kubernetes':
+        return 4
+    return 1
+
+
+def get_max_workers_for_file_mounts(common_file_mounts: Dict[str, str],
+                                    cloud_str: Optional[str] = None) -> int:
+    global _fd_limit_warning_shown
+    fd_limit, _ = resource.getrlimit(resource.RLIMIT_NOFILE)
+
+    # Raise warning for low fd_limit (only once)
+    if fd_limit < 1024 and not _fd_limit_warning_shown:
+        logger.warning(
+            f'Open file descriptor limit ({fd_limit}) is low. File sync to '
+            'remote clusters may be slow. Consider increasing the limit using '
+            '`ulimit -n <number>` or modifying system limits.')
+        _fd_limit_warning_shown = True
+
+    fd_per_rsync = 5
+    for src in common_file_mounts.values():
+        if os.path.isdir(src):
+            # Assume that each file/folder under src takes 5 file descriptors
+            # on average.
+            fd_per_rsync = max(fd_per_rsync, len(os.listdir(src)) * 5)
+
+    # Reserve some file descriptors for the system and other processes
+    fd_reserve = 100
+
+    max_workers = (fd_limit - fd_reserve) // fd_per_rsync
+    # At least 1 worker, and avoid too many workers overloading the system.
+    num_threads = get_parallel_threads(cloud_str)
+    max_workers = min(max(max_workers, 1), num_threads)
+    logger.debug(f'Using {max_workers} workers for file mounts.')
+    return max_workers
+
+
+def get_parallel_threads(cloud_str: Optional[str] = None) -> int:
+    """Returns the number of threads to use for parallel execution.
+
+    Args:
+        cloud_str: The cloud
+    """
     cpu_count = os.cpu_count()
     if cpu_count is None:
         cpu_count = 1
-    return max(4, cpu_count - 1)
+    return max(4, cpu_count - 1) * _get_thread_multiplier(cloud_str)
 
 
-def run_in_parallel(func: Callable, args: Iterable[Any]) -> List[Any]:
+def run_in_parallel(func: Callable,
+                    args: List[Any],
+                    num_threads: Optional[int] = None) -> List[Any]:
     """Run a function in parallel on a list of arguments.
 
-    The function 'func' should raise a CommandError if the command fails.
+    Args:
+        func: The function to run in parallel
+        args: Iterable of arguments to pass to func
+        num_threads: Number of threads to use. If None, uses
+          get_parallel_threads()
 
     Returns:
       A list of the return values of the function func, in the same order as the
-      arguments.
+        arguments.
+
+    Raises:
+        Exception: The first exception encountered.
     """
-    # Reference: https://stackoverflow.com/questions/25790279/python-multiprocessing-early-termination # pylint: disable=line-too-long
-    with pool.ThreadPool(processes=get_parallel_threads()) as p:
-        # Run the function in parallel on the arguments, keeping the order.
-        return list(p.imap(func, args))
+    # Short-circuit for short lists
+    if len(args) == 0:
+        return []
+    if len(args) == 1:
+        return [func(args[0])]
+
+    processes = (num_threads
+                 if num_threads is not None else get_parallel_threads())
+
+    with pool.ThreadPool(processes=processes) as p:
+        ordered_iterators = p.imap(func, args)
+        return list(ordered_iterators)
 
 
 def handle_returncode(returncode: int,
@@ -77,8 +141,9 @@ def handle_returncode(returncode: int,
         command: The command that was run.
         error_msg: The error message to print.
         stderr: The stderr of the command.
+        stream_logs: Whether to stream logs.
     """
-    echo = logger.error if stream_logs else lambda _: None
+    echo = logger.error if stream_logs else logger.debug
     if returncode != 0:
         if stderr is not None:
             echo(stderr)
@@ -92,9 +157,9 @@ def handle_returncode(returncode: int,
                                           stderr)
 
 
-def kill_children_processes(
-        first_pid_to_kill: Optional[Union[int, List[Optional[int]]]] = None,
-        force: bool = False):
+def kill_children_processes(parent_pids: Optional[Union[
+    int, List[Optional[int]]]] = None,
+                            force: bool = False) -> None:
     """Kill children processes recursively.
 
     We need to kill the children, so that
@@ -104,41 +169,57 @@ def kill_children_processes(
        etc. while we are cleaning up the clusters.
 
     Args:
-        first_pid_to_kill: Optional PID of a process, or PIDs of a series of
-         processes to be killed first. If a list of PID is specified, it is
-         killed by the order in the list.
-         This is for guaranteeing the order of cleaning up and suppress
-         flaky errors.
+        parent_pids: Optional PIDs of a series of processes. The processes and
+          their children will be killed.  If a list of PID is specified, it is
+          killed by the order in the list. This is for guaranteeing the order
+          of cleaning up and suppress flaky errors.
+        force: bool, send SIGKILL if force, otherwise, use SIGTERM for
+          gracefully kill the process.
     """
-    pid_to_proc = dict()
-    child_processes = []
-    if isinstance(first_pid_to_kill, int):
-        first_pid_to_kill = [first_pid_to_kill]
-    elif first_pid_to_kill is None:
-        first_pid_to_kill = []
-
-    def _kill_processes(processes: List[psutil.Process]) -> None:
-        for process in processes:
+    if isinstance(parent_pids, int):
+        parent_pids = [parent_pids]
+
+    def kill(proc: psutil.Process):
+        if not proc.is_running():
+            # Skip if the process is not running.
+            return
+        logger.debug(f'Killing process {proc.pid}')
+        try:
+            if force:
+                proc.kill()
+            else:
+                proc.terminate()
+            proc.wait(timeout=10)
+        except psutil.NoSuchProcess:
+            # The child process may have already been terminated.
+            pass
+        except psutil.TimeoutExpired:
+            logger.debug(
+                f'Process {proc.pid} did not terminate after 10 seconds')
+            # Attempt to force kill if the normal termination fails
+            if not force:
+                logger.debug(f'Force killing process {proc.pid}')
+                proc.kill()
+                proc.wait(timeout=5)  # Shorter timeout after force kill
+
+    parent_processes = []
+    if parent_pids is None:
+        parent_processes = [psutil.Process()]
+    else:
+        for pid in parent_pids:
             try:
-                if force:
-                    process.kill()
-                else:
-                    process.terminate()
+                process = psutil.Process(pid)
             except psutil.NoSuchProcess:
-                # The process may have already been terminated.
-                pass
-
-    parent_process = psutil.Process()
-    for child in parent_process.children(recursive=True):
-        if child.pid in first_pid_to_kill:
-            pid_to_proc[child.pid] = child
-        else:
-            child_processes.append(child)
+                continue
+            parent_processes.append(process)
 
-    _kill_processes([
-        pid_to_proc[proc] for proc in first_pid_to_kill if proc in pid_to_proc
-    ])
-    _kill_processes(child_processes)
+    for parent_process in parent_processes:
+        child_processes = parent_process.children(recursive=True)
+        if parent_pids is not None:
+            kill(parent_process)
+        logger.debug(f'Killing child processes: {child_processes}')
+        for child in child_processes:
+            kill(child)
 
 
 def run_with_retries(
@@ -187,3 +268,88 @@ def run_with_retries(
                 continue
         break
     return returncode, stdout, stderr
+
+
+def kill_process_daemon(process_pid: int) -> None:
+    """Start a daemon as a safety net to kill the process.
+
+    Args:
+        process_pid: The PID of the process to kill.
+    """
+    # Get initial children list
+    try:
+        process = psutil.Process(process_pid)
+        initial_children = [p.pid for p in process.children(recursive=True)]
+    except psutil.NoSuchProcess:
+        initial_children = []
+
+    parent_pid = os.getpid()
+    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,
+        daemon_script,
+        '--parent-pid',
+        str(parent_pid),
+        '--proc-pid',
+        str(process_pid),
+        # We pass the initial children list to avoid the race condition where
+        # the process_pid is terminated before the daemon starts and gets the
+        # children list.
+        '--initial-children',
+        ','.join(map(str, initial_children)),
+    ]
+
+    # 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
+    # daemonize the process in `sky/skylet/subprocess_daemon.py`.
+    subprocess.Popen(
+        daemon_cmd,
+        # Suppress output
+        stdout=subprocess.DEVNULL,
+        stderr=subprocess.DEVNULL,
+        # Disable input
+        stdin=subprocess.DEVNULL,
+    )
+
+
+def launch_new_process_tree(cmd: str, log_output: str = '/dev/null') -> int:
+    """Launch a new process that will not be a child of the current process.
+
+    This will launch bash in a new session, which will launch the given cmd.
+    This will ensure that cmd is in its own process tree, and once bash exits,
+    will not be an ancestor of the current process. This is useful for job
+    launching.
+
+    Returns the pid of the launched cmd.
+    """
+    # Use nohup to ensure the job driver process is a separate process tree,
+    # instead of being a child of the current process. This is important to
+    # avoid a chain of driver processes (job driver can call schedule_step() to
+    # submit new jobs, and the new job can also call schedule_step()
+    # recursively).
+    #
+    # echo $! will output the PID of the last background process started in the
+    # current shell, so we can retrieve it and record in the DB.
+    #
+    # TODO(zhwu): A more elegant solution is to use another daemon process to be
+    # in charge of starting these driver processes, instead of starting them in
+    # the current process.
+    wrapped_cmd = (f'nohup bash -c {shlex.quote(cmd)} '
+                   f'</dev/null >{log_output} 2>&1 & echo $!')
+    proc = subprocess.run(wrapped_cmd,
+                          stdout=subprocess.PIPE,
+                          stderr=subprocess.PIPE,
+                          stdin=subprocess.DEVNULL,
+                          start_new_session=True,
+                          check=True,
+                          shell=True,
+                          text=True)
+    # Get the PID of the detached process
+    return int(proc.stdout.strip())

sky/utils/timeline.py

@@ -9,6 +9,7 @@ import json
 import os
 import threading
 import time
+import traceback
 from typing import Callable, Optional, Union
 
 import filelock
@@ -48,8 +49,9 @@ class Event:
             'ph': 'B',
             'ts': f'{time.time() * 10 ** 6: .3f}',
         })
+        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
         _events.append(event_begin)
 
     def end(self):
@@ -77,11 +79,11 @@ def event(name_or_fn: Union[str, Callable], message: Optional[str] = None):
 class FileLockEvent:
     """Serve both as a file lock and event for the lock."""
 
-    def __init__(self, lockfile: Union[str, os.PathLike]):
+    def __init__(self, lockfile: Union[str, os.PathLike], timeout: float = -1):
         self._lockfile = lockfile
-        # TODO(mraheja): remove pylint disabling when filelock version updated
-        # pylint: disable=abstract-class-instantiated
-        self._lock = filelock.FileLock(self._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):
@@ -116,7 +118,10 @@ class FileLockEvent:
         return wrapper
 
 
-def _save_timeline(file_path: str):
+def save_timeline():
+    file_path = os.environ.get('SKYPILOT_TIMELINE_FILE_PATH')
+    if not file_path:
+        return
     json_output = {
         'traceEvents': _events,
         'displayTimeUnit': 'ms',
@@ -130,4 +135,4 @@ def _save_timeline(file_path: str):
 
 
 if os.environ.get('SKYPILOT_TIMELINE_FILE_PATH'):
-    atexit.register(_save_timeline, os.environ['SKYPILOT_TIMELINE_FILE_PATH'])
+    atexit.register(save_timeline)

sky/utils/ux_utils.py

@@ -1,21 +1,42 @@
 """Utility functions for UX."""
 import contextlib
+import enum
+import os
 import sys
 import traceback
-from typing import Callable
+import typing
+from typing import Callable, Optional, Union
 
+import colorama
 import rich.console as rich_console
 
 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 ux_utils
+
+if typing.TYPE_CHECKING:
+    import pathlib
 
 console = rich_console.Console()
 
+INDENT_SYMBOL = f'{colorama.Style.DIM}├── {colorama.Style.RESET_ALL}'
+INDENT_LAST_SYMBOL = f'{colorama.Style.DIM}└── {colorama.Style.RESET_ALL}'
+
+# Console formatting constants
+BOLD = '\033[1m'
+RESET_BOLD = '\033[0m'
+
+# Log path hint in the spinner during launching
+_LOG_PATH_HINT = (f'{colorama.Style.DIM}View logs: sky api logs -l '
+                  '{log_path}'
+                  f'{colorama.Style.RESET_ALL}')
+_LOG_PATH_HINT_LOCAL = (f'{colorama.Style.DIM}View logs: '
+                        '{log_path}'
+                        f'{colorama.Style.RESET_ALL}')
+
 
 def console_newline():
-    """Print a newline to the console using rich.
+    """Prints a newline to the console using rich.
 
     Useful when catching exceptions inside console.status()
     """
@@ -38,19 +59,15 @@ def print_exception_no_traceback():
             if error():
                 raise ValueError('...')
     """
-    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
+    original_tracelimit = getattr(sys, 'tracebacklimit', 1000)
+    sys.tracebacklimit = 0
+    yield
+    sys.tracebacklimit = original_tracelimit
 
 
 @contextlib.contextmanager
 def enable_traceback():
-    """Revert the effect of print_exception_no_traceback().
+    """Reverts the effect of print_exception_no_traceback().
 
     This is used for usage_lib to collect the full traceback.
     """
@@ -61,7 +78,7 @@ def enable_traceback():
 
 
 class RedirectOutputForProcess:
-    """Redirect stdout and stderr to a file.
+    """Redirects stdout and stderr to a file.
 
     This class enabled output redirect for multiprocessing.Process.
     Example usage:
@@ -99,6 +116,142 @@ class RedirectOutputForProcess:
             except Exception as e:  # pylint: disable=broad-except
                 logger.error(f'Failed to run {self.func.__name__}. '
                              f'Details: {common_utils.format_exception(e)}')
-                with ux_utils.enable_traceback():
+                with enable_traceback():
                     logger.error(f'  Traceback:\n{traceback.format_exc()}')
                 raise
+
+
+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."""
+    log_path = str(log_path)
+    expanded_home = os.path.expanduser('~')
+    if log_path.startswith(expanded_home):
+        log_path = '~' + log_path[len(expanded_home):]
+    if is_local:
+        return _LOG_PATH_HINT_LOCAL.format(log_path=log_path)
+    if log_path.startswith(constants.SKY_LOGS_DIRECTORY):
+        log_path = log_path[len(constants.SKY_LOGS_DIRECTORY):]
+    log_path = log_path.lstrip(os.path.sep)
+    return _LOG_PATH_HINT.format(log_path=log_path)
+
+
+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
+    # previous spinner with dimmed color overflows in a narrow terminal, the
+    # color might be messed up.
+    return f'{colorama.Style.RESET_ALL}⚙︎ {message}'
+
+
+def finishing_message(message: str,
+                      log_path: Optional[Union[str, 'pathlib.Path']] = None,
+                      is_local: bool = False,
+                      follow_up_message: Optional[str] = None) -> str:
+    """Gets the finishing message for the given message.
+
+    Args:
+        message: The main message to be displayed.
+        log_path: The log path to be displayed in the message.
+        is_local: Whether the log path is local or on remote API server.
+        follow_up_message: A message to be displayed after the main message.
+          The follow up message is not colored.
+    """
+    # 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.
+    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}')
+    if log_path is None:
+        return success_prefix
+    path_hint = log_path_hint(log_path, is_local)
+    return f'{success_prefix}  {path_hint}'
+
+
+def error_message(message: str,
+                  log_path: Optional[Union[str, 'pathlib.Path']] = None,
+                  is_local: bool = False) -> 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}')
+    if log_path is None:
+        return error_prefix
+    path_hint = log_path_hint(log_path, is_local)
+    return f'{error_prefix}  {path_hint}'
+
+
+def retry_message(message: str) -> str:
+    """Gets the retry 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.
+    return (f'{colorama.Style.RESET_ALL}{colorama.Fore.YELLOW}↺'
+            f'{colorama.Style.RESET_ALL} {message}')
+
+
+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."""
+    colored_spinner = f'[bold cyan]{message}[/]'
+    if log_path is None:
+        return colored_spinner
+    path_hint = log_path_hint(log_path, is_local)
+    return f'{colored_spinner}  {path_hint}'
+
+
+class CommandHintType(enum.Enum):
+    CLUSTER_JOB = 'cluster_job'
+    MANAGED_JOB = 'managed_job'
+
+
+def command_hint_messages(hint_type: CommandHintType,
+                          job_id: Optional[str] = None,
+                          cluster_name: Optional[str] = None) -> str:
+    """Gets the command hint messages for the given job id."""
+    if hint_type == CommandHintType.CLUSTER_JOB:
+        job_hint_str = (f'\nJob ID: {job_id}'
+                        f'\n{INDENT_SYMBOL}To cancel the job:\t\t'
+                        f'{BOLD}sky cancel {cluster_name} {job_id}{RESET_BOLD}'
+                        f'\n{INDENT_SYMBOL}To stream job logs:\t\t'
+                        f'{BOLD}sky logs {cluster_name} {job_id}{RESET_BOLD}'
+                        f'\n{INDENT_LAST_SYMBOL}To view job queue:\t\t'
+                        f'{BOLD}sky queue {cluster_name}{RESET_BOLD}')
+        cluster_hint_str = (f'\nCluster name: {cluster_name}'
+                            f'\n{INDENT_SYMBOL}To log into the head VM:\t'
+                            f'{BOLD}ssh {cluster_name}'
+                            f'{RESET_BOLD}'
+                            f'\n{INDENT_SYMBOL}To submit a job:'
+                            f'\t\t{BOLD}sky exec {cluster_name} yaml_file'
+                            f'{RESET_BOLD}'
+                            f'\n{INDENT_SYMBOL}To stop the cluster:'
+                            f'\t{BOLD}sky stop {cluster_name}'
+                            f'{RESET_BOLD}'
+                            f'\n{INDENT_LAST_SYMBOL}To teardown the cluster:'
+                            f'\t{BOLD}sky down {cluster_name}'
+                            f'{RESET_BOLD}')
+        hint_str = '\n📋 Useful Commands'
+        if job_id is not None:
+            hint_str += f'{job_hint_str}'
+        hint_str += f'{cluster_hint_str}'
+        return hint_str
+    elif hint_type == CommandHintType.MANAGED_JOB:
+        return (f'\n📋 Useful Commands'
+                f'\nManaged Job ID: {job_id}'
+                f'\n{INDENT_SYMBOL}To cancel the job:\t\t'
+                f'{BOLD}sky jobs cancel {job_id}{RESET_BOLD}'
+                f'\n{INDENT_SYMBOL}To stream job logs:\t\t'
+                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}')
+    else:
+        raise ValueError(f'Invalid hint type: {hint_type}')