rayssh.py

#!/usr/bin/env python3
"""
RaySSH - Ray-native terminal tool
WebSocket-based terminal communication using Ray actors.
"""

import asyncio
import argparse
import os
import signal
import subprocess
import sys
import json
from datetime import datetime

from terminal import RaySSHTerminal
from cli import (
    get_random_worker_node,
    get_node_by_index,
    print_nodes_table,
    interactive_node_selector,
    submit_file_job,
    submit_shell_command,
    handle_lab_command,
    handle_code_command,
    handle_debug_command,
    handle_tell_cursor_command,
    handle_push_command,
    handle_pull_command,
    handle_sync_command,
)
from utils import (
    ensure_ray_initialized,
    load_last_session_preferred_ip,
    find_node_by_ip,
    parse_require_constraints_from_env,
    select_worker_node,
    parse_n_gpus_from_env,
)


def signal_handler(signum, frame):
    """Handle Ctrl+C gracefully."""
    print(f"\nReceived signal {signum}, initiating graceful shutdown...")
    # The RaySSHTerminal class has its own signal handlers
    # This is a fallback for any unhandled signals
    sys.exit(0)


def main():
    """Main entry point."""
    # Check if RAY_ADDRESS is set for remote mode
    if not (ray_address := os.environ.get("RAY_ADDRESS")):
        print(
            "Error: RAY_ADDRESS env var is not set. RaySSH is intended for remote-cluster use. To ensure consistency between local and remote processes, please set RAY_ADDRESS"
            "in your environment. It can either be 'ray://<ip>:<port>' (Ray client server) or '<ip>:<port>' (Ray GCS server).",
            file=sys.stderr,
        )
        return 1

    working_dir = None
    node_arg = None

    # Get require constraints from environment
    require_constraints = parse_require_constraints_from_env()

    # Handle special commands first
    if len(sys.argv) == 1:
        working_dir = None  # No working_dir means HOME
        try:
            ensure_ray_initialized(ray_address=ray_address)
            # Try previous session first
            prefer_ip = load_last_session_preferred_ip()
            if prefer_ip:
                try:
                    node_info = find_node_by_ip(
                        prefer_ip, resource_constraints=require_constraints
                    )
                    if node_info and node_info.get("Alive"):
                        node_arg = prefer_ip
                    else:
                        prefer_ip = None
                except Exception:
                    prefer_ip = None

            if not prefer_ip:
                # Use our custom node selection logic
                n_gpus = parse_n_gpus_from_env()
                selected_node = select_worker_node(n_gpus=n_gpus)
                node_arg = selected_node.get("NodeManagerAddress")
        except ValueError as e:
            print(f"Error: {e}")
            return 1
        except Exception as e:
            print(f"Error selecting random worker node: {e}", file=sys.stderr)
            return 1

    elif len(sys.argv) == 2:
        argument = sys.argv[1]

        # Handle help command
        if argument in ["--help", "-h"]:
            print_help()
            return 0

        # Handle lab subcommand
        elif argument == "lab":
            return handle_lab_command(["lab"] + sys.argv[2:])

        # Handle code subcommand
        elif argument == "code":
            return handle_code_command(["code"] + sys.argv[2:])

        # Handle debug subcommand
        elif argument == "debug":
            return handle_debug_command(["debug"] + sys.argv[2:])

        # Handle sync subcommand
        elif argument == "sync":
            return handle_sync_command(["sync"] + sys.argv[2:])

        # Handle tell-cursor subcommand
        elif argument == "tell-cursor":
            return handle_tell_cursor_command(["tell-cursor"] + sys.argv[2:])

        # Handle push/pull subcommands
        elif argument == "push":
            return handle_push_command(["push"] + sys.argv[2:])
        elif argument == "pull":
            return handle_pull_command(["pull"] + sys.argv[2:])

        # Handle special commands
        elif argument in ["--ls"]:
            return print_nodes_table()
        elif argument in ["--list", "--show", "-l"]:
            selected_node_ip = interactive_node_selector()
            if selected_node_ip is None:
                print("\n❌ Cancelled.")
                return 0
            node_arg = selected_node_ip

        # Check if it's a file for job submission
        elif os.path.exists(argument) and os.path.isfile(argument) and "." in argument:
            # It's a file - submit as Ray job
            return submit_file_job(argument, no_wait=False)

        # Check if it's a directory and RAY_ADDRESS is set
        elif os.path.exists(argument) and os.path.isdir(argument):
            # It's a directory and we're in remote mode - upload and connect
            working_dir = argument
            # print(f"📁 Directory specified: {argument}")
        # Check if it's a directory but no RAY_ADDRESS

        # Handle node index argument (-0, -1, -2, etc.)
        elif argument.startswith("-") and argument[1:].isdigit():
            try:
                ensure_ray_initialized(ray_address=ray_address)
                index = int(argument[1:])  # Extract number after '-'
                node = get_node_by_index(index)
                # Use the node's IP address as the connection target
                node_arg = node.get("NodeManagerAddress")
                print(f"🔗 Connecting to node -{index}: {node_arg}")
            except ValueError as e:
                print(f"Error: {e}", file=sys.stderr)
                print("Use 'rayssh --ls' to see available nodes", file=sys.stderr)
                return 1
            except Exception as e:
                print(f"Error getting node by index: {e}", file=sys.stderr)
                return 1

        # Otherwise, treat as node argument
        else:
            node_arg = argument

    elif len(sys.argv) == 3:
        # Handle -q file pattern for quick job submission
        if sys.argv[1] == "-q":
            potential_file = sys.argv[2]
            if (
                os.path.exists(potential_file)
                and os.path.isfile(potential_file)
                and "." in potential_file
            ):
                return submit_file_job(potential_file, no_wait=True)
            else:
                print(
                    f"Error: File '{potential_file}' not found or not a valid file",
                    file=sys.stderr,
                )
                return 1

        # Handle lab, code, debug, sync, and tell-cursor subcommands with a single extra argument
        elif sys.argv[1] == "lab":
            return handle_lab_command(["lab"] + sys.argv[2:])
        elif sys.argv[1] == "code":
            return handle_code_command(["code"] + sys.argv[2:])
        elif sys.argv[1] == "debug":
            return handle_debug_command(["debug"] + sys.argv[2:])
        elif sys.argv[1] == "sync":
            return handle_sync_command(["sync"] + sys.argv[2:])
        elif sys.argv[1] == "tell-cursor":
            return handle_tell_cursor_command(["tell-cursor"] + sys.argv[2:])
        elif sys.argv[1] == "push":
            return handle_push_command(["push"] + sys.argv[2:])
        elif sys.argv[1] == "pull":
            return handle_pull_command(["pull"] + sys.argv[2:])
        # Handle -- <command>
        elif sys.argv[1] == "--":
            import shlex

            return submit_shell_command(shlex.join(sys.argv[2:]))

        # Handle -0 lab, -0 code, and -0 debug patterns
        elif sys.argv[1] == "-0":
            if sys.argv[2] == "lab":
                return handle_lab_command(["-0", "lab"] + sys.argv[3:])
            elif sys.argv[2] == "code":
                return handle_code_command(["-0", "code"] + sys.argv[3:])
            elif sys.argv[2] == "debug":
                return handle_debug_command(["-0", "debug"] + sys.argv[3:])
            else:
                # Fall through to node index handling
                pass
        else:
            # treat it as a node argument followed by a path
            node_arg = sys.argv[1]
            working_dir = sys.argv[2]
            if not os.path.exists(working_dir) or not os.path.isdir(working_dir):
                print(
                    f"Error: Directory '{working_dir}' not found or not a valid directory",
                    file=sys.stderr,
                )
                return 1

    elif len(sys.argv) > 3:
        # Handle lab, code, debug, sync, and tell-cursor commands with additional arguments (like paths)
        if sys.argv[1] == "lab":
            return handle_lab_command(["lab"] + sys.argv[2:])
        elif sys.argv[1] == "code":
            return handle_code_command(["code"] + sys.argv[2:])
        elif sys.argv[1] == "debug":
            return handle_debug_command(["debug"] + sys.argv[2:])
        elif sys.argv[1] == "sync":
            return handle_sync_command(["sync"] + sys.argv[2:])
        elif sys.argv[1] == "tell-cursor":
            return handle_tell_cursor_command(["tell-cursor"] + sys.argv[2:])
        elif sys.argv[1] == "push":
            return handle_push_command(["push"] + sys.argv[2:])
        elif sys.argv[1] == "pull":
            return handle_pull_command(["pull"] + sys.argv[2:])
        elif sys.argv[1] == "--":
            # Join the rest as a properly quoted command string
            import shlex

            return submit_shell_command(shlex.join(sys.argv[2:]))
        elif sys.argv[1] == "-0" and len(sys.argv) >= 4:
            if sys.argv[2] == "lab":
                return handle_lab_command(["-0", "lab"] + sys.argv[3:])
            elif sys.argv[2] == "code":
                return handle_code_command(["-0", "code"] + sys.argv[3:])
            elif sys.argv[2] == "debug":
                return handle_debug_command(["-0", "debug"] + sys.argv[3:])
            else:
                print_help()
                return 1
        else:
            # More than 2 arguments - for now, just show help
            print_help()
            return 1

    try:
        ensure_ray_initialized(ray_address=ray_address, working_dir=working_dir)
    except Exception as e:
        print(f"Error initializing Ray: {e}", file=sys.stderr)
        return 1

    # Set up fallback signal handlers
    signal.signal(signal.SIGINT, signal_handler)
    signal.signal(signal.SIGTERM, signal_handler)

    terminal = RaySSHTerminal(
        node_arg,
        ray_address=ray_address,
        working_dir=working_dir,
        resource_constraints=require_constraints,
    )

    # Persist last session info for cluster node connections
    # If this changes, terminal will overwrite it
    try:
        if node_arg and not ray_address:
            os.makedirs(os.path.expanduser("~/.rayssh"), exist_ok=True)
            last_path = os.path.expanduser("~/.rayssh/last_session.json")
            payload = {
                "node_ip": node_arg,
                "timestamp": datetime.now().isoformat(timespec="seconds"),
            }
            with open(last_path, "w", encoding="utf-8") as f:
                json.dump(payload, f)
    except Exception:
        pass

    try:
        asyncio.run(terminal.run())
    except KeyboardInterrupt:
        print("\nSession interrupted by user.")
    except Exception as e:
        print(f"Fatal error: {e}")
        sys.exit(1)
    finally:
        print("👋 Goodbye!")


def print_help():
    """Print help information."""
    help_text = """
RaySSH: Ray-native terminal tool

Usage:
    rayssh                          # Randomly connect to a worker node at remote HOME
    rayssh <ip|node_id|-index>      # Connect to specific node
    rayssh <dir>                    # Remote mode with directory upload (requires RAY_ADDRESS)
    rayssh [-q] <file>              # Submit file as Ray job, -q for no-wait
    rayssh -l                       # Interactive node selection
    rayssh --ls                     # Print nodes table
    rayssh [lab|code|debug] [path]  # Launch Jupyter Lab / code-server / debug code-server on remote
    rayssh code <job_or_submission_id>  # Launch code-server inheriting that job's runtime_env
    rayssh sync <dir> [node]        # Upload directory and start terminal with file sync
    rayssh tell-cursor [dir]        # Create .cursorrules file for debugging guidance
    rayssh -- <command>             # Submit shell command as job

Options:
    -h, --help                      # Show help
    -l, --list, --show              # Interactive node selection
    --ls                            # Print nodes table
    -q                              # Quick mode (no-wait for jobs)
    -- <command>                    # Submit shell command as job

Examples:
    rayssh                          # Random worker node
    rayssh 192.168.1.100            # Connect by IP
    rayssh -1                       # Connect to first worker
    rayssh -l                       # Interactive node selection
    rayssh --ls                     # Show nodes table
    rayssh ./myproject              # Upload and work in directory (remote mode)
    rayssh [-q] script.py           # Submit Python job and wait. "-q" for no-wait.
    rayssh lab                      # Launch Jupyter Lab on worker node
    rayssh code ./src               # Launch code-server with uploaded directory
    rayssh code job_20240914_abc    # Launch code-server using job's runtime_env working_dir
    rayssh debug ./src              # Launch debug-enabled code-server with Ray debugging
    rayssh sync ./myproject         # Upload directory and start terminal with live file sync
    rayssh tell-cursor              # Create .cursorrules in current directory
    rayssh tell-cursor ./myproject  # Create .cursorrules in specific directory
    rayssh -- nvidia-smi            # Submit shell command as job and tail logs
    n_gpus=8 rayssh train.py        # GPUs to request for job (through --entrypoint-num-gpus)

Environment Variables:
    RAY_ADDRESS=ray://host:port     # Enable remote mode

🖥️  Terminal features: Real-time shell via WebSockets, graceful shutdown
🌐 Remote mode: Upload local directories, work on remote clusters  
🚀 Job submission: Python/Bash files, with working dir upload.
🔬 Lab features: Jupyter Lab on Ray nodes with optional working dir upload
💻 Code features: VS Code server on Ray nodes with working dir upload
🐛 Debug features: Ray-enabled VS Code server with debugging extensions and environment
 """
    print(help_text.strip())


if __name__ == "__main__":
    main()