skypilot-nightly 1.0.0.dev20250609__py3-none-any.whl → 1.0.0.dev20250611__py3-none-any.whl

sky/server/stream_utils.py

@@ -3,7 +3,7 @@
 import asyncio
 import collections
 import pathlib
-from typing import AsyncGenerator, Deque, Optional
+from typing import AsyncGenerator, Deque, List, Optional
 
 import aiofiles
 import fastapi
@@ -15,6 +15,12 @@ from sky.utils import rich_utils
 
 logger = sky_logging.init_logger(__name__)
 
+# When streaming log lines, buffer the lines in memory and flush them in chunks
+# to improve log tailing throughput. Buffer size is the max size bytes of each
+# chunk and the timeout threshold for flushing the buffer to ensure
+# responsiveness.
+_BUFFER_SIZE = 8 * 1024  # 8KB
+_BUFFER_TIMEOUT = 0.02  # 20ms
 _HEARTBEAT_INTERVAL = 30
 
 
@@ -36,7 +42,16 @@ async def log_streamer(request_id: Optional[str],
                        plain_logs: bool = False,
                        tail: Optional[int] = None,
                        follow: bool = True) -> AsyncGenerator[str, None]:
-    """Streams the logs of a request."""
+    """Streams the logs of a request.
+
+    Args:
+        request_id: The request ID to check whether the log tailing process
+            should be stopped.
+        log_path: The path to the log file.
+        plain_logs: Whether to show plain logs.
+        tail: The number of lines to tail. If None, tail the whole file.
+        follow: Whether to follow the log file.
+    """
 
     if request_id is not None:
         status_msg = rich_utils.EncodedStatusMessage(
@@ -80,65 +95,106 @@ async def log_streamer(request_id: Optional[str],
         if show_request_waiting_spinner:
             yield status_msg.stop()
 
-    # Find last n lines of the log file. Do not read the whole file into memory.
     async with aiofiles.open(log_path, 'rb') as f:
-        if tail is not None:
-            # TODO(zhwu): this will include the control lines for rich status,
-            # which may not lead to exact tail lines when showing on the client
-            # side.
-            lines: Deque[str] = collections.deque(maxlen=tail)
-            async for line_str in _yield_log_file_with_payloads_skipped(f):
-                lines.append(line_str)
-            for line_str in lines:
-                yield line_str
-
-        last_heartbeat_time = asyncio.get_event_loop().time()
+        async for chunk in _tail_log_file(f, request_id, plain_logs, tail,
+                                          follow):
+            yield chunk
+
+
+async def _tail_log_file(f: aiofiles.threadpool.binary.AsyncBufferedReader,
+                         request_id: Optional[str] = None,
+                         plain_logs: bool = False,
+                         tail: Optional[int] = None,
+                         follow: bool = True) -> AsyncGenerator[str, None]:
+    """Tail the opened log file, buffer the lines and flush in chunks."""
+
+    if tail is not None:
+        # Find last n lines of the log file. Do not read the whole file into
+        # memory.
+        # TODO(zhwu): this will include the control lines for rich status,
+        # which may not lead to exact tail lines when showing on the client
+        # side.
+        lines: Deque[str] = collections.deque(maxlen=tail)
+        async for line_str in _yield_log_file_with_payloads_skipped(f):
+            lines.append(line_str)
+        for line_str in lines:
+            yield line_str
 
-        while True:
-            # Sleep 0 to yield control to allow other coroutines to run,
-            # while keeps the loop tight to make log stream responsive.
-            await asyncio.sleep(0)
-            line: Optional[bytes] = await f.readline()
-            if not line:
-                if request_id is not None:
-                    request_task = requests_lib.get_request(request_id)
-                    if request_task.status > requests_lib.RequestStatus.RUNNING:
-                        if (request_task.status ==
-                                requests_lib.RequestStatus.CANCELLED):
-                            yield (f'{request_task.name!r} request {request_id}'
-                                   ' cancelled\n')
-                        break
-                if not follow:
+    last_heartbeat_time = asyncio.get_event_loop().time()
+
+    # Buffer the lines in memory and flush them in chunks to improve log
+    # tailing throughput.
+    buffer: List[str] = []
+    buffer_bytes = 0
+    last_flush_time = asyncio.get_event_loop().time()
+
+    async def flush_buffer() -> AsyncGenerator[str, None]:
+        nonlocal buffer, buffer_bytes, last_flush_time
+        if buffer:
+            yield ''.join(buffer)
+            buffer.clear()
+            buffer_bytes = 0
+            last_flush_time = asyncio.get_event_loop().time()
+
+    while True:
+        # Sleep 0 to yield control to allow other coroutines to run,
+        # while keeps the loop tight to make log stream responsive.
+        await asyncio.sleep(0)
+        current_time = asyncio.get_event_loop().time()
+        # Flush the buffer when it is not empty and the buffer is full or the
+        # flush timeout is reached.
+        if buffer and (buffer_bytes >= _BUFFER_SIZE or
+                       (current_time - last_flush_time) >= _BUFFER_TIMEOUT):
+            async for chunk in flush_buffer():
+                yield chunk
+
+        line: Optional[bytes] = await f.readline()
+        if not line:
+            if request_id is not None:
+                request_task = requests_lib.get_request(request_id)
+                if request_task.status > requests_lib.RequestStatus.RUNNING:
+                    if (request_task.status ==
+                            requests_lib.RequestStatus.CANCELLED):
+                        buffer.append(
+                            f'{request_task.name!r} request {request_id}'
+                            ' cancelled\n')
                     break
+            if not follow:
+                break
+
+            if current_time - last_heartbeat_time >= _HEARTBEAT_INTERVAL:
+                # Currently just used to keep the connection busy, refer to
+                # https://github.com/skypilot-org/skypilot/issues/5750 for
+                # more details.
+                buffer.append(
+                    message_utils.encode_payload(
+                        rich_utils.Control.HEARTBEAT.encode('')))
+                last_heartbeat_time = current_time
+
+            # Sleep shortly to avoid storming the DB and CPU, this has
+            # little impact on the responsivness here since we are waiting
+            # for a new line to come in.
+            await asyncio.sleep(0.1)
+            continue
 
-                current_time = asyncio.get_event_loop().time()
-                if current_time - last_heartbeat_time >= _HEARTBEAT_INTERVAL:
-                    # Currently just used to keep the connection busy, refer to
-                    # https://github.com/skypilot-org/skypilot/issues/5750 for
-                    # more details.
-                    yield message_utils.encode_payload(
-                        rich_utils.Control.HEARTBEAT.encode(''))
-                    last_heartbeat_time = current_time
-
-                # Sleep shortly to avoid storming the DB and CPU, this has
-                # little impact on the responsivness here since we are waiting
-                # for a new line to come in.
-                await asyncio.sleep(0.1)
+        # Refresh the heartbeat time, this is a trivial optimization for
+        # performance but it helps avoid unnecessary heartbeat strings
+        # being printed when the client runs in an old version.
+        last_heartbeat_time = asyncio.get_event_loop().time()
+        line_str = line.decode('utf-8')
+        if plain_logs:
+            is_payload, line_str = message_utils.decode_payload(
+                line_str, raise_for_mismatch=False)
+            # TODO(aylei): implement heartbeat mechanism for plain logs,
+            # sending invisible characters might be okay.
+            if is_payload:
                 continue
+        buffer.append(line_str)
+        buffer_bytes += len(line_str.encode('utf-8'))
 
-            # Refresh the heartbeat time, this is a trivial optimization for
-            # performance but it helps avoid unnecessary heartbeat strings
-            # being printed when the client runs in an old version.
-            last_heartbeat_time = asyncio.get_event_loop().time()
-            line_str = line.decode('utf-8')
-            if plain_logs:
-                is_payload, line_str = message_utils.decode_payload(
-                    line_str, raise_for_mismatch=False)
-                # TODO(aylei): implement heartbeat mechanism for plain logs,
-                # sending invisible characters might be okay.
-                if is_payload:
-                    continue
-            yield line_str
+    # Flush remaining lines in the buffer.
+    async for chunk in flush_buffer():
+        yield chunk
 
 
 def stream_response(

sky/skylet/constants.py

@@ -377,8 +377,7 @@ OVERRIDEABLE_CONFIG_KEYS_IN_TASK: List[Tuple[str, ...]] = [
 ]
 # When overriding the SkyPilot configs on the API server with the client one,
 # we skip the following keys because they are meant to be client-side configs.
-SKIPPED_CLIENT_OVERRIDE_KEYS: List[Tuple[str, ...]] = [('admin_policy',),
-                                                       ('api_server',),
+SKIPPED_CLIENT_OVERRIDE_KEYS: List[Tuple[str, ...]] = [('api_server',),
                                                        ('allowed_clouds',),
                                                        ('workspaces',), ('db',)]
 

sky/skylet/job_lib.py

@@ -3,6 +3,7 @@
 This is a remote utility module that provides job queue functionality.
 """
 import enum
+import functools
 import getpass
 import json
 import os
@@ -10,6 +11,7 @@ import pathlib
 import shlex
 import signal
 import sqlite3
+import threading
 import time
 import typing
 from typing import Any, Dict, List, Optional, Sequence
@@ -119,9 +121,25 @@ def create_table(cursor, conn):
     conn.commit()
 
 
-_DB = db_utils.SQLiteConn(_DB_PATH, create_table)
-_CURSOR = _DB.cursor
-_CONN = _DB.conn
+_DB = None
+_db_init_lock = threading.Lock()
+
+
+def init_db(func):
+    """Initialize the database."""
+
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        global _DB
+        if _DB is not None:
+            return func(*args, **kwargs)
+
+        with _db_init_lock:
+            if _DB is None:
+                _DB = db_utils.SQLiteConn(_DB_PATH, create_table)
+        return func(*args, **kwargs)
+
+    return wrapper
 
 
 class JobStatus(enum.Enum):
@@ -210,30 +228,37 @@ _PRE_RESOURCE_STATUSES = [JobStatus.PENDING]
 class JobScheduler:
     """Base class for job scheduler"""
 
+    @init_db
     def queue(self, job_id: int, cmd: str) -> None:
-        _CURSOR.execute('INSERT INTO pending_jobs VALUES (?,?,?,?)',
-                        (job_id, cmd, 0, int(time.time())))
-        _CONN.commit()
+        assert _DB is not None
+        _DB.cursor.execute('INSERT INTO pending_jobs VALUES (?,?,?,?)',
+                           (job_id, cmd, 0, int(time.time())))
+        _DB.conn.commit()
         set_status(job_id, JobStatus.PENDING)
         self.schedule_step()
 
+    @init_db
     def remove_job_no_lock(self, job_id: int) -> None:
-        _CURSOR.execute(f'DELETE FROM pending_jobs WHERE job_id={job_id!r}')
-        _CONN.commit()
+        assert _DB is not None
+        _DB.cursor.execute(f'DELETE FROM pending_jobs WHERE job_id={job_id!r}')
+        _DB.conn.commit()
 
+    @init_db
     def _run_job(self, job_id: int, run_cmd: str):
-        _CURSOR.execute((f'UPDATE pending_jobs SET submit={int(time.time())} '
-                         f'WHERE job_id={job_id!r}'))
-        _CONN.commit()
+        assert _DB is not None
+        _DB.cursor.execute(
+            (f'UPDATE pending_jobs SET submit={int(time.time())} '
+             f'WHERE job_id={job_id!r}'))
+        _DB.conn.commit()
         pid = subprocess_utils.launch_new_process_tree(run_cmd)
         # TODO(zhwu): Backward compatibility, remove this check after 0.10.0.
         # This is for the case where the job is submitted with SkyPilot older
         # than #4318, using ray job submit.
         if 'job submit' in run_cmd:
             pid = -1
-        _CURSOR.execute((f'UPDATE jobs SET pid={pid} '
-                         f'WHERE job_id={job_id!r}'))
-        _CONN.commit()
+        _DB.cursor.execute((f'UPDATE jobs SET pid={pid} '
+                            f'WHERE job_id={job_id!r}'))
+        _DB.conn.commit()
 
     def schedule_step(self, force_update_jobs: bool = False) -> None:
         if force_update_jobs:
@@ -282,8 +307,10 @@ class JobScheduler:
 class FIFOScheduler(JobScheduler):
     """First in first out job scheduler"""
 
+    @init_db
     def _get_pending_job_ids(self) -> List[int]:
-        rows = _CURSOR.execute(
+        assert _DB is not None
+        rows = _DB.cursor.execute(
             'SELECT job_id FROM pending_jobs ORDER BY job_id').fetchall()
         return [row[0] for row in rows]
 
@@ -308,26 +335,30 @@ def make_job_command_with_user_switching(username: str,
     return ['sudo', '-H', 'su', '--login', username, '-c', command]
 
 
+@init_db
 def add_job(job_name: str, username: str, run_timestamp: str,
             resources_str: str) -> int:
     """Atomically reserve the next available job id for the user."""
+    assert _DB is not None
     job_submitted_at = time.time()
     # job_id will autoincrement with the null value
-    _CURSOR.execute(
+    _DB.cursor.execute(
         'INSERT INTO jobs VALUES (null, ?, ?, ?, ?, ?, ?, null, ?, 0)',
         (job_name, username, job_submitted_at, JobStatus.INIT.value,
          run_timestamp, None, resources_str))
-    _CONN.commit()
-    rows = _CURSOR.execute('SELECT job_id FROM jobs WHERE run_timestamp=(?)',
-                           (run_timestamp,))
+    _DB.conn.commit()
+    rows = _DB.cursor.execute('SELECT job_id FROM jobs WHERE run_timestamp=(?)',
+                              (run_timestamp,))
     for row in rows:
         job_id = row[0]
     assert job_id is not None
     return job_id
 
 
+@init_db
 def _set_status_no_lock(job_id: int, status: JobStatus) -> None:
     """Setting the status of the job in the database."""
+    assert _DB is not None
     assert status != JobStatus.RUNNING, (
         'Please use set_job_started() to set job status to RUNNING')
     if status.is_terminal():
@@ -339,15 +370,15 @@ def _set_status_no_lock(job_id: int, status: JobStatus) -> None:
         check_end_at_str = ' AND end_at IS NULL'
         if status != JobStatus.FAILED_SETUP:
             check_end_at_str = ''
-        _CURSOR.execute(
+        _DB.cursor.execute(
             'UPDATE jobs SET status=(?), end_at=(?) '
             f'WHERE job_id=(?) {check_end_at_str}',
             (status.value, end_at, job_id))
     else:
-        _CURSOR.execute(
+        _DB.cursor.execute(
             'UPDATE jobs SET status=(?), end_at=NULL '
             'WHERE job_id=(?)', (status.value, job_id))
-    _CONN.commit()
+    _DB.conn.commit()
 
 
 def set_status(job_id: int, status: JobStatus) -> None:
@@ -357,16 +388,19 @@ def set_status(job_id: int, status: JobStatus) -> None:
         _set_status_no_lock(job_id, status)
 
 
+@init_db
 def set_job_started(job_id: int) -> None:
     # TODO(mraheja): remove pylint disabling when filelock version updated.
     # pylint: disable=abstract-class-instantiated
+    assert _DB is not None
     with filelock.FileLock(_get_lock_path(job_id)):
-        _CURSOR.execute(
+        _DB.cursor.execute(
             'UPDATE jobs SET status=(?), start_at=(?), end_at=NULL '
             'WHERE job_id=(?)', (JobStatus.RUNNING.value, time.time(), job_id))
-        _CONN.commit()
+        _DB.conn.commit()
 
 
+@init_db
 def get_status_no_lock(job_id: int) -> Optional[JobStatus]:
     """Get the status of the job with the given id.
 
@@ -375,8 +409,9 @@ def get_status_no_lock(job_id: int) -> Optional[JobStatus]:
     the status in a while loop as in `log_lib._follow_job_logs`. Otherwise, use
     `get_status`.
     """
-    rows = _CURSOR.execute('SELECT status FROM jobs WHERE job_id=(?)',
-                           (job_id,))
+    assert _DB is not None
+    rows = _DB.cursor.execute('SELECT status FROM jobs WHERE job_id=(?)',
+                              (job_id,))
     for (status,) in rows:
         if status is None:
             return None
@@ -391,11 +426,13 @@ def get_status(job_id: int) -> Optional[JobStatus]:
         return get_status_no_lock(job_id)
 
 
+@init_db
 def get_statuses_payload(job_ids: List[Optional[int]]) -> str:
+    assert _DB is not None
     # Per-job lock is not required here, since the staled job status will not
     # affect the caller.
     query_str = ','.join(['?'] * len(job_ids))
-    rows = _CURSOR.execute(
+    rows = _DB.cursor.execute(
         f'SELECT job_id, status FROM jobs WHERE job_id IN ({query_str})',
         job_ids)
     statuses = {job_id: None for job_id in job_ids}
@@ -419,14 +456,17 @@ def load_statuses_payload(
     return statuses
 
 
+@init_db
 def get_latest_job_id() -> Optional[int]:
-    rows = _CURSOR.execute(
+    assert _DB is not None
+    rows = _DB.cursor.execute(
         'SELECT job_id FROM jobs ORDER BY job_id DESC LIMIT 1')
     for (job_id,) in rows:
         return job_id
     return None
 
 
+@init_db
 def get_job_submitted_or_ended_timestamp_payload(job_id: int,
                                                  get_ended_time: bool) -> str:
     """Get the job submitted/ended timestamp.
@@ -440,9 +480,10 @@ def get_job_submitted_or_ended_timestamp_payload(job_id: int,
     `format_job_queue()`), because the job may stay in PENDING if the cluster is
     busy.
     """
+    assert _DB is not None
     field = 'end_at' if get_ended_time else 'submitted_at'
-    rows = _CURSOR.execute(f'SELECT {field} FROM jobs WHERE job_id=(?)',
-                           (job_id,))
+    rows = _DB.cursor.execute(f'SELECT {field} FROM jobs WHERE job_id=(?)',
+                              (job_id,))
     for (timestamp,) in rows:
         return message_utils.encode_payload(timestamp)
     return message_utils.encode_payload(None)
@@ -496,10 +537,12 @@ def _get_records_from_rows(rows) -> List[Dict[str, Any]]:
     return records
 
 
+@init_db
 def _get_jobs(
         user_hash: Optional[str],
         status_list: Optional[List[JobStatus]] = None) -> List[Dict[str, Any]]:
     """Returns jobs with the given fields, sorted by job_id, descending."""
+    assert _DB is not None
     if status_list is None:
         status_list = list(JobStatus)
     status_str_list = [repr(status.value) for status in status_list]
@@ -509,14 +552,16 @@ def _get_jobs(
         # We use the old username field for compatibility.
         filter_str += ' AND username=(?)'
         params.append(user_hash)
-    rows = _CURSOR.execute(
+    rows = _DB.cursor.execute(
         f'SELECT * FROM jobs {filter_str} ORDER BY job_id DESC', params)
     records = _get_records_from_rows(rows)
     return records
 
 
+@init_db
 def _get_jobs_by_ids(job_ids: List[int]) -> List[Dict[str, Any]]:
-    rows = _CURSOR.execute(
+    assert _DB is not None
+    rows = _DB.cursor.execute(
         f"""\
         SELECT * FROM jobs
         WHERE job_id IN ({','.join(['?'] * len(job_ids))})
@@ -527,8 +572,10 @@ def _get_jobs_by_ids(job_ids: List[int]) -> List[Dict[str, Any]]:
     return records
 
 
+@init_db
 def _get_pending_job(job_id: int) -> Optional[Dict[str, Any]]:
-    rows = _CURSOR.execute(
+    assert _DB is not None
+    rows = _DB.cursor.execute(
         'SELECT created_time, submit, run_cmd FROM pending_jobs '
         f'WHERE job_id={job_id!r}')
     for row in rows:
@@ -698,16 +745,18 @@ def update_job_status(job_ids: List[int],
     return statuses
 
 
+@init_db
 def fail_all_jobs_in_progress() -> None:
+    assert _DB is not None
     in_progress_status = [
         status.value for status in JobStatus.nonterminal_statuses()
     ]
-    _CURSOR.execute(
+    _DB.cursor.execute(
         f"""\
         UPDATE jobs SET status=(?)
         WHERE status IN ({','.join(['?'] * len(in_progress_status))})
         """, (JobStatus.FAILED_DRIVER.value, *in_progress_status))
-    _CONN.commit()
+    _DB.conn.commit()
 
 
 def update_status() -> None:
@@ -720,12 +769,14 @@ def update_status() -> None:
     update_job_status(nonterminal_job_ids)
 
 
+@init_db
 def is_cluster_idle() -> bool:
     """Returns if the cluster is idle (no in-flight jobs)."""
+    assert _DB is not None
     in_progress_status = [
         status.value for status in JobStatus.nonterminal_statuses()
     ]
-    rows = _CURSOR.execute(
+    rows = _DB.cursor.execute(
         f"""\
         SELECT COUNT(*) FROM jobs
         WHERE status IN ({','.join(['?'] * len(in_progress_status))})
@@ -905,27 +956,31 @@ def cancel_jobs_encoded_results(jobs: Optional[List[int]],
     return message_utils.encode_payload(cancelled_ids)
 
 
+@init_db
 def get_run_timestamp(job_id: Optional[int]) -> Optional[str]:
     """Returns the relative path to the log file for a job."""
-    _CURSOR.execute(
+    assert _DB is not None
+    _DB.cursor.execute(
         """\
             SELECT * FROM jobs
             WHERE job_id=(?)""", (job_id,))
-    row = _CURSOR.fetchone()
+    row = _DB.cursor.fetchone()
     if row is None:
         return None
     run_timestamp = row[JobInfoLoc.RUN_TIMESTAMP.value]
     return run_timestamp
 
 
+@init_db
 def run_timestamp_with_globbing_payload(job_ids: List[Optional[str]]) -> str:
     """Returns the relative paths to the log files for job with globbing."""
+    assert _DB is not None
     query_str = ' OR '.join(['job_id GLOB (?)'] * len(job_ids))
-    _CURSOR.execute(
+    _DB.cursor.execute(
         f"""\
             SELECT * FROM jobs
             WHERE {query_str}""", job_ids)
-    rows = _CURSOR.fetchall()
+    rows = _DB.cursor.fetchall()
     run_timestamps = {}
     for row in rows:
         job_id = row[JobInfoLoc.JOB_ID.value]