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

sky/skylet/job_lib.py

@@ -66,6 +66,7 @@ class JobInfoLoc(enum.IntEnum):
     PID = 9
     LOG_PATH = 10
     METADATA = 11
+    EXIT_CODES = 12
 
 
 def create_table(cursor, conn):
@@ -124,6 +125,8 @@ def create_table(cursor, conn):
                                  'metadata',
                                  'TEXT DEFAULT \'{}\'',
                                  value_to_replace_existing_entries='{}')
+    db_utils.add_column_to_table(cursor, conn, 'jobs', 'exit_codes',
+                                 'TEXT DEFAULT NULL')
     conn.commit()
 
 
@@ -388,10 +391,16 @@ def add_job(job_name: str,
     assert _DB is not None
     job_submitted_at = time.time()
     # job_id will autoincrement with the null value
-    _DB.cursor.execute(
-        'INSERT INTO jobs VALUES (null, ?, ?, ?, ?, ?, ?, null, ?, 0, null, ?)',
-        (job_name, username, job_submitted_at, JobStatus.INIT.value,
-         run_timestamp, None, resources_str, metadata))
+    if int(constants.SKYLET_VERSION) >= 28:
+        _DB.cursor.execute(
+            'INSERT INTO jobs VALUES (null, ?, ?, ?, ?, ?, ?, null, ?, 0, null, ?, null)',  # pylint: disable=line-too-long
+            (job_name, username, job_submitted_at, JobStatus.INIT.value,
+             run_timestamp, None, resources_str, metadata))
+    else:
+        _DB.cursor.execute(
+            'INSERT INTO jobs VALUES (null, ?, ?, ?, ?, ?, ?, null, ?, 0, null, ?)',  # pylint: disable=line-too-long
+            (job_name, username, job_submitted_at, JobStatus.INIT.value,
+             run_timestamp, None, resources_str, metadata))
     _DB.conn.commit()
     rows = _DB.cursor.execute('SELECT job_id FROM jobs WHERE run_timestamp=(?)',
                               (run_timestamp,))
@@ -468,6 +477,41 @@ def set_status(job_id: int, status: JobStatus) -> None:
         _set_status_no_lock(job_id, status)
 
 
+@init_db
+def set_exit_codes(job_id: int, exit_codes: List[int]) -> None:
+    """Set exit codes for a job as comma-separated string.
+
+    Args:
+        job_id: The job ID to update.
+        exit_codes: A list of exit codes to store.
+    """
+    assert _DB is not None
+    exit_codes_str = ','.join(str(code) for code in exit_codes)
+    with filelock.FileLock(_get_lock_path(job_id)):
+        _DB.cursor.execute('UPDATE jobs SET exit_codes=(?) WHERE job_id=(?)',
+                           (exit_codes_str, job_id))
+        _DB.conn.commit()
+
+
+@init_db
+def get_exit_codes(job_id: int) -> Optional[List[int]]:
+    """Get exit codes for a job from comma-separated string.
+
+    Args:
+        job_id: The job ID to retrieve exit codes for.
+
+    Returns:
+        A list of exit codes, or None if not found.
+    """
+    assert _DB is not None
+    rows = _DB.cursor.execute('SELECT exit_codes FROM jobs WHERE job_id=(?)',
+                              (job_id,))
+    row = rows.fetchone()
+    if row is None or row[0] is None:
+        return None
+    return [int(code) for code in row[0].split(',')]
+
+
 @init_db
 def set_job_started(job_id: int) -> None:
     # TODO(mraheja): remove pylint disabling when filelock version updated.
@@ -506,6 +550,20 @@ def get_status(job_id: int) -> Optional[JobStatus]:
         return get_status_no_lock(job_id)
 
 
+def wait_for_job_completion(job_id: int, poll_interval: float = 1.0) -> None:
+    """Wait for a job to reach a terminal state.
+
+    Args:
+        job_id: The job ID to wait for.
+        poll_interval: How often to poll the job status in seconds.
+    """
+    while True:
+        status = get_status(job_id)
+        if status is None or status.is_terminal():
+            break
+        time.sleep(poll_interval)
+
+
 @init_db
 def get_statuses_payload(job_ids: List[Optional[int]]) -> str:
     return message_utils.encode_payload(get_statuses(job_ids))
@@ -674,6 +732,14 @@ def _get_records_from_rows(rows) -> List[Dict[str, Any]]:
             'pid': row[JobInfoLoc.PID.value],
             'metadata': json.loads(row[JobInfoLoc.METADATA.value]),
         })
+        if int(constants.SKYLET_VERSION) >= 28:
+            exit_code_str = row[JobInfoLoc.EXIT_CODES.value]
+            if not isinstance(exit_code_str, str):
+                records[-1]['exit_codes'] = None
+            else:
+                records[-1]['exit_codes'] = ([
+                    int(code) for code in exit_code_str.split(',')
+                ])
     return records
 
 
@@ -1152,6 +1218,15 @@ class JobLibCodeGen:
         ]
         return cls._build(code)
 
+    @classmethod
+    def wait_for_job(cls, job_id: int) -> str:
+        code = [
+            # TODO(kevin): backward compatibility, remove in 0.13.0.
+            (f'job_lib.wait_for_job_completion({job_id!r}) if '
+             'hasattr(job_lib, "wait_for_job_completion") else None'),
+        ]
+        return cls._build(code)
+
     @classmethod
     def update_status(cls) -> str:
         code = ['job_lib.update_status()']
@@ -1269,8 +1344,19 @@ class JobLibCodeGen:
         ]
         return cls._build(code)
 
+    @classmethod
+    def get_job_exit_codes(cls, job_id: Optional[int] = None) -> str:
+        """Generate shell command to retrieve exit codes."""
+        code = [
+            f'job_id = {job_id} if {job_id} is not None else job_lib.get_latest_job_id()',  # pylint: disable=line-too-long
+            'exit_codes = job_lib.get_exit_codes(job_id) if job_id is not None and int(constants.SKYLET_VERSION) >= 28 else {}',  # pylint: disable=line-too-long
+            'print(exit_codes, flush=True)',
+        ]
+        return cls._build(code)
+
     @classmethod
     def _build(cls, code: List[str]) -> str:
         code = cls._PREFIX + code
         code = ';'.join(code)
-        return f'{constants.SKY_PYTHON_CMD} -u -c {shlex.quote(code)}'
+        return (f'{constants.ACTIVATE_SKY_REMOTE_PYTHON_ENV}; '
+                f'{constants.SKY_PYTHON_CMD} -u -c {shlex.quote(code)}')

sky/skylet/log_lib.py

@@ -172,7 +172,7 @@ def run_with_log(
     streaming_prefix: Optional[str] = None,
     log_cmd: bool = False,
     **kwargs,
-) -> Union[int, Tuple[int, str, str]]:
+) -> Union[int, Tuple[int, str, str], Tuple[int, int]]:
     """Runs a command and logs its output to a file.
 
     Args:
@@ -183,6 +183,8 @@ def run_with_log(
         process_stream: Whether to post-process the stdout/stderr of the
             command, such as replacing or skipping lines on the fly. If
             enabled, lines are printed only when '\r' or '\n' is found.
+        streaming_prefix: Optional prefix for each log line. Can contain {pid}
+            placeholder which will be replaced with the subprocess PID.
 
     Returns the returncode or returncode, stdout and stderr of the command.
       Note that the stdout and stderr is already decoded.
@@ -228,6 +230,13 @@ def run_with_log(
                 # For backward compatibility, do not specify use_kill_pg by
                 # default.
                 subprocess_utils.kill_process_daemon(proc.pid)
+
+            # Format streaming_prefix with subprocess PID if it contains {pid}
+            formatted_streaming_prefix = streaming_prefix
+            if streaming_prefix and '{pid}' in streaming_prefix:
+                formatted_streaming_prefix = streaming_prefix.format(
+                    pid=proc.pid)
+
             stdout = ''
             stderr = ''
             stdout_stream_handler = None
@@ -256,7 +265,7 @@ def run_with_log(
                     line_processor=line_processor,
                     # Replace CRLF when the output is logged to driver by ray.
                     replace_crlf=with_ray,
-                    streaming_prefix=streaming_prefix,
+                    streaming_prefix=formatted_streaming_prefix,
                 )
                 stdout_stream_handler = functools.partial(
                     _handle_io_stream,
@@ -349,7 +358,8 @@ def run_bash_command_with_log(bash_command: str,
                               log_path: str,
                               env_vars: Optional[Dict[str, str]] = None,
                               stream_logs: bool = False,
-                              with_ray: bool = False):
+                              with_ray: bool = False,
+                              streaming_prefix: Optional[str] = None):
     with tempfile.NamedTemporaryFile('w', prefix='sky_app_',
                                      delete=False) as fp:
         bash_command = make_task_bash_script(bash_command, env_vars=env_vars)
@@ -364,6 +374,7 @@ def run_bash_command_with_log(bash_command: str,
                             log_path,
                             stream_logs=stream_logs,
                             with_ray=with_ray,
+                            streaming_prefix=streaming_prefix,
                             shell=True)
 
 
@@ -372,9 +383,14 @@ def run_bash_command_with_log_and_return_pid(
         log_path: str,
         env_vars: Optional[Dict[str, str]] = None,
         stream_logs: bool = False,
-        with_ray: bool = False):
-    return_code = run_bash_command_with_log(bash_command, log_path, env_vars,
-                                            stream_logs, with_ray)
+        with_ray: bool = False,
+        streaming_prefix: Optional[str] = None):
+    return_code = run_bash_command_with_log(bash_command,
+                                            log_path,
+                                            env_vars,
+                                            stream_logs,
+                                            with_ray,
+                                            streaming_prefix=streaming_prefix)
     return {'return_code': return_code, 'pid': os.getpid()}
 
 

sky/skylet/log_lib.pyi

@@ -68,7 +68,7 @@ def run_with_log(cmd: Union[List[str], str],
                  process_stream: bool = ...,
                  line_processor: Optional[log_utils.LineProcessor] = ...,
                  streaming_prefix: Optional[str] = ...,
-                 ray_job_id: Optional[str] = ...,
+                 log_cmd: bool = ...,
                  **kwargs) -> int:
     ...
 
@@ -87,7 +87,7 @@ def run_with_log(cmd: Union[List[str], str],
                  process_stream: bool = ...,
                  line_processor: Optional[log_utils.LineProcessor] = ...,
                  streaming_prefix: Optional[str] = ...,
-                 ray_job_id: Optional[str] = ...,
+                 log_cmd: bool = ...,
                  **kwargs) -> Tuple[int, str, str]:
     ...
 
@@ -106,8 +106,8 @@ def run_with_log(cmd: Union[List[str], str],
                  process_stream: bool = ...,
                  line_processor: Optional[log_utils.LineProcessor] = ...,
                  streaming_prefix: Optional[str] = ...,
-                 ray_job_id: Optional[str] = ...,
-                 **kwargs) -> Union[int, Tuple[int, str, str]]:
+                 log_cmd: bool = ...,
+                 **kwargs) -> Tuple[int, int]:
     ...
 
 
@@ -125,7 +125,8 @@ def run_bash_command_with_log(bash_command: str,
                               log_path: str,
                               env_vars: Optional[Dict[str, str]] = ...,
                               stream_logs: bool = ...,
-                              with_ray: bool = ...):
+                              with_ray: bool = ...,
+                              streaming_prefix: Optional[str] = ...) -> int:
     ...
 
 
@@ -134,7 +135,8 @@ def run_bash_command_with_log_and_return_pid(
         log_path: str,
         env_vars: Optional[Dict[str, str]] = ...,
         stream_logs: bool = ...,
-        with_ray: bool = ...):
+        with_ray: bool = ...,
+        streaming_prefix: Optional[str] = ...) -> Dict[str, Union[int, str]]:
     ...
 
 

sky/skylet/services.py

@@ -197,12 +197,11 @@ class JobsServiceImpl(jobsv1_pb2_grpc.JobsServiceServicer):
                     f.write(request.codegen)
                 os.chmod(script_path, 0o755)
 
-            cd = f'cd {constants.SKY_REMOTE_WORKDIR}'
             job_submit_cmd = (
                 # JOB_CMD_IDENTIFIER is used for identifying the process
                 # retrieved with pid is the same driver process.
                 f'{job_lib.JOB_CMD_IDENTIFIER.format(job_id)} && '
-                f'{cd} && {constants.SKY_PYTHON_CMD} -u {script_path}'
+                f'{constants.SKY_PYTHON_CMD} -u {script_path}'
                 # Do not use &>, which is not POSIX and may not work.
                 # Note that the order of ">filename 2>&1" matters.
                 f' > {remote_log_path} 2>&1')
@@ -387,6 +386,21 @@ class JobsServiceImpl(jobsv1_pb2_grpc.JobsServiceServicer):
         except Exception as e:  # pylint: disable=broad-except
             context.abort(grpc.StatusCode.INTERNAL, str(e))
 
+    def GetJobExitCodes(  # type: ignore[return]
+            self, request: jobsv1_pb2.GetJobExitCodesRequest,
+            context: grpc.ServicerContext
+    ) -> jobsv1_pb2.GetJobExitCodesResponse:
+        try:
+            job_id = request.job_id if request.HasField(
+                'job_id') else job_lib.get_latest_job_id()
+            exit_codes: Optional[List[int]] = None
+            if job_id:
+                exit_codes_list = job_lib.get_exit_codes(job_id)
+                exit_codes = exit_codes_list if exit_codes_list else []
+            return jobsv1_pb2.GetJobExitCodesResponse(exit_codes=exit_codes)
+        except Exception as e:  # pylint: disable=broad-except
+            context.abort(grpc.StatusCode.INTERNAL, str(e))
+
 
 class ManagedJobsServiceImpl(managed_jobsv1_pb2_grpc.ManagedJobsServiceServicer
                             ):
@@ -488,7 +502,8 @@ class ManagedJobsServiceImpl(managed_jobsv1_pb2_grpc.ManagedJobsServiceServicer
                     entrypoint=job.get('entrypoint'),
                     metadata=converted_metadata,
                     pool=job.get('pool'),
-                    pool_hash=job.get('pool_hash'))
+                    pool_hash=job.get('pool_hash'),
+                    links=job.get('links'))
                 jobs_info.append(job_info)
 
             return managed_jobsv1_pb2.GetJobTableResponse(

sky/skylet/skylet.py

@@ -48,8 +48,12 @@ def start_grpc_server(port: int = constants.SKYLET_GRPC_PORT) -> grpc.Server:
     # putting it here for visibility.
     # TODO(kevin): Determine the optimal max number of threads.
     max_workers = min(32, (os.cpu_count() or 1) + 4)
+    # There's only a single skylet process per cluster, so disable
+    # SO_REUSEPORT to raise an error if the port is already in use.
+    options = (('grpc.so_reuseport', 0),)
     server = grpc.server(
-        concurrent.futures.ThreadPoolExecutor(max_workers=max_workers))
+        concurrent.futures.ThreadPoolExecutor(max_workers=max_workers),
+        options=options)
 
     autostopv1_pb2_grpc.add_AutostopServiceServicer_to_server(
         services.AutostopServiceImpl(), server)

sky/skylet/subprocess_daemon.py

@@ -110,7 +110,8 @@ def kill_process_tree(process: psutil.Process,
 
 
 def main():
-    # daemonize()
+    daemonize()
+
     parser = argparse.ArgumentParser()
     parser.add_argument('--parent-pid', type=int, required=True)
     parser.add_argument('--proc-pid', type=int, required=True)

sky/ssh_node_pools/constants.py

@@ -0,0 +1,12 @@
+"""Constants for SSH Node Pools"""
+# pylint: disable=line-too-long
+import os
+
+DEFAULT_KUBECONFIG_PATH = os.path.expanduser('~/.kube/config')
+SSH_CONFIG_PATH = os.path.expanduser('~/.ssh/config')
+NODE_POOLS_INFO_DIR = os.path.expanduser('~/.sky/ssh_node_pools_info')
+NODE_POOLS_KEY_DIR = os.path.expanduser('~/.sky/ssh_keys')
+DEFAULT_SSH_NODE_POOLS_PATH = os.path.expanduser('~/.sky/ssh_node_pools.yaml')
+
+# TODO (kyuds): make this configurable?
+K3S_TOKEN = 'mytoken'  # Any string can be used as the token

sky/ssh_node_pools/core.py

@@ -1,10 +1,15 @@
 """SSH Node Pool management core functionality."""
 import os
 from pathlib import Path
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Optional, Tuple
 
 import yaml
 
+from sky import clouds
+from sky.ssh_node_pools import constants
+from sky.ssh_node_pools import deploy
+from sky.usage import usage_lib
+from sky.utils import common_utils
 from sky.utils import yaml_utils
 
 
@@ -12,8 +17,8 @@ class SSHNodePoolManager:
     """Manager for SSH Node Pool configurations."""
 
     def __init__(self):
-        self.config_path = Path.home() / '.sky' / 'ssh_node_pools.yaml'
-        self.keys_dir = Path.home() / '.sky' / 'ssh_keys'
+        self.config_path = Path(constants.DEFAULT_SSH_NODE_POOLS_PATH)
+        self.keys_dir = Path(constants.NODE_POOLS_KEY_DIR)
         self.keys_dir.mkdir(parents=True, exist_ok=True)
 
     def get_all_pools(self) -> Dict[str, Any]:
@@ -133,3 +138,35 @@ def list_ssh_keys() -> List[str]:
     """List available SSH keys."""
     manager = SSHNodePoolManager()
     return manager.list_ssh_keys()
+
+
+@usage_lib.entrypoint
+def ssh_up(infra: Optional[str] = None, cleanup: bool = False) -> None:
+    """Deploys or tears down a Kubernetes cluster on SSH targets.
+
+    Args:
+        infra: Name of the cluster configuration in ssh_node_pools.yaml.
+            If None, the first cluster in the file is used.
+        cleanup: If True, clean up the cluster instead of deploying.
+    """
+    deploy.run(cleanup=cleanup, infra=infra)
+
+
+@usage_lib.entrypoint
+def ssh_status(context_name: str) -> Tuple[bool, str]:
+    """Check the status of an SSH Node Pool context.
+
+    Args:
+        context_name: The SSH context name (e.g., 'ssh-my-cluster')
+
+    Returns:
+        Tuple[bool, str]: (is_ready, reason)
+            - is_ready: True if the SSH Node Pool is ready, False otherwise
+            - reason: Explanation of the status
+    """
+    try:
+        is_ready, reason = clouds.SSH.check_single_context(context_name)
+        return is_ready, reason
+    except Exception as e:  # pylint: disable=broad-except
+        return False, ('Failed to check SSH context: '
+                       f'{common_utils.format_exception(e)}')

sky/ssh_node_pools/deploy/__init__.py

@@ -0,0 +1,4 @@
+"""Module for Deploying SSH Node Pools"""
+from sky.ssh_node_pools.deploy.deploy import run
+
+__all__ = ['run']