From 36413b6760bd51f93df148f8ecf9be8df05b1ed7 Mon Sep 17 00:00:00 2001 From: Eugene Volen Date: Fri, 3 Jul 2026 11:36:08 +0000 Subject: [PATCH 1/3] Add `rocm remote serve` (Phase 1: foreground) Bring up a managed model server on a remote GPU host with one command: connect over SSH, ensure the CLI + ROCm are present (auto-push the CLI if missing, detect-and-fail if ROCm is absent), start a managed `rocm serve` on the remote, open a loopback port-forward, and print an OpenAI-compatible base URL. The command holds the tunnel open until interrupted; the remote server is managed and persists. The new remote module reuses the remote's existing serve/services machinery rather than reimplementing engine/runtime/device selection. Transport is a trait (SshTransport shells out to system ssh/scp) so a Tailscale-native transport can slot in later. Adds `rocm services list --json` as the machine-readable listing the remote orchestration reads back. Also registers `remote` in the natural-language planner's structured-command allowlist so `rocm remote ...` routes to clap instead of the planner. Phase 2 (session registry, detached tunnel keeper, remote list/status/stop with reconciliation) is not included. Signed-off-by: Eugene Volen --- apps/rocm/src/main.rs | 93 +++++++++- apps/rocm/src/remote/bootstrap.rs | 233 +++++++++++++++++++++++++ apps/rocm/src/remote/mod.rs | 259 ++++++++++++++++++++++++++++ apps/rocm/src/remote/transport.rs | 278 ++++++++++++++++++++++++++++++ 4 files changed, 860 insertions(+), 3 deletions(-) create mode 100644 apps/rocm/src/remote/bootstrap.rs create mode 100644 apps/rocm/src/remote/mod.rs create mode 100644 apps/rocm/src/remote/transport.rs diff --git a/apps/rocm/src/main.rs b/apps/rocm/src/main.rs index d427c0a4..ebd65a69 100644 --- a/apps/rocm/src/main.rs +++ b/apps/rocm/src/main.rs @@ -9,6 +9,7 @@ mod dash; mod dash_seam; mod provider_keys; mod providers; +mod remote; mod serve_summary; mod therock; mod uninstall; @@ -331,6 +332,11 @@ rocm serve qwen2.5-7b-instruct --verbose --device gpu_required")] #[command(subcommand)] command: Option, }, + /// Serve a model on a remote GPU host and forward it to this machine. + Remote { + #[command(subcommand)] + command: remote::RemoteCommand, + }, /// Manage optional background checks and review requests. Automations { #[command(subcommand)] @@ -602,6 +608,9 @@ enum ServicesCommand { /// Include failed, stopped, and old service records. #[arg(short, long)] all: bool, + /// Emit the machine-readable service records as JSON. + #[arg(long)] + json: bool, }, /// Show logs for a local model server. Logs { @@ -1505,6 +1514,7 @@ fn dispatch(cli: Cli) -> Result<()> { }), Some(Command::Comfyui { command }) => comfyui(command), Some(Command::Services { command }) => services(command), + Some(Command::Remote { command }) => remote::remote(command), Some(Command::Automations { command }) => automations(command), Some(Command::Config { command }) => config(command), Some(Command::Logs { @@ -4444,9 +4454,20 @@ fn run_foreground_service( fn services(command: Option) -> Result<()> { let paths = AppPaths::discover()?; - match command.unwrap_or(ServicesCommand::List { all: false }) { - ServicesCommand::List { all } => { - print!("{}", render_services_text(&paths, all)?); + match command.unwrap_or(ServicesCommand::List { + all: false, + json: false, + }) { + ServicesCommand::List { all, json } => { + if json { + let records = load_managed_services(&paths)? + .into_iter() + .filter(|record| all || managed_service_is_live(record)) + .collect::>(); + println!("{}", serde_json::to_string_pretty(&records)?); + } else { + print!("{}", render_services_text(&paths, all)?); + } Ok(()) } ServicesCommand::Logs { service_id } => { @@ -15446,6 +15467,7 @@ fn treat_as_natural_language(args: &[String]) -> bool { "comfyui", "comfy", "services", + "remote", "automations", "config", "logs", @@ -16246,6 +16268,21 @@ mod tests { assert!(!should_treat_as_freeform(&structured)); } + #[test] + fn remote_subcommand_routes_to_clap_not_the_planner() { + // `rocm remote serve …` must reach the structured command, not be + // swallowed by the natural-language planner. Guards the STRUCTURED + // allowlist against dropping `remote`. + let invocation = parse_freeform_invocation(&[ + "remote".to_owned(), + "serve".to_owned(), + "gpu-box".to_owned(), + "qwen2.5-7b-instruct".to_owned(), + ]); + assert!(!treat_as_natural_language(&invocation.request_args)); + assert!(!should_treat_as_freeform(&invocation)); + } + #[test] fn freeform_invocation_rejects_unquoted_structured_command_names_after_yes() { let invalid_install = parse_freeform_invocation(&[ @@ -19063,6 +19100,56 @@ install therock"; Ok(()) } + #[test] + fn services_list_json_round_trips_records() -> Result<()> { + // `rocm services list --json` is the machine-readable source of truth + // that `rocm remote` reconciliation reads over SSH. Guard that the + // emitted JSON deserializes back into ManagedServiceRecord with the + // fields consumers depend on (service_id, endpoint_url, port, status). + let (root, paths) = test_paths("services-list-json"); + paths.ensure()?; + let current_pid = std::process::id(); + for (service_id, status, port) in [ + ("svc-a", "ready", 11440_u16), + ("svc-b", "failed", 11441_u16), + ] { + let mut record = ManagedServiceRecord::new( + &paths, + service_id, + "pytorch", + "qwen", + "Qwen/Qwen3.5", + "127.0.0.1", + port, + "managed", + current_pid, + Some("therock-release".to_owned()), + None, + Some("gpu_required".to_owned()), + ); + record.status = status.to_owned(); + record.write()?; + } + + // Mirror the handler's `--json --all` path: serialize every record. + let records = load_managed_services(&paths)?; + let json = serde_json::to_string_pretty(&records)?; + let parsed: Vec = serde_json::from_str(&json)?; + let _ = fs::remove_dir_all(root); + + assert_eq!(parsed.len(), 2); + let a = parsed + .iter() + .find(|r| r.service_id == "svc-a") + .expect("svc-a present in JSON"); + assert_eq!(a.port, 11440); + assert_eq!(a.engine, "pytorch"); + assert!(!a.status.is_empty()); + assert!(a.endpoint_url.contains("11440")); + assert!(parsed.iter().any(|r| r.service_id == "svc-b")); + Ok(()) + } + #[test] fn render_services_text_demotes_stale_ready_record() -> Result<()> { let (root, paths) = test_paths("services-stale-ready"); diff --git a/apps/rocm/src/remote/bootstrap.rs b/apps/rocm/src/remote/bootstrap.rs new file mode 100644 index 00000000..7767a67d --- /dev/null +++ b/apps/rocm/src/remote/bootstrap.rs @@ -0,0 +1,233 @@ +//! Ensuring a remote host is ready to serve: probe for the `rocm` CLI and ROCm, +//! and push the CLI if it is missing. +//! +//! Scope decisions (see `plans/rocm-remote-serve.md`): +//! - **CLI missing → auto-push.** The whole `rocm` runtime, including the managed +//! daemon, is a single self-contained binary (`rocm daemon` is the supervisor; +//! there is no separate `rocmd` to install), so making the remote ready is just +//! copying one file. Because the release repo is currently private (no public +//! `curl` URL), v1 pushes the *local* `rocm` binary over `scp`. This works when +//! the local machine is itself linux-amd64 (the demo case); other hosts get a +//! clear error pointing at the manual install. +//! - **ROCm missing → detect and fail.** We never attempt to install the +//! multi-GB ROCm SDK remotely; we report it clearly and stop. + +use anyhow::{Context, Result, bail}; + +use super::transport::Transport; + +/// Where a freshly-pushed CLI lands on the remote. Chosen to match the +/// documented manual install location (`~/.local/bin`). +const REMOTE_CLI_PATH: &str = "~/.local/bin/rocm"; + +/// What a probe found on the remote host. +#[derive(Debug, Clone)] +pub struct RemoteReadiness { + pub cli_present: bool, + pub cli_version: Option, + pub rocm_present: bool, +} + +/// Probe the remote for the `rocm` CLI and a ROCm installation. Neither absence +/// is an error here — that judgement belongs to [`ensure_ready`]. +pub fn probe(transport: &dyn Transport) -> Result { + let cli = transport.exec("rocm --version")?; + let (cli_present, cli_version) = if cli.success { + let version = cli.stdout.lines().next().map(|l| l.trim().to_owned()); + (true, version.filter(|v| !v.is_empty())) + } else { + (false, None) + }; + + // ROCm detection independent of our CLI: any of the usual markers is enough. + let rocm = transport.exec( + "command -v rocminfo >/dev/null 2>&1 || command -v amd-smi >/dev/null 2>&1 || test -d /opt/rocm", + )?; + + Ok(RemoteReadiness { + cli_present, + cli_version, + rocm_present: rocm.success, + }) +} + +/// How to invoke `rocm` on the remote once we know it is present — either the +/// CLI already on `PATH`, or the explicit path we just installed to. +#[derive(Debug, Clone)] +pub struct RemoteCli { + invocation: String, +} + +impl RemoteCli { + /// The command prefix to use when building remote `rocm …` invocations. + pub fn invocation(&self) -> &str { + &self.invocation + } +} + +/// Make the remote ready to serve, pushing the CLI if needed. Returns how to +/// invoke `rocm` on the remote. Fails (without attempting any install) if ROCm +/// is absent. +pub fn ensure_ready(transport: &dyn Transport) -> Result { + let readiness = probe(transport)?; + + if !readiness.rocm_present { + bail!( + "ROCm was not detected on the remote host (no rocminfo/amd-smi and no /opt/rocm).\n\ + Install ROCm on the remote first, then re-run. `rocm remote serve` does not \ + install the ROCm SDK for you." + ); + } + + if readiness.cli_present { + match readiness.cli_version.as_deref() { + Some(version) => println!(" remote rocm: {version}"), + None => println!(" remote rocm: present"), + } + return Ok(RemoteCli { + invocation: "rocm".to_owned(), + }); + } + + // CLI absent: push the local binary. + push_local_cli(transport)?; + Ok(RemoteCli { + invocation: REMOTE_CLI_PATH.to_owned(), + }) +} + +/// Copy the locally-running `rocm` binary to the remote and verify it runs. +fn push_local_cli(transport: &dyn Transport) -> Result<()> { + let local_exe = std::env::current_exe() + .context("failed to resolve the local rocm executable to push to the remote")?; + + transport + .run("mkdir -p \"$HOME/.local/bin\"") + .context("failed to create ~/.local/bin on the remote")?; + transport.push_file(&local_exe, REMOTE_CLI_PATH).context( + "failed to copy the local rocm binary to the remote. If your local machine is not \ + linux-amd64, install rocm on the remote manually and re-run.", + )?; + transport + .run("chmod +x \"$HOME/.local/bin/rocm\"") + .context("failed to mark the pushed rocm binary executable")?; + + // Verify the pushed binary actually runs on the remote (catches arch/libc + // mismatch early, before we try to serve with it). + let verify = transport.exec("\"$HOME/.local/bin/rocm\" --version")?; + if !verify.success { + bail!( + "pushed rocm binary does not run on the remote (exit {}): {}\n\ + This usually means an architecture or libc mismatch. Install rocm on the remote \ + manually and re-run.", + verify + .code + .map_or_else(|| "signal".to_owned(), |c| c.to_string()), + verify.stderr.trim(), + ); + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::remote::transport::RemoteOutcome; + use std::cell::RefCell; + + /// A scripted transport: each `exec` pops the next queued outcome and records + /// the command it was asked to run. `run` reuses the trait default over + /// `exec`, so queued outcomes must account for `run` calls too. + struct ScriptedTransport { + outcomes: RefCell>, + commands: RefCell>, + pushed: RefCell>, + } + + impl ScriptedTransport { + fn new(outcomes: Vec) -> Self { + Self { + outcomes: RefCell::new(outcomes), + commands: RefCell::new(Vec::new()), + pushed: RefCell::new(Vec::new()), + } + } + } + + fn ok(stdout: &str) -> RemoteOutcome { + RemoteOutcome { + success: true, + code: Some(0), + stdout: stdout.to_owned(), + stderr: String::new(), + } + } + + fn fail() -> RemoteOutcome { + RemoteOutcome { + success: false, + code: Some(1), + stdout: String::new(), + stderr: "not found".to_owned(), + } + } + + impl Transport for ScriptedTransport { + fn exec(&self, command: &str) -> Result { + self.commands.borrow_mut().push(command.to_owned()); + Ok(self + .outcomes + .borrow_mut() + .drain(..1) + .next() + .unwrap_or_else(fail)) + } + fn push_file(&self, _local: &std::path::Path, remote_path: &str) -> Result<()> { + self.pushed.borrow_mut().push(remote_path.to_owned()); + Ok(()) + } + fn forward( + &self, + _local_port: u16, + _remote_host: &str, + _remote_port: u16, + ) -> Result { + unreachable!("bootstrap tests do not forward") + } + } + + #[test] + fn probe_reports_present_cli_and_rocm() { + let t = ScriptedTransport::new(vec![ok("rocm 0.3.0"), ok("")]); + let r = probe(&t).unwrap(); + assert!(r.cli_present); + assert_eq!(r.cli_version.as_deref(), Some("rocm 0.3.0")); + assert!(r.rocm_present); + } + + #[test] + fn ensure_ready_bails_when_rocm_absent() { + // cli --version fails, rocm detection fails. + let t = ScriptedTransport::new(vec![fail(), fail()]); + let err = ensure_ready(&t).unwrap_err().to_string(); + assert!(err.contains("ROCm was not detected"), "got: {err}"); + } + + #[test] + fn ensure_ready_returns_path_invocation_after_push() { + // probe: cli absent (fail), rocm present (ok); + // push_local_cli: mkdir ok, chmod ok, verify ok. + let t = ScriptedTransport::new(vec![fail(), ok(""), ok(""), ok(""), ok("rocm 0.3.0")]); + let cli = ensure_ready(&t).unwrap(); + assert_eq!(cli.invocation(), REMOTE_CLI_PATH); + assert!(t.pushed.borrow().iter().any(|p| p == REMOTE_CLI_PATH)); + } + + #[test] + fn ensure_ready_uses_path_rocm_when_present() { + let t = ScriptedTransport::new(vec![ok("rocm 0.3.0"), ok("")]); + let cli = ensure_ready(&t).unwrap(); + assert_eq!(cli.invocation(), "rocm"); + assert!(t.pushed.borrow().is_empty()); + } +} diff --git a/apps/rocm/src/remote/mod.rs b/apps/rocm/src/remote/mod.rs new file mode 100644 index 00000000..ea40d5aa --- /dev/null +++ b/apps/rocm/src/remote/mod.rs @@ -0,0 +1,259 @@ +//! `rocm remote` — bring up a managed model server on a remote GPU host and make +//! it callable from this machine. +//! +//! Phase 1 (this module) implements the foreground happy path: connect over SSH, +//! ensure the CLI + ROCm are present (pushing the CLI if missing), start a +//! *managed* `rocm serve` on the remote (reusing all of the remote's engine / +//! runtime / device selection), open a loopback port-forward, and print an +//! OpenAI-compatible base URL. The command holds the tunnel open until +//! interrupted; the remote server is managed and persists. +//! +//! Phase 2 will add a local session registry and `rocm remote list/status/stop` +//! that reconcile against the remote's `rocm services list --json`. + +mod bootstrap; +mod transport; + +use std::fmt::Write as _; + +use anyhow::{Context, Result}; +use clap::Subcommand; +use rocm_core::ManagedServiceRecord; + +use transport::{SshTransport, Transport}; + +#[derive(Subcommand, Debug)] +pub(crate) enum RemoteCommand { + /// Start a managed model server on a remote GPU host and forward it locally. + #[command(after_help = "EXAMPLES:\n \ +rocm remote serve gpu-box qwen2.5-7b-instruct\n \ +rocm remote serve user@10.0.0.5 qwen2.5-7b-instruct --engine llama.cpp\n \ +rocm remote serve gpu-box qwen2.5-7b-instruct --local-port 8000")] + Serve { + /// SSH destination: `user@host` or a `~/.ssh/config` host alias. + host: String, + /// Model name, alias, or model file path as seen on the remote host. + model: String, + /// SSH port on the remote host (defaults to the ssh/config default). + #[arg(long)] + ssh_port: Option, + /// Port the model server listens on, on the remote side. + #[arg(long, default_value_t = rocm_core::DEFAULT_LOCAL_PORT)] + remote_port: u16, + /// Local loopback port to forward to (defaults to the remote port). + #[arg(long)] + local_port: Option, + /// Engine to use on the remote [e.g. lemonade, llama.cpp, vllm]. + #[arg(long)] + engine: Option, + /// Device policy [possible values: gpu_required, gpu_preferred, cpu_only]. + #[arg(long)] + device: Option, + /// GPU device to serve on: `auto` or a single index like `1`. + #[arg(long, value_name = "INDEX|auto")] + gpu: Option, + }, +} + +pub(crate) fn remote(command: RemoteCommand) -> Result<()> { + match command { + RemoteCommand::Serve { + host, + model, + ssh_port, + remote_port, + local_port, + engine, + device, + gpu, + } => serve(ServeArgs { + host, + model, + ssh_port, + remote_port, + local_port, + engine, + device, + gpu, + }), + } +} + +struct ServeArgs { + host: String, + model: String, + ssh_port: Option, + remote_port: u16, + local_port: Option, + engine: Option, + device: Option, + gpu: Option, +} + +fn serve(args: ServeArgs) -> Result<()> { + let transport = SshTransport::new(args.host.clone(), args.ssh_port); + let local_port = args.local_port.unwrap_or(args.remote_port); + + // 1. Reachability — fail fast with a clear message before doing anything else. + println!("remote serve"); + println!(" host: {}", args.host); + transport + .run("true") + .with_context(|| format!("cannot reach {} over SSH", args.host))?; + + // 2. Ensure the remote is ready (push the CLI if missing; require ROCm). + let cli = bootstrap::ensure_ready(&transport)?; + let rocm = cli.invocation(); + + // 3. Start a managed server on the remote, bound to remote loopback. + let serve_cmd = build_remote_serve_command(rocm, &args); + println!( + " starting managed server (remote port {}) ...", + args.remote_port + ); + transport + .run(&serve_cmd) + .context("failed to start the managed server on the remote host")?; + + // 4. Discover the server we just started via the machine-readable listing, + // matching on the port we asked it to bind. + let record = discover_started_service(&transport, rocm, args.remote_port)?; + println!(" remote model: {}", record.canonical_model_id); + println!(" remote status: {}", record.status); + + // 5. Open the loopback port-forward. + let mut guard = transport + .forward(local_port, "127.0.0.1", args.remote_port) + .context("failed to open the SSH port-forward")?; + + // 6. Tell the user how to call it. + println!(); + println!("Ready. OpenAI-compatible base URL:"); + println!(" http://127.0.0.1:{local_port}/v1"); + println!(); + println!( + " try: rocm chat --provider local --model {} --prompt \"hello\"", + record.canonical_model_id + ); + println!(" (set OPENAI_BASE_URL=http://127.0.0.1:{local_port}/v1 for other OpenAI clients)"); + println!(); + println!("The remote server is managed and will keep running after you disconnect."); + println!("Press Ctrl-C to close the local forward."); + + // 7. Hold the tunnel open until interrupted; the guard tears it down on drop. + guard.wait()?; + Ok(()) +} + +/// Build the remote `rocm serve … --managed` invocation, quoting user-supplied +/// values so they survive the remote shell. +fn build_remote_serve_command(rocm: &str, args: &ServeArgs) -> String { + let mut cmd = format!( + "{rocm} serve {} --managed --host 127.0.0.1 --port {}", + shell_quote(&args.model), + args.remote_port, + ); + if let Some(engine) = args.engine.as_deref() { + let _ = write!(cmd, " --engine {}", shell_quote(engine)); + } + if let Some(device) = args.device.as_deref() { + let _ = write!(cmd, " --device {}", shell_quote(device)); + } + if let Some(gpu) = args.gpu.as_deref() { + let _ = write!(cmd, " --gpu {}", shell_quote(gpu)); + } + cmd +} + +/// Read `rocm services list --json` on the remote and pick the record bound to +/// `remote_port` (the one we just started). If several match, the most recently +/// created wins. +fn discover_started_service( + transport: &dyn Transport, + rocm: &str, + remote_port: u16, +) -> Result { + let listing = transport + .run(&format!("{rocm} services list --json")) + .context("failed to list managed services on the remote host")?; + let records: Vec = serde_json::from_str(listing.trim()).context( + "could not parse `rocm services list --json` from the remote — the remote CLI may be too \ + old to support --json", + )?; + records + .into_iter() + .filter(|r| r.port == remote_port) + .max_by_key(|r| r.created_at_unix_ms) + .ok_or_else(|| { + anyhow::anyhow!( + "the managed server did not appear on the remote (no service bound to port {remote_port})" + ) + }) +} + +/// Minimal POSIX single-quote quoting for values interpolated into a remote +/// shell command: wrap in single quotes and escape embedded single quotes. +fn shell_quote(value: &str) -> String { + if !value.is_empty() + && value + .bytes() + .all(|b| b.is_ascii_alphanumeric() || matches!(b, b'.' | b'_' | b'-' | b'/' | b':')) + { + return value.to_owned(); + } + let escaped = value.replace('\'', "'\\''"); + format!("'{escaped}'") +} + +#[cfg(test)] +mod tests { + use super::*; + + fn base_args() -> ServeArgs { + ServeArgs { + host: "gpu-box".to_owned(), + model: "qwen2.5-7b-instruct".to_owned(), + ssh_port: None, + remote_port: 11435, + local_port: None, + engine: None, + device: None, + gpu: None, + } + } + + #[test] + fn remote_serve_command_is_managed_and_loopback() { + let cmd = build_remote_serve_command("rocm", &base_args()); + assert_eq!( + cmd, + "rocm serve qwen2.5-7b-instruct --managed --host 127.0.0.1 --port 11435" + ); + } + + #[test] + fn remote_serve_command_threads_passthrough_flags() { + let mut args = base_args(); + args.engine = Some("llama.cpp".to_owned()); + args.device = Some("gpu_required".to_owned()); + args.gpu = Some("auto".to_owned()); + let cmd = build_remote_serve_command("$HOME/.local/bin/rocm", &args); + assert!(cmd.starts_with("$HOME/.local/bin/rocm serve qwen2.5-7b-instruct --managed")); + assert!(cmd.contains(" --engine llama.cpp")); + assert!(cmd.contains(" --device gpu_required")); + assert!(cmd.contains(" --gpu auto")); + } + + #[test] + fn shell_quote_leaves_simple_values_bare() { + assert_eq!(shell_quote("qwen2.5-7b-instruct"), "qwen2.5-7b-instruct"); + assert_eq!(shell_quote("./models/m.gguf"), "./models/m.gguf"); + } + + #[test] + fn shell_quote_wraps_dangerous_values() { + assert_eq!(shell_quote("a b"), "'a b'"); + assert_eq!(shell_quote("x; rm -rf /"), "'x; rm -rf /'"); + assert_eq!(shell_quote("it's"), "'it'\\''s'"); + } +} diff --git a/apps/rocm/src/remote/transport.rs b/apps/rocm/src/remote/transport.rs new file mode 100644 index 00000000..885846fd --- /dev/null +++ b/apps/rocm/src/remote/transport.rs @@ -0,0 +1,278 @@ +//! Transport abstraction for reaching a remote GPU host. +//! +//! `rocm remote` needs three things from a remote host: run a command and read +//! its output, copy a file over, and open a local port-forward to a service the +//! remote is listening on. The [`Transport`] trait captures exactly those three +//! capabilities so the orchestration in `remote::serve` stays agnostic to *how* +//! we reach the host. v1 ships a single implementation, [`SshTransport`], that +//! shells out to the system `ssh`/`scp`; a future Tailscale-native transport can +//! slot in behind the same trait without touching the callers. +//! +//! Why shell out to `ssh` instead of a Rust SSH crate: it reuses the user's +//! existing keys/agent/`~/.ssh/config` and connection multiplexing for free, +//! keeps this code synchronous (the surrounding command handlers are sync), and +//! makes the port-forward a plain child process that plugs into the existing +//! detach + PID-tracking machinery the managed-service supervisor already uses. + +use std::process::{Child, Command, Stdio}; + +use anyhow::{Context, Result, bail}; + +/// A running local port-forward. Dropping it tears the forward down by killing +/// the underlying `ssh -N -L` child. +#[derive(Debug)] +pub struct ForwardGuard { + child: Child, +} + +impl ForwardGuard { + /// Block until the forward's `ssh` child exits (e.g. the connection drops or + /// the process group receives SIGINT). Used to hold a foreground session + /// open until the user interrupts it. + pub fn wait(&mut self) -> Result<()> { + self.child + .wait() + .context("failed while waiting on the ssh port-forward")?; + Ok(()) + } +} + +impl Drop for ForwardGuard { + fn drop(&mut self) { + // Best-effort teardown for the foreground path (Ctrl-C / normal return). + let _ = self.child.kill(); + let _ = self.child.wait(); + } +} + +/// Captured result of running a command on the remote host. Distinguishes a +/// clean-but-non-zero exit (e.g. `rocm --version` when the CLI is absent) from a +/// transport failure (couldn't reach the host at all) — the bootstrap probes +/// depend on that distinction. +#[derive(Debug, Clone)] +pub struct RemoteOutcome { + pub success: bool, + pub code: Option, + pub stdout: String, + pub stderr: String, +} + +/// Ways to reach a remote host. See the module docs for the design rationale. +pub trait Transport { + /// Run `command` on the remote host and capture its outcome. Returns `Err` + /// only when the command could not be launched / the host is unreachable; a + /// non-zero remote exit is reported via [`RemoteOutcome::success`]. + fn exec(&self, command: &str) -> Result; + + /// Run `command` and return its stdout, failing if it exits non-zero. A + /// convenience over [`exec`](Transport::exec) for commands expected to + /// succeed. + fn run(&self, command: &str) -> Result { + let outcome = self.exec(command)?; + if !outcome.success { + bail!( + "remote command failed (exit {}): {}\n command: {command}", + outcome + .code + .map_or_else(|| "signal".to_owned(), |c| c.to_string()), + outcome.stderr.trim(), + ); + } + Ok(outcome.stdout) + } + + /// Copy a local file to `remote_path` on the host. + fn push_file(&self, local_path: &std::path::Path, remote_path: &str) -> Result<()>; + + /// Open a local port-forward: `127.0.0.1:local_port` on this machine is + /// forwarded to `remote_host:remote_port` as seen from the remote side + /// (typically `127.0.0.1:`). Returns a guard that keeps the + /// forward alive until dropped. + fn forward(&self, local_port: u16, remote_host: &str, remote_port: u16) + -> Result; +} + +/// SSH-backed transport that shells out to the system `ssh`/`scp`. +#[derive(Debug, Clone)] +pub struct SshTransport { + /// SSH destination as accepted by `ssh` (e.g. `user@host`, or a + /// `~/.ssh/config` host alias). + destination: String, + /// Optional explicit port (`ssh -p`); `None` uses the ssh default / config. + port: Option, +} + +impl SshTransport { + pub fn new(destination: impl Into, port: Option) -> Self { + Self { + destination: destination.into(), + port, + } + } + + /// Common `ssh` options applied to every invocation: fail fast on connect, + /// stay non-interactive (no password prompts hanging a pipeline), and reuse + /// a multiplexed master connection so repeated `run` calls are cheap. + fn base_ssh_args(&self) -> Vec { + let mut args = vec![ + "-o".to_owned(), + "BatchMode=yes".to_owned(), + "-o".to_owned(), + "ConnectTimeout=10".to_owned(), + "-o".to_owned(), + "ControlMaster=auto".to_owned(), + "-o".to_owned(), + "ControlPersist=60s".to_owned(), + ]; + if let Some(port) = self.port { + args.push("-p".to_owned()); + args.push(port.to_string()); + } + args + } + + /// Argument vector for running `command` on the remote (excludes the `ssh` + /// program name). Split out for unit testing. + fn exec_argv(&self, command: &str) -> Vec { + let mut args = self.base_ssh_args(); + args.push(self.destination.clone()); + args.push("--".to_owned()); + args.push(command.to_owned()); + args + } + + /// Argument vector for the port-forward child (excludes the `ssh` program + /// name). `-N` = no remote command, `-T` = no pty; the process stays alive + /// holding the forward. Split out for unit testing. + fn forward_argv(&self, local_port: u16, remote_host: &str, remote_port: u16) -> Vec { + let mut args = self.base_ssh_args(); + args.push("-N".to_owned()); + args.push("-T".to_owned()); + args.push("-L".to_owned()); + args.push(format!( + "127.0.0.1:{local_port}:{remote_host}:{remote_port}" + )); + args.push(self.destination.clone()); + args + } + + /// Argument vector for `scp` to copy a local file to the remote (excludes the + /// `scp` program name). Split out for unit testing. + fn scp_argv(&self, local_path: &str, remote_path: &str) -> Vec { + let mut args = vec![ + "-o".to_owned(), + "BatchMode=yes".to_owned(), + "-o".to_owned(), + "ConnectTimeout=10".to_owned(), + ]; + // scp takes the port with a capital -P, unlike ssh's lowercase -p. + if let Some(port) = self.port { + args.push("-P".to_owned()); + args.push(port.to_string()); + } + args.push(local_path.to_owned()); + args.push(format!("{}:{remote_path}", self.destination)); + args + } +} + +impl Transport for SshTransport { + fn exec(&self, command: &str) -> Result { + let output = Command::new("ssh") + .args(self.exec_argv(command)) + .stdin(Stdio::null()) + .output() + .with_context(|| { + format!( + "failed to launch ssh to run a command on {}", + self.destination + ) + })?; + Ok(RemoteOutcome { + success: output.status.success(), + code: output.status.code(), + stdout: String::from_utf8_lossy(&output.stdout).into_owned(), + stderr: String::from_utf8_lossy(&output.stderr).into_owned(), + }) + } + + fn push_file(&self, local_path: &std::path::Path, remote_path: &str) -> Result<()> { + let local = local_path.to_string_lossy(); + let output = Command::new("scp") + .args(self.scp_argv(&local, remote_path)) + .stdin(Stdio::null()) + .output() + .with_context(|| format!("failed to launch scp to {}", self.destination))?; + if !output.status.success() { + let stderr = String::from_utf8_lossy(&output.stderr); + bail!( + "failed to copy {} to {}:{remote_path}: {}", + local, + self.destination, + stderr.trim(), + ); + } + Ok(()) + } + + fn forward( + &self, + local_port: u16, + remote_host: &str, + remote_port: u16, + ) -> Result { + let child = Command::new("ssh") + .args(self.forward_argv(local_port, remote_host, remote_port)) + .stdin(Stdio::null()) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .spawn() + .with_context(|| { + format!("failed to launch ssh port-forward to {}", self.destination) + })?; + Ok(ForwardGuard { child }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn exec_argv_includes_batchmode_and_destination() { + let t = SshTransport::new("user@gpubox", None); + let argv = t.exec_argv("rocm --version"); + assert!(argv.windows(2).any(|w| w == ["-o", "BatchMode=yes"])); + assert_eq!(argv.last().unwrap(), "rocm --version"); + // The command follows a `--` guard so remote args aren't parsed by ssh. + let dashdash = argv.iter().position(|a| a == "--").unwrap(); + assert_eq!(argv[dashdash - 1], "user@gpubox"); + } + + #[test] + fn exec_argv_threads_explicit_port_lowercase() { + let t = SshTransport::new("user@gpubox", Some(2222)); + let argv = t.exec_argv("echo hi"); + assert!(argv.windows(2).any(|w| w == ["-p", "2222"])); + } + + #[test] + fn forward_argv_maps_loopback_local_to_remote() { + let t = SshTransport::new("gpubox", None); + let argv = t.forward_argv(11435, "127.0.0.1", 11435); + assert!(argv.contains(&"-N".to_owned())); + assert!( + argv.windows(2) + .any(|w| { w[0] == "-L" && w[1] == "127.0.0.1:11435:127.0.0.1:11435" }) + ); + assert_eq!(argv.last().unwrap(), "gpubox"); + } + + #[test] + fn scp_argv_uses_uppercase_port_and_remote_colon_path() { + let t = SshTransport::new("user@gpubox", Some(2222)); + let argv = t.scp_argv("/tmp/rocm.tar.gz", "/tmp/rocm.tar.gz"); + assert!(argv.windows(2).any(|w| w == ["-P", "2222"])); + assert_eq!(argv.last().unwrap(), "user@gpubox:/tmp/rocm.tar.gz"); + } +} From e0d6ebcb7bf841741751648ca830675b910680e3 Mon Sep 17 00:00:00 2001 From: Eugene Volen Date: Fri, 3 Jul 2026 13:00:58 +0000 Subject: [PATCH 2/3] Make `rocm remote` sessions persistent and observable (Phase 2) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Remote serving is now async and always-persistent. `serve` starts the managed server on the remote (which loads the model in the background), opens a detached SSH tunnel that survives the command, records the session locally, and returns immediately with a Pending status and the OpenAI base URL — no more foreground hold. New verbs reconcile two independent lifecycles: `status [] [--watch]` shows a table of sessions with remote-server health (polled over SSH via `services list --json`) and local tunnel liveness (PID check) as separate columns; `attach` rebuilds a dropped tunnel without redeploying; `stop` closes the tunnel and stops the remote server. `list` is reserved for configured hosts. The Pending->Healthy progression comes from the remote's own managed-service lifecycle, so there is no local reconciler daemon — status polls on demand. The tunnel keeper is a standalone detached ssh process using ControlMaster=no so killing its PID deterministically closes the forward. Bootstrap now reports progress when pushing the CLI to a remote that lacks it. State lives in a new RemoteSessionRecord (one JSON file per session under remote-sessions/), mirroring ManagedServiceRecord. Signed-off-by: Eugene Volen --- apps/rocm/src/remote/bootstrap.rs | 11 +- apps/rocm/src/remote/mod.rs | 466 ++++++++++++++++++++++++++---- apps/rocm/src/remote/transport.rs | 140 +++++---- crates/rocm-core/src/lib.rs | 127 ++++++++ 4 files changed, 638 insertions(+), 106 deletions(-) diff --git a/apps/rocm/src/remote/bootstrap.rs b/apps/rocm/src/remote/bootstrap.rs index 7767a67d..d51ec996 100644 --- a/apps/rocm/src/remote/bootstrap.rs +++ b/apps/rocm/src/remote/bootstrap.rs @@ -90,7 +90,9 @@ pub fn ensure_ready(transport: &dyn Transport) -> Result { } // CLI absent: push the local binary. + println!(" remote rocm: not found — installing the local rocm binary..."); push_local_cli(transport)?; + println!(" remote rocm: installed to {REMOTE_CLI_PATH}"); Ok(RemoteCli { invocation: REMOTE_CLI_PATH.to_owned(), }) @@ -100,6 +102,11 @@ pub fn ensure_ready(transport: &dyn Transport) -> Result { fn push_local_cli(transport: &dyn Transport) -> Result<()> { let local_exe = std::env::current_exe() .context("failed to resolve the local rocm executable to push to the remote")?; + let size_hint = std::fs::metadata(&local_exe) + .ok() + .map(|m| format!(" (~{} MB)", m.len() / 1_000_000)) + .unwrap_or_default(); + println!(" copying {}{size_hint} ...", local_exe.display()); transport .run("mkdir -p \"$HOME/.local/bin\"") @@ -186,12 +193,12 @@ mod tests { self.pushed.borrow_mut().push(remote_path.to_owned()); Ok(()) } - fn forward( + fn open_detached_forward( &self, _local_port: u16, _remote_host: &str, _remote_port: u16, - ) -> Result { + ) -> Result { unreachable!("bootstrap tests do not forward") } } diff --git a/apps/rocm/src/remote/mod.rs b/apps/rocm/src/remote/mod.rs index ea40d5aa..0bff4076 100644 --- a/apps/rocm/src/remote/mod.rs +++ b/apps/rocm/src/remote/mod.rs @@ -1,27 +1,35 @@ //! `rocm remote` — bring up a managed model server on a remote GPU host and make //! it callable from this machine. //! -//! Phase 1 (this module) implements the foreground happy path: connect over SSH, -//! ensure the CLI + ROCm are present (pushing the CLI if missing), start a -//! *managed* `rocm serve` on the remote (reusing all of the remote's engine / -//! runtime / device selection), open a loopback port-forward, and print an -//! OpenAI-compatible base URL. The command holds the tunnel open until -//! interrupted; the remote server is managed and persists. +//! `serve` connects over SSH, ensures the CLI + ROCm are present (pushing the CLI +//! if missing), starts a *managed* `rocm serve` on the remote (reusing all of the +//! remote's engine/runtime/device selection), opens a **detached** loopback +//! tunnel that survives this command, records the session locally, and returns +//! immediately printing the OpenAI base URL and a `Pending` status. //! -//! Phase 2 will add a local session registry and `rocm remote list/status/stop` -//! that reconcile against the remote's `rocm services list --json`. +//! The `Pending → Healthy` progression comes for free from the remote's own +//! managed-service lifecycle (it loads the model in the background). There is no +//! local reconciler daemon: `status` polls the remote over SSH on demand and +//! checks the local tunnel PID for liveness, surfacing the two lifecycles +//! independently. `attach` rebuilds a dropped tunnel; `stop` tears both down. mod bootstrap; mod transport; use std::fmt::Write as _; +use std::io::Write as _; +use std::thread::sleep; +use std::time::Duration; -use anyhow::{Context, Result}; +use anyhow::{Context, Result, bail}; use clap::Subcommand; -use rocm_core::ManagedServiceRecord; +use rocm_core::{AppPaths, ManagedServiceRecord, RemoteSessionRecord, load_remote_sessions}; use transport::{SshTransport, Transport}; +/// How often `status --watch` refreshes. +const WATCH_INTERVAL: Duration = Duration::from_secs(2); + #[derive(Subcommand, Debug)] pub(crate) enum RemoteCommand { /// Start a managed model server on a remote GPU host and forward it locally. @@ -40,7 +48,7 @@ rocm remote serve gpu-box qwen2.5-7b-instruct --local-port 8000")] /// Port the model server listens on, on the remote side. #[arg(long, default_value_t = rocm_core::DEFAULT_LOCAL_PORT)] remote_port: u16, - /// Local loopback port to forward to (defaults to the remote port). + /// Local loopback port to forward to (defaults to an auto-picked free port). #[arg(long)] local_port: Option, /// Engine to use on the remote [e.g. lemonade, llama.cpp, vllm]. @@ -53,9 +61,32 @@ rocm remote serve gpu-box qwen2.5-7b-instruct --local-port 8000")] #[arg(long, value_name = "INDEX|auto")] gpu: Option, }, + /// Show remote serving sessions and their health. + /// + /// With no id, lists every session; with an id (or an unambiguous host), + /// shows that one. `--watch` keeps the view live. (`rocm remote list` is + /// reserved for listing configured hosts, distinct from active sessions.) + Status { + /// Session id, or a host that resolves to exactly one session. + session: Option, + /// Refresh the view continuously until interrupted. + #[arg(long)] + watch: bool, + }, + /// Rebuild a dropped local tunnel for a session without redeploying. + Attach { + /// Session id, or a host that resolves to exactly one session. + session: String, + }, + /// Stop a session: close the local tunnel and stop the remote server. + Stop { + /// Session id, or a host that resolves to exactly one session. + session: String, + }, } pub(crate) fn remote(command: RemoteCommand) -> Result<()> { + let paths = AppPaths::discover()?; match command { RemoteCommand::Serve { host, @@ -66,16 +97,22 @@ pub(crate) fn remote(command: RemoteCommand) -> Result<()> { engine, device, gpu, - } => serve(ServeArgs { - host, - model, - ssh_port, - remote_port, - local_port, - engine, - device, - gpu, - }), + } => serve( + &paths, + ServeArgs { + host, + model, + ssh_port, + remote_port, + local_port, + engine, + device, + gpu, + }, + ), + RemoteCommand::Status { session, watch } => status(&paths, session.as_deref(), watch), + RemoteCommand::Attach { session } => attach(&paths, &session), + RemoteCommand::Stop { session } => stop(&paths, &session), } } @@ -90,13 +127,11 @@ struct ServeArgs { gpu: Option, } -fn serve(args: ServeArgs) -> Result<()> { +fn serve(paths: &AppPaths, args: ServeArgs) -> Result<()> { let transport = SshTransport::new(args.host.clone(), args.ssh_port); - let local_port = args.local_port.unwrap_or(args.remote_port); // 1. Reachability — fail fast with a clear message before doing anything else. - println!("remote serve"); - println!(" host: {}", args.host); + println!("Connecting to {} over SSH...", args.host); transport .run("true") .with_context(|| format!("cannot reach {} over SSH", args.host))?; @@ -105,43 +140,118 @@ fn serve(args: ServeArgs) -> Result<()> { let cli = bootstrap::ensure_ready(&transport)?; let rocm = cli.invocation(); - // 3. Start a managed server on the remote, bound to remote loopback. + // 3. Start a managed server on the remote. This returns quickly: the remote + // daemon loads the model in the background and tracks starting -> ready. let serve_cmd = build_remote_serve_command(rocm, &args); - println!( - " starting managed server (remote port {}) ...", - args.remote_port - ); transport .run(&serve_cmd) .context("failed to start the managed server on the remote host")?; - // 4. Discover the server we just started via the machine-readable listing, - // matching on the port we asked it to bind. + // 4. Discover the server we just started (match on the port we asked it to + // bind) so we can record its remote service id and current status. let record = discover_started_service(&transport, rocm, args.remote_port)?; - println!(" remote model: {}", record.canonical_model_id); - println!(" remote status: {}", record.status); - // 5. Open the loopback port-forward. - let mut guard = transport - .forward(local_port, "127.0.0.1", args.remote_port) + // 5. Open a detached loopback tunnel that survives this command exiting. + let local_port = match args.local_port { + Some(port) => port, + None => pick_free_local_port()?, + }; + let tunnel_pid = transport + .open_detached_forward(local_port, "127.0.0.1", args.remote_port) .context("failed to open the SSH port-forward")?; - // 6. Tell the user how to call it. - println!(); - println!("Ready. OpenAI-compatible base URL:"); - println!(" http://127.0.0.1:{local_port}/v1"); - println!(); + // 6. Register the session locally so status/attach/stop can find it later. + let session_id = session_id_for(&args.host, args.remote_port); + let session = RemoteSessionRecord::new( + session_id, + &args.host, + args.ssh_port, + &args.model, + &record.service_id, + rocm, + args.remote_port, + local_port, + tunnel_pid, + ); + session.write(paths)?; + + // 7. Return immediately with a pending status and a ready-to-use endpoint. + println!("Deploying {} on remote host {}...", args.model, args.host); + println!("Status: {}", remote_status_label(&record.status)); + println!("Local endpoint: {}", session.base_url); + println!( + "Track progress with `rocm remote status {}` (add --watch to keep it live).", + session.session_id + ); + Ok(()) +} + +/// Show one or all sessions; with `watch`, refresh until interrupted. +fn status(paths: &AppPaths, session: Option<&str>, watch: bool) -> Result<()> { + if !watch { + print!("{}", render_status_table(paths, session)?); + return Ok(()); + } + loop { + // Clear screen + home cursor, then redraw. + print!("\x1b[2J\x1b[H"); + print!("{}", render_status_table(paths, session)?); + println!( + "\n(watching every {}s — Ctrl-C to exit)", + WATCH_INTERVAL.as_secs() + ); + let _ = std::io::stdout().flush(); + sleep(WATCH_INTERVAL); + } +} + +/// Rebuild a dropped tunnel for a session, leaving the remote server untouched. +fn attach(paths: &AppPaths, needle: &str) -> Result<()> { + let mut session = resolve_session(paths, needle)?; + if rocm_core::process_is_running(session.tunnel_pid) { + println!( + "Tunnel for {} is already up (pid {}).", + session.session_id, session.tunnel_pid + ); + return Ok(()); + } + let transport = SshTransport::new(session.host.clone(), session.ssh_port); + let tunnel_pid = transport + .open_detached_forward(session.local_port, "127.0.0.1", session.remote_port) + .context("failed to reopen the SSH port-forward")?; + session.tunnel_pid = tunnel_pid; + session.write(paths)?; println!( - " try: rocm chat --provider local --model {} --prompt \"hello\"", - record.canonical_model_id + "Reattached {}: {} (tunnel pid {}).", + session.session_id, session.base_url, tunnel_pid ); - println!(" (set OPENAI_BASE_URL=http://127.0.0.1:{local_port}/v1 for other OpenAI clients)"); - println!(); - println!("The remote server is managed and will keep running after you disconnect."); - println!("Press Ctrl-C to close the local forward."); + Ok(()) +} - // 7. Hold the tunnel open until interrupted; the guard tears it down on drop. - guard.wait()?; +/// Stop a session: close the local tunnel and stop the remote server. +fn stop(paths: &AppPaths, needle: &str) -> Result<()> { + let session = resolve_session(paths, needle)?; + + // Close the local tunnel (best effort — it may already be gone). + if rocm_core::process_is_running(session.tunnel_pid) { + let _ = rocm_core::terminate_process(session.tunnel_pid); + } + + // Stop the remote server (best effort — surface but don't abort on failure). + let transport = SshTransport::new(session.host.clone(), session.ssh_port); + let stop_cmd = format!( + "{} services stop {} --yes", + session.remote_cli, session.remote_service_id + ); + if let Err(err) = transport.run(&stop_cmd) { + eprintln!( + "warning: could not stop the remote server for {}: {err:#}", + session.session_id + ); + } + + RemoteSessionRecord::remove(paths, &session.session_id)?; + println!("Stopped {} and removed the session.", session.session_id); Ok(()) } @@ -191,6 +301,160 @@ fn discover_started_service( }) } +/// Pick a free loopback port by binding to `:0` and reading the assigned port. +/// (Small TOCTOU window before `ssh` binds it; acceptable for v1.) +fn pick_free_local_port() -> Result { + let listener = std::net::TcpListener::bind("127.0.0.1:0") + .context("failed to find a free local port for the tunnel")?; + let port = listener + .local_addr() + .context("failed to read the assigned local port")? + .port(); + Ok(port) +} + +/// Deterministic, filename-safe session id from host + remote port. Re-serving +/// the same host+port intentionally reuses the same id (overwrites the record). +fn session_id_for(host: &str, remote_port: u16) -> String { + let slug: String = host + .chars() + .map(|c| if c.is_ascii_alphanumeric() { c } else { '-' }) + .collect(); + format!("{}-{remote_port}", slug.trim_matches('-')) +} + +/// Resolve a user handle — a session id, or a host that maps to exactly one +/// session — to its record. +fn resolve_session(paths: &AppPaths, needle: &str) -> Result { + let sessions = load_remote_sessions(paths)?; + if let Some(exact) = sessions.iter().find(|s| s.session_id == needle) { + return Ok(exact.clone()); + } + let by_host: Vec<&RemoteSessionRecord> = sessions.iter().filter(|s| s.host == needle).collect(); + match by_host.as_slice() { + [] => bail!( + "no remote session matches '{needle}'. Run `rocm remote status` to list sessions." + ), + [only] => Ok((*only).clone()), + many => { + let ids = many + .iter() + .map(|s| s.session_id.as_str()) + .collect::>() + .join(", "); + bail!("'{needle}' matches multiple sessions ({ids}); specify a session id.") + } + } +} + +/// Map a remote server status word to a display label. +fn remote_status_label(raw: &str) -> String { + match raw { + "ready" | "running" => "Healthy".to_owned(), + "starting" | "recovering" => "Pending".to_owned(), + "failed" | "stopped" => "Failed".to_owned(), + other => format!("Unknown ({other})"), + } +} + +/// Query the remote over SSH for this session's server status. +/// `Unreachable` = SSH failed; `Error` = the CLI ran but errored; `Gone` = the +/// service is no longer present on the remote. +fn query_remote_status(session: &RemoteSessionRecord) -> String { + let transport = SshTransport::new(session.host.clone(), session.ssh_port); + let listing = match transport.exec(&format!("{} services list --json", session.remote_cli)) { + Ok(outcome) if outcome.success => outcome.stdout, + Ok(_) => return "Error".to_owned(), + Err(_) => return "Unreachable".to_owned(), + }; + let records: Vec = match serde_json::from_str(listing.trim()) { + Ok(records) => records, + Err(_) => return "Error".to_owned(), + }; + records + .iter() + .find(|r| r.service_id == session.remote_service_id) + .map_or_else(|| "Gone".to_owned(), |r| remote_status_label(&r.status)) +} + +/// Render the status table for all sessions, or one when `filter` is given. +/// Wires the live remote-status probe; formatting lives in `format_sessions`. +fn render_status_table(paths: &AppPaths, filter: Option<&str>) -> Result { + let sessions = match filter { + Some(needle) => vec![resolve_session(paths, needle)?], + None => load_remote_sessions(paths)?, + }; + Ok(format_sessions( + &sessions, + filter.is_some(), + query_remote_status, + )) +} + +/// Pure formatter for the status table. `remote_status` supplies each session's +/// remote label (injected so tests don't hit the network); tunnel liveness is a +/// local PID check. `show_detail` adds the attach/stop footer for single-session +/// views. +fn format_sessions String>( + sessions: &[RemoteSessionRecord], + show_detail: bool, + remote_status: F, +) -> String { + if sessions.is_empty() { + return "No remote sessions. Start one with `rocm remote serve `.\n" + .to_owned(); + } + + let headers = ["Host", "Model", "Local endpoint", "Remote", "Tunnel"]; + let mut rows: Vec<[String; 5]> = Vec::new(); + for session in sessions { + let tunnel = if rocm_core::process_is_running(session.tunnel_pid) { + format!("Up (pid {})", session.tunnel_pid) + } else { + "Down".to_owned() + }; + rows.push([ + session.host.clone(), + session.model.clone(), + session.base_url.clone(), + remote_status(session), + tunnel, + ]); + } + + // Column widths = max of header and cells. + let mut widths = headers.map(str::len); + for row in &rows { + for (i, cell) in row.iter().enumerate() { + widths[i] = widths[i].max(cell.len()); + } + } + let mut out = String::new(); + let render_row = |out: &mut String, cells: &[String; 5]| { + for (i, cell) in cells.iter().enumerate() { + if i > 0 { + out.push_str(" "); + } + let _ = write!(out, "{cell: String { @@ -256,4 +520,106 @@ mod tests { assert_eq!(shell_quote("x; rm -rf /"), "'x; rm -rf /'"); assert_eq!(shell_quote("it's"), "'it'\\''s'"); } + + #[test] + fn session_id_is_filename_safe_and_deterministic() { + assert_eq!(session_id_for("gpu-box", 11435), "gpu-box-11435"); + assert_eq!(session_id_for("user@10.0.0.5", 8001), "user-10-0-0-5-8001"); + // Same host+port always yields the same id (record reuse on re-serve). + assert_eq!( + session_id_for("gpu-box", 11435), + session_id_for("gpu-box", 11435) + ); + } + + #[test] + fn remote_status_label_maps_lifecycle_words() { + assert_eq!(remote_status_label("ready"), "Healthy"); + assert_eq!(remote_status_label("starting"), "Pending"); + assert_eq!(remote_status_label("failed"), "Failed"); + assert_eq!(remote_status_label("weird"), "Unknown (weird)"); + } + + // Isolated AppPaths rooted at a temp dir, so registry tests don't touch the + // real data dir. + fn temp_paths(tag: &str) -> (std::path::PathBuf, AppPaths) { + let root = + std::env::temp_dir().join(format!("rocm-remote-test-{tag}-{}", std::process::id())); + let _ = std::fs::remove_dir_all(&root); + let paths = AppPaths { + config_dir: root.join("config"), + data_dir: root.join("data"), + cache_dir: root.join("cache"), + }; + (root, paths) + } + + fn sample_session(host: &str, remote_port: u16) -> RemoteSessionRecord { + RemoteSessionRecord::new( + session_id_for(host, remote_port), + host, + None, + "qwen2.5-7b", + "lemonade-qwen-abc", + "rocm", + remote_port, + 8001, + 999_999_999, // pid that is not running + ) + } + + #[test] + fn resolve_session_by_id_and_by_unique_host() { + let (root, paths) = temp_paths("resolve"); + sample_session("gpu-box", 11435).write(&paths).unwrap(); + + // By exact session id. + let by_id = resolve_session(&paths, "gpu-box-11435").unwrap(); + assert_eq!(by_id.host, "gpu-box"); + // By host (unique). + let by_host = resolve_session(&paths, "gpu-box").unwrap(); + assert_eq!(by_host.session_id, "gpu-box-11435"); + // Unknown handle errors. + assert!(resolve_session(&paths, "nope").is_err()); + + let _ = std::fs::remove_dir_all(root); + } + + #[test] + fn resolve_session_ambiguous_host_errors_with_ids() { + let (root, paths) = temp_paths("ambiguous"); + sample_session("gpu-box", 11435).write(&paths).unwrap(); + sample_session("gpu-box", 11436).write(&paths).unwrap(); + + let err = resolve_session(&paths, "gpu-box").unwrap_err().to_string(); + assert!(err.contains("multiple sessions"), "got: {err}"); + assert!(err.contains("gpu-box-11435")); + assert!(err.contains("gpu-box-11436")); + + let _ = std::fs::remove_dir_all(root); + } + + #[test] + fn status_table_shows_both_lifecycles() { + // Dead pid -> Tunnel: Down; remote status injected as Healthy. This is + // the degraded case (remote up, tunnel down) surfaced as two columns. + let session = sample_session("gpu-box", 11435); + let table = format_sessions(std::slice::from_ref(&session), true, |_| { + "Healthy".to_owned() + }); + assert!(table.contains("Host")); + assert!(table.contains("Local endpoint")); + assert!(table.contains("Healthy")); + assert!(table.contains("Down")); + assert!(table.contains("http://127.0.0.1:8001/v1")); + // Single-session view exposes the actionable handle. + assert!(table.contains("attach: rocm remote attach gpu-box-11435")); + assert!(table.contains("stop: rocm remote stop gpu-box-11435")); + } + + #[test] + fn status_table_empty_when_no_sessions() { + let table = format_sessions(&[], false, |_| "Healthy".to_owned()); + assert!(table.contains("No remote sessions")); + } } diff --git a/apps/rocm/src/remote/transport.rs b/apps/rocm/src/remote/transport.rs index 885846fd..2811a5c5 100644 --- a/apps/rocm/src/remote/transport.rs +++ b/apps/rocm/src/remote/transport.rs @@ -11,40 +11,13 @@ //! Why shell out to `ssh` instead of a Rust SSH crate: it reuses the user's //! existing keys/agent/`~/.ssh/config` and connection multiplexing for free, //! keeps this code synchronous (the surrounding command handlers are sync), and -//! makes the port-forward a plain child process that plugs into the existing -//! detach + PID-tracking machinery the managed-service supervisor already uses. +//! makes the port-forward a plain detached child process tracked by PID — the +//! same detach + liveness machinery the managed-service supervisor already uses. -use std::process::{Child, Command, Stdio}; +use std::process::{Command, Stdio}; use anyhow::{Context, Result, bail}; -/// A running local port-forward. Dropping it tears the forward down by killing -/// the underlying `ssh -N -L` child. -#[derive(Debug)] -pub struct ForwardGuard { - child: Child, -} - -impl ForwardGuard { - /// Block until the forward's `ssh` child exits (e.g. the connection drops or - /// the process group receives SIGINT). Used to hold a foreground session - /// open until the user interrupts it. - pub fn wait(&mut self) -> Result<()> { - self.child - .wait() - .context("failed while waiting on the ssh port-forward")?; - Ok(()) - } -} - -impl Drop for ForwardGuard { - fn drop(&mut self) { - // Best-effort teardown for the foreground path (Ctrl-C / normal return). - let _ = self.child.kill(); - let _ = self.child.wait(); - } -} - /// Captured result of running a command on the remote host. Distinguishes a /// clean-but-non-zero exit (e.g. `rocm --version` when the CLI is absent) from a /// transport failure (couldn't reach the host at all) — the bootstrap probes @@ -84,12 +57,18 @@ pub trait Transport { /// Copy a local file to `remote_path` on the host. fn push_file(&self, local_path: &std::path::Path, remote_path: &str) -> Result<()>; - /// Open a local port-forward: `127.0.0.1:local_port` on this machine is - /// forwarded to `remote_host:remote_port` as seen from the remote side - /// (typically `127.0.0.1:`). Returns a guard that keeps the - /// forward alive until dropped. - fn forward(&self, local_port: u16, remote_host: &str, remote_port: u16) - -> Result; + /// Open a **detached** local port-forward: `127.0.0.1:local_port` on this + /// machine is forwarded to `remote_host:remote_port` as seen from the remote + /// side (typically `127.0.0.1:`). The forward runs in its own + /// detached process that survives this command exiting; the returned PID is + /// recorded so the session can later be checked for liveness and torn down. + /// Killing the PID closes the forward. + fn open_detached_forward( + &self, + local_port: u16, + remote_host: &str, + remote_port: u16, + ) -> Result; } /// SSH-backed transport that shells out to the system `ssh`/`scp`. @@ -141,11 +120,33 @@ impl SshTransport { args } - /// Argument vector for the port-forward child (excludes the `ssh` program - /// name). `-N` = no remote command, `-T` = no pty; the process stays alive - /// holding the forward. Split out for unit testing. + /// Argument vector for the detached port-forward child (excludes the `ssh` + /// program name). `-N` = no remote command, `-T` = no pty. + /// + /// Unlike the exec path this uses **`ControlMaster=no`** — a dedicated + /// connection, not a shared/persisted master — so that killing this process + /// deterministically closes the forward (a persisted master could otherwise + /// outlive the PID). `ExitOnForwardFailure=yes` makes a failed bind exit + /// instead of silently staying up with no forward, and `ServerAliveInterval` + /// lets a dropped link terminate the process rather than hang. Split out for + /// unit testing. fn forward_argv(&self, local_port: u16, remote_host: &str, remote_port: u16) -> Vec { - let mut args = self.base_ssh_args(); + let mut args = vec![ + "-o".to_owned(), + "BatchMode=yes".to_owned(), + "-o".to_owned(), + "ConnectTimeout=10".to_owned(), + "-o".to_owned(), + "ControlMaster=no".to_owned(), + "-o".to_owned(), + "ExitOnForwardFailure=yes".to_owned(), + "-o".to_owned(), + "ServerAliveInterval=15".to_owned(), + ]; + if let Some(port) = self.port { + args.push("-p".to_owned()); + args.push(port.to_string()); + } args.push("-N".to_owned()); args.push("-T".to_owned()); args.push("-L".to_owned()); @@ -215,25 +216,42 @@ impl Transport for SshTransport { Ok(()) } - fn forward( + fn open_detached_forward( &self, local_port: u16, remote_host: &str, remote_port: u16, - ) -> Result { - let child = Command::new("ssh") - .args(self.forward_argv(local_port, remote_host, remote_port)) - .stdin(Stdio::null()) - .stdout(Stdio::null()) - .stderr(Stdio::null()) - .spawn() - .with_context(|| { - format!("failed to launch ssh port-forward to {}", self.destination) - })?; - Ok(ForwardGuard { child }) + ) -> Result { + let argv = self.forward_argv(local_port, remote_host, remote_port); + spawn_detached_ssh(&argv) + .with_context(|| format!("failed to launch ssh port-forward to {}", self.destination)) } } +/// Spawn `ssh ` as a detached process (its own session/process group) so +/// it survives the parent command exiting, and return its PID. Mirrors the +/// managed-service supervisor's platform-split detach. +#[cfg(not(windows))] +fn spawn_detached_ssh(argv: &[String]) -> Result { + let mut command = Command::new("ssh"); + command + .args(argv) + .stdin(Stdio::null()) + .stdout(Stdio::null()) + .stderr(Stdio::null()); + rocm_core::detach_command_session(&mut command); + let child = command.spawn().context("failed to spawn detached ssh")?; + // Intentionally do not hold the Child handle: on Unix, dropping it does not + // kill the process, and the setsid detach keeps it running past our exit. + Ok(child.id()) +} + +#[cfg(windows)] +fn spawn_detached_ssh(argv: &[String]) -> Result { + let program = std::path::Path::new("ssh"); + rocm_core::spawn_detached_no_inherit(program, argv, &[]).context("failed to spawn detached ssh") +} + #[cfg(test)] mod tests { use super::*; @@ -259,15 +277,29 @@ mod tests { #[test] fn forward_argv_maps_loopback_local_to_remote() { let t = SshTransport::new("gpubox", None); - let argv = t.forward_argv(11435, "127.0.0.1", 11435); + let argv = t.forward_argv(8001, "127.0.0.1", 11435); assert!(argv.contains(&"-N".to_owned())); assert!( argv.windows(2) - .any(|w| { w[0] == "-L" && w[1] == "127.0.0.1:11435:127.0.0.1:11435" }) + .any(|w| { w[0] == "-L" && w[1] == "127.0.0.1:8001:127.0.0.1:11435" }) ); assert_eq!(argv.last().unwrap(), "gpubox"); } + #[test] + fn forward_argv_uses_dedicated_connection_for_deterministic_teardown() { + // ControlMaster=no ensures killing the forward's PID actually closes it + // (a shared/persisted master could outlive the process). + let t = SshTransport::new("gpubox", None); + let argv = t.forward_argv(8001, "127.0.0.1", 11435); + assert!(argv.windows(2).any(|w| w == ["-o", "ControlMaster=no"])); + assert!( + argv.windows(2) + .any(|w| w == ["-o", "ExitOnForwardFailure=yes"]) + ); + assert!(!argv.iter().any(|a| a.starts_with("ControlPersist"))); + } + #[test] fn scp_argv_uses_uppercase_port_and_remote_colon_path() { let t = SshTransport::new("user@gpubox", Some(2222)); diff --git a/crates/rocm-core/src/lib.rs b/crates/rocm-core/src/lib.rs index 1bf707f0..629180e0 100644 --- a/crates/rocm-core/src/lib.rs +++ b/crates/rocm-core/src/lib.rs @@ -847,6 +847,17 @@ impl AppPaths { self.data_dir.join("services") } + /// Directory holding one JSON record per `rocm remote` serving session. + pub fn remote_sessions_dir(&self) -> PathBuf { + self.data_dir.join("remote-sessions") + } + + /// Path to a single remote-session record. + pub fn remote_session_path(&self, session_id: &str) -> PathBuf { + self.remote_sessions_dir() + .join(format!("{session_id}.json")) + } + pub fn audit_dir(&self) -> PathBuf { self.data_dir.join("audit") } @@ -5782,6 +5793,122 @@ impl ManagedServiceRecord { } } +/// One `rocm remote` serving session. +/// +/// A managed server on a remote host plus the local loopback tunnel that reaches +/// it. Persisted as one JSON file per session under +/// [`AppPaths::remote_sessions_dir`], mirroring [`ManagedServiceRecord`]. +/// +/// The record captures enough to reconcile two independent lifecycles later: +/// the remote server (via `remote_service_id` over SSH) and the local tunnel +/// (via `tunnel_pid`). +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct RemoteSessionRecord { + pub session_id: String, + /// SSH destination as given on the command line (`user@host` or alias). + pub host: String, + #[serde(default)] + pub ssh_port: Option, + pub model: String, + /// Service id of the managed server on the remote (`rocm services` there). + pub remote_service_id: String, + /// How to invoke `rocm` on the remote (`"rocm"` if on PATH, or an explicit + /// path like `~/.local/bin/rocm` when we pushed it — a non-interactive SSH + /// shell may not have `~/.local/bin` on PATH, so the exact form is stored). + pub remote_cli: String, + /// Port the server listens on, on the remote side. + pub remote_port: u16, + /// Loopback port on this machine that forwards to the remote server. + pub local_port: u16, + /// OpenAI-compatible base URL clients use locally. + pub base_url: String, + /// PID of the local `ssh -L` process holding the forward open. + pub tunnel_pid: u32, + /// Transport backing the session (`"ssh"` today; `"mesh"` later). + pub transport: String, + pub created_at_unix_ms: u128, +} + +impl RemoteSessionRecord { + #[allow(clippy::too_many_arguments)] + #[must_use] + pub fn new( + session_id: impl Into, + host: impl Into, + ssh_port: Option, + model: impl Into, + remote_service_id: impl Into, + remote_cli: impl Into, + remote_port: u16, + local_port: u16, + tunnel_pid: u32, + ) -> Self { + Self { + session_id: session_id.into(), + host: host.into(), + ssh_port, + model: model.into(), + remote_service_id: remote_service_id.into(), + remote_cli: remote_cli.into(), + remote_port, + local_port, + base_url: format!("{}/v1", format_http_base_url("127.0.0.1", local_port)), + tunnel_pid, + transport: "ssh".to_owned(), + created_at_unix_ms: unix_time_millis(), + } + } + + /// Persist this record as `/.json`. + pub fn write(&self, paths: &AppPaths) -> Result<()> { + let path = paths.remote_session_path(&self.session_id); + let parent = path + .parent() + .context("remote session path must have a parent directory")?; + fs::create_dir_all(parent) + .with_context(|| format!("failed to create {}", parent.display()))?; + fs::write( + &path, + serde_json::to_vec_pretty(self).context("failed to serialize remote session record")?, + ) + .with_context(|| format!("failed to write {}", path.display()))?; + Ok(()) + } + + /// Delete this session's record file. Missing file is not an error. + pub fn remove(paths: &AppPaths, session_id: &str) -> Result<()> { + let path = paths.remote_session_path(session_id); + match fs::remove_file(&path) { + Ok(()) => Ok(()), + Err(err) if err.kind() == std::io::ErrorKind::NotFound => Ok(()), + Err(err) => Err(err).with_context(|| format!("failed to remove {}", path.display())), + } + } +} + +/// Load all remote-session records, newest first. Unreadable/foreign files are +/// skipped rather than failing the whole listing. +pub fn load_remote_sessions(paths: &AppPaths) -> Result> { + let dir = paths.remote_sessions_dir(); + if !dir.is_dir() { + return Ok(Vec::new()); + } + let mut records = Vec::new(); + for entry in fs::read_dir(&dir).with_context(|| format!("failed to read {}", dir.display()))? { + let path = entry?.path(); + if path.extension().and_then(|value| value.to_str()) != Some("json") { + continue; + } + let bytes = + fs::read(&path).with_context(|| format!("failed to read {}", path.display()))?; + if let Ok(record) = serde_json::from_slice::(&bytes) { + records.push(record); + } + } + records.sort_by_key(|record| std::cmp::Reverse(record.created_at_unix_ms)); + Ok(records) +} + #[derive(Debug, Clone, Serialize, Deserialize)] pub struct CodexBridgeSnapshot { pub protocol: String, From 2533b457912db80622d5fad405d2ba5f9df453fd Mon Sep 17 00:00:00 2001 From: Eugene Volen Date: Fri, 3 Jul 2026 14:05:33 +0000 Subject: [PATCH 3/3] Add MIT license headers to the new remote module sources Signed-off-by: Eugene Volen --- apps/rocm/src/remote/bootstrap.rs | 4 ++++ apps/rocm/src/remote/mod.rs | 4 ++++ apps/rocm/src/remote/transport.rs | 4 ++++ 3 files changed, 12 insertions(+) diff --git a/apps/rocm/src/remote/bootstrap.rs b/apps/rocm/src/remote/bootstrap.rs index d51ec996..e90d8ff7 100644 --- a/apps/rocm/src/remote/bootstrap.rs +++ b/apps/rocm/src/remote/bootstrap.rs @@ -1,3 +1,7 @@ +// Copyright © Advanced Micro Devices, Inc., or its affiliates. +// +// SPDX-License-Identifier: MIT + //! Ensuring a remote host is ready to serve: probe for the `rocm` CLI and ROCm, //! and push the CLI if it is missing. //! diff --git a/apps/rocm/src/remote/mod.rs b/apps/rocm/src/remote/mod.rs index 0bff4076..97e56683 100644 --- a/apps/rocm/src/remote/mod.rs +++ b/apps/rocm/src/remote/mod.rs @@ -1,3 +1,7 @@ +// Copyright © Advanced Micro Devices, Inc., or its affiliates. +// +// SPDX-License-Identifier: MIT + //! `rocm remote` — bring up a managed model server on a remote GPU host and make //! it callable from this machine. //! diff --git a/apps/rocm/src/remote/transport.rs b/apps/rocm/src/remote/transport.rs index 2811a5c5..8081dd1b 100644 --- a/apps/rocm/src/remote/transport.rs +++ b/apps/rocm/src/remote/transport.rs @@ -1,3 +1,7 @@ +// Copyright © Advanced Micro Devices, Inc., or its affiliates. +// +// SPDX-License-Identifier: MIT + //! Transport abstraction for reaching a remote GPU host. //! //! `rocm remote` needs three things from a remote host: run a command and read