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

sky/ssh_node_pools/server.py

@@ -0,0 +1,232 @@
+"""SSH Node Pool management API endpoints."""
+import re
+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 requests as requests_lib
+from sky.ssh_node_pools import core as ssh_node_pools_core
+from sky.utils import common_utils
+
+router = fastapi.APIRouter()
+
+
+@router.get('')
+async def get_ssh_node_pools() -> Dict[str, Any]:
+    """Get all SSH Node Pool configurations."""
+    try:
+        return ssh_node_pools_core.get_all_pools()
+    except Exception as e:
+        raise fastapi.HTTPException(
+            status_code=500,
+            detail=
+            f'Failed to get SSH Node Pools: {common_utils.format_exception(e)}')
+
+
+@router.post('')
+async 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)
+        return {'status': 'success'}
+    except Exception as e:
+        raise fastapi.HTTPException(status_code=400,
+                                    detail=f'Failed to update SSH Node Pools:'
+                                    f' {common_utils.format_exception(e)}')
+
+
+@router.delete('/{pool_name}')
+async 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):
+            return {'status': 'success'}
+        else:
+            raise fastapi.HTTPException(
+                status_code=404,
+                detail=f'SSH Node Pool `{pool_name}` not found')
+    except fastapi.HTTPException:
+        raise
+    except Exception as e:
+        raise fastapi.HTTPException(status_code=500,
+                                    detail='Failed to delete SSH Node Pool: '
+                                    f'{common_utils.format_exception(e)}')
+
+
+@router.post('/keys')
+async def upload_ssh_key(request: fastapi.Request) -> Dict[str, str]:
+    """Upload SSH private key."""
+    try:
+        form = await request.form()
+        key_name = form.get('key_name')
+        key_file = form.get('key_file')
+
+        if not key_name or not key_file:
+            raise fastapi.HTTPException(status_code=400,
+                                        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())
+
+        return {'status': 'success', 'key_path': key_path}
+    except fastapi.HTTPException:
+        raise
+    except Exception as e:
+        raise fastapi.HTTPException(
+            status_code=500,
+            detail=
+            f'Failed to upload SSH key: {common_utils.format_exception(e)}')
+
+
+@router.get('/keys')
+async def list_ssh_keys() -> List[str]:
+    """List available SSH keys."""
+    try:
+        return ssh_node_pools_core.list_ssh_keys()
+    except Exception as e:
+        exception_msg = common_utils.format_exception(e)
+        raise fastapi.HTTPException(
+            status_code=500, detail=f'Failed to list SSH keys: {exception_msg}')
+
+
+@router.post('/{pool_name}/deploy')
+async def deploy_ssh_node_pool(request: fastapi.Request,
+                               pool_name: str) -> Dict[str, str]:
+    """Deploy SSH Node Pool using existing ssh_up functionality."""
+    try:
+        ssh_up_body = payloads.SSHUpBody(infra=pool_name, cleanup=False)
+        executor.schedule_request(
+            request_id=request.state.request_id,
+            request_name='ssh_up',
+            request_body=ssh_up_body,
+            func=sky_core.ssh_up,
+            schedule_type=requests_lib.ScheduleType.LONG,
+        )
+
+        return {
+            'status': 'success',
+            'request_id': request.state.request_id,
+            'message': f'SSH Node Pool `{pool_name}` deployment started'
+        }
+    except Exception as e:
+        raise fastapi.HTTPException(status_code=500,
+                                    detail=f'Failed to deploy SSH Node Pool: '
+                                    f'{common_utils.format_exception(e)}')
+
+
+@router.post('/deploy')
+async def deploy_ssh_node_pool_general(
+        request: fastapi.Request,
+        ssh_up_body: payloads.SSHUpBody) -> Dict[str, str]:
+    """Deploys all SSH Node Pools."""
+    try:
+        executor.schedule_request(
+            request_id=request.state.request_id,
+            request_name='ssh_up',
+            request_body=ssh_up_body,
+            func=sky_core.ssh_up,
+            schedule_type=requests_lib.ScheduleType.LONG,
+        )
+
+        pool_name = ssh_up_body.infra or 'default'
+        return {
+            'status': 'success',
+            'request_id': request.state.request_id,
+            'message': f'SSH Node Pool `{pool_name}` deployment started'
+        }
+    except Exception as e:
+        raise fastapi.HTTPException(status_code=500,
+                                    detail=f'Failed to deploy SSH Node Pool: '
+                                    f'{common_utils.format_exception(e)}')
+
+
+@router.post('/{pool_name}/down')
+async def down_ssh_node_pool(request: fastapi.Request,
+                             pool_name: str) -> Dict[str, str]:
+    """Cleans up a SSH Node Pools."""
+    try:
+        ssh_up_body = payloads.SSHUpBody(infra=pool_name, cleanup=True)
+        executor.schedule_request(
+            request_id=request.state.request_id,
+            request_name='ssh_down',
+            request_body=ssh_up_body,
+            func=sky_core.ssh_up,  # Reuse ssh_up function with cleanup=True
+            schedule_type=requests_lib.ScheduleType.LONG,
+        )
+
+        return {
+            'status': 'success',
+            'request_id': request.state.request_id,
+            'message': f'SSH Node Pool `{pool_name}` teardown started'
+        }
+    except Exception as e:
+        raise fastapi.HTTPException(
+            status_code=500,
+            detail=f'Failed to tear down SSH Node Pool: '
+            f'{common_utils.format_exception(e)}')
+
+
+@router.post('/down')
+async def down_ssh_node_pool_general(
+        request: fastapi.Request,
+        ssh_up_body: payloads.SSHUpBody) -> Dict[str, str]:
+    """Cleans up all SSH Node Pools."""
+    try:
+        # Set cleanup=True for down operation
+        ssh_up_body.cleanup = True
+        executor.schedule_request(
+            request_id=request.state.request_id,
+            request_name='ssh_down',
+            request_body=ssh_up_body,
+            func=sky_core.ssh_up,  # Reuse ssh_up function with cleanup=True
+            schedule_type=requests_lib.ScheduleType.LONG,
+        )
+
+        pool_name = ssh_up_body.infra or 'default'
+        return {
+            'status': 'success',
+            'request_id': request.state.request_id,
+            'message': f'SSH Node Pool `{pool_name}` teardown started'
+        }
+    except Exception as e:
+        raise fastapi.HTTPException(
+            status_code=500,
+            detail=f'Failed to tear down SSH Node Pool: '
+            f'{common_utils.format_exception(e)}')
+
+
+@router.get('/{pool_name}/status')
+async def get_ssh_node_pool_status(pool_name: str) -> Dict[str, str]:
+    """Get the status of a specific SSH Node Pool."""
+    try:
+        # Call ssh_status to check the context
+        context_name = f'ssh-{pool_name}'
+        is_ready, reason = sky_core.ssh_status(context_name)
+
+        # Strip ANSI escape codes from the reason text
+        def strip_ansi_codes(text):
+            if not text:
+                return text
+            # Remove ANSI escape sequences (color codes, formatting, etc.)
+            text = re.sub(r'\x1b\[[0-9;]*m', '', text)
+            # Remove 'disabled. Reason: ' prefix if present
+            text = text.replace('disabled. Reason: ', '')
+            return text
+
+        cleaned_reason = strip_ansi_codes(reason) if reason else reason
+
+        return {
+            'pool_name': pool_name,
+            'context_name': context_name,
+            'status': 'Ready' if is_ready else 'Not Ready',
+            'reason': cleaned_reason
+        }
+    except Exception as e:
+        raise fastapi.HTTPException(
+            status_code=500,
+            detail=f'Failed to get SSH Node Pool status: '
+            f'{common_utils.format_exception(e)}')

