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

sky/skylet/attempt_skylet.py

@@ -1,51 +1,133 @@
 """Restarts skylet if version does not match"""
 
 import os
+import signal
 import subprocess
+from typing import List, Optional, Tuple
+
+import psutil
 
 from sky.skylet import constants
+from sky.skylet import runtime_utils
+
+VERSION_FILE = runtime_utils.get_runtime_dir_path(constants.SKYLET_VERSION_FILE)
+SKYLET_LOG_FILE = runtime_utils.get_runtime_dir_path(constants.SKYLET_LOG_FILE)
+PID_FILE = runtime_utils.get_runtime_dir_path(constants.SKYLET_PID_FILE)
+PORT_FILE = runtime_utils.get_runtime_dir_path(constants.SKYLET_PORT_FILE)
+
+
+def _is_running_skylet_process(pid: int) -> bool:
+    if pid <= 0:
+        return False
+    try:
+        process = psutil.Process(pid)
+        if not process.is_running():
+            return False
+        # Check if command line contains the skylet module identifier
+        cmdline = process.cmdline()
+        return any('sky.skylet.skylet' in arg for arg in cmdline)
+    except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess,
+            OSError) as e:
+        print(f'Error checking if skylet process {pid} is running: {e}')
+        return False
+
+
+def _find_running_skylet_pids() -> List[int]:
+    if os.path.exists(PID_FILE):
+        try:
+            with open(PID_FILE, 'r', encoding='utf-8') as pid_file:
+                pid = int(pid_file.read().strip())
+            if _is_running_skylet_process(pid):
+                return [pid]
+        except (OSError, ValueError, IOError) as e:
+            # Don't fallback to grep-based detection as the existence of the
+            # PID file implies that we are on the new version, and there is
+            # possibility of there being multiple skylet processes running,
+            # and we don't want to accidentally kill the wrong skylet(s).
+            print(f'Error reading PID file {PID_FILE}: {e}')
+        return []
+    else:
+        # Fall back to grep-based detection for backward compatibility.
+        pids = []
+        # We use -m to grep instead of {constants.SKY_PYTHON_CMD} -m to grep
+        # because need to handle the backward compatibility of the old skylet
+        # started before #3326, which does not use the full path to python.
+        proc = subprocess.run(
+            'ps aux | grep -v "grep" | grep "sky.skylet.skylet" | grep " -m"',
+            shell=True,
+            check=False,
+            capture_output=True,
+            text=True)
+        if proc.returncode == 0:
+            # Parse the output to extract PIDs (column 2)
+            for line in proc.stdout.strip().split('\n'):
+                if line:
+                    parts = line.split()
+                    if len(parts) >= 2:
+                        try:
+                            pids.append(int(parts[1]))
+                        except ValueError:
+                            continue
+        return pids
+
+
+def _check_version_match() -> Tuple[bool, Optional[str]]:
+    """Check if the version file matches the current skylet version.
 
-VERSION_FILE = os.path.expanduser(constants.SKYLET_VERSION_FILE)
+    Returns:
+        Tuple of (version_match: bool, version: str or None)
+    """
+    version: Optional[str] = None
+    if os.path.exists(VERSION_FILE):
+        try:
+            with open(VERSION_FILE, 'r', encoding='utf-8') as f:
+                version = f.read().strip()
+                return version == constants.SKYLET_VERSION, version
+        except (OSError, IOError):
+            pass
+    return False, version
 
 
 def restart_skylet():
     # Kills old skylet if it is running.
     # TODO(zhwu): make the killing graceful, e.g., use a signal to tell
     # skylet to exit, instead of directly killing it.
-    subprocess.run(
-        # We use -m to grep instead of {constants.SKY_PYTHON_CMD} -m to grep
-        # because need to handle the backward compatibility of the old skylet
-        # started before #3326, which does not use the full path to python.
-        'ps aux | grep "sky.skylet.skylet" | grep " -m "'
-        '| awk \'{print $2}\' | xargs kill >> ~/.sky/skylet.log 2>&1',
-        shell=True,
-        check=False)
+
+    # Find and kill running skylet processes
+    for pid in _find_running_skylet_pids():
+        try:
+            os.kill(pid, signal.SIGKILL)
+        except (OSError, ProcessLookupError):
+            # Process died between detection and kill
+            pass
+    # Clean up the PID file
+    try:
+        os.remove(PID_FILE)
+    except OSError:
+        pass  # Best effort cleanup
+
+    port = constants.SKYLET_GRPC_PORT
     subprocess.run(
         # We have made sure that `attempt_skylet.py` is executed with the
         # skypilot runtime env activated, so that skylet can access the cloud
         # CLI tools.
-        f'nohup {constants.SKY_PYTHON_CMD} -m sky.skylet.skylet'
-        ' >> ~/.sky/skylet.log 2>&1 &',
+        f'nohup {constants.SKY_PYTHON_CMD} -m sky.skylet.skylet '
+        f'--port={port} '
+        f'>> {SKYLET_LOG_FILE} 2>&1 & echo $! > {PID_FILE}',
         shell=True,
         check=True)
