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

sky/ssh_node_pools/deploy/tunnel/ssh-tunnel.sh

@@ -0,0 +1,379 @@
+#!/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
+      echo "jq is not installed. Please install jq to use this script." >&2
+      exit 1
+    fi
+    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
+    # 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 

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}')