[API server] handle logs request in event loop

Signed-off-by: Aylei <[email protected]>

Author: aylei

sky/backends/backend_utils.py

@@ -40,6 +40,7 @@
 from sky.utils import command_runner
 from sky.utils import common
 from sky.utils import common_utils
+from sky.utils import context
 from sky.utils import controller_utils
 from sky.utils import env_options
 from sky.utils import registry
@@ -2185,6 +2186,7 @@ def refresh_cluster_record(
 
 
 @timeline.event
+@context.cancellation_guard
 def refresh_cluster_status_handle(
     cluster_name: str,
     *,
@@ -2234,6 +2236,7 @@ def check_cluster_available(
     ...
 
 
+@context.cancellation_guard
 def check_cluster_available(
     cluster_name: str,
     *,

sky/backends/cloud_vm_ray_backend.py

@@ -55,7 +55,7 @@
 from sky.skylet import job_lib
 from sky.skylet import log_lib
 from sky.usage import usage_lib
-from sky.utils import accelerator_registry
+from sky.utils import accelerator_registry, context
 from sky.utils import annotations
 from sky.utils import cluster_utils
 from sky.utils import command_runner
@@ -2410,6 +2410,7 @@ def is_provided_ips_valid(
             internal_external_ips[1:], key=lambda x: x[1])
         self.stable_internal_external_ips = stable_internal_external_ips
 
+    @context.cancellation_guard
     @annotations.lru_cache(scope='global')
     @timeline.event
     def get_command_runners(self,
@@ -3817,6 +3818,7 @@ def _rsync_down(args) -> None:
         subprocess_utils.run_in_parallel(_rsync_down, parallel_args)
         return dict(zip(job_ids, local_log_dirs))
 
+    @context.cancellation_guard
     def tail_logs(self,
                   handle: CloudVmRayResourceHandle,
                   job_id: Optional[int],
@@ -4534,6 +4536,7 @@ def is_definitely_autostopping(self,
     # TODO(zhwu): Refactor this to a CommandRunner class, so different backends
     # can support its own command runner.
     @timeline.event
+    @context.cancellation_guard
     def run_on_head(
         self,
         handle: CloudVmRayResourceHandle,

sky/client/sdk.py

@@ -83,6 +83,9 @@ def stream_response(request_id: Optional[str],
         logger.debug(f'To stream request logs: sky api logs {request_id}')
         raise
 
+def decode_stream(response: requests.Response,
+                  output_stream: Optional['io.TextIOBase'] = None) -> str:
+
 
 @usage_lib.entrypoint
 @server_common.check_server_healthy_or_start

sky/core.py

@@ -31,6 +31,7 @@
 from sky.utils import admin_policy_utils
 from sky.utils import common
 from sky.utils import common_utils
+from sky.utils import context
 from sky.utils import controller_utils
 from sky.utils import rich_utils
 from sky.utils import status_lib
@@ -851,6 +852,7 @@ def tail_logs(cluster_name: str,
           script.
 
     """
+    logger.info('tail_log: go to backend_utils.check_cluster_available')
     # Check the status of the cluster.
     handle = backend_utils.check_cluster_available(
         cluster_name,
@@ -859,6 +861,7 @@ def tail_logs(cluster_name: str,
     backend = backend_utils.get_backend_from_handle(handle)
 
     usage_lib.record_cluster_name_for_current_operation(cluster_name)
+    logger.info('tail_log: go to backend.tail_logs')
     return backend.tail_logs(handle, job_id, follow=follow, tail=tail)
 
 

sky/global_user_state.py

@@ -18,7 +18,7 @@
 
 from sky import models
 from sky import sky_logging
-from sky.utils import common_utils
+from sky.utils import common_utils, context
 from sky.utils import db_utils
 from sky.utils import registry
 from sky.utils import status_lib
@@ -670,6 +670,7 @@ def _load_storage_mounts_metadata(
     return pickle.loads(record_storage_mounts_metadata)
 
 
+@context.cancellation_guard
 def get_cluster_from_name(
         cluster_name: Optional[str]) -> Optional[Dict[str, Any]]:
     rows = _DB.cursor.execute(

sky/server/requests/executor.py

@@ -20,8 +20,10 @@
 """
 import contextlib
 import enum
+import logging
 import multiprocessing
 import os
+import pathlib
 import queue as queue_lib
 import signal
 import sys
@@ -45,7 +47,7 @@
 from sky.server.requests.queues import local_queue
 from sky.server.requests.queues import mp_queue
 from sky.skylet import constants
-from sky.utils import annotations
+from sky.utils import annotations, context
 from sky.utils import common_utils
 from sky.utils import subprocess_utils
 from sky.utils import timeline
@@ -60,7 +62,6 @@
     from typing_extensions import ParamSpec
 
 P = ParamSpec('P')
-
 logger = sky_logging.init_logger(__name__)
 
 # On macOS, the default start method for multiprocessing is 'fork', which
@@ -325,7 +326,6 @@ def _restore_output(original_stdout: int, original_stderr: int) -> None:
 def _sigterm_handler(signum: int, frame: Optional['types.FrameType']) -> None:
     raise KeyboardInterrupt
 
-
 def _request_execution_wrapper(request_id: str,
                                ignore_return_value: bool) -> None:
     """Wrapper for a request execution.
@@ -389,6 +389,57 @@ def _request_execution_wrapper(request_id: str,
             _restore_output(original_stdout, original_stderr)
             logger.info(f'Request {request_id} finished')
 
+async def execute_request(request: api_requests.Request):
+    """Execute a request in current event loop.
+    
+    Similar to _request_execution_wrapper, but executed as coroutine in current
+    event loop. This is designed for executing tasks that are not CPU
+    intensive, e.g. sky logs.
+    """
+    ctx = context.get()
+    if ctx is None:
+        raise ValueError('Context is not initialized')
+    log_path = request.log_path
+    func = request.entrypoint
+    request_body = request.request_body
+    with api_requests.update_request(request.request_id) as request_task:
+        request_task.status = api_requests.RequestStatus.RUNNING
+    ctx.log_handler = logging.FileHandler(log_path.absolute())
+    sky_logging.reload_logger()
+        
+def prepare_request(
+    request_id: str,
+    request_name: str,
+    request_body: payloads.RequestBody,
+    func: Callable[P, Any],
+    request_cluster_name: Optional[str] = None,
+    schedule_type: api_requests.ScheduleType = (
+        api_requests.ScheduleType.LONG),
+    is_skypilot_system: bool = False,
+) -> Optional[api_requests.Request]:
+    """Prepare a request for execution."""
+    user_id = request_body.env_vars[constants.USER_ID_ENV_VAR]
+    if is_skypilot_system:
+        user_id = server_constants.SKYPILOT_SYSTEM_USER_ID
+        global_user_state.add_or_update_user(
+            models.User(id=user_id, name=user_id))
+    request = api_requests.Request(request_id=request_id,
+                                   name=server_constants.REQUEST_NAME_PREFIX +
+                                   request_name,
+                                   entrypoint=func,
+                                   request_body=request_body,
+                                   status=api_requests.RequestStatus.PENDING,
+                                   created_at=time.time(),
+                                   schedule_type=schedule_type,
+                                   user_id=user_id,
+                                   cluster_name=request_cluster_name)
+
+    if not api_requests.create_if_not_exists(request):
+        logger.debug(f'Request {request_id} already exists.')
+        return None
+
+    request.log_path.touch()
+    return request
 
 def schedule_request(
         request_id: str,
@@ -421,28 +472,16 @@ def schedule_request(
             The precondition is waited asynchronously and does not block the
             caller.
     """
-    user_id = request_body.env_vars[constants.USER_ID_ENV_VAR]
-    if is_skypilot_system:
-        user_id = server_constants.SKYPILOT_SYSTEM_USER_ID
-        global_user_state.add_or_update_user(
-            models.User(id=user_id, name=user_id))
-    request = api_requests.Request(request_id=request_id,
-                                   name=server_constants.REQUEST_NAME_PREFIX +
-                                   request_name,
-                                   entrypoint=func,
-                                   request_body=request_body,
-                                   status=api_requests.RequestStatus.PENDING,
-                                   created_at=time.time(),
-                                   schedule_type=schedule_type,
-                                   user_id=user_id,
-                                   cluster_name=request_cluster_name)
-
-    if not api_requests.create_if_not_exists(request):
-        logger.debug(f'Request {request_id} already exists.')
+    request = prepare_request(request_id,
+                              request_name,
+                              request_body,
+                              func,
+                              request_cluster_name,
+                              schedule_type,
+                              is_skypilot_system)
+    if request is None:
         return
 
-    request.log_path.touch()
-
     def enqueue():
         input_tuple = (request_id, ignore_return_value)
         logger.info(f'Queuing request: {request_id}')

sky/server/requests/requests.py

@@ -11,7 +11,7 @@
 import sqlite3
 import time
 import traceback
-from typing import Any, Callable, Dict, List, Optional, Tuple
+from typing import Any, Callable, Dict, List, Optional, ParamSpec, Tuple
 
 import colorama
 import filelock

sky/server/server.py

@@ -5,6 +5,7 @@
 import contextlib
 import dataclasses
 import datetime
+import functools
 import logging
 import multiprocessing
 import os
@@ -43,7 +44,7 @@
 from sky.server.requests import requests as requests_lib
 from sky.skylet import constants
 from sky.usage import usage_lib
-from sky.utils import admin_policy_utils
+from sky.utils import admin_policy_utils, context
 from sky.utils import common as common_lib
 from sky.utils import common_utils
 from sky.utils import dag_utils
@@ -101,7 +102,6 @@ async def dispatch(self, request: fastapi.Request, call_next):
         response.headers['X-Request-ID'] = request_id
         return response
 
-
 # Default expiration time for upload ids before cleanup.
 _DEFAULT_UPLOAD_EXPIRATION_TIME = datetime.timedelta(hours=1)
 # Key: (upload_id, user_hash), Value: the time when the upload id needs to be
@@ -666,31 +666,41 @@ async def logs(
     background_tasks: fastapi.BackgroundTasks
 ) -> fastapi.responses.StreamingResponse:
     """Tails the logs of a job."""
+    del request
     # TODO(zhwu): This should wait for the request on the cluster, e.g., async
     # launch, to finish, so that a user does not need to manually pull the
     # request status.
-    executor.schedule_request(
-        request_id=request.state.request_id,
-        request_name='logs',
-        request_body=cluster_job_body,
-        func=core.tail_logs,
-        # TODO(aylei): We have tail logs scheduled as SHORT request, because it
-        # should be responsive. However, it can be long running if the user's
-        # job keeps running, and we should avoid it taking the SHORT worker.
-        schedule_type=requests_lib.ScheduleType.SHORT,
-        request_cluster_name=cluster_job_body.cluster_name,
-    )
-
-    request_task = requests_lib.get_request(request.state.request_id)
-
+    # Only initialize the context in logs handler to limit the scope of this
+    # experimental change.
+    # TODO(aylei): init in lifespan() to enable SkyPilot context in all APIs.
+    logger.info('Initializing context')
+    context.initialize()
+    ctx = context.get()
+    # ctx.env_vars = cluster_job_body.env_vars
+    # buffer = stream_utils.StreamingBuffer()
+    # ctx.log_stream = buffer
+    logger.info('Starting tail logs')
+    task = asyncio.to_thread(core.tail_logs, **cluster_job_body.to_kwargs())
+    logger.info('Tail logs started')
+
+    async def on_disconnect():
+        ctx.cancel()
+        # buffer.close()
+        await task
+
+    background_tasks.add_task(on_disconnect)
+
+    await task
     # TODO(zhwu): This makes viewing logs in browser impossible. We should adopt
     # the same approach as /stream.
-    return stream_utils.stream_response(
-        request_id=request_task.request_id,
-        logs_path=request_task.log_path,
-        background_tasks=background_tasks,
-    )
-
+    return fastapi.responses.StreamingResponse(
+        "Hello",
+        media_type='text/plain',
+        headers={
+            'Cache-Control': 'no-cache, no-transform',
+            'X-Accel-Buffering': 'no',
+            'Transfer-Encoding': 'chunked'
+        })
 
 @app.get('/users')
 async def users() -> List[Dict[str, Any]]:

sky/server/stream_utils.py

@@ -2,6 +2,7 @@
 
 import asyncio
 import collections
+import io
 import pathlib
 from typing import AsyncGenerator, Deque, Optional
 
@@ -142,3 +143,48 @@ async def on_disconnect():
             'X-Accel-Buffering': 'no',
             'Transfer-Encoding': 'chunked'
         })
+
+class StreamingBuffer:
+    """Memory backed streaming buffer."""
+    def __init__(self):
+        self._buffer = io.StringIO()
+        self._event = asyncio.Event()
+        self._closed = False
+        
+    def write(self, data: str) -> int:
+        """Write data to the buffer."""
+        if self._closed:
+            raise ValueError('Buffer is closed')
+        n = self._buffer.write(data)
+        self._event.set()
+        return n
+        
+    def flush(self) -> None:
+        """Flush the buffer."""
+        self._buffer.flush()
+        
+    def close(self) -> None:
+        """Close the buffer."""
+        self._closed = True
+        self._event.set()
+        
+    async def read(self) -> AsyncGenerator[str, None]:
+        """Read from the buffer as a stream."""
+        # Start position in the buffer
+        pos = 0
+        
+        while True:
+            # Get current buffer contents
+            current = self._buffer.getvalue()
+            if pos < len(current):
+                # New data available, yield it
+                chunk = current[pos:]
+                pos = len(current)
+                yield chunk
+            elif self._closed:
+                # Buffer is closed and no more data
+                break
+            else:
+                # Wait for new data
+                self._event.clear()
+                await self._event.wait()