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

sky/client/sdk.py

@@ -10,75 +10,165 @@ Usage example:
     statuses = sky.get(request_id)
 
 """
-import getpass
+from http import cookiejar
 import json
 import logging
 import os
-import pathlib
 import subprocess
 import typing
-from typing import Any, Dict, List, Optional, Tuple, Union
-import webbrowser
+from typing import (Any, Dict, Iterator, List, Literal, Optional, Tuple,
+                    TypeVar, Union)
+from urllib import parse as urlparse
 
 import click
 import colorama
 import filelock
 
 from sky import admin_policy
-from sky import backends
 from sky import exceptions
 from sky import sky_logging
 from sky import skypilot_config
 from sky.adaptors import common as adaptors_common
 from sky.client import common as client_common
+from sky.client import oauth as oauth_lib
+from sky.jobs import scheduler
+from sky.jobs import utils as managed_job_utils
+from sky.schemas.api import responses
 from sky.server import common as server_common
+from sky.server import rest
+from sky.server import versions
 from sky.server.requests import payloads
+from sky.server.requests import request_names
 from sky.server.requests import requests as requests_lib
+from sky.skylet import autostop_lib
 from sky.skylet import constants
 from sky.usage import usage_lib
+from sky.utils import admin_policy_utils
 from sky.utils import annotations
 from sky.utils import cluster_utils
 from sky.utils import common
 from sky.utils import common_utils
+from sky.utils import context as sky_context
 from sky.utils import dag_utils
 from sky.utils import env_options
+from sky.utils import infra_utils
 from sky.utils import rich_utils
 from sky.utils import status_lib
 from sky.utils import subprocess_utils
 from sky.utils import ux_utils
+from sky.utils import yaml_utils
+from sky.utils.kubernetes import ssh_utils
 
 if typing.TYPE_CHECKING:
+    import base64
+    import binascii
     import io
+    import pathlib
+    import time
+    import webbrowser
 
     import psutil
     import requests
 
     import sky
+    from sky import backends
+    from sky import catalog
+    from sky import models
+    from sky.provision.kubernetes import utils as kubernetes_utils
+    from sky.skylet import job_lib
 else:
+    # only used in api_login()
+    base64 = adaptors_common.LazyImport('base64')
+    binascii = adaptors_common.LazyImport('binascii')
+    pathlib = adaptors_common.LazyImport('pathlib')
+    time = adaptors_common.LazyImport('time')
+    # only used in dashboard() and api_login()
+    webbrowser = adaptors_common.LazyImport('webbrowser')
+    # only used in api_stop()
     psutil = adaptors_common.LazyImport('psutil')
-    requests = adaptors_common.LazyImport('requests')
 
 logger = sky_logging.init_logger(__name__)
 logging.getLogger('httpx').setLevel(logging.CRITICAL)
 
+_LINE_PROCESSED_KEY = 'line_processed'
 
-def stream_response(request_id: Optional[str],
+T = TypeVar('T')
+
+
+def reload_config() -> None:
+    """Reloads the client-side config."""
+    skypilot_config.safe_reload_config()
+
+
+# The overloads are not comprehensive - e.g. get_result Literal[False] could be
+# specified to return None. We can add more overloads if needed. To do that see
+# https://github.com/python/mypy/issues/8634#issuecomment-609411104
+@typing.overload
+def stream_response(request_id: None,
                     response: 'requests.Response',
-                    output_stream: Optional['io.TextIOBase'] = None) -> Any:
+                    output_stream: Optional['io.TextIOBase'] = None,
+                    resumable: bool = False,
+                    get_result: bool = True) -> None:
+    ...
+
+
+@typing.overload
+def stream_response(request_id: server_common.RequestId[T],
+                    response: 'requests.Response',
+                    output_stream: Optional['io.TextIOBase'] = None,
+                    resumable: bool = False,
+                    get_result: Literal[True] = True) -> T:
+    ...
+
+
+@typing.overload
+def stream_response(request_id: server_common.RequestId[T],
+                    response: 'requests.Response',
+                    output_stream: Optional['io.TextIOBase'] = None,
+                    resumable: bool = False,
+                    get_result: bool = True) -> Optional[T]:
+    ...
+
+
+def stream_response(request_id: Optional[server_common.RequestId[T]],
+                    response: 'requests.Response',
+                    output_stream: Optional['io.TextIOBase'] = None,
+                    resumable: bool = False,
+                    get_result: bool = True) -> Optional[T]:
     """Streams the response to the console.
 
     Args:
-        request_id: The request ID.
+        request_id: The request ID of the request to stream. May be a full
+            request ID or a prefix.
+            If None, the latest request submitted to the API server is streamed.
+            Using None request_id is not recommended in multi-user environments.
         response: The HTTP response.
         output_stream: The output stream to write to. If None, print to the
             console.
+        resumable: Whether the response is resumable on retry. If True, the
+            streaming will start from the previous failure point on retry.
+        get_result: Whether to get the result of the request. This will
+            typically be set to False for `--no-follow` flags as requests may
+            continue to run for long periods of time without further streaming.
     """
 
+    retry_context: Optional[rest.RetryContext] = None
+    if resumable:
+        retry_context = rest.get_retry_context()
     try:
+        line_count = 0
         for line in rich_utils.decode_rich_status(response):
             if line is not None:
-                print(line, flush=True, end='', file=output_stream)
-        return get(request_id)
+                line_count += 1
+                if retry_context is None:
+                    print(line, flush=True, end='', file=output_stream)
+                elif line_count > retry_context.line_processed:
+                    print(line, flush=True, end='', file=output_stream)
+                    retry_context.line_processed = line_count
+        if request_id is not None and get_result:
+            return get(request_id)
+        else:
+            return None
     except Exception:  # pylint: disable=broad-except
         logger.debug(f'To stream request logs: sky api logs {request_id}')
         raise
@@ -87,13 +177,18 @@ def stream_response(request_id: Optional[str],
 @usage_lib.entrypoint
 @server_common.check_server_healthy_or_start
 @annotations.client_api
-def check(clouds: Optional[Tuple[str]],
-          verbose: bool) -> server_common.RequestId:
+def check(
+    infra_list: Optional[Tuple[str, ...]],
+    verbose: bool,
+    workspace: Optional[str] = None
+) -> server_common.RequestId[Dict[str, List[str]]]:
     """Checks the credentials to enable clouds.
 
     Args:
-        clouds: The clouds to check.
+        infra: The infra to check.
         verbose: Whether to show verbose output.
+        workspace: The workspace to check. If None, all workspaces will be
+        checked.
 
     Returns:
         The request ID of the check request.
@@ -101,41 +196,69 @@ def check(clouds: Optional[Tuple[str]],
     Request Returns:
         None
     """
-    body = payloads.CheckBody(clouds=clouds, verbose=verbose)
-    response = requests.post(f'{server_common.get_server_url()}/check',
-                             json=json.loads(body.model_dump_json()),
-                             cookies=server_common.get_api_cookie_jar())
+    if infra_list is None:
+        clouds = None
+    else:
+        specified_clouds = []
+        for infra_str in infra_list:
+            infra = infra_utils.InfraInfo.from_str(infra_str)
+            if infra.cloud is None:
+                with ux_utils.print_exception_no_traceback():
+                    raise ValueError(f'Invalid infra to check: {infra_str}')
+            if infra.region is not None or infra.zone is not None:
+                region_zone = infra_str.partition('/')[-1]
+                logger.warning(f'Infra {infra_str} is specified, but `check` '
+                               f'only supports checking {infra.cloud}, '
+                               f'ignoring {region_zone}')
+            specified_clouds.append(infra.cloud)
+        clouds = tuple(specified_clouds)
+    body = payloads.CheckBody(clouds=clouds,
+                              verbose=verbose,
+                              workspace=workspace)
+    response = server_common.make_authenticated_request(
+        'POST', '/check', json=json.loads(body.model_dump_json()))
     return server_common.get_request_id(response)
 
 
 @usage_lib.entrypoint
 @server_common.check_server_healthy_or_start
 @annotations.client_api
-def enabled_clouds() -> server_common.RequestId:
+def enabled_clouds(workspace: Optional[str] = None,
+                   expand: bool = False) -> server_common.RequestId[List[str]]:
     """Gets the enabled clouds.
 
+    Args:
+        workspace: The workspace to get the enabled clouds for. If None, the
+        active workspace will be used.
+        expand: Whether to expand Kubernetes and SSH to list of resource pools.
+
     Returns:
         The request ID of the enabled clouds request.
 
     Request Returns:
         A list of enabled clouds in string format.
     """
-    response = requests.get(f'{server_common.get_server_url()}/enabled_clouds',
-                            cookies=server_common.get_api_cookie_jar())
+    if workspace is None:
+        workspace = skypilot_config.get_active_workspace()
+    response = server_common.make_authenticated_request(
+        'GET', f'/enabled_clouds?workspace={workspace}&expand={expand}')
     return server_common.get_request_id(response)
 
 
 @usage_lib.entrypoint
 @server_common.check_server_healthy_or_start
 @annotations.client_api
-def list_accelerators(gpus_only: bool = True,
-                      name_filter: Optional[str] = None,
-                      region_filter: Optional[str] = None,
-                      quantity_filter: Optional[int] = None,
-                      clouds: Optional[Union[List[str], str]] = None,
-                      all_regions: bool = False,
-                      require_price: bool = True,
-                      case_sensitive: bool = True) -> server_common.RequestId:
+def list_accelerators(
+    gpus_only: bool = True,
+    name_filter: Optional[str] = None,
+    region_filter: Optional[str] = None,
+    quantity_filter: Optional[int] = None,
+    clouds: Optional[Union[List[str], str]] = None,
+    all_regions: bool = False,
+    require_price: bool = True,
+    case_sensitive: bool = True
+) -> server_common.RequestId[Dict[str,
+                                  List['catalog.common.InstanceTypeInfo']]]:
     """Lists the names of all accelerators offered by Sky.
 
     This will include all accelerators offered by Sky, including those
@@ -169,10 +292,8 @@ def list_accelerators(gpus_only: bool = True,
         require_price=require_price,
         case_sensitive=case_sensitive,
     )
-    response = requests.post(
-        f'{server_common.get_server_url()}/list_accelerators',
-        json=json.loads(body.model_dump_json()),
-        cookies=server_common.get_api_cookie_jar())
+    response = server_common.make_authenticated_request(
+        'POST', '/list_accelerators', json=json.loads(body.model_dump_json()))
     return server_common.get_request_id(response)
 
 
@@ -180,12 +301,12 @@ def list_accelerators(gpus_only: bool = True,
 @server_common.check_server_healthy_or_start
 @annotations.client_api
 def list_accelerator_counts(
-        gpus_only: bool = True,
-        name_filter: Optional[str] = None,
-        region_filter: Optional[str] = None,
-        quantity_filter: Optional[int] = None,
-        clouds: Optional[Union[List[str],
-                               str]] = None) -> server_common.RequestId:
+    gpus_only: bool = True,
+    name_filter: Optional[str] = None,
+    region_filter: Optional[str] = None,
+    quantity_filter: Optional[int] = None,
+    clouds: Optional[Union[List[str], str]] = None
+) -> server_common.RequestId[Dict[str, List[float]]]:
     """Lists all accelerators offered by Sky and available counts.
 
     Args:
@@ -203,17 +324,17 @@ def list_accelerator_counts(
             accelerator names mapped to a list of available counts. See usage
             in cli.py.
     """
-    body = payloads.ListAcceleratorsBody(
+    body = payloads.ListAcceleratorCountsBody(
         gpus_only=gpus_only,
         name_filter=name_filter,
         region_filter=region_filter,
         quantity_filter=quantity_filter,
         clouds=clouds,
     )
-    response = requests.post(
-        f'{server_common.get_server_url()}/list_accelerator_counts',
-        json=json.loads(body.model_dump_json()),
-        cookies=server_common.get_api_cookie_jar())
+    response = server_common.make_authenticated_request(
+        'POST',
+        '/list_accelerator_counts',
+        json=json.loads(body.model_dump_json()))
     return server_common.get_request_id(response)
 
 
@@ -224,7 +345,7 @@ def optimize(
     dag: 'sky.Dag',
     minimize: common.OptimizeTarget = common.OptimizeTarget.COST,
     admin_policy_request_options: Optional[admin_policy.RequestOptions] = None
-) -> server_common.RequestId:
+) -> server_common.RequestId['sky.Dag']:
     """Finds the best execution plan for the given DAG.
 
     Args:
@@ -250,12 +371,27 @@ def optimize(
     body = payloads.OptimizeBody(dag=dag_str,
                                  minimize=minimize,
                                  request_options=admin_policy_request_options)
-    response = requests.post(f'{server_common.get_server_url()}/optimize',
-                             json=json.loads(body.model_dump_json()),
-                             cookies=server_common.get_api_cookie_jar())
+    response = server_common.make_authenticated_request(
+        'POST', '/optimize', json=json.loads(body.model_dump_json()))
     return server_common.get_request_id(response)
 
 
+def workspaces() -> server_common.RequestId[Dict[str, Any]]:
+    """Gets the workspaces."""
+    response = server_common.make_authenticated_request('GET', '/workspaces')
+    return server_common.get_request_id(response)
+
+
+def _raise_exception_object_on_client(e: BaseException) -> None:
+    """Raise the exception object on the client."""
+    if env_options.Options.SHOW_DEBUG_INFO.get():
+        stacktrace = getattr(e, 'stacktrace', str(e))
+        logger.error('=== Traceback on SkyPilot API Server ===\n'
+                     f'{stacktrace}')
+    with ux_utils.print_exception_no_traceback():
+        raise e
+
+
 @usage_lib.entrypoint
 @server_common.check_server_healthy_or_start
 @annotations.client_api
@@ -279,29 +415,35 @@ def validate(
             validation. This is only required when a admin policy is in use,
             see: https://docs.skypilot.co/en/latest/cloud-setup/policy.html
     """
+    remote_api_version = versions.get_remote_api_version()
+    # TODO(kevin): remove this in v0.13.0
+    omit_user_specified_yaml = (remote_api_version is None or
+                                remote_api_version < 15)
     for task in dag.tasks:
+        if omit_user_specified_yaml:
+            # pylint: disable=protected-access
+            task._user_specified_yaml = None
         task.expand_and_validate_workdir()
         if not workdir_only:
             task.expand_and_validate_file_mounts()
     dag_str = dag_utils.dump_chain_dag_to_yaml_str(dag)
     body = payloads.ValidateBody(dag=dag_str,
                                  request_options=admin_policy_request_options)
-    response = requests.post(f'{server_common.get_server_url()}/validate',
-                             json=json.loads(body.model_dump_json()),
-                             cookies=server_common.get_api_cookie_jar())
+    response = server_common.make_authenticated_request(
+        'POST', '/validate', json=json.loads(body.model_dump_json()))
     if response.status_code == 400:
-        with ux_utils.print_exception_no_traceback():
-            raise exceptions.deserialize_exception(
-                response.json().get('detail'))
+        _raise_exception_object_on_client(
+            exceptions.deserialize_exception(response.json().get('detail')))
 
 
 @usage_lib.entrypoint
 @server_common.check_server_healthy_or_start
 @annotations.client_api
-def dashboard() -> None:
+def dashboard(starting_page: Optional[str] = None) -> None:
     """Starts the dashboard for SkyPilot."""
     api_server_url = server_common.get_server_url()
-    url = server_common.get_dashboard_url(api_server_url)
+    url = server_common.get_dashboard_url(api_server_url,
+                                          starting_page=starting_page)
     logger.info(f'Opening dashboard in browser: {url}')
     webbrowser.open(url)
 
@@ -309,14 +451,16 @@ def dashboard() -> None:
 @usage_lib.entrypoint
 @server_common.check_server_healthy_or_start
 @annotations.client_api
+@sky_context.contextual
 def launch(
     task: Union['sky.Task', 'sky.Dag'],
     cluster_name: Optional[str] = None,
     retry_until_up: bool = False,
     idle_minutes_to_autostop: Optional[int] = None,
+    wait_for: Optional[autostop_lib.AutostopWaitFor] = None,
     dryrun: bool = False,
     down: bool = False,  # pylint: disable=redefined-outer-name
-    backend: Optional[backends.Backend] = None,
+    backend: Optional['backends.Backend'] = None,
     optimize_target: common.OptimizeTarget = common.OptimizeTarget.COST,
     no_setup: bool = False,
     clone_disk_from: Optional[str] = None,
@@ -327,7 +471,8 @@ def launch(
     _is_launched_by_jobs_controller: bool = False,
     _is_launched_by_sky_serve_controller: bool = False,
     _disable_controller_check: bool = False,
-) -> server_common.RequestId:
+) -> server_common.RequestId[Tuple[Optional[int],
+                                   Optional['backends.ResourceHandle']]]:
     """Launches a cluster or task.
 
     The task's setup and run commands are executed under the task's workdir
@@ -344,7 +489,7 @@ def launch(
             import sky
             task = sky.Task(run='echo hello SkyPilot')
             task.set_resources(
-                sky.Resources(cloud=sky.AWS(), accelerators='V100:4'))
+                sky.Resources(infra='aws', accelerators='V100:4'))
             sky.launch(task, cluster_name='my-cluster')
 
 
@@ -355,18 +500,31 @@ def launch(
         retry_until_up: whether to retry launching the cluster until it is
           up.
         idle_minutes_to_autostop: automatically stop the cluster after this
-          many minute of idleness, i.e., no running or pending jobs in the
-          cluster's job queue. Idleness gets reset whenever setting-up/
-          running/pending jobs are found in the job queue. Setting this
-          flag is equivalent to running ``sky.launch()`` and then
-          ``sky.autostop(idle_minutes=<minutes>)``. If not set, the cluster
-          will not be autostopped.
+            many minute of idleness, i.e., no running or pending jobs in the
+            cluster's job queue. Idleness gets reset whenever setting-up/
+            running/pending jobs are found in the job queue. Setting this
+            flag is equivalent to running
+            ``sky.launch(...)`` and then
+            ``sky.autostop(idle_minutes=<minutes>)``. If set, the autostop
+            config specified in the task' resources will be overridden by
+            this parameter.
+        wait_for: determines the condition for resetting the idleness timer.
+            This option works in conjunction with ``idle_minutes_to_autostop``.
+            Choices:
+
+            1. "jobs_and_ssh" (default) - Wait for in-progress jobs and SSH
+               connections to finish.
+            2. "jobs" - Only wait for in-progress jobs.
+            3. "none" - Wait for nothing; autostop right after
+               ``idle_minutes_to_autostop``.
         dryrun: if True, do not actually launch the cluster.
         down: Tear down the cluster after all jobs finish (successfully or
-          abnormally). If --idle-minutes-to-autostop is also set, the
-          cluster will be torn down after the specified idle time.
-          Note that if errors occur during provisioning/data syncing/setting
-          up, the cluster will not be torn down for debugging purposes.
+            abnormally). If --idle-minutes-to-autostop is also set, the
+            cluster will be torn down after the specified idle time.
+            Note that if errors occur during provisioning/data syncing/setting
+            up, the cluster will not be torn down for debugging purposes. If
+            set, the autostop config specified in the task' resources will be
+            overridden by this parameter.
         backend: backend to use.  If None, use the default backend
           (CloudVMRayBackend).
         optimize_target: target to optimize for. Choices: OptimizeTarget.COST,
@@ -422,35 +580,115 @@ def launch(
             raise NotImplementedError('clone_disk_from is not implemented yet. '
                                       'Please contact the SkyPilot team if you '
                                       'need this feature at slack.skypilot.co.')
+
+    remote_api_version = versions.get_remote_api_version()
+    if wait_for is not None and (remote_api_version is None or
+                                 remote_api_version < 13):
+        logger.warning('wait_for is not supported in your API server. '
+                       'Please upgrade to a newer API server to use it.')
+
     dag = dag_utils.convert_entrypoint_to_dag(task)
+    # Override the autostop config from command line flags to task YAML.
+    for task in dag.tasks:
+        for resource in task.resources:
+            if remote_api_version is None or remote_api_version < 13:
+                # An older server would not recognize the wait_for field
+                # in the schema, so we need to omit it.
+                resource.override_autostop_config(
+                    down=down, idle_minutes=idle_minutes_to_autostop)
+            else:
+                resource.override_autostop_config(
+                    down=down,
+                    idle_minutes=idle_minutes_to_autostop,
+                    wait_for=wait_for)
+            if resource.autostop_config is not None:
+                # For backward-compatibility, get the final autostop config for
+                # admin policy.
+                # TODO(aylei): remove this after 0.12.0
+                down = resource.autostop_config.down
+                idle_minutes_to_autostop = resource.autostop_config.idle_minutes
+
     request_options = admin_policy.RequestOptions(
         cluster_name=cluster_name,
         idle_minutes_to_autostop=idle_minutes_to_autostop,
         down=down,
         dryrun=dryrun)
+    with admin_policy_utils.apply_and_use_config_in_current_request(
+            dag,
+            request_name=request_names.AdminPolicyRequestName.CLUSTER_LAUNCH,
+            request_options=request_options,
+            at_client_side=True) as dag:
+        return _launch(
+            dag,
+            cluster_name,
+            request_options,
+            retry_until_up,
+            idle_minutes_to_autostop,
+            dryrun,
+            down,
+            backend,
+            optimize_target,
+            no_setup,
+            clone_disk_from,
+            fast,
+            _need_confirmation,
+            _is_launched_by_jobs_controller,
+            _is_launched_by_sky_serve_controller,
+            _disable_controller_check,
+        )
+
+
+def _launch(
+    dag: 'sky.Dag',
+    cluster_name: str,
+    request_options: admin_policy.RequestOptions,
+    retry_until_up: bool = False,
+    idle_minutes_to_autostop: Optional[int] = None,
+    dryrun: bool = False,
+    down: bool = False,  # pylint: disable=redefined-outer-name
+    backend: Optional['backends.Backend'] = None,
+    optimize_target: common.OptimizeTarget = common.OptimizeTarget.COST,
+    no_setup: bool = False,
+    clone_disk_from: Optional[str] = None,
+    fast: bool = False,
+    # Internal only:
+    # pylint: disable=invalid-name
+    _need_confirmation: bool = False,
+    _is_launched_by_jobs_controller: bool = False,
+    _is_launched_by_sky_serve_controller: bool = False,
+    _disable_controller_check: bool = False,
+) -> server_common.RequestId[Tuple[Optional[int],
+                                   Optional['backends.ResourceHandle']]]:
+    """Auxiliary function for launch(), refer to launch() for details."""
+
     validate(dag, admin_policy_request_options=request_options)
+    # The flags have been applied to the task YAML and the backward
+    # compatibility of admin policy has been handled. We should no longer use
+    # these flags.
+    del down, idle_minutes_to_autostop
 
     confirm_shown = False
     if _need_confirmation:
         cluster_status = None
         # TODO(SKY-998): we should reduce RTTs before launching the cluster.
-        request_id = status([cluster_name], all_users=True)
-        clusters = get(request_id)
+        status_request_id = status([cluster_name], all_users=True)
+        clusters = get(status_request_id)
         cluster_user_hash = common_utils.get_user_hash()
         cluster_user_hash_str = ''
-        cluster_user_name = getpass.getuser()
+        current_user = common_utils.get_current_user_name()
+        cluster_user_name = current_user
         if not clusters:
             # Show the optimize log before the prompt if the cluster does not
             # exist.
-            request_id = optimize(dag,
-                                  admin_policy_request_options=request_options)
-            stream_and_get(request_id)
+            optimize_request_id = optimize(
+                dag, admin_policy_request_options=request_options)
+            stream_and_get(optimize_request_id)
         else:
             cluster_record = clusters[0]
             cluster_status = cluster_record['status']
             cluster_user_hash = cluster_record['user_hash']
             cluster_user_name = cluster_record['user_name']
-            if cluster_user_name == getpass.getuser():
+            if cluster_user_name == current_user:
                 # Only show the hash if the username is the same as the local
                 # username, to avoid confusion.
                 cluster_user_hash_str = f' (hash: {cluster_user_hash})'
@@ -492,9 +730,7 @@ def launch(
         task=dag_str,
         cluster_name=cluster_name,
         retry_until_up=retry_until_up,
-        idle_minutes_to_autostop=idle_minutes_to_autostop,
         dryrun=dryrun,
-        down=down,
         backend=backend.NAME if backend else None,
         optimize_target=optimize_target,
         no_setup=no_setup,
@@ -507,12 +743,8 @@ def launch(
             _is_launched_by_sky_serve_controller),
         disable_controller_check=_disable_controller_check,
     )
-    response = requests.post(
-        f'{server_common.get_server_url()}/launch',
-        json=json.loads(body.model_dump_json()),
-        timeout=5,
-        cookies=server_common.get_api_cookie_jar(),
-    )
+    response = server_common.make_authenticated_request(
+        'POST', '/launch', json=json.loads(body.model_dump_json()), timeout=5)
     return server_common.get_request_id(response)
 
 
@@ -524,8 +756,9 @@ def exec(  # pylint: disable=redefined-builtin
     cluster_name: Optional[str] = None,
     dryrun: bool = False,
     down: bool = False,  # pylint: disable=redefined-outer-name
-    backend: Optional[backends.Backend] = None,
-) -> server_common.RequestId:
+    backend: Optional['backends.Backend'] = None,
+) -> server_common.RequestId[Tuple[Optional[int],
+                                   Optional['backends.ResourceHandle']]]:
     """Executes a task on an existing cluster.
 
     This function performs two actions:
@@ -591,23 +824,49 @@ def exec(  # pylint: disable=redefined-builtin
         backend=backend.NAME if backend else None,
     )
 
-    response = requests.post(
-        f'{server_common.get_server_url()}/exec',
-        json=json.loads(body.model_dump_json()),
-        timeout=5,
-        cookies=server_common.get_api_cookie_jar(),
-    )
+    response = server_common.make_authenticated_request(
+        'POST', '/exec', json=json.loads(body.model_dump_json()), timeout=5)
     return server_common.get_request_id(response)
 
 
-@usage_lib.entrypoint
-@server_common.check_server_healthy_or_start
-@annotations.client_api
+@typing.overload
+def tail_logs(
+        cluster_name: str,
+        job_id: Optional[int],
+        follow: bool,
+        tail: int = 0,
+        output_stream: Optional['io.TextIOBase'] = None,
+        *,  # keyword only separator
+        preload_content: Literal[True] = True) -> int:
+    ...
+
+
+@typing.overload
 def tail_logs(cluster_name: str,
               job_id: Optional[int],
               follow: bool,
               tail: int = 0,
-              output_stream: Optional['io.TextIOBase'] = None) -> int:
+              output_stream: None = None,
+              *,
+              preload_content: Literal[False]) -> Iterator[Optional[str]]:
+    ...
+
+
+# TODO(aylei): when retry logs request, there will be duplicated log entries.
+# We should fix this.
+@usage_lib.entrypoint
+@server_common.check_server_healthy_or_start
+@annotations.client_api
+@rest.retry_transient_errors()
+def tail_logs(
+    cluster_name: str,
+    job_id: Optional[int],
+    follow: bool,
+    tail: int = 0,
+    output_stream: Optional['io.TextIOBase'] = None,
+    *,  # keyword only separator
+    preload_content: bool = True
+) -> Union[int, Iterator[Optional[str]]]:
     """Tails the logs of a job.
 
     Args:
@@ -617,12 +876,21 @@ def tail_logs(cluster_name: str,
             immediately.
         tail: if > 0, tail the last N lines of the logs.
         output_stream: the stream to write the logs to. If None, print to the
-            console.
+            console. Cannot be used with preload_content=False.
+        preload_content: if False, returns an Iterator[str | None] containing
+            the logs without the function blocking on the retrieval of entire
+            log. Iterator returns None when the log has been completely
+            streamed. Default True. Cannot be used with output_stream.
 
     Returns:
-        Exit code based on success or failure of the job. 0 if success,
-        100 if the job failed. See exceptions.JobExitCode for possible exit
-        codes.
+        If preload_content is True:
+            Exit code based on success or failure of the job. 0 if success,
+            100 if the job failed. See exceptions.JobExitCode for possible exit
+            codes.
+        If preload_content is False:
+            Iterator[str | None] containing the logs without the function
+            blocking on the retrieval of entire log. Iterator returns None
+            when the log has been completely streamed.
 
     Request Raises:
         ValueError: if arguments are invalid or the cluster is not supported.
@@ -635,21 +903,110 @@ def tail_logs(cluster_name: str,
         sky.exceptions.CloudUserIdentityError: if we fail to get the current
           user identity.
     """
+    if output_stream is not None and not preload_content:
+        raise ValueError(
+            'output_stream cannot be specified when preload_content is False')
+
     body = payloads.ClusterJobBody(
         cluster_name=cluster_name,
         job_id=job_id,
         follow=follow,
         tail=tail,
     )
-    response = requests.post(
-        f'{server_common.get_server_url()}/logs',
+    response = server_common.make_authenticated_request(
+        'POST',
+        '/logs',
         json=json.loads(body.model_dump_json()),
         stream=True,
         timeout=(client_common.API_SERVER_REQUEST_CONNECTION_TIMEOUT_SECONDS,
-                 None),
-        cookies=server_common.get_api_cookie_jar())
-    request_id = server_common.get_request_id(response)
-    return stream_response(request_id, response, output_stream)
+                 None))
+    request_id: server_common.RequestId[int] = server_common.get_request_id(
+        response)
+    if preload_content:
+        # Log request is idempotent when tail is 0, thus can resume previous
+        # streaming point on retry.
+        return stream_response(request_id=request_id,
+                               response=response,
+                               output_stream=output_stream,
+                               resumable=(tail == 0))
+    else:
+        return rich_utils.decode_rich_status(response)
+
+
+@usage_lib.entrypoint
+@server_common.check_server_healthy_or_start
+@versions.minimal_api_version(17)
+@annotations.client_api
+@rest.retry_transient_errors()
+def tail_provision_logs(cluster_name: str,
+                        worker: Optional[int] = None,
+                        follow: bool = True,
+                        tail: int = 0,
+                        output_stream: Optional['io.TextIOBase'] = None) -> int:
+    """Tails the provisioning logs (provision.log) for a cluster.
+
+    Args:
+        cluster_name: name of the cluster.
+        worker: worker id in multi-node cluster.
+             If None, stream the logs of the head node.
+        follow: follow the logs.
+        tail: lines from end to tail.
+        output_stream: optional stream to write logs.
+    Returns:
+        Exit code 0 on streaming success; raises on HTTP error.
+    """
+    body = payloads.ProvisionLogsBody(cluster_name=cluster_name)
+
+    if worker is not None:
+        remote_api_version = versions.get_remote_api_version()
+        if remote_api_version is not None and remote_api_version >= 21:
+            if worker < 1:
+                raise ValueError('Worker must be a positive integer.')
+            body.worker = worker
+        else:
+            raise exceptions.APINotSupportedError(
+                'Worker node provision logs are not supported in your API '
+                'server. Please upgrade to a newer API server to use it.')
+    params = {
+        'follow': str(follow).lower(),
+        'tail': tail,
+    }
+
+    response = server_common.make_authenticated_request(
+        'POST',
+        '/provision_logs',
+        json=json.loads(body.model_dump_json()),
+        params=params,
+        stream=True,
+        timeout=(client_common.API_SERVER_REQUEST_CONNECTION_TIMEOUT_SECONDS,
+                 None))
+    # Check for HTTP errors before streaming the response
+    if response.status_code != 200:
+        with ux_utils.print_exception_no_traceback():
+            raise exceptions.CommandError(response.status_code,
+                                          'tail_provision_logs',
+                                          'Failed to stream provision logs',
+                                          response.text)
+
+    # Log request is idempotent when tail is 0, thus can resume previous
+    # streaming point on retry.
+    # request_id=None here because /provision_logs does not create an async
+    # request. Instead, it streams a plain file from the server. This does NOT
+    # violate the stream_response doc warning about None in multi-user
+    # environments: we are not asking stream_response to select "the latest
+    # request". We already have the HTTP response to stream; request_id=None
+    # merely disables the follow-up GET. It is also necessary for --no-follow
+    # to return cleanly after printing the tailed lines. If we provided a
+    # non-None request_id here, the get(request_id) in stream_response(
+    # would fail since /provision_logs does not create a request record.
+    # By virtue of this, we set get_result to False to block get() from
+    # running.
+    stream_response(request_id=None,
+                    response=response,
+                    output_stream=output_stream,
+                    resumable=(tail == 0),
+                    get_result=False)
+    return 0
 
 
 @usage_lib.entrypoint
@@ -683,11 +1040,11 @@ def download_logs(cluster_name: str,
         cluster_name=cluster_name,
         job_ids=job_ids,
     )
-    response = requests.post(f'{server_common.get_server_url()}/download_logs',
-                             json=json.loads(body.model_dump_json()),
-                             cookies=server_common.get_api_cookie_jar())
-    job_id_remote_path_dict = stream_and_get(
-        server_common.get_request_id(response))
+    response = server_common.make_authenticated_request(
+        'POST', '/download_logs', json=json.loads(body.model_dump_json()))
+    request_id: server_common.RequestId[Dict[
+        str, str]] = server_common.get_request_id(response)
+    job_id_remote_path_dict = stream_and_get(request_id)
     remote2local_path_dict = client_common.download_logs_from_api_server(
         job_id_remote_path_dict.values())
     return {
@@ -702,10 +1059,11 @@ def download_logs(cluster_name: str,
 def start(
     cluster_name: str,
     idle_minutes_to_autostop: Optional[int] = None,
+    wait_for: Optional[autostop_lib.AutostopWaitFor] = None,
     retry_until_up: bool = False,
     down: bool = False,  # pylint: disable=redefined-outer-name
     force: bool = False,
-) -> server_common.RequestId:
+) -> server_common.RequestId['backends.CloudVmRayResourceHandle']:
     """Restart a cluster.
 
     If a cluster is previously stopped (status is STOPPED) or failed in
@@ -728,6 +1086,15 @@ def start(
             flag is equivalent to running ``sky.launch()`` and then
             ``sky.autostop(idle_minutes=<minutes>)``. If not set, the
             cluster will not be autostopped.
+        wait_for: determines the condition for resetting the idleness timer.
+            This option works in conjunction with ``idle_minutes_to_autostop``.
+            Choices:
+
+            1. "jobs_and_ssh" (default) - Wait for in-progress jobs and SSH
+               connections to finish.
+            2. "jobs" - Only wait for in-progress jobs.
+            3. "none" - Wait for nothing; autostop right after
+               ``idle_minutes_to_autostop``.
         retry_until_up: whether to retry launching the cluster until it is
             up.
         down: Autodown the cluster: tear down the cluster after specified
@@ -756,26 +1123,30 @@ def start(
         sky.exceptions.ClusterOwnerIdentitiesMismatchError: if the cluster to
             restart was launched by a different user.
     """
+    remote_api_version = versions.get_remote_api_version()
+    if wait_for is not None and (remote_api_version is None or
+                                 remote_api_version < 13):
+        logger.warning('wait_for is not supported in your API server. '
+                       'Please upgrade to a newer API server to use it.')
+
     body = payloads.StartBody(
         cluster_name=cluster_name,
         idle_minutes_to_autostop=idle_minutes_to_autostop,
+        wait_for=wait_for,
         retry_until_up=retry_until_up,
         down=down,
         force=force,
     )
-    response = requests.post(
-        f'{server_common.get_server_url()}/start',
-        json=json.loads(body.model_dump_json()),
-        timeout=5,
-        cookies=server_common.get_api_cookie_jar(),
-    )
+    response = server_common.make_authenticated_request(
+        'POST', '/start', json=json.loads(body.model_dump_json()), timeout=5)
     return server_common.get_request_id(response)
 
 
 @usage_lib.entrypoint
 @server_common.check_server_healthy_or_start
 @annotations.client_api
-def down(cluster_name: str, purge: bool = False) -> server_common.RequestId:
+def down(cluster_name: str,
+         purge: bool = False) -> server_common.RequestId[None]:
     """Tears down a cluster.
 
     Tearing down a cluster will delete all associated resources (all billing
@@ -809,19 +1180,16 @@ def down(cluster_name: str, purge: bool = False) -> server_common.RequestId:
         cluster_name=cluster_name,
         purge=purge,
     )
-    response = requests.post(
-        f'{server_common.get_server_url()}/down',
-        json=json.loads(body.model_dump_json()),
-        timeout=5,
-        cookies=server_common.get_api_cookie_jar(),
-    )
+    response = server_common.make_authenticated_request(
+        'POST', '/down', json=json.loads(body.model_dump_json()), timeout=5)
     return server_common.get_request_id(response)
 
 
 @usage_lib.entrypoint
 @server_common.check_server_healthy_or_start
 @annotations.client_api
-def stop(cluster_name: str, purge: bool = False) -> server_common.RequestId:
+def stop(cluster_name: str,
+         purge: bool = False) -> server_common.RequestId[None]:
     """Stops a cluster.
 
     Data on attached disks is not lost when a cluster is stopped.  Billing for
@@ -858,12 +1226,8 @@ def stop(cluster_name: str, purge: bool = False) -> server_common.RequestId:
         cluster_name=cluster_name,
         purge=purge,
     )
-    response = requests.post(
-        f'{server_common.get_server_url()}/stop',
-        json=json.loads(body.model_dump_json()),
-        timeout=5,
-        cookies=server_common.get_api_cookie_jar(),
-    )
+    response = server_common.make_authenticated_request(
+        'POST', '/stop', json=json.loads(body.model_dump_json()), timeout=5)
     return server_common.get_request_id(response)
 
 
@@ -871,10 +1235,11 @@ def stop(cluster_name: str, purge: bool = False) -> server_common.RequestId:
 @server_common.check_server_healthy_or_start
 @annotations.client_api
 def autostop(
-    cluster_name: str,
-    idle_minutes: int,
-    down: bool = False  # pylint: disable=redefined-outer-name
-) -> server_common.RequestId:
+        cluster_name: str,
+        idle_minutes: int,
+        wait_for: Optional[autostop_lib.AutostopWaitFor] = None,
+        down: bool = False,  # pylint: disable=redefined-outer-name
+) -> server_common.RequestId[None]:
     """Schedules an autostop/autodown for a cluster.
 
     Autostop/autodown will automatically stop or teardown a cluster when it
@@ -904,6 +1269,14 @@ def autostop(
         idle_minutes: the number of minutes of idleness (no pending/running
             jobs) after which the cluster will be stopped automatically. Setting
             to a negative number cancels any autostop/autodown setting.
+        wait_for: determines the condition for resetting the idleness timer.
+            This option works in conjunction with ``idle_minutes``.
+            Choices:
+
+            1. "jobs_and_ssh" (default) - Wait for in-progress jobs and SSH
+               connections to finish.
+            2. "jobs" - Only wait for in-progress jobs.
+            3. "none" - Wait for nothing; autostop right after ``idle_minutes``.
         down: if true, use autodown (tear down the cluster; non-restartable),
             rather than autostop (restartable).
 
@@ -923,26 +1296,31 @@ def autostop(
         sky.exceptions.CloudUserIdentityError: if we fail to get the current
             user identity.
     """
+    remote_api_version = versions.get_remote_api_version()
+    if wait_for is not None and (remote_api_version is None or
+                                 remote_api_version < 13):
+        logger.warning('wait_for is not supported in your API server. '
+                       'Please upgrade to a newer API server to use it.')
+
     body = payloads.AutostopBody(
         cluster_name=cluster_name,
         idle_minutes=idle_minutes,
+        wait_for=wait_for,
         down=down,
     )
-    response = requests.post(
-        f'{server_common.get_server_url()}/autostop',
-        json=json.loads(body.model_dump_json()),
-        timeout=5,
-        cookies=server_common.get_api_cookie_jar(),
-    )
+    response = server_common.make_authenticated_request(
+        'POST', '/autostop', json=json.loads(body.model_dump_json()), timeout=5)
     return server_common.get_request_id(response)
 
 
 @usage_lib.entrypoint
 @server_common.check_server_healthy_or_start
 @annotations.client_api
-def queue(cluster_name: str,
-          skip_finished: bool = False,
-          all_users: bool = False) -> server_common.RequestId:
+def queue(
+    cluster_name: str,
+    skip_finished: bool = False,
+    all_users: bool = False
+) -> server_common.RequestId[List[responses.ClusterJobRecord]]:
     """Gets the job queue of a cluster.
 
     Args:
@@ -955,8 +1333,8 @@ def queue(cluster_name: str,
         The request ID of the queue request.
 
     Request Returns:
-        job_records (List[Dict[str, Any]]): A list of dicts for each job in the
-            queue.
+        job_records (List[responses.ClusterJobRecord]): A list of job records
+            for each job in the queue.
 
             .. code-block:: python
 
@@ -991,17 +1369,19 @@ def queue(cluster_name: str,
         skip_finished=skip_finished,
         all_users=all_users,
     )
-    response = requests.post(f'{server_common.get_server_url()}/queue',
-                             json=json.loads(body.model_dump_json()),
-                             cookies=server_common.get_api_cookie_jar())
+    response = server_common.make_authenticated_request(
+        'POST', '/queue', json=json.loads(body.model_dump_json()))
     return server_common.get_request_id(response)
 
 
 @usage_lib.entrypoint
 @server_common.check_server_healthy_or_start
 @annotations.client_api
-def job_status(cluster_name: str,
-               job_ids: Optional[List[int]] = None) -> server_common.RequestId:
+def job_status(
+    cluster_name: str,
+    job_ids: Optional[List[int]] = None
+) -> server_common.RequestId[Dict[Optional[int],
+                                  Optional['job_lib.JobStatus']]]:
     """Gets the status of jobs on a cluster.
 
     Args:
@@ -1033,9 +1413,8 @@ def job_status(cluster_name: str,
         cluster_name=cluster_name,
         job_ids=job_ids,
     )
-    response = requests.post(f'{server_common.get_server_url()}/job_status',
-                             json=json.loads(body.model_dump_json()),
-                             cookies=server_common.get_api_cookie_jar())
+    response = server_common.make_authenticated_request(
+        'POST', '/job_status', json=json.loads(body.model_dump_json()))
     return server_common.get_request_id(response)
 
 
@@ -1049,7 +1428,7 @@ def cancel(
     job_ids: Optional[List[int]] = None,
     # pylint: disable=invalid-name
     _try_cancel_if_cluster_is_init: bool = False
-) -> server_common.RequestId:
+) -> server_common.RequestId[None]:
     """Cancels jobs on a cluster.
 
     Args:
@@ -1087,9 +1466,8 @@ def cancel(
         job_ids=job_ids,
         try_cancel_if_cluster_is_init=_try_cancel_if_cluster_is_init,
     )
-    response = requests.post(f'{server_common.get_server_url()}/cancel',
-                             json=json.loads(body.model_dump_json()),
-                             cookies=server_common.get_api_cookie_jar())
+    response = server_common.make_authenticated_request(
+        'POST', '/cancel', json=json.loads(body.model_dump_json()))
     return server_common.get_request_id(response)
 
 
@@ -1100,7 +1478,10 @@ def status(
     cluster_names: Optional[List[str]] = None,
     refresh: common.StatusRefreshMode = common.StatusRefreshMode.NONE,
     all_users: bool = False,
-) -> server_common.RequestId:
+    *,
+    _include_credentials: bool = False,
+    _summary_response: bool = False,
+) -> server_common.RequestId[List[responses.StatusResponse]]:
     """Gets cluster statuses.
 
     If cluster_names is given, return those clusters. Otherwise, return all
@@ -1148,6 +1529,8 @@ def status(
             provider(s).
         all_users: whether to include all users' clusters. By default, only
             the current user's clusters are included.
+        _include_credentials: (internal) whether to include cluster ssh
+            credentials in the response (default: False).
 
     Returns:
         The request ID of the status request.
@@ -1182,10 +1565,11 @@ def status(
         cluster_names=cluster_names,
         refresh=refresh,
         all_users=all_users,
+        include_credentials=_include_credentials,
+        summary_response=_summary_response,
     )
-    response = requests.post(f'{server_common.get_server_url()}/status',
-                             json=json.loads(body.model_dump_json()),
-                             cookies=server_common.get_api_cookie_jar())
+    response = server_common.make_authenticated_request(
+        'POST', '/status', json=json.loads(body.model_dump_json()))
     return server_common.get_request_id(response)
 
 
@@ -1193,10 +1577,19 @@ def status(
 @server_common.check_server_healthy_or_start
 @annotations.client_api
 def endpoints(
-        cluster: str,
-        port: Optional[Union[int, str]] = None) -> server_common.RequestId:
+    cluster: str,
+    port: Optional[Union[int, str]] = None
+) -> server_common.RequestId[Dict[int, str]]:
     """Gets the endpoint for a given cluster and port number (endpoint).
 
+    Example:
+        .. code-block:: python
+
+            import sky
+            request_id = sky.endpoints('test-cluster')
+            sky.get(request_id)
+
+
     Args:
         cluster: The name of the cluster.
         port: The port number to get the endpoint for. If None, endpoints
@@ -1206,8 +1599,9 @@ def endpoints(
         The request ID of the endpoints request.
 
     Request Returns:
-        A dictionary of port numbers to endpoints. If port is None,
-        the dictionary will contain all ports:endpoints exposed on the cluster.
+        A dictionary of port numbers to endpoints.
+        If port is None, the dictionary contains all
+            ports:endpoints exposed on the cluster.
 
     Request Raises:
         ValueError: if the cluster is not UP or the endpoint is not exposed.
@@ -1218,16 +1612,17 @@ def endpoints(
         cluster=cluster,
         port=port,
     )
-    response = requests.post(f'{server_common.get_server_url()}/endpoints',
-                             json=json.loads(body.model_dump_json()),
-                             cookies=server_common.get_api_cookie_jar())
+    response = server_common.make_authenticated_request(
+        'POST', '/endpoints', json=json.loads(body.model_dump_json()))
     return server_common.get_request_id(response)
 
 
 @usage_lib.entrypoint
 @server_common.check_server_healthy_or_start
 @annotations.client_api
-def cost_report() -> server_common.RequestId:  # pylint: disable=redefined-builtin
+def cost_report(
+    days: Optional[int] = None
+) -> server_common.RequestId[List[Dict[str, Any]]]:  # pylint: disable=redefined-builtin
     """Gets all cluster cost reports, including those that have been downed.
 
     The estimated cost column indicates price for the cluster based on the type
@@ -1237,6 +1632,10 @@ def cost_report() -> server_common.RequestId:  # pylint: disable=redefined-built
     cache of the cluster status, and may not be accurate for the cluster with
     autostop/use_spot set or terminated/stopped on the cloud console.
 
+    Args:
+        days: The number of days to get the cost report for. If not provided,
+            the default is 30 days.
+
     Returns:
         The request ID of the cost report request.
 
@@ -1258,8 +1657,9 @@ def cost_report() -> server_common.RequestId:  # pylint: disable=redefined-built
               'total_cost': (float) cost given resources and usage intervals,
             }
     """
-    response = requests.get(f'{server_common.get_server_url()}/cost_report',
-                            cookies=server_common.get_api_cookie_jar())
+    body = payloads.CostReportBody(days=days)
+    response = server_common.make_authenticated_request(
+        'POST', '/cost_report', json=json.loads(body.model_dump_json()))
     return server_common.get_request_id(response)
 
 
@@ -1267,36 +1667,24 @@ def cost_report() -> server_common.RequestId:  # pylint: disable=redefined-built
 @usage_lib.entrypoint
 @server_common.check_server_healthy_or_start
 @annotations.client_api
-def storage_ls() -> server_common.RequestId:
+def storage_ls() -> server_common.RequestId[List[responses.StorageRecord]]:
     """Gets the storages.
 
     Returns:
         The request ID of the storage list request.
 
     Request Returns:
-        storage_records (List[Dict[str, Any]]): A list of dicts, with each dict
-            containing the information of a storage.
-
-            .. code-block:: python
-
-                {
-                    'name': (str) storage name,
-                    'launched_at': (int) timestamp of creation,
-                    'store': (List[sky.StoreType]) storage type,
-                    'last_use': (int) timestamp of last use,
-                    'status': (sky.StorageStatus) storage status,
-                }
-        ]
+        storage_records (List[responses.StorageRecord]):
+            A list of storage records.
     """
-    response = requests.get(f'{server_common.get_server_url()}/storage/ls',
-                            cookies=server_common.get_api_cookie_jar())
+    response = server_common.make_authenticated_request('GET', '/storage/ls')
     return server_common.get_request_id(response)
 
 
 @usage_lib.entrypoint
 @server_common.check_server_healthy_or_start
 @annotations.client_api
-def storage_delete(name: str) -> server_common.RequestId:
+def storage_delete(name: str) -> server_common.RequestId[None]:
     """Deletes a storage.
 
     Args:
@@ -1312,9 +1700,8 @@ def storage_delete(name: str) -> server_common.RequestId:
         ValueError: If the storage does not exist.
     """
     body = payloads.StorageBody(name=name)
-    response = requests.post(f'{server_common.get_server_url()}/storage/delete',
-                             json=json.loads(body.model_dump_json()),
-                             cookies=server_common.get_api_cookie_jar())
+    response = server_common.make_authenticated_request(
+        'POST', '/storage/delete', json=json.loads(body.model_dump_json()))
     return server_common.get_request_id(response)
 
 
@@ -1325,12 +1712,8 @@ def storage_delete(name: str) -> server_common.RequestId:
 @server_common.check_server_healthy_or_start
 @annotations.client_api
 def local_up(gpus: bool,
-             ips: Optional[List[str]],
-             ssh_user: Optional[str],
-             ssh_key: Optional[str],
-             cleanup: bool,
-             context_name: Optional[str] = None,
-             password: Optional[str] = None) -> server_common.RequestId:
+             name: Optional[str] = None,
+             port_start: Optional[int] = None) -> server_common.RequestId[None]:
     """Launches a Kubernetes cluster on local machines.
 
     Returns:
@@ -1341,36 +1724,151 @@ def local_up(gpus: bool,
     # TODO: move this check to server.
     if not server_common.is_api_server_local():
         with ux_utils.print_exception_no_traceback():
-            raise ValueError(
-                'sky local up is only supported when running SkyPilot locally.')
-
-    body = payloads.LocalUpBody(gpus=gpus,
-                                ips=ips,
-                                ssh_user=ssh_user,
-                                ssh_key=ssh_key,
-                                cleanup=cleanup,
-                                context_name=context_name,
-                                password=password)
-    response = requests.post(f'{server_common.get_server_url()}/local_up',
-                             json=json.loads(body.model_dump_json()),
-                             cookies=server_common.get_api_cookie_jar())
+            raise ValueError('`sky local up` is only supported when '
+                             'running SkyPilot locally.')
+
+    body = payloads.LocalUpBody(gpus=gpus, name=name, port_start=port_start)
+    response = server_common.make_authenticated_request(
+        'POST', '/local_up', json=json.loads(body.model_dump_json()))
     return server_common.get_request_id(response)
 
 
 @usage_lib.entrypoint
 @server_common.check_server_healthy_or_start
 @annotations.client_api
-def local_down() -> server_common.RequestId:
+def local_down(name: Optional[str]) -> server_common.RequestId[None]:
     """Tears down the Kubernetes cluster started by local_up."""
     # We do not allow local up when the API server is running remotely since it
     # will modify the kubeconfig.
     # TODO: move this check to remote server.
     if not server_common.is_api_server_local():
         with ux_utils.print_exception_no_traceback():
-            raise ValueError('sky local down is only supported when running '
+            raise ValueError('`sky local down` is only supported when running '
                              'SkyPilot locally.')
-    response = requests.post(f'{server_common.get_server_url()}/local_down',
-                             cookies=server_common.get_api_cookie_jar())
+
+    body = payloads.LocalDownBody(name=name)
+    response = server_common.make_authenticated_request(
+        'POST', '/local_down', json=json.loads(body.model_dump_json()))
+    return server_common.get_request_id(response)
+
+
+def _update_remote_ssh_node_pools(file: str,
+                                  infra: Optional[str] = None) -> None:
+    """Update the SSH node pools on the remote server.
+
+    This function will also upload the local SSH key to the remote server, and
+    replace the file path to the remote SSH key file path.
+
+    Args:
+        file: The path to the local SSH node pools config file.
+        infra: The name of the cluster configuration in the local SSH node
+            pools config file. If None, all clusters in the file are updated.
+    """
+    file = os.path.expanduser(file)
+    if not os.path.exists(file):
+        with ux_utils.print_exception_no_traceback():
+            raise ValueError(
+                f'SSH Node Pool config file {file} does not exist. '
+                'Please check if the file exists and the path is correct.')
+    config = ssh_utils.load_ssh_targets(file)
+    config = ssh_utils.get_cluster_config(config, infra)
+    pools_config = {}
+    for name, pool_config in config.items():
+        hosts_info = ssh_utils.prepare_hosts_info(
+            name, pool_config, upload_ssh_key_func=_upload_ssh_key_and_wait)
+        pools_config[name] = {'hosts': hosts_info}
+    server_common.make_authenticated_request('POST',
+                                             '/ssh_node_pools',
+                                             json=pools_config)
+
+
+def _upload_ssh_key_and_wait(key_name: str, key_file_path: str) -> str:
+    """Upload the SSH key to the remote server and wait for the key to be
+    uploaded.
+
+    Args:
+        key_name: The name of the SSH key.
+        key_file_path: The path to the local SSH key file.
+
+    Returns:
+        The path for the remote SSH key file on the API server.
+    """
+    if not os.path.exists(os.path.expanduser(key_file_path)):
+        with ux_utils.print_exception_no_traceback():
+            raise ValueError(f'SSH key file not found: {key_file_path}')
+
+    with open(os.path.expanduser(key_file_path), 'rb') as key_file:
+        response = server_common.make_authenticated_request(
+            'POST',
+            '/ssh_node_pools/keys',
+            files={
+                'key_file': (key_name, key_file, 'application/octet-stream')
+            },
+            data={'key_name': key_name},
+            cookies=server_common.get_api_cookie_jar())
+
+    return response.json()['key_path']
+
+
+@usage_lib.entrypoint
+@server_common.check_server_healthy_or_start
+@annotations.client_api
+def ssh_up(infra: Optional[str] = None,
+           file: Optional[str] = None) -> server_common.RequestId[None]:
+    """Deploys the SSH Node Pools defined in ~/.sky/ssh_targets.yaml.
+
+    Args:
+        infra: Name of the cluster configuration in ssh_targets.yaml.
+            If None, the first cluster in the file is used.
+        file: Name of the ssh node pool configuration file to use. If
+            None, the default path, ~/.sky/ssh_node_pools.yaml is used.
+
+    Returns:
+        request_id: The request ID of the SSH cluster deployment request.
+    """
+    if file is not None:
+        _update_remote_ssh_node_pools(file, infra)
+
+    # Use SSH node pools router endpoint
+    body = payloads.SSHUpBody(infra=infra, cleanup=False)
+    if infra is not None:
+        # Call the specific pool deployment endpoint
+        response = server_common.make_authenticated_request(
+            'POST', f'/ssh_node_pools/{infra}/deploy')
+    else:
+        # Call the general deployment endpoint
+        response = server_common.make_authenticated_request(
+            'POST',
+            '/ssh_node_pools/deploy',
+            json=json.loads(body.model_dump_json()))
+    return server_common.get_request_id(response)
+
+
+@usage_lib.entrypoint
+@server_common.check_server_healthy_or_start
+@annotations.client_api
+def ssh_down(infra: Optional[str] = None) -> server_common.RequestId[None]:
+    """Tears down a Kubernetes cluster on SSH targets.
+
+    Args:
+        infra: Name of the cluster configuration in ssh_targets.yaml.
+            If None, the first cluster in the file is used.
+
+    Returns:
+        request_id: The request ID of the SSH cluster teardown request.
+    """
+    # Use SSH node pools router endpoint
+    body = payloads.SSHUpBody(infra=infra, cleanup=True)
+    if infra is not None:
+        # Call the specific pool down endpoint
+        response = server_common.make_authenticated_request(
+            'POST', f'/ssh_node_pools/{infra}/down')
+    else:
+        # Call the general down endpoint
+        response = server_common.make_authenticated_request(
+            'POST',
+            '/ssh_node_pools/down',
+            json=json.loads(body.model_dump_json()))
     return server_common.get_request_id(response)
 
 
@@ -1378,9 +1876,12 @@ def local_down() -> server_common.RequestId:
 @server_common.check_server_healthy_or_start
 @annotations.client_api
 def realtime_kubernetes_gpu_availability(
-        context: Optional[str] = None,
-        name_filter: Optional[str] = None,
-        quantity_filter: Optional[int] = None) -> server_common.RequestId:
+    context: Optional[str] = None,
+    name_filter: Optional[str] = None,
+    quantity_filter: Optional[int] = None,
+    is_ssh: Optional[bool] = None
+) -> server_common.RequestId[List[Tuple[
+        str, List['models.RealtimeGpuAvailability']]]]:
     """Gets the real-time Kubernetes GPU availability.
 
     Returns:
@@ -1390,12 +1891,12 @@ def realtime_kubernetes_gpu_availability(
         context=context,
         name_filter=name_filter,
         quantity_filter=quantity_filter,
+        is_ssh=is_ssh,
     )
-    response = requests.post(
-        f'{server_common.get_server_url()}/'
-        'realtime_kubernetes_gpu_availability',
-        json=json.loads(body.model_dump_json()),
-        cookies=server_common.get_api_cookie_jar())
+    response = server_common.make_authenticated_request(
+        'POST',
+        '/realtime_kubernetes_gpu_availability',
+        json=json.loads(body.model_dump_json()))
     return server_common.get_request_id(response)
 
 
@@ -1403,7 +1904,8 @@ def realtime_kubernetes_gpu_availability(
 @server_common.check_server_healthy_or_start
 @annotations.client_api
 def kubernetes_node_info(
-        context: Optional[str] = None) -> server_common.RequestId:
+    context: Optional[str] = None
+) -> server_common.RequestId['models.KubernetesNodesInfo']:
     """Gets the resource information for all the nodes in the cluster.
 
     Currently only GPU resources are supported. The function returns the total
@@ -1424,18 +1926,22 @@ def kubernetes_node_info(
             information.
     """
     body = payloads.KubernetesNodeInfoRequestBody(context=context)
-    response = requests.post(
-        f'{server_common.get_server_url()}/kubernetes_node_info',
-        json=json.loads(body.model_dump_json()),
-        cookies=server_common.get_api_cookie_jar())
+    response = server_common.make_authenticated_request(
+        'POST',
+        '/kubernetes_node_info',
+        json=json.loads(body.model_dump_json()))
     return server_common.get_request_id(response)
 
 
 @usage_lib.entrypoint
 @server_common.check_server_healthy_or_start
 @annotations.client_api
-def status_kubernetes() -> server_common.RequestId:
-    """Gets all SkyPilot clusters and jobs in the Kubernetes cluster.
+def status_kubernetes() -> server_common.RequestId[
+    Tuple[List['kubernetes_utils.KubernetesSkyPilotClusterInfoPayload'],
+          List['kubernetes_utils.KubernetesSkyPilotClusterInfoPayload'],
+          List[responses.ManagedJobRecord], Optional[str]]]:
+    """[Experimental] Gets all SkyPilot clusters and jobs
+    in the Kubernetes cluster.
 
     Managed jobs and services are also included in the clusters returned.
     The caller must parse the controllers to identify which clusters are run
@@ -1455,21 +1961,24 @@ def status_kubernetes() -> server_common.RequestId:
             dictionary job info, see jobs.queue_from_kubernetes_pod for details.
         - context: Kubernetes context used to fetch the cluster information.
     """
-    response = requests.get(
-        f'{server_common.get_server_url()}/status_kubernetes',
-        cookies=server_common.get_api_cookie_jar())
+    response = server_common.make_authenticated_request('GET',
+                                                        '/status_kubernetes')
     return server_common.get_request_id(response)
 
 
 # === API request APIs ===
 @usage_lib.entrypoint
-@server_common.check_server_healthy_or_start
 @annotations.client_api
-def get(request_id: str) -> Any:
+def get(request_id: server_common.RequestId[T]) -> T:
     """Waits for and gets the result of a request.
 
+    This function will not check the server health since /api/get is typically
+    not the first API call in an SDK session and checking the server health
+    may cause GET /api/get being sent to a restarted API server.
+
     Args:
-        request_id: The request ID of the request to get.
+        request_id: The request ID of the request to get. May be a full request
+            ID or a prefix.
 
     Returns:
         The ``Request Returns`` of the specified request. See the documentation
@@ -1480,19 +1989,20 @@ def get(request_id: str) -> Any:
             see ``Request Raises`` in the documentation of the specific requests
             above.
     """
-    response = requests.get(
-        f'{server_common.get_server_url()}/api/get?request_id={request_id}',
+    response = server_common.make_authenticated_request(
+        'GET',
+        f'/api/get?request_id={request_id}',
+        retry=False,
         timeout=(client_common.API_SERVER_REQUEST_CONNECTION_TIMEOUT_SECONDS,
-                 None),
-        cookies=server_common.get_api_cookie_jar())
+                 None))
     request_task = None
     if response.status_code == 200:
         request_task = requests_lib.Request.decode(
-            requests_lib.RequestPayload(**response.json()))
+            payloads.RequestPayload(**response.json()))
     elif response.status_code == 500:
         try:
             request_task = requests_lib.Request.decode(
-                requests_lib.RequestPayload(**response.json().get('detail')))
+                payloads.RequestPayload(**response.json().get('detail')))
             logger.debug(f'Got request with error: {request_task.name}')
         except Exception:  # pylint: disable=broad-except
             request_task = None
@@ -1503,12 +2013,7 @@ def get(request_id: str) -> Any:
     error = request_task.get_error()
     if error is not None:
         error_obj = error['object']
-        if env_options.Options.SHOW_DEBUG_INFO.get():
-            stacktrace = getattr(error_obj, 'stacktrace', str(error_obj))
-            logger.error('=== Traceback on SkyPilot API Server ===\n'
-                         f'{stacktrace}')
-        with ux_utils.print_exception_no_traceback():
-            raise error_obj
+        _raise_exception_object_on_client(error_obj)
     if request_task.status == requests_lib.RequestStatus.CANCELLED:
         with ux_utils.print_exception_no_traceback():
             raise exceptions.RequestCancelled(
@@ -1518,23 +2023,45 @@ def get(request_id: str) -> Any:
     return request_task.get_return_value()
 
 
+@typing.overload
+def stream_and_get(request_id: server_common.RequestId[T],
+                   log_path: Optional[str] = None,
+                   tail: Optional[int] = None,
+                   follow: bool = True,
+                   output_stream: Optional['io.TextIOBase'] = None) -> T:
+    ...
+
+
+@typing.overload
+def stream_and_get(request_id: None = None,
+                   log_path: Optional[str] = None,
+                   tail: Optional[int] = None,
+                   follow: bool = True,
+                   output_stream: Optional['io.TextIOBase'] = None) -> None:
+    ...
+
+
 @usage_lib.entrypoint
 @server_common.check_server_healthy_or_start
 @annotations.client_api
+@rest.retry_transient_errors()
 def stream_and_get(
-    request_id: Optional[str] = None,
+    request_id: Optional[server_common.RequestId[T]] = None,
     log_path: Optional[str] = None,
     tail: Optional[int] = None,
     follow: bool = True,
     output_stream: Optional['io.TextIOBase'] = None,
-) -> Any:
+) -> Optional[T]:
     """Streams the logs of a request or a log file and gets the final result.
 
     This will block until the request is finished. The request id can be a
     prefix of the full request id.
 
     Args:
-        request_id: The prefix of the request ID of the request to stream.
+        request_id: The request ID of the request to stream. May be a full
+            request ID or a prefix.
+            If None, the latest request submitted to the API server is streamed.
+            Using None request_id is not recommended in multi-user environments.
         log_path: The path to the log file to stream.
         tail: The number of lines to show from the end of the logs.
             If None, show all logs.
@@ -1545,6 +2072,8 @@ def stream_and_get(
     Returns:
         The ``Request Returns`` of the specified request. See the documentation
         of the specific requests above for more details.
+        If follow is False, will always return None. See note on
+        stream_response.
 
     Raises:
         Exception: It raises the same exceptions as the specific requests,
@@ -1558,27 +2087,44 @@ def stream_and_get(
         'follow': follow,
         'format': 'console',
     }
-    response = requests.get(
-        f'{server_common.get_server_url()}/api/stream',
+    response = server_common.make_authenticated_request(
+        'GET',
+        '/api/stream',
         params=params,
+        retry=False,
         timeout=(client_common.API_SERVER_REQUEST_CONNECTION_TIMEOUT_SECONDS,
                  None),
-        stream=True,
-        cookies=server_common.get_api_cookie_jar())
+        stream=True)
     if response.status_code in [404, 400]:
         detail = response.json().get('detail')
         with ux_utils.print_exception_no_traceback():
-            raise RuntimeError(f'Failed to stream logs: {detail}')
+            raise exceptions.ClientError(f'Failed to stream logs: {detail}')
+    stream_request_id: Optional[server_common.RequestId[
+        T]] = server_common.get_stream_request_id(response)
+    if request_id is not None and stream_request_id is not None:
+        assert request_id == stream_request_id
+    if request_id is None:
+        request_id = stream_request_id
     elif response.status_code != 200:
+        # TODO(syang): handle the case where the requestID is not provided
+        # see https://github.com/skypilot-org/skypilot/issues/6549
+        if request_id is None:
+            return None
         return get(request_id)
-    return stream_response(request_id, response, output_stream)
+    return stream_response(request_id,
+                           response,
+                           output_stream,
+                           resumable=True,
+                           get_result=follow)
 
 
 @usage_lib.entrypoint
 @annotations.client_api
-def api_cancel(request_ids: Optional[Union[str, List[str]]] = None,
+def api_cancel(request_ids: Optional[Union[server_common.RequestId[T],
+                                           List[server_common.RequestId[T]],
+                                           str, List[str]]] = None,
                all_users: bool = False,
-               silent: bool = False) -> server_common.RequestId:
+               silent: bool = False) -> server_common.RequestId[List[str]]:
     """Aborts a request or all requests.
 
     Args:
@@ -1618,20 +2164,35 @@ def api_cancel(request_ids: Optional[Union[str, List[str]]] = None,
         echo(f'Cancelling {len(request_ids)} request{plural}: '
              f'{request_id_str}...')
 
-    response = requests.post(f'{server_common.get_server_url()}/api/cancel',
-                             json=json.loads(body.model_dump_json()),
-                             timeout=5,
-                             cookies=server_common.get_api_cookie_jar())
+    response = server_common.make_authenticated_request(
+        'POST',
+        '/api/cancel',
+        json=json.loads(body.model_dump_json()),
+        timeout=5)
     return server_common.get_request_id(response)
 
 
+def _local_api_server_running(kill: bool = False) -> bool:
+    """Checks if the local api server is running."""
+    for process in psutil.process_iter(attrs=['pid', 'cmdline']):
+        cmdline = process.info['cmdline']
+        if cmdline and server_common.API_SERVER_CMD in ' '.join(cmdline):
+            if kill:
+                subprocess_utils.kill_children_processes(
+                    parent_pids=[process.pid], force=True)
+            return True
+    return False
+
+
 @usage_lib.entrypoint
 @annotations.client_api
 def api_status(
-    request_ids: Optional[List[str]] = None,
+    request_ids: Optional[List[Union[server_common.RequestId[T], str]]] = None,
     # pylint: disable=redefined-builtin
-    all_status: bool = False
-) -> List[requests_lib.RequestPayload]:
+    all_status: bool = False,
+    limit: Optional[int] = None,
+    fields: Optional[List[str]] = None,
+) -> List[payloads.RequestPayload]:
     """Lists all requests.
 
     Args:
@@ -1639,29 +2200,37 @@ def api_status(
             If None, all requests are queried.
         all_status: Whether to list all finished requests as well. This argument
             is ignored if request_ids is not None.
+        limit: The number of requests to show. If None, show all requests.
+        fields: The fields to get. If None, get all fields.
 
     Returns:
         A list of request payloads.
     """
-    body = payloads.RequestStatusBody(request_ids=request_ids,
-                                      all_status=all_status)
-    response = requests.get(
-        f'{server_common.get_server_url()}/api/status',
+    if server_common.is_api_server_local() and not _local_api_server_running():
+        logger.info('SkyPilot API server is not running.')
+        return []
+
+    body = payloads.RequestStatusBody(
+        request_ids=request_ids,
+        all_status=all_status,
+        limit=limit,
+        fields=fields,
+    )
+    response = server_common.make_authenticated_request(
+        'GET',
+        '/api/status',
         params=server_common.request_body_to_params(body),
         timeout=(client_common.API_SERVER_REQUEST_CONNECTION_TIMEOUT_SECONDS,
-                 None),
-        cookies=server_common.get_api_cookie_jar())
+                 None))
     server_common.handle_request_error(response)
-    return [
-        requests_lib.RequestPayload(**request) for request in response.json()
-    ]
+    return [payloads.RequestPayload(**request) for request in response.json()]
 
 
 # === API server management APIs ===
 @usage_lib.entrypoint
 @server_common.check_server_healthy_or_start
 @annotations.client_api
-def api_info() -> Dict[str, str]:
+def api_info() -> responses.APIHealthResponse:
     """Gets the server's status, commit and version.
 
     Returns:
@@ -1674,13 +2243,19 @@ def api_info() -> Dict[str, str]:
                 'api_version': '1',
                 'commit': 'abc1234567890',
                 'version': '1.0.0',
+                'version_on_disk': '1.0.0',
+                'user': {
+                    'name': 'test@example.com',
+                    'id': '12345abcd',
+                },
             }
 
+        Note that user may be None if we are not using an auth proxy.
+
     """
-    response = requests.get(f'{server_common.get_server_url()}/api/health',
-                            cookies=server_common.get_api_cookie_jar())
+    response = server_common.make_authenticated_request('GET', '/api/health')
     response.raise_for_status()
-    return response.json()
+    return responses.APIHealthResponse(**response.json())
 
 
 @usage_lib.entrypoint
@@ -1690,6 +2265,9 @@ def api_start(
     deploy: bool = False,
     host: str = '127.0.0.1',
     foreground: bool = False,
+    metrics: bool = False,
+    metrics_port: Optional[int] = None,
+    enable_basic_auth: bool = False,
 ) -> None:
     """Starts the API server.
 
@@ -1703,6 +2281,10 @@ def api_start(
             if deploy is True, to allow remote access.
         foreground: Whether to run the API server in the foreground (run in
             the current process).
+        metrics: Whether to export metrics of the API server.
+        metrics_port: The port to export metrics of the API server.
+        enable_basic_auth: Whether to enable basic authentication
+            in the API server.
     Returns:
         None
     """
@@ -1721,15 +2303,15 @@ def api_start(
                              'from the config file and/or unset the '
                              'SKYPILOT_API_SERVER_ENDPOINT environment '
                              'variable.')
-    server_common.check_server_healthy_or_start_fn(deploy, host, foreground)
+    server_common.check_server_healthy_or_start_fn(deploy, host, foreground,
+                                                   metrics, metrics_port,
+                                                   enable_basic_auth)
     if foreground:
         # Explain why current process exited
         logger.info('API server is already running:')
     api_server_url = server_common.get_server_url(host)
-    dashboard_url = server_common.get_dashboard_url(api_server_url)
-    dashboard_msg = f'Dashboard: {dashboard_url}'
-    logger.info(f'{ux_utils.INDENT_SYMBOL}SkyPilot API server: '
-                f'{api_server_url} {dashboard_msg}\n'
+    logger.info(f'{ux_utils.INDENT_SYMBOL}SkyPilot API server and dashboard: '
+                f'{api_server_url}\n'
                 f'{ux_utils.INDENT_LAST_SYMBOL}'
                 f'View API server logs at: {constants.API_SERVER_LOGS}')
 
@@ -1752,16 +2334,32 @@ def api_stop() -> None:
                 f'Cannot kill the API server at {server_url} because it is not '
                 f'the default SkyPilot API server started locally.')
 
