skypilot-nightly 1.0.0.dev20250627__py3-none-any.whl → 1.0.0.dev20250628__py3-none-any.whl

sky/client/sdk.py

@@ -147,9 +147,8 @@ def check(infra_list: Optional[Tuple[str, ...]],
     body = payloads.CheckBody(clouds=clouds,
                               verbose=verbose,
                               workspace=workspace)
-    response = rest.post(f'{server_common.get_server_url()}/check',
-                         json=json.loads(body.model_dump_json()),
-                         cookies=server_common.get_api_cookie_jar())
+    response = server_common.make_authenticated_request(
+        'POST', '/check', json=json.loads(body.model_dump_json()))
     return server_common.get_request_id(response)
 
 
@@ -173,9 +172,8 @@ def enabled_clouds(workspace: Optional[str] = None,
     """
     if workspace is None:
         workspace = skypilot_config.get_active_workspace()
-    response = rest.get((f'{server_common.get_server_url()}/enabled_clouds?'
-                         f'workspace={workspace}&expand={expand}'),
-                        cookies=server_common.get_api_cookie_jar())
+    response = server_common.make_authenticated_request(
+        'GET', f'/enabled_clouds?workspace={workspace}&expand={expand}')
     return server_common.get_request_id(response)
 
 
@@ -223,9 +221,8 @@ def list_accelerators(gpus_only: bool = True,
         require_price=require_price,
         case_sensitive=case_sensitive,
     )
-    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())
+    response = server_common.make_authenticated_request(
+        'POST', '/list_accelerators', json=json.loads(body.model_dump_json()))
     return server_common.get_request_id(response)
 
 
@@ -263,10 +260,10 @@ def list_accelerator_counts(
         quantity_filter=quantity_filter,
         clouds=clouds,
     )
-    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())
+    response = server_common.make_authenticated_request(
+        'POST',
+        '/list_accelerator_counts',
+        json=json.loads(body.model_dump_json()))
     return server_common.get_request_id(response)
 
 
@@ -303,16 +300,14 @@ def optimize(
     body = payloads.OptimizeBody(dag=dag_str,
                                  minimize=minimize,
                                  request_options=admin_policy_request_options)
-    response = rest.post(f'{server_common.get_server_url()}/optimize',
-                         json=json.loads(body.model_dump_json()),
-                         cookies=server_common.get_api_cookie_jar())
+    response = server_common.make_authenticated_request(
+        'POST', '/optimize', json=json.loads(body.model_dump_json()))
     return server_common.get_request_id(response)
 
 
 def workspaces() -> server_common.RequestId:
     """Gets the workspaces."""
-    response = rest.get(f'{server_common.get_server_url()}/workspaces',
-                        cookies=server_common.get_api_cookie_jar())
+    response = server_common.make_authenticated_request('GET', '/workspaces')
     return server_common.get_request_id(response)
 
 
@@ -346,9 +341,8 @@ 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 = rest.post(f'{server_common.get_server_url()}/validate',
-                         json=json.loads(body.model_dump_json()),
-                         cookies=server_common.get_api_cookie_jar())
+    response = server_common.make_authenticated_request(
+        'POST', '/validate', json=json.loads(body.model_dump_json()))
     if response.status_code == 400:
         with ux_utils.print_exception_no_traceback():
             raise exceptions.deserialize_exception(
@@ -632,12 +626,8 @@ def _launch(
             _is_launched_by_sky_serve_controller),
         disable_controller_check=_disable_controller_check,
     )
-    response = rest.post(
-        f'{server_common.get_server_url()}/launch',
-        json=json.loads(body.model_dump_json()),
-        timeout=5,
-        cookies=server_common.get_api_cookie_jar(),
-    )
+    response = server_common.make_authenticated_request(
+        'POST', '/launch', json=json.loads(body.model_dump_json()), timeout=5)
     return server_common.get_request_id(response)
 
 
@@ -716,12 +706,8 @@ def exec(  # pylint: disable=redefined-builtin
         backend=backend.NAME if backend else None,
     )
 
-    response = rest.post(
-        f'{server_common.get_server_url()}/exec',
-        json=json.loads(body.model_dump_json()),
-        timeout=5,
-        cookies=server_common.get_api_cookie_jar(),
-    )
+    response = server_common.make_authenticated_request(
+        'POST', '/exec', json=json.loads(body.model_dump_json()), timeout=5)
     return server_common.get_request_id(response)
 
 
@@ -769,13 +755,13 @@ def tail_logs(cluster_name: str,
         follow=follow,
         tail=tail,
     )
-    response = rest.post(
-        f'{server_common.get_server_url()}/logs',
+    response = server_common.make_authenticated_request(
+        'POST',
+        '/logs',
         json=json.loads(body.model_dump_json()),
         stream=True,
         timeout=(client_common.API_SERVER_REQUEST_CONNECTION_TIMEOUT_SECONDS,
-                 None),
-        cookies=server_common.get_api_cookie_jar())
+                 None))
     request_id = server_common.get_request_id(response)
     # Log request is idempotent when tail is 0, thus can resume previous
     # streaming point on retry.
@@ -816,9 +802,8 @@ def download_logs(cluster_name: str,
         cluster_name=cluster_name,
         job_ids=job_ids,
     )
-    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())
+    response = server_common.make_authenticated_request(
+        'POST', '/download_logs', json=json.loads(body.model_dump_json()))
     job_id_remote_path_dict = stream_and_get(
         server_common.get_request_id(response))
     remote2local_path_dict = client_common.download_logs_from_api_server(
@@ -896,12 +881,8 @@ def start(
         down=down,
         force=force,
     )
-    response = rest.post(
-        f'{server_common.get_server_url()}/start',
-        json=json.loads(body.model_dump_json()),
-        timeout=5,
-        cookies=server_common.get_api_cookie_jar(),
-    )
+    response = server_common.make_authenticated_request(
+        'POST', '/start', json=json.loads(body.model_dump_json()), timeout=5)
     return server_common.get_request_id(response)
 
 
@@ -942,12 +923,8 @@ def down(cluster_name: str, purge: bool = False) -> server_common.RequestId:
         cluster_name=cluster_name,
         purge=purge,
     )
-    response = rest.post(
-        f'{server_common.get_server_url()}/down',
-        json=json.loads(body.model_dump_json()),
-        timeout=5,
-        cookies=server_common.get_api_cookie_jar(),
-    )
+    response = server_common.make_authenticated_request(
+        'POST', '/down', json=json.loads(body.model_dump_json()), timeout=5)
     return server_common.get_request_id(response)
 
 
@@ -991,12 +968,8 @@ def stop(cluster_name: str, purge: bool = False) -> server_common.RequestId:
         cluster_name=cluster_name,
         purge=purge,
     )
-    response = rest.post(
-        f'{server_common.get_server_url()}/stop',
-        json=json.loads(body.model_dump_json()),
-        timeout=5,
-        cookies=server_common.get_api_cookie_jar(),
-    )
+    response = server_common.make_authenticated_request(
+        'POST', '/stop', json=json.loads(body.model_dump_json()), timeout=5)
     return server_common.get_request_id(response)
 
 
@@ -1061,12 +1034,8 @@ def autostop(
         idle_minutes=idle_minutes,
         down=down,
     )
-    response = rest.post(
-        f'{server_common.get_server_url()}/autostop',
-        json=json.loads(body.model_dump_json()),
-        timeout=5,
-        cookies=server_common.get_api_cookie_jar(),
-    )
+    response = server_common.make_authenticated_request(
+        'POST', '/autostop', json=json.loads(body.model_dump_json()), timeout=5)
     return server_common.get_request_id(response)
 
 
@@ -1124,9 +1093,8 @@ def queue(cluster_name: str,
         skip_finished=skip_finished,
         all_users=all_users,
     )
-    response = rest.post(f'{server_common.get_server_url()}/queue',
-                         json=json.loads(body.model_dump_json()),
-                         cookies=server_common.get_api_cookie_jar())
+    response = server_common.make_authenticated_request(
+        'POST', '/queue', json=json.loads(body.model_dump_json()))
     return server_common.get_request_id(response)
 
 
@@ -1166,9 +1134,8 @@ def job_status(cluster_name: str,
         cluster_name=cluster_name,
         job_ids=job_ids,
     )
-    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())
+    response = server_common.make_authenticated_request(
+        'POST', '/job_status', json=json.loads(body.model_dump_json()))
     return server_common.get_request_id(response)
 
 
@@ -1220,9 +1187,8 @@ def cancel(
         job_ids=job_ids,
         try_cancel_if_cluster_is_init=_try_cancel_if_cluster_is_init,
     )
-    response = rest.post(f'{server_common.get_server_url()}/cancel',
-                         json=json.loads(body.model_dump_json()),
-                         cookies=server_common.get_api_cookie_jar())
+    response = server_common.make_authenticated_request(
+        'POST', '/cancel', json=json.loads(body.model_dump_json()))
     return server_common.get_request_id(response)
 
 
@@ -1316,9 +1282,8 @@ def status(
         refresh=refresh,
         all_users=all_users,
     )
-    response = rest.post(f'{server_common.get_server_url()}/status',
-                         json=json.loads(body.model_dump_json()),
-                         cookies=server_common.get_api_cookie_jar())
+    response = server_common.make_authenticated_request(
+        'POST', '/status', json=json.loads(body.model_dump_json()))
     return server_common.get_request_id(response)
 
 
@@ -1351,9 +1316,8 @@ def endpoints(
         cluster=cluster,
         port=port,
     )
-    response = rest.post(f'{server_common.get_server_url()}/endpoints',
-                         json=json.loads(body.model_dump_json()),
-                         cookies=server_common.get_api_cookie_jar())
+    response = server_common.make_authenticated_request(
+        'POST', '/endpoints', json=json.loads(body.model_dump_json()))
     return server_common.get_request_id(response)
 
 
@@ -1396,9 +1360,8 @@ def cost_report(days: Optional[int] = None) -> server_common.RequestId:  # pylin
             }
     """
     body = payloads.CostReportBody(days=days)
-    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())
+    response = server_common.make_authenticated_request(
+        'POST', '/cost_report', json=json.loads(body.model_dump_json()))
     return server_common.get_request_id(response)
 
 
@@ -1427,8 +1390,7 @@ def storage_ls() -> server_common.RequestId:
                 }
         ]
     """
-    response = rest.get(f'{server_common.get_server_url()}/storage/ls',
-                        cookies=server_common.get_api_cookie_jar())
+    response = server_common.make_authenticated_request('GET', '/storage/ls')
     return server_common.get_request_id(response)
 
 
@@ -1451,9 +1413,8 @@ def storage_delete(name: str) -> server_common.RequestId:
         ValueError: If the storage does not exist.
     """
     body = payloads.StorageBody(name=name)
-    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())
+    response = server_common.make_authenticated_request(
+        'POST', '/storage/delete', json=json.loads(body.model_dump_json()))
     return server_common.get_request_id(response)
 
 
@@ -1490,9 +1451,8 @@ def local_up(gpus: bool,
                                 cleanup=cleanup,
                                 context_name=context_name,
                                 password=password)
-    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())
+    response = server_common.make_authenticated_request(
+        'POST', '/local_up', json=json.loads(body.model_dump_json()))
     return server_common.get_request_id(response)
 
 
@@ -1508,8 +1468,7 @@ 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 = rest.post(f'{server_common.get_server_url()}/local_down',
-                         cookies=server_common.get_api_cookie_jar())
+    response = server_common.make_authenticated_request('POST', '/local_down')
     return server_common.get_request_id(response)
 
 
@@ -1538,9 +1497,9 @@ def _update_remote_ssh_node_pools(file: str,
         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())
+    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:
@@ -1559,8 +1518,9 @@ def _upload_ssh_key_and_wait(key_name: str, key_file_path: str) -> str:
             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',
+        response = server_common.make_authenticated_request(
+            'POST',
+            '/ssh_node_pools/keys',
             files={
                 'key_file': (key_name, key_file, 'application/octet-stream')
             },
@@ -1593,15 +1553,14 @@ def ssh_up(infra: Optional[str] = None,
     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())
+        response = server_common.make_authenticated_request(
+            'POST', f'/ssh_node_pools/{infra}/deploy')
     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())
+        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)
 
 
@@ -1622,15 +1581,14 @@ def ssh_down(infra: Optional[str] = None) -> server_common.RequestId:
     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())
+        response = server_common.make_authenticated_request(
+            'POST', f'/ssh_node_pools/{infra}/down')
     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())
+        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)
 
 
@@ -1653,11 +1611,10 @@ def realtime_kubernetes_gpu_availability(
         quantity_filter=quantity_filter,
         is_ssh=is_ssh,
     )
-    response = rest.post(
-        f'{server_common.get_server_url()}/'
-        'realtime_kubernetes_gpu_availability',
-        json=json.loads(body.model_dump_json()),
-        cookies=server_common.get_api_cookie_jar())
+    response = server_common.make_authenticated_request(
+        'POST',
+        '/realtime_kubernetes_gpu_availability',
+        json=json.loads(body.model_dump_json()))
     return server_common.get_request_id(response)
 
 
@@ -1686,10 +1643,10 @@ def kubernetes_node_info(
             information.
     """
     body = payloads.KubernetesNodeInfoRequestBody(context=context)
-    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())
+    response = server_common.make_authenticated_request(
+        'POST',
+        '/kubernetes_node_info',
+        json=json.loads(body.model_dump_json()))
     return server_common.get_request_id(response)
 
 
@@ -1717,8 +1674,8 @@ 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 = rest.get(f'{server_common.get_server_url()}/status_kubernetes',
-                        cookies=server_common.get_api_cookie_jar())
+    response = server_common.make_authenticated_request('GET',
+                                                        '/status_kubernetes')
     return server_common.get_request_id(response)
 
 
@@ -1744,11 +1701,12 @@ def get(request_id: str) -> Any:
             see ``Request Raises`` in the documentation of the specific requests
             above.
     """
-    response = rest.get_without_retry(
-        f'{server_common.get_server_url()}/api/get?request_id={request_id}',
+    response = server_common.make_authenticated_request(
+        'GET',
+        f'/api/get?request_id={request_id}',
+        retry=False,
         timeout=(client_common.API_SERVER_REQUEST_CONNECTION_TIMEOUT_SECONDS,
-                 None),
-        cookies=server_common.get_api_cookie_jar())
+                 None))
     request_task = None
     if response.status_code == 200:
         request_task = requests_lib.Request.decode(
@@ -1822,13 +1780,14 @@ def stream_and_get(
         'follow': follow,
         'format': 'console',
     }
-    response = rest.get_without_retry(
-        f'{server_common.get_server_url()}/api/stream',
+    response = server_common.make_authenticated_request(
+        'GET',
+        '/api/stream',
         params=params,
+        retry=False,
         timeout=(client_common.API_SERVER_REQUEST_CONNECTION_TIMEOUT_SECONDS,
                  None),
-        stream=True,
-        cookies=server_common.get_api_cookie_jar())
+        stream=True)
     if response.status_code in [404, 400]:
         detail = response.json().get('detail')
         with ux_utils.print_exception_no_traceback():
@@ -1882,10 +1841,11 @@ def api_cancel(request_ids: Optional[Union[str, List[str]]] = None,
         echo(f'Cancelling {len(request_ids)} request{plural}: '
              f'{request_id_str}...')
 
-    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())
+    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)
 
 
@@ -1909,12 +1869,12 @@ def api_status(
     """
     body = payloads.RequestStatusBody(request_ids=request_ids,
                                       all_status=all_status)
-    response = rest.get(
-        f'{server_common.get_server_url()}/api/status',
+    response = server_common.make_authenticated_request(
+        'GET',
+        '/api/status',
         params=server_common.request_body_to_params(body),
         timeout=(client_common.API_SERVER_REQUEST_CONNECTION_TIMEOUT_SECONDS,
-                 None),
-        cookies=server_common.get_api_cookie_jar())
+                 None))
     server_common.handle_request_error(response)
     return [
         requests_lib.RequestPayload(**request) for request in response.json()
@@ -1948,8 +1908,7 @@ def api_info() -> Dict[str, Any]:
         Note that user may be None if we are not using an auth proxy.
 
     """
-    response = rest.get(f'{server_common.get_server_url()}/api/health',
-                        cookies=server_common.get_api_cookie_jar())
+    response = server_common.make_authenticated_request('GET', '/api/health')
     response.raise_for_status()
     return response.json()
 
@@ -2074,9 +2033,53 @@ def api_server_logs(follow: bool = True, tail: Optional[int] = None) -> None:
         stream_and_get(log_path=constants.API_SERVER_LOGS, tail=tail)
 
 
+def _save_config_updates(endpoint: Optional[str] = None,
+                         service_account_token: Optional[str] = None) -> None:
+    """Save endpoint and/or service account token to config file."""
+    config_path = pathlib.Path(
+        skypilot_config.get_user_config_path()).expanduser()
+    with filelock.FileLock(config_path.with_suffix('.lock')):
+        if not config_path.exists():
+            config_path.touch()
+            config: Dict[str, Any] = {}
+        else:
+            config = skypilot_config.get_user_config()
+            config = dict(config)
+
+        # Update endpoint if provided
+        if endpoint is not None:
+            # We should always reset the api_server config to avoid legacy
+            # service account token.
+            config['api_server'] = {}
+            config['api_server']['endpoint'] = endpoint
+
+        # Update service account token if provided
+        if service_account_token is not None:
+            if 'api_server' not in config:
+                config['api_server'] = {}
+            config['api_server'][
+                'service_account_token'] = service_account_token
+
+        common_utils.dump_yaml(str(config_path), config)
+        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('/')
+
+
 @usage_lib.entrypoint
 @annotations.client_api
-def api_login(endpoint: Optional[str] = None, get_token: bool = False) -> None:
+def api_login(endpoint: Optional[str] = None,
+              relogin: bool = False,
+              service_account_token: Optional[str] = None) -> None:
     """Logs into a SkyPilot API server.
 
     This sets the endpoint globally, i.e., all SkyPilot CLI and SDK calls will
@@ -2088,25 +2091,80 @@ def api_login(endpoint: Optional[str] = None, get_token: bool = False) -> None:
     Args:
         endpoint: The endpoint of the SkyPilot API server, e.g.,
             http://1.2.3.4:46580 or https://skypilot.mydomain.com.
-        get_token: Whether to force getting a new token even if not needed.
+        relogin: Whether to force relogin with OAuth2 when enabled.
+        service_account_token: Service account token for authentication.
 
     Returns:
         None
     """
+    # 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')
+
+    # 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)
+            _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.
-    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.')
-    endpoint = endpoint.rstrip('/')
 
     server_status, api_server_info = server_common.check_server_healthy(
         endpoint)
-    if server_status == server_common.ApiServerStatus.NEEDS_AUTH or get_token:
+    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
@@ -2253,20 +2311,12 @@ def api_login(endpoint: Optional[str] = None, get_token: bool = False) -> None:
                 f.write(user_hash)
 
     # Set the endpoint in the config file
-    config_path = pathlib.Path(
-        skypilot_config.get_user_config_path()).expanduser()
-    with filelock.FileLock(config_path.with_suffix('.lock')):
-        if not config_path.exists():
-            config_path.touch()
-            config = {'api_server': {'endpoint': endpoint}}
-        else:
-            config = skypilot_config.get_user_config()
-            config.set_nested(('api_server', 'endpoint'), endpoint)
-        common_utils.dump_yaml(str(config_path), dict(config))
-        dashboard_url = server_common.get_dashboard_url(endpoint)
-        dashboard_msg = f'Dashboard: {dashboard_url}'
-        click.secho(
-            f'Logged into SkyPilot API server at: {endpoint}'
-            f'\n{ux_utils.INDENT_LAST_SYMBOL}{colorama.Fore.GREEN}'
-            f'{dashboard_msg}',
-            fg='green')
+    _save_config_updates(endpoint=endpoint)
+    dashboard_url = server_common.get_dashboard_url(endpoint)
+
+    # 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)