+
+    with open(PORT_FILE, 'w', encoding='utf-8') as pf:
+        pf.write(str(port))
+
     with open(VERSION_FILE, 'w', encoding='utf-8') as v_f:
         v_f.write(constants.SKYLET_VERSION)
 
 
-proc = subprocess.run(
-    'ps aux | grep -v "grep" | grep "sky.skylet.skylet" | grep " -m"',
-    shell=True,
-    check=False)
-
-running = (proc.returncode == 0)
+# Check if our skylet is running
+running = bool(_find_running_skylet_pids())
 
-version_match = False
-found_version = None
-if os.path.exists(VERSION_FILE):
-    with open(VERSION_FILE, 'r', encoding='utf-8') as f:
-        found_version = f.read().strip()
-        if found_version == constants.SKYLET_VERSION:
-            version_match = True
+version_match, found_version = _check_version_match()
 
 version_string = (f' (found version {found_version}, new version '
                   f'{constants.SKYLET_VERSION})')

sky/skylet/autostop_lib.py

@@ -1,6 +1,8 @@
 """Autostop utilities."""
+import enum
 import pickle
 import shlex
+import subprocess
 import time
 import typing
 from typing import List, Optional
@@ -10,11 +12,17 @@ from sky.adaptors import common as adaptors_common
 from sky.skylet import configs
 from sky.skylet import constants
 from sky.utils import message_utils
+from sky.utils import ux_utils
 
 if typing.TYPE_CHECKING:
     import psutil
+
+    from sky.schemas.generated import autostopv1_pb2
 else:
     psutil = adaptors_common.LazyImport('psutil')
+    # To avoid requiring protobuf to be installed on the client side.
+    autostopv1_pb2 = adaptors_common.LazyImport(
+        'sky.schemas.generated.autostopv1_pb2')
 
 logger = sky_logging.init_logger(__name__)
 
@@ -30,6 +38,83 @@ _AUTOSTOP_LAST_ACTIVE_TIME = 'autostop_last_active_time'
 _AUTOSTOP_INDICATOR = 'autostop_indicator'
 
 
+class AutostopWaitFor(enum.Enum):
+    """Enum for the Autostop behaviour.
+
+    JOBS: Wait for jobs to finish.
+    JOBS_AND_SSH: Wait for jobs to finish and all SSH sessions to be closed.
+    NONE: Unconditionally stop the cluster after the idle time.
+    """
+    JOBS_AND_SSH = 'jobs_and_ssh'
+    JOBS = 'jobs'
+    NONE = 'none'
+
+    @classmethod
+    def supported_modes(cls) -> List[str]:
+        return [mode.value for mode in cls]
+
+    @classmethod
+    def cli_help_message(cls, pair: str) -> str:
+        return f"""\
+Determines the condition for resetting the idleness timer.
+This option works in conjunction with ``--{pair}``. Options:
+
+\b
+1. ``jobs_and_ssh`` (default): Wait for in-progress jobs and SSH connections to finish.
+2. ``jobs``: Only wait for in-progress jobs.
+3. ``none``: Wait for nothing; autostop right after ``{pair}``."""
+
+    @classmethod
+    def from_str(cls, mode: str) -> 'AutostopWaitFor':
+        """Returns the enum value for the given string."""
+        if mode.lower() == cls.JOBS.value:
+            return cls.JOBS
+        elif mode.lower() == cls.JOBS_AND_SSH.value:
+            return cls.JOBS_AND_SSH
+        elif mode.lower() == cls.NONE.value:
+            return cls.NONE
+        else:
+            with ux_utils.print_exception_no_traceback():
+                raise ValueError(f'Unsupported autostop wait mode: '
+                                 f'{mode}. The mode must be either '
+                                 f'\'{cls.JOBS_AND_SSH.value}\', '
+                                 f'\'{cls.JOBS.value}\', or '
+                                 f'\'{cls.NONE.value}\'. ')
+
+    @classmethod
+    def from_protobuf(
+        cls, protobuf_value: 'autostopv1_pb2.AutostopWaitFor'
+    ) -> Optional['AutostopWaitFor']:
+        """Convert protobuf AutostopWaitFor enum to Python enum value."""
+        protobuf_to_enum = {
+            autostopv1_pb2.AUTOSTOP_WAIT_FOR_JOBS_AND_SSH: cls.JOBS_AND_SSH,
+            autostopv1_pb2.AUTOSTOP_WAIT_FOR_JOBS: cls.JOBS,
+            autostopv1_pb2.AUTOSTOP_WAIT_FOR_NONE: cls.NONE,
+            autostopv1_pb2.AUTOSTOP_WAIT_FOR_UNSPECIFIED: None,
+        }
+        if protobuf_value not in protobuf_to_enum:
+            with ux_utils.print_exception_no_traceback():
+                raise ValueError(
+                    f'Unknown protobuf AutostopWaitFor value: {protobuf_value}')
+        return protobuf_to_enum[protobuf_value]
+
+    def to_protobuf(self) -> 'autostopv1_pb2.AutostopWaitFor':
+        """Convert this Python enum value to protobuf enum value."""
+        enum_to_protobuf = {
+            AutostopWaitFor.JOBS_AND_SSH:
+                autostopv1_pb2.AUTOSTOP_WAIT_FOR_JOBS_AND_SSH,
+            AutostopWaitFor.JOBS: autostopv1_pb2.AUTOSTOP_WAIT_FOR_JOBS,
+            AutostopWaitFor.NONE: autostopv1_pb2.AUTOSTOP_WAIT_FOR_NONE,
+        }
+        if self not in enum_to_protobuf:
+            with ux_utils.print_exception_no_traceback():
+                raise ValueError(f'Unknown AutostopWaitFor value: {self}')
+        return enum_to_protobuf[self]
+
+
+DEFAULT_AUTOSTOP_WAIT_FOR: AutostopWaitFor = AutostopWaitFor.JOBS_AND_SSH
+
+
 class AutostopConfig:
     """Autostop configuration."""
 
@@ -37,12 +122,14 @@ class AutostopConfig:
                  autostop_idle_minutes: int,
                  boot_time: float,
                  backend: Optional[str],
+                 wait_for: AutostopWaitFor,
                  down: bool = False):
         assert autostop_idle_minutes < 0 or backend is not None, (
             autostop_idle_minutes, backend)
         self.autostop_idle_minutes = autostop_idle_minutes
         self.boot_time = boot_time
         self.backend = backend
+        self.wait_for = wait_for
         self.down = down
 
     def __setstate__(self, state: dict):
@@ -53,15 +140,18 @@ class AutostopConfig:
 def get_autostop_config() -> AutostopConfig:
     config_str = configs.get_config(_AUTOSTOP_CONFIG_KEY)
     if config_str is None:
-        return AutostopConfig(-1, -1, None)
+        return AutostopConfig(-1, -1, None, DEFAULT_AUTOSTOP_WAIT_FOR)
     return pickle.loads(config_str)
 
 
-def set_autostop(idle_minutes: int, backend: Optional[str], down: bool) -> None:
+def set_autostop(idle_minutes: int, backend: Optional[str],
+                 wait_for: AutostopWaitFor, down: bool) -> None:
     boot_time = psutil.boot_time()
-    autostop_config = AutostopConfig(idle_minutes, boot_time, backend, down)
+    autostop_config = AutostopConfig(idle_minutes, boot_time, backend, wait_for,
+                                     down)
     configs.set_config(_AUTOSTOP_CONFIG_KEY, pickle.dumps(autostop_config))
-    logger.debug(f'set_autostop(): idle_minutes {idle_minutes}, down {down}.')
+    logger.debug(f'set_autostop(): idle_minutes {idle_minutes}, down {down}, '
+                 f'wait_for {wait_for.value}.')
     # Reset timer whenever an autostop setting is submitted, i.e. the idle
     # time will be counted from now.
     set_last_active_time_to_now()
@@ -107,6 +197,28 @@ def set_last_active_time_to_now() -> None:
     configs.set_config(_AUTOSTOP_LAST_ACTIVE_TIME, str(time.time()))
 
 
+def has_active_ssh_sessions() -> bool:
+    """Returns True if there are any active SSH sessions on the node."""
+    try:
+        # /dev/pts is a virtual filesystem that contains the pseudo-terminal
+        # devices. ptmx is the pseudo-terminal multiplexer, which is the
+        # "master" device that creates new pseudo-terminal devices, so we
+        # exclude it from the count.
+        proc = subprocess.run('ls /dev/pts | grep -v ptmx | wc -l',
+                              capture_output=True,
+                              text=True,
+                              check=False,
+                              shell=True)
+        if proc.returncode != 0:
+            logger.warning(f'SSH session check command failed with return code '
+                           f'{proc.returncode}.')
+            return False
+        return int(proc.stdout.strip()) > 0
+    except Exception as e:  # pylint: disable=broad-except
+        logger.warning(f'Error checking active SSH sessions: {e}.')
+        return False
+
+
 class AutostopCodeGen:
     """Code generator for autostop utility functions.
 
@@ -114,13 +226,22 @@ class AutostopCodeGen:
 
       >> codegen = AutostopCodeGen.set_autostop(...)
     """
-    _PREFIX = ['from sky.skylet import autostop_lib']
+    _PREFIX = ['from sky.skylet import autostop_lib, constants']
 
     @classmethod
-    def set_autostop(cls, idle_minutes: int, backend: str, down: bool) -> str:
+    def set_autostop(cls,
+                     idle_minutes: int,
+                     backend: str,
+                     wait_for: Optional[AutostopWaitFor],
+                     down: bool = False) -> str:
+        if wait_for is None:
+            wait_for = DEFAULT_AUTOSTOP_WAIT_FOR
         code = [
-            f'autostop_lib.set_autostop({idle_minutes}, {backend!r},'
-            f' {down})',
+            f'\nif getattr(constants, "SKYLET_LIB_VERSION", 1) < 4: '
+            f'\n autostop_lib.set_autostop({idle_minutes}, {backend!r}, {down})'
+            f'\nelse: '
+            f'\n autostop_lib.set_autostop({idle_minutes}, {backend!r}, '
+            f'autostop_lib.{wait_for}, {down})',
         ]
         return cls._build(code)
 

sky/skylet/configs.py

@@ -2,17 +2,17 @@
 import functools
 import os
 import pathlib
+import threading
 from typing import Callable, Optional, Union
 
-from sky.utils import db_utils
+from sky.skylet import runtime_utils
+from sky.utils.db import db_utils
 
-_DB_PATH = os.path.expanduser('~/.sky/skylet_config.db')
-os.makedirs(pathlib.Path(_DB_PATH).parents[0], exist_ok=True)
+_DB_PATH = None
+_db_init_lock = threading.Lock()
 
-_table_created = False
 
-
-def ensure_table(func: Callable):
+def init_db(func: Callable):
     """Ensure the table exists before calling the function.
 
     Since this module will be imported whenever `sky` is imported (due to
@@ -24,25 +24,33 @@ def ensure_table(func: Callable):
 
     @functools.wraps(func)
     def wrapper(*args, **kwargs):
-        global _table_created
-        if not _table_created:
-            with db_utils.safe_cursor(
-                    _DB_PATH) as c:  # Call it 'c' to avoid pylint complaining.
-                # Use WAL mode to avoid locking problem in #1507.
-                # Reference: https://stackoverflow.com/a/39265148
-                c.execute('PRAGMA journal_mode=WAL')
-                c.execute("""\
-                    CREATE TABLE IF NOT EXISTS config (
-                        key TEXT PRIMARY KEY,
-                        value TEXT)""")
-        _table_created = True
+        global _DB_PATH
+        if _DB_PATH is not None:
+            return func(*args, **kwargs)
+
+        with _db_init_lock:
+            if _DB_PATH is None:
+                _DB_PATH = runtime_utils.get_runtime_dir_path(
+                    '.sky/skylet_config.db')
+                os.makedirs(pathlib.Path(_DB_PATH).parents[0], exist_ok=True)
+                with db_utils.safe_cursor(
+                        _DB_PATH
+                ) as c:  # Call it 'c' to avoid pylint complaining.
+                    # Use WAL mode to avoid locking problem in #1507.
+                    # Reference: https://stackoverflow.com/a/39265148
+                    c.execute('PRAGMA journal_mode=WAL')
+                    c.execute("""\
+                        CREATE TABLE IF NOT EXISTS config (
+                            key TEXT PRIMARY KEY,
+                            value TEXT)""")
         return func(*args, **kwargs)
 
     return wrapper
 
 
-@ensure_table
+@init_db
 def get_config(key: str) -> Optional[bytes]:
+    assert _DB_PATH is not None
     with db_utils.safe_cursor(_DB_PATH) as cursor:
         rows = cursor.execute('SELECT value FROM config WHERE key = ?', (key,))
         for (value,) in rows:
@@ -50,8 +58,9 @@ def get_config(key: str) -> Optional[bytes]:
         return None
 
 
-@ensure_table
+@init_db
 def set_config(key: str, value: Union[bytes, str]) -> None:
+    assert _DB_PATH is not None
     with db_utils.safe_cursor(_DB_PATH) as cursor:
         cursor.execute(
             """\