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

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
@@ -42,6 +43,7 @@ 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
@@ -57,7 +59,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
@@ -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:
@@ -675,7 +683,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
@@ -2744,3 +2752,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

@@ -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:

sky/clouds/__init__.py

@@ -31,6 +31,7 @@ from sky.clouds.runpod import RunPod
 from sky.clouds.scp import SCP
 from sky.clouds.seeweb import Seeweb
 from sky.clouds.shadeform import Shadeform
+from sky.clouds.slurm import Slurm
 from sky.clouds.ssh import SSH
 from sky.clouds.vast import Vast
 from sky.clouds.vsphere import Vsphere
@@ -48,6 +49,7 @@ __all__ = [
     'Paperspace',
     'PrimeIntellect',
     'SCP',
+    'Slurm',
     'RunPod',
     'Shadeform',
     'Vast',

sky/clouds/aws.py

@@ -12,6 +12,7 @@ import typing
 from typing import (Any, Callable, Dict, Iterator, List, Literal, Optional, Set,
                     Tuple, TypeVar, Union)
 
+import colorama
 from typing_extensions import ParamSpec
 
 from sky import catalog
@@ -758,6 +759,36 @@ class AWS(clouds.Cloud):
             max_efa_interfaces = 0
             enable_efa = False
 
+        use_internal_ips = skypilot_config.get_effective_region_config(
+            cloud='aws',
+            region=region_name,
+            keys=('use_internal_ips',),
+            default_value=False)
+        if max_efa_interfaces > 1 and not use_internal_ips:
+            logger.warning(
+                f'{colorama.Fore.YELLOW}'
+                f'Instance type {resources.instance_type} supports up to '
+                f'{max_efa_interfaces} EFA interfaces, but '
+                '`use_internal_ips` is not enabled.\nLaunching with the '
+                'current configuration will use only 1 EFA interface.\n'
+                f'To use all {max_efa_interfaces} EFA interfaces, enable '
+                'internal IPs by adding one of the following '
+                'configurations to SkyPilot config:\n'
+                'Option 1 (with SSM):\n'
+                '  aws:\n'
+                '    use_internal_ips: true\n'
+                '    use_ssm: true\n'
+                'Option 2 (with SSH proxy):\n'
+                '  aws:\n'
+                '    use_internal_ips: true\n'
+                '    ssh_proxy_command: ssh -W %h:%p -i <ssh key path> '
+                '-o StrictHostKeyChecking=no <user>@<jump server public'
+                ' ip>\n'
+                'Refer to '
+                'https://docs.skypilot.co/en/latest/reference/config.html'
+                '#aws-use-internal-ips for more details.'
+                f'{colorama.Style.RESET_ALL}')
+
         docker_run_options = []
         if resources.extract_docker_image() is not None:
             image_id_to_use = None
@@ -1005,8 +1036,10 @@ class AWS(clouds.Cloud):
             hints = 'AWS SSO is set.'
             if static_credential_exists:
                 hints += (
-                    ' To ensure multiple clouds work correctly, please use SkyPilot '
-                    'with static credentials (e.g., ~/.aws/credentials) by unsetting '
+                    ' To ensure S3 mounting and other features work correctly '
+                    'on Kubernetes and other clouds, '
+                    'please use SkyPilot with static AWS credentials '
+                    '(e.g., ~/.aws/credentials) by unsetting '
                     'the AWS_PROFILE environment variable.')
             else:
                 hints += single_cloud_hint
@@ -1081,6 +1114,31 @@ class AWS(clouds.Cloud):
                 return identity_type
         return AWSIdentityType.SHARED_CREDENTIALS_FILE
 
+    @classmethod
+    def should_use_env_auth_for_s3(cls) -> bool:
+        """Returns True if S3 should use environment-based auth.
+
+        When using non-static AWS credentials (SSO, IAM role, container role),
+        we should not embed credentials into rclone config. Instead, we should
+        use env_auth=true so that rclone uses the AWS SDK credential chain,
+        which properly handles temporary credentials and IAM roles.
+
+        Returns:
+            True if environment-based auth should be used, False for static
+            credentials that can be embedded.
+        """
+        identity_type = cls._current_identity_type()
+        if identity_type is None:
+            return False
+        # These credential types use temporary credentials that should not be
+        # embedded in config files. They rely on the AWS SDK credential chain.
+        non_static_types = {
+            AWSIdentityType.SSO,
+            AWSIdentityType.IAM_ROLE,
+            AWSIdentityType.CONTAINER_ROLE,
+        }
+        return identity_type in non_static_types
+
     @classmethod
     @aws_profile_aware_lru_cache(scope='request',
                                  maxsize=_AWS_PROFILE_SCOPED_FUNC_CACHE_SIZE)

sky/clouds/azure.py

@@ -97,6 +97,8 @@ class Azure(clouds.Cloud):
             clouds.CloudImplementationFeatures.HIGH_AVAILABILITY_CONTROLLERS: (
                 f'High availability controllers are not supported on {cls._REPR}.'
             ),
+            clouds.CloudImplementationFeatures.CUSTOM_NETWORK_TIER:
+                (f'Custom network tier is not supported on {cls._REPR}.'),
             clouds.CloudImplementationFeatures.CUSTOM_MULTI_NETWORK: (
                 f'Customized multiple network interfaces are not supported on {cls._REPR}.'
             ),

sky/clouds/cloud.py

@@ -182,6 +182,13 @@ class Cloud:
         """
         return cls._SUPPORTS_SERVICE_ACCOUNT_ON_REMOTE
 
+    @classmethod
+    def uses_ray(cls) -> bool:
+        """Returns whether this cloud uses Ray as the distributed
+        execution framework.
+        """
+        return True
+
     #### Regions/Zones ####
 
     @classmethod

sky/clouds/kubernetes.py

@@ -766,6 +766,8 @@ class Kubernetes(clouds.Cloud):
             'ha_recovery_log_path':
                 constants.HA_PERSISTENT_RECOVERY_LOG_PATH.format(''),
             'sky_python_cmd': constants.SKY_PYTHON_CMD,
+            'sky_unset_pythonpath_and_set_cwd':
+                constants.SKY_UNSET_PYTHONPATH_AND_SET_CWD,
             'k8s_high_availability_storage_class_name':
                 (k8s_ha_storage_class_name),
             'avoid_label_keys': avoid_label_keys,

sky/clouds/runpod.py

@@ -7,6 +7,7 @@ from typing import Dict, Iterator, List, Optional, Tuple, Union
 
 from sky import catalog
 from sky import clouds
+from sky.utils import common_utils
 from sky.utils import registry
 from sky.utils import resources_utils
 
@@ -312,18 +313,48 @@ class RunPod(clouds.Cloud):
             # If that happens to be set to None, then ValueError is raised.
             return False, dependency_error_msg
 
+        hint_msg = (
+            'Credentials can be set up by running: \n'
+            '        $ pip install runpod \n'
+            '        $ runpod config\n'
+            '    For more information, see https://docs.skypilot.co/en/latest/getting-started/installation.html#runpod'  # pylint: disable=line-too-long
+        )
+
         valid, error = cls._check_runpod_credentials()
         if not valid:
-            return False, (
-                f'{error} \n'  # First line is indented by 4 spaces
-                '    Credentials can be set up by running: \n'
-                f'        $ pip install runpod \n'
-                f'        $ runpod config\n'
-                '    For more information, see https://docs.skypilot.co/en/latest/getting-started/installation.html#runpod'  # pylint: disable=line-too-long
-            )
+            return False, (f'{error} \n    {hint_msg}')
+
+        # Validate credentials by making an actual API call
+        valid, error = cls._validate_api_key()
+        if not valid:
+            return False, (f'{error} \n    {hint_msg}')
 
         return True, None
 
+    @classmethod
+    def _validate_api_key(cls) -> Tuple[bool, Optional[str]]:
+        """Validate RunPod API key by making an actual API call."""
+        # Import here to avoid circular imports and ensure runpod is configured
+        # pylint: disable=import-outside-toplevel
+        from sky.provision.runpod import utils as runpod_utils
+        try:
+            # Try to list instances to validate the API key works
+            runpod_utils.list_instances()
+            return True, None
+        except Exception as e:  # pylint: disable=broad-except
+            from sky.adaptors import runpod
+            error_msg = common_utils.format_exception(e, use_bracket=True)
+            if isinstance(e, runpod.runpod.error.QueryError):
+                error_msg_lower = str(e).lower()
+                auth_keywords = ['unauthorized', 'forbidden', '401', '403']
+                if any(keyword in error_msg_lower for keyword in auth_keywords):
+                    return False, (
+                        'RunPod API key is invalid or lacks required '
+                        f'permissions. {error_msg}')
+                return False, (f'Failed to verify RunPod API key. {error_msg}')
+            return False, ('An unexpected error occurred during RunPod API '
+                           f'key validation. {error_msg}')
+
     @classmethod
     def _check_runpod_credentials(cls, profile: str = 'default'):
         """Checks if the credentials file exists and is valid."""