-    found = False
-    for process in psutil.process_iter(attrs=['pid', 'cmdline']):
-        cmdline = process.info['cmdline']
-        if cmdline and server_common.API_SERVER_CMD in ' '.join(cmdline):
-            subprocess_utils.kill_children_processes(parent_pids=[process.pid],
-                                                     force=True)
-            found = True
-
-    # Remove the database for requests.
-    server_common.clear_local_api_server_database()
+    # Acquire the api server creation lock to prevent multiple processes from
+    # stopping and starting the API server at the same time.
+    with filelock.FileLock(
+            os.path.expanduser(constants.API_SERVER_CREATION_LOCK_PATH)):
+        try:
+            records = scheduler.get_controller_process_records()
+            if records is not None:
+                for record in records:
+                    try:
+                        if managed_job_utils.controller_process_alive(
+                                record, quiet=False):
+                            subprocess_utils.kill_children_processes(
+                                parent_pids=[record.pid], force=True)
+                    except (psutil.NoSuchProcess, psutil.ZombieProcess):
+                        continue
+                os.remove(os.path.expanduser(scheduler.JOB_CONTROLLER_PID_PATH))
+        except FileNotFoundError:
+            # its fine we will create it
+            pass
+        except Exception as e:  # pylint: disable=broad-except
+            # in case we get perm issues or something is messed up, just ignore
+            # it and assume the process is dead
+            logger.error(f'Error looking at job controller pid file: {e}')
+            pass
+
+        found = _local_api_server_running(kill=True)
 
     if found:
         logger.info(f'{colorama.Fore.GREEN}SkyPilot API server stopped.'
@@ -1796,9 +2394,86 @@ def api_server_logs(follow: bool = True, tail: Optional[int] = None) -> None:
         stream_and_get(log_path=constants.API_SERVER_LOGS, tail=tail)
 
 
+def _save_config_updates(endpoint: Optional[str] = None,
+                         service_account_token: Optional[str] = None) -> None:
+    """Save endpoint and/or service account token to config file."""
+    config_path = pathlib.Path(
+        skypilot_config.get_user_config_path()).expanduser()
+    with filelock.FileLock(config_path.with_suffix('.lock')):
+        if not config_path.exists():
+            config_path.touch()
+            config: Dict[str, Any] = {}
+        else:
+            config = skypilot_config.get_user_config()
+            config = dict(config)
+
+        # Update endpoint if provided
+        if endpoint is not None:
+            # We should always reset the api_server config to avoid legacy
+            # service account token.
+            config['api_server'] = {}
+            config['api_server']['endpoint'] = endpoint
+
+        # Update service account token if provided
+        if service_account_token is not None:
+            if 'api_server' not in config:
+                config['api_server'] = {}
+            config['api_server'][
+                'service_account_token'] = service_account_token
+
+        yaml_utils.dump_yaml(str(config_path), config)
+        skypilot_config.reload_config()
+
+
+def _clear_api_server_config() -> None:
+    """Clear endpoint and service account token from config file."""
+    config_path = pathlib.Path(
+        skypilot_config.get_user_config_path()).expanduser()
+    with filelock.FileLock(config_path.with_suffix('.lock')):
+        if not config_path.exists():
+            return
+
+        config = skypilot_config.get_user_config()
+        config = dict(config)
+        if 'api_server' in config:
+            # We might not have set the endpoint in the config file, so we
+            # need to check before deleting.
+            del config['api_server']
+
+        yaml_utils.dump_yaml(str(config_path), config, blank=True)
+        skypilot_config.reload_config()
+
+
+def _validate_endpoint(endpoint: Optional[str]) -> str:
+    """Validate and normalize the endpoint URL."""
+    if endpoint is None:
+        endpoint = click.prompt('Enter your SkyPilot API server endpoint')
+    # Check endpoint is a valid URL
+    if (endpoint is not None and not endpoint.startswith('http://') and
+            not endpoint.startswith('https://')):
+        raise click.BadParameter('Endpoint must be a valid URL.')
+    return endpoint.rstrip('/')
+
+
+def _check_endpoint_in_env_var(is_login: bool) -> None:
+    # If the user has set the endpoint via the environment variable, we should
+    # not do anything as we can't disambiguate between the env var and the
+    # config file.
+    """Check if the endpoint is set in the environment variable."""
+    if constants.SKY_API_SERVER_URL_ENV_VAR in os.environ:
+        with ux_utils.print_exception_no_traceback():
+            action = 'login to' if is_login else 'logout of'
+            raise RuntimeError(f'Cannot {action} API server when the endpoint '
+                               'is set via the environment variable. Run unset '
+                               f'{constants.SKY_API_SERVER_URL_ENV_VAR} to '
+                               'clear the environment variable.')
+
+
 @usage_lib.entrypoint
 @annotations.client_api
-def api_login(endpoint: Optional[str] = None) -> None:
+def api_login(endpoint: Optional[str] = None,
+              relogin: bool = False,
+              service_account_token: Optional[str] = None) -> None:
     """Logs into a SkyPilot API server.
 
     This sets the endpoint globally, i.e., all SkyPilot CLI and SDK calls will
@@ -1810,37 +2485,262 @@ def api_login(endpoint: Optional[str] = None) -> None:
     Args:
         endpoint: The endpoint of the SkyPilot API server, e.g.,
             http://1.2.3.4:46580 or https://skypilot.mydomain.com.
+        relogin: Whether to force relogin with OAuth2 when enabled.
+        service_account_token: Service account token for authentication.
 
     Returns:
         None
     """
-    # TODO(zhwu): this SDK sets global endpoint, which may not be the best
-    # design as a user may expect this is only effective for the current
-    # session. We should consider using env var for specifying endpoint.
-    if endpoint is None:
-        endpoint = click.prompt('Enter your SkyPilot API server endpoint')
-    # Check endpoint is a valid URL
-    if (endpoint is not None and not endpoint.startswith('http://') and
-            not endpoint.startswith('https://')):
-        raise click.BadParameter('Endpoint must be a valid URL.')
-
-    server_common.check_server_healthy(endpoint)
-
-    # Set the endpoint in the config file
-    config_path = pathlib.Path(
-        skypilot_config.get_user_config_path()).expanduser()
-    with filelock.FileLock(config_path.with_suffix('.lock')):
-        if not config_path.exists():
-            config_path.touch()
-            config = {'api_server': {'endpoint': endpoint}}
+    _check_endpoint_in_env_var(is_login=True)
+
+    # Validate and normalize endpoint
+    endpoint = _validate_endpoint(endpoint)
+
+    def _show_logged_in_message(
+            endpoint: str, dashboard_url: str, user: Optional[Dict[str, Any]],
+            server_status: server_common.ApiServerStatus) -> None:
+        """Show the logged in message."""
+        if server_status != server_common.ApiServerStatus.HEALTHY:
+            with ux_utils.print_exception_no_traceback():
+                raise ValueError(f'Cannot log in API server at '
+                                 f'{endpoint} (status: {server_status.value})')
+
+        identity_info = f'\n{ux_utils.INDENT_SYMBOL}{colorama.Fore.GREEN}User: '
+        if user:
+            user_name = user.get('name')
+            user_id = user.get('id')
+            if user_name and user_id:
+                identity_info += f'{user_name} ({user_id})'
+            elif user_id:
+                identity_info += user_id
         else:
-            config = skypilot_config.get_user_config()
-            config.set_nested(('api_server', 'endpoint'), endpoint)
-        common_utils.dump_yaml(str(config_path), dict(config))
-        dashboard_url = server_common.get_dashboard_url(endpoint)
+            identity_info = ''
         dashboard_msg = f'Dashboard: {dashboard_url}'
         click.secho(
             f'Logged into SkyPilot API server at: {endpoint}'
+            f'{identity_info}'
             f'\n{ux_utils.INDENT_LAST_SYMBOL}{colorama.Fore.GREEN}'
             f'{dashboard_msg}',
             fg='green')
+
+    def _set_user_hash(user_hash: Optional[str]) -> None:
+        if user_hash is not None:
+            if not common_utils.is_valid_user_hash(user_hash):
+                raise ValueError(f'Invalid user hash: {user_hash}')
+            common_utils.set_user_hash_locally(user_hash)
+
+    # Handle service account token authentication
+    if service_account_token:
+        if not service_account_token.startswith('sky_'):
+            raise ValueError('Invalid service account token format. '
+                             'Token must start with "sky_"')
+
+        # Save both endpoint and token to config in a single operation
+        _save_config_updates(endpoint=endpoint,
+                             service_account_token=service_account_token)
+
+        # Test the authentication by checking server health
+        try:
+            server_status, api_server_info = server_common.check_server_healthy(
+                endpoint)
+            dashboard_url = server_common.get_dashboard_url(endpoint)
+            if api_server_info.user is not None:
+                _set_user_hash(api_server_info.user.get('id'))
+            _show_logged_in_message(endpoint, dashboard_url,
+                                    api_server_info.user, server_status)
+
+            return
+        except exceptions.ApiServerConnectionError as e:
+            with ux_utils.print_exception_no_traceback():
+                raise RuntimeError(
+                    f'Failed to connect to API server at {endpoint}: {e}'
+                ) from e
+        except Exception as e:  # pylint: disable=broad-except
+            with ux_utils.print_exception_no_traceback():
+                raise RuntimeError(
+                    f'{colorama.Fore.RED}Service account token authentication '
+                    f'failed:{colorama.Style.RESET_ALL} {e}') from None
+
+    # OAuth2/cookie-based authentication flow
+    # TODO(zhwu): this SDK sets global endpoint, which may not be the best
+    # design as a user may expect this is only effective for the current
+    # session. We should consider using env var for specifying endpoint.
+
+    server_status, api_server_info = server_common.check_server_healthy(
+        endpoint)
+    if server_status == server_common.ApiServerStatus.NEEDS_AUTH or relogin:
+        # We detected an auth proxy, so go through the auth proxy cookie flow.
+        token: Optional[str] = None
+        server: Optional[oauth_lib.HTTPServer] = None
+        try:
+            callback_port = common_utils.find_free_port(8000)
+
+            token_container: Dict[str, Optional[str]] = {'token': None}
+            logger.debug('Starting local authentication server...')
+            server = oauth_lib.start_local_auth_server(callback_port,
+                                                       token_container,
+                                                       endpoint)
+
+            token_url = (f'{endpoint}/token?local_port={callback_port}')
+            if webbrowser.open(token_url):
+                click.echo(f'{colorama.Fore.GREEN}A web browser has been '
+                           f'opened at {token_url}. Please continue the login '
+                           f'in the web browser.{colorama.Style.RESET_ALL}\n'
+                           f'{colorama.Style.DIM}To manually copy the token, '
+                           f'press ctrl+c.{colorama.Style.RESET_ALL}')
+            else:
+                raise ValueError('Failed to open browser.')
+
+            start_time = time.time()
+
+            while (token_container['token'] is None and
+                   time.time() - start_time < oauth_lib.AUTH_TIMEOUT):
+                time.sleep(1)
+
+            if token_container['token'] is None:
+                click.echo(f'{colorama.Fore.YELLOW}Authentication timed out '
+                           f'after {oauth_lib.AUTH_TIMEOUT} seconds.')
+            else:
+                token = token_container['token']
+
+        except (Exception, KeyboardInterrupt) as e:  # pylint: disable=broad-except
+            logger.debug(f'Automatic authentication failed: {e}, '
+                         'falling back to manual token entry.')
+            if isinstance(e, KeyboardInterrupt):
+                click.echo(f'\n{colorama.Style.DIM}Interrupted. Press ctrl+c '
+                           f'again to exit.{colorama.Style.RESET_ALL}')
+            # Fall back to manual token entry
+            token_url = f'{endpoint}/token'
+            click.echo('Authentication is needed. Please visit this URL '
+                       f'to set up the token:{colorama.Style.BRIGHT}\n\n'
+                       f'{token_url}\n{colorama.Style.RESET_ALL}')
+            token = click.prompt('Paste the token')
+        finally:
+            if server is not None:
+                try:
+                    server.server_close()
+                except Exception:  # pylint: disable=broad-except
+                    pass
+            if not token:
+                with ux_utils.print_exception_no_traceback():
+                    raise ValueError('Authentication failed.')
+
+        # Parse the token.
+        # b64decode will ignore invalid characters, but does some length and
+        # padding checks.
+        try:
+            data = base64.b64decode(token)
+        except binascii.Error as e:
+            raise ValueError(f'Malformed token: {token}') from e
+        logger.debug(f'Token data: {data!r}')
+        try:
+            json_data = json.loads(data)
+        except (json.JSONDecodeError, UnicodeDecodeError) as e:
+            raise ValueError(f'Malformed token data: {data!r}') from e
+        if not isinstance(json_data, dict):
+            raise ValueError(f'Malformed token JSON: {json_data}')
+
+        if json_data.get('v') == 1:
+            user_hash = json_data.get('user')
+            cookie_dict = json_data['cookies']
+        elif 'v' not in json_data:
+            user_hash = None
+            cookie_dict = json_data
+        else:
+            raise ValueError(f'Unsupported token version: {json_data.get("v")}')
+
+        parsed_url = urlparse.urlparse(endpoint)
+        cookie_jar = cookiejar.MozillaCookieJar()
+        for (name, value) in cookie_dict.items():
+            # dict keys in JSON must be strings
+            assert isinstance(name, str)
+            if not isinstance(value, str):
+                raise ValueError('Malformed token - bad key/value: '
+                                 f'{name}: {value}')
+
+            # See CookieJar._cookie_from_cookie_tuple
+            # oauth2proxy default is Max-Age 604800
+            expires = int(time.time()) + 604800
+            domain = str(parsed_url.hostname)
+            domain_initial_dot = domain.startswith('.')
+            secure = parsed_url.scheme == 'https'
+            if not domain_initial_dot:
+                domain = '.' + domain
+
+            cookie_jar.set_cookie(
+                cookiejar.Cookie(
+                    version=0,
+                    name=name,
+                    value=value,
+                    port=None,
+                    port_specified=False,
+                    domain=domain,
+                    domain_specified=True,
+                    domain_initial_dot=domain_initial_dot,
+                    path='',
+                    path_specified=False,
+                    secure=secure,
+                    expires=expires,
+                    discard=False,
+                    comment=None,
+                    comment_url=None,
+                    rest=dict(),
+                ))
+
+        # Now that the cookies are parsed, save them to the cookie jar.
+        server_common.set_api_cookie_jar(cookie_jar)
+
+        # Set the user hash in the local file.
+        # If the server already has a token for this user set it to the local
+        # file, otherwise use the new user hash.
+        if (api_server_info.user is not None and
+                api_server_info.user.get('id') is not None):
+            _set_user_hash(api_server_info.user.get('id'))
+        else:
+            _set_user_hash(user_hash)
+    else:
+        # Check if basic auth is enabled
+        if api_server_info.basic_auth_enabled:
+            if api_server_info.user is None:
+                with ux_utils.print_exception_no_traceback():
+                    raise ValueError(
+                        'Basic auth is enabled but no valid user is found')
+
+        # Set the user hash in the local file.
+        if api_server_info.user is not None:
+            _set_user_hash(api_server_info.user.get('id'))
+
+    # Set the endpoint in the config file
+    _save_config_updates(endpoint=endpoint)
+    dashboard_url = server_common.get_dashboard_url(endpoint)
+
+    # see https://github.com/python/mypy/issues/5107 on why
+    # typing is disabled on this line
+    server_common.get_api_server_status.cache_clear()  # type: ignore
+    # After successful authentication, check server health again to get user
+    # identity
+    server_status, final_api_server_info = server_common.check_server_healthy(
+        endpoint)
+    _show_logged_in_message(endpoint, dashboard_url, final_api_server_info.user,
+                            server_status)
+
+
+@usage_lib.entrypoint
+@annotations.client_api
+def api_logout() -> None:
+    """Logout of the API server.
+
+    Clears all cookies and settings stored in ~/.sky/config.yaml"""
+    _check_endpoint_in_env_var(is_login=False)
+
+    if server_common.is_api_server_local():
+        with ux_utils.print_exception_no_traceback():
+            raise RuntimeError('Local api server cannot be logged out. '
+                               'Use `sky api stop` instead.')
+
+    # no need to clear cookies if it doesn't exist.
+    server_common.set_api_cookie_jar(cookiejar.MozillaCookieJar(),
+                                     create_if_not_exists=False)
+    _clear_api_server_config()
+    logger.info(f'{colorama.Fore.GREEN}Logged out of SkyPilot API server.'
+                f'{colorama.Style.RESET_ALL}')