skypilot-nightly 1.0.0.dev20250521__py3-none-any.whl → 1.0.0.dev20250523__py3-none-any.whl

sky/utils/kubernetes/kubernetes_deploy_utils.py

@@ -2,9 +2,12 @@
 import os
 import shlex
 import subprocess
+import sys
 import tempfile
 from typing import List, Optional
 
+import colorama
+
 from sky import check as sky_check
 from sky import sky_logging
 from sky.backends import backend_utils
@@ -19,6 +22,103 @@ from sky.utils import ux_utils
 
 logger = sky_logging.init_logger(__name__)
 
+# Default path for Kubernetes configuration file
+DEFAULT_KUBECONFIG_PATH = os.path.expanduser('~/.kube/config')
+
+
+def deploy_ssh_cluster(cleanup: bool = False,
+                       infra: Optional[str] = None,
+                       kubeconfig_path: Optional[str] = None):
+    """Deploy a Kubernetes cluster on SSH targets.
+
+    This function reads ~/.sky/ssh_node_pools.yaml and uses it to deploy a
+    Kubernetes cluster on the specified machines.
+
+    Args:
+        cleanup: Whether to clean up the cluster instead of deploying.
+        infra: Name of the cluster in ssh_node_pools.yaml to use.
+            If None, the first cluster in the file will be used.
+        kubeconfig_path: Path to save the Kubernetes configuration file.
+            If None, the default ~/.kube/config will be used.
+    """
+    # Prepare command to call deploy_remote_cluster.py script
+    # TODO(romilb): We should move this to a native python method/class call
+    #  instead of invoking a script with subprocess.
+    path_to_package = os.path.dirname(__file__)
+    up_script_path = os.path.join(path_to_package, 'deploy_remote_cluster.py')
+    cwd = os.path.dirname(os.path.abspath(up_script_path))
+
+    deploy_command = [sys.executable, up_script_path]
+
+    if cleanup:
+        deploy_command.append('--cleanup')
+
+    if infra:
+        deploy_command.extend(['--infra', infra])
+
+    # Use the default kubeconfig path if none is provided
+    kubeconfig_path = kubeconfig_path or DEFAULT_KUBECONFIG_PATH
+    deploy_command.extend(['--kubeconfig-path', kubeconfig_path])
+
+    # Setup logging paths
+    run_timestamp = sky_logging.get_run_timestamp()
+    log_path = os.path.join(constants.SKY_LOGS_DIRECTORY, run_timestamp,
+                            'ssh_up.log')
+
+    if cleanup:
+        msg_str = 'Cleaning up SSH Node Pools...'
+    else:
+        msg_str = 'Initializing deployment to SSH Node Pools...'
+
+    # Create environment with PYTHONUNBUFFERED=1 to ensure unbuffered output
+    env = os.environ.copy()
+    env['PYTHONUNBUFFERED'] = '1'
+
+    with rich_utils.safe_status(
+            ux_utils.spinner_message(msg_str, log_path=log_path,
+                                     is_local=True)):
+        returncode, _, stderr = log_lib.run_with_log(
+            cmd=deploy_command,
+            log_path=log_path,
+            require_outputs=True,
+            stream_logs=False,  # TODO: Fixme to False after we fix the logging
+            line_processor=log_utils.SkySSHUpLineProcessor(log_path=log_path,
+                                                           is_local=True),
+            cwd=cwd,
+            env=env)
+
+    if returncode == 0:
+        success = True
+    else:
+        with ux_utils.print_exception_no_traceback():
+            log_hint = ux_utils.log_path_hint(log_path, is_local=True)
+            raise RuntimeError('Failed to deploy SkyPilot on SSH targets. '
+                               f'Full log: {log_hint}'
+                               f'\nError: {stderr}')
+
+    if success:
+        # Add an empty line to separate the deployment logs from the final
+        # message
+        logger.info('')
+        if cleanup:
+            logger.info(
+                ux_utils.finishing_message(
+                    '🎉 SSH Node Pools cleaned up successfully.',
+                    log_path=log_path,
+                    is_local=True))
+        else:
+            logger.info(
+                ux_utils.finishing_message(
+                    '🎉 SSH Node Pools set up successfully. ',
+                    follow_up_message=(
+                        f'Run `{colorama.Style.BRIGHT}'
+                        f'sky check ssh'
+                        f'{colorama.Style.RESET_ALL}` to verify access, '
+                        f'`{colorama.Style.BRIGHT}sky launch --infra ssh'
+                        f'{colorama.Style.RESET_ALL}` to launch a cluster. '),
+                    log_path=log_path,
+                    is_local=True))
+
 
 def deploy_remote_cluster(ip_list: List[str],
                           ssh_user: str,
@@ -28,7 +128,7 @@ def deploy_remote_cluster(ip_list: List[str],
                           password: Optional[str] = None):
     success = False
     path_to_package = os.path.dirname(__file__)
-    up_script_path = os.path.join(path_to_package, 'deploy_remote_cluster.sh')
+    up_script_path = os.path.join(path_to_package, 'deploy_remote_cluster.py')
     # Get directory of script and run it from there
     cwd = os.path.dirname(os.path.abspath(up_script_path))
 
@@ -44,17 +144,18 @@ def deploy_remote_cluster(ip_list: List[str],
         key_file.flush()
         os.chmod(key_file.name, 0o600)
 
-        deploy_command = (f'{up_script_path} {ip_file.name} '
-                          f'{ssh_user} {key_file.name}')
+        # Use the legacy mode command line arguments for backward compatibility
+        deploy_command = [
+            sys.executable, up_script_path, '--ips-file', ip_file.name,
+            '--user', ssh_user, '--ssh-key', key_file.name
+        ]
+
         if context_name is not None:
-            deploy_command += f' {context_name}'
+            deploy_command.extend(['--context-name', context_name])
         if password is not None:
-            deploy_command += f' --password {password}'
+            deploy_command.extend(['--password', password])
         if cleanup:
-            deploy_command += ' --cleanup'
-
-        # Convert the command to a format suitable for subprocess
-        deploy_command = shlex.split(deploy_command)
+            deploy_command.append('--cleanup')
 
         # Setup logging paths
         run_timestamp = sky_logging.get_run_timestamp()
@@ -65,6 +166,11 @@ def deploy_remote_cluster(ip_list: List[str],
             msg_str = 'Cleaning up remote cluster...'
         else:
             msg_str = 'Deploying remote cluster...'
+
+        # Create environment with PYTHONUNBUFFERED=1 to ensure unbuffered output
+        env = os.environ.copy()
+        env['PYTHONUNBUFFERED'] = '1'
+
         with rich_utils.safe_status(
                 ux_utils.spinner_message(msg_str,
                                          log_path=log_path,
@@ -76,7 +182,8 @@ def deploy_remote_cluster(ip_list: List[str],
                 stream_logs=False,
                 line_processor=log_utils.SkyRemoteUpLineProcessor(
                     log_path=log_path, is_local=True),
-                cwd=cwd)
+                cwd=cwd,
+                env=env)
         if returncode == 0:
             success = True
         else:

sky/utils/kubernetes/ssh-tunnel.sh

@@ -0,0 +1,387 @@
+#!/bin/bash
+# ssh-tunnel.sh - SSH tunnel script for Kubernetes API access
+# Used as kubectl exec credential plugin to establish SSH tunnel on demand.
+# Returns a valid credential format for kubectl with expiration. The expiration
+# is calculated based on the TTL argument and is required to force kubectl to
+# check the tunnel status frequently.
+
+# Usage: ssh-tunnel.sh --host HOST [--user USER] [--use-ssh-config] [--ssh-key KEY] [--context CONTEXT] [--port PORT] [--ttl SECONDS]
+
+# Default time-to-live for credential in seconds
+# This forces kubectl to check the tunnel status frequently
+TTL_SECONDS=30
+
+# Parse arguments
+USE_SSH_CONFIG=0
+SSH_KEY=""
+CONTEXT=""
+HOST=""
+USER=""
+PORT=6443  # Default port if not specified
+
+# Debug log to ~/.sky/ssh_node_pools_info/$CONTEXT-tunnel.log
+debug_log() {
+    local message="$(date): $1"
+    echo "$message" >> "$LOG_FILE"
+}
+
+# Generate expiration timestamp for credential
+generate_expiration_timestamp() {
+  # Try macOS date format first, fallback to Linux format
+  date -u -v+${TTL_SECONDS}S +"%Y-%m-%dT%H:%M:%SZ" 2>/dev/null || date -u -d "+${TTL_SECONDS} seconds" +"%Y-%m-%dT%H:%M:%SZ"
+}
+
+# Acquire the lock, return 0 if successful, 1 if another process is already holding the lock
+acquire_lock() {
+  # Check for flock command
+  if ! command -v flock >/dev/null 2>&1; then
+    debug_log "flock command not available, using alternative lock mechanism"
+    # Simple file-based locking
+    if [ -f "$LOCK_FILE" ]; then
+      lock_pid=$(cat "$LOCK_FILE" 2>/dev/null)
+      if [ -n "$lock_pid" ] && kill -0 "$lock_pid" 2>/dev/null; then
+        debug_log "Another process ($lock_pid) is starting the tunnel, waiting briefly"
+        return 1
+      else
+        # Stale lock file
+        debug_log "Removing stale lock file"
+        rm -f "$LOCK_FILE"
+      fi
+    fi
+    # Create our lock
+    echo $$ > "$LOCK_FILE"
+    return 0
+  else
+    # Use flock for better locking
+    exec 9>"$LOCK_FILE"
+    if ! flock -n 9; then
+      debug_log "Another process is starting the tunnel, waiting briefly"
+      return 1
+    fi
+    return 0
+  fi
+}
+
+# Release the lock
+release_lock() {
+  if command -v flock >/dev/null 2>&1; then
+    # Using flock
+    exec 9>&- # Close file descriptor to release lock
+  else
+    # Using simple lock
+    rm -f "$LOCK_FILE"
+  fi
+  debug_log "Lock released"
+}
+
+# Generate SSH command based on available tools and parameters
+generate_ssh_command() {
+  # Check for autossh
+  if ! command -v autossh >/dev/null 2>&1; then
+    debug_log "WARNING: autossh is not installed but recommended for reliable SSH tunnels"
+    debug_log "Install autossh: brew install autossh (macOS), apt-get install autossh (Ubuntu/Debian)"
+    
+    # Fall back to regular ssh
+    if [[ $USE_SSH_CONFIG -eq 1 ]]; then
+      SSH_CMD=("ssh" "-o" "ServerAliveInterval=30" "-o" "ServerAliveCountMax=3" "-o" "ExitOnForwardFailure=yes" "-L" "$PORT:127.0.0.1:6443" "-N" "$HOST")
+    else
+      SSH_CMD=("ssh" "-o" "StrictHostKeyChecking=no" "-o" "IdentitiesOnly=yes" "-o" "ServerAliveInterval=30" "-o" "ServerAliveCountMax=3" "-o" "ExitOnForwardFailure=yes" "-L" "$PORT:127.0.0.1:6443" "-N")
+      
+      # Add SSH key if provided
+      if [[ -n "$SSH_KEY" ]]; then
+        SSH_CMD+=("-i" "$SSH_KEY")
+      fi
+      
+      # Add user@host
+      SSH_CMD+=("$USER@$HOST")
+    fi
+  else
+    # Configure autossh
+    if [[ $USE_SSH_CONFIG -eq 1 ]]; then
+      SSH_CMD=("autossh" "-M" "0" "-o" "ServerAliveInterval=30" "-o" "ServerAliveCountMax=3" "-o" "ExitOnForwardFailure=yes" "-L" "$PORT:127.0.0.1:6443" "-N" "$HOST")
+    else
+      SSH_CMD=("autossh" "-M" "0" "-o" "StrictHostKeyChecking=no" "-o" "IdentitiesOnly=yes" "-o" "ServerAliveInterval=30" "-o" "ServerAliveCountMax=3" "-o" "ExitOnForwardFailure=yes" "-L" "$PORT:127.0.0.1:6443" "-N")
+      
+      # Add SSH key if provided
+      if [[ -n "$SSH_KEY" ]]; then
+        SSH_CMD+=("-i" "$SSH_KEY")
+      fi
+      
+      # Add user@host
+      SSH_CMD+=("$USER@$HOST")
+    fi
+  fi
+}
+
+# Function to read certificate files if they exist
+read_certificate_data() {
+  local client_cert_file="$TUNNEL_DIR/$CONTEXT-cert.pem"
+  local client_key_file="$TUNNEL_DIR/$CONTEXT-key.pem"
+  local cert_data=""
+  local key_data=""
+  
+  if [[ -f "$client_cert_file" ]]; then
+    # Read the certificate file as is - it's already in PEM format
+    cert_data=$(cat "$client_cert_file")
+    debug_log "Found client certificate data for context $CONTEXT"
+    
+    # Log the first and last few characters to verify PEM format
+    local cert_start=$(head -1 "$client_cert_file")
+    local cert_end=$(tail -1 "$client_cert_file")
+    debug_log "Certificate starts with: $cert_start"
+    debug_log "Certificate ends with: $cert_end"
+    
+    # Check if it has proper PEM format
+    if ! grep -q "BEGIN CERTIFICATE" "$client_cert_file" || ! grep -q "END CERTIFICATE" "$client_cert_file"; then
+      debug_log "WARNING: Certificate file may not be in proper PEM format"
+      # Try to fix it if needed
+      if ! grep -q "BEGIN CERTIFICATE" "$client_cert_file"; then
+        echo "-----BEGIN CERTIFICATE-----" > "$client_cert_file.fixed"
+        cat "$client_cert_file" >> "$client_cert_file.fixed"
+        echo "-----END CERTIFICATE-----" >> "$client_cert_file.fixed"
+        mv "$client_cert_file.fixed" "$client_cert_file"
+        cert_data=$(cat "$client_cert_file")
+        debug_log "Fixed certificate format by adding BEGIN/END markers"
+      fi
+    fi
+  fi
+  
+  if [[ -f "$client_key_file" ]]; then
+    # Read the key file as is - it's already in PEM format
+    key_data=$(cat "$client_key_file")
+    debug_log "Found client key data for context $CONTEXT"
+    
+    # Log the first and last few characters to verify PEM format
+    local key_start=$(head -1 "$client_key_file")
+    local key_end=$(tail -1 "$client_key_file")
+    debug_log "Key starts with: $key_start"
+    debug_log "Key ends with: $key_end"
+    
+    # Check if it has proper PEM format
+    if ! grep -q "BEGIN" "$client_key_file" || ! grep -q "END" "$client_key_file"; then
+      debug_log "WARNING: Key file may not be in proper PEM format"
+      # Try to fix it if needed
+      if ! grep -q "BEGIN" "$client_key_file"; then
+        echo "-----BEGIN PRIVATE KEY-----" > "$client_key_file.fixed"
+        cat "$client_key_file" >> "$client_key_file.fixed"
+        echo "-----END PRIVATE KEY-----" >> "$client_key_file.fixed"
+        mv "$client_key_file.fixed" "$client_key_file"
+        key_data=$(cat "$client_key_file")
+        debug_log "Fixed key format by adding BEGIN/END markers"
+      fi
+    fi
+  fi
+  
+  echo "$cert_data:$key_data"
+}
+
+# Function to generate credentials JSON
+generate_credentials_json() {
+  local expiration_time=$(generate_expiration_timestamp)
+  local cert_bundle=$(read_certificate_data)
+  local client_cert_data=${cert_bundle%:*}
+  local client_key_data=${cert_bundle#*:}
+  
+  if [[ -n "$client_cert_data" && -n "$client_key_data" ]]; then
+    # Debug the certificate data
+    debug_log "Certificate data length: $(echo -n "$client_cert_data" | wc -c) bytes"
+    debug_log "Key data length: $(echo -n "$client_key_data" | wc -c) bytes"
+    
+    # Check if we can create proper JSON with `jq`
+    if command -v jq &>/dev/null; then
+      debug_log "Using jq for JSON formatting"
+      
+      # Create a temporary file for the JSON output to avoid shell escaping issues
+      local TEMP_JSON_FILE=$(mktemp)
+      
+      # Write the JSON to the temporary file using jq for proper JSON formatting
+      cat > "$TEMP_JSON_FILE" << EOL
+{
+  "apiVersion": "client.authentication.k8s.io/v1beta1",
+  "kind": "ExecCredential",
+  "status": {
+    "clientCertificateData": $(printf '%s' "$client_cert_data" | jq -R -s .),
+    "clientKeyData": $(printf '%s' "$client_key_data" | jq -R -s .),
+    "expirationTimestamp": "$expiration_time"
+  }
+}
+EOL
+      
+      # Read the JSON from the file
+      local json_response=$(cat "$TEMP_JSON_FILE")
+      
+      # Clean up
+      rm -f "$TEMP_JSON_FILE"
+      
+      # Output the JSON
+      echo "$json_response"
+    else
+      debug_log "jq is not available, using simpler formatting method"
+      
+      # Alternative approach: encode with base64 and use the token field instead
+      # This works because kubectl will decode token data properly
+      local combined_data=$(echo -n "${client_cert_data}:${client_key_data}" | base64 | tr -d '\n')
+      
+      echo "{\"apiVersion\":\"client.authentication.k8s.io/v1beta1\",\"kind\":\"ExecCredential\",\"status\":{\"token\":\"$combined_data\",\"expirationTimestamp\":\"$expiration_time\"}}"
+      
+      debug_log "Sent certificate data as encoded token instead of direct certificate fields"
+    fi
+  else
+    # Fallback to token-based credential for tunnel-only authentication
+    echo "{\"apiVersion\":\"client.authentication.k8s.io/v1beta1\",\"kind\":\"ExecCredential\",\"status\":{\"token\":\"k8s-ssh-tunnel-token\",\"expirationTimestamp\":\"$expiration_time\"}}"
+  fi
+}
+
+while [[ $# -gt 0 ]]; do
+  case $1 in
+    --use-ssh-config)
+      USE_SSH_CONFIG=1
+      shift
+      ;;
+    --ssh-key)
+      SSH_KEY="$2"
+      shift 2
+      ;;
+    --context)
+      CONTEXT="$2"
+      shift 2
+      ;;
+    --port)
+      PORT="$2"
+      shift 2
+      ;;
+    --host)
+      HOST="$2"
+      shift 2
+      ;;
+    --user)
+      USER="$2"
+      shift 2
+      ;;
+    --ttl)
+      TTL_SECONDS="$2"
+      shift 2
+      ;;
+    *)
+      echo "Unknown parameter: $1" >&2
+      exit 1
+      ;;
+  esac
+done
+
+# Validate required parameters
+if [[ -z "$HOST" ]]; then
+  echo "Error: --host parameter is required" >&2
+  exit 1
+fi
+
+# Setup directories
+TUNNEL_DIR="$HOME/.sky/ssh_node_pools_info"
+mkdir -p "$TUNNEL_DIR"
+
+# Get context name for PID file
+if [[ -z "$CONTEXT" ]]; then
+  CONTEXT="default"
+fi
+
+PID_FILE="$TUNNEL_DIR/$CONTEXT-tunnel.pid"
+LOG_FILE="$TUNNEL_DIR/$CONTEXT-tunnel.log"
+LOCK_FILE="$TUNNEL_DIR/$CONTEXT-tunnel.lock"
+
+debug_log "Starting ssh-tunnel.sh for context $CONTEXT, host $HOST, port $PORT"
+debug_log "SSH Config: $USE_SSH_CONFIG, User: $USER, TTL: ${TTL_SECONDS}s"
+
+# Check if specified port is already in use (tunnel may be running)
+if nc -z 127.0.0.1 "$PORT" 2>/dev/null; then
+  debug_log "Port $PORT already in use, checking if it's our tunnel"
+  
+  # Check if there's a PID file and if that process is running
+  if [[ -f "$PID_FILE" ]]; then
+    OLD_PID=$(cat "$PID_FILE")
+    if kill -0 "$OLD_PID" 2>/dev/null; then
+      debug_log "Tunnel appears to be running with PID $OLD_PID"
+    else
+      debug_log "PID file exists but process $OLD_PID is not running"
+    fi
+  else
+    debug_log "Port $PORT is in use but no PID file exists"
+  fi
+  
+  # Return valid credential format for kubectl with expiration
+  generate_credentials_json
+  exit 0
+fi
+
+# Try to acquire the lock
+if ! acquire_lock; then
+  # Wait briefly for the tunnel to be established
+  for i in {1..10}; do
+    if nc -z 127.0.0.1 "$PORT" 2>/dev/null; then
+      debug_log "Tunnel is now active"
+      
+      # Return valid credential format for kubectl with expiration
+      generate_credentials_json
+      exit 0
+    fi
+    sleep 0.2
+  done
+  debug_log "Waited for tunnel but port $PORT still not available"
+fi
+
+# Check if we have a PID file with running process
+if [[ -f "$PID_FILE" ]]; then
+  OLD_PID=$(cat "$PID_FILE")
+  if kill -0 "$OLD_PID" 2>/dev/null; then
+    # Process exists but port isn't open - something's wrong, kill it
+    kill "$OLD_PID" 2>/dev/null
+    debug_log "Killed stale tunnel process $OLD_PID"
+  else
+    debug_log "PID file exists but process $OLD_PID is not running anymore"
+  fi
+  # Remove the stale PID file
+  rm -f "$PID_FILE"
+fi
+
+# Generate the SSH command
+generate_ssh_command
+
+debug_log "Starting SSH tunnel: ${SSH_CMD[*]}"
+
+# Start the tunnel in foreground and wait for it to establish
+"${SSH_CMD[@]}" >> "$LOG_FILE" 2>&1 &
+TUNNEL_PID=$!
+
+# Save PID
+echo $TUNNEL_PID > "$PID_FILE"
+debug_log "Tunnel started with PID $TUNNEL_PID"
+
+# Wait for tunnel to establish
+tunnel_up=0
+for i in {1..20}; do
+  if nc -z 127.0.0.1 "$PORT" 2>/dev/null; then
+    debug_log "Tunnel established successfully on port $PORT"
+    tunnel_up=1
+    break
+  fi
+  sleep 0.2
+done
+
+# Clean up lock file
+release_lock
+
+# Check if the tunnel process is still running
+if ! kill -0 $TUNNEL_PID 2>/dev/null; then
+  debug_log "ERROR: Tunnel process exited unexpectedly! Check logs for details"
+  if [[ -f "$PID_FILE" ]]; then
+    rm -f "$PID_FILE"
+  fi
+  # Return error in case of tunnel failure
+  echo "Failed to establish SSH tunnel. See $TUNNEL_DIR/$CONTEXT-tunnel.log for details." >&2
+  exit 1
+elif [[ $tunnel_up -eq 0 ]]; then
+  debug_log "WARNING: Tunnel process is running but port $PORT is not responding"
+fi
+
+# Return valid credential format with certificates if available
+generate_credentials_json
+exit 0