skypilot-nightly 1.0.0.dev20250804__py3-none-any.whl → 1.0.0.dev20250807__py3-none-any.whl

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/server/rest.py

@@ -1,5 +1,6 @@
 """REST API client of SkyPilot API server"""
 
+import asyncio
 import contextlib
 import contextvars
 import functools
@@ -21,9 +22,11 @@ from sky.utils import ux_utils
 logger = sky_logging.init_logger(__name__)
 
 if typing.TYPE_CHECKING:
+    import aiohttp
     import requests
 
 else:
+    aiohttp = adaptors_common.LazyImport('aiohttp')
     requests = adaptors_common.LazyImport('requests')
 
 F = TypeVar('F', bound=Callable[..., Any])
@@ -204,3 +207,114 @@ def request_without_retry(method, url, **kwargs) -> 'requests.Response':
     if remote_version is not None:
         versions.set_remote_version(remote_version)
     return response
+
+
+# Async versions of the above functions
+
+
+async def request_async(session: 'aiohttp.ClientSession', method: str, url: str,
+                        **kwargs) -> 'aiohttp.ClientResponse':
+    """Send an async request to the API server, retry on server temporarily
+    unavailable."""
+    max_retries = 3
+    initial_backoff = 1.0
+    max_backoff_factor = 5
+
+    backoff = common_utils.Backoff(initial_backoff, max_backoff_factor)
+    last_exception = Exception('Uknown Exception')  # this will be replaced by e
+
+    for retry_count in range(max_retries):
+        try:
+            return await request_without_retry_async(session, method, url,
+                                                     **kwargs)
+        except exceptions.RequestInterruptedError:
+            logger.debug('Request interrupted. Retry immediately.')
+            continue
+        except Exception as e:  # pylint: disable=broad-except
+            last_exception = e
+            if retry_count >= max_retries - 1:
+                # Retries exhausted
+                raise
+
+            # Check if this is a transient error (similar to sync version logic)
+            is_transient = _is_transient_error_async(e)
+            if not is_transient:
+                # Permanent error, no need to retry
+                raise
+
+            logger.debug(f'Retry async request due to {e}, '
+                         f'attempt {retry_count + 1}/{max_retries}')
+            await asyncio.sleep(backoff.current_backoff())
+
+    # This should never be reached, but just in case
+    raise last_exception
+
+
+async def request_without_retry_async(session: 'aiohttp.ClientSession',
+                                      method: str, url: str,
+                                      **kwargs) -> 'aiohttp.ClientResponse':
+    """Send an async request to the API server without retry."""
+    # Add API version headers for compatibility (like sync version does)
+    if 'headers' not in kwargs:
+        kwargs['headers'] = {}
+    kwargs['headers'][constants.API_VERSION_HEADER] = str(constants.API_VERSION)
+    kwargs['headers'][constants.VERSION_HEADER] = (
+        versions.get_local_readable_version())
+
+    try:
+        response = await session.request(method, url, **kwargs)
+
+        # Handle server unavailability (503 status) - same as sync version
+        if response.status == 503:
+            with ux_utils.print_exception_no_traceback():
+                raise exceptions.ServerTemporarilyUnavailableError(
+                    'SkyPilot API server is temporarily unavailable. '
+                    'Please try again later.')
+
+        # Set remote API version and version from headers - same as sync version
+        remote_api_version = response.headers.get(constants.API_VERSION_HEADER)
+        remote_version = response.headers.get(constants.VERSION_HEADER)
+        if remote_api_version is not None:
+            versions.set_remote_api_version(int(remote_api_version))
+        if remote_version is not None:
+            versions.set_remote_version(remote_version)
+
+        return response
+
+    except aiohttp.ClientError as e:
+        # Convert aiohttp errors to appropriate SkyPilot exceptions
+        if isinstance(e, aiohttp.ClientConnectorError):
+            raise exceptions.RequestInterruptedError(
+                f'Connection failed: {e}') from e
+        elif isinstance(e, aiohttp.ClientTimeout):
+            raise exceptions.RequestInterruptedError(
+                f'Request timeout: {e}') from e
+        else:
+            raise
+
+
+def _is_transient_error_async(e: Exception) -> bool:
+    """Check if an exception from async request is transient and should be
+    retried.
+
+    Mirrors the logic from the sync version's is_transient_error().
+    """
+    if isinstance(e, aiohttp.ClientError):
+        # For response errors, check status code if available
+        if isinstance(e, aiohttp.ClientResponseError):
+            # Only server error is considered as transient (same as sync
+            # version)
+            return e.status >= 500
+        # Consider connection errors and timeouts as transient
+        if isinstance(e, (aiohttp.ClientConnectorError, aiohttp.ClientTimeout)):
+            return True
+
+    # Consider server temporarily unavailable as transient
+    if isinstance(e, exceptions.ServerTemporarilyUnavailableError):
+        return True
+
+    # It is hard to enumerate all other errors that are transient, e.g.
+    # broken pipe, connection refused, etc. Instead, it is safer to assume
+    # all other errors might be transient since we only retry for 3 times
+    # by default. (Same comment as in sync version)
+    return True

sky/server/server.py

@@ -51,6 +51,8 @@ from sky.server import metrics
 from sky.server import state
 from sky.server import stream_utils
 from sky.server import versions
+from sky.server.auth import authn
+from sky.server.auth import oauth2_proxy
 from sky.server.requests import executor
 from sky.server.requests import payloads
 from sky.server.requests import preconditions
@@ -120,41 +122,6 @@ def _basic_auth_401_response(content: str):
         content=content)
 
 
-# TODO(hailong): Remove this function and use request.state.auth_user instead.
-async def _override_user_info_in_request_body(request: fastapi.Request,
-                                              auth_user: Optional[models.User]):
-    if auth_user is None:
-        return
-
-    body = await request.body()
-    if body:
-        try:
-            original_json = await request.json()
-        except (json.JSONDecodeError, UnicodeDecodeError) as e:
-            logger.error(f'Error parsing request JSON: {e}')
-        else:
-            logger.debug(f'Overriding user for {request.state.request_id}: '
-                         f'{auth_user.name}, {auth_user.id}')
-            if 'env_vars' in original_json:
-                if isinstance(original_json.get('env_vars'), dict):
-                    original_json['env_vars'][
-                        constants.USER_ID_ENV_VAR] = auth_user.id
-                    original_json['env_vars'][
-                        constants.USER_ENV_VAR] = auth_user.name
-                else:
-                    logger.warning(
-                        f'"env_vars" in request body is not a dictionary '
-                        f'for request {request.state.request_id}. '
-                        'Skipping user info injection into body.')
-            else:
-                original_json['env_vars'] = {}
-                original_json['env_vars'][
-                    constants.USER_ID_ENV_VAR] = auth_user.id
-                original_json['env_vars'][
-                    constants.USER_ENV_VAR] = auth_user.name
-            request._body = json.dumps(original_json).encode('utf-8')  # pylint: disable=protected-access
-
-
 def _try_set_basic_auth_user(request: fastapi.Request):
     auth_header = request.headers.get('authorization')
     if not auth_header or not auth_header.lower().startswith('basic '):
@@ -281,7 +248,7 @@ class BasicAuthMiddleware(starlette.middleware.base.BaseHTTPMiddleware):
                     apr_md5_crypt.verify(password, user.password)):
                 valid_user = True
                 request.state.auth_user = user
-                await _override_user_info_in_request_body(request, user)
+                await authn.override_user_info_in_request_body(request, user)
                 break
         if not valid_user:
             return _basic_auth_401_response('Invalid credentials')
@@ -400,7 +367,7 @@ class BearerTokenMiddleware(starlette.middleware.base.BaseHTTPMiddleware):
             request.state.auth_user = auth_user
 
             # Override user info in request body for service account requests
-            await _override_user_info_in_request_body(request, auth_user)
+            await authn.override_user_info_in_request_body(request, auth_user)
 
             logger.debug(f'Authenticated service account: {user_id}')
 
@@ -445,7 +412,7 @@ class AuthProxyMiddleware(starlette.middleware.base.BaseHTTPMiddleware):
         if auth_user is not None:
             request.state.auth_user = auth_user
 
-        await _override_user_info_in_request_body(request, auth_user)
+        await authn.override_user_info_in_request_body(request, auth_user)
         return await call_next(request)
 
 
@@ -484,6 +451,8 @@ async def lifespan(app: fastapi.FastAPI):  # pylint: disable=redefined-outer-nam
     del app  # unused
     # Startup: Run background tasks
     for event in daemons.INTERNAL_REQUEST_DAEMONS:
+        if event.should_skip():
+            continue
         try:
             executor.schedule_request(
                 request_id=event.id,
@@ -625,6 +594,8 @@ app.add_middleware(
 # RBACMiddleware must precede all the auth middleware, so it can access
 # request.state.auth_user.
 app.add_middleware(RBACMiddleware)
+# Authentication based on oauth2-proxy.
+app.add_middleware(oauth2_proxy.OAuth2ProxyMiddleware)
 # AuthProxyMiddleware should precede BasicAuthMiddleware and
 # BearerTokenMiddleware, since it should be skipped if either of those set the
 # auth user.
@@ -1574,9 +1545,42 @@ async def health(request: fastapi.Request) -> Dict[str, Any]:
         - commit: str; The commit hash of SkyPilot used for API server.
     """
     user = request.state.auth_user
-    logger.info(f'Health endpoint: request.state.auth_user = {user}')
+    server_status = common.ApiServerStatus.HEALTHY
+    if getattr(request.state, 'anonymous_user', False):
+        # API server authentication is enabled, but the request is not
+        # authenticated. We still have to serve the request because the
+        # /api/health endpoint has two different usage:
+        # 1. For health check from `api start` and external ochestration
+        #    tools (k8s), which does not require authentication and user info.
+        # 2. Return server info to client and hint client to login if required.
+        # Separating these two usage to different APIs will break backward
+        # compatibility for existing ochestration solutions (e.g. helm chart).
+        # So we serve these two usages in a backward compatible manner below.
+        client_version = versions.get_remote_api_version()
+        # - For Client with API version >= 14, we return 200 response with
+        #   status=NEEDS_AUTH, new client will handle the login process.
+        # - For health check from `sky api start`, the client code always uses
+        #   the same API version with the server, thus there is no compatibility
+        #   issue.
+        server_status = common.ApiServerStatus.NEEDS_AUTH
+        if client_version is None:
+            # - For health check from ochestration tools (e.g. k8s), we also
+            #   return 200 with status=NEEDS_AUTH, which passes HTTP probe
+            #   check.
+            # - There is no harm when an malicious client calls /api/health
+            #   without authentication since no sensitive information is
+            #   returned.
+            return {'status': common.ApiServerStatus.HEALTHY}
+        # TODO(aylei): remove this after min_compatible_api_version >= 14.
+        if client_version < 14:
+            # For Client with API version < 14, the NEEDS_AUTH status is not
+            # honored. Return 401 to trigger the login process.
+            raise fastapi.HTTPException(status_code=401,
+                                        detail='Authentication required')
+
+    logger.debug(f'Health endpoint: request.state.auth_user = {user}')
     return {
-        'status': common.ApiServerStatus.HEALTHY.value,
+        'status': server_status,
         # Kept for backward compatibility, clients before 0.11.0 will read this
         # field to check compatibility and hint the user to upgrade the CLI.
         # TODO(aylei): remove this field after 0.13.0

sky/setup_files/dependencies.py

@@ -69,6 +69,7 @@ install_requires = [
     'gitpython',
     'types-paramiko',
     'alembic',
+    'aiohttp',
 ]
 
 server_dependencies = [
@@ -76,6 +77,7 @@ server_dependencies = [
     'sqlalchemy_adapter',
     'passlib',
     'pyjwt',
+    'aiohttp',
 ]
 
 local_ray = [

sky/skylet/constants.py

@@ -406,7 +406,7 @@ PERSISTENT_RUN_SCRIPT_DIR = '~/.sky/.controller_recovery_task_run'
 PERSISTENT_RUN_RESTARTING_SIGNAL_FILE = (
     '~/.sky/.controller_recovery_restarting_signal')
 
-HA_PERSISTENT_RECOVERY_LOG_PATH = '/tmp/ha_recovery.log'
+HA_PERSISTENT_RECOVERY_LOG_PATH = '/tmp/{}ha_recovery.log'
 
 # The placeholder for the local skypilot config path in file mounts for
 # controllers.

sky/skylet/events.py

@@ -96,8 +96,12 @@ class ServiceUpdateEvent(SkyletEvent):
     """
     EVENT_INTERVAL_SECONDS = 300
 
+    def __init__(self, pool: bool) -> None:
+        super().__init__()
+        self._pool = pool
+
     def _run(self):
-        serve_utils.update_service_status()
+        serve_utils.update_service_status(self._pool)
 
 
 class UsageHeartbeatReportEvent(SkyletEvent):

sky/skylet/skylet.py

@@ -24,7 +24,9 @@ EVENTS = [
     # This is for monitoring controller job status. If it becomes
     # unhealthy, this event will correctly update the controller
     # status to CONTROLLER_FAILED.
-    events.ServiceUpdateEvent(),
+    events.ServiceUpdateEvent(pool=False),
+    # Status refresh for pool.
+    events.ServiceUpdateEvent(pool=True),
     # Report usage heartbeat every 10 minutes.
     events.UsageHeartbeatReportEvent(),
 ]

sky/task.py

@@ -241,21 +241,26 @@ 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,
         num_nodes: Optional[int] = None,
+        file_mounts: Optional[Dict[str, str]] = None,
+        storage_mounts: Optional[Dict[str, storage_lib.Storage]] = None,
         volumes: Optional[Dict[str, str]] = None,
+        resources: Optional[Union['resources_lib.Resources',
+                                  List['resources_lib.Resources'],
+                                  Set['resources_lib.Resources']]] = None,
         # Advanced:
         docker_image: Optional[str] = None,
         event_callback: Optional[str] = None,
         blocked_resources: Optional[Iterable['resources_lib.Resources']] = None,
         # Internal use only.
-        file_mounts_mapping: Optional[Dict[str, str]] = None,
-        volume_mounts: Optional[List[volume_lib.VolumeMount]] = None,
-        metadata: Optional[Dict[str, Any]] = None,
+        _file_mounts_mapping: Optional[Dict[str, str]] = None,
+        _volume_mounts: Optional[List[volume_lib.VolumeMount]] = None,
+        _metadata: Optional[Dict[str, Any]] = None,
         _user_specified_yaml: Optional[str] = None,
     ):
         """Initializes a Task.
@@ -288,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
@@ -315,22 +320,49 @@ class Task:
             setup/run command, where ``run`` can either be a str, meaning all
             nodes get the same command, or a lambda, with the semantics
             documented above.
+          file_mounts: An optional dict of ``{remote_path: (local_path|cloud
+            URI)}``, where remote means the VM(s) on which this Task will
+            eventually run on, and local means the node from which the task is
+            launched.
+          storage_mounts: an optional dict of ``{mount_path: sky.Storage
+            object}``, where mount_path is the path inside the remote VM(s)
+            where the Storage object will be mounted on.
+          volumes: A dict of volumes to be mounted for the task. The dict has
+            the form of ``{mount_path: volume_name}``.
+          resources: either a sky.Resources, a set of them, or a list of them.
+            A set or a list of resources asks the optimizer to "pick the
+            best of these resources" to run this task.
           docker_image: (EXPERIMENTAL: Only in effect when LocalDockerBackend
             is used.) The base docker image that this Task will be built on.
             Defaults to 'gpuci/miniforge-cuda:11.4-devel-ubuntu18.04'.
+          event_callback: A bash script that will be executed when the task
+            changes state.
           blocked_resources: A set of resources that this task cannot run on.
-          metadata: A dictionary of metadata to be added to the task.
+          _file_mounts_mapping: (Internal use only) A dictionary of file mounts
+            mapping.
+          _volume_mounts: (Internal use only) A list of volume mounts.
+          _metadata: (Internal use only) A dictionary of metadata to be added to
+            the task.
+          _user_specified_yaml: (Internal use only) A string of user-specified
+            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:
@@ -372,11 +404,19 @@ class Task:
         self.best_resources: Optional[sky.Resources] = None
 
         # For internal use only.
-        self.file_mounts_mapping: Optional[Dict[str, str]] = file_mounts_mapping
+        self.file_mounts_mapping: Optional[Dict[str,
+                                                str]] = _file_mounts_mapping
         self.volume_mounts: Optional[List[volume_lib.VolumeMount]] = (
-            volume_mounts)
+            _volume_mounts)
 
-        self._metadata = metadata if metadata is not None else {}
+        self._metadata = _metadata if _metadata is not None else {}
+
+        if resources is not None:
+            self.set_resources(resources)
+        if storage_mounts is not None:
+            self.set_storage_mounts(storage_mounts)
+        if file_mounts is not None:
+            self.set_file_mounts(file_mounts)
 
         dag = sky.dag.get_current_dag()
         if dag is not None:
@@ -621,10 +661,10 @@ class Task:
             num_nodes=config.pop('num_nodes', None),
             envs=config.pop('envs', None),
             secrets=config.pop('secrets', None),
-            event_callback=config.pop('event_callback', None),
-            file_mounts_mapping=config.pop('file_mounts_mapping', None),
             volumes=config.pop('volumes', None),
-            metadata=config.pop('_metadata', None),
+            event_callback=config.pop('event_callback', None),
+            _file_mounts_mapping=config.pop('file_mounts_mapping', None),
+            _metadata=config.pop('_metadata', None),
             _user_specified_yaml=user_specified_yaml,
         )
 

sky/templates/kubernetes-ray.yml.j2

@@ -777,6 +777,15 @@ 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; 
+                }
                 touch /tmp/ray_skypilot_installation_complete
                 echo "=== Ray and skypilot installation completed ==="
 

sky/templates/nebius-ray.yml.j2

@@ -46,6 +46,7 @@ available_node_types:
       InstanceType: {{instance_type}}
       ImageId: {{image_id}}
       DiskSize: {{disk_size}}
+      use_spot: {{ use_spot }}
       network_tier: {{network_tier}}
       filesystems:
         {%- for fs in filesystems %}

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/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:

sky/utils/controller_utils.py

@@ -209,6 +209,13 @@ class Controllers(enum.Enum):
         return None
 
 
+def get_controller_for_pool(pool: bool) -> Controllers:
+    """Get the controller type."""
+    if pool:
+        return Controllers.JOBS_CONTROLLER
+    return Controllers.SKY_SERVE_CONTROLLER
+
+
 def high_availability_specified(cluster_name: Optional[str]) -> bool:
     """Check if the controller high availability is specified in user config.
     """

sky/utils/db/migration_utils.py

@@ -19,11 +19,11 @@ logger = sky_logging.init_logger(__name__)
 DB_INIT_LOCK_TIMEOUT_SECONDS = 10
 
 GLOBAL_USER_STATE_DB_NAME = 'state_db'
-GLOBAL_USER_STATE_VERSION = '001'
+GLOBAL_USER_STATE_VERSION = '002'
 GLOBAL_USER_STATE_LOCK_PATH = '~/.sky/locks/.state_db.lock'
 
 SPOT_JOBS_DB_NAME = 'spot_jobs_db'
-SPOT_JOBS_VERSION = '002'
+SPOT_JOBS_VERSION = '003'
 SPOT_JOBS_LOCK_PATH = '~/.sky/locks/.spot_jobs_db.lock'