skypilot-nightly 1.0.0.dev20250624__py3-none-any.whl → 1.0.0.dev20250626__py3-none-any.whl

sky/client/sdk.py

@@ -37,6 +37,7 @@ 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.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 constants
@@ -54,6 +55,7 @@ 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.kubernetes import ssh_utils
 
 if typing.TYPE_CHECKING:
     import io
@@ -64,15 +66,17 @@ if typing.TYPE_CHECKING:
     import sky
 else:
     psutil = adaptors_common.LazyImport('psutil')
-    requests = adaptors_common.LazyImport('requests')
 
 logger = sky_logging.init_logger(__name__)
 logging.getLogger('httpx').setLevel(logging.CRITICAL)
 
+_LINE_PROCESSED_KEY = 'line_processed'
+
 
 def stream_response(request_id: Optional[str],
                     response: 'requests.Response',
-                    output_stream: Optional['io.TextIOBase'] = None) -> Any:
+                    output_stream: Optional['io.TextIOBase'] = None,
+                    resumable: bool = False) -> Any:
     """Streams the response to the console.
 
     Args:
@@ -80,12 +84,23 @@ def stream_response(request_id: Optional[str],
         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.
     """
 
+    retry_context: Optional[rest.RetryContext] = None
+    if resumable:
+        retry_context = rest.get_retry_context()
     try:
+        line_count = 0
         for line in rich_utils.decode_rich_status(response):
             if line is not None:
-                print(line, flush=True, end='', file=output_stream)
+                line_count += 1
+                if retry_context is None:
+                    print(line, flush=True, end='', file=output_stream)
+                elif line_count > retry_context.line_processed:
+                    print(line, flush=True, end='', file=output_stream)
+                    retry_context.line_processed = line_count
         if request_id is not None:
             return get(request_id)
     except Exception:  # pylint: disable=broad-except
@@ -132,9 +147,9 @@ def check(infra_list: Optional[Tuple[str, ...]],
     body = payloads.CheckBody(clouds=clouds,
                               verbose=verbose,
                               workspace=workspace)
-    response = requests.post(f'{server_common.get_server_url()}/check',
-                             json=json.loads(body.model_dump_json()),
-                             cookies=server_common.get_api_cookie_jar())
+    response = rest.post(f'{server_common.get_server_url()}/check',
+                         json=json.loads(body.model_dump_json()),
+                         cookies=server_common.get_api_cookie_jar())
     return server_common.get_request_id(response)
 
 
@@ -158,9 +173,9 @@ def enabled_clouds(workspace: Optional[str] = None,
     """
     if workspace is None:
         workspace = skypilot_config.get_active_workspace()
-    response = requests.get((f'{server_common.get_server_url()}/enabled_clouds?'
-                             f'workspace={workspace}&expand={expand}'),
-                            cookies=server_common.get_api_cookie_jar())
+    response = rest.get((f'{server_common.get_server_url()}/enabled_clouds?'
+                         f'workspace={workspace}&expand={expand}'),
+                        cookies=server_common.get_api_cookie_jar())
     return server_common.get_request_id(response)
 
 
@@ -208,10 +223,9 @@ def list_accelerators(gpus_only: bool = True,
         require_price=require_price,
         case_sensitive=case_sensitive,
     )
-    response = requests.post(
-        f'{server_common.get_server_url()}/list_accelerators',
-        json=json.loads(body.model_dump_json()),
-        cookies=server_common.get_api_cookie_jar())
+    response = rest.post(f'{server_common.get_server_url()}/list_accelerators',
+                         json=json.loads(body.model_dump_json()),
+                         cookies=server_common.get_api_cookie_jar())
     return server_common.get_request_id(response)
 
 
@@ -249,7 +263,7 @@ def list_accelerator_counts(
         quantity_filter=quantity_filter,
         clouds=clouds,
     )
-    response = requests.post(
+    response = rest.post(
         f'{server_common.get_server_url()}/list_accelerator_counts',
         json=json.loads(body.model_dump_json()),
         cookies=server_common.get_api_cookie_jar())
@@ -289,16 +303,16 @@ def optimize(
     body = payloads.OptimizeBody(dag=dag_str,
                                  minimize=minimize,
                                  request_options=admin_policy_request_options)
-    response = requests.post(f'{server_common.get_server_url()}/optimize',
-                             json=json.loads(body.model_dump_json()),
-                             cookies=server_common.get_api_cookie_jar())
+    response = rest.post(f'{server_common.get_server_url()}/optimize',
+                         json=json.loads(body.model_dump_json()),
+                         cookies=server_common.get_api_cookie_jar())
     return server_common.get_request_id(response)
 
 
 def workspaces() -> server_common.RequestId:
     """Gets the workspaces."""
-    response = requests.get(f'{server_common.get_server_url()}/workspaces',
-                            cookies=server_common.get_api_cookie_jar())
+    response = rest.get(f'{server_common.get_server_url()}/workspaces',
+                        cookies=server_common.get_api_cookie_jar())
     return server_common.get_request_id(response)
 
 
@@ -332,9 +346,9 @@ def validate(
     dag_str = dag_utils.dump_chain_dag_to_yaml_str(dag)
     body = payloads.ValidateBody(dag=dag_str,
                                  request_options=admin_policy_request_options)
-    response = requests.post(f'{server_common.get_server_url()}/validate',
-                             json=json.loads(body.model_dump_json()),
-                             cookies=server_common.get_api_cookie_jar())
+    response = rest.post(f'{server_common.get_server_url()}/validate',
+                         json=json.loads(body.model_dump_json()),
+                         cookies=server_common.get_api_cookie_jar())
     if response.status_code == 400:
         with ux_utils.print_exception_no_traceback():
             raise exceptions.deserialize_exception(
@@ -618,7 +632,7 @@ def _launch(
             _is_launched_by_sky_serve_controller),
         disable_controller_check=_disable_controller_check,
     )
-    response = requests.post(
+    response = rest.post(
         f'{server_common.get_server_url()}/launch',
         json=json.loads(body.model_dump_json()),
         timeout=5,
@@ -702,7 +716,7 @@ def exec(  # pylint: disable=redefined-builtin
         backend=backend.NAME if backend else None,
     )
 
-    response = requests.post(
+    response = rest.post(
         f'{server_common.get_server_url()}/exec',
         json=json.loads(body.model_dump_json()),
         timeout=5,
@@ -711,9 +725,12 @@ def exec(  # pylint: disable=redefined-builtin
     return server_common.get_request_id(response)
 
 
+# TODO(aylei): when retry logs request, there will be duplciated log entries.
+# We should fix this.
 @usage_lib.entrypoint
 @server_common.check_server_healthy_or_start
 @annotations.client_api
+@rest.retry_on_server_unavailable()
 def tail_logs(cluster_name: str,
               job_id: Optional[int],
               follow: bool,
@@ -752,7 +769,7 @@ def tail_logs(cluster_name: str,
         follow=follow,
         tail=tail,
     )
-    response = requests.post(
+    response = rest.post(
         f'{server_common.get_server_url()}/logs',
         json=json.loads(body.model_dump_json()),
         stream=True,
@@ -760,7 +777,12 @@ def tail_logs(cluster_name: str,
                  None),
         cookies=server_common.get_api_cookie_jar())
     request_id = server_common.get_request_id(response)
-    return stream_response(request_id, response, output_stream)
+    # 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))
 
 
 @usage_lib.entrypoint
@@ -794,9 +816,9 @@ def download_logs(cluster_name: str,
         cluster_name=cluster_name,
         job_ids=job_ids,
     )
-    response = requests.post(f'{server_common.get_server_url()}/download_logs',
-                             json=json.loads(body.model_dump_json()),
-                             cookies=server_common.get_api_cookie_jar())
+    response = rest.post(f'{server_common.get_server_url()}/download_logs',
+                         json=json.loads(body.model_dump_json()),
+                         cookies=server_common.get_api_cookie_jar())
     job_id_remote_path_dict = stream_and_get(
         server_common.get_request_id(response))
     remote2local_path_dict = client_common.download_logs_from_api_server(
@@ -874,7 +896,7 @@ def start(
         down=down,
         force=force,
     )
-    response = requests.post(
+    response = rest.post(
         f'{server_common.get_server_url()}/start',
         json=json.loads(body.model_dump_json()),
         timeout=5,
@@ -920,7 +942,7 @@ def down(cluster_name: str, purge: bool = False) -> server_common.RequestId:
         cluster_name=cluster_name,
         purge=purge,
     )
-    response = requests.post(
+    response = rest.post(
         f'{server_common.get_server_url()}/down',
         json=json.loads(body.model_dump_json()),
         timeout=5,
@@ -969,7 +991,7 @@ def stop(cluster_name: str, purge: bool = False) -> server_common.RequestId:
         cluster_name=cluster_name,
         purge=purge,
     )
-    response = requests.post(
+    response = rest.post(
         f'{server_common.get_server_url()}/stop',
         json=json.loads(body.model_dump_json()),
         timeout=5,
@@ -1039,7 +1061,7 @@ def autostop(
         idle_minutes=idle_minutes,
         down=down,
     )
-    response = requests.post(
+    response = rest.post(
         f'{server_common.get_server_url()}/autostop',
         json=json.loads(body.model_dump_json()),
         timeout=5,
@@ -1102,9 +1124,9 @@ def queue(cluster_name: str,
         skip_finished=skip_finished,
         all_users=all_users,
     )
-    response = requests.post(f'{server_common.get_server_url()}/queue',
-                             json=json.loads(body.model_dump_json()),
-                             cookies=server_common.get_api_cookie_jar())
+    response = rest.post(f'{server_common.get_server_url()}/queue',
+                         json=json.loads(body.model_dump_json()),
+                         cookies=server_common.get_api_cookie_jar())
     return server_common.get_request_id(response)
 
 
@@ -1144,9 +1166,9 @@ def job_status(cluster_name: str,
         cluster_name=cluster_name,
         job_ids=job_ids,
     )
-    response = requests.post(f'{server_common.get_server_url()}/job_status',
-                             json=json.loads(body.model_dump_json()),
-                             cookies=server_common.get_api_cookie_jar())
+    response = rest.post(f'{server_common.get_server_url()}/job_status',
+                         json=json.loads(body.model_dump_json()),
+                         cookies=server_common.get_api_cookie_jar())
     return server_common.get_request_id(response)
 
 
@@ -1198,9 +1220,9 @@ def cancel(
         job_ids=job_ids,
         try_cancel_if_cluster_is_init=_try_cancel_if_cluster_is_init,
     )
-    response = requests.post(f'{server_common.get_server_url()}/cancel',
-                             json=json.loads(body.model_dump_json()),
-                             cookies=server_common.get_api_cookie_jar())
+    response = rest.post(f'{server_common.get_server_url()}/cancel',
+                         json=json.loads(body.model_dump_json()),
+                         cookies=server_common.get_api_cookie_jar())
     return server_common.get_request_id(response)
 
 
@@ -1294,9 +1316,9 @@ def status(
         refresh=refresh,
         all_users=all_users,
     )
-    response = requests.post(f'{server_common.get_server_url()}/status',
-                             json=json.loads(body.model_dump_json()),
-                             cookies=server_common.get_api_cookie_jar())
+    response = rest.post(f'{server_common.get_server_url()}/status',
+                         json=json.loads(body.model_dump_json()),
+                         cookies=server_common.get_api_cookie_jar())
     return server_common.get_request_id(response)
 
 
@@ -1329,9 +1351,9 @@ def endpoints(
         cluster=cluster,
         port=port,
     )
-    response = requests.post(f'{server_common.get_server_url()}/endpoints',
-                             json=json.loads(body.model_dump_json()),
-                             cookies=server_common.get_api_cookie_jar())
+    response = rest.post(f'{server_common.get_server_url()}/endpoints',
+                         json=json.loads(body.model_dump_json()),
+                         cookies=server_common.get_api_cookie_jar())
     return server_common.get_request_id(response)
 
 
@@ -1374,9 +1396,9 @@ def cost_report(days: Optional[int] = None) -> server_common.RequestId:  # pylin
             }
     """
     body = payloads.CostReportBody(days=days)
-    response = requests.post(f'{server_common.get_server_url()}/cost_report',
-                             json=json.loads(body.model_dump_json()),
-                             cookies=server_common.get_api_cookie_jar())
+    response = rest.post(f'{server_common.get_server_url()}/cost_report',
+                         json=json.loads(body.model_dump_json()),
+                         cookies=server_common.get_api_cookie_jar())
     return server_common.get_request_id(response)
 
 
@@ -1405,8 +1427,8 @@ def storage_ls() -> server_common.RequestId:
                 }
         ]
     """
-    response = requests.get(f'{server_common.get_server_url()}/storage/ls',
-                            cookies=server_common.get_api_cookie_jar())
+    response = rest.get(f'{server_common.get_server_url()}/storage/ls',
+                        cookies=server_common.get_api_cookie_jar())
     return server_common.get_request_id(response)
 
 
@@ -1429,9 +1451,9 @@ def storage_delete(name: str) -> server_common.RequestId:
         ValueError: If the storage does not exist.
     """
     body = payloads.StorageBody(name=name)
-    response = requests.post(f'{server_common.get_server_url()}/storage/delete',
-                             json=json.loads(body.model_dump_json()),
-                             cookies=server_common.get_api_cookie_jar())
+    response = rest.post(f'{server_common.get_server_url()}/storage/delete',
+                         json=json.loads(body.model_dump_json()),
+                         cookies=server_common.get_api_cookie_jar())
     return server_common.get_request_id(response)
 
 
@@ -1468,9 +1490,9 @@ def local_up(gpus: bool,
                                 cleanup=cleanup,
                                 context_name=context_name,
                                 password=password)
-    response = requests.post(f'{server_common.get_server_url()}/local_up',
-                             json=json.loads(body.model_dump_json()),
-                             cookies=server_common.get_api_cookie_jar())
+    response = rest.post(f'{server_common.get_server_url()}/local_up',
+                         json=json.loads(body.model_dump_json()),
+                         cookies=server_common.get_api_cookie_jar())
     return server_common.get_request_id(response)
 
 
@@ -1486,31 +1508,100 @@ def local_down() -> server_common.RequestId:
         with ux_utils.print_exception_no_traceback():
             raise ValueError('sky local down is only supported when running '
                              'SkyPilot locally.')
-    response = requests.post(f'{server_common.get_server_url()}/local_down',
-                             cookies=server_common.get_api_cookie_jar())
+    response = rest.post(f'{server_common.get_server_url()}/local_down',
+                         cookies=server_common.get_api_cookie_jar())
     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}
+    rest.post(f'{server_common.get_server_url()}/ssh_node_pools',
+              json=pools_config,
+              cookies=server_common.get_api_cookie_jar())
+
+
+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 = rest.post(
+            f'{server_common.get_server_url()}/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) -> server_common.RequestId:
+def ssh_up(infra: Optional[str] = None,
+           file: Optional[str] = None) -> server_common.RequestId:
     """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.
     """
-    body = payloads.SSHUpBody(
-        infra=infra,
-        cleanup=False,
-    )
-    response = requests.post(f'{server_common.get_server_url()}/ssh_up',
-                             json=json.loads(body.model_dump_json()),
-                             cookies=server_common.get_api_cookie_jar())
+    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 = rest.post(
+            f'{server_common.get_server_url()}/ssh_node_pools/{infra}/deploy',
+            cookies=server_common.get_api_cookie_jar())
+    else:
+        # Call the general deployment endpoint
+        response = rest.post(
+            f'{server_common.get_server_url()}/ssh_node_pools/deploy',
+            json=json.loads(body.model_dump_json()),
+            cookies=server_common.get_api_cookie_jar())
     return server_common.get_request_id(response)
 
 
@@ -1527,13 +1618,19 @@ def ssh_down(infra: Optional[str] = None) -> server_common.RequestId:
     Returns:
         request_id: The request ID of the SSH cluster teardown request.
     """
-    body = payloads.SSHUpBody(
-        infra=infra,
-        cleanup=True,
-    )
-    response = requests.post(f'{server_common.get_server_url()}/ssh_down',
-                             json=json.loads(body.model_dump_json()),
-                             cookies=server_common.get_api_cookie_jar())
+    # 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 = rest.post(
+            f'{server_common.get_server_url()}/ssh_node_pools/{infra}/down',
+            cookies=server_common.get_api_cookie_jar())
+    else:
+        # Call the general down endpoint
+        response = rest.post(
+            f'{server_common.get_server_url()}/ssh_node_pools/down',
+            json=json.loads(body.model_dump_json()),
+            cookies=server_common.get_api_cookie_jar())
     return server_common.get_request_id(response)
 
 
@@ -1556,7 +1653,7 @@ def realtime_kubernetes_gpu_availability(
         quantity_filter=quantity_filter,
         is_ssh=is_ssh,
     )
-    response = requests.post(
+    response = rest.post(
         f'{server_common.get_server_url()}/'
         'realtime_kubernetes_gpu_availability',
         json=json.loads(body.model_dump_json()),
@@ -1589,7 +1686,7 @@ def kubernetes_node_info(
             information.
     """
     body = payloads.KubernetesNodeInfoRequestBody(context=context)
-    response = requests.post(
+    response = rest.post(
         f'{server_common.get_server_url()}/kubernetes_node_info',
         json=json.loads(body.model_dump_json()),
         cookies=server_common.get_api_cookie_jar())
@@ -1620,19 +1717,21 @@ def status_kubernetes() -> server_common.RequestId:
             dictionary job info, see jobs.queue_from_kubernetes_pod for details.
         - context: Kubernetes context used to fetch the cluster information.
     """
-    response = requests.get(
-        f'{server_common.get_server_url()}/status_kubernetes',
-        cookies=server_common.get_api_cookie_jar())
+    response = rest.get(f'{server_common.get_server_url()}/status_kubernetes',
+                        cookies=server_common.get_api_cookie_jar())
     return server_common.get_request_id(response)
 
 
 # === API request APIs ===
 @usage_lib.entrypoint
-@server_common.check_server_healthy_or_start
 @annotations.client_api
 def get(request_id: str) -> Any:
     """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.
 
@@ -1645,7 +1744,7 @@ def get(request_id: str) -> Any:
             see ``Request Raises`` in the documentation of the specific requests
             above.
     """
-    response = requests.get(
+    response = rest.get_without_retry(
         f'{server_common.get_server_url()}/api/get?request_id={request_id}',
         timeout=(client_common.API_SERVER_REQUEST_CONNECTION_TIMEOUT_SECONDS,
                  None),
@@ -1723,7 +1822,7 @@ def stream_and_get(
         'follow': follow,
         'format': 'console',
     }
-    response = requests.get(
+    response = rest.get_without_retry(
         f'{server_common.get_server_url()}/api/stream',
         params=params,
         timeout=(client_common.API_SERVER_REQUEST_CONNECTION_TIMEOUT_SECONDS,
@@ -1783,10 +1882,10 @@ def api_cancel(request_ids: Optional[Union[str, List[str]]] = None,
         echo(f'Cancelling {len(request_ids)} request{plural}: '
              f'{request_id_str}...')
 
-    response = requests.post(f'{server_common.get_server_url()}/api/cancel',
-                             json=json.loads(body.model_dump_json()),
-                             timeout=5,
-                             cookies=server_common.get_api_cookie_jar())
+    response = rest.post(f'{server_common.get_server_url()}/api/cancel',
+                         json=json.loads(body.model_dump_json()),
+                         timeout=5,
+                         cookies=server_common.get_api_cookie_jar())
     return server_common.get_request_id(response)
 
 
@@ -1810,7 +1909,7 @@ def api_status(
     """
     body = payloads.RequestStatusBody(request_ids=request_ids,
                                       all_status=all_status)
-    response = requests.get(
+    response = rest.get(
         f'{server_common.get_server_url()}/api/status',
         params=server_common.request_body_to_params(body),
         timeout=(client_common.API_SERVER_REQUEST_CONNECTION_TIMEOUT_SECONDS,
@@ -1849,8 +1948,8 @@ def api_info() -> Dict[str, Any]:
         Note that user may be None if we are not using an auth proxy.
 
     """
-    response = requests.get(f'{server_common.get_server_url()}/api/health',
-                            cookies=server_common.get_api_cookie_jar())
+    response = rest.get(f'{server_common.get_server_url()}/api/health',
+                        cookies=server_common.get_api_cookie_jar())
     response.raise_for_status()
     return response.json()
 
@@ -1862,6 +1961,8 @@ 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.
@@ -1876,6 +1977,8 @@ def api_start(
             if deploy is True, to allow remote access.
         foreground: Whether to run the API server in the foreground (run in
             the current process).
+        metrics: Whether to export metrics of the API server.
+        metrics_port: The port to export metrics of the API server.
         enable_basic_auth: Whether to enable basic authentication
             in the API server.
     Returns:
@@ -1897,6 +2000,7 @@ def api_start(
                              '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

sky/clouds/aws.py

@@ -32,6 +32,7 @@ if typing.TYPE_CHECKING:
     # renaming to avoid shadowing variables
     from sky import resources as resources_lib
     from sky.utils import status_lib
+    from sky.volumes import volume as volume_lib
 
 logger = sky_logging.init_logger(__name__)
 
@@ -428,13 +429,15 @@ class AWS(clouds.Cloud):
                                                         clouds='aws')
 
     def make_deploy_resources_variables(
-            self,
-            resources: 'resources_lib.Resources',
-            cluster_name: resources_utils.ClusterName,
-            region: 'clouds.Region',
-            zones: Optional[List['clouds.Zone']],
-            num_nodes: int,
-            dryrun: bool = False) -> Dict[str, Any]:
+        self,
+        resources: 'resources_lib.Resources',
+        cluster_name: resources_utils.ClusterName,
+        region: 'clouds.Region',
+        zones: Optional[List['clouds.Zone']],
+        num_nodes: int,
+        dryrun: bool = False,
+        volume_mounts: Optional[List['volume_lib.VolumeMount']] = None,
+    ) -> Dict[str, Any]:
         del dryrun  # unused
         assert zones is not None, (region, zones)
 

sky/clouds/azure.py

@@ -24,6 +24,7 @@ from sky.utils import ux_utils
 
 if typing.TYPE_CHECKING:
     from sky import resources
+    from sky.volumes import volume as volume_lib
 
 logger = sky_logging.init_logger(__name__)
 
@@ -313,13 +314,15 @@ class Azure(clouds.Cloud):
         return None
 
     def make_deploy_resources_variables(
-            self,
-            resources: 'resources.Resources',
-            cluster_name: resources_utils.ClusterName,
-            region: 'clouds.Region',
-            zones: Optional[List['clouds.Zone']],
-            num_nodes: int,
-            dryrun: bool = False) -> Dict[str, Any]:
+        self,
+        resources: 'resources.Resources',
+        cluster_name: resources_utils.ClusterName,
+        region: 'clouds.Region',
+        zones: Optional[List['clouds.Zone']],
+        num_nodes: int,
+        dryrun: bool = False,
+        volume_mounts: Optional[List['volume_lib.VolumeMount']] = None,
+    ) -> Dict[str, Any]:
         assert zones is None, ('Azure does not support zones', zones)
 
         region_name = region.name