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

sky/server/stream_utils.py

@@ -8,10 +8,13 @@ from typing import AsyncGenerator, Deque, List, Optional
 import aiofiles
 import fastapi
 
+from sky import global_user_state
 from sky import sky_logging
 from sky.server.requests import requests as requests_lib
+from sky.utils import common_utils
 from sky.utils import message_utils
 from sky.utils import rich_utils
+from sky.utils import status_lib
 
 logger = sky_logging.init_logger(__name__)
 
@@ -22,6 +25,14 @@ logger = sky_logging.init_logger(__name__)
 _BUFFER_SIZE = 8 * 1024  # 8KB
 _BUFFER_TIMEOUT = 0.02  # 20ms
 _HEARTBEAT_INTERVAL = 30
+_READ_CHUNK_SIZE = 256 * 1024  # 256KB chunks for file reading
+
+# If a SHORT request has been stuck in pending for
+# _SHORT_REQUEST_SPINNER_TIMEOUT seconds, we show the waiting spinner
+_SHORT_REQUEST_SPINNER_TIMEOUT = 2
+
+LONG_REQUEST_POLL_INTERVAL = 1
+DEFAULT_POLL_INTERVAL = 0.1
 
 
 async def _yield_log_file_with_payloads_skipped(
@@ -37,34 +48,51 @@ async def _yield_log_file_with_payloads_skipped(
         yield line_str
 
 
-async def log_streamer(request_id: Optional[str],
-                       log_path: pathlib.Path,
-                       plain_logs: bool = False,
-                       tail: Optional[int] = None,
-                       follow: bool = True) -> AsyncGenerator[str, None]:
+async def log_streamer(
+    request_id: Optional[str],
+    log_path: Optional[pathlib.Path] = None,
+    plain_logs: bool = False,
+    tail: Optional[int] = None,
+    follow: bool = True,
+    cluster_name: Optional[str] = None,
+    polling_interval: float = DEFAULT_POLL_INTERVAL
+) -> AsyncGenerator[str, None]:
     """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.
+        log_path: The path to the log file or directory containing the log
+        files. If it is a directory, all *.log files in the directory will be
+        streamed.
         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.
+        cluster_name: The cluster name to check status for provision logs.
+            If provided and cluster status is UP, streaming will terminate.
     """
 
     if request_id is not None:
+        start_time = asyncio.get_event_loop().time()
         status_msg = rich_utils.EncodedStatusMessage(
             f'[dim]Checking request: {request_id}[/dim]')
-        request_task = await requests_lib.get_request_async(request_id)
+        request_task = await requests_lib.get_request_async(request_id,
+                                                            fields=[
+                                                                'request_id',
+                                                                'name',
+                                                                'schedule_type',
+                                                                'status',
+                                                                'status_msg'
+                                                            ])
 
         if request_task is None:
             raise fastapi.HTTPException(
                 status_code=404, detail=f'Request {request_id} not found')
         request_id = request_task.request_id
 
-        # Do not show the waiting spinner if the request is a fast, non-blocking
-        # request.
+        # By default, do not show the waiting spinner for SHORT requests.
+        # If the request has been stuck in pending for
+        # _SHORT_REQUEST_SPINNER_TIMEOUT seconds, we show the waiting spinner
         show_request_waiting_spinner = (not plain_logs and
                                         request_task.schedule_type
                                         == requests_lib.ScheduleType.LONG)
@@ -77,9 +105,23 @@ async def log_streamer(request_id: Optional[str],
                        f'scheduled: {request_id}')
         req_status = request_task.status
         req_msg = request_task.status_msg
+        del request_task
+        # Slowly back off the database polling up to every 1 second, to avoid
+        # overloading the CPU and DB.
+        backoff = common_utils.Backoff(initial_backoff=polling_interval,
+                                       max_backoff_factor=10,
+                                       multiplier=1.2)
         while req_status < requests_lib.RequestStatus.RUNNING:
+            current_time = asyncio.get_event_loop().time()
+            # Show the waiting spinner for a SHORT request if it has been stuck
+            # in pending for _SHORT_REQUEST_SPINNER_TIMEOUT seconds
+            if not show_request_waiting_spinner and (
+                    current_time - start_time > _SHORT_REQUEST_SPINNER_TIMEOUT):
+                show_request_waiting_spinner = True
+                yield status_msg.init()
+                yield status_msg.start()
             if req_msg is not None:
-                waiting_msg = request_task.status_msg
+                waiting_msg = req_msg
             if show_request_waiting_spinner:
                 yield status_msg.update(f'[dim]{waiting_msg}[/dim]')
             elif plain_logs and waiting_msg != last_waiting_msg:
@@ -92,7 +134,7 @@ async def log_streamer(request_id: Optional[str],
             # TODO(aylei): we should use a better mechanism to avoid busy
             # polling the DB, which can be a bottleneck for high-concurrency
             # requests.
-            await asyncio.sleep(0.1)
+            await asyncio.sleep(backoff.current_backoff())
             status_with_msg = await requests_lib.get_request_status_async(
                 request_id, include_msg=True)
             req_status = status_with_msg.status
@@ -102,17 +144,42 @@ async def log_streamer(request_id: Optional[str],
         if show_request_waiting_spinner:
             yield status_msg.stop()
 
-    async with aiofiles.open(log_path, 'rb') as f:
-        async for chunk in _tail_log_file(f, request_id, plain_logs, tail,
-                                          follow):
-            yield chunk
+    # worker node provision logs
+    if log_path is not None and log_path.is_dir():
+        # Get all *.log files in the log_path dir
+        log_files = sorted(log_path.glob('*.log'))
+
+        for log_file_path in log_files:
+            # Add header before each file (similar to tail -f behavior)
+            header = f'\n==> {log_file_path} <==\n\n'
+            yield header
+
+            async with aiofiles.open(log_file_path, 'rb') as f:
+                async for chunk in _tail_log_file(f, request_id, plain_logs,
+                                                  tail, follow, cluster_name,
+                                                  polling_interval):
+                    yield chunk
+
+    # api server request logs (if request_id is provided) or
+    # head node provision logs (if cluster_name is provided)
+    else:
+        assert log_path is not None, (request_id, cluster_name)
+        async with aiofiles.open(log_path, 'rb') as f:
+            async for chunk in _tail_log_file(f, request_id, plain_logs, tail,
+                                              follow, cluster_name,
+                                              polling_interval):
+                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]:
+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,
+    cluster_name: Optional[str] = None,
+    polling_interval: float = DEFAULT_POLL_INTERVAL
+) -> AsyncGenerator[str, None]:
     """Tail the opened log file, buffer the lines and flush in chunks."""
 
     if tail is not None:
@@ -128,6 +195,7 @@ async def _tail_log_file(f: aiofiles.threadpool.binary.AsyncBufferedReader,
             yield line_str
 
     last_heartbeat_time = asyncio.get_event_loop().time()
+    last_status_check_time = asyncio.get_event_loop().time()
 
     # Buffer the lines in memory and flush them in chunks to improve log
     # tailing throughput.
@@ -135,6 +203,9 @@ async def _tail_log_file(f: aiofiles.threadpool.binary.AsyncBufferedReader,
     buffer_bytes = 0
     last_flush_time = asyncio.get_event_loop().time()
 
+    # Read file in chunks instead of line-by-line for better performance
+    incomplete_line = b''  # Buffer for incomplete lines across chunks
+
     async def flush_buffer() -> AsyncGenerator[str, None]:
         nonlocal buffer, buffer_bytes, last_flush_time
         if buffer:
@@ -155,16 +226,41 @@ async def _tail_log_file(f: aiofiles.threadpool.binary.AsyncBufferedReader,
             async for chunk in flush_buffer():
                 yield chunk
 
-        line: Optional[bytes] = await f.readline()
-        if not line:
-            if request_id is not None:
+        # Read file in chunks for better I/O performance
+        file_chunk: bytes = await f.read(_READ_CHUNK_SIZE)
+        if not file_chunk:
+            # Process any remaining incomplete line
+            if incomplete_line:
+                line_str = incomplete_line.decode('utf-8')
+                if plain_logs:
+                    is_payload, line_str = message_utils.decode_payload(
+                        line_str, raise_for_mismatch=False)
+                    if not is_payload:
+                        buffer.append(line_str)
+                        buffer_bytes += len(line_str.encode('utf-8'))
+                else:
+                    buffer.append(line_str)
+                    buffer_bytes += len(line_str.encode('utf-8'))
+                incomplete_line = b''
+
+            # Avoid checking the status too frequently to avoid overloading the
+            # DB.
+            should_check_status = (current_time -
+                                   last_status_check_time) >= polling_interval
+            if not follow:
+                # We will only hit this path once, but we should make sure to
+                # check the status so that we display the final request status
+                # if the request is complete.
+                should_check_status = True
+            if request_id is not None and should_check_status:
+                last_status_check_time = current_time
                 req_status = await requests_lib.get_request_status_async(
                     request_id)
                 if req_status.status > requests_lib.RequestStatus.RUNNING:
                     if (req_status.status ==
                             requests_lib.RequestStatus.CANCELLED):
                         request_task = await requests_lib.get_request_async(
-                            request_id)
+                            request_id, fields=['name', 'should_retry'])
                         if request_task.should_retry:
                             buffer.append(
                                 message_utils.encode_payload(
@@ -173,10 +269,44 @@ async def _tail_log_file(f: aiofiles.threadpool.binary.AsyncBufferedReader,
                             buffer.append(
                                 f'{request_task.name!r} request {request_id}'
                                 ' cancelled\n')
+                        del request_task
                     break
             if not follow:
+                # The below checks (cluster status, heartbeat) are not needed
+                # for non-follow logs.
                 break
-
+            # Provision logs pass in cluster_name, check cluster status
+            # periodically to see if provisioning is done.
+            if cluster_name is not None:
+                if should_check_status:
+                    last_status_check_time = current_time
+                    cluster_status = await (
+                        global_user_state.get_status_from_cluster_name_async(
+                            cluster_name))
+                    if cluster_status is None:
+                        logger.debug(
+                            'Stop tailing provision logs for cluster'
+                            f' status for cluster {cluster_name} not found')
+                        break
+                    # if the cluster is not in INIT state (UP or STOPPED),
+                    # stop tailing provision logs
+                    if cluster_status != status_lib.ClusterStatus.INIT:
+                        logger.debug(
+                            f'Stop tailing provision logs for cluster'
+                            f' {cluster_name} has status {cluster_status} '
+                            '(not in INIT state)')
+                        break
+                    req_filter = requests_lib.RequestTaskFilter(
+                        status=[requests_lib.RequestStatus.RUNNING],
+                        cluster_names=[cluster_name],
+                        include_request_names=['sky.launch'],
+                        fields=['cluster_name'])
+                    req_tasks = await requests_lib.get_request_tasks_async(
+                        req_filter)
+                    # if the cluster is in INIT state and there is no ongoing
+                    # launch request, stop tailing provision logs
+                    if len(req_tasks) == 0:
+                        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
@@ -196,38 +326,82 @@ async def _tail_log_file(f: aiofiles.threadpool.binary.AsyncBufferedReader,
         # 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'))
+
+        # Combine with any incomplete line from previous chunk
+        file_chunk = incomplete_line + file_chunk
+        incomplete_line = b''
+
+        # Split chunk into lines, preserving line structure
+        lines_bytes = file_chunk.split(b'\n')
+
+        # If chunk doesn't end with newline, the last element is incomplete
+        if file_chunk and not file_chunk.endswith(b'\n'):
+            incomplete_line = lines_bytes[-1]
+            lines_bytes = lines_bytes[:-1]
+        else:
+            # If ends with \n, split creates an empty last element we should
+            # ignore
+            if lines_bytes and lines_bytes[-1] == b'':
+                lines_bytes = lines_bytes[:-1]
+
+        # Process all complete lines in this chunk
+        for line_bytes in lines_bytes:
+            # Reconstruct line with newline (since split removed it)
+            line_str = line_bytes.decode('utf-8') + '\n'
+
+            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'))
 
     # Flush remaining lines in the buffer.
     async for chunk in flush_buffer():
         yield chunk
 
 
+def stream_response_for_long_request(
+    request_id: str,
+    logs_path: pathlib.Path,
+    background_tasks: fastapi.BackgroundTasks,
+    kill_request_on_disconnect: bool = True,
+) -> fastapi.responses.StreamingResponse:
+    """Stream the logs of a long request."""
+    return stream_response(
+        request_id,
+        logs_path,
+        background_tasks,
+        polling_interval=LONG_REQUEST_POLL_INTERVAL,
+        kill_request_on_disconnect=kill_request_on_disconnect,
+    )
+
+
 def stream_response(
-    request_id: str, logs_path: pathlib.Path,
-    background_tasks: fastapi.BackgroundTasks
+    request_id: str,
+    logs_path: pathlib.Path,
+    background_tasks: fastapi.BackgroundTasks,
+    polling_interval: float = DEFAULT_POLL_INTERVAL,
+    kill_request_on_disconnect: bool = True,
 ) -> fastapi.responses.StreamingResponse:
 
-    async def on_disconnect():
-        logger.info(f'User terminated the connection for request '
-                    f'{request_id}')
-        requests_lib.kill_requests([request_id])
+    if kill_request_on_disconnect:
+
+        async def on_disconnect():
+            logger.info(f'User terminated the connection for request '
+                        f'{request_id}')
+            await requests_lib.kill_request_async(request_id)
 
-    # The background task will be run after returning a response.
-    # https://fastapi.tiangolo.com/tutorial/background-tasks/
-    background_tasks.add_task(on_disconnect)
+        # The background task will be run after returning a response.
+        # https://fastapi.tiangolo.com/tutorial/background-tasks/
+        background_tasks.add_task(on_disconnect)
 
     return fastapi.responses.StreamingResponse(
-        log_streamer(request_id, logs_path),
+        log_streamer(request_id, logs_path, polling_interval=polling_interval),
         media_type='text/plain',
         headers={
             'Cache-Control': 'no-cache, no-transform',

sky/server/uvicorn.py

@@ -19,6 +19,7 @@ from uvicorn.supervisors import multiprocess
 
 from sky import sky_logging
 from sky.server import daemons
+from sky.server import metrics as metrics_lib
 from sky.server import state
 from sky.server.requests import requests as requests_lib
 from sky.skylet import constants
@@ -45,11 +46,11 @@ except ValueError:
 
 # TODO(aylei): use decorator to register requests that need to be proactively
 # cancelled instead of hardcoding here.
-_RETRIABLE_REQUEST_NAMES = [
+_RETRIABLE_REQUEST_NAMES = {
     'sky.logs',
     'sky.jobs.logs',
     'sky.serve.logs',
-]
+}
 
 
 def add_timestamp_prefix_for_server_logs() -> None:
@@ -150,37 +151,38 @@ class Server(uvicorn.Server):
                 requests_lib.RequestStatus.PENDING,
                 requests_lib.RequestStatus.RUNNING,
             ]
-            reqs = requests_lib.get_request_tasks(
-                req_filter=requests_lib.RequestTaskFilter(status=statuses))
-            if not reqs:
+            requests = [(request_task.request_id, request_task.name)
+                        for request_task in requests_lib.get_request_tasks(
+                            req_filter=requests_lib.RequestTaskFilter(
+                                status=statuses, fields=['request_id', 'name']))
+                       ]
+            if not requests:
                 break
-            logger.info(f'{len(reqs)} on-going requests '
+            logger.info(f'{len(requests)} on-going requests '
                         'found, waiting for them to finish...')
             # Proactively cancel internal requests and logs requests since
             # they can run for infinite time.
-            internal_request_ids = [
+            internal_request_ids = {
                 d.id for d in daemons.INTERNAL_REQUEST_DAEMONS
-            ]
+            }
             if time.time() - start_time > _WAIT_REQUESTS_TIMEOUT_SECONDS:
                 logger.warning('Timeout waiting for on-going requests to '
                                'finish, cancelling all on-going requests.')
-                for req in reqs:
-                    self.interrupt_request_for_retry(req.request_id)
+                for request_id, _ in requests:
+                    self.interrupt_request_for_retry(request_id)
                 break
             interrupted = 0
-            for req in reqs:
-                if req.request_id in internal_request_ids:
-                    self.interrupt_request_for_retry(req.request_id)
-                    interrupted += 1
-                elif req.name in _RETRIABLE_REQUEST_NAMES:
-                    self.interrupt_request_for_retry(req.request_id)
+            for request_id, name in requests:
+                if (name in _RETRIABLE_REQUEST_NAMES or
+                        request_id in internal_request_ids):
+                    self.interrupt_request_for_retry(request_id)
                     interrupted += 1
                 # TODO(aylei): interrupt pending requests to accelerate the
                 # shutdown.
             # If some requests are not interrupted, wait for them to finish,
             # otherwise we just check again immediately to accelerate the
             # shutdown process.
-            if interrupted < len(reqs):
+            if interrupted < len(requests):
                 time.sleep(_WAIT_REQUESTS_INTERVAL_SECONDS)
 
     def interrupt_request_for_retry(self, request_id: str) -> None:
@@ -212,8 +214,17 @@ class Server(uvicorn.Server):
             # Same as set PYTHONASYNCIODEBUG=1, but with custom threshold.
             event_loop.set_debug(True)
             event_loop.slow_callback_duration = lag_threshold
-        with self.capture_signals():
-            asyncio.run(self.serve(*args, **kwargs))
+        stop_monitor = threading.Event()
+        monitor = threading.Thread(target=metrics_lib.process_monitor,
+                                   args=('server', stop_monitor),
+                                   daemon=True)
+        monitor.start()
+        try:
+            with self.capture_signals():
+                asyncio.run(self.serve(*args, **kwargs))
+        finally:
+            stop_monitor.set()
+            monitor.join()
 
 
 def run(config: uvicorn.Config, max_db_connections: Optional[int] = None):

sky/setup_files/MANIFEST.in

@@ -1,5 +1,5 @@
 include sky/backends/monkey_patches/*.py
-exclude sky/clouds/service_catalog/data_fetchers/analyze.py
+exclude sky/catalog/data_fetchers/analyze.py
 include sky/provision/kubernetes/manifests/*
 include sky/provision/azure/*
 include sky/setup_files/*
@@ -21,3 +21,8 @@ include sky/users/*.conf
 include sky/utils/*.sh
 include sky/setup_files/alembic.ini
 recursive-include sky/schemas/db *
+
+# SkyPilot templates package
+recursive-include sky_templates/ray *
+recursive-include sky_templates *.py
+include sky_templates/README.md

sky/setup_files/alembic.ini

@@ -98,6 +98,14 @@ version_table = alembic_version_spot_jobs_db
 version_locations = %(here)s/../schemas/db/serve_state
 version_table = alembic_version_serve_state_db
 
+[sky_config_db]
+version_locations = %(here)s/../schemas/db/skypilot_config
+version_table = alembic_version_sky_config_db
+
+[kv_cache_db]
+version_locations = %(here)s/../schemas/db/kv_cache
+version_table = alembic_version_kv_cache_db
+
 [post_write_hooks]
 # post_write_hooks defines scripts or Python functions that are run
 # on newly generated revision scripts.  See the documentation for further