sky/utils/kubernetes/deploy_remote_cluster.py

@@ -11,11 +11,12 @@ import shutil
 import subprocess
 import sys
 import tempfile
-from typing import Any, Dict, List, Optional, Set
+from typing import List, Set
 
 import yaml
 
 from sky.utils import ux_utils
+from sky.utils.kubernetes import ssh_utils
 
 # Colors for nicer UX
 RED = '\033[0;31m'
@@ -24,7 +25,6 @@ YELLOW = '\033[1;33m'
 WARNING_YELLOW = '\x1b[33m'
 NC = '\033[0m'  # No color
 
-DEFAULT_SSH_NODE_POOLS_PATH = os.path.expanduser('~/.sky/ssh_node_pools.yaml')
 DEFAULT_KUBECONFIG_PATH = os.path.expanduser('~/.kube/config')
 SSH_CONFIG_PATH = os.path.expanduser('~/.ssh/config')
 NODE_POOLS_INFO_DIR = os.path.expanduser('~/.sky/ssh_node_pools_info')
@@ -33,29 +33,6 @@ NODE_POOLS_INFO_DIR = os.path.expanduser('~/.sky/ssh_node_pools_info')
 SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
 
 
-class UniqueKeySafeLoader(yaml.SafeLoader):
-    """Custom YAML loader that raises an error if there are duplicate keys."""
-
-    def construct_mapping(self, node, deep=False):
-        mapping = {}
-        for key_node, value_node in node.value:
-            key = self.construct_object(key_node, deep=deep)
-            if key in mapping:
-                raise yaml.constructor.ConstructorError(
-                    note=(f'Duplicate cluster config for cluster {key!r}.\n'
-                          'Please remove one of them from: '
-                          f'{DEFAULT_SSH_NODE_POOLS_PATH}'))
-            value = self.construct_object(value_node, deep=deep)
-            mapping[key] = value
-        return mapping
-
-
-# Register the custom constructor inside the class
-UniqueKeySafeLoader.add_constructor(
-    yaml.resolver.BaseResolver.DEFAULT_MAPPING_TAG,
-    UniqueKeySafeLoader.construct_mapping)
-
-
 def parse_args():
     parser = argparse.ArgumentParser(
         description='Deploy a Kubernetes cluster on remote machines.')
@@ -64,9 +41,9 @@ def parse_args():
     parser.add_argument(
         '--ssh-node-pools-file',
         dest='ssh_node_pools_file',
-        default=DEFAULT_SSH_NODE_POOLS_PATH,
+        default=ssh_utils.DEFAULT_SSH_NODE_POOLS_PATH,
         help=
-        f'Path to SSH node pools YAML file (default: {DEFAULT_SSH_NODE_POOLS_PATH})'
+        f'Path to SSH node pools YAML file (default: {ssh_utils.DEFAULT_SSH_NODE_POOLS_PATH})'
     )
     parser.add_argument(
         '--kubeconfig-path',
@@ -117,156 +94,6 @@ def parse_args():
     return parser.parse_args()
 
 
-def load_ssh_targets(file_path: str) -> Dict[str, Any]:
-    """Load SSH targets from YAML file."""
-    if not os.path.exists(file_path):
-        with ux_utils.print_exception_no_traceback():
-            raise ValueError(f'SSH Node Pools file not found: {file_path}')
-
-    try:
-        with open(file_path, 'r', encoding='utf-8') as f:
-            targets = yaml.load(f, Loader=UniqueKeySafeLoader)
-        return targets
-    except yaml.constructor.ConstructorError as e:
-        with ux_utils.print_exception_no_traceback():
-            raise ValueError(e.note) from e
-    except (yaml.YAMLError, IOError, OSError) as e:
-        with ux_utils.print_exception_no_traceback():
-            raise ValueError(f'Error loading SSH Node Pools file: {e}') from e
-
-
-def check_host_in_ssh_config(hostname: str) -> bool:
-    """Return True iff *hostname* matches at least one `Host`/`Match` stanza
-    in the user's OpenSSH client configuration (including anything pulled in
-    via Include).
-
-    It calls:  ssh -vvG <hostname> -o ConnectTimeout=0
-    which:
-      • -G  expands the effective config without connecting
-      • -vv prints debug lines that show which stanzas are applied
-      • ConnectTimeout=0 avoids a DNS lookup if <hostname> is a FQDN/IP
-
-    No config files are opened or parsed manually.
-
-    Parameters
-    ----------
-    hostname : str
-        The alias/IP/FQDN you want to test.
-
-    Returns
-    -------
-    bool
-        True  – a specific stanza matched the host
-        False – nothing but the global defaults (`Host *`) applied
-    """
-    # We direct stderr→stdout because debug output goes to stderr.
-    proc = subprocess.run(
-        ['ssh', '-vvG', hostname, '-o', 'ConnectTimeout=0'],
-        text=True,
-        stdout=subprocess.PIPE,
-        stderr=subprocess.STDOUT,
-        check=False,  # we only want the text, not to raise
-    )
-
-    # Look for lines like:
-    #   debug1: ~/.ssh/config line 42: Applying options for <hostname>
-    # Anything other than "*"
-    pattern = re.compile(r'^debug\d+: .*Applying options for ([^*].*)$',
-                         re.MULTILINE)
-
-    return bool(pattern.search(proc.stdout))
-
-
-def get_cluster_config(targets: Dict[str, Any],
-                       cluster_name: Optional[str] = None,
-                       file_path: Optional[str] = None) -> Dict[str, Any]:
-    """Get configuration for specific clusters or all clusters."""
-    if not targets:
-        with ux_utils.print_exception_no_traceback():
-            raise ValueError(
-                f'No clusters defined in SSH Node Pools file {file_path}')
-
-    if cluster_name:
-        if cluster_name not in targets:
-            with ux_utils.print_exception_no_traceback():
-                raise ValueError(f'Cluster {cluster_name!r} not found in '
-                                 f'SSH Node Pools file {file_path}')
-        return {cluster_name: targets[cluster_name]}
-
-    # Return all clusters if no specific cluster is specified
-    return targets
-
-
-def prepare_hosts_info(cluster_name: str,
-                       cluster_config: Dict[str, Any]) -> List[Dict[str, str]]:
-    """Prepare list of hosts with resolved user, identity_file, and password."""
-    if 'hosts' not in cluster_config or not cluster_config['hosts']:
-        with ux_utils.print_exception_no_traceback():
-            raise ValueError(
-                f'No hosts defined in cluster {cluster_name} configuration')
-
-    # Get cluster-level defaults
-    cluster_user = cluster_config.get('user', '')
-    cluster_identity_file = os.path.expanduser(
-        cluster_config.get('identity_file', ''))
-    cluster_password = cluster_config.get('password', '')
-
-    # Check if cluster identity file exists
-    if cluster_identity_file and not os.path.isfile(cluster_identity_file):
-        with ux_utils.print_exception_no_traceback():
-            raise ValueError(
-                f'SSH Identity File Missing: {cluster_identity_file}')
-
-    hosts_info = []
-    for host in cluster_config['hosts']:
-        # Host can be a string (IP or SSH config hostname) or a dict
-        if isinstance(host, str):
-            # Check if this is an SSH config hostname
-            is_ssh_config_host = check_host_in_ssh_config(host)
-
-            hosts_info.append({
-                'ip': host,
-                'user': '' if is_ssh_config_host else cluster_user,
-                'identity_file': '' if is_ssh_config_host else
-                                 cluster_identity_file,
-                'password': cluster_password,
-                'use_ssh_config': is_ssh_config_host
-            })
-        else:
-            # It's a dict with potential overrides
-            if 'ip' not in host:
-                print(
-                    f'{RED}Warning: Host missing \'ip\' field, skipping: {host}{NC}'
-                )
-                continue
-
-            # Check if this is an SSH config hostname
-            is_ssh_config_host = check_host_in_ssh_config(host['ip'])
-
-            # Use host-specific values or fall back to cluster defaults
-            host_user = '' if is_ssh_config_host else host.get(
-                'user', cluster_user)
-            host_identity_file = os.path.expanduser(
-                '' if is_ssh_config_host else host.
-                get('identity_file', cluster_identity_file))
-            host_password = host.get('password', cluster_password)
-
-            if host_identity_file and not os.path.isfile(host_identity_file):
-                with ux_utils.print_exception_no_traceback():
-                    raise ValueError(
-                        f'SSH Identity File Missing: {host_identity_file}')
-
-            hosts_info.append({
-                'ip': host['ip'],
-                'user': host_user,
-                'identity_file': host_identity_file,
-                'password': host_password,
-                'use_ssh_config': is_ssh_config_host
-            })
-
-    return hosts_info
-
-
 def run_command(cmd, shell=False):
     """Run a local command and return the output."""
     process = subprocess.run(cmd,
@@ -675,10 +502,10 @@ def main():
         password = args.password
 
         # Check if hosts are in SSH config
-        head_use_ssh_config = global_use_ssh_config or check_host_in_ssh_config(
+        head_use_ssh_config = global_use_ssh_config or ssh_utils.check_host_in_ssh_config(
             head_node)
         worker_use_ssh_config = [
-            global_use_ssh_config or check_host_in_ssh_config(node)
+            global_use_ssh_config or ssh_utils.check_host_in_ssh_config(node)
             for node in worker_nodes
         ]
 
@@ -688,10 +515,9 @@ def main():
                        kubeconfig_path, args.cleanup)
     else:
         # Using YAML configuration
-        targets = load_ssh_targets(args.ssh_node_pools_file)
-        clusters_config = get_cluster_config(targets,
-                                             args.infra,
-                                             file_path=args.ssh_node_pools_file)
+        targets = ssh_utils.load_ssh_targets(args.ssh_node_pools_file)
+        clusters_config = ssh_utils.get_cluster_config(
+            targets, args.infra, file_path=args.ssh_node_pools_file)
 
         # Print information about clusters being processed
         num_clusters = len(clusters_config)
@@ -705,7 +531,8 @@ def main():
                 print(f'SKYPILOT_CURRENT_CLUSTER: {cluster_name}')
                 print(
                     f'{YELLOW}==== Deploying cluster: {cluster_name} ====${NC}')
-                hosts_info = prepare_hosts_info(cluster_name, cluster_config)
+                hosts_info = ssh_utils.prepare_hosts_info(
+                    cluster_name, cluster_config)
 
                 if not hosts_info:
                     print(
@@ -744,7 +571,7 @@ def main():
                                 f'Cluster configuration has changed for field {key!r}. '
                                 f'Previous value: {history.get(key)}, '
                                 f'Current value: {cluster_config.get(key)}')
-                    history_hosts_info = prepare_hosts_info(
+                    history_hosts_info = ssh_utils.prepare_hosts_info(
                         cluster_name, history)
                     if not args.cleanup and history_hosts_info[0] != hosts_info[
                             0]: