skypilot-nightly 1.0.0.dev20250905__py3-none-any.whl → 1.0.0.dev20251210__py3-none-any.whl

sky/client/sdk.py

@@ -31,14 +31,18 @@ from sky import skypilot_config
 from sky.adaptors import common as adaptors_common
 from sky.client import common as client_common
 from sky.client import oauth as oauth_lib
+from sky.jobs import scheduler
+from sky.jobs import utils as managed_job_utils
 from sky.schemas.api import responses
 from sky.server import common as server_common
 from sky.server import rest
 from sky.server import versions
 from sky.server.requests import payloads
+from sky.server.requests import request_names
 from sky.server.requests import requests as requests_lib
 from sky.skylet import autostop_lib
 from sky.skylet import constants
+from sky.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
@@ -54,7 +58,6 @@ from sky.utils import status_lib
 from sky.utils import subprocess_utils
 from sky.utils import ux_utils
 from sky.utils import yaml_utils
-from sky.utils.kubernetes import ssh_utils
 
 if typing.TYPE_CHECKING:
     import base64
@@ -97,6 +100,9 @@ def reload_config() -> None:
     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',
@@ -111,7 +117,16 @@ 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) -> T:
+                    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]:
     ...
 
 
@@ -367,6 +382,16 @@ def workspaces() -> server_common.RequestId[Dict[str, Any]]:
     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
@@ -407,9 +432,8 @@ def validate(
     response = server_common.make_authenticated_request(
         'POST', '/validate', json=json.loads(body.model_dump_json()))
     if response.status_code == 400:
-        with ux_utils.print_exception_no_traceback():
-            raise exceptions.deserialize_exception(
-                response.json().get('detail'))
+        _raise_exception_object_on_client(
+            exceptions.deserialize_exception(response.json().get('detail')))
 
 
 @usage_lib.entrypoint
@@ -590,7 +614,10 @@ def launch(
         down=down,
         dryrun=dryrun)
     with admin_policy_utils.apply_and_use_config_in_current_request(
-            dag, request_options=request_options, at_client_side=True) as dag:
+            dag,
+            request_name=request_names.AdminPolicyRequestName.CLUSTER_LAUNCH,
+            request_options=request_options,
+            at_client_side=True) as dag:
         return _launch(
             dag,
             cluster_name,
@@ -648,7 +675,7 @@ def _launch(
         clusters = get(status_request_id)
         cluster_user_hash = common_utils.get_user_hash()
         cluster_user_hash_str = ''
-        current_user = common_utils.get_current_user_name()
+        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
@@ -912,6 +939,7 @@ def tail_logs(
 @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:
@@ -919,17 +947,31 @@ def tail_provision_logs(cluster_name: str,
 
     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.ClusterNameBody(cluster_name=cluster_name)
+    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',
@@ -938,13 +980,21 @@ def tail_provision_logs(cluster_name: str,
         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
+    # 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(
@@ -1266,9 +1316,11 @@ def autostop(
 @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[dict]]:
+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:
@@ -1281,8 +1333,8 @@ def queue(cluster_name: str,
         The request ID of the queue request.
 
     Request Returns:
-        job_records (List[Dict[str, Any]]): A list of dicts for each job in the
-            queue.
+        job_records (List[responses.ClusterJobRecord]): A list of job records
+            for each job in the queue.
 
             .. code-block:: python
 
@@ -1428,6 +1480,7 @@ def status(
     all_users: bool = False,
     *,
     _include_credentials: bool = False,
+    _summary_response: bool = False,
 ) -> server_common.RequestId[List[responses.StatusResponse]]:
     """Gets cluster statuses.
 
@@ -1513,6 +1566,7 @@ def status(
         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()))
@@ -1613,26 +1667,15 @@ def cost_report(
 @usage_lib.entrypoint
 @server_common.check_server_healthy_or_start
 @annotations.client_api
-def storage_ls() -> server_common.RequestId[List[Dict[str, Any]]]:
+def storage_ls() -> server_common.RequestId[List[responses.StorageRecord]]:
     """Gets the storages.
 
     Returns:
         The request ID of the storage list request.
 
     Request Returns:
-        storage_records (List[Dict[str, Any]]): A list of dicts, with each dict
-            containing the information of a storage.
-
-            .. code-block:: python
-
-                {
-                    'name': (str) storage name,
-                    'launched_at': (int) timestamp of creation,
-                    'store': (List[sky.StoreType]) storage type,
-                    'last_use': (int) timestamp of last use,
-                    'status': (sky.StorageStatus) storage status,
-                }
-        ]
+        storage_records (List[responses.StorageRecord]):
+            A list of storage records.
     """
     response = server_common.make_authenticated_request('GET', '/storage/ls')
     return server_common.get_request_id(response)
@@ -1669,12 +1712,8 @@ def storage_delete(name: str) -> server_common.RequestId[None]:
 @server_common.check_server_healthy_or_start
 @annotations.client_api
 def local_up(gpus: bool,
-             ips: Optional[List[str]],
-             ssh_user: Optional[str],
-             ssh_key: Optional[str],
-             cleanup: bool,
-             context_name: Optional[str] = None,
-             password: Optional[str] = None) -> server_common.RequestId[None]:
+             name: Optional[str] = None,
+             port_start: Optional[int] = None) -> server_common.RequestId[None]:
     """Launches a Kubernetes cluster on local machines.
 
     Returns:
@@ -1685,16 +1724,10 @@ def local_up(gpus: bool,
     # TODO: move this check to server.
     if not server_common.is_api_server_local():
         with ux_utils.print_exception_no_traceback():
-            raise ValueError(
-                'sky local up is only supported when running SkyPilot locally.')
-
-    body = payloads.LocalUpBody(gpus=gpus,
-                                ips=ips,
-                                ssh_user=ssh_user,
-                                ssh_key=ssh_key,
-                                cleanup=cleanup,
-                                context_name=context_name,
-                                password=password)
+            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)
@@ -1703,16 +1736,19 @@ def local_up(gpus: bool,
 @usage_lib.entrypoint
 @server_common.check_server_healthy_or_start
 @annotations.client_api
-def local_down() -> server_common.RequestId[None]:
+def local_down(name: Optional[str]) -> server_common.RequestId[None]:
     """Tears down the Kubernetes cluster started by local_up."""
     # We do not allow local up when the API server is running remotely since it
     # will modify the kubeconfig.
     # TODO: move this check to remote server.
     if not server_common.is_api_server_local():
         with ux_utils.print_exception_no_traceback():
-            raise ValueError('sky local down is only supported when running '
+            raise ValueError('`sky local down` is only supported when running '
                              'SkyPilot locally.')
-    response = server_common.make_authenticated_request('POST', '/local_down')
+
+    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)
 
 
@@ -1900,11 +1936,12 @@ def kubernetes_node_info(
 @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[Dict[
-        str, Any]], Optional[str]]]:
-    """Gets all SkyPilot clusters and jobs in the Kubernetes cluster.
+def status_kubernetes() -> server_common.RequestId[
+    Tuple[List['kubernetes_utils.KubernetesSkyPilotClusterInfoPayload'],
+          List['kubernetes_utils.KubernetesSkyPilotClusterInfoPayload'],
+          List[responses.ManagedJobRecord], Optional[str]]]:
+    """[Experimental] Gets all SkyPilot clusters and jobs
+    in the Kubernetes cluster.
 
     Managed jobs and services are also included in the clusters returned.
     The caller must parse the controllers to identify which clusters are run
@@ -1976,12 +2013,7 @@ def get(request_id: server_common.RequestId[T]) -> T:
     error = request_task.get_error()
     if error is not None:
         error_obj = error['object']
-        if env_options.Options.SHOW_DEBUG_INFO.get():
-            stacktrace = getattr(error_obj, 'stacktrace', str(error_obj))
-            logger.error('=== Traceback on SkyPilot API Server ===\n'
-                         f'{stacktrace}')
-        with ux_utils.print_exception_no_traceback():
-            raise error_obj
+        _raise_exception_object_on_client(error_obj)
     if request_task.status == requests_lib.RequestStatus.CANCELLED:
         with ux_utils.print_exception_no_traceback():
             raise exceptions.RequestCancelled(
@@ -2067,6 +2099,12 @@ def stream_and_get(
         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
@@ -2076,6 +2114,7 @@ def stream_and_get(
     return stream_response(request_id,
                            response,
                            output_stream,
+                           resumable=True,
                            get_result=follow)
 
 
@@ -2150,7 +2189,9 @@ def _local_api_server_running(kill: bool = False) -> bool:
 def api_status(
     request_ids: Optional[List[Union[server_common.RequestId[T], str]]] = None,
     # pylint: disable=redefined-builtin
-    all_status: bool = False
+    all_status: bool = False,
+    limit: Optional[int] = None,
+    fields: Optional[List[str]] = None,
 ) -> List[payloads.RequestPayload]:
     """Lists all requests.
 
@@ -2159,6 +2200,8 @@ def api_status(
             If None, all requests are queried.
         all_status: Whether to list all finished requests as well. This argument
             is ignored if request_ids is not None.
+        limit: The number of requests to show. If None, show all requests.
+        fields: The fields to get. If None, get all fields.
 
     Returns:
         A list of request payloads.
@@ -2167,8 +2210,12 @@ def api_status(
         logger.info('SkyPilot API server is not running.')
         return []
 
-    body = payloads.RequestStatusBody(request_ids=request_ids,
-                                      all_status=all_status)
+    body = payloads.RequestStatusBody(
+        request_ids=request_ids,
+        all_status=all_status,
+        limit=limit,
+        fields=fields,
+    )
     response = server_common.make_authenticated_request(
         'GET',
         '/api/status',
@@ -2287,10 +2334,32 @@ def api_stop() -> None:
                 f'Cannot kill the API server at {server_url} because it is not '
                 f'the default SkyPilot API server started locally.')
 
-    found = _local_api_server_running(kill=True)
+    # 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
 
-    # Remove the database for requests.
-    server_common.clear_local_api_server_database()
+        found = _local_api_server_running(kill=True)
 
     if found:
         logger.info(f'{colorama.Fore.GREEN}SkyPilot API server stopped.'
@@ -2675,3 +2744,57 @@ def api_logout() -> None:
     _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) -> 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.
+
+    Returns:
+        The request ID of the Slurm GPU availability request.
+    """
+    body = payloads.SlurmGpuAvailabilityRequestBody(
+        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)

sky/client/sdk_async.py

@@ -19,20 +19,16 @@ import aiohttp
 import colorama
 
 from sky import admin_policy
-from sky import backends
 from sky import catalog
 from sky import exceptions
-from sky import models
 from sky import sky_logging
 from sky.client import common as client_common
 from sky.client import sdk
-from sky.provision.kubernetes import utils as kubernetes_utils
 from sky.schemas.api import responses
 from sky.server import common as server_common
 from sky.server import rest
 from sky.server.requests import payloads
 from sky.server.requests import requests as requests_lib
-from sky.skylet import job_lib
 from sky.usage import usage_lib
 from sky.utils import annotations
 from sky.utils import common
@@ -45,6 +41,11 @@ if typing.TYPE_CHECKING:
     import io
 
     import sky
+    from sky import backends
+    from sky import models
+    from sky.provision.kubernetes import utils as kubernetes_utils
+    from sky.skylet import autostop_lib
+    from sky.skylet import job_lib
 
 logger = sky_logging.init_logger(__name__)
 logging.getLogger('httpx').setLevel(logging.CRITICAL)
@@ -381,9 +382,10 @@ async def launch(
     cluster_name: Optional[str] = None,
     retry_until_up: bool = False,
     idle_minutes_to_autostop: Optional[int] = None,
+    wait_for: Optional['autostop_lib.AutostopWaitFor'] = None,
     dryrun: bool = False,
     down: bool = False,  # pylint: disable=redefined-outer-name
-    backend: Optional[backends.Backend] = None,
+    backend: Optional['backends.Backend'] = None,
     optimize_target: common.OptimizeTarget = common.OptimizeTarget.COST,
     no_setup: bool = False,
     clone_disk_from: Optional[str] = None,
@@ -395,12 +397,12 @@ async def launch(
     _is_launched_by_sky_serve_controller: bool = False,
     _disable_controller_check: bool = False,
     stream_logs: Optional[StreamConfig] = DEFAULT_STREAM_CONFIG,
-) -> Tuple[Optional[int], Optional[backends.ResourceHandle]]:
+) -> Tuple[Optional[int], Optional['backends.ResourceHandle']]:
     """Async version of launch() that launches a cluster or task."""
     request_id = await context_utils.to_thread(
         sdk.launch, task, cluster_name, retry_until_up,
-        idle_minutes_to_autostop, dryrun, down, backend, optimize_target,
-        no_setup, clone_disk_from, fast, _need_confirmation,
+        idle_minutes_to_autostop, wait_for, 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)
     if stream_logs is not None:
@@ -416,9 +418,9 @@ async def exec(  # pylint: disable=redefined-builtin
     cluster_name: Optional[str] = None,
     dryrun: bool = False,
     down: bool = False,  # pylint: disable=redefined-outer-name
-    backend: Optional[backends.Backend] = None,
+    backend: Optional['backends.Backend'] = None,
     stream_logs: Optional[StreamConfig] = DEFAULT_STREAM_CONFIG,
-) -> Tuple[Optional[int], Optional[backends.ResourceHandle]]:
+) -> Tuple[Optional[int], Optional['backends.ResourceHandle']]:
     """Async version of exec() that executes a task on an existing cluster."""
     request_id = await context_utils.to_thread(sdk.exec, task, cluster_name,
                                                dryrun, down, backend)
@@ -454,15 +456,17 @@ async def download_logs(cluster_name: str,
 async 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,
     stream_logs: Optional[StreamConfig] = DEFAULT_STREAM_CONFIG,
-) -> backends.CloudVmRayResourceHandle:
+) -> 'backends.CloudVmRayResourceHandle':
     """Async version of start() that restarts a cluster."""
     request_id = await context_utils.to_thread(sdk.start, cluster_name,
                                                idle_minutes_to_autostop,
-                                               retry_until_up, down, force)
+                                               wait_for, retry_until_up, down,
+                                               force)
     if stream_logs is not None:
         return await _stream_and_get(request_id, stream_logs)
     else:
@@ -502,13 +506,14 @@ async def stop(
 async def autostop(
     cluster_name: str,
     idle_minutes: int,
+    wait_for: Optional['autostop_lib.AutostopWaitFor'] = None,
     down: bool = False,  # pylint: disable=redefined-outer-name
     stream_logs: Optional[StreamConfig] = DEFAULT_STREAM_CONFIG
 ) -> None:
     """Async version of autostop() that schedules an autostop/autodown for a
       cluster."""
     request_id = await context_utils.to_thread(sdk.autostop, cluster_name,
-                                               idle_minutes, down)
+                                               idle_minutes, wait_for, down)
     if stream_logs is not None:
         return await _stream_and_get(request_id, stream_logs)
     else:
@@ -518,11 +523,11 @@ async def autostop(
 @usage_lib.entrypoint
 @annotations.client_api
 async def queue(
-        cluster_name: str,
-        skip_finished: bool = False,
-        all_users: bool = False,
-        stream_logs: Optional[StreamConfig] = DEFAULT_STREAM_CONFIG
-) -> List[dict]:
+    cluster_name: str,
+    skip_finished: bool = False,
+    all_users: bool = False,
+    stream_logs: Optional[StreamConfig] = DEFAULT_STREAM_CONFIG
+) -> List[responses.ClusterJobRecord]:
     """Async version of queue() that gets the job queue of a cluster."""
     request_id = await context_utils.to_thread(sdk.queue, cluster_name,
                                                skip_finished, all_users)
@@ -538,7 +543,7 @@ async def job_status(
     cluster_name: str,
     job_ids: Optional[List[int]] = None,
     stream_logs: Optional[StreamConfig] = DEFAULT_STREAM_CONFIG
-) -> Dict[Optional[int], Optional[job_lib.JobStatus]]:
+) -> Dict[Optional[int], Optional['job_lib.JobStatus']]:
     """Async version of job_status() that gets the status of jobs on a
       cluster."""
     request_id = await context_utils.to_thread(sdk.job_status, cluster_name,
@@ -651,18 +656,13 @@ async def storage_delete(
 @annotations.client_api
 async def local_up(
         gpus: bool,
-        ips: Optional[List[str]],
-        ssh_user: Optional[str],
-        ssh_key: Optional[str],
-        cleanup: bool,
-        context_name: Optional[str] = None,
-        password: Optional[str] = None,
+        name: Optional[str] = None,
+        port_start: Optional[int] = None,
         stream_logs: Optional[StreamConfig] = DEFAULT_STREAM_CONFIG) -> None:
     """Async version of local_up() that launches a Kubernetes cluster on
     local machines."""
-    request_id = await context_utils.to_thread(sdk.local_up, gpus, ips,
-                                               ssh_user, ssh_key, cleanup,
-                                               context_name, password)
+    request_id = await context_utils.to_thread(sdk.local_up, gpus, name,
+                                               port_start)
     if stream_logs is not None:
         return await _stream_and_get(request_id, stream_logs)
     else:
@@ -672,10 +672,11 @@ async def local_up(
 @usage_lib.entrypoint
 @annotations.client_api
 async def local_down(
+        name: Optional[str] = None,
         stream_logs: Optional[StreamConfig] = DEFAULT_STREAM_CONFIG) -> None:
     """Async version of local_down() that tears down the Kubernetes cluster
     started by local_up."""
-    request_id = await context_utils.to_thread(sdk.local_down)
+    request_id = await context_utils.to_thread(sdk.local_down, name)
     if stream_logs is not None:
         return await _stream_and_get(request_id, stream_logs)
     else:
@@ -718,7 +719,7 @@ async def realtime_kubernetes_gpu_availability(
     quantity_filter: Optional[int] = None,
     is_ssh: Optional[bool] = None,
     stream_logs: Optional[StreamConfig] = DEFAULT_STREAM_CONFIG
-) -> List[Tuple[str, List[models.RealtimeGpuAvailability]]]:
+) -> List[Tuple[str, List['models.RealtimeGpuAvailability']]]:
     """Async version of realtime_kubernetes_gpu_availability() that gets the
       real-time Kubernetes GPU availability."""
     request_id = await context_utils.to_thread(
@@ -735,7 +736,7 @@ async def realtime_kubernetes_gpu_availability(
 async def kubernetes_node_info(
     context: Optional[str] = None,
     stream_logs: Optional[StreamConfig] = DEFAULT_STREAM_CONFIG
-) -> models.KubernetesNodesInfo:
+) -> 'models.KubernetesNodesInfo':
     """Async version of kubernetes_node_info() that gets the resource
     information for all the nodes in the cluster."""
     request_id = await context_utils.to_thread(sdk.kubernetes_node_info,
@@ -750,8 +751,8 @@ async def kubernetes_node_info(
 @annotations.client_api
 async def status_kubernetes(
     stream_logs: Optional[StreamConfig] = DEFAULT_STREAM_CONFIG
-) -> Tuple[List[kubernetes_utils.KubernetesSkyPilotClusterInfoPayload],
-           List[kubernetes_utils.KubernetesSkyPilotClusterInfoPayload],
+) -> Tuple[List['kubernetes_utils.KubernetesSkyPilotClusterInfoPayload'],
+           List['kubernetes_utils.KubernetesSkyPilotClusterInfoPayload'],
            List[Dict[str, Any]], Optional[str]]:
     """Async version of status_kubernetes() that gets all SkyPilot clusters
       and jobs in the Kubernetes cluster."""