skypilot-nightly 1.0.0.dev20250806__py3-none-any.whl → 1.0.0.dev20250808__py3-none-any.whl

sky/serve/server/impl.py

@@ -1,7 +1,12 @@
 """Implementation of the SkyServe core APIs."""
+import pathlib
 import re
+import shlex
+import signal
 import tempfile
-from typing import Any, Dict, List, Optional, Tuple, Union
+import threading
+from typing import Any, Dict, List, Optional, Set, Tuple, Union
+import uuid
 
 import colorama
 import filelock
@@ -21,6 +26,7 @@ from sky.serve import serve_state
 from sky.serve import serve_utils
 from sky.skylet import constants
 from sky.utils import admin_policy_utils
+from sky.utils import command_runner
 from sky.utils import common
 from sky.utils import common_utils
 from sky.utils import controller_utils
@@ -189,7 +195,12 @@ def up(
             task_resources=task.resources)
         controller_job_id = None
         if serve_utils.is_consolidation_mode(pool):
-            controller_job_id = 0
+            # We need a unique integer per sky.serve.up call to avoid name
+            # conflict. Originally in non-consolidation mode, this is the ray
+            # job id; now we use the request id hash instead. Here we also
+            # make sure it is a 32-bit integer to avoid overflow on sqlalchemy.
+            rid = common_utils.get_current_request_id()
+            controller_job_id = hash(uuid.UUID(rid).int) & 0x7FFFFFFF
 
         vars_to_fill = {
             'remote_task_yaml_path': remote_tmp_task_yaml_path,
@@ -201,6 +212,7 @@ def up(
             'modified_catalogs':
                 service_catalog_common.get_modified_catalog_file_mounts(),
             'consolidation_mode_job_id': controller_job_id,
+            'entrypoint': shlex.quote(common_utils.get_current_command()),
             **tls_template_vars,
             **controller_utils.shared_controller_vars_to_fill(
                 controller=controller_utils.Controllers.SKY_SERVE_CONTROLLER,
@@ -708,3 +720,228 @@ def status(
                 service_record['endpoint'] = f'{protocol}://{endpoint}'
 
     return service_records
+
+
+ServiceComponentOrStr = Union[str, serve_utils.ServiceComponent]
+
+
+def tail_logs(
+    service_name: str,
+    *,
+    target: ServiceComponentOrStr,
+    replica_id: Optional[int] = None,
+    follow: bool = True,
+    tail: Optional[int] = None,
+    pool: bool = False,
+) -> None:
+    """Tail logs of a service or pool."""
+    if isinstance(target, str):
+        target = serve_utils.ServiceComponent(target)
+
+    if pool and target == serve_utils.ServiceComponent.LOAD_BALANCER:
+        raise ValueError(f'Target {target} is not supported for pool.')
+
+    if target == serve_utils.ServiceComponent.REPLICA:
+        if replica_id is None:
+            with ux_utils.print_exception_no_traceback():
+                raise ValueError(
+                    '`replica_id` must be specified when using target=REPLICA.')
+    else:
+        if replica_id is not None:
+            with ux_utils.print_exception_no_traceback():
+                raise ValueError('`replica_id` must be None when using '
+                                 'target=CONTROLLER/LOAD_BALANCER.')
+
+    controller_type = controller_utils.get_controller_for_pool(pool)
+    handle = backend_utils.is_controller_accessible(
+        controller=controller_type,
+        stopped_message=controller_type.value.default_hint_if_non_existent)
+
+    backend = backend_utils.get_backend_from_handle(handle)
+    assert isinstance(backend, backends.CloudVmRayBackend), backend
+
+    if target != serve_utils.ServiceComponent.REPLICA:
+        code = serve_utils.ServeCodeGen.stream_serve_process_logs(
+            service_name,
+            stream_controller=(
+                target == serve_utils.ServiceComponent.CONTROLLER),
+            follow=follow,
+            tail=tail,
+            pool=pool)
+    else:
+        assert replica_id is not None, service_name
+        code = serve_utils.ServeCodeGen.stream_replica_logs(service_name,
+                                                            replica_id,
+                                                            follow,
+                                                            tail=tail,
+                                                            pool=pool)
+
+    # With the stdin=subprocess.DEVNULL, the ctrl-c will not directly
+    # kill the process, so we need to handle it manually here.
+    if threading.current_thread() is threading.main_thread():
+        signal.signal(signal.SIGINT, backend_utils.interrupt_handler)
+        signal.signal(signal.SIGTSTP, backend_utils.stop_handler)
+
+    # Refer to the notes in
+    # sky/backends/cloud_vm_ray_backend.py::CloudVmRayBackend::tail_logs.
+    backend.run_on_head(handle,
+                        code,
+                        stream_logs=True,
+                        process_stream=False,
+                        ssh_mode=command_runner.SshMode.INTERACTIVE)
+
+
+def _get_all_replica_targets(
+        service_name: str, backend: backends.CloudVmRayBackend,
+        handle: backends.CloudVmRayResourceHandle,
+        pool: bool) -> Set[serve_utils.ServiceComponentTarget]:
+    """Helper function to get targets for all live replicas."""
+    code = serve_utils.ServeCodeGen.get_service_status([service_name],
+                                                       pool=pool)
+    returncode, serve_status_payload, stderr = backend.run_on_head(
+        handle,
+        code,
+        require_outputs=True,
+        stream_logs=False,
+        separate_stderr=True)
+
+    try:
+        subprocess_utils.handle_returncode(returncode,
+                                           code,
+                                           'Failed to fetch services',
+                                           stderr,
+                                           stream_logs=True)
+    except exceptions.CommandError as e:
+        raise RuntimeError(e.error_msg) from e
+
+    service_records = serve_utils.load_service_status(serve_status_payload)
+    if not service_records:
+        raise ValueError(f'Service {service_name!r} not found.')
+    assert len(service_records) == 1
+    service_record = service_records[0]
+
+    return {
+        serve_utils.ServiceComponentTarget(serve_utils.ServiceComponent.REPLICA,
+                                           replica_info['replica_id'])
+        for replica_info in service_record['replica_info']
+    }
+
+
+def sync_down_logs(
+    service_name: str,
+    *,
+    local_dir: str,
+    targets: Union[ServiceComponentOrStr, List[ServiceComponentOrStr],
+                   None] = None,
+    replica_ids: Optional[List[int]] = None,
+    tail: Optional[int] = None,
+    pool: bool = False,
+) -> str:
+    """Sync down logs of a service or pool."""
+    noun = 'pool' if pool else 'service'
+    repnoun = 'worker' if pool else 'replica'
+    caprepnoun = repnoun.capitalize()
+
+    # Step 0) get the controller handle
+    with rich_utils.safe_status(
+            ux_utils.spinner_message(f'Checking {noun} status...')):
+        controller_type = controller_utils.get_controller_for_pool(pool)
+        handle = backend_utils.is_controller_accessible(
+            controller=controller_type,
+            stopped_message=controller_type.value.default_hint_if_non_existent)
+        backend: backends.CloudVmRayBackend = (
+            backend_utils.get_backend_from_handle(handle))
+
+    requested_components: Set[serve_utils.ServiceComponent] = set()
+    if not targets:
+        # No targets specified -> request all components
+        requested_components = {
+            serve_utils.ServiceComponent.CONTROLLER,
+            serve_utils.ServiceComponent.LOAD_BALANCER,
+            serve_utils.ServiceComponent.REPLICA
+        }
+    else:
+        # Parse provided targets
+        if isinstance(targets, (str, serve_utils.ServiceComponent)):
+            requested_components = {serve_utils.ServiceComponent(targets)}
+        else:  # list
+            requested_components = {
+                serve_utils.ServiceComponent(t) for t in targets
+            }
+
+    normalized_targets: Set[serve_utils.ServiceComponentTarget] = set()
+    if serve_utils.ServiceComponent.CONTROLLER in requested_components:
+        normalized_targets.add(
+            serve_utils.ServiceComponentTarget(
+                serve_utils.ServiceComponent.CONTROLLER))
+    if serve_utils.ServiceComponent.LOAD_BALANCER in requested_components:
+        normalized_targets.add(
+            serve_utils.ServiceComponentTarget(
+                serve_utils.ServiceComponent.LOAD_BALANCER))
+    if serve_utils.ServiceComponent.REPLICA in requested_components:
+        with rich_utils.safe_status(
+                ux_utils.spinner_message(f'Getting live {repnoun} infos...')):
+            replica_targets = _get_all_replica_targets(service_name, backend,
+                                                       handle, pool)
+        if not replica_ids:
+            # Replica target requested but no specific IDs
+            # -> Get all replica logs
+            normalized_targets.update(replica_targets)
+        else:
+            # Replica target requested with specific IDs
+            requested_replica_targets = [
+                serve_utils.ServiceComponentTarget(
+                    serve_utils.ServiceComponent.REPLICA, rid)
+                for rid in replica_ids
+            ]
+            for target in requested_replica_targets:
+                if target not in replica_targets:
+                    logger.warning(f'{caprepnoun} ID {target.replica_id} not '
+                                   f'found for {service_name}. Skipping...')
+                else:
+                    normalized_targets.add(target)
+
+    def sync_down_logs_by_target(target: serve_utils.ServiceComponentTarget):
+        component = target.component
+        # We need to set one side of the pipe to a logs stream, and the other
+        # side to a file.
+        log_path = str(pathlib.Path(local_dir) / f'{target}.log')
+        stream_logs_code: str
+
+        if component == serve_utils.ServiceComponent.CONTROLLER:
+            stream_logs_code = (
+                serve_utils.ServeCodeGen.stream_serve_process_logs(
+                    service_name,
+                    stream_controller=True,
+                    follow=False,
+                    tail=tail,
+                    pool=pool))
+        elif component == serve_utils.ServiceComponent.LOAD_BALANCER:
+            stream_logs_code = (
+                serve_utils.ServeCodeGen.stream_serve_process_logs(
+                    service_name,
+                    stream_controller=False,
+                    follow=False,
+                    tail=tail,
+                    pool=pool))
+        elif component == serve_utils.ServiceComponent.REPLICA:
+            replica_id = target.replica_id
+            assert replica_id is not None, service_name
+            stream_logs_code = serve_utils.ServeCodeGen.stream_replica_logs(
+                service_name, replica_id, follow=False, tail=tail, pool=pool)
+        else:
+            assert False, component
+
+        # Refer to the notes in
+        # sky/backends/cloud_vm_ray_backend.py::CloudVmRayBackend::tail_logs.
+        backend.run_on_head(handle,
+                            stream_logs_code,
+                            stream_logs=False,
+                            process_stream=False,
+                            ssh_mode=command_runner.SshMode.INTERACTIVE,
+                            log_path=log_path)
+
+    subprocess_utils.run_in_parallel(sync_down_logs_by_target,
+                                     list(normalized_targets))
+
+    return local_dir

sky/serve/service.py

@@ -176,7 +176,7 @@ def _cleanup_task_run_script(job_id: int) -> None:
             logger.warning(f'Task run script {this_task_run_script} not found')
 
 
-def _start(service_name: str, tmp_task_yaml: str, job_id: int):
+def _start(service_name: str, tmp_task_yaml: str, job_id: int, entrypoint: str):
     """Starts the service.
     This including the controller and load balancer.
     """
@@ -228,7 +228,8 @@ def _start(service_name: str, tmp_task_yaml: str, job_id: int):
             status=serve_state.ServiceStatus.CONTROLLER_INIT,
             tls_encrypted=service_spec.tls_credential is not None,
             pool=service_spec.pool,
-            controller_pid=os.getpid())
+            controller_pid=os.getpid(),
+            entrypoint=entrypoint)
         # Directly throw an error here. See sky/serve/api.py::up
         # for more details.
         if not success:
@@ -365,8 +366,12 @@ if __name__ == '__main__':
                         required=True,
                         type=int,
                         help='Job id for the service job.')
+    parser.add_argument('--entrypoint',
+                        type=str,
+                        help='Entrypoint to launch the service',
+                        required=True)
     args = parser.parse_args()
     # We start process with 'spawn', because 'fork' could result in weird
     # behaviors; 'spawn' is also cross-platform.
     multiprocessing.set_start_method('spawn', force=True)
-    _start(args.service_name, args.task_yaml, args.job_id)
+    _start(args.service_name, args.task_yaml, args.job_id, args.entrypoint)

sky/server/common.py

@@ -16,13 +16,15 @@ import tempfile
 import threading
 import time
 import typing
-from typing import Any, Dict, Literal, Optional, Tuple, Union
+from typing import (Any, Callable, cast, Dict, Generic, Literal, Optional,
+                    Tuple, TypeVar, Union)
 from urllib import parse
 import uuid
 
 import cachetools
 import colorama
 import filelock
+from typing_extensions import ParamSpec
 
 from sky import exceptions
 from sky import sky_logging
@@ -87,7 +89,14 @@ _SERVER_INSTALL_VERSION_MISMATCH_WARNING = (
     'restarting the API server.'
     f'{colorama.Style.RESET_ALL}')
 
-RequestId = str
+T = TypeVar('T')
+P = ParamSpec('P')
+
+
+class RequestId(str, Generic[T]):
+    pass
+
+
 ApiVersion = Optional[str]
 
 logger = sky_logging.init_logger(__name__)
@@ -486,7 +495,7 @@ def handle_request_error(response: 'requests.Response') -> None:
                 f'{response.text}')
 
 
-def get_request_id(response: 'requests.Response') -> RequestId:
+def get_request_id(response: 'requests.Response') -> RequestId[T]:
     handle_request_error(response)
     request_id = response.headers.get('X-Skypilot-Request-ID')
     if request_id is None:
@@ -497,7 +506,7 @@ def get_request_id(response: 'requests.Response') -> RequestId:
                 'Failed to get request ID from SkyPilot API server at '
                 f'{get_server_url()}. Response: {response.status_code} '
                 f'{response.text}')
-    return request_id
+    return RequestId[T](request_id)
 
 
 def _start_api_server(deploy: bool = False,
@@ -753,14 +762,14 @@ def check_server_healthy_or_start_fn(deploy: bool = False,
                                   metrics_port, enable_basic_auth)
 
 
-def check_server_healthy_or_start(func):
+def check_server_healthy_or_start(func: Callable[P, T]) -> Callable[P, T]:
 
     @functools.wraps(func)
     def wrapper(*args, deploy: bool = False, host: str = '127.0.0.1', **kwargs):
         check_server_healthy_or_start_fn(deploy, host)
         return func(*args, **kwargs)
 
-    return wrapper
+    return cast(Callable[P, T], wrapper)
 
 
 def process_mounts_in_task_on_api_server(task: str, env_vars: Dict[str, str],
@@ -878,7 +887,8 @@ def request_body_to_params(body: 'pydantic.BaseModel') -> Dict[str, Any]:
 
 def reload_for_new_request(client_entrypoint: Optional[str],
                            client_command: Optional[str],
-                           using_remote_api_server: bool, user: 'models.User'):
+                           using_remote_api_server: bool, user: 'models.User',
+                           request_id: str) -> None:
     """Reload modules, global variables, and usage message for a new request."""
     # This should be called first to make sure the logger is up-to-date.
     sky_logging.reload_logger()
@@ -892,6 +902,7 @@ def reload_for_new_request(client_entrypoint: Optional[str],
         client_command=client_command,
         using_remote_api_server=using_remote_api_server,
         user=user,
+        request_id=request_id,
     )
 
     # Clear cache should be called before reload_logger and usage reset,

sky/server/constants.py

@@ -10,7 +10,7 @@ from sky.skylet import constants
 # based on version info is needed.
 # For more details and code guidelines, refer to:
 # https://docs.skypilot.co/en/latest/developers/CONTRIBUTING.html#backward-compatibility-guidelines
-API_VERSION = 14
+API_VERSION = 16
 
 # The minimum peer API version that the code should still work with.
 # Notes (dev):

sky/server/requests/executor.py

@@ -271,7 +271,8 @@ def _get_queue(schedule_type: api_requests.ScheduleType) -> RequestQueue:
 
 @contextlib.contextmanager
 def override_request_env_and_config(
-        request_body: payloads.RequestBody) -> Generator[None, None, None]:
+        request_body: payloads.RequestBody,
+        request_id: str) -> Generator[None, None, None]:
     """Override the environment and SkyPilot config for a request."""
     original_env = os.environ.copy()
     os.environ.update(request_body.env_vars)
@@ -292,7 +293,8 @@ def override_request_env_and_config(
         client_entrypoint=request_body.entrypoint,
         client_command=request_body.entrypoint_command,
         using_remote_api_server=request_body.using_remote_api_server,
-        user=user)
+        user=user,
+        request_id=request_id)
     try:
         logger.debug(
             f'override path: {request_body.override_skypilot_config_path}')
@@ -376,7 +378,7 @@ def _request_execution_wrapper(request_id: str,
         # config, as there can be some logs during override that needs to be
         # captured in the log file.
         try:
-            with override_request_env_and_config(request_body), \
+            with override_request_env_and_config(request_body, request_id), \
                 tempstore.tempdir():
                 if sky_logging.logging_enabled(logger, sky_logging.DEBUG):
                     config = skypilot_config.to_dict()

sky/server/requests/payloads.py

@@ -707,6 +707,25 @@ class JobsPoolStatusBody(RequestBody):
     pool_names: Optional[Union[str, List[str]]]
 
 
+class JobsPoolLogsBody(RequestBody):
+    """The request body for the jobs pool logs endpoint."""
+    pool_name: str
+    target: Union[str, serve.ServiceComponent]
+    worker_id: Optional[int] = None
+    follow: bool = True
+    tail: Optional[int] = None
+
+
+class JobsPoolDownloadLogsBody(RequestBody):
+    """The request body for the jobs pool download logs endpoint."""
+    pool_name: str
+    local_dir: str
+    targets: Optional[Union[str, serve.ServiceComponent,
+                            List[Union[str, serve.ServiceComponent]]]]
+    worker_ids: Optional[List[int]] = None
+    tail: Optional[int] = None
+
+
 class UploadZipFileResponse(pydantic.BaseModel):
     """The response body for the upload zip file endpoint."""
     status: str

sky/setup_files/alembic.ini

@@ -94,6 +94,10 @@ version_table = alembic_version_state_db
 version_locations = %(here)s/../schemas/db/spot_jobs
 version_table = alembic_version_spot_jobs_db
 
+[serve_db]
+version_locations = %(here)s/../schemas/db/serve_state
+version_table = alembic_version_serve_state_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

sky/task.py

@@ -241,8 +241,8 @@ class Task:
         self,
         name: Optional[str] = None,
         *,
-        setup: Optional[str] = None,
-        run: Optional[CommandOrCommandGen] = None,
+        setup: Optional[Union[str, List[str]]] = None,
+        run: Optional[Union[CommandOrCommandGen, List[str]]] = None,
         envs: Optional[Dict[str, str]] = None,
         secrets: Optional[Dict[str, str]] = None,
         workdir: Optional[Union[str, Dict[str, Any]]] = None,
@@ -293,15 +293,15 @@ class Task:
 
         Args:
           name: A string name for the Task for display purposes.
-          setup: A setup command, which will be run before executing the run
+          setup: A setup command(s), which will be run before executing the run
             commands ``run``, and executed under ``workdir``.
           run: The actual command for the task. If not None, either a shell
-            command (str) or a command generator (callable).  If latter, it
-            must take a node rank and a list of node addresses as input and
-            return a shell command (str) (valid to return None for some nodes,
-            in which case no commands are run on them).  Run commands will be
-            run under ``workdir``. Note the command generator should be a
-            self-contained lambda.
+            command(s) (str, list(str)) or a command generator (callable). If
+            latter, it must take a node rank and a list of node addresses as
+            input and return a shell command (str) (valid to return None for
+            some nodes, in which case no commands are run on them). Run
+            commands will be run under ``workdir``. Note the command generator
+            should be a self-contained lambda.
           envs: A dictionary of environment variables to set before running the
             setup and run commands.
           secrets: A dictionary of secret environment variables to set before
@@ -347,15 +347,22 @@ class Task:
             YAML config.
         """
         self.name = name
-        self.run = run
         self.storage_mounts: Dict[str, storage_lib.Storage] = {}
         self.storage_plans: Dict[storage_lib.Storage,
                                  storage_lib.StoreType] = {}
-        self.setup = setup
         self._envs = envs or {}
         self._secrets = secrets or {}
         self._volumes = volumes or {}
 
+        # concatenate commands if given as list
+        def _concat(commands):
+            if isinstance(commands, list):
+                return '\n'.join(commands)
+            return commands
+
+        self.run = _concat(run)
+        self.setup = _concat(setup)
+
         # Validate Docker login configuration early if both envs and secrets
         # contain Docker variables
         if self._envs or self._secrets:

sky/templates/kubernetes-ray.yml.j2

@@ -777,6 +777,11 @@ available_node_types:
                 {{ ray_installation_commands }}
 
                 VIRTUAL_ENV=~/skypilot-runtime ~/.local/bin/uv pip install skypilot[kubernetes,remote]
+                # Wait for `patch` package to be installed before applying ray patches
+                until dpkg -l | grep -q "^ii  patch "; do
+                  sleep 0.1
+                  echo "Waiting for patch package to be installed..."
+                done
                 # Apply Ray patches for progress bar fix
                 ~/.local/bin/uv pip list | grep "ray " | grep 2.9.3 2>&1 > /dev/null && { 
                   VIRTUAL_ENV=~/skypilot-runtime python -c "from sky.skylet.ray_patches import patch; patch()" || exit 1; 

sky/templates/sky-serve-controller.yaml.j2

@@ -57,6 +57,7 @@ run: |
    -u -m sky.serve.service \
     --service-name {{service_name}} \
     --task-yaml {{remote_task_yaml_path}} \
+    --entrypoint {{entrypoint}} \
     {%- if consolidation_mode_job_id is not none %}
     --job-id {{consolidation_mode_job_id}} \
     {%- else %}

sky/usage/usage_lib.py

@@ -10,6 +10,8 @@ import traceback
 import typing
 from typing import Any, Callable, Dict, List, Optional, Union
 
+from typing_extensions import ParamSpec
+
 import sky
 from sky import sky_logging
 from sky.adaptors import common as adaptors_common
@@ -517,26 +519,26 @@ def entrypoint_context(name: str, fallback: bool = False):
 
 
 T = typing.TypeVar('T')
+P = ParamSpec('P')
 
 
 @typing.overload
 def entrypoint(
         name_or_fn: str,
-        fallback: bool = False
-) -> Callable[[Callable[..., T]], Callable[..., T]]:
+        fallback: bool = False) -> Callable[[Callable[P, T]], Callable[P, T]]:
     ...
 
 
 @typing.overload
-def entrypoint(name_or_fn: Callable[..., T],
-               fallback: bool = False) -> Callable[..., T]:
+def entrypoint(name_or_fn: Callable[P, T],
+               fallback: bool = False) -> Callable[P, T]:
     ...
 
 
 def entrypoint(
-    name_or_fn: Union[str, Callable[..., T]],
+    name_or_fn: Union[str, Callable[P, T]],
     fallback: bool = False
-) -> Union[Callable[..., T], Callable[[Callable[..., T]], Callable[..., T]]]:
+) -> Union[Callable[P, T], Callable[[Callable[P, T]], Callable[P, T]]]:
     return common_utils.make_decorator(entrypoint_context,
                                        name_or_fn,
                                        fallback=fallback)

sky/utils/annotations.py

@@ -1,14 +1,19 @@
 """Annotations for public APIs."""
 
 import functools
-from typing import Callable, Literal
+from typing import Callable, Literal, TypeVar
+
+from typing_extensions import ParamSpec
 
 # Whether the current process is a SkyPilot API server process.
 is_on_api_server = True
 FUNCTIONS_NEED_RELOAD_CACHE = []
 
+T = TypeVar('T')
+P = ParamSpec('P')
+
 
-def client_api(func):
+def client_api(func: Callable[P, T]) -> Callable[P, T]:
     """Mark a function as a client-side API.
 
     Code invoked by server-side functions will find annotations.is_on_api_server
@@ -38,7 +43,7 @@ def lru_cache(scope: Literal['global', 'request'], *lru_cache_args,
         lru_cache_kwargs: Keyword arguments for functools.lru_cache.
     """
 
-    def decorator(func: Callable) -> Callable:
+    def decorator(func: Callable[P, T]) -> Callable[P, T]:
         if scope == 'global':
             return functools.lru_cache(*lru_cache_args,
                                        **lru_cache_kwargs)(func)

sky/utils/cli_utils/status_utils.py

@@ -401,7 +401,7 @@ def _get_estimated_cost_for_cost_report(
 
 
 def show_kubernetes_cluster_status_table(
-        clusters: List['kubernetes_utils.KubernetesSkyPilotClusterInfo'],
+        clusters: List['kubernetes_utils.KubernetesSkyPilotClusterInfoPayload'],
         show_all: bool) -> None:
     """Compute cluster table values and display for Kubernetes clusters."""
     status_columns = [

sky/utils/common_utils.py

@@ -271,12 +271,13 @@ _current_command: Optional[str] = None
 _current_client_entrypoint: Optional[str] = None
 _using_remote_api_server: Optional[bool] = None
 _current_user: Optional['models.User'] = None
+_current_request_id: Optional[str] = None
 
 
 def set_request_context(client_entrypoint: Optional[str],
                         client_command: Optional[str],
                         using_remote_api_server: bool,
-                        user: Optional['models.User']):
+                        user: Optional['models.User'], request_id: str) -> None:
     """Override the current client entrypoint and command.
 
     This is useful when we are on the SkyPilot API server side and we have a
@@ -286,10 +287,19 @@ def set_request_context(client_entrypoint: Optional[str],
     global _current_client_entrypoint
     global _using_remote_api_server
     global _current_user
+    global _current_request_id
     _current_command = client_command
     _current_client_entrypoint = client_entrypoint
     _using_remote_api_server = using_remote_api_server
     _current_user = user
+    _current_request_id = request_id
+
+
+def get_current_request_id() -> str:
+    """Returns the current request id."""
+    if _current_request_id is not None:
+        return _current_request_id
+    return 'dummy-request-id'
 
 
 def get_current_command() -> str: