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

sky/catalog/slurm_catalog.py

@@ -112,9 +112,6 @@ def list_accelerators_realtime(
     else:
         slurm_cluster = region_filter
 
-    partition_filter = slurm_utils.get_cluster_default_partition(slurm_cluster)
-
-    # Call the helper function to get node info
     slurm_nodes_info = slurm_utils.slurm_node_info(
         slurm_cluster_name=slurm_cluster)
 
@@ -126,8 +123,6 @@ def list_accelerators_realtime(
             filters_applied.append(f'gpu_name={name_filter!r}')
         if quantity_filter:
             filters_applied.append(f'quantity>={quantity_filter}')
-        if region_filter:
-            filters_applied.append(f'cluster={region_filter!r}')
         if filters_applied:
             err_msg += f' with filters ({", ".join(filters_applied)})'
         err_msg += '.'
@@ -214,8 +209,6 @@ def list_accelerators_realtime(
             filters_applied.append(f'gpu_name={name_filter!r}')
         if quantity_filter:
             filters_applied.append(f'quantity>={quantity_filter}')
-        if partition_filter:
-            filters_applied.append(f'partition={partition_filter!r}')
         if filters_applied:
             err_msg += f' with filters ({", ".join(filters_applied)})'
         err_msg += '.'

sky/catalog/vast_catalog.py

@@ -7,7 +7,10 @@ query instance types and pricing information for Vast.ai.
 import typing
 from typing import Dict, List, Optional, Tuple, Union
 
+import pandas as pd
+
 from sky.catalog import common
+from sky.utils import resources_utils
 from sky.utils import ux_utils
 
 if typing.TYPE_CHECKING:
@@ -16,6 +19,17 @@ if typing.TYPE_CHECKING:
 _df = common.read_catalog('vast/vms.csv')
 
 
+def _apply_datacenter_filter(df: pd.DataFrame,
+                             datacenter_only: bool) -> pd.DataFrame:
+    """Filter dataframe by hosting_type if datacenter_only is True.
+
+    hosting_type: 0 = Consumer hosted, 1 = Datacenter hosted
+    """
+    if not datacenter_only or 'HostingType' not in df.columns:
+        return df
+    return df[df['HostingType'] >= 1]
+
+
 def instance_type_exists(instance_type: str) -> bool:
     return common.instance_type_exists_impl(_df, instance_type)
 
@@ -48,13 +62,16 @@ def get_vcpus_mem_from_instance_type(
 
 def get_default_instance_type(cpus: Optional[str] = None,
                               memory: Optional[str] = None,
-                              disk_tier: Optional[str] = None,
+                              disk_tier: Optional[
+                                  resources_utils.DiskTier] = None,
                               region: Optional[str] = None,
-                              zone: Optional[str] = None) -> Optional[str]:
+                              zone: Optional[str] = None,
+                              datacenter_only: bool = False) -> Optional[str]:
     del disk_tier
     # NOTE: After expanding catalog to multiple entries, you may
     # want to specify a default instance type or family.
-    return common.get_instance_type_for_cpus_mem_impl(_df, cpus, memory, region,
+    df = _apply_datacenter_filter(_df, datacenter_only)
+    return common.get_instance_type_for_cpus_mem_impl(df, cpus, memory, region,
                                                       zone)
 
 
@@ -70,12 +87,19 @@ def get_instance_type_for_accelerator(
         memory: Optional[str] = None,
         use_spot: bool = False,
         region: Optional[str] = None,
-        zone: Optional[str] = None) -> Tuple[Optional[List[str]], List[str]]:
-    """Returns a list of instance types that have the given accelerator."""
+        zone: Optional[str] = None,
+        datacenter_only: bool = False) -> Tuple[Optional[List[str]], List[str]]:
+    """Returns a list of instance types that have the given accelerator.
+
+    Args:
+        datacenter_only: If True, only return instances hosted in datacenters
+            (hosting_type >= 1).
+    """
     if zone is not None:
         with ux_utils.print_exception_no_traceback():
             raise ValueError('Vast does not support zones.')
-    return common.get_instance_type_for_accelerator_impl(df=_df,
+    df = _apply_datacenter_filter(_df, datacenter_only)
+    return common.get_instance_type_for_accelerator_impl(df=df,
                                                          acc_name=acc_name,
                                                          acc_count=acc_count,
                                                          cpus=cpus,

sky/check.py

@@ -528,8 +528,9 @@ def _print_checked_cloud(
         # `dict` reasons for K8s and SSH will be printed in detail in
         # _format_enabled_cloud. Skip here unless the cloud is disabled.
         if not isinstance(reason, str):
-            if not ok and isinstance(cloud_tuple[1],
-                                     (sky_clouds.SSH, sky_clouds.Kubernetes)):
+            if not ok and isinstance(
+                    cloud_tuple[1],
+                (sky_clouds.SSH, sky_clouds.Kubernetes, sky_clouds.Slurm)):
                 if reason is not None:
                     reason_str = _format_context_details(cloud_tuple[1],
                                                          show_details=True,
@@ -555,7 +556,9 @@ def _print_checked_cloud(
         capability_string = f'[{", ".join(enabled_capabilities)}]'
         if verbose and cloud is not cloudflare and cloud is not coreweave:
             activated_account = cloud.get_active_user_identity_str()
-        if isinstance(cloud_tuple[1], (sky_clouds.SSH, sky_clouds.Kubernetes)):
+        if isinstance(
+                cloud_tuple[1],
+            (sky_clouds.SSH, sky_clouds.Kubernetes, sky_clouds.Slurm)):
             detail_string = _format_context_details(cloud_tuple[1],
                                                     show_details=True,
                                                     ctx2text=ctx2text)
@@ -653,11 +656,11 @@ def _format_context_details(cloud: Union[str, sky_clouds.Cloud],
                                               'configuration.'))
                 else:
                     # Default case - not set up
-                    text_suffix = (': ' + _red_color('disabled. ') +
-                                   _dim_color('Reason: Not set up. Use '
-                                              '`sky ssh up --infra '
-                                              f'{context.lstrip("ssh-")}` '
-                                              'to set up.'))
+                    text_suffix = (': ' + _red_color('disabled. ') + _dim_color(
+                        'Reason: Not set up. Use '
+                        '`sky ssh up --infra '
+                        f'{common_utils.removeprefix(context, "ssh-")}` '
+                        'to set up.'))
         contexts_formatted.append(
             f'\n    {symbol}{cleaned_context}{text_suffix}')
     if isinstance(cloud_type, sky_clouds.SSH):

sky/client/cli/command.py

@@ -216,45 +216,27 @@ def _get_cluster_records_and_set_ssh_config(
                 f'\"{escaped_executable_path} '
                 f'{escaped_websocket_proxy_path} '
                 f'{server_common.get_server_url()} '
-                f'{handle.cluster_name}\"')
+                f'{handle.cluster_name} '
+                f'kubernetes-pod-ssh-proxy\"')
             credentials['ssh_proxy_command'] = proxy_command
         elif isinstance(handle.launched_resources.cloud, clouds.Slurm):
-            # TODO(kevin): This is a temporary workaround, ideally we want to
-            # get a shell through srun --pty bash on the existing sbatch job.
-
-            # Proxy through the controller/login node to reach the worker node.
-            if (handle.cached_internal_ips is None or
-                    not handle.cached_internal_ips):
-                logger.debug(
-                    f'Cluster {name} does not have cached internal IPs. '
-                    'Skipping SSH config update.')
-                cluster_utils.SSHConfigHelper.remove_cluster(name)
-                continue
-
-            escaped_key_path = shlex.quote(
-                cluster_utils.SSHConfigHelper.generate_local_key_file(
-                    handle.cluster_name, credentials))
-            controller_host = handle.cached_external_ips[0]
-
-            # Build jump proxy: ssh to worker via controller/login node
-            proxy_command = (f'ssh -tt -i {escaped_key_path} '
-                             '-o StrictHostKeyChecking=no '
-                             '-o UserKnownHostsFile=/dev/null '
-                             '-o IdentitiesOnly=yes '
-                             '-W %h:%p '
-                             f'{handle.ssh_user}@{controller_host}')
-            original_proxy = credentials.get('ssh_proxy_command')
-            if original_proxy:
-                proxy_command += (
-                    f' -o ProxyCommand={shlex.quote(original_proxy)}')
-
+            # Replace the proxy command to proxy through the SkyPilot API
+            # server with websocket.
+            escaped_executable_path = shlex.quote(sys.executable)
+            escaped_websocket_proxy_path = shlex.quote(
+                f'{directory_utils.get_sky_dir()}/templates/websocket_proxy.py')
+            # %w is a placeholder for the node index, substituted per-node
+            # in cluster_utils.SSHConfigHelper.add_cluster().
+            proxy_command = (f'{escaped_executable_path} '
+                             f'{escaped_websocket_proxy_path} '
+                             f'{server_common.get_server_url()} '
+                             f'{handle.cluster_name} '
+                             f'slurm-job-ssh-proxy %w')
             credentials['ssh_proxy_command'] = proxy_command
 
-            # For Slurm, use the worker's internal IP as the SSH target
-            ips = handle.cached_internal_ips
-
         cluster_utils.SSHConfigHelper.add_cluster(
             handle.cluster_name,
+            handle.cluster_name_on_cloud,
             ips,
             credentials,
             handle.cached_external_ssh_ports,
@@ -3471,7 +3453,12 @@ def _down_or_stop_clusters(
                     click.echo(f'      {name} ({first})')
 
     if failures:
-        click.echo('Cluster(s) failed. See details above.')
+        failure_str = 'Cluster(s) failed. See details above.'
+        if down:
+            failure_str += (
+                ' If you want to ignore the errors and remove the '
+                'cluster(s) from the status table, use `sky down --purge`.')
+        click.echo(failure_str)
 
 
 @cli.command(cls=_DocumentedCodeCommand)
@@ -3898,8 +3885,10 @@ def show_gpus(
             contexts_info: List[Tuple[str, 'models.KubernetesNodesInfo']],
             cloud_str: str = 'Kubernetes',
             context_title_str: str = 'CONTEXT') -> str:
-        node_table = log_utils.create_table(
-            [context_title_str, 'NODE', 'GPU', 'UTILIZATION'])
+        node_table = log_utils.create_table([
+            context_title_str, 'NODE', 'vCPU', 'Memory (GB)', 'GPU',
+            'GPU UTILIZATION'
+        ])
 
         no_permissions_str = '<no permissions>'
         hints = []
@@ -3916,6 +3905,44 @@ def show_gpus(
                 acc_type = node_info.accelerator_type
                 if acc_type is None:
                     acc_type = '-'
+
+                # Format CPU and memory: "X of Y free" or just "Y" if
+                # free is unknown
+                cpu_str = '-'
+                if node_info.cpu_count is not None:
+                    cpu_total_str = common_utils.format_float(
+                        node_info.cpu_count, precision=0)
+
+                    # Check if we have free CPU info (use hasattr to
+                    # check if field exists, then access directly)
+                    cpu_free = None
+                    if hasattr(node_info, 'cpu_free'):
+                        cpu_free = node_info.cpu_free
+                    if cpu_free is not None:
+                        cpu_free_str = common_utils.format_float(cpu_free,
+                                                                 precision=0)
+                        cpu_str = f'{cpu_free_str} of {cpu_total_str} free'
+                    else:
+                        cpu_str = cpu_total_str
+
+                memory_str = '-'
+                if node_info.memory_gb is not None:
+                    memory_total_str = common_utils.format_float(
+                        node_info.memory_gb, precision=0)
+
+                    # Check if we have free memory info (use hasattr
+                    # to check if field exists, then access directly)
+                    memory_free_gb = None
+                    if hasattr(node_info, 'memory_free_gb'):
+                        memory_free_gb = node_info.memory_free_gb
+                    if memory_free_gb is not None:
+                        memory_free_str = common_utils.format_float(
+                            memory_free_gb, precision=0)
+                        memory_str = (
+                            f'{memory_free_str} of {memory_total_str} free')
+                    else:
+                        memory_str = memory_total_str
+
                 utilization_str = (
                     f'{available} of '
                     f'{node_info.total["accelerator_count"]} free')
@@ -3924,8 +3951,11 @@ def show_gpus(
                 node_is_ready = getattr(node_info, 'is_ready', True)
                 if not node_is_ready:
                     utilization_str += ' (Node NotReady)'
-                node_table.add_row(
-                    [context_name, node_name, acc_type, utilization_str])
+
+                node_table.add_row([
+                    context_name, node_name, cpu_str, memory_str, acc_type,
+                    utilization_str
+                ])
 
         k8s_per_node_acc_message = (f'{cloud_str} per-node GPU availability')
         if hints:
@@ -3936,7 +3966,7 @@ def show_gpus(
                 f'{colorama.Style.RESET_ALL}\n'
                 f'{node_table.get_string()}')
 
-    def _format_slurm_node_info() -> str:
+    def _format_slurm_node_info(slurm_cluster_names: List[str]) -> str:
         node_table = log_utils.create_table([
             'CLUSTER',
             'NODE',
@@ -3946,13 +3976,12 @@ def show_gpus(
             'UTILIZATION',
         ])
 
-        # Get all cluster names
-        slurm_cluster_names = clouds.Slurm.existing_allowed_clusters()
+        request_ids = [(cluster_name,
+                        sdk.slurm_node_info(slurm_cluster_name=cluster_name))
+                       for cluster_name in slurm_cluster_names]
 
-        # Query each cluster
-        for cluster_name in slurm_cluster_names:
-            nodes_info = sdk.stream_and_get(
-                sdk.slurm_node_info(slurm_cluster_name=cluster_name))
+        for cluster_name, request_id in request_ids:
+            nodes_info = sdk.stream_and_get(request_id)
 
             for node_info in nodes_info:
                 node_table.add_row([
@@ -4122,7 +4151,8 @@ def show_gpus(
             yield from slurm_realtime_table.get_string()
             yield '\n'
         if show_node_info:
-            yield _format_slurm_node_info()
+            cluster_names = [cluster for cluster, _ in slurm_realtime_infos]
+            yield _format_slurm_node_info(cluster_names)
 
     def _output() -> Generator[str, None, None]:
         gpu_table = log_utils.create_table(
@@ -4705,6 +4735,13 @@ def volumes_ls(verbose: bool):
               is_flag=True,
               required=False,
               help='Delete all volumes.')
+@click.option('--purge',
+              '-p',
+              default=False,
+              is_flag=True,
+              required=False,
+              help=('Forcibly delete the volume from the volumes table even '
+                    'if the deletion API fails.'))
 @click.option('--yes',
               '-y',
               default=False,
@@ -4713,7 +4750,12 @@ def volumes_ls(verbose: bool):
               help='Skip confirmation prompt.')
 @_add_click_options(flags.COMMON_OPTIONS)
 @usage_lib.entrypoint
-def volumes_delete(names: List[str], all: bool, yes: bool, async_call: bool):  # pylint: disable=redefined-builtin
+def volumes_delete(
+        names: List[str],
+        all: bool,  # pylint: disable=redefined-builtin
+        purge: bool,
+        yes: bool,
+        async_call: bool):
     """Delete volumes.
 
     Examples:
@@ -4728,6 +4770,9 @@ def volumes_delete(names: List[str], all: bool, yes: bool, async_call: bool):  #
         \b
         # Delete all volumes.
         sky volumes delete -a
+        \b
+        # Forcibly delete a volume.
+        sky volumes delete pvc1 -p
     """
     if sum([bool(names), all]) != 1:
         raise click.UsageError('Either --all or a name must be specified.')
@@ -4754,8 +4799,8 @@ def volumes_delete(names: List[str], all: bool, yes: bool, async_call: bool):  #
                 show_default=True)
 
         try:
-            _async_call_or_wait(volumes_sdk.delete(names), async_call,
-                                'sky.volumes.delete')
+            _async_call_or_wait(volumes_sdk.delete(names, purge=purge),
+                                async_call, 'sky.volumes.delete')
         except Exception as e:  # pylint: disable=broad-except
             logger.error(f'{colorama.Fore.RED}Error deleting volumes {names}: '
                          f'{str(e)}{colorama.Style.RESET_ALL}')
@@ -5427,9 +5472,14 @@ def jobs_pool_apply(
 @flags.config_option(expose_value=False)
 @flags.verbose_option()
 @click.argument('pool_names', required=False, type=str, nargs=-1)
+@click.option('--all',
+              '-a',
+              'show_all',
+              is_flag=True,
+              default=False,
+              help='Show all workers.')
 @usage_lib.entrypoint
-# pylint: disable=redefined-builtin
-def jobs_pool_status(verbose: bool, pool_names: List[str]):
+def jobs_pool_status(verbose: bool, pool_names: List[str], show_all: bool):
     """Show statuses of pools.
 
     Show detailed statuses of one or more pools. If POOL_NAME is not
@@ -5442,7 +5492,7 @@ def jobs_pool_status(verbose: bool, pool_names: List[str]):
         pool_status_request_id = managed_jobs.pool_status(pool_names_to_query)
         _, msg = _handle_services_request(pool_status_request_id,
                                           service_names=pool_names_to_query,
-                                          show_all=verbose,
+                                          show_all=verbose or show_all,
                                           show_endpoint=False,
                                           pool=True,
                                           is_called_by_user=True)
@@ -6745,9 +6795,11 @@ def api_status(request_id_prefixes: Optional[List[str]], all_status: bool,
             if not verbose:
                 r_id = common_utils.truncate_long_string(r_id, 36)
             req_status = requests.RequestStatus(request.status)
-            row = [r_id, request.user_name, request.name]
+            user_display = status_utils.get_user_display_name(
+                request.user_name or '-', request.user_id)
+            row = [r_id, user_display, request.name]
             if verbose:
-                row.append(request.cluster_name)
+                row.append(request.cluster_name or '-')
             row.extend([
                 log_utils.readable_time_duration(request.created_at),
                 req_status.colored_str()

sky/client/interactive_utils.py

@@ -0,0 +1,190 @@
+"""Utilities for handling interactive SSH authentication."""
+import asyncio
+import fcntl
+import os
+import re
+import sys
+import termios
+import tty
+import typing
+
+from sky import sky_logging
+from sky.adaptors import common as adaptors_common
+from sky.client import service_account_auth
+from sky.server import common as server_common
+from sky.utils import rich_utils
+
+if typing.TYPE_CHECKING:
+    import websockets
+else:
+    websockets = adaptors_common.LazyImport('websockets')
+
+logger = sky_logging.init_logger(__name__)
+
+SKY_INTERACTIVE_PATTERN = re.compile(r'<sky-interactive session="([^"]+)"/>')
+
+
+# TODO(kevin): Refactor to share code with websocket_proxy.py.
+async def _handle_interactive_auth_websocket(session_id: str) -> None:
+    """Handle interactive SSH authentication via websocket.
+
+    This establishes a websocket connection to the API server and bridges
+    the user's terminal I/O bidirectionally with the PTY on the server,
+    allowing interactive authentication (e.g., 2FA).
+
+    Args:
+        session_id: The session identifier from the <sky-interactive> signal.
+    """
+    # Get HTTP server URL and convert to websocket URL
+    server_url = server_common.get_server_url()
+    server_proto, server_fqdn = server_url.split('://')
+    websocket_proto = 'wss' if server_proto == 'https' else 'ws'
+    ws_url = (f'{websocket_proto}://{server_fqdn}'
+              f'/ssh-interactive-auth?session_id={session_id}')
+
+    logger.info('Starting interactive SSH authentication...')
+
+    headers = {}
+    # Add service account auth if available
+    headers.update(service_account_auth.get_service_account_headers())
+    # Add cookie auth with URL-aware filtering
+    headers.update(server_common.get_cookie_header_for_url(ws_url))
+
+    # Set terminal to raw mode if stdin is a tty
+    old_settings = None
+    if os.isatty(sys.stdin.fileno()):
+        old_settings = termios.tcgetattr(sys.stdin.fileno())
+        tty.setraw(sys.stdin.fileno())
+
+    stdin_dup_fd = None
+    stdout_dup_fd = None
+    try:
+        # Duplicate stdin/stdout fds before passing to asyncio.
+        # When asyncio's loop.connect_read/write_pipe() is called,
+        # it creates a transport that takes ownership of the file passed to it.
+        # By duplicating the fds, we give asyncio independent copies that it can
+        # safely close, while preserving the original sys.stdin/stdout.
+        stdin_dup_fd = os.dup(sys.stdin.fileno())
+        stdout_dup_fd = os.dup(sys.stdout.fileno())
+
+        async with websockets.connect(ws_url,
+                                      additional_headers=headers,
+                                      ping_interval=None) as ws:
+            loop = asyncio.get_running_loop()
+
+            stdin_reader = asyncio.StreamReader()
+            stdin_protocol = asyncio.StreamReaderProtocol(stdin_reader)
+            stdin_dup_file = os.fdopen(stdin_dup_fd, 'rb', buffering=0)
+            stdin_dup_fd = None  # File object now owns the FD
+            await loop.connect_read_pipe(lambda: stdin_protocol, stdin_dup_file)
+
+            stdout_dup_file = os.fdopen(stdout_dup_fd, 'wb', buffering=0)
+            stdout_dup_fd = None  # File object now owns the FD
+            stdout_transport, stdout_protocol = await loop.connect_write_pipe(
+                asyncio.streams.FlowControlMixin,
+                stdout_dup_file)  # type: ignore
+            stdout_writer = asyncio.StreamWriter(stdout_transport,
+                                                 stdout_protocol, None, loop)
+
+            async def stdin_to_websocket():
+                """Forward stdin to websocket."""
+                try:
+                    while True:
+                        data = await stdin_reader.read(4096)
+                        if not data:
+                            break
+                        await ws.send(data)
+                except asyncio.CancelledError:
+                    # Task was cancelled - auth complete
+                    pass
+                except Exception as e:  # pylint: disable=broad-except
+                    logger.debug(f'Error in stdin_to_websocket: {e}')
+
+            async def websocket_to_stdout():
+                """Forward websocket to stdout."""
+                try:
+                    async for message in ws:
+                        stdout_writer.write(message)
+                        await stdout_writer.drain()
+                except Exception as e:  # pylint: disable=broad-except
+                    logger.debug(f'Error in websocket_to_stdout: {e}')
+
+            # Run both directions concurrently
+            # Use tasks so we can cancel stdin reader when websocket closes
+            stdin_task = asyncio.create_task(stdin_to_websocket())
+            stdout_task = asyncio.create_task(websocket_to_stdout())
+
+            # Wait for websocket to close (auth complete)
+            await stdout_task
+            # Cancel stdin reader so it doesn't consume the next keystroke
+            stdin_task.cancel()
+            try:
+                await stdin_task
+            except asyncio.CancelledError:
+                pass
+    except Exception as e:  # pylint: disable=broad-except
+        logger.error(f'Failed to handle interactive authentication: {e}')
+        raise
+    finally:
+        # Restore terminal settings if they were changed
+        if old_settings:
+            termios.tcsetattr(sys.stdin.fileno(), termios.TCSADRAIN,
+                              old_settings)
+            # Flush any buffered input from stdin
+            termios.tcflush(sys.stdin.fileno(), termios.TCIFLUSH)
+            # Ensure stdout is in blocking mode (can be non-blocking after
+            # asyncio transport operations)
+            flags = fcntl.fcntl(sys.stdout.fileno(), fcntl.F_GETFL)
+            fcntl.fcntl(sys.stdout.fileno(), fcntl.F_SETFL,
+                        flags & ~os.O_NONBLOCK)
+
+        for fd in [stdin_dup_fd, stdout_dup_fd]:
+            if fd is not None:
+                try:
+                    os.close(fd)
+                except OSError:
+                    # Already closed by asyncio or never opened
+                    pass
+
+
+def handle_interactive_auth(line: str) -> typing.Optional[str]:
+    """Handle interactive SSH authentication signals (sync version).
+
+    Args:
+        line: The log line to check for interactive auth markers.
+
+    Returns:
+        The line with the marker removed, or None if this was an interactive
+        auth signal (meaning the line was consumed).
+    """
+    match = SKY_INTERACTIVE_PATTERN.search(line)
+    if not match:
+        return line
+
+    session_id = match.group(1)
+    # Temporarily stop any spinners to allow terminal I/O
+    with rich_utils.safe_logger():
+        asyncio.run(_handle_interactive_auth_websocket(session_id))
+
+    return None
+
+
+async def handle_interactive_auth_async(line: str) -> typing.Optional[str]:
+    """Handle interactive SSH authentication signals (async version).
+
+    Args:
+        line: The log line to check for interactive auth markers.
+
+    Returns:
+        The line with the marker removed, or None if this was an interactive
+        auth signal (meaning the line was consumed).
+    """
+    match = SKY_INTERACTIVE_PATTERN.search(line)
+    if not match:
+        return line
+
+    session_id = match.group(1)
+    with rich_utils.safe_logger():
+        await _handle_interactive_auth_websocket(session_id)
+
+    return None

sky/client/sdk.py

@@ -30,6 +30,7 @@ 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
@@ -157,9 +158,16 @@ def stream_response(request_id: Optional[server_common.RequestId[T]],
         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:

sky/client/sdk_async.py

@@ -23,6 +23,7 @@ from sky import catalog
 from sky import exceptions
 from sky import sky_logging
 from sky.client import common as client_common
+from sky.client import interactive_utils
 from sky.client import sdk
 from sky.schemas.api import responses
 from sky.server import common as server_common
@@ -167,9 +168,17 @@ async def stream_response_async(request_id: Optional[str],
         retry_context = rest.get_retry_context()
     try:
         line_count = 0
+
         async for line in rich_utils.decode_rich_status_async(response):
             if line is not None:
                 line_count += 1
+
+                line = await interactive_utils.handle_interactive_auth_async(
+                    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: