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

sky/ssh_node_pools/deploy/tunnel_utils.py

@@ -0,0 +1,199 @@
+"""Utilities to setup SSH Tunnel"""
+import os
+import random
+import re
+import subprocess
+import sys
+from typing import Set
+
+import colorama
+
+from sky import sky_logging
+from sky.ssh_node_pools import constants
+from sky.ssh_node_pools.deploy import utils as deploy_utils
+
+logger = sky_logging.init_logger(__name__)
+
+# Get the directory of this script
+SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
+
+
+def _get_used_localhost_ports() -> Set[int]:
+    """Get SSH port forwardings already in use on localhost"""
+    used_ports = set()
+
+    # Get ports from netstat (works on macOS and Linux)
+    try:
+        if sys.platform == 'darwin':
+            # macOS
+            result = subprocess.run(['netstat', '-an', '-p', 'tcp'],
+                                    capture_output=True,
+                                    text=True,
+                                    check=False)
+        else:
+            # Linux and other Unix-like systems
+            result = subprocess.run(['netstat', '-tln'],
+                                    capture_output=True,
+                                    text=True,
+                                    check=False)
+
+        if result.returncode == 0:
+            # Look for lines with 'localhost:<port>' or '127.0.0.1:<port>'
+            for line in result.stdout.splitlines():
+                if '127.0.0.1:' in line or 'localhost:' in line:
+                    match = re.search(r':(64\d\d)\s', line)
+                    if match:
+                        port = int(match.group(1))
+                        if 6400 <= port <= 6500:  # Only consider our range
+                            used_ports.add(port)
+    except (subprocess.SubprocessError, FileNotFoundError):
+        # If netstat fails, try another approach
+        pass
+
+    # Also check ports from existing kubeconfig entries
+    try:
+        result = subprocess.run([
+            'kubectl', 'config', 'view', '-o',
+            'jsonpath=\'{.clusters[*].cluster.server}\''
+        ],
+                                capture_output=True,
+                                text=True,
+                                check=False)
+
+        if result.returncode == 0:
+            # Look for localhost URLs with ports
+            for url in result.stdout.split():
+                if 'localhost:' in url or '127.0.0.1:' in url:
+                    match = re.search(r':(\d+)', url)
+                    if match:
+                        port = int(match.group(1))
+                        if 6400 <= port <= 6500:  # Only consider our range
+                            used_ports.add(port)
+    except subprocess.SubprocessError:
+        pass
+
+    return used_ports
+
+
+def get_available_port(start: int = 6443, end: int = 6499) -> int:
+    """Get an available port in the given range not used by other tunnels"""
+    used_ports = _get_used_localhost_ports()
+
+    # Try to use port 6443 first if available for the first cluster
+    if start == 6443 and start not in used_ports:
+        return start
+
+    # Otherwise find any available port in the range
+    available_ports = list(set(range(start, end + 1)) - used_ports)
+
+    if not available_ports:
+        # If all ports are used, pick a random one from our range
+        # (we'll terminate any existing connection in the setup)
+        return random.randint(start, end)
+
+    # Sort to get deterministic allocation
+    available_ports.sort()
+    return available_ports[0]
+
+
+def setup_kubectl_ssh_tunnel(head_node,
+                             ssh_user,
+                             ssh_key,
+                             context_name,
+                             use_ssh_config=False):
+    """Set up kubeconfig exec credential plugin for SSH tunnel"""
+    logger.info(f'{colorama.Fore.YELLOW}➜ Setting up SSH tunnel for '
+                f'Kubernetes API access...{colorama.Style.RESET_ALL}')
+
+    # Get an available port for this cluster
+    port = get_available_port()
+
+    # Paths to scripts
+    tunnel_script = os.path.join(SCRIPT_DIR, 'tunnel', 'ssh-tunnel.sh')
+
+    # Make sure scripts are executable
+    os.chmod(tunnel_script, 0o755)
+
+    # Certificate files
+    client_cert_file = os.path.join(constants.NODE_POOLS_INFO_DIR,
+                                    f'{context_name}-cert.pem')
+    client_key_file = os.path.join(constants.NODE_POOLS_INFO_DIR,
+                                   f'{context_name}-key.pem')
+
+    # Update kubeconfig to use localhost with the selected port
+    deploy_utils.run_command([
+        'kubectl', 'config', 'set-cluster', context_name,
+        f'--server=https://127.0.0.1:{port}', '--insecure-skip-tls-verify=true'
+    ])
+
+    # Build the exec args list based on auth method
+    exec_args = [
+        '--exec-command', tunnel_script, '--exec-api-version',
+        'client.authentication.k8s.io/v1beta1'
+    ]
+
+    # Set credential TTL to force frequent tunnel checks
+    ttl_seconds = 30
+
+    # Verify if we have extracted certificate data files
+    has_cert_files = os.path.isfile(client_cert_file) and os.path.isfile(
+        client_key_file)
+    if has_cert_files:
+        logger.info(f'{colorama.Fore.GREEN}Client certificate data extracted '
+                    'and will be used for authentication'
+                    f'{colorama.Style.RESET_ALL}')
+
+    if use_ssh_config:
+        deploy_utils.run_command(
+            ['kubectl', 'config', 'set-credentials', context_name] + exec_args +
+            [
+                '--exec-arg=--context', f'--exec-arg={context_name}',
+                '--exec-arg=--port', f'--exec-arg={port}', '--exec-arg=--ttl',
+                f'--exec-arg={ttl_seconds}', '--exec-arg=--use-ssh-config',
+                '--exec-arg=--host', f'--exec-arg={head_node}'
+            ])
+    else:
+        deploy_utils.run_command(
+            ['kubectl', 'config', 'set-credentials', context_name] + exec_args +
+            [
+                '--exec-arg=--context', f'--exec-arg={context_name}',
+                '--exec-arg=--port', f'--exec-arg={port}', '--exec-arg=--ttl',
+                f'--exec-arg={ttl_seconds}', '--exec-arg=--host',
+                f'--exec-arg={head_node}', '--exec-arg=--user',
+                f'--exec-arg={ssh_user}', '--exec-arg=--ssh-key',
+                f'--exec-arg={ssh_key}'
+            ])
+
+    logger.info(f'{colorama.Fore.GREEN}✔ SSH tunnel configured through '
+                'kubectl credential plugin on port '
+                f'{port}{colorama.Style.RESET_ALL}')
+    logger.info('Your kubectl connection is now tunneled through SSH '
+                f'(port {port}).')
+    logger.info('This tunnel will be automatically established when needed.')
+    logger.info(f'Credential TTL set to {ttl_seconds}s to ensure tunnel '
+                'health is checked frequently.')
+    return port
+
+
+def cleanup_kubectl_ssh_tunnel(cluster_name, context_name):
+    """Clean up the SSH tunnel for a specific context"""
+    logger.info(f'{colorama.Fore.YELLOW}➜ Cleaning up SSH tunnel for '
+                f'`{cluster_name}`...{colorama.Style.RESET_ALL}')
+
+    # Path to cleanup script
+    cleanup_script = os.path.join(SCRIPT_DIR, 'tunnel', 'cleanup-tunnel.sh')
+
+    # Make sure script is executable
+    if os.path.exists(cleanup_script):
+        os.chmod(cleanup_script, 0o755)
+
+        # Run the cleanup script
+        subprocess.run([cleanup_script, context_name],
+                       stdout=subprocess.DEVNULL,
+                       stderr=subprocess.DEVNULL,
+                       check=False)
+        logger.info(f'{colorama.Fore.GREEN}✔ SSH tunnel for `{cluster_name}` '
+                    f'cleaned up.{colorama.Style.RESET_ALL}')
+    else:
+        logger.error(f'{colorama.Fore.YELLOW}Cleanup script not found: '
+                     f'{cleanup_script}{colorama.Style.RESET_ALL}')

sky/ssh_node_pools/deploy/utils.py

@@ -0,0 +1,173 @@
+"""Utilities for SSH Node Pools Deployment"""
+import os
+import subprocess
+from typing import List, Optional
+
+import colorama
+
+from sky import sky_logging
+from sky.utils import ux_utils
+
+logger = sky_logging.init_logger(__name__)
+
+
+def check_ssh_cluster_dependencies(
+        raise_error: bool = True) -> Optional[List[str]]:
+    """Checks if the dependencies for ssh cluster are installed.
+
+    Args:
+        raise_error: set to true when the dependency needs to be present.
+            set to false for `sky check`, where reason strings are compiled
+            at the end.
+
+    Returns: the reasons list if there are missing dependencies.
+    """
+    # error message
+    jq_message = ('`jq` is required to setup ssh cluster.')
+
+    # save
+    reasons = []
+    required_binaries = []
+
+    # Ensure jq is installed
+    try:
+        subprocess.run(['jq', '--version'],
+                       stdout=subprocess.DEVNULL,
+                       stderr=subprocess.DEVNULL,
+                       check=True)
+    except (FileNotFoundError, subprocess.CalledProcessError):
+        required_binaries.append('jq')
+        reasons.append(jq_message)
+
+    if required_binaries:
+        reasons.extend([
+            'On Debian/Ubuntu, install the missing dependenc(ies) with:',
+            f'  $ sudo apt install {" ".join(required_binaries)}',
+            'On MacOS, install with: ',
+            f'  $ brew install {" ".join(required_binaries)}',
+        ])
+        if raise_error:
+            with ux_utils.print_exception_no_traceback():
+                raise RuntimeError('\n'.join(reasons))
+        return reasons
+    return None
+
+
+def run_command(cmd, shell=False, silent=False):
+    """Run a local command and return the output."""
+    process = subprocess.run(cmd,
+                             shell=shell,
+                             capture_output=True,
+                             text=True,
+                             check=False)
+    if process.returncode != 0:
+        if not silent:
+            logger.error(f'{colorama.Fore.RED}Error executing command: {cmd}\n'
+                         f'{colorama.Style.RESET_ALL}STDOUT: {process.stdout}\n'
+                         f'STDERR: {process.stderr}')
+        return None
+    return process.stdout.strip()
+
+
+def get_effective_host_ip(hostname: str) -> str:
+    """Get the effective IP for a hostname from SSH config."""
+    try:
+        result = subprocess.run(['ssh', '-G', hostname],
+                                capture_output=True,
+                                text=True,
+                                check=False)
+        if result.returncode == 0:
+            for line in result.stdout.splitlines():
+                if line.startswith('hostname '):
+                    return line.split(' ', 1)[1].strip()
+    except Exception:  # pylint: disable=broad-except
+        pass
+    return hostname  # Return the original hostname if lookup fails
+
+
+def run_remote(node,
+               cmd,
+               user='',
+               ssh_key='',
+               connect_timeout=30,
+               use_ssh_config=False,
+               print_output=False,
+               use_shell=False,
+               silent=False):
+    """Run a command on a remote machine via SSH."""
+    ssh_cmd: List[str]
+    if use_ssh_config:
+        # Use SSH config for connection parameters
+        ssh_cmd = ['ssh', node, cmd]
+    else:
+        # Use explicit parameters
+        ssh_cmd = [
+            'ssh', '-o', 'StrictHostKeyChecking=no', '-o', 'IdentitiesOnly=yes',
+            '-o', f'ConnectTimeout={connect_timeout}', '-o',
+            'ServerAliveInterval=10', '-o', 'ServerAliveCountMax=3'
+        ]
+
+        if ssh_key:
+            if not os.path.isfile(ssh_key):
+                raise ValueError(f'SSH key not found: {ssh_key}')
+            ssh_cmd.extend(['-i', ssh_key])
+
+        ssh_cmd.append(f'{user}@{node}' if user else node)
+        ssh_cmd.append(cmd)
+
+    subprocess_cmd = ' '.join(ssh_cmd) if use_shell else ssh_cmd
+    process = subprocess.run(subprocess_cmd,
+                             capture_output=True,
+                             text=True,
+                             check=False,
+                             shell=use_shell)
+    if process.returncode != 0:
+        if not silent:
+            logger.error(f'{colorama.Fore.RED}Error executing command {cmd} on '
+                         f'{node}:{colorama.Style.RESET_ALL} {process.stderr}')
+        return None
+    if print_output:
+        logger.info(process.stdout)
+    return process.stdout.strip()
+
+
+def ensure_directory_exists(path):
+    """Ensure the directory for the specified file path exists."""
+    directory = os.path.dirname(path)
+    if directory and not os.path.exists(directory):
+        os.makedirs(directory, exist_ok=True)
+
+
+def check_gpu(node, user, ssh_key, use_ssh_config=False, is_head=False):
+    """Check if a node has a GPU."""
+    cmd = ('command -v nvidia-smi &> /dev/null && '
+           'nvidia-smi --query-gpu=gpu_name --format=csv,noheader')
+    result = run_remote(node,
+                        cmd,
+                        user,
+                        ssh_key,
+                        use_ssh_config=use_ssh_config,
+                        silent=True)
+    if result is not None:
+        # Check that all GPUs have the same type.
+        # Currently, SkyPilot does not support heterogeneous GPU node
+        # (i.e. more than one GPU type on the same node).
+        gpu_names = {
+            line.strip() for line in result.splitlines() if line.strip()
+        }
+        if not gpu_names:
+            # This can happen if nvidia-smi returns only whitespace.
+            # Set result to None to ensure this function returns False.
+            result = None
+        elif len(gpu_names) > 1:
+            # Sort for a deterministic error message.
+            sorted_gpu_names = sorted(list(gpu_names))
+            raise RuntimeError(
+                f'Node {node} has more than one GPU types '
+                f'({", ".join(sorted_gpu_names)}). '
+                'SkyPilot does not support a node with multiple GPU types.')
+        else:
+            logger.info(f'{colorama.Fore.YELLOW}➜ GPU {list(gpu_names)[0]} '
+                        f'detected on {"head" if is_head else "worker"} '
+                        f'node ({node}).{colorama.Style.RESET_ALL}')
+    return result is not None

sky/ssh_node_pools/server.py

@@ -4,11 +4,11 @@ from typing import Any, Dict, List
 
 import fastapi
 
-from sky import core as sky_core
 from sky.server.requests import executor
 from sky.server.requests import payloads
+from sky.server.requests import request_names
 from sky.server.requests import requests as requests_lib
-from sky.ssh_node_pools import core as ssh_node_pools_core
+from sky.ssh_node_pools import core
 from sky.utils import common_utils
 
 router = fastapi.APIRouter()
@@ -18,7 +18,7 @@ router = fastapi.APIRouter()
 def get_ssh_node_pools() -> Dict[str, Any]:
     """Get all SSH Node Pool configurations."""
     try:
-        return ssh_node_pools_core.get_all_pools()
+        return core.get_all_pools()
     except Exception as e:
         raise fastapi.HTTPException(
             status_code=500,
@@ -30,7 +30,7 @@ def get_ssh_node_pools() -> Dict[str, Any]:
 def update_ssh_node_pools(pools_config: Dict[str, Any]) -> Dict[str, str]:
     """Update SSH Node Pool configurations."""
     try:
-        ssh_node_pools_core.update_pools(pools_config)
+        core.update_pools(pools_config)
         return {'status': 'success'}
     except Exception as e:
         raise fastapi.HTTPException(status_code=400,
@@ -42,7 +42,7 @@ def update_ssh_node_pools(pools_config: Dict[str, Any]) -> Dict[str, str]:
 def delete_ssh_node_pool(pool_name: str) -> Dict[str, str]:
     """Delete a SSH Node Pool configuration."""
     try:
-        if ssh_node_pools_core.delete_pool(pool_name):
+        if core.delete_pool(pool_name):
             return {'status': 'success'}
         else:
             raise fastapi.HTTPException(
@@ -69,8 +69,7 @@ async def upload_ssh_key(request: fastapi.Request) -> Dict[str, str]:
                                         detail='Missing key_name or key_file')
 
         key_content = await key_file.read()
-        key_path = ssh_node_pools_core.upload_ssh_key(key_name,
-                                                      key_content.decode())
+        key_path = core.upload_ssh_key(key_name, key_content.decode())
 
         return {'status': 'success', 'key_path': key_path}
     except fastapi.HTTPException:
@@ -86,7 +85,7 @@ async def upload_ssh_key(request: fastapi.Request) -> Dict[str, str]:
 def list_ssh_keys() -> List[str]:
     """List available SSH keys."""
     try:
-        return ssh_node_pools_core.list_ssh_keys()
+        return core.list_ssh_keys()
     except Exception as e:
         exception_msg = common_utils.format_exception(e)
         raise fastapi.HTTPException(
@@ -99,11 +98,11 @@ async def deploy_ssh_node_pool(request: fastapi.Request,
     """Deploy SSH Node Pool using existing ssh_up functionality."""
     try:
         ssh_up_body = payloads.SSHUpBody(infra=pool_name, cleanup=False)
-        executor.schedule_request(
+        await executor.schedule_request_async(
             request_id=request.state.request_id,
-            request_name='ssh_up',
+            request_name=request_names.RequestName.SSH_NODE_POOLS_UP,
             request_body=ssh_up_body,
-            func=sky_core.ssh_up,
+            func=core.ssh_up,
             schedule_type=requests_lib.ScheduleType.LONG,
         )
 
@@ -124,11 +123,11 @@ async def deploy_ssh_node_pool_general(
         ssh_up_body: payloads.SSHUpBody) -> Dict[str, str]:
     """Deploys all SSH Node Pools."""
     try:
-        executor.schedule_request(
+        await executor.schedule_request_async(
             request_id=request.state.request_id,
-            request_name='ssh_up',
+            request_name=request_names.RequestName.SSH_NODE_POOLS_UP,
             request_body=ssh_up_body,
-            func=sky_core.ssh_up,
+            func=core.ssh_up,
             schedule_type=requests_lib.ScheduleType.LONG,
         )
 
@@ -150,11 +149,11 @@ async def down_ssh_node_pool(request: fastapi.Request,
     """Cleans up a SSH Node Pools."""
     try:
         ssh_up_body = payloads.SSHUpBody(infra=pool_name, cleanup=True)
-        executor.schedule_request(
+        await executor.schedule_request_async(
             request_id=request.state.request_id,
-            request_name='ssh_down',
+            request_name=request_names.RequestName.SSH_NODE_POOLS_DOWN,
             request_body=ssh_up_body,
-            func=sky_core.ssh_up,  # Reuse ssh_up function with cleanup=True
+            func=core.ssh_up,  # Reuse ssh_up function with cleanup=True
             schedule_type=requests_lib.ScheduleType.LONG,
         )
 
@@ -178,11 +177,11 @@ async def down_ssh_node_pool_general(
     try:
         # Set cleanup=True for down operation
         ssh_up_body.cleanup = True
-        executor.schedule_request(
+        await executor.schedule_request_async(
             request_id=request.state.request_id,
-            request_name='ssh_down',
+            request_name=request_names.RequestName.SSH_NODE_POOLS_DOWN,
             request_body=ssh_up_body,
-            func=sky_core.ssh_up,  # Reuse ssh_up function with cleanup=True
+            func=core.ssh_up,  # Reuse ssh_up function with cleanup=True
             schedule_type=requests_lib.ScheduleType.LONG,
         )
 
@@ -205,7 +204,7 @@ def get_ssh_node_pool_status(pool_name: str) -> Dict[str, str]:
     try:
         # Call ssh_status to check the context
         context_name = f'ssh-{pool_name}'
-        is_ready, reason = sky_core.ssh_status(context_name)
+        is_ready, reason = core.ssh_status(context_name)
 
         # Strip ANSI escape codes from the reason text
         def strip_ansi_codes(text):

sky/{utils/kubernetes/ssh_utils.py → ssh_node_pools/utils.py}

@@ -5,13 +5,14 @@ import subprocess
 from typing import Any, Callable, Dict, List, Optional
 import uuid
 
+import colorama
 import yaml
 
+from sky import sky_logging
+from sky.ssh_node_pools import constants
 from sky.utils import ux_utils
 
-DEFAULT_SSH_NODE_POOLS_PATH = os.path.expanduser('~/.sky/ssh_node_pools.yaml')
-RED = '\033[0;31m'
-NC = '\033[0m'  # No color
+logger = sky_logging.init_logger(__name__)
 
 
 def check_host_in_ssh_config(hostname: str) -> bool:
@@ -92,7 +93,8 @@ def load_ssh_targets(file_path: str) -> Dict[str, Any]:
 def get_cluster_config(
         targets: Dict[str, Any],
         cluster_name: Optional[str] = None,
-        file_path: str = DEFAULT_SSH_NODE_POOLS_PATH) -> Dict[str, Any]:
+        file_path: str = constants.DEFAULT_SSH_NODE_POOLS_PATH
+) -> Dict[str, Any]:
     """Get configuration for specific clusters or all clusters."""
     if not targets:
         with ux_utils.print_exception_no_traceback():
@@ -186,8 +188,9 @@ def prepare_hosts_info(
         else:
             # It's a dict with potential overrides
             if 'ip' not in host:
-                print(f'{RED}Warning: Host missing \'ip\' field, '
-                      f'skipping: {host}{NC}')
+                logger.warning(f'{colorama.Fore.RED}Warning: Host missing'
+                               f'\'ip\' field, skipping: {host}'
+                               f'{colorama.Style.RESET_ALL}')
                 continue
 
             # Check if this is an SSH config hostname