skypilot-nightly 1.0.0.dev20251022__py3-none-any.whl → 1.0.0.dev20251024__py3-none-any.whl

sky/jobs/client/sdk.py

@@ -130,8 +130,11 @@ def queue(
     refresh: bool,
     skip_finished: bool = False,
     all_users: bool = False,
-    job_ids: Optional[List[int]] = None
-) -> server_common.RequestId[List[responses.ManagedJobRecord]]:
+    job_ids: Optional[List[int]] = None,
+    limit: Optional[int] = None,
+    fields: Optional[List[str]] = None,
+) -> server_common.RequestId[Union[List[responses.ManagedJobRecord], Tuple[
+        List[responses.ManagedJobRecord], int, Dict[str, int], int]]]:
     """Gets statuses of managed jobs.
 
     Please refer to sky.cli.job_queue for documentation.
@@ -141,6 +144,8 @@ def queue(
         skip_finished: Whether to skip finished jobs.
         all_users: Whether to show all users' jobs.
         job_ids: IDs of the managed jobs to show.
+        limit: Number of jobs to show.
+        fields: Fields to get for the managed jobs.
 
     Returns:
         The request ID of the queue request.
@@ -173,15 +178,29 @@ def queue(
           does not exist.
         RuntimeError: if failed to get the managed jobs with ssh.
     """
-    body = payloads.JobsQueueBody(
-        refresh=refresh,
-        skip_finished=skip_finished,
-        all_users=all_users,
-        job_ids=job_ids,
-    )
+    remote_api_version = versions.get_remote_api_version()
+    if remote_api_version and remote_api_version >= 18:
+        body = payloads.JobsQueueV2Body(
+            refresh=refresh,
+            skip_finished=skip_finished,
+            all_users=all_users,
+            job_ids=job_ids,
+            limit=limit,
+            fields=fields,
+        )
+        path = '/jobs/queue/v2'
+    else:
+        body = payloads.JobsQueueBody(
+            refresh=refresh,
+            skip_finished=skip_finished,
+            all_users=all_users,
+            job_ids=job_ids,
+        )
+        path = '/jobs/queue'
+
     response = server_common.make_authenticated_request(
         'POST',
-        '/jobs/queue',
+        path,
         json=json.loads(body.model_dump_json()),
         timeout=(5, None))
     return server_common.get_request_id(response=response)

sky/jobs/client/sdk_async.py

@@ -1,12 +1,13 @@
 """Async SDK functions for managed jobs."""
 import typing
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Dict, List, Optional, Tuple, Union
 
 from sky import backends
 from sky import sky_logging
 from sky.adaptors import common as adaptors_common
 from sky.client import sdk_async
 from sky.jobs.client import sdk
+from sky.schemas.api import responses
 from sky.skylet import constants
 from sky.usage import usage_lib
 from sky.utils import common_utils
@@ -50,12 +51,17 @@ async def queue(
     refresh: bool,
     skip_finished: bool = False,
     all_users: bool = False,
+    job_ids: Optional[List[int]] = None,
+    limit: Optional[int] = None,
+    fields: Optional[List[str]] = None,
     stream_logs: Optional[
         sdk_async.StreamConfig] = sdk_async.DEFAULT_STREAM_CONFIG
-) -> List[Dict[str, Any]]:
+) -> Union[List[responses.ManagedJobRecord], Tuple[
+        List[responses.ManagedJobRecord], int, Dict[str, int], int]]:
     """Async version of queue() that gets statuses of managed jobs."""
     request_id = await context_utils.to_thread(sdk.queue, refresh,
-                                               skip_finished, all_users)
+                                               skip_finished, all_users,
+                                               job_ids, limit, fields)
     if stream_logs is not None:
         return await sdk_async._stream_and_get(request_id, stream_logs)  # pylint: disable=protected-access
     else:

sky/jobs/server/core.py

@@ -337,6 +337,7 @@ def launch(
     def _submit_one(
         consolidation_mode_job_id: Optional[int] = None,
         job_rank: Optional[int] = None,
+        num_jobs: Optional[int] = None,
     ) -> Tuple[Optional[int], Optional[backends.ResourceHandle]]:
         rank_suffix = '' if job_rank is None else f'-{job_rank}'
         remote_original_user_yaml_path = (
@@ -359,6 +360,7 @@ def launch(
             for task_ in dag.tasks:
                 if job_rank is not None:
                     task_.update_envs({'SKYPILOT_JOB_RANK': str(job_rank)})
+                    task_.update_envs({'SKYPILOT_NUM_JOBS': str(num_jobs)})
 
             dag_utils.dump_chain_dag_to_yaml(dag, f.name)
 
@@ -475,7 +477,7 @@ def launch(
     for job_rank in range(num_jobs):
         job_id = (consolidation_mode_job_ids[job_rank]
                   if consolidation_mode_job_ids is not None else None)
-        jid, handle = _submit_one(job_id, job_rank)
+        jid, handle = _submit_one(job_id, job_rank, num_jobs=num_jobs)
         assert jid is not None, (job_id, handle)
         ids.append(jid)
         all_handle = handle

sky/jobs/utils.py

@@ -1710,29 +1710,37 @@ def _get_job_status_from_tasks(
 
 
 @typing.overload
-def format_job_table(tasks: List[Dict[str, Any]],
-                     show_all: bool,
-                     show_user: bool,
-                     return_rows: Literal[False] = False,
-                     max_jobs: Optional[int] = None) -> str:
+def format_job_table(
+    tasks: List[Dict[str, Any]],
+    show_all: bool,
+    show_user: bool,
+    return_rows: Literal[False] = False,
+    max_jobs: Optional[int] = None,
+    job_status_counts: Optional[Dict[str, int]] = None,
+) -> str:
     ...
 
 
 @typing.overload
-def format_job_table(tasks: List[Dict[str, Any]],
-                     show_all: bool,
-                     show_user: bool,
-                     return_rows: Literal[True],
-                     max_jobs: Optional[int] = None) -> List[List[str]]:
+def format_job_table(
+    tasks: List[Dict[str, Any]],
+    show_all: bool,
+    show_user: bool,
+    return_rows: Literal[True],
+    max_jobs: Optional[int] = None,
+    job_status_counts: Optional[Dict[str, int]] = None,
+) -> List[List[str]]:
     ...
 
 
 def format_job_table(
-        tasks: List[Dict[str, Any]],
-        show_all: bool,
-        show_user: bool,
-        return_rows: bool = False,
-        max_jobs: Optional[int] = None) -> Union[str, List[List[str]]]:
+    tasks: List[Dict[str, Any]],
+    show_all: bool,
+    show_user: bool,
+    return_rows: bool = False,
+    max_jobs: Optional[int] = None,
+    job_status_counts: Optional[Dict[str, int]] = None,
+) -> Union[str, List[List[str]]]:
     """Returns managed jobs as a formatted string.
 
     Args:
@@ -1741,6 +1749,7 @@ def format_job_table(
         max_jobs: The maximum number of jobs to show in the table.
         return_rows: If True, return the rows as a list of strings instead of
           all rows concatenated into a single string.
+        job_status_counts: The counts of each job status.
 
     Returns: A formatted string of managed jobs, if not `return_rows`; otherwise
       a list of "rows" (each of which is a list of str).
@@ -1762,12 +1771,8 @@ def format_job_table(
         # by the task_id.
         jobs[get_hash(task)].append(task)
 
-    status_counts: Dict[str, int] = collections.defaultdict(int)
     workspaces = set()
     for job_tasks in jobs.values():
-        managed_job_status = _get_job_status_from_tasks(job_tasks)[0]
-        if not managed_job_status.is_terminal():
-            status_counts[managed_job_status.value] += 1
         workspaces.add(job_tasks[0].get('workspace',
                                         constants.SKYPILOT_DEFAULT_WORKSPACE))
 
@@ -1810,9 +1815,15 @@ def format_job_table(
     job_table = log_utils.create_table(columns)
 
     status_counts: Dict[str, int] = collections.defaultdict(int)
-    for task in tasks:
-        if not task['status'].is_terminal():
-            status_counts[task['status'].value] += 1
+    if job_status_counts:
+        for status_value, count in job_status_counts.items():
+            status = managed_job_state.ManagedJobStatus(status_value)
+            if not status.is_terminal():
+                status_counts[status_value] = count
+    else:
+        for task in tasks:
+            if not task['status'].is_terminal():
+                status_counts[task['status'].value] += 1
 
     all_tasks = tasks
     if max_jobs is not None:

sky/server/auth/oauth2_proxy.py

@@ -126,13 +126,10 @@ class OAuth2ProxyMiddleware(starlette.middleware.base.BaseHTTPMiddleware):
 
     async def _authenticate(self, request: fastapi.Request, call_next,
                             session: aiohttp.ClientSession):
-        forwarded_headers = dict(request.headers)
+        forwarded_headers = {}
         auth_url = f'{self.proxy_base}/oauth2/auth'
         forwarded_headers['X-Forwarded-Uri'] = str(request.url).rstrip('/')
-        # Remove content-length and content-type headers and drop request body
-        # to reduce the auth overhead.
-        forwarded_headers.pop('content-length', None)
-        forwarded_headers.pop('content-type', None)
+        forwarded_headers['Host'] = request.url.hostname
         logger.debug(f'authenticate request: {auth_url}, '
                      f'headers: {forwarded_headers}')
 

sky/server/requests/requests.py

@@ -578,6 +578,26 @@ def reset_db_and_logs():
                  f'{server_common.API_SERVER_CLIENT_DIR.expanduser()}')
     shutil.rmtree(server_common.API_SERVER_CLIENT_DIR.expanduser(),
                   ignore_errors=True)
+    with _init_db_lock:
+        _init_db_within_lock()
+    assert _DB is not None
+    with _DB.conn:
+        cursor = _DB.conn.cursor()
+        cursor.execute('SELECT sqlite_version()')
+        row = cursor.fetchone()
+        if row is None:
+            raise RuntimeError('Failed to get SQLite version')
+        version_str = row[0]
+        version_parts = version_str.split('.')
+        assert len(version_parts) >= 2, \
+            f'Invalid version string: {version_str}'
+        major, minor = int(version_parts[0]), int(version_parts[1])
+        # SQLite 3.35.0+ supports RETURNING statements.
+        # 3.35.0 was released in March 2021.
+        if not ((major > 3) or (major == 3 and minor >= 35)):
+            raise RuntimeError(
+                f'SQLite version {version_str} is not supported. '
+                'Please upgrade to SQLite 3.35.0 or later.')
 
 
 def request_lock_path(request_id: str) -> str:
@@ -733,12 +753,25 @@ async def get_request_status_async(
 @metrics_lib.time_me_async
 @asyncio_utils.shield
 async def create_if_not_exists_async(request: Request) -> bool:
-    """Async version of create_if_not_exists."""
-    async with filelock.AsyncFileLock(request_lock_path(request.request_id)):
-        if await _get_request_no_lock_async(request.request_id) is not None:
-            return False
-        await _add_or_update_request_no_lock_async(request)
-        return True
+    """Create a request if it does not exist, otherwise do nothing.
+
+    Returns:
+        True if a new request is created, False if the request already exists.
+    """
+    assert _DB is not None
+    request_columns = ', '.join(REQUEST_COLUMNS)
+    values_str = ', '.join(['?'] * len(REQUEST_COLUMNS))
+    sql_statement = (
+        f'INSERT INTO {REQUEST_TABLE} '
+        f'({request_columns}) VALUES '
+        f'({values_str}) ON CONFLICT(request_id) DO NOTHING RETURNING ROWID')
+    request_row = request.to_row()
+    # Execute the SQL statement without getting the request lock.
+    # The request lock is used to prevent racing with cancellation codepath,
+    # but a request cannot be cancelled before it is created.
+    row = await _DB.execute_get_returning_value_async(sql_statement,
+                                                      request_row)
+    return True if row else False
 
 
 @dataclasses.dataclass

sky/server/requests/serializers/decoders.py

@@ -2,7 +2,7 @@
 import base64
 import pickle
 import typing
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 from sky import jobs as managed_jobs
 from sky import models
@@ -116,22 +116,35 @@ def decode_jobs_queue(return_value: List[dict],) -> List[Dict[str, Any]]:
 
 
 @register_decoders('jobs.queue_v2')
-def decode_jobs_queue_v2(return_value) -> List[responses.ManagedJobRecord]:
+def decode_jobs_queue_v2(
+    return_value
+) -> Union[Tuple[List[responses.ManagedJobRecord], int, Dict[str, int], int],
+           List[responses.ManagedJobRecord]]:
     """Decode jobs queue response.
 
-    Supports legacy list, or a dict {jobs, total}.
-    - Returns list[job]
+    Supports legacy list, or a dict {jobs, total, total_no_filter,
+    status_counts}.
+
+    - Returns either list[job] or tuple(list[job], total, status_counts,
+      total_no_filter)
     """
-    # Case 1: dict shape {jobs, total}
-    if isinstance(return_value, dict) and 'jobs' in return_value:
+    # Case 1: dict shape {jobs, total, total_no_filter, status_counts}
+    if isinstance(return_value, dict):
         jobs = return_value.get('jobs', [])
+        total = return_value.get('total', len(jobs))
+        total_no_filter = return_value.get('total_no_filter', total)
+        status_counts = return_value.get('status_counts', {})
+        for job in jobs:
+            job['status'] = managed_jobs.ManagedJobStatus(job['status'])
+        jobs = [responses.ManagedJobRecord(**job) for job in jobs]
+        return jobs, total, status_counts, total_no_filter
     else:
         # Case 2: legacy list
         jobs = return_value
-    for job in jobs:
-        job['status'] = managed_jobs.ManagedJobStatus(job['status'])
-    jobs = [responses.ManagedJobRecord(**job) for job in jobs]
-    return jobs
+        for job in jobs:
+            job['status'] = managed_jobs.ManagedJobStatus(job['status'])
+        jobs = [responses.ManagedJobRecord(**job) for job in jobs]
+        return jobs
 
 
 def _decode_serve_status(

sky/server/requests/serializers/encoders.py

@@ -148,12 +148,13 @@ def encode_jobs_queue_v2(
     else:
         jobs = jobs_or_tuple
         total = None
-    for job in jobs:
+    jobs_dict = [job.model_dump(by_alias=True) for job in jobs]
+    for job in jobs_dict:
         job['status'] = job['status'].value
     if total is None:
-        return [job.model_dump(by_alias=True) for job in jobs]
+        return jobs_dict
     return {
-        'jobs': [job.model_dump(by_alias=True) for job in jobs],
+        'jobs': jobs_dict,
         'total': total,
         'total_no_filter': total_no_filter,
         'status_counts': status_counts

sky/server/rest.py

@@ -256,6 +256,40 @@ def handle_server_unavailable(response: 'requests.Response') -> None:
         raise exceptions.ServerTemporarilyUnavailableError(error_msg)
 
 
+async def handle_server_unavailable_async(
+        response: 'aiohttp.ClientResponse') -> None:
+    """Async version: Handle 503 (Service Unavailable) error
+
+    The client get 503 error in the following cases:
+    1. The reverse proxy cannot find any ready backend endpoints to serve the
+       request, e.g. when there is and rolling-update.
+    2. The skypilot API server has temporary resource issue, e.g. when the
+       cucurrency of the handling process is exhausted.
+
+    We expect the caller (CLI or SDK) retry on these cases and show clear wait
+    message to the user to let user decide whether keep waiting or abort the
+    request.
+    """
+    if response.status != 503:
+        return
+
+    error_msg = ''
+    try:
+        response_data = await response.json()
+        if 'detail' in response_data:
+            error_msg = response_data['detail']
+    except Exception:  # pylint: disable=broad-except
+        try:
+            text = await response.text()
+            if text:
+                error_msg = text
+        except Exception:  # pylint: disable=broad-except
+            pass
+
+    with ux_utils.print_exception_no_traceback():
+        raise exceptions.ServerTemporarilyUnavailableError(error_msg)
+
+
 @_retry_on_server_unavailable()
 def request(method, url, **kwargs) -> 'requests.Response':
     """Send a request to the API server, retry on server temporarily
@@ -332,7 +366,7 @@ async def request_without_retry_async(session: 'aiohttp.ClientSession',
         response = await session.request(method, url, **kwargs)
 
         # Handle server unavailability (503 status) - same as sync version
-        handle_server_unavailable(response)
+        await handle_server_unavailable_async(response)
 
         # Set remote API version and version from headers - same as sync version
         remote_api_version = response.headers.get(constants.API_VERSION_HEADER)

sky/skylet/log_lib.py

@@ -220,7 +220,14 @@ def run_with_log(
                           stdin=stdin,
                           **kwargs) as proc:
         try:
-            subprocess_utils.kill_process_daemon(proc.pid)
+            if ctx is not None:
+                # When runs in coroutine, use kill_pg if available to avoid
+                # the overhead of refreshing the process tree in the daemon.
+                subprocess_utils.kill_process_daemon(proc.pid, use_kill_pg=True)
+            else:
+                # For backward compatibility, do not specify use_kill_pg by
+                # default.
+                subprocess_utils.kill_process_daemon(proc.pid)
             stdout = ''
             stderr = ''
             stdout_stream_handler = None

sky/skylet/subprocess_daemon.py

@@ -4,11 +4,16 @@ processes of proc_pid.
 """
 import argparse
 import os
+import signal
 import sys
 import time
+from typing import List, Optional
 
 import psutil
 
+# Environment variable to enable kill_pg in subprocess daemon.
+USE_KILL_PG_ENV_VAR = 'SKYPILOT_SUBPROCESS_DAEMON_KILL_PG'
+
 
 def daemonize():
     """Detaches the process from its parent process with double-forking.
@@ -38,8 +43,74 @@ def daemonize():
     # This process is now fully detached from the original parent and terminal
 
 
-if __name__ == '__main__':
-    daemonize()
+def get_pgid_if_leader(pid) -> Optional[int]:
+    """Get the process group ID of the target process if it is the leader."""
+    try:
+        pgid = os.getpgid(pid)
+        # Only use process group if the target process is the leader. This is
+        # to avoid killing the entire process group while the target process is
+        # just a subprocess in the group.
+        if pgid == pid:
+            print(f'Process group {pgid} is the leader.')
+            return pgid
+        return None
+    except Exception:  # pylint: disable=broad-except
+        # Process group is only available in UNIX.
+        return None
+
+
+def kill_process_group(pgid: int) -> bool:
+    """Kill the target process group."""
+    try:
+        print(f'Terminating process group {pgid}...')
+        os.killpg(pgid, signal.SIGTERM)
+    except Exception:  # pylint: disable=broad-except
+        return False
+
+    # Wait 30s for the process group to exit gracefully.
+    time.sleep(30)
+
+    try:
+        print(f'Force killing process group {pgid}...')
+        os.killpg(pgid, signal.SIGKILL)
+    except Exception:  # pylint: disable=broad-except
+        pass
+
+    return True
+
+
+def kill_process_tree(process: psutil.Process,
+                      children: List[psutil.Process]) -> bool:
+    """Kill the process tree of the target process."""
+    if process is not None:
+        # Kill the target process first to avoid having more children, or fail
+        # the process due to the children being defunct.
+        children = [process] + children
+
+    if not children:
+        sys.exit()
+
+    for child in children:
+        try:
+            child.terminate()
+        except psutil.NoSuchProcess:
+            continue
+
+    # Wait 30s for the processes to exit gracefully.
+    time.sleep(30)
+
+    # SIGKILL if they're still running.
+    for child in children:
+        try:
+            child.kill()
+        except psutil.NoSuchProcess:
+            continue
+
+    return True
+
+
+def main():
+    # daemonize()
     parser = argparse.ArgumentParser()
     parser.add_argument('--parent-pid', type=int, required=True)
     parser.add_argument('--proc-pid', type=int, required=True)
@@ -72,37 +143,40 @@ if __name__ == '__main__':
             except (psutil.NoSuchProcess, ValueError):
                 pass
 
+    pgid: Optional[int] = None
+    if os.environ.get(USE_KILL_PG_ENV_VAR) == '1':
+        # Use kill_pg on UNIX system if allowed to reduce the resource usage.
+        # Note that both implementations might leave subprocessed uncancelled:
+        # - kill_process_tree(default): a subprocess is able to detach itself
+        #   from the process tree use the same technique as daemonize(). Also,
+        #   since we refresh the process tree per second, if the subprocess is
+        #   launched between the [last_poll, parent_die] interval, the
+        #   subprocess will not be captured will not be killed.
+        # - kill_process_group: kill_pg will kill all the processed in the group
+        #   but if a subprocess calls setpgid(0, 0) to detach itself from the
+        #   process group (usually to daemonize itself), the subprocess will
+        #   not be killed.
+        pgid = get_pgid_if_leader(process.pid)
+
     if process is not None and parent_process is not None:
         # Wait for either parent or target process to exit
         while process.is_running() and parent_process.is_running():
-            try:
-                tmp_children = process.children(recursive=True)
-                if tmp_children:
-                    children = tmp_children
-            except psutil.NoSuchProcess:
-                pass
+            if pgid is None:
+                # Refresh process tree for cleanup if process group is not
+                # available.
+                try:
+                    tmp_children = process.children(recursive=True)
+                    if tmp_children:
+                        children = tmp_children
+                except psutil.NoSuchProcess:
+                    pass
             time.sleep(1)
 
-    if process is not None:
-        # Kill the target process first to avoid having more children, or fail
-        # the process due to the children being defunct.
-        children = [process] + children
+    if pgid is not None:
+        kill_process_group(pgid)
+    else:
+        kill_process_tree(process, children)
 
-    if not children:
-        sys.exit()
 
-    for child in children:
-        try:
-            child.terminate()
-        except psutil.NoSuchProcess:
-            continue
-
-    # Wait 30s for the processes to exit gracefully.
-    time.sleep(30)
-
-    # SIGKILL if they're still running.
-    for child in children:
-        try:
-            child.kill()
-        except psutil.NoSuchProcess:
-            continue
+if __name__ == '__main__':
+    main()

sky/utils/db/db_utils.py

@@ -358,6 +358,27 @@ class SQLiteConn(threading.local):
         conn = await self._get_async_conn()
         return await conn.execute_fetchall(sql, parameters)
 
+    async def execute_get_returning_value_async(
+            self,
+            sql: str,
+            parameters: Optional[Iterable[Any]] = None
+    ) -> Optional[sqlite3.Row]:
+        conn = await self._get_async_conn()
+
+        if parameters is None:
+            parameters = []
+
+        def exec_and_get_returning_value(sql: str,
+                                         parameters: Optional[Iterable[Any]]):
+            # pylint: disable=protected-access
+            row = conn._conn.execute(sql, parameters).fetchone()
+            conn._conn.commit()
+            return row
+
+        # pylint: disable=protected-access
+        return await conn._execute(exec_and_get_returning_value, sql,
+                                   parameters)
+
     async def close(self):
         if self._async_conn is not None:
             await self._async_conn.close()

sky/utils/subprocess_utils.py

@@ -19,6 +19,7 @@ from sky import exceptions
 from sky import sky_logging
 from sky.adaptors import common as adaptors_common
 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
@@ -306,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:
@@ -337,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
@@ -348,6 +359,7 @@ def kill_process_daemon(process_pid: int) -> None:
         stderr=subprocess.DEVNULL,
         # Disable input
         stdin=subprocess.DEVNULL,
+        env=env,
     )