"""Client-side Python SDK for SkyPilot.

All functions will return a future that can be awaited on with the `get` method.

Usage example:

.. code-block:: python

    request_id = sky.status()
    statuses = sky.get(request_id)

"""
from http import cookiejar
import json
import logging
import os
import subprocess
import typing
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 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 interactive_utils
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.ssh_node_pools import utils as ssh_utils
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

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')

logger = sky_logging.init_logger(__name__)
logging.getLogger('httpx').setLevel(logging.CRITICAL)

_LINE_PROCESSED_KEY = 'line_processed'

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,
                    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 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:
                line_count += 1

                line = interactive_utils.handle_interactive_auth(line)
                if line is None:
                    # Line was consumed by interactive auth handler
                    continue

                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


@usage_lib.entrypoint
@server_common.check_server_healthy_or_start
@annotations.client_api
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:
        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.

    Request Returns:
        None
    """
    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(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.
    """
    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[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
    that may not be available in the user's account.

    Args:
        gpus_only: Whether to only list GPU accelerators.
        name_filter: The name filter.
        region_filter: The region filter.
        quantity_filter: The quantity filter.
        clouds: The clouds to list.
        all_regions: Whether to list all regions.
        require_price: Whether to require price.
        case_sensitive: Whether to case sensitive.

    Returns:
        The request ID of the list accelerator counts request.

    Request Returns:
        acc_to_instance_type_dict (Dict[str, List[InstanceTypeInfo]]): A
            dictionary of canonical accelerator names mapped to a list of
            instance type offerings. See usage in cli.py.
    """
    body = payloads.ListAcceleratorsBody(
        gpus_only=gpus_only,
        name_filter=name_filter,
        region_filter=region_filter,
        quantity_filter=quantity_filter,
        clouds=clouds,
        all_regions=all_regions,
        require_price=require_price,
        case_sensitive=case_sensitive,
    )
    response = server_common.make_authenticated_request(
        'POST', '/list_accelerators', 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 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[Dict[str, List[float]]]:
    """Lists all accelerators offered by Sky and available counts.

    Args:
        gpus_only: Whether to only list GPU accelerators.
        name_filter: The name filter.
        region_filter: The region filter.
        quantity_filter: The quantity filter.
        clouds: The clouds to list.

    Returns:
        The request ID of the list accelerator counts request.

    Request Returns:
        acc_to_acc_num_dict (Dict[str, List[int]]): A dictionary of canonical
            accelerator names mapped to a list of available counts. See usage
            in cli.py.
    """
    body = payloads.ListAcceleratorCountsBody(
        gpus_only=gpus_only,
        name_filter=name_filter,
        region_filter=region_filter,
        quantity_filter=quantity_filter,
        clouds=clouds,
    )
    response = server_common.make_authenticated_request(
        'POST',
        '/list_accelerator_counts',
        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 optimize(
    dag: 'sky.Dag',
    minimize: common.OptimizeTarget = common.OptimizeTarget.COST,
    admin_policy_request_options: Optional[admin_policy.RequestOptions] = None
) -> server_common.RequestId['sky.Dag']:
    """Finds the best execution plan for the given DAG.

    Args:
        dag: the DAG to optimize.
        minimize: whether to minimize cost or time.
        admin_policy_request_options: Request options used for admin policy
            validation. This is only required when a admin policy is in use,
            see: https://docs.skypilot.co/en/latest/cloud-setup/policy.html

    Returns:
        The request ID of the optimize request.

    Request Returns:
        optimized_dag (str): The optimized DAG in YAML format.

    Request Raises:
        exceptions.ResourcesUnavailableError: if no resources are available
            for a task.
        exceptions.NoCloudAccessError: if no public clouds are enabled.
    """
    dag_str = dag_utils.dump_chain_dag_to_yaml_str(dag)

    body = payloads.OptimizeBody(dag=dag_str,
                                 minimize=minimize,
                                 request_options=admin_policy_request_options)
    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
def validate(
    dag: 'sky.Dag',
    workdir_only: bool = False,
    admin_policy_request_options: Optional[admin_policy.RequestOptions] = None
) -> None:
    """Validates the tasks.

    The file paths (workdir and file_mounts) are validated on the client side
    while the rest (e.g. resource) are validated on server side.

    Raises exceptions if the DAG is invalid.

    Args:
        dag: the DAG to validate.
        workdir_only: whether to only validate the workdir. This is used for
            `exec` as it does not need other files/folders in file_mounts.
        admin_policy_request_options: Request options used for admin policy
            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 = server_common.make_authenticated_request(
        'POST', '/validate', json=json.loads(body.model_dump_json()))
    if response.status_code == 400:
        _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(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,
                                          starting_page=starting_page)
    logger.info(f'Opening dashboard in browser: {url}')
    webbrowser.open(url)


@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,
    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']]]:
    """Launches a cluster or task.

    The task's setup and run commands are executed under the task's workdir
    (when specified, it is synced to remote cluster).  The task undergoes job
    queue scheduling on the cluster.

    Currently, the first argument must be a sky.Task, or (EXPERIMENTAL advanced
    usage) a sky.Dag. In the latter case, currently it must contain a single
    task; support for pipelines/general DAGs are in experimental branches.

    Example:
        .. code-block:: python

            import sky
            task = sky.Task(run='echo hello SkyPilot')
            task.set_resources(
                sky.Resources(infra='aws', accelerators='V100:4'))
            sky.launch(task, cluster_name='my-cluster')


    Args:
        task: sky.Task, or sky.Dag (experimental; 1-task only) to launch.
        cluster_name: name of the cluster to create/reuse.  If None,
          auto-generate a name.
        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 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. 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,
          OptimizeTarget.TIME.
        no_setup: if True, do not re-run setup commands.
        clone_disk_from: [Experimental] if set, clone the disk from the
          specified cluster. This is useful to migrate the cluster to a
          different availability zone or region.
        fast: [Experimental] If the cluster is already up and available,
          skip provisioning and setup steps.
        _need_confirmation: (Internal only) If True, show the confirmation
            prompt.

    Returns:
        The request ID of the launch request.

    Request Returns:
        job_id (Optional[int]): the job ID of the submitted job. None if the
          backend is not ``CloudVmRayBackend``, or no job is submitted to the
          cluster.
        handle (Optional[backends.ResourceHandle]): the handle to the cluster.
          None if dryrun.

    Request Raises:
        exceptions.ClusterOwnerIdentityMismatchError: if the cluster is owned
          by another user.
        exceptions.InvalidClusterNameError: if the cluster name is invalid.
        exceptions.ResourcesMismatchError: if the requested resources
          do not match the existing cluster.
        exceptions.NotSupportedError: if required features are not supported
          by the backend/cloud/cluster.
        exceptions.ResourcesUnavailableError: if the requested resources
          cannot be satisfied. The failover_history of the exception will be set
          as:

          1. Empty: iff the first-ever sky.optimize() fails to find a feasible
             resource; no pre-check or actual launch is attempted.

          2. Non-empty: iff at least 1 exception from either our pre-checks
             (e.g., cluster name invalid) or a region/zone throwing resource
             unavailability.

        exceptions.CommandError: any ssh command error.
        exceptions.NoCloudAccessError: if all clouds are disabled.

    Other exceptions may be raised depending on the backend.
    """
    if cluster_name is None:
        cluster_name = cluster_utils.generate_cluster_name()

    if clone_disk_from is not None:
        with ux_utils.print_exception_no_traceback():
            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.
        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 = ''
        current_user = common_utils.get_local_user_name()
        cluster_user_name = current_user
        if not clusters:
            # Show the optimize log before the prompt if the cluster does not
            # exist.
            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 == 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})'

        # Prompt if (1) --cluster is None, or (2) cluster doesn't exist, or (3)
        # it exists but is STOPPED.
        prompt = None
        if cluster_status is None:
            prompt = (
                f'Launching a new cluster {cluster_name!r}. '
                # '{clone_source_str}. '
                'Proceed?')
        elif cluster_status == status_lib.ClusterStatus.STOPPED:
            user_name_str = ''
            if cluster_user_hash != common_utils.get_user_hash():
                user_name_str = (' created by another user '
                                 f'{cluster_user_name!r}'
                                 f'{cluster_user_hash_str}')
            prompt = (f'Restarting the stopped cluster {cluster_name!r}'
                      f'{user_name_str}. Proceed?')
        elif cluster_user_hash != common_utils.get_user_hash():
            # Prompt if the cluster was created by a different user.
            prompt = (f'Cluster {cluster_name!r} was created by another user '
                      f'{cluster_user_name!r}{cluster_user_hash_str}. '
                      'Reusing the cluster. Proceed?')
        if prompt is not None:
            confirm_shown = True
            click.confirm(prompt, default=True, abort=True, show_default=True)

    if not confirm_shown:
        click.secho('Running on cluster: ', fg='cyan', nl=False)
        click.secho(cluster_name)

    dag = client_common.upload_mounts_to_api_server(dag)

    dag_str = dag_utils.dump_chain_dag_to_yaml_str(dag)

    body = payloads.LaunchBody(
        task=dag_str,
        cluster_name=cluster_name,
        retry_until_up=retry_until_up,
        dryrun=dryrun,
        backend=backend.NAME if backend else None,
        optimize_target=optimize_target,
        no_setup=no_setup,
        clone_disk_from=clone_disk_from,
        fast=fast,
        # For internal use
        quiet_optimizer=_need_confirmation,
        is_launched_by_jobs_controller=_is_launched_by_jobs_controller,
        is_launched_by_sky_serve_controller=(
            _is_launched_by_sky_serve_controller),
        disable_controller_check=_disable_controller_check,
    )
    response = server_common.make_authenticated_request(
        'POST', '/launch', 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 exec(  # pylint: disable=redefined-builtin
    task: Union['sky.Task', 'sky.Dag'],
    cluster_name: Optional[str] = None,
    dryrun: bool = False,
    down: bool = False,  # pylint: disable=redefined-outer-name
    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:

    (1) workdir syncing, if the task has a workdir specified;
    (2) executing the task's ``run`` commands.

    All other steps (provisioning, setup commands, file mounts syncing) are
    skipped.  If any of those specifications changed in the task, this function
    will not reflect those changes.  To ensure a cluster's setup is up to date,
    use ``sky.launch()`` instead.

    Execution and scheduling behavior:

    - The task will undergo job queue scheduling, respecting any specified
      resource requirement. It can be executed on any node of the cluster with
      enough resources.
    - The task is run under the workdir (if specified).
    - The task is run non-interactively (without a pseudo-terminal or
      pty), so interactive commands such as ``htop`` do not work.
      Use ``ssh my_cluster`` instead.

    Args:
        task: sky.Task, or sky.Dag (experimental; 1-task only) containing the
          task to execute.
        cluster_name: name of an existing cluster to execute the task.
        dryrun: if True, do not actually execute the task.
        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.
        backend: backend to use.  If None, use the default backend
          (CloudVMRayBackend).

    Returns:
        The request ID of the exec request.


    Request Returns:
        job_id (Optional[int]): the job ID of the submitted job. None if the
          backend is not CloudVmRayBackend, or no job is submitted to
          the cluster.
        handle (Optional[backends.ResourceHandle]): the handle to the cluster.
          None if dryrun.

    Request Raises:
        ValueError: if the specified cluster is not in UP status.
        sky.exceptions.ClusterDoesNotExist: if the specified cluster does not
          exist.
        sky.exceptions.NotSupportedError: if the specified cluster is a
          controller that does not support this operation.
    """
    dag = dag_utils.convert_entrypoint_to_dag(task)
    validate(dag, workdir_only=True)
    dag = client_common.upload_mounts_to_api_server(dag, workdir_only=True)
    dag_str = dag_utils.dump_chain_dag_to_yaml_str(dag)
    body = payloads.ExecBody(
        task=dag_str,
        cluster_name=cluster_name,
        dryrun=dryrun,
        down=down,
        backend=backend.NAME if backend else None,
    )

    response = server_common.make_authenticated_request(
        'POST', '/exec', json=json.loads(body.model_dump_json()), timeout=5)
    return server_common.get_request_id(response)


@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: 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:
        cluster_name: name of the cluster.
        job_id: job id.
        follow: if True, follow the logs. Otherwise, return the logs
            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. 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:
        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.
        sky.exceptions.ClusterDoesNotExist: if the cluster does not exist.
        sky.exceptions.ClusterNotUpError: if the cluster is not UP.
        sky.exceptions.NotSupportedError: if the cluster is not based on
          CloudVmRayBackend.
        sky.exceptions.ClusterOwnerIdentityMismatchError: if the current user is
          not the same as the user who created the cluster.
        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 = 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))
    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
@server_common.check_server_healthy_or_start
@annotations.client_api
def download_logs(cluster_name: str,
                  job_ids: Optional[List[str]]) -> Dict[str, str]:
    """Downloads the logs of jobs.

    Args:
        cluster_name: (str) name of the cluster.
        job_ids: (List[str]) job ids.

    Returns:
        The request ID of the download_logs request.

    Request Returns:
        job_log_paths (Dict[str, str]): a mapping of job_id to local log path.

    Request Raises:
        sky.exceptions.ClusterDoesNotExist: if the cluster does not exist.
        sky.exceptions.ClusterNotUpError: if the cluster is not UP.
        sky.exceptions.NotSupportedError: if the cluster is not based on
          CloudVmRayBackend.
        sky.exceptions.ClusterOwnerIdentityMismatchError: if the current user is
          not the same as the user who created the cluster.
        sky.exceptions.CloudUserIdentityError: if we fail to get the current
          user identity.
    """
    body = payloads.ClusterJobsDownloadLogsBody(
        cluster_name=cluster_name,
        job_ids=job_ids,
    )
    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 {
        job_id: remote2local_path_dict[remote_path]
        for job_id, remote_path in job_id_remote_path_dict.items()
    }


@usage_lib.entrypoint
@server_common.check_server_healthy_or_start
@annotations.client_api
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['backends.CloudVmRayResourceHandle']:
    """Restart a cluster.

    If a cluster is previously stopped (status is STOPPED) or failed in
    provisioning/runtime installation (status is INIT), this function will
    attempt to start the cluster.  In the latter case, provisioning and runtime
    installation will be retried.

    Auto-failover provisioning is not used when restarting a stopped
    cluster. It will be started on the same cloud, region, and zone that were
    chosen before.

    If a cluster is already in the UP status, this function has no effect.

    Args:
        cluster_name: name of the cluster to start.
        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.
        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
            minutes of idle time after all jobs finish (successfully or
            abnormally). Requires ``idle_minutes_to_autostop`` to be set.
        force: whether to force start the cluster even if it is already up.
            Useful for upgrading SkyPilot runtime.

    Returns:
        The request ID of the start request.

    Request Returns:
        None

    Request Raises:
        ValueError: argument values are invalid: (1) if ``down`` is set to True
            but ``idle_minutes_to_autostop`` is None; (2) if the specified
            cluster is the managed jobs controller, and either
            ``idle_minutes_to_autostop`` is not None or ``down`` is True (omit
            them to use the default autostop settings).
        sky.exceptions.ClusterDoesNotExist: the specified cluster does not
            exist.
        sky.exceptions.NotSupportedError: if the cluster to restart was
            launched using a non-default backend that does not support this
            operation.
        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 = 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[None]:
    """Tears down a cluster.

    Tearing down a cluster will delete all associated resources (all billing
    stops), and any data on the attached disks will be lost.  Accelerators
    (e.g., TPUs) that are part of the cluster will be deleted too.

    Args:
        cluster_name: name of the cluster to down.
        purge: (Advanced) Forcefully remove the cluster from SkyPilot's cluster
            table, even if the actual cluster termination failed on the cloud.
            WARNING: This flag should only be set sparingly in certain manual
            troubleshooting scenarios; with it set, it is the user's
            responsibility to ensure there are no leaked instances and related
            resources.

    Returns:
        The request ID of the down request.

    Request Returns:
        None

    Request Raises:
        sky.exceptions.ClusterDoesNotExist: the specified cluster does not
            exist.
        RuntimeError: failed to tear down the cluster.
        sky.exceptions.NotSupportedError: the specified cluster is the managed
            jobs controller.

    """
    body = payloads.StopOrDownBody(
        cluster_name=cluster_name,
        purge=purge,
    )
    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[None]:
    """Stops a cluster.

    Data on attached disks is not lost when a cluster is stopped.  Billing for
    the instances will stop, while the disks will still be charged.  Those
    disks will be reattached when restarting the cluster.

    Currently, spot instance clusters cannot be stopped (except for GCP, which
    does allow disk contents to be preserved when stopping spot VMs).

    Args:
        cluster_name: name of the cluster to stop.
        purge: (Advanced) Forcefully mark the cluster as stopped in SkyPilot's
            cluster table, even if the actual cluster stop operation failed on
            the cloud. WARNING: This flag should only be set sparingly in
            certain manual troubleshooting scenarios; with it set, it is the
            user's responsibility to ensure there are no leaked instances and
            related resources.

    Returns:
        The request ID of the stop request.

    Request Returns:
        None

    Request Raises:
        sky.exceptions.ClusterDoesNotExist: the specified cluster does not
            exist.
        RuntimeError: failed to stop the cluster.
        sky.exceptions.NotSupportedError: if the specified cluster is a spot
            cluster, or a TPU VM Pod cluster, or the managed jobs controller.

    """
    body = payloads.StopOrDownBody(
        cluster_name=cluster_name,
        purge=purge,
    )
    response = server_common.make_authenticated_request(
        'POST', '/stop', 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 autostop(
        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
    becomes idle for a specified duration.  Idleness means there are no
    in-progress (pending/running) jobs in a cluster's job queue.

    Idleness time of a cluster is reset to zero, whenever:

    - A job is submitted (``sky.launch()`` or ``sky.exec()``).

    - The cluster has restarted.

    - An autostop is set when there is no active setting. (Namely, either
      there's never any autostop setting set, or the previous autostop setting
      was canceled.) This is useful for restarting the autostop timer.

    Example: say a cluster without any autostop set has been idle for 1 hour,
    then an autostop of 30 minutes is set. The cluster will not be immediately
    autostopped. Instead, the idleness timer only starts counting after the
    autostop setting was set.

    When multiple autostop settings are specified for the same cluster, the
    last setting takes precedence.

    Args:
        cluster_name: name of the cluster.
        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).

    Returns:
        The request ID of the autostop request.

    Request Returns:
        None

    Request Raises:
        sky.exceptions.ClusterDoesNotExist: if the cluster does not exist.
        sky.exceptions.ClusterNotUpError: if the cluster is not UP.
        sky.exceptions.NotSupportedError: if the cluster is not based on
            CloudVmRayBackend or the cluster is TPU VM Pod.
        sky.exceptions.ClusterOwnerIdentityMismatchError: if the current user is
            not the same as the user who created the cluster.
        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 = 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[List[responses.ClusterJobRecord]]:
    """Gets the job queue of a cluster.

    Args:
        cluster_name: name of the cluster.
        skip_finished: if True, skip finished jobs.
        all_users: if True, return jobs from all users.


    Returns:
        The request ID of the queue request.

    Request Returns:
        job_records (List[responses.ClusterJobRecord]): A list of job records
            for each job in the queue.

            .. code-block:: python

                [
                    {
                        'job_id': (int) job id,
                        'job_name': (str) job name,
                        'username': (str) username,
                        'user_hash': (str) user hash,
                        'submitted_at': (int) timestamp of submitted,
                        'start_at': (int) timestamp of started,
                        'end_at': (int) timestamp of ended,
                        'resources': (str) resources,
                        'status': (job_lib.JobStatus) job status,
                        'log_path': (str) log path,
                    }
                ]

    Request Raises:
        sky.exceptions.ClusterDoesNotExist: if the cluster does not exist.
        sky.exceptions.ClusterNotUpError: if the cluster is not UP.
        sky.exceptions.NotSupportedError: if the cluster is not based on
            ``CloudVmRayBackend``.
        sky.exceptions.ClusterOwnerIdentityMismatchError: if the current user is
            not the same as the user who created the cluster.
        sky.exceptions.CloudUserIdentityError: if we fail to get the current
            user identity.
        sky.exceptions.CommandError: if failed to get the job queue with ssh.
    """
    body = payloads.QueueBody(
        cluster_name=cluster_name,
        skip_finished=skip_finished,
        all_users=all_users,
    )
    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[Dict[Optional[int],
                                  Optional['job_lib.JobStatus']]]:
    """Gets the status of jobs on a cluster.

    Args:
        cluster_name: name of the cluster.
        job_ids: job ids. If None, get the status of the last job.

    Returns:
        The request ID of the job status request.

    Request Returns:
        job_statuses (Dict[Optional[int], Optional[job_lib.JobStatus]]): A
            mapping of job_id to job statuses. The status will be None if the
            job does not exist. If job_ids is None and there is no job on the
            cluster, it will return {None: None}.

    Request Raises:
        sky.exceptions.ClusterDoesNotExist: if the cluster does not exist.
        sky.exceptions.ClusterNotUpError: if the cluster is not UP.
        sky.exceptions.NotSupportedError: if the cluster is not based on
            ``CloudVmRayBackend``.
        sky.exceptions.ClusterOwnerIdentityMismatchError: if the current user is
            not the same as the user who created the cluster.
        sky.exceptions.CloudUserIdentityError: if we fail to get the current
            user identity.
    """
    # TODO: merge this into the queue endpoint, i.e., let the queue endpoint
    # take job_ids to filter the returned jobs.
    body = payloads.JobStatusBody(
        cluster_name=cluster_name,
        job_ids=job_ids,
    )
    response = server_common.make_authenticated_request(
        'POST', '/job_status', 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 cancel(
    cluster_name: str,
    all: bool = False,  # pylint: disable=redefined-builtin
    all_users: bool = False,
    job_ids: Optional[List[int]] = None,
    # pylint: disable=invalid-name
    _try_cancel_if_cluster_is_init: bool = False
) -> server_common.RequestId[None]:
    """Cancels jobs on a cluster.

    Args:
        cluster_name: name of the cluster.
        all: if True, cancel all jobs.
        all_users: if True, cancel all jobs from all users.
        job_ids: a list of job IDs to cancel.
        _try_cancel_if_cluster_is_init: (bool) whether to try cancelling the job
            even if the cluster is not UP, but the head node is still alive.
            This is used by the jobs controller to cancel the job when the
            worker node is preempted in the spot cluster.

    Returns:
        The request ID of the cancel request.

    Request Returns:
        None

    Request Raises:
        ValueError: if arguments are invalid.
        sky.exceptions.ClusterDoesNotExist: if the cluster does not exist.
        sky.exceptions.ClusterNotUpError: if the cluster is not UP.
        sky.exceptions.NotSupportedError: if the specified cluster is a
            controller that does not support this operation.
        sky.exceptions.ClusterOwnerIdentityMismatchError: if the current user is
            not the same as the user who created the cluster.
        sky.exceptions.CloudUserIdentityError: if we fail to get the current
            user identity.

    """
    body = payloads.CancelBody(
        cluster_name=cluster_name,
        all=all,
        all_users=all_users,
        job_ids=job_ids,
        try_cancel_if_cluster_is_init=_try_cancel_if_cluster_is_init,
    )
    response = server_common.make_authenticated_request(
        'POST', '/cancel', 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(
    cluster_names: Optional[List[str]] = None,
    refresh: common.StatusRefreshMode = common.StatusRefreshMode.NONE,
    all_users: bool = False,
    *,
    _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
    clusters.

    Each cluster can have one of the following statuses:

    - ``INIT``: The cluster may be live or down. It can happen in the following
      cases:

      - Ongoing provisioning or runtime setup. (A ``sky.launch()`` has started
        but has not completed.)
      - Or, the cluster is in an abnormal state, e.g., some cluster nodes are
        down, or the SkyPilot runtime is unhealthy. (To recover the cluster,
        try ``sky launch`` again on it.)

    - ``UP``: Provisioning and runtime setup have succeeded and the cluster is
      live.  (The most recent ``sky.launch()`` has completed successfully.)

    - ``STOPPED``: The cluster is stopped and the storage is persisted. Use
      ``sky.start()`` to restart the cluster.

    Autostop column:

    - The autostop column indicates how long the cluster will be autostopped
      after minutes of idling (no jobs running). If ``to_down`` is True, the
      cluster will be autodowned, rather than autostopped.

    Getting up-to-date cluster statuses:

    - In normal cases where clusters are entirely managed by SkyPilot (i.e., no
      manual operations in cloud consoles) and no autostopping is used, the
      table returned by this command will accurately reflect the cluster
      statuses.

    - In cases where the clusters are changed outside of SkyPilot (e.g., manual
      operations in cloud consoles; unmanaged spot clusters getting preempted)
      or for autostop-enabled clusters, use ``refresh=True`` to query the
      latest cluster statuses from the cloud providers.

    Args:
        cluster_names: a list of cluster names to query. If not
            provided, all clusters will be queried.
        refresh: whether to query the latest cluster statuses from the cloud
            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.

    Request Returns:
        cluster_records (List[Dict[str, Any]]): A list of dicts, with each dict
          containing the information of a cluster. If a cluster is found to be
          terminated or not found, it will be omitted from the returned list.

          .. code-block:: python

            {
              'name': (str) cluster name,
              'launched_at': (int) timestamp of last launch on this cluster,
              'handle': (ResourceHandle) an internal handle to the cluster,
              'last_use': (str) the last command/entrypoint that affected this
              cluster,
              'status': (sky.ClusterStatus) cluster status,
              'autostop': (int) idle time before autostop,
              'to_down': (bool) whether autodown is used instead of autostop,
              'metadata': (dict) metadata of the cluster,
              'user_hash': (str) user hash of the cluster owner,
              'user_name': (str) user name of the cluster owner,
              'resources_str': (str) the resource string representation of the
                cluster,
            }

    """
    # TODO(zhwu): this does not stream the logs output by logger back to the
    # user, due to the rich progress implementation.
    body = payloads.StatusBody(
        cluster_names=cluster_names,
        refresh=refresh,
        all_users=all_users,
        include_credentials=_include_credentials,
        summary_response=_summary_response,
    )
    response = server_common.make_authenticated_request(
        'POST', '/status', 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 endpoints(
    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
            for all ports are returned.

    Returns:
        The request ID of the endpoints request.

    Request Returns:
        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.
        RuntimeError: if the cluster has no ports to be exposed or no endpoints
            are exposed yet.
    """
    body = payloads.EndpointsBody(
        cluster=cluster,
        port=port,
    )
    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(
    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
    of resources being used and the duration of use up until the call to
    status. This means if the cluster is UP, successive calls to report will
    show increasing price. The estimated cost is calculated based on the local
    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.

    Request Returns:
        cluster_cost_records (List[Dict[str, Any]]): A list of dicts, with each
          dict containing the cost information of a cluster.

          .. code-block:: python

            {
              'name': (str) cluster name,
              'launched_at': (int) timestamp of last launch on this cluster,
              'duration': (int) total seconds that cluster was up and running,
              'last_use': (str) the last command/entrypoint that affected this
              'num_nodes': (int) number of nodes launched for cluster,
              'resources': (resources.Resources) type of resource launched,
              'cluster_hash': (str) unique hash identifying cluster,
              'usage_intervals': (List[Tuple[int, int]]) cluster usage times,
              'total_cost': (float) cost given resources and usage intervals,
            }
    """
    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)


# === Storage APIs ===
@usage_lib.entrypoint
@server_common.check_server_healthy_or_start
@annotations.client_api
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[responses.StorageRecord]):
            A list of storage records.
    """
    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[None]:
    """Deletes a storage.

    Args:
        name: The name of the storage to delete.

    Returns:
        The request ID of the storage delete request.

    Request Returns:
        None

    Request Raises:
        ValueError: If the storage does not exist.
    """
    body = payloads.StorageBody(name=name)
    response = server_common.make_authenticated_request(
        'POST', '/storage/delete', json=json.loads(body.model_dump_json()))
    return server_common.get_request_id(response)


# === Kubernetes ===


@usage_lib.entrypoint
@server_common.check_server_healthy_or_start
@annotations.client_api
def local_up(gpus: bool,
             name: Optional[str] = None,
             port_start: Optional[int] = None) -> server_common.RequestId[None]:
    """Launches a Kubernetes cluster on local machines.

    Returns:
        request_id: The request ID of the local up request.
    """
    # We do not allow local up when the API server is running remotely since it
    # will modify the kubeconfig.
    # 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, 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(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 '
                             'SkyPilot locally.')

    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)


@usage_lib.entrypoint
@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,
    is_ssh: Optional[bool] = None
) -> server_common.RequestId[List[Tuple[
        str, List['models.RealtimeGpuAvailability']]]]:
    """Gets the real-time Kubernetes GPU availability.

    Returns:
        The request ID of the real-time Kubernetes GPU availability request.
    """
    body = payloads.RealtimeGpuAvailabilityRequestBody(
        context=context,
        name_filter=name_filter,
        quantity_filter=quantity_filter,
        is_ssh=is_ssh,
    )
    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)


@usage_lib.entrypoint
@server_common.check_server_healthy_or_start
@annotations.client_api
def kubernetes_node_info(
    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
    number of GPUs available on the node and the number of free GPUs on the
    node.

    If the user does not have sufficient permissions to list pods in all
    namespaces, the function will return free GPUs as -1.

    Args:
        context: The Kubernetes context. If None, the default context is used.

    Returns:
        The request ID of the Kubernetes node info request.

    Request Returns:
        KubernetesNodesInfo: A model that contains the node info map and other
            information.
    """
    body = payloads.KubernetesNodeInfoRequestBody(context=context)
    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[
    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
    as managed jobs or services.

    Returns:
        The request ID of the status request.

    Request Returns:
        A tuple containing:
        - all_clusters: List of KubernetesSkyPilotClusterInfoPayload with info
            for all clusters, including managed jobs, services and controllers.
        - unmanaged_clusters: List of KubernetesSkyPilotClusterInfoPayload with
            info for all clusters excluding managed jobs and services.
            Controllers are included.
        - all_jobs: List of managed jobs from all controllers. Each entry is a
            dictionary job info, see jobs.queue_from_kubernetes_pod for details.
        - context: Kubernetes context used to fetch the cluster information.
    """
    response = server_common.make_authenticated_request('GET',
                                                        '/status_kubernetes')
    return server_common.get_request_id(response)


# === API request APIs ===
@usage_lib.entrypoint
@annotations.client_api
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. May be a full request
            ID or a prefix.

    Returns:
        The ``Request Returns`` of the specified request. See the documentation
        of the specific requests above for more details.

    Raises:
        Exception: It raises the same exceptions as the specific requests,
            see ``Request Raises`` in the documentation of the specific requests
            above.
    """
    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))
    request_task = None
    if response.status_code == 200:
        request_task = requests_lib.Request.decode(
            payloads.RequestPayload(**response.json()))
    elif response.status_code == 500:
        try:
            request_task = requests_lib.Request.decode(
                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
    if request_task is None:
        with ux_utils.print_exception_no_traceback():
            raise RuntimeError(f'Failed to get request {request_id}: '
                               f'{response.status_code} {response.text}')
    error = request_task.get_error()
    if error is not None:
        error_obj = error['object']
        _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(
                f'{colorama.Fore.YELLOW}Current {request_task.name!r} request '
                f'({request_task.request_id}) is cancelled by another process.'
                f'{colorama.Style.RESET_ALL}')
    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[server_common.RequestId[T]] = None,
    log_path: Optional[str] = None,
    tail: Optional[int] = None,
    follow: bool = True,
    output_stream: Optional['io.TextIOBase'] = None,
) -> 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 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.
        follow: Whether to follow the logs.
        output_stream: The output stream to write to. If None, print to the
            console.

    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,
            see ``Request Raises`` in the documentation of the specific requests
            above.
    """
    params = {
        'request_id': request_id,
        'log_path': log_path,
        'tail': str(tail) if tail is not None else None,
        'follow': follow,
        'format': 'console',
    }
    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)
    if response.status_code in [404, 400]:
        detail = response.json().get('detail')
        with ux_utils.print_exception_no_traceback():
            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,
                           resumable=True,
                           get_result=follow)


@usage_lib.entrypoint
@annotations.client_api
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[List[str]]:
    """Aborts a request or all requests.

    Args:
        request_ids: The request ID(s) to abort. Can be a single string or a
            list of strings.
        all_users: Whether to abort all requests from all users.
        silent: Whether to suppress the output.

    Returns:
        The request ID of the abort request itself.

    Request Returns:
        A list of request IDs that were cancelled.

    Raises:
        click.BadParameter: If no request ID is specified and not all or
            all_users is not set.
    """
    echo = logger.info if not silent else logger.debug
    user_id = None
    if not all_users:
        user_id = common_utils.get_user_hash()

    # Convert single request ID to list if needed
    if isinstance(request_ids, str):
        request_ids = [request_ids]

    body = payloads.RequestCancelBody(request_ids=request_ids, user_id=user_id)
    if all_users:
        echo('Cancelling all users\' requests...')
    elif request_ids is None:
        echo(f'Cancelling all requests for user {user_id!r}...')
    else:
        request_id_str = ', '.join(
            repr(request_id) for request_id in request_ids)
        plural = 's' if len(request_ids) > 1 else ''
        echo(f'Cancelling {len(request_ids)} request{plural}: '
             f'{request_id_str}...')

    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[Union[server_common.RequestId[T], str]]] = None,
    # pylint: disable=redefined-builtin
    all_status: bool = False,
    limit: Optional[int] = None,
    fields: Optional[List[str]] = None,
) -> List[payloads.RequestPayload]:
    """Lists all requests.

    Args:
        request_ids: The prefixes of the request IDs of the requests to query.
            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.
    """
    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))
    server_common.handle_request_error(response)
    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() -> responses.APIHealthResponse:
    """Gets the server's status, commit and version.

    Returns:
        A dictionary containing the server's status, commit and version.

        .. code-block:: python

            {
                'status': 'healthy',
                '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 = server_common.make_authenticated_request('GET', '/api/health')
    response.raise_for_status()
    api_health_response = responses.APIHealthResponse(**response.json())

    return api_health_response


@usage_lib.entrypoint
@annotations.client_api
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.

    It checks the existence of the API server and starts it if it does not
    exist.

    Args:
        deploy: Whether to deploy the API server, i.e. fully utilize the
            resources of the machine.
        host: The host to deploy the API server. It will be set to 0.0.0.0
            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
    """
    if deploy:
        host = '0.0.0.0'
    if host not in server_common.AVAILBLE_LOCAL_API_SERVER_HOSTS:
        raise ValueError(f'Invalid host: {host}. Should be one of: '
                         f'{server_common.AVAILBLE_LOCAL_API_SERVER_HOSTS}')
    is_local_api_server = server_common.is_api_server_local()
    if not is_local_api_server:
        server_url = server_common.get_server_url()
        with ux_utils.print_exception_no_traceback():
            raise ValueError(f'Unable to start local API server: '
                             f'server endpoint is set to {server_url}. '
                             'To start a local API server, remove the endpoint '
                             '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,
                                                   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)
    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}')


@usage_lib.entrypoint
@annotations.client_api
def api_stop() -> None:
    """Stops the API server.

    It will do nothing if the API server is remotely hosted.

    Returns:
        None
    """
    # Kill the uvicorn process by name: uvicorn sky.server.server:app
    server_url = server_common.get_server_url()
    if not server_common.is_api_server_local():
        with ux_utils.print_exception_no_traceback():
            raise RuntimeError(
                f'Cannot kill the API server at {server_url} because it is not '
                f'the default SkyPilot API server started locally.')

    # 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.'
                    f'{colorama.Style.RESET_ALL}')
    else:
        logger.info('SkyPilot API server is not running.')


# Use the same args as `docker logs`
@usage_lib.entrypoint
@annotations.client_api
def api_server_logs(follow: bool = True, tail: Optional[int] = None) -> None:
    """Streams the API server logs.

    Args:
        follow: Whether to follow the logs.
        tail: the number of lines to show from the end of the logs.
            If None, show all logs.

    Returns:
        None
    """
    if server_common.is_api_server_local():
        tail_args = ['-f'] if follow else []
        if tail is None:
            tail_args.extend(['-n', '+1'])
        else:
            tail_args.extend(['-n', f'{tail}'])
        log_path = os.path.expanduser(constants.API_SERVER_LOGS)
        subprocess.run(['tail', *tail_args, f'{log_path}'], check=False)
    else:
        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,
              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
    use this endpoint.

    To temporarily override the endpoint, use the environment variable
    `SKYPILOT_API_SERVER_ENDPOINT` instead.

    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
    """
    _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:
            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}')


@usage_lib.entrypoint
@server_common.check_server_healthy_or_start
@versions.minimal_api_version(24)
@annotations.client_api
def realtime_slurm_gpu_availability(
        name_filter: Optional[str] = None,
        quantity_filter: Optional[int] = None,
        slurm_cluster_name: Optional[str] = None) -> server_common.RequestId:
    """Gets the real-time Slurm GPU availability.

    Args:
        name_filter: Optional name filter for GPUs.
        quantity_filter: Optional quantity filter for GPUs.
        slurm_cluster_name: Optional Slurm cluster name to filter by.

    Returns:
        The request ID of the Slurm GPU availability request.
    """
    remote_api_version = versions.get_remote_api_version()
    # TODO(kevin): remove this in v0.13.0
    if (slurm_cluster_name is not None and remote_api_version is not None and
            remote_api_version < 27):
        logger.warning(
            'The Slurm cluster filter is not supported in your API server; '
            'the server will ignore it and show all Slurm clusters. '
            'Please upgrade the API server to enable it.')

    body = payloads.SlurmGpuAvailabilityRequestBody(
        slurm_cluster_name=slurm_cluster_name,
        name_filter=name_filter,
        quantity_filter=quantity_filter,
    )
    response = server_common.make_authenticated_request(
        'POST',
        '/slurm_gpu_availability',
        json=json.loads(body.model_dump_json()),
    )
    return server_common.get_request_id(response)


@usage_lib.entrypoint
@server_common.check_server_healthy_or_start
@versions.minimal_api_version(24)
@annotations.client_api
def slurm_node_info(
        slurm_cluster_name: Optional[str] = None) -> server_common.RequestId:
    """Gets the resource information for all nodes in the Slurm cluster.

    Returns:
        The request ID of the Slurm node info request.

    Request Returns:
        List[Dict[str, Any]]: A list of dictionaries, each containing info
            for a single Slurm node (node_name, partition, node_state,
            gpu_type, total_gpus, free_gpus, vcpu_count, memory_gb).
    """
    body = payloads.SlurmNodeInfoRequestBody(
        slurm_cluster_name=slurm_cluster_name)
    response = server_common.make_authenticated_request(
        'GET',
        '/slurm_node_info',
        json=json.loads(body.model_dump_json()),
    )
    return server_common.get_request_id(response)
