From 2d45544da02b5ff4ba6342970d37fbf0f2f4dd1c Mon Sep 17 00:00:00 2001 From: Sami Khan Date: Tue, 13 Jan 2026 02:46:25 +0500 Subject: [PATCH] use rsh server instead of ssh --- pyproject.toml | 1 + src/exo/main.py | 3 + src/exo/rsh/__init__.py | 13 +++ src/exo/rsh/client.py | 91 ++++++++++++++++ src/exo/rsh/server.py | 149 ++++++++++++++++++++++++++ src/exo/worker/runner/flash_runner.py | 18 +++- 6 files changed, 271 insertions(+), 4 deletions(-) create mode 100644 src/exo/rsh/__init__.py create mode 100644 src/exo/rsh/client.py create mode 100644 src/exo/rsh/server.py diff --git a/pyproject.toml b/pyproject.toml index e51c9ec2..c3d7fd6d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -41,6 +41,7 @@ dependencies = [ exo-master = "exo.master.main:main" exo-worker = "exo.worker.main:main" exo = "exo.main:main" +exo-rsh = "exo.rsh.client:main" # dependencies only required for development [dependency-groups] diff --git a/src/exo/main.py b/src/exo/main.py index 43d69aa5..66ed22fc 100644 --- a/src/exo/main.py +++ b/src/exo/main.py @@ -13,6 +13,7 @@ from pydantic import PositiveInt import exo.routing.topics as topics from exo.master.api import API # TODO: should API be in master? from exo.master.main import Master +from exo.rsh.server import run_rsh_server, RSH_PORT from exo.routing.router import Router, get_node_id_keypair from exo.shared.constants import EXO_LOG from exo.shared.election import Election, ElectionResult @@ -112,6 +113,8 @@ class Node: if self.api: tg.start_soon(self.api.run) tg.start_soon(self._elect_loop) + # Start RSH server for remote execution (used by MPI) + tg.start_soon(run_rsh_server, RSH_PORT) def shutdown(self): # if this is our second call to shutdown, just sys.exit diff --git a/src/exo/rsh/__init__.py b/src/exo/rsh/__init__.py new file mode 100644 index 00000000..43ada8b9 --- /dev/null +++ b/src/exo/rsh/__init__.py @@ -0,0 +1,13 @@ +"""Exo RSH - Remote Shell for MPI without SSH. + +This module provides a remote execution mechanism that allows mpirun to spawn +processes on remote nodes without requiring SSH setup. It works by: + +1. Each Exo node runs a small HTTP server (RSH server) on port 52416 +2. The exo-rsh script acts as a drop-in replacement for ssh +3. When mpirun calls "exo-rsh hostname command", it HTTP POSTs to the target +4. The target executes the command and streams output back + +Usage: + mpirun --mca plm_rsh_agent exo-rsh -np 4 --hostfile hosts.txt ./program +""" diff --git a/src/exo/rsh/client.py b/src/exo/rsh/client.py new file mode 100644 index 00000000..7cdb3a4d --- /dev/null +++ b/src/exo/rsh/client.py @@ -0,0 +1,91 @@ +#!/usr/bin/env python3 +"""exo-rsh - Remote shell client for MPI. + +This script is called by mpirun as a replacement for ssh. +Usage: exo-rsh [ssh-options...] hostname command [args...] + +It connects to the target node's RSH server (port 52416) and executes the command. +""" + +import socket +import sys +from urllib.request import urlopen, Request +from urllib.error import URLError +import json + +RSH_PORT = 52416 + + +def resolve_hostname(hostname: str) -> str: + """Resolve hostname to IP address.""" + try: + return socket.gethostbyname(hostname) + except socket.gaierror: + # If resolution fails, try using the hostname directly + return hostname + + +def main(): + # Parse arguments - mpirun calls us like: exo-rsh [options] hostname command [args...] + # SSH options we might see: -x (disable X11), -o options, etc. + args = sys.argv[1:] + + # Skip SSH-style options + hostname = None + command_start = 0 + + i = 0 + while i < len(args): + arg = args[i] + if arg.startswith("-"): + # Skip option and its value if needed + if arg in ("-o", "-i", "-l", "-p", "-F"): + i += 2 # Skip option and its argument + continue + i += 1 + continue + else: + # First non-option is the hostname + hostname = arg + command_start = i + 1 + break + i += 1 + + if hostname is None or command_start >= len(args): + print("Usage: exo-rsh [options] hostname command [args...]", file=sys.stderr) + sys.exit(1) + + command = args[command_start:] + + # Resolve hostname to IP + ip = resolve_hostname(hostname) + + # Make request to RSH server + url = f"http://{ip}:{RSH_PORT}/execute" + data = json.dumps({"command": command}).encode("utf-8") + + try: + req = Request(url, data=data, headers={"Content-Type": "application/json"}) + with urlopen(req, timeout=300) as response: + result = json.loads(response.read().decode("utf-8")) + + # Output stdout/stderr + if result.get("stdout"): + sys.stdout.write(result["stdout"]) + sys.stdout.flush() + if result.get("stderr"): + sys.stderr.write(result["stderr"]) + sys.stderr.flush() + + sys.exit(result.get("exit_code", 0)) + + except URLError as e: + print(f"exo-rsh: Failed to connect to {hostname}:{RSH_PORT}: {e}", file=sys.stderr) + sys.exit(255) + except Exception as e: + print(f"exo-rsh: Error: {e}", file=sys.stderr) + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/src/exo/rsh/server.py b/src/exo/rsh/server.py new file mode 100644 index 00000000..cc2b6d5d --- /dev/null +++ b/src/exo/rsh/server.py @@ -0,0 +1,149 @@ +"""RSH Server - runs on each Exo node to accept remote execution requests.""" + +import asyncio +import subprocess +from typing import Optional + +from fastapi import FastAPI, HTTPException +from fastapi.responses import StreamingResponse +from hypercorn.asyncio import serve +from hypercorn.config import Config +from loguru import logger +from pydantic import BaseModel + +RSH_PORT = 52416 + + +class ExecuteRequest(BaseModel): + """Request to execute a command.""" + command: list[str] + cwd: Optional[str] = None + env: Optional[dict[str, str]] = None + + +class ExecuteResponse(BaseModel): + """Response from command execution.""" + exit_code: int + stdout: str + stderr: str + + +def create_rsh_app() -> FastAPI: + """Create the RSH FastAPI application.""" + app = FastAPI(title="Exo RSH Server") + + @app.get("/health") + async def health(): + """Health check endpoint.""" + return {"status": "ok"} + + @app.post("/execute") + async def execute(request: ExecuteRequest) -> ExecuteResponse: + """Execute a command and return the result.""" + cmd_str = " ".join(request.command) + logger.info(f"RSH executing: {cmd_str}") + + try: + # Build environment + import os + env = os.environ.copy() + if request.env: + env.update(request.env) + + # Check if command contains shell metacharacters (semicolons, pipes, etc.) + # If so, run through shell. mpirun sends complex commands like: + # "VAR=value;export VAR;/path/to/prted --args" + needs_shell = any(c in cmd_str for c in ";|&$`") + + if needs_shell: + # Run through shell + process = await asyncio.create_subprocess_shell( + cmd_str, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + cwd=request.cwd, + env=env, + ) + else: + # Execute directly + process = await asyncio.create_subprocess_exec( + *request.command, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + cwd=request.cwd, + env=env, + ) + + stdout, stderr = await process.communicate() + exit_code = process.returncode or 0 + + logger.info(f"RSH command completed with exit code {exit_code}") + + return ExecuteResponse( + exit_code=exit_code, + stdout=stdout.decode("utf-8", errors="replace"), + stderr=stderr.decode("utf-8", errors="replace"), + ) + + except FileNotFoundError as e: + logger.error(f"RSH command not found: {e}") + return ExecuteResponse( + exit_code=127, + stdout="", + stderr=f"Command not found: {request.command[0]}", + ) + except Exception as e: + logger.error(f"RSH execution error: {e}") + return ExecuteResponse( + exit_code=1, + stdout="", + stderr=str(e), + ) + + @app.post("/execute_streaming") + async def execute_streaming(request: ExecuteRequest): + """Execute a command and stream the output.""" + logger.info(f"RSH streaming execute: {' '.join(request.command)}") + + async def stream_output(): + try: + env = None + if request.env: + import os + env = os.environ.copy() + env.update(request.env) + + process = await asyncio.create_subprocess_exec( + *request.command, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.STDOUT, + cwd=request.cwd, + env=env, + ) + + if process.stdout: + async for line in process.stdout: + yield line + + await process.wait() + + except Exception as e: + yield f"Error: {e}\n".encode() + + return StreamingResponse( + stream_output(), + media_type="application/octet-stream", + ) + + return app + + +async def run_rsh_server(port: int = RSH_PORT): + """Run the RSH server.""" + app = create_rsh_app() + config = Config() + config.bind = [f"0.0.0.0:{port}"] + config.accesslog = None # Disable access logs for cleaner output + + logger.info(f"Starting RSH server on port {port}") + await serve(app, config) # type: ignore diff --git a/src/exo/worker/runner/flash_runner.py b/src/exo/worker/runner/flash_runner.py index 59c44d93..b876192a 100644 --- a/src/exo/worker/runner/flash_runner.py +++ b/src/exo/worker/runner/flash_runner.py @@ -3,7 +3,8 @@ Exo-native distributed MPI: - Exo handles node discovery and coordination - Coordinator generates hostfile from Exo topology -- mpirun uses SSH (keys already set up) to spawn on remote nodes +- mpirun uses exo-rsh (no SSH required) to spawn on remote nodes +- Each Exo node runs an RSH server on port 52416 for remote execution - Workers just report ready and wait """ @@ -42,6 +43,11 @@ from exo.worker.runner.bootstrap import logger # Find mpirun in PATH, fallback to common locations MPIRUN_PATH = shutil.which("mpirun") or "/opt/homebrew/bin/mpirun" +# exo-rsh is installed as console script by exo package +EXO_RSH_PATH = shutil.which("exo-rsh") +if not EXO_RSH_PATH: + raise RuntimeError("exo-rsh not found in PATH - this should be installed with exo") + def get_my_rank(instance: FLASHInstance, my_node_id: str) -> int: """Determine this node's rank based on position in hosts_by_node.""" @@ -186,9 +192,13 @@ def main( "--mca", "btl_tcp_if_include", iface, "--mca", "oob_tcp_if_include", iface, "--mca", "plm_rsh_no_tree_spawn", "1", - instance.flash_executable_path, ] + # Use exo-rsh for remote execution (no SSH needed) + cmd.extend(["--mca", "plm_rsh_agent", EXO_RSH_PATH]) + + cmd.append(instance.flash_executable_path) + logger.info(f"FLASH distributed launch: {' '.join(cmd)}") process = subprocess.Popen( @@ -207,8 +217,8 @@ def main( logger.info(f"FLASH running on {world_size} nodes with {instance.total_ranks} ranks") else: - # Worker: mpirun on coordinator will SSH to spawn processes here - logger.info(f"Worker {my_rank}: Ready for mpirun to spawn processes") + # Worker: mpirun on coordinator will use exo-rsh to spawn processes here + logger.info(f"Worker {my_rank}: Ready for mpirun to spawn processes via exo-rsh") current_status = RunnerRunning() except Exception as e: