From 698d2c7ce96ce490746ea7894bdb80f289aac392 Mon Sep 17 00:00:00 2001 From: Eugene Volen Date: Mon, 14 Sep 2026 08:57:36 +0000 Subject: [PATCH 1/5] feat(remote): serve a model on a tailnet GPU machine MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds `rocm remote`: discover GPU machines on a tailnet, check their health, install what they are missing, serve a model on one, and reach it from any machine on the tailnet. SSH is the control channel, not the data path. Everything that inspects or changes the remote goes over SSH; the inference traffic does not. `rocm serve` binds loopback on the GPU machine as it always has, and the machine then tells its own Tailscale daemon to forward a tailnet port to it. Nothing runs locally, so the endpoint outlives the command that created it and answers from any machine rather than only the one that started it. Two touch points with existing behaviour: - `rocm serve --require-api-key` makes a loopback bind authenticated anyway. Publishing the port makes "loopback means only this machine" false while leaving the bind address unchanged, which would otherwise put an unauthenticated model endpoint on the tailnet. The key travels to the remote on stdin, never in a command line, since both machines expose command arguments in their process tables. - `install.sh` grows download-only and install-from-archive modes. Provisioning never copies the local binary — that only works when both machines share an OS and CPU, and when they do not the copy still lands and still looks installed. The remote fetches its own build; if it cannot reach the release host, this machine fetches one for the remote's platform and pushes it with its checksum and signature so the remote repeats every check. Signing-key selection matches install.sh's own order exactly, and a forwarded key blanks the remote's own path so the two machines cannot end up on different trust roots. `rocm services list --json` is the machine-readable listing the orchestration reads back, applying the same liveness filter as the table. Signed-off-by: Eugene Volen --- apps/rocm/src/main.rs | 668 +++++++- apps/rocm/src/remote/bootstrap.rs | 402 +++++ apps/rocm/src/remote/doctor.rs | 716 ++++++++ apps/rocm/src/remote/install.rs | 417 +++++ apps/rocm/src/remote/mod.rs | 2235 +++++++++++++++++++++++++ apps/rocm/src/remote/provision.rs | 729 ++++++++ apps/rocm/src/remote/publish.rs | 942 +++++++++++ apps/rocm/src/remote/session.rs | 718 ++++++++ apps/rocm/src/remote/tailnet.rs | 714 ++++++++ apps/rocm/src/remote/transport.rs | 918 ++++++++++ apps/rocmd/src/lib.rs | 322 +++- crates/e2e-report/src/lib.rs | 6 + crates/rocm-core/src/lib.rs | 23 + install.sh | 81 +- tests/e2e-cucumber/src/expectation.rs | 14 +- 15 files changed, 8841 insertions(+), 64 deletions(-) create mode 100644 apps/rocm/src/remote/bootstrap.rs create mode 100644 apps/rocm/src/remote/doctor.rs create mode 100644 apps/rocm/src/remote/install.rs create mode 100644 apps/rocm/src/remote/mod.rs create mode 100644 apps/rocm/src/remote/provision.rs create mode 100644 apps/rocm/src/remote/publish.rs create mode 100644 apps/rocm/src/remote/session.rs create mode 100644 apps/rocm/src/remote/tailnet.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 8b9e399bd..6c6f25439 100644 --- a/apps/rocm/src/main.rs +++ b/apps/rocm/src/main.rs @@ -14,6 +14,7 @@ mod endpoint_keys; mod logging; mod provider_keys; mod providers; +mod remote; mod serve_summary; mod storage; mod therock; @@ -409,6 +410,16 @@ rocm serve qwen --verbose --device gpu_required")] /// Allow binding to a non-local address. #[arg(long)] allow_public_bind: bool, + /// Require an API key even on a loopback bind. + /// + /// Loopback serving is credential-free because only this machine can + /// reach it. That stops being true when something else republishes the + /// port — a tailnet publish, a reverse proxy, a container port map — at + /// which point the bind address no longer describes who can call it. + /// Pass this to keep the endpoint authenticated anyway. `rocm remote` + /// sets it on every session it starts. + #[arg(long)] + require_api_key: bool, /// vLLM tool-call parser to enable OpenAI tool calling for this model /// (e.g. `hermes`, `llama3_json`, `mistral`). Overrides any catalog default /// and implies `--enable-auto-tool-choice`. Applies to vLLM only. @@ -444,8 +455,9 @@ rocm serve qwen --verbose --device gpu_required")] /// When binding a public interface and this is omitted, a strong key is /// generated automatically. Prefer the `ROCM_SERVE_API_KEY` environment /// variable over this flag so the secret does not appear in shell history - /// or the process table. Ignored for loopback binds, which stay - /// credential-free. + /// or the process table. Ignored for a loopback bind unless + /// `--require-api-key` is also passed, which makes a loopback endpoint + /// authenticated too. #[arg(long)] api_key: Option, }, @@ -460,6 +472,11 @@ rocm serve qwen --verbose --device gpu_required")] #[command(subcommand)] command: Option, }, + /// [preview] Work with GPU machines on your tailnet. + Remote { + #[command(subcommand)] + command: remote::RemoteCommand, + }, /// [preview] Manage optional background checks and review requests. Automations { #[command(subcommand)] @@ -855,6 +872,13 @@ enum ServicesCommand { /// Include failed, stopped, and old service records. #[arg(short, long)] all: bool, + /// Emit the service records as JSON instead of a table. + /// + /// This is the machine-readable form `rocm remote` reads back over its + /// control channel to discover which service a remote `rocm serve` just + /// started, rather than scraping the human table. + #[arg(long)] + json: bool, }, /// Show logs for a local model server. Logs { @@ -2327,6 +2351,7 @@ fn dispatch(cli: Cli) -> Result<()> { verbose, no_smoke_test, allow_public_bind, + require_api_key, tool_call_parser, gpu_memory_utilization, temperature, @@ -2347,6 +2372,7 @@ fn dispatch(cli: Cli) -> Result<()> { verbose, no_smoke_test, allow_public_bind, + require_api_key, tool_call_parser, gpu_memory_utilization, temperature, @@ -2356,6 +2382,7 @@ fn dispatch(cli: Cli) -> Result<()> { }), Some(Command::Comfyui { command }) => comfyui(command), Some(Command::Services { command }) => services(command), + Some(Command::Remote { command }) => remote::run(command), Some(Command::Automations { command }) => automations(command), Some(Command::Config { command }) => config(command), Some(Command::Logs { @@ -5475,6 +5502,7 @@ struct ServeArgs { verbose: bool, no_smoke_test: bool, allow_public_bind: bool, + require_api_key: bool, tool_call_parser: Option, gpu_memory_utilization: Option, temperature: Option, @@ -5498,6 +5526,7 @@ fn serve(args: ServeArgs) -> Result<()> { verbose, no_smoke_test, allow_public_bind, + require_api_key, tool_call_parser, gpu_memory_utilization, temperature, @@ -5517,7 +5546,7 @@ fn serve(args: ServeArgs) -> Result<()> { .ok() .filter(|value| !value.trim().is_empty()) }); - let endpoint_auth = resolve_endpoint_auth(&host, supplied_key.as_deref())?; + let endpoint_auth = resolve_endpoint_auth(&host, supplied_key.as_deref(), require_api_key)?; let paths = AppPaths::discover()?; let mut config = RocmCliConfig::load(&paths)?; // Host GPU detection can involve sysfs/WSL probing, so only run it when engine @@ -5834,6 +5863,7 @@ fn serve(args: ServeArgs) -> Result<()> { resolve.engine_recipe.as_ref(), endpoint_auth.as_deref(), launch_lock, + require_api_key, &mut |_elapsed| spinner.tick(), )?; ensure_background_helper_running_quiet(summary_mode)?; @@ -5913,6 +5943,7 @@ fn serve(args: ServeArgs) -> Result<()> { resolved_selection.env_id.as_deref(), endpoint_auth.as_deref(), launch_lock, + require_api_key, ) } @@ -5984,8 +6015,20 @@ fn is_loopback_host(host: &str) -> bool { /// otherwise generate a strong random one so a public endpoint can never come /// up anonymous. An empty/whitespace supplied key is rejected rather than /// silently treated as "no auth". -fn resolve_endpoint_auth(host: &str, supplied: Option<&str>) -> Result> { - if is_loopback_host(host) { +/// - **`required`** → treat a loopback bind as public for this purpose. +/// +/// That last case exists because "loopback" is a statement about the bind +/// address, not about who can reach the port. Publishing the port onto a +/// tailnet, proxying it, or mapping it out of a container all leave the bind +/// loopback while widening the audience — and the policy above would then hand +/// out an unauthenticated endpoint. Whoever widens the reach is responsible for +/// asking for the credential, so this is an explicit flag rather than a guess. +fn resolve_endpoint_auth( + host: &str, + supplied: Option<&str>, + required: bool, +) -> Result> { + if is_loopback_host(host) && !required { return Ok(None); } match supplied { @@ -6070,7 +6113,27 @@ fn ensure_public_bind_engine_supported( /// `key_present` is a plain `bool` rather than a path so both branches are /// unit-testable without touching the filesystem, mirroring `is_windows` in /// [`ensure_public_bind_engine_supported`]. -fn ensure_public_service_has_endpoint_key(host: &str, key_present: bool) -> Result<()> { +fn ensure_public_service_has_endpoint_key( + host: &str, + key_present: bool, + requires_api_key: bool, +) -> Result<()> { + // Two ways a service can need a key. A public bind is the obvious one. The + // other is a service that asked for auth on a loopback bind, because + // something outside this process republishes the port — a tailnet publish + // survives a reboot, let alone a restart, so "loopback" stops meaning + // "only this machine" and the bind address can no longer be trusted to + // answer the question on its own. + if requires_api_key && !key_present { + bail!( + "managed service was launched with `--require-api-key` but has no endpoint API key, \ + so restarting it would reopen it without authentication. Something outside this \ + machine may still be publishing its port. The key is dropped when a service stops \ + and cannot be recovered. Launch it again with \ + `rocm serve --require-api-key` (add `--api-key `, or set ROCM_SERVE_API_KEY, \ + to choose the key instead of generating one)." + ); + } if rocm_engine_protocol::is_public_bind_host(host) && !key_present { bail!( "managed service is bound to the public host `{host}` but has no endpoint API key, \ @@ -6214,6 +6277,7 @@ fn spawn_managed_engine_child( runtime_id: Option<&str>, env_id: Option<&str>, engine_recipe: Option<&EngineRecipeHint>, + require_api_key: bool, ) -> Result { paths.ensure()?; fs::create_dir_all(paths.services_dir())?; @@ -6276,6 +6340,19 @@ fn spawn_managed_engine_child( ); record.gpu_indices = gpu_indices.to_vec(); record.engine_recipe_json = requested_recipe_json; + // The flag the user actually passed, carried through rather than re-derived. + // + // Deriving it from key-file presence looked equivalent and was not: + // `resolve_endpoint_auth` mints a key for *every* non-loopback bind whether or + // not auth was demanded, so a plain `--host 0.0.0.0 --allow-public-bind` + // recorded `true` here. The guard below tests this field before the bind + // address, so that service was then refused with a message naming a flag it + // never used and a relaunch command that drops `--allow-public-bind` — the + // public-bind branch, which carries the right command, became unreachable. + // + // This field means "the user demanded auth on a bind that would not otherwise + // require it". A public bind needs no such record; its address still says so. + record.requires_api_key = require_api_key; record.write()?; if let Some(parent) = record.engine_state_path.parent() { @@ -6311,11 +6388,24 @@ fn spawn_managed_engine_child( // still produce an unauthenticated public listener. let endpoint_key_file = endpoint_keys::endpoint_key_file_if_present(paths, service_id) .filter(|path| rocm_engine_protocol::endpoint_api_key_from_file(path).is_some()); - // `serve()` already resolved and stored the key for a public bind, so this - // cannot fire on the fresh-launch path today. It is the shared choke point - // for managed spawns, so enforce the invariant here too rather than relying - // on every future caller having done so. - ensure_public_service_has_endpoint_key(host, endpoint_key_file.is_some())?; + // `serve()` already resolved and stored the key for a public bind, so the + // public-bind branch cannot fire on the fresh-launch path today. It is the + // shared choke point for managed spawns, so enforce the invariant here too + // rather than relying on every future caller having done so. + // + // `record.requires_api_key` is passed, not a literal, and the two arguments + // are deliberately different things: that field is the `--require-api-key` + // flag the caller passed, `endpoint_key_file` is whether a *usable* key is on + // disk. A present but empty or malformed key file is where they disagree, and + // is exactly what the `requires_api_key` branch exists to refuse. + // + // The field is threaded, never derived from the key file. Deriving it marked + // every public bind as having demanded auth — see the assignment above. + ensure_public_service_has_endpoint_key( + host, + endpoint_key_file.is_some(), + record.requires_api_key, + )?; #[cfg(windows)] let child_pid = { let env_values = app_path_env_var_values(paths, engine_envs_root.as_deref()); @@ -6384,6 +6474,7 @@ fn start_managed_service( engine_recipe: Option<&EngineRecipeHint>, endpoint_api_key: Option<&str>, launch_lock: rocm_core::FileLock, + require_api_key: bool, on_wait_tick: &mut dyn FnMut(Duration), ) -> Result { let paths = AppPaths::discover()?; @@ -6400,6 +6491,7 @@ fn start_managed_service( runtime_id, env_id, engine_recipe, + require_api_key, )? { ManagedSpawn::AlreadyRunning(report) => return Ok(report), ManagedSpawn::Spawned { record, child_pid } => (*record, child_pid), @@ -6697,6 +6789,7 @@ fn run_attached_service( env_id: Option<&str>, endpoint_api_key: Option<&str>, launch_lock: rocm_core::FileLock, + require_api_key: bool, ) -> Result<()> { let paths = AppPaths::discover()?; @@ -6713,6 +6806,7 @@ fn run_attached_service( runtime_id, env_id, resolve.engine_recipe.as_ref(), + require_api_key, )?; // The claiming record is persisted (or an existing service was found), so the // selected GPU is now visible to concurrent auto-selection. Release the launch @@ -6963,9 +7057,16 @@ fn stream_attached_logs_no_tty(log_path: &Path, child_pid: u32) -> Result) -> 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 { + print!("{}", render_services_json(&paths, all)?); + } else { + print!("{}", render_services_text(&paths, all)?); + } Ok(()) } ServicesCommand::Logs { service_id } => { @@ -13586,6 +13687,31 @@ fn chat_rocm_command_action_from_args(mut args: Vec) -> Result + { + Ok(ChatRocmCommandAction::ReadOnly(args)) + } + // `serve`, `attach` and `stop` start, publish or tear down something on + // another machine, so they go through approval like any other mutation. + Some("remote") + if second + .as_deref() + .is_some_and(|value| matches!(value, "serve" | "attach" | "stop")) => + { + let verb = second.as_deref().unwrap_or_default().to_owned(); + Ok(ChatRocmCommandAction::Approval { + args, + pending_title: format!("Remote {verb}"), + command_title: "Remote".to_owned(), + }) + } Some(command) => bail!("local assistant cannot use unsupported rocm command `{command}`"), None => bail!("rocm_command requires at least one argument"), } @@ -17010,6 +17136,25 @@ pub(crate) fn render_services_text(paths: &AppPaths, all: bool) -> Result Result { + let records = load_managed_services(paths)? + .into_iter() + .filter(|record| all || managed_service_is_live(record)) + .collect::>(); + let mut output = serde_json::to_string_pretty(&records) + .context("failed to serialize the managed service records as JSON")?; + output.push('\n'); + Ok(output) +} + fn render_services_tool_result_text(records: &[ManagedServiceRecord]) -> String { let mut output = String::new(); let _ = writeln!(output, "managed_services: {}", records.len()); @@ -17380,7 +17525,11 @@ fn restart_internal_managed_service( let preserved_endpoint_key = endpoint_keys::endpoint_api_key(paths, service_id); // Checked before the stop, so a refused restart leaves a running service // running instead of stopping it and then failing to bring it back. - ensure_public_service_has_endpoint_key(&record.host, preserved_endpoint_key.is_some())?; + ensure_public_service_has_endpoint_key( + &record.host, + preserved_endpoint_key.is_some(), + record.requires_api_key, + )?; let _ = stop_internal_managed_service(paths, service_id); if let Some(key) = preserved_endpoint_key.as_deref() { endpoint_keys::store_endpoint_api_key(paths, service_id, key)?; @@ -19646,6 +19795,45 @@ fn build_uninstall_plan(paths: &AppPaths, options: &UninstallOptions) -> Result< )); } + // Remote sessions are worse than local ones to drop silently. The model runs + // on someone else's machine and its endpoint is published there, so removing + // the record here does not stop either — it only destroys the last thing that + // knew they existed. Name them and the command that tears them down properly. + // + // A read that fails is reported, not defaulted away. `load_all` already warns + // past an individual unreadable record and keeps going, so an `Err` here is a + // directory-level I/O failure, and a missing directory is `Ok`. Taking the + // default would state "there are none" on the one path where we do not know — + // and this plan goes on to delete the data directory those records live in, + // so the warning that something is still published elsewhere would be lost at + // exactly the moment it was the last copy. + let remote_sessions = match remote::session::load_all(paths) { + Ok(sessions) => sessions, + Err(error) => { + plan.warnings.push(format!( + "could not read the remote session records under {}: {error:#}\n\ + This pass cannot tell whether models are still running on other machines with \ + their endpoints published. Check with `rocm remote status` before continuing.", + paths.remote_sessions_dir().display() + )); + Vec::new() + } + }; + if !remote_sessions.is_empty() { + plan.warnings.push(format!( + "{} remote session record(s) exist under {}; their models keep running on the \ + remote machines and their endpoints stay published. Removing these records only \ + loses track of them — run `rocm remote stop ` for each first: {}", + remote_sessions.len(), + paths.remote_sessions_dir().display(), + remote_sessions + .iter() + .map(|session| session.session_id.as_str()) + .collect::>() + .join(", ") + )); + } + plan.actions .sort_by(|left, right| left.path.cmp(&right.path)); plan.actions.dedup_by(|left, right| left.path == right.path); @@ -21245,6 +21433,7 @@ fn treat_as_natural_language(args: &[String]) -> bool { "comfyui", "comfy", "services", + "remote", "automations", "config", "logs", @@ -25261,6 +25450,17 @@ model recipes "old-runtime".to_owned(), "--dry-run".to_owned(), ], + // Must agree with the daemon's `ensure_rocm_command_is_read_only`. + // Its comment claims the two are mirrored, and nothing enforced + // that — the arm was added there and not here, and the classifiers + // disagreed until a reviewer noticed. + vec!["remote".to_owned(), "targets".to_owned()], + vec![ + "remote".to_owned(), + "doctor".to_owned(), + "gpu-box".to_owned(), + ], + vec!["remote".to_owned(), "status".to_owned()], ]; for args in read_only { let action = chat_rocm_command_action_from_args(args.clone()) @@ -25271,32 +25471,67 @@ model recipes ); } + // Paired with whether the command has a `--yes` to inject. The flag + // exists to keep a consent prompt from hanging a null-stdin spawn, so + // the demand only makes sense for commands that would prompt — and + // injecting it where clap defines no such flag would make the spawn fail + // to parse rather than succeed unattended. let mutating = [ - vec!["update".to_owned(), "--apply".to_owned()], - vec!["comfyui".to_owned(), "install".to_owned()], - vec!["comfyui".to_owned(), "start".to_owned()], - vec!["comfyui".to_owned(), "stop".to_owned()], - vec!["uninstall".to_owned()], - vec!["setup".to_owned(), "reset".to_owned()], - vec![ - "runtimes".to_owned(), - "uninstall".to_owned(), - "old-runtime".to_owned(), - ], - vec![ - "runtimes".to_owned(), - "remove".to_owned(), - "old-runtime".to_owned(), - ], + (vec!["update".to_owned(), "--apply".to_owned()], true), + // These start, publish or tear down on another machine. They take no + // `--yes` because they never prompt: consent is the approval step + // itself, and every destructive choice they make is already settled + // by an explicit flag (`remote stop --force`). Give any of them an + // interactive prompt and it needs a consent flag here too. + ( + vec![ + "remote".to_owned(), + "serve".to_owned(), + "gpu-box".to_owned(), + "m".to_owned(), + ], + false, + ), + ( + vec!["remote".to_owned(), "attach".to_owned(), "sess".to_owned()], + false, + ), + ( + vec!["remote".to_owned(), "stop".to_owned(), "sess".to_owned()], + false, + ), + (vec!["comfyui".to_owned(), "install".to_owned()], true), + (vec!["comfyui".to_owned(), "start".to_owned()], true), + (vec!["comfyui".to_owned(), "stop".to_owned()], true), + (vec!["uninstall".to_owned()], true), + (vec!["setup".to_owned(), "reset".to_owned()], true), + ( + vec![ + "runtimes".to_owned(), + "uninstall".to_owned(), + "old-runtime".to_owned(), + ], + true, + ), + ( + vec![ + "runtimes".to_owned(), + "remove".to_owned(), + "old-runtime".to_owned(), + ], + true, + ), ]; - for args in mutating { + for (args, expects_yes) in mutating { let action = chat_rocm_command_action_from_args(args.clone()) .unwrap_or_else(|err| panic!("{args:?} should classify: {err}")); match &action { ChatRocmCommandAction::Approval { args, .. } => { - assert!( + assert_eq!( args.iter().any(|arg| arg == "--yes"), - "{args:?} should have --yes injected for the approval path" + expects_yes, + "{args:?} disagrees with whether the approval path should \ + carry --yes" ); } other @ ChatRocmCommandAction::ReadOnly(_) => { @@ -26938,6 +27173,126 @@ install therock"; Ok(()) } + #[test] + fn services_json_round_trips_and_applies_the_same_liveness_filter_as_the_table() -> Result<()> { + // This JSON is a contract, not a convenience: `rocm remote` parses it + // back over its control channel to learn which service a remote serve + // just started. Two things have to hold — every record survives the + // round trip, and `--json` agrees with the table about what is live. If + // they disagreed, the remote orchestration would act on a different set + // of services than the operator sees. + let (root, paths) = test_paths("services-json"); + paths.ensure()?; + let current_pid = std::process::id(); + for (service_id, status, port) in [ + ("svc-live", "starting", 11440_u16), + ("svc-past", "failed", 11441_u16), + ] { + let mut record = ManagedServiceRecord::new( + &paths, + service_id, + "vllm", + "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()?; + } + + let live = render_services_json(&paths, false)?; + let every = render_services_json(&paths, true)?; + let _ = fs::remove_dir_all(root); + + let live: Vec = serde_json::from_str(&live)?; + let every: Vec = serde_json::from_str(&every)?; + + assert_eq!( + live.iter() + .map(|r| r.service_id.as_str()) + .collect::>(), + vec!["svc-live"], + "the default listing must hide past attempts, exactly as the table does" + ); + let mut every_ids = every + .iter() + .map(|r| r.service_id.as_str()) + .collect::>(); + every_ids.sort_unstable(); + assert_eq!(every_ids, vec!["svc-live", "svc-past"]); + + // The fields remote orchestration actually reads must survive intact. + let record = &live[0]; + assert_eq!(record.port, 11440); + assert_eq!(record.status, "starting"); + // Note for remote orchestration: the recorded endpoint is already the + // OpenAI-compatible base, `/v1` suffix included — not a bare origin. + assert_eq!(record.endpoint_url, "http://127.0.0.1:11440/v1"); + assert_eq!(record.canonical_model_id, "Qwen/Qwen3.5"); + Ok(()) + } + + #[test] + fn service_records_tolerate_unknown_fields_but_not_missing_required_ones() -> Result<()> { + // A remote may run a different CLI version than the machine driving it. + // Newer fields it emits must not break an older parser, or a version skew + // turns every remote command into a parse error; a genuinely absent + // required field must still fail, and name itself when it does. + // Built from a real record rather than hand-written JSON, so the fixture + // cannot drift out of step with the struct and quietly stop testing the + // thing it claims to. + let (root, paths) = test_paths("services-json-contract"); + paths.ensure()?; + let record = ManagedServiceRecord::new( + &paths, + "svc-a", + "vllm", + "qwen", + "Qwen/Qwen3.5", + "127.0.0.1", + 11440, + "managed", + 4242, + None, + None, + None, + ); + let mut value = serde_json::to_value(&record)?; + let _ = fs::remove_dir_all(root); + let fields = value + .as_object_mut() + .expect("a service record serializes as a JSON object"); + + fields.insert( + "a_field_from_a_newer_release".to_owned(), + serde_json::Value::Bool(true), + ); + let parsed: ManagedServiceRecord = serde_json::from_value(value.clone()) + .context("a newer remote's extra fields must not break an older parser")?; + assert_eq!(parsed.service_id, "svc-a"); + assert_eq!(parsed.port, 11440); + + value + .as_object_mut() + .expect("still an object") + .remove("port") + .expect("port was present before removal"); + let error = serde_json::from_value::(value) + .expect_err("a missing required field must be rejected, not defaulted") + .to_string(); + assert!( + error.contains("port"), + "the error should name the missing field, got: {error}" + ); + Ok(()) + } + #[test] fn render_services_text_demotes_stale_ready_record() -> Result<()> { let (root, paths) = test_paths("services-stale-ready"); @@ -28059,6 +28414,7 @@ install therock"; None, None, Some(&requested_recipe), + false, ); let _ = fs::remove_dir_all(root); @@ -28077,6 +28433,76 @@ install therock"; Ok(()) } + #[test] + fn a_managed_spawn_refuses_an_invalid_key_file_on_a_service_that_requires_one() -> Result<()> { + // Drives the real call site, not the guard's own arguments. The service + // was launched with `--require-api-key`, and its key file is present but + // empty — so `requires_api_key` is true while `key_present` is false, + // which is the only way to reach the first branch. A test calling + // `ensure_public_service_has_endpoint_key` directly cannot catch a + // mis-wired call site, which is the defect that has occurred here twice. + // + // The flag is passed explicitly rather than inferred from the key file. + // Inferring it marked every public bind as having demanded auth, because + // a public bind always has a key file whether or not it asked for one. + let (root, paths) = test_paths("managed-spawn-invalid-key"); + paths.ensure()?; + fs::create_dir_all(paths.services_dir())?; + + // Empty, so the file exists (requires_api_key = true) but yields no + // usable key (key_present = false). + endpoint_keys::store_endpoint_api_key(&paths, "lemonade-qwen-3000", "")?; + + let resolve = ResolveModelResponse { + canonical_model_id: "qwen-canonical".to_owned(), + task: "chat".to_owned(), + source: "hf".to_owned(), + revision: "main".to_owned(), + loader: "llama.cpp".to_owned(), + trust_remote_code: false, + chat_template_mode: "auto".to_owned(), + dtype: "auto".to_owned(), + device_policy: DevicePolicy::GpuPreferred, + estimated_memory: "unknown".to_owned(), + launch_defaults: serde_json::json!({}), + engine_recipe: None, + warnings: Vec::new(), + }; + + let result = spawn_managed_engine_child( + &paths, + "lemonade", + "lemonade-qwen-3000", + "qwen", + &resolve, + // Loopback on purpose: the public-bind branch must not be what + // refuses this, or the test would pass with the guard disabled. + "127.0.0.1", + 11512, + &resolve.device_policy, + &[], + None, + None, + None, + true, + ); + let _ = fs::remove_dir_all(root); + + let Err(error) = result else { + panic!("a service requiring a key must not spawn with an unusable key file") + }; + let message = error.to_string(); + assert!( + message.contains("--require-api-key"), + "the refusal must name the flag the service was launched with: {message}" + ); + assert!( + message.contains("without authentication"), + "the refusal must say what the risk is: {message}" + ); + Ok(()) + } + #[test] fn services_tool_result_text_includes_running_interpretation() { let (_root, paths) = test_paths("services-tool-text"); @@ -28188,6 +28614,31 @@ install therock"; assert!(update_should_preview_or_apply(true, true)); } + #[test] + fn remote_is_structured_not_freeform() { + // `rocm remote …` reads like a plain-English request, so without an + // entry in the structured allowlist the natural-language planner + // swallows it and the real command becomes unreachable. This guards the + // allowlist against losing `remote`. + let invocation = parse_freeform_invocation(&[ + "remote".to_owned(), + "targets".to_owned(), + "--tag".to_owned(), + "gpu".to_owned(), + ]); + assert!(!treat_as_natural_language(&invocation.request_args)); + assert!(!should_treat_as_freeform(&invocation)); + + Cli::try_parse_from(["rocm", "remote", "targets"]) + .expect("remote targets should be a real command"); + Cli::try_parse_from(["rocm", "remote", "targets", "--tag", "gpu"]) + .expect("remote targets should accept a tag filter"); + // The group has no useful default action, so a bare `rocm remote` must + // show help rather than silently doing something. + Cli::try_parse_from(["rocm", "remote"]) + .expect_err("bare `rocm remote` should require a subcommand"); + } + #[test] fn install_sdk_accepts_family_override() { Cli::try_parse_from([ @@ -28395,14 +28846,42 @@ install therock"; fn resolve_endpoint_auth_loopback_stays_credential_free() { // Loopback binds never require auth, even if a key is supplied. for host in ["127.0.0.1", "localhost", "::1"] { - assert_eq!(resolve_endpoint_auth(host, None).unwrap(), None); - assert_eq!(resolve_endpoint_auth(host, Some("ignored")).unwrap(), None); + assert_eq!(resolve_endpoint_auth(host, None, false).unwrap(), None); + assert_eq!( + resolve_endpoint_auth(host, Some("ignored"), false).unwrap(), + None + ); } } + #[test] + fn resolve_endpoint_auth_loopback_can_be_required_when_something_republishes_it() { + // "Loopback" describes the bind address, not who can reach the port. A + // tailnet publish, a proxy, or a container port map all leave the bind + // loopback while widening the audience, and the default policy would + // hand out an unauthenticated endpoint. Whoever widens the reach asks + // for the credential explicitly. + for host in ["127.0.0.1", "localhost", "::1"] { + let generated = resolve_endpoint_auth(host, None, true) + .unwrap() + .expect("a required key must be generated, not skipped"); + assert!(!generated.trim().is_empty()); + + assert_eq!( + resolve_endpoint_auth(host, Some("supplied-key"), true).unwrap(), + Some("supplied-key".to_owned()), + "a supplied key must be honoured rather than ignored as it is by default" + ); + } + + // The same validation a public bind gets: an empty key is a refusal, not + // a silent downgrade to no auth. + assert!(resolve_endpoint_auth("127.0.0.1", Some(" "), true).is_err()); + } + #[test] fn resolve_endpoint_auth_public_uses_supplied_key_trimmed() { - let key = resolve_endpoint_auth("0.0.0.0", Some(" my-key ")) + let key = resolve_endpoint_auth("0.0.0.0", Some(" my-key "), false) .unwrap() .expect("public bind must have a key"); assert_eq!(key, "my-key"); @@ -28410,7 +28889,7 @@ install therock"; #[test] fn resolve_endpoint_auth_public_generates_key_when_absent() { - let key = resolve_endpoint_auth("0.0.0.0", None) + let key = resolve_endpoint_auth("0.0.0.0", None, false) .unwrap() .expect("public bind must generate a key"); assert_eq!(key.len(), 48); @@ -28419,7 +28898,7 @@ install therock"; #[test] fn resolve_endpoint_auth_public_rejects_empty_supplied_key() { - let error = resolve_endpoint_auth("0.0.0.0", Some(" ")).unwrap_err(); + let error = resolve_endpoint_auth("0.0.0.0", Some(" "), false).unwrap_err(); assert!(error.to_string().contains("non-empty"), "{error:#}"); } @@ -28433,7 +28912,7 @@ install therock"; "good-key\nmore", "line\rreturn", ] { - let error = resolve_endpoint_auth("0.0.0.0", Some(supplied)).unwrap_err(); + let error = resolve_endpoint_auth("0.0.0.0", Some(supplied), false).unwrap_err(); assert!(error.to_string().contains("control character"), "{error:#}"); } } @@ -28476,11 +28955,93 @@ install therock"; ensure_public_bind_engine_supported("lemonade", false, true).unwrap(); // loopback needs no key } + #[test] + fn a_public_bind_is_refused_with_the_command_that_restores_it() { + // `resolve_endpoint_auth` mints a key for every non-loopback bind whether + // or not auth was demanded, so deriving `requires_api_key` from key-file + // presence marked ordinary public binds as having asked for it. The guard + // tests that field first, so those services were refused with a message + // naming a flag they never passed and a relaunch command that drops + // `--allow-public-bind` — coming back on loopback instead. + // + // A plain `--host 0.0.0.0 --allow-public-bind` launch: key file present, + // `--require-api-key` never passed. + let error = ensure_public_service_has_endpoint_key("0.0.0.0", false, false) + .expect_err("a public bind with no key must be refused"); + let rendered = format!("{error:#}"); + assert!( + rendered.contains("--allow-public-bind"), + "the refusal must name the command that restores the public bind: {rendered}" + ); + assert!( + !rendered.contains("--require-api-key"), + "a service that never passed the flag must not be told it did: {rendered}" + ); + + // And the loopback-with-forced-auth case still reports its own reason. + let error = ensure_public_service_has_endpoint_key("127.0.0.1", false, true) + .expect_err("a service that demanded auth must not come back without it"); + assert!( + format!("{error:#}").contains("--require-api-key"), + "{error:#}" + ); + } + + #[test] + fn restarting_a_keyless_public_service_names_the_public_bind_not_the_key_flag() -> Result<()> { + // The same defect at a real call site rather than through the guard's own + // arguments. `restart` reads `requires_api_key` off the record on disk, so + // a record written by an ordinary public-bind launch must not claim the + // service asked for `--require-api-key` — the branch order means a record + // that claims it wins, and its remediation drops `--allow-public-bind`. + let (root, paths) = test_paths("restart-public-no-key"); + paths.ensure()?; + let mut record = ManagedServiceRecord::new( + &paths, + "vllm-public-2000", + "vllm", + "qwen", + "qwen-canonical", + "0.0.0.0", + 12000, + "managed", + std::process::id(), + None, + None, + None, + ); + // What a plain `--host 0.0.0.0 --allow-public-bind` launch records: the + // bind is public, and the flag was never passed. + record.requires_api_key = false; + record.write()?; + + // No key file: the situation after a stop, which drops it. + let error = restart_internal_managed_service(&paths, "vllm-public-2000") + .expect_err("a keyless public service must not be restarted"); + let rendered = format!("{error:#}"); + assert!( + rendered.contains("--allow-public-bind"), + "the refusal must name the command that brings it back public: {rendered}" + ); + assert!( + !rendered.contains("--require-api-key"), + "a service that never passed the flag must not be told it did: {rendered}" + ); + + // The refusal happens before the stop, so the service is left alone. + assert!( + load_managed_service(&paths, "vllm-public-2000").is_ok(), + "a refused restart must not have removed the record" + ); + let _ = fs::remove_dir_all(root); + Ok(()) + } + #[test] fn respawn_fails_closed_for_a_public_service_whose_key_is_gone() { // A stop deletes the key file, so a later restart of a public service // would otherwise respawn it with no auth at all. - let error = ensure_public_service_has_endpoint_key("0.0.0.0", false).unwrap_err(); + let error = ensure_public_service_has_endpoint_key("0.0.0.0", false, false).unwrap_err(); let message = error.to_string(); assert!(message.contains("0.0.0.0"), "{error:#}"); assert!(message.contains("without authentication"), "{error:#}"); @@ -28488,19 +29049,38 @@ install therock"; assert!(message.contains("--allow-public-bind"), "{error:#}"); // A public service that still has its key restarts normally. - ensure_public_service_has_endpoint_key("0.0.0.0", true).unwrap(); + ensure_public_service_has_endpoint_key("0.0.0.0", true, false).unwrap(); } #[test] fn respawn_allows_loopback_services_without_an_endpoint_key() { // Loopback stays credential-free, so every accepted spelling must pass - // the guard with no key present. + // the guard with no key present — when nothing asked for auth. for host in ["127.0.0.1", "localhost", "::1"] { - ensure_public_service_has_endpoint_key(host, false) + ensure_public_service_has_endpoint_key(host, false, false) .unwrap_or_else(|error| panic!("{host} must not require a key: {error:#}")); } } + #[test] + fn respawn_refuses_a_loopback_service_that_was_launched_with_a_key() { + // The hole this closes: a loopback bind that something else republishes + // — a tailnet publish, a proxy, a container port map. The publish + // outlives the process, so a restart after the key was dropped would + // reopen a reachable endpoint with no authentication, and the bind + // address gives the guard no way to notice. + for host in ["127.0.0.1", "localhost", "::1"] { + let error = ensure_public_service_has_endpoint_key(host, false, true) + .expect_err("a service launched with a key must not restart without one"); + let message = format!("{error:#}"); + assert!(message.contains("without authentication"), "{message}"); + assert!(message.contains("--require-api-key"), "{message}"); + } + + // With its key still present it restarts normally. + ensure_public_service_has_endpoint_key("127.0.0.1", true, true).unwrap(); + } + #[test] fn restart_refuses_a_public_service_without_a_key_before_stopping_it() { // The guard runs before the stop, so a refused restart must leave the diff --git a/apps/rocm/src/remote/bootstrap.rs b/apps/rocm/src/remote/bootstrap.rs new file mode 100644 index 000000000..c546dc437 --- /dev/null +++ b/apps/rocm/src/remote/bootstrap.rs @@ -0,0 +1,402 @@ +// Copyright © Advanced Micro Devices, Inc., or its affiliates. +// +// SPDX-License-Identifier: MIT + +//! Deciding whether a remote machine can host a model, before we ask it to. +//! +//! Four things have to be true: the machine has a GPU stack (ROCm), it has the +//! ROCm CLI to drive it, it has Tailscale so the endpoint can be published, and +//! we know what kind of machine it is so anything we install matches. +//! +//! Every failure here is refused up front rather than discovered halfway +//! through. Starting a model server and *then* finding the endpoint cannot be +//! published leaves a process running on someone's GPU with no way for them to +//! reach it and no record that it exists. + +use anyhow::{Result, bail}; + +use super::transport::Transport; + +/// What a remote machine looks like right now. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct RemoteReadiness { + /// How to invoke the CLI, when it is present. + pub(crate) cli: Option, + pub(crate) cli_version: Option, + pub(crate) rocm_present: bool, + pub(crate) tailscale_present: bool, + /// Operating system and CPU architecture, as the machine reports them. + /// Needed before anything is installed onto it. + pub(crate) platform: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct RemotePlatform { + pub(crate) os: String, + pub(crate) arch: String, +} + +/// Where a provisioned CLI lands, matching the documented manual install path. +pub(crate) const REMOTE_CLI_PATH: &str = "$HOME/.local/bin/rocm"; + +/// Inspect the remote. Nothing here is an error: absence is a finding, and +/// [`ensure_ready`] is where it becomes a refusal. +pub(crate) fn probe(transport: &dyn Transport) -> Result { + // Look on PATH first, then where we would have installed it. A CLI put + // there by an earlier run is often not on a non-interactive shell's PATH, + // and missing it would mean provisioning over a perfectly good install. + let mut cli = None; + let mut cli_version = None; + for candidate in ["rocm", REMOTE_CLI_PATH] { + let outcome = transport.exec(&format!("{candidate} --version"))?; + if outcome.success { + cli = Some(candidate.to_owned()); + cli_version = outcome + .stdout + .lines() + .next() + .map(str::trim) + .filter(|line| !line.is_empty()) + .map(ToOwned::to_owned); + break; + } + } + + // Detect ROCm independently of our own CLI: the machine may have a working + // GPU stack and no CLI at all, which is a provisioning job, not a refusal. + let rocm = transport.exec( + "command -v rocminfo >/dev/null 2>&1 || command -v amd-smi >/dev/null 2>&1 || test -d /opt/rocm", + )?; + let tailscale = transport.exec("command -v tailscale >/dev/null 2>&1")?; + + let platform = transport.exec("uname -s && uname -m")?; + let platform = if platform.success { + let mut lines = platform.stdout.lines().map(str::trim); + match (lines.next(), lines.next()) { + (Some(os), Some(arch)) if !os.is_empty() && !arch.is_empty() => Some(RemotePlatform { + os: os.to_ascii_lowercase(), + arch: arch.to_ascii_lowercase(), + }), + _ => None, + } + } else { + None + }; + + Ok(RemoteReadiness { + cli, + cli_version, + rocm_present: rocm.success, + tailscale_present: tailscale.success, + platform, + }) +} + +/// Find the CLI on the remote without installing anything. +/// +/// For read-only callers. `ensure_ready` provisions a missing CLI, which is +/// right when the user asked to serve on the machine and wrong when they only +/// asked a question about it: checking a machine's health must not change it, +/// and the assistant path treats that check as read-only precisely because it +/// is supposed to be. +pub(crate) fn locate_cli(transport: &dyn Transport, target: &str) -> Result { + let readiness = probe(transport)?; + readiness.cli.ok_or_else(|| { + anyhow::anyhow!( + "{target} has no ROCm CLI, so there is nothing there to report its state.\n\ + Installing one is a change, and this command only reads. Put it there with:\n \ + ssh {target} -- 'curl -fsSL https://raw.githubusercontent.com/ROCm/rocm-cli/main/install.sh | sh'\n\ + Or run `rocm remote serve`, which provisions it as part of starting a model." + ) + }) +} + +/// Confirm the remote can host a model, or explain what is missing, and install +/// ROCm when explicitly asked. +/// +/// Returns how to invoke the CLI there, provisioning one if the machine has +/// none — which is why only the serve path may call this, and why the read-only +/// health check uses [`locate_cli`] instead. +/// +/// Installing ROCm is opt-in and never implied. It can run for minutes, may +/// need a reboot, and is happening on a machine nobody is looking at — so the +/// default stays "tell the user what is missing" rather than "fix it while they +/// wait". +pub(crate) fn ensure_ready_with( + transport: &dyn Transport, + target: &str, + channel: &str, + install_rocm: bool, +) -> Result { + let readiness = probe(transport)?; + + // Checked before anything else runs, including a ROCm install below — not + // just before the model starts. The endpoint is published by the remote's + // own Tailscale; without it, a `--install-rocm` run would still install + // ROCm (minutes, possibly a reboot), mint and store a credential, and start + // a model, only to fail at `publish::publish` afterwards. Hoisted above the + // `rocm_present` branch so it guards both paths, not just the one where + // ROCm is already there. + if !readiness.tailscale_present { + bail!( + "{target} has no Tailscale, so it cannot publish an endpoint.\n\ + `rocm remote` reaches a model over the tailnet, not through this machine.\n\ + Install Tailscale there and run `tailscale up`, then try again." + ); + } + + if !readiness.rocm_present { + if !install_rocm { + bail!( + "no ROCm installation was found on {target}.\n\ + Install it there first:\n \ + ssh {target} -- rocm bootstrap setup\n\ + Or pass --install-rocm to have this command do it." + ); + } + // Installing ROCm provisions the CLI first, since the CLI is what runs + // the install. Return the invocation it settled on rather than falling + // through to the readiness snapshot below, which was taken *before* any + // of that and still says there is no CLI — reading it here provisioned + // the machine a second time, and a hiccup during that redundant install + // failed the whole command after ROCm was already in place. + return install_rocm_on(transport, target, channel, &readiness); + } + + match readiness.cli { + Some(cli) => { + match readiness.cli_version.as_deref() { + Some(version) => println!(" remote CLI: {version}"), + None => println!(" remote CLI: present"), + } + Ok(cli) + } + // A missing CLI is a job, not a refusal: unlike ROCm it is one small + // artifact, and the machine can usually fetch its own. + None => { + super::provision::install_cli(transport, target, readiness.platform.as_ref(), channel) + .map(|(cli, _)| cli) + } + } +} + +/// Install ROCm on the remote, once the CLI is there to do it with. +/// +/// The CLI has to come first: it is what runs the install, and it is the +/// smaller, safer artifact of the two. +fn install_rocm_on( + transport: &dyn Transport, + target: &str, + channel: &str, + readiness: &RemoteReadiness, +) -> Result { + let remote_cli = match &readiness.cli { + Some(cli) => cli.clone(), + None => { + super::provision::install_cli(transport, target, readiness.platform.as_ref(), channel)? + .0 + } + }; + + // Look before installing. The catalog knows the states that need a person, + // and the wizard that walks a person through them cannot run over a + // connection with nobody watching. + let (_, report) = super::doctor::examine_remote(transport, &remote_cli, None)?; + let passwordless_sudo = super::install::has_passwordless_sudo(transport)?; + if let Some(refusal) = super::install::assess(&report, passwordless_sudo) { + bail!( + "{}", + super::install::describe_refusal(&refusal, target, &report, 5) + ); + } + + super::install::install(transport, target, &remote_cli)?; + Ok(remote_cli) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::remote::transport::{ScriptedStep, ScriptedTransport}; + + /// A machine with everything present. + fn ready_steps() -> Vec { + vec![ + ScriptedStep::ok("rocm --version", "rocm 1.2.3"), + ScriptedStep::ok("command -v rocminfo", ""), + ScriptedStep::ok("command -v tailscale", ""), + ScriptedStep::ok("uname -s", "Linux\nx86_64\n"), + ] + } + + #[test] + fn a_ready_machine_reports_its_cli_rocm_tailscale_and_platform() { + let transport = ScriptedTransport::new(ready_steps()); + let readiness = probe(&transport).expect("probe"); + + assert_eq!(readiness.cli.as_deref(), Some("rocm")); + assert_eq!(readiness.cli_version.as_deref(), Some("rocm 1.2.3")); + assert!(readiness.rocm_present); + assert!(readiness.tailscale_present); + assert_eq!( + readiness.platform, + Some(RemotePlatform { + os: "linux".to_owned(), + arch: "x86_64".to_owned() + }) + ); + } + + #[test] + fn a_cli_outside_path_is_found_where_we_would_have_installed_it() { + // A non-interactive shell often has no ~/.local/bin on PATH. Missing + // this would reinstall over a perfectly good CLI on every run. + let transport = ScriptedTransport::new(vec![ + ScriptedStep::fails("rocm --version", 127, "not found"), + ScriptedStep::ok("/.local/bin/rocm --version", "rocm 1.2.3"), + ScriptedStep::ok("command -v rocminfo", ""), + ScriptedStep::ok("command -v tailscale", ""), + ScriptedStep::ok("uname -s", "Linux\nx86_64\n"), + ]); + let readiness = probe(&transport).expect("probe"); + assert_eq!(readiness.cli.as_deref(), Some(REMOTE_CLI_PATH)); + } + + #[test] + fn a_machine_without_rocm_is_refused_with_the_command_that_fixes_it() { + let transport = ScriptedTransport::new(vec![ + ScriptedStep::ok("rocm --version", "rocm 1.2.3"), + ScriptedStep::fails("command -v rocminfo", 1, ""), + ScriptedStep::ok("command -v tailscale", ""), + ScriptedStep::ok("uname -s", "Linux\nx86_64\n"), + ]); + let error = ensure_ready_with(&transport, "gpu-box", "release", false) + .unwrap_err() + .to_string(); + assert!(error.contains("no ROCm installation"), "{error}"); + assert!(error.contains("ssh gpu-box --"), "{error}"); + } + + #[test] + fn a_machine_without_tailscale_is_refused_before_a_model_is_started() { + // Discovering this after starting the server leaves a process on + // someone's GPU that nobody can reach and nothing records. + let transport = ScriptedTransport::new(vec![ + ScriptedStep::ok("rocm --version", "rocm 1.2.3"), + ScriptedStep::ok("command -v rocminfo", ""), + ScriptedStep::fails("command -v tailscale", 1, ""), + ScriptedStep::ok("uname -s", "Linux\nx86_64\n"), + ]); + let error = ensure_ready_with(&transport, "gpu-box", "release", false) + .unwrap_err() + .to_string(); + assert!(error.contains("cannot publish an endpoint"), "{error}"); + } + + #[test] + fn install_rocm_does_not_skip_the_tailscale_check() { + // On a machine with neither ROCm nor Tailscale, `--install-rocm` must + // still refuse here rather than running a multi-minute ROCm install, + // minting a credential, and starting a model, only to fail afterwards + // at `publish::publish`. + let transport = ScriptedTransport::new(vec![ + ScriptedStep::fails("rocm --version", 127, "not found"), + ScriptedStep::fails("command -v rocminfo", 1, ""), + ScriptedStep::fails("command -v tailscale", 1, ""), + ScriptedStep::ok("uname -s", "Linux\nx86_64\n"), + ]); + let error = ensure_ready_with(&transport, "gpu-box", "release", true) + .unwrap_err() + .to_string(); + assert!(error.contains("cannot publish an endpoint"), "{error}"); + assert!( + !transport.calls().iter().any(|call| matches!( + call, + crate::remote::transport::TransportCall::Exec { command, .. } + if command.contains("install.sh") + )), + "ROCm install must not run before the Tailscale check: {:?}", + transport.calls() + ); + } + + /// What `rocm examine --json` prints on a host the failure catalog has no + /// entries for. + /// + /// Built by serializing a real [`rocm_core::Examination`] with its + /// `os_family` changed, so the fixture cannot drift out of step with the + /// struct the remote side deserializes into — the same reason + /// `doctor::tests::examine_json` is built that way rather than written out. + fn out_of_scope_examine_json() -> String { + let mut examination = rocm_core::Examination::probe(rocm_core::FrameworkProbe::Skip); + examination.is_wsl = false; + examination.os_family = "darwin".to_owned(); + serde_json::to_string(&examination).expect("an examination serializes") + } + + #[test] + fn install_rocm_stops_on_a_machine_the_catalog_never_scored() { + // `--install-rocm` is gated on the failure catalog having looked at the + // machine. On a platform it carries no entries for it looks at nothing + // and says so, and an empty findings list then reads as "nothing else is + // wrong" — which would run a multi-minute privileged install, unattended, + // on a machine no check has ever been applied to. + // + // The install steps are programmed deliberately. Leaving them out would + // make the scripted transport error on the first one, and this test would + // then pass for the wrong reason on a machine that had no guard at all. + let transport = ScriptedTransport::new(vec![ + ScriptedStep::ok("rocm --version", "rocm 1.2.3"), + ScriptedStep::fails("command -v rocminfo", 1, ""), + ScriptedStep::ok("command -v tailscale", ""), + ScriptedStep::ok("uname -s", "Darwin\narm64\n"), + ScriptedStep::ok("examine --json", &out_of_scope_examine_json()), + ScriptedStep::ok("sudo -n true", ""), + ScriptedStep::ok("install driver --yes", "status: completed"), + ScriptedStep::ok("install sdk --yes", "status: completed"), + ]); + + let error = ensure_ready_with(&transport, "mac-mini", "release", true) + .expect_err("a platform the catalog never scored must not be installed on"); + assert!( + format!("{error:#}").contains("nothing was checked"), + "{error:#}" + ); + assert!( + !transport.calls().iter().any(|call| matches!( + call, + crate::remote::transport::TransportCall::Exec { command, .. } + if command.contains("install driver") + )), + "nothing may be installed on a machine the catalog never evaluated: {:?}", + transport.calls() + ); + } + + #[test] + fn a_machine_without_the_cli_has_one_installed_rather_than_being_refused() { + // Unlike ROCm, the CLI is one small artifact the machine can usually + // fetch itself, so its absence is a job rather than a dead end. + let transport = ScriptedTransport::new(vec![ + ScriptedStep::fails("rocm --version", 127, "not found"), + ScriptedStep::ok("command -v rocminfo", ""), + ScriptedStep::ok("command -v tailscale", ""), + ScriptedStep::ok("uname -s", "Linux\nx86_64\n"), + ScriptedStep::ok("install.sh | sh", ""), + ScriptedStep::ok("/.local/bin/rocm --version", "rocm 1.2.3"), + ]); + assert_eq!( + ensure_ready_with(&transport, "gpu-box", "release", false).expect("provisioned"), + REMOTE_CLI_PATH + ); + } + + #[test] + fn a_ready_machine_returns_how_to_invoke_its_cli() { + let transport = ScriptedTransport::new(ready_steps()); + assert_eq!( + ensure_ready_with(&transport, "gpu-box", "release", false).expect("ready"), + "rocm" + ); + } +} diff --git a/apps/rocm/src/remote/doctor.rs b/apps/rocm/src/remote/doctor.rs new file mode 100644 index 000000000..06ca50bad --- /dev/null +++ b/apps/rocm/src/remote/doctor.rs @@ -0,0 +1,716 @@ +// Copyright © Advanced Micro Devices, Inc., or its affiliates. +// +// SPDX-License-Identifier: MIT + +//! Checking a remote machine's health from here. +//! +//! This needs almost no new logic, because of how the local checks are already +//! built: gathering facts about a machine produces a plain serializable +//! snapshot, and scoring that snapshot against the failure-mode catalog reads +//! nothing but the snapshot. Neither half touches the local filesystem while +//! deciding. So the fetch happens on the remote and the scoring happens here, +//! with the same catalog the local command uses and no remote-side code at all. +//! +//! The snapshot is the contract, deliberately — not the human report, which +//! mixes in local paths, cache directories and engine inventory that describe +//! whichever machine rendered it. Reading that from a remote and printing it +//! here would produce a report that is subtly about the wrong computer. +//! +//! Fixes are rewritten to name the target. A command that repairs a machine you +//! are not sitting at is not a command you can paste, and printing it bare +//! invites running it against your own. + +use std::fmt::Write as _; + +use anyhow::{Context, Result}; +use rocm_core::Examination; +use rocm_core::diagnose::{DiagnoseReport, diagnose}; + +use super::transport::Transport; + +/// Fetch the remote's own view of itself and score it here. +pub(crate) fn examine_remote( + transport: &dyn Transport, + remote_cli: &str, + symptom: Option<&str>, +) -> Result<(Examination, DiagnoseReport)> { + let json = transport + .run(&format!("{remote_cli} examine --json")) + .context("could not read the remote machine's system state")?; + let examination = parse_examination(&json)?; + let report = diagnose(&examination, symptom.unwrap_or_default()); + Ok((examination, report)) +} + +/// Read an examination out of what the remote printed. +/// +/// The remote wraps its examination in a document carrying extra rendering +/// fields; ignoring what we do not recognise is what lets a remote on a +/// different CLI version still be understood. A field we *do* need being absent +/// is the opposite case, and says so. +fn parse_examination(json: &str) -> Result { + serde_json::from_str::(json.trim()).context( + "could not understand the remote machine's system state. The remote CLI is \ + probably a different version than this one — update whichever is older.", + ) +} + +/// Render the findings, with every fix aimed at the machine they are about. +/// +/// The aiming happens on the *report*, before rendering, not on the rendered +/// text afterwards. A `Fix.commands` entry is one command, but it is not one +/// line: the catalog holds backslash-continued `pip install` and `docker run` +/// entries spanning up to seven lines apiece. Rewriting line by line cut those +/// at the first newline and left the remainder as a bare local command — for the +/// `pip` entry, that silently dropped the `--index-url` and installed the wrong +/// wheels. Rewriting per entry keeps a command whole by construction. +pub(crate) fn render_report(target: &str, report: &DiagnoseReport, top: usize) -> String { + let local = rocm_core::diagnose::render_report_text(&aim_at(target, report), top); + let mut output = format!("Health of {target}\n\n"); + output.push_str(&redirect_quoted_commands( + &redirect_apply_with(&local, target), + target, + )); + output +} + +/// Prefix of the one runnable command the renderer synthesises itself. +/// +/// Everything else under a fix comes from a `Fix` field that [`aim_at`] can +/// rewrite before rendering. This line does not: the renderer builds it from +/// `fix_id` (` apply with: rocm fix `), so the only place to aim it is +/// after the fact. Kept deliberately narrow — one prefix, one synthesised shape +/// — rather than the general line-matching this module used to do, which is what +/// cut multi-line catalog commands in half. +const APPLY_WITH: &str = "apply with: "; + +/// Aim the renderer's own `apply with:` line at the target machine. +fn redirect_apply_with(rendered: &str, target: &str) -> String { + let mut output = rendered + .lines() + .map(|line| { + let trimmed = line.trim_start(); + let Some(command) = trimmed.strip_prefix(APPLY_WITH) else { + return line.to_owned(); + }; + let command = command.trim(); + if command.is_empty() { + return line.to_owned(); + } + let lead = &line[..line.len() - trimmed.len() + APPLY_WITH.len()]; + format!("{lead}{}", remote_invocation(target, command)) + }) + .collect::>() + .join("\n"); + output.push('\n'); + output +} + +/// Copy `report` with every runnable command aimed at `target`. +fn aim_at(target: &str, report: &DiagnoseReport) -> DiagnoseReport { + let mut aimed = report.clone(); + for diagnosis in &mut aimed.matched { + let Some(fix) = diagnosis.fix.as_mut() else { + continue; + }; + // A sequence that only works as a sequence cannot be split across + // separate ssh invocations, each with its own shell and no shared state. + // Rewriting it would hand the user commands that look runnable and are + // not, which is worse than saying so. + if is_stateful_sequence(fix) { + fix.notes.push(format!( + "These steps share one shell session, so they cannot be run over separate \ + connections. Open a session first with `ssh {target}`, then run them there." + )); + continue; + } + fix.commands = join_continuations(&fix.commands) + .into_iter() + .map(|command| { + if is_comment(&command) { + command + } else { + remote_invocation(target, &command) + } + }) + .collect(); + if !fix.verify.is_empty() { + fix.verify = remote_invocation(target, &fix.verify); + } + } + aimed +} + +/// Merge backslash-continued entries into the single command they spell. +/// +/// The catalog stores a multi-line command in two different shapes, and they are +/// not interchangeable. `fix-1-arch`'s `pip install` is **one** entry with an +/// embedded newline. `fix-10-container`'s `docker run` is **seven** entries, +/// each ending in a trailing `\`, meant to read as one continued invocation. +/// +/// Rewriting per entry is right for the first shape and wrong for the second: it +/// turns one `docker run` into seven independent `ssh` calls, each carrying a +/// dangling backslash inside its quoting. Joining first makes both shapes the +/// same thing — one entry holding the whole command — before anything rewrites +/// it, and the far shell then reads the continuation exactly as a local one +/// would. +/// +/// A line that does not end in `\` closes the group, so an entry that stands +/// alone passes through untouched. +fn join_continuations(commands: &[String]) -> Vec { + let mut joined: Vec = Vec::new(); + let mut pending: Option = None; + + for command in commands { + let continues = command.trim_end().ends_with('\\'); + match pending.as_mut() { + Some(open) => { + open.push('\n'); + open.push_str(command); + } + None => pending = Some(command.clone()), + } + if !continues && let Some(complete) = pending.take() { + joined.push(complete); + } + } + // A trailing `\` with nothing after it is malformed, but dropping the text + // would be worse than emitting it as it stands. + if let Some(unterminated) = pending { + joined.push(unterminated); + } + joined +} + +/// Whether a catalog entry is prose rather than something to run. +/// +/// The renderer prints every `commands` element behind `$ `, including the +/// `# Recommended: …` annotations the catalog uses to explain a choice. Wrapping +/// one in `ssh … -- '# …'` presents a comment as a command to run. +fn is_comment(command: &str) -> bool { + command.trim_start().starts_with('#') +} + +/// Whether a fix's steps depend on each other's shell state. +/// +/// The catalog has one such entry today — a three-stage fix whose later steps +/// run *inside* a subshell the first one opens — and it labels its own steps. +/// Matching on that label rather than guessing at the commands keeps this +/// honest: an entry that stops saying so stops being treated as one. +fn is_stateful_sequence(fix: &rocm_core::diagnose::Fix) -> bool { + fix.commands + .iter() + .any(|command| command.contains("step 2 of") || command.contains("INSIDE the subshell")) +} + +/// The command a user should paste to run `command` on `target`. +/// +/// The quoting is load-bearing. A suggested fix is copied into the user's own +/// shell, and an interpolated command splits across the two machines at its +/// first metacharacter: `ssh box -- echo x | sudo tee /etc/f` runs the `echo` +/// there and the privileged `tee` **here**, rewriting the operator's own system +/// while they believe they are repairing someone else's. +/// +/// That is the common case, not an exotic one. Most commands in the fix catalog +/// contain a pipe, a `&&` or a redirect — `rocminfo | grep …`, `sudo apt update +/// 2>&1 | tail …`, `echo … | sudo tee …`. +/// +/// One layer of quoting, not two, and no `sh -c`: the local shell strips the +/// quotes, `ssh` sends what is left as the command, and the *remote* login shell +/// is what parses the pipeline — on the machine it is meant to run on. Wrapping +/// it in a second layer would only make an already-correct line unreadable, and +/// these lines exist to be read and pasted. +fn remote_invocation(target: &str, command: &str) -> String { + format!("ssh {target} -- {}", super::shell_quote(command)) +} + +/// Rewrite every backtick-quoted `rocm …` so it names the target machine. +/// +/// Matched by the quoting rather than by listing each phrase the renderer might +/// wrap around one. The phrases have changed before and will again; what does +/// not change is that a command a user is meant to run is put in backticks. +fn redirect_quoted_commands(rendered: &str, target: &str) -> String { + let mut output = String::with_capacity(rendered.len()); + let mut rest = rendered; + + while let Some(open) = rest.find('`') { + let (before, from_open) = rest.split_at(open); + output.push_str(before); + let after_open = &from_open[1..]; + let Some(close) = after_open.find('`') else { + // An unpaired backtick is prose, not a quote. Leave the remainder be. + output.push_str(from_open); + return output; + }; + let (quoted, remainder) = after_open.split_at(close); + if quoted.starts_with("rocm ") { + // Quoted the same way as a `$ ` line: these carry `#` comments and + // shell operators too, and a backticked command is copied just as + // readily as one on its own line. + let _ = write!(output, "`{}`", remote_invocation(target, quoted)); + } else { + let _ = write!(output, "`{quoted}`"); + } + rest = &remainder[1..]; + } + output.push_str(rest); + output +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::remote::transport::{ScriptedStep, ScriptedTransport}; + + /// The shape `rocm examine --json` prints: an examination, flattened + /// together with a rendering summary that belongs to whoever printed it. + fn examine_json() -> String { + let examination = Examination::probe(rocm_core::FrameworkProbe::Skip); + let mut value = serde_json::to_value(&examination).unwrap(); + value.as_object_mut().unwrap().insert( + "summary".to_owned(), + serde_json::json!({"default_engine": "vllm"}), + ); + serde_json::to_string(&value).unwrap() + } + + #[test] + fn a_remote_examination_is_read_through_the_snapshot_not_the_human_report() { + // The extra rendering fields describe the machine that printed them, so + // they are ignored rather than adopted. + let transport = + ScriptedTransport::new(vec![ScriptedStep::ok("examine --json", &examine_json())]); + let (examination, _) = examine_remote(&transport, "rocm", None).expect("examined"); + // Round-tripping the snapshot is the contract; the value itself is + // whatever this machine happens to be. + assert!(!examination.os_family.is_empty()); + } + + #[test] + fn a_remote_on_another_version_says_so_instead_of_failing_obscurely() { + let error = parse_examination(r#"{"unexpected": true}"#) + .unwrap_err() + .to_string(); + assert!(error.contains("different version"), "{error}"); + } + + #[test] + fn the_container_fixture_is_something_this_code_can_actually_read() { + // The stub that stands in for a remote CLI has to answer with a document + // this deserializer accepts. An earlier version returned a two-field + // fragment that looked plausible and could never have parsed — the + // container lane never noticed, because nothing there ran this code. + // Pinning it here means a drift in either direction fails a unit test. + let fixture = include_str!("../../../../tests/remote-ssh/examination.json"); + let parsed = + parse_examination(fixture).expect("the container stub's examination must deserialize"); + assert_eq!(parsed.os_family, "linux"); + } + + #[test] + fn a_newer_remote_adding_fields_is_still_understood() { + // Version skew between the machine driving and the machine driven is + // normal; unknown fields must not break the read. + let mut value: serde_json::Value = serde_json::from_str(&examine_json()).unwrap(); + value + .as_object_mut() + .unwrap() + .insert("added_in_a_later_release".to_owned(), serde_json::json!(1)); + assert!(parse_examination(&value.to_string()).is_ok()); + } + + /// A report the renderer will render fully, built without asking this + /// machine anything. + /// + /// Deliberately not `diagnose(&Examination::probe(..))`: on a host the + /// catalog considers out of scope — WSL2, for one — that returns an + /// out-of-scope report and the renderer short-circuits before printing a + /// single command. Tests built that way pass or fail depending on the + /// developer's machine, which is the opposite of what these need to prove. + fn report(matched: Vec) -> DiagnoseReport { + DiagnoseReport { + has_match: matched.iter().any(|d| d.score >= 50), + matched, + min_score_for_match: rocm_core::diagnose::MIN_SCORE_FOR_MATCH, + high_confidence_threshold: rocm_core::diagnose::HIGH_CONFIDENCE, + route_when_no_match: rocm_core::diagnose::Route { + target: "rocm-cli".to_owned(), + url: "https://example.invalid/issues".to_owned(), + }, + out_of_scope: None, + } + } + + /// A report containing a fix, in the shape the renderer prints in full. + /// + /// The previous version of these tests invented an output shape the renderer + /// never produces, so they passed while the rewriting matched nothing at all. + /// + /// The commands are taken verbatim from the real catalog rather than + /// invented, for the same class of reason one step in: the invented ones were + /// all single words with no shell operators, while most catalog entries + /// contain a pipe, a `&&` or a redirect. A fixture that cannot express the + /// failing input cannot fail, which is how the unquoted rewriting survived + /// five review rounds. + fn real_report_with_a_fix() -> DiagnoseReport { + use rocm_core::diagnose::{Diagnosis, Fix}; + report(vec![Diagnosis { + id: "dkms-mismatch".to_owned(), + title: "DKMS built against another kernel".to_owned(), + score: 90, + evidence: vec!["dkms status reports a stale build".to_owned()], + fix: Some(Fix { + summary: "rebuild the module".to_owned(), + commands: vec![ + "sudo dkms autoinstall".to_owned(), + // `fix-wsl-2-dxcore-missing`, verbatim. The privileged half + // is behind the pipe, which is what made the unquoted form + // rewrite the operator's own machine. + "echo /usr/lib/wsl/lib | sudo tee /etc/ld.so.conf.d/wsl.conf".to_owned(), + ], + needs_sudo: true, + needs_reboot: true, + fix_id: "dkms-mismatch".to_owned(), + verify: "lsmod | grep amdgpu && rocminfo | head -n 5".to_owned(), + ..Fix::default() + }), + }]) + } + + /// Paste `line` into a real shell, with `ssh` replaced by a stub, and report + /// the command text `ssh` was handed. + /// + /// Asking a shell rather than inspecting the string is the whole point. What + /// decides which machine a command runs on is what the *local* shell does + /// with it before `ssh` ever sees it, and no assertion about the rendered + /// text can observe that. + fn what_ssh_would_send(line: &str) -> String { + // Drop the destination and the `--` guard; what remains is the command + // ssh transmits. `$*` rejoins it exactly as ssh does. + let stub = r#"ssh() { shift 2; printf %s "$*"; }"#; + let output = std::process::Command::new("sh") + .arg("-c") + .arg(format!("{stub}\n{line}")) + .output() + .expect("sh should run"); + String::from_utf8_lossy(&output.stdout).into_owned() + } + + #[test] + fn a_piped_fix_command_is_sent_whole_rather_than_split_across_two_machines() { + // The defect this replaces: `ssh box -- echo x | sudo tee /etc/f` runs + // the echo on the remote and the privileged tee *locally*. Nothing about + // the rendered string reveals that — only running it through a shell. + const PIPED: &str = "echo /usr/lib/wsl/lib | sudo tee /etc/ld.so.conf.d/wsl.conf"; + let rendered = render_report("gpu-box", &real_report_with_a_fix(), 5); + let line = rendered + .lines() + .map(str::trim) + .find(|line| line.starts_with("$ ") && line.contains("ld.so.conf.d")) + .expect("the piped fix command must still be rendered") + .trim_start_matches("$ ") + .to_owned(); + + assert_eq!( + what_ssh_would_send(&line), + PIPED, + "ssh must be handed the whole pipeline; anything missing here is a \ + command the local machine ran instead: {line}" + ); + + // The same property for the `verify after fix:` line, which carries `&&` + // and a second pipe and is copied just as readily. + let verify = rendered + .lines() + .map(str::trim) + .find(|line| line.starts_with("verify after fix: ")) + .expect("the verify command must still be rendered") + .trim_start_matches("verify after fix: ") + .to_owned(); + assert_eq!( + what_ssh_would_send(&verify), + "lsmod | grep amdgpu && rocminfo | head -n 5", + "the verify command split across two machines: {verify}" + ); + } + + #[test] + fn every_command_the_real_renderer_emits_is_aimed_at_the_remote() { + let report = real_report_with_a_fix(); + let rendered = render_report("gpu-box", &report, 5); + + // The three shapes the renderer actually produces: a fix command, the + // handle that applies it, and the check to run afterwards. Each is a + // command a user would otherwise paste into their own terminal. + assert!( + rendered.contains("$ ssh gpu-box -- 'sudo dkms autoinstall'"), + "fix command not redirected:\n{rendered}" + ); + assert!( + rendered.contains("apply with: ssh gpu-box -- 'rocm fix dkms-mismatch'"), + "apply-with command not redirected:\n{rendered}" + ); + assert!( + rendered.contains( + "verify after fix: ssh gpu-box -- 'lsmod | grep amdgpu && rocminfo | head -n 5'" + ), + "verify command not redirected:\n{rendered}" + ); + } + + #[test] + fn the_rewriting_is_not_silently_a_no_op() { + // The failure this guards against is the one that already happened: the + // renderer's layout and this module's expectations drifted apart, and + // nothing noticed because the tests supplied their own input. If the + // renderer stops emitting these prefixes, this fails. + let local = rocm_core::diagnose::render_report_text(&real_report_with_a_fix(), 5); + let redirected = render_report("gpu-box", &real_report_with_a_fix(), 5); + assert!( + !redirected.contains(&format!( + "$ {}", + local + .lines() + .find_map(|line| line.trim().strip_prefix("$ ")) + .expect("a real report emits at least one command") + )), + "a command in a real report was left aimed at this machine:\n{redirected}" + ); + } + + #[test] + fn prose_that_opens_with_a_command_name_is_left_intact() { + // The report's own headers start with `rocm diagnose: …`. A rule that + // recognised commands by their first word would turn each into + // `ssh gpu-box -- rocm diagnose: no known misconfiguration matched.` + let rendered = render_report("gpu-box", &report(vec![]), 5); + assert!( + !rendered.contains("ssh gpu-box -- rocm diagnose:"), + "a sentence was rewritten as a command:\n{rendered}" + ); + } + + #[test] + fn a_command_quoted_inside_a_sentence_is_redirected_too() { + // The renderer closes with "Next step: run `rocm fix `." — no prefix, + // so a line-prefix rule misses it entirely, and it is exactly the line a + // user acts on. + let rendered = render_report("gpu-box", &real_report_with_a_fix(), 5); + assert!( + rendered.contains("run `ssh gpu-box -- 'rocm fix dkms-mismatch'`"), + "the closing instruction still points at the local machine:\n{rendered}" + ); + assert!( + !rendered.contains("run `rocm fix"), + "no bare local command should remain:\n{rendered}" + ); + } + + #[test] + fn a_below_threshold_report_redirects_its_closing_advice_as_well() { + // The other trailing branch, reached when nothing clears the confidence + // threshold. It names a command too. + let mut low = real_report_with_a_fix(); + low.matched[0].score = 40; + low.has_match = false; + let rendered = render_report("gpu-box", &low, 5); + assert!( + !rendered.contains("run `rocm fix"), + "the low-confidence branch still points locally:\n{rendered}" + ); + } + + #[test] + fn ordinary_quoted_text_is_left_alone() { + // Only commands get redirected; backticks around anything else stay put. + let left = redirect_quoted_commands("see the `apply with:` line and `/dev/kfd`", "gpu-box"); + assert_eq!(left, "see the `apply with:` line and `/dev/kfd`"); + // An unpaired backtick is prose, not a quote. + assert_eq!( + redirect_quoted_commands("a ` stray tick", "gpu-box"), + "a ` stray tick" + ); + } + + #[test] + fn the_quoted_examine_command_asks_for_the_remote_machines_state() { + // Run locally it reports the wrong computer, and the user never learns + // why the answer looked irrelevant. + let rendered = render_report("gpu-box", &report(vec![]), 5); + assert!( + rendered.contains("`ssh gpu-box -- 'rocm examine --json'`"), + "{rendered}" + ); + } + + #[test] + fn indentation_survives_so_the_report_still_reads_as_one() { + let rendered = render_report("gpu-box", &real_report_with_a_fix(), 5); + assert!( + rendered + .lines() + .any(|line| line.starts_with(" $ ssh gpu-box -- ")), + "{rendered}" + ); + } + + #[test] + fn a_comment_in_a_fix_is_not_presented_as_a_command_to_run() { + // The catalog annotates its steps with `#` lines, and the renderer puts + // every element behind `$ `. Wrapping one in ssh tells the user to run a + // comment on another machine. + use rocm_core::diagnose::{Diagnosis, Fix}; + let rendered = render_report( + "gpu-box", + &report(vec![Diagnosis { + id: "annotated".to_owned(), + title: "a fix that explains itself".to_owned(), + score: 90, + evidence: vec!["something".to_owned()], + fix: Some(Fix { + summary: "do the thing".to_owned(), + commands: vec![ + "# Recommended: the nightly wheels".to_owned(), + "sudo dkms autoinstall".to_owned(), + ], + fix_id: "annotated".to_owned(), + ..Fix::default() + }), + }]), + 5, + ); + assert!( + rendered.contains("$ # Recommended: the nightly wheels"), + "the comment should be left as prose:\n{rendered}" + ); + assert!( + !rendered.contains("ssh gpu-box -- '#"), + "a comment was presented as a command:\n{rendered}" + ); + assert!( + rendered.contains("$ ssh gpu-box -- 'sudo dkms autoinstall'"), + "the real command must still be aimed at the remote:\n{rendered}" + ); + } + + #[test] + fn a_backslash_continued_command_stored_as_many_entries_becomes_one_invocation() { + // `fix-10-container` stores one `docker run` as seven entries, each + // ending in a trailing `\`. Rewriting per entry — correct for the + // *other* multi-line shape — turns it into seven independent ssh calls, + // each with a dangling backslash inside its quoting. + use rocm_core::diagnose::{Diagnosis, Fix}; + let rendered = render_report( + "gpu-box", + &report(vec![Diagnosis { + id: "container".to_owned(), + title: "container missing devices".to_owned(), + score: 90, + evidence: vec!["no /dev/kfd in the container".to_owned()], + fix: Some(Fix { + summary: "re-launch with the devices passed through".to_owned(), + commands: vec![ + "# Docker / Podman flags AMD-recommends:".to_owned(), + "docker run --rm -it \\".to_owned(), + " --device=/dev/kfd \\".to_owned(), + " --group-add render \\".to_owned(), + " rocm/pytorch:latest".to_owned(), + ], + fix_id: "container".to_owned(), + ..Fix::default() + }), + }]), + 5, + ); + + // Exactly one rewritten command line for the whole docker run, not four. + // Counted over `$ `-prefixed lines only, so the `apply with:` and the + // closing "Next step" line — both legitimately rewritten — do not count. + assert_eq!( + rendered + .lines() + .filter(|line| line.trim_start().starts_with("$ ssh gpu-box --")) + .count(), + 1, + "the docker run must be one ssh call, not one per entry:\n{rendered}" + ); + // The comment stays prose. + assert!(rendered.contains("$ # Docker / Podman flags"), "{rendered}"); + // And the far shell receives the whole continued command as one piece. + let line = rendered + .lines() + .map(str::trim_start) + .find(|line| line.starts_with("$ ssh gpu-box -- 'docker run")) + .expect("the docker run must be rewritten as a single command") + .trim_start_matches("$ ") + .to_owned(); + // The remaining fragments follow inside the same quoting, so the line + // the report prints is only the first physical line of one command. + let whole = rendered + .split_once("ssh gpu-box -- '") + .expect("rewritten") + .1 + .split_once("'\n") + .map_or_else(String::new, |(body, _)| body.to_owned()); + for fragment in [ + "--device=/dev/kfd", + "--group-add render", + "rocm/pytorch:latest", + ] { + assert!( + whole.contains(fragment), + "`{fragment}` was split into its own ssh call:\n{rendered}" + ); + } + assert!(line.starts_with("ssh gpu-box -- 'docker run"), "{line}"); + } + + #[test] + fn a_multi_line_fix_command_is_kept_whole_rather_than_cut_at_its_first_newline() { + // `fix-1-arch` is one `commands` element spanning two lines via a + // backslash continuation. Rewriting line by line aimed the first half at + // the remote and left `--index-url …` behind as a bare local fragment — + // so the user installed wheels from the default index. + use rocm_core::diagnose::{Diagnosis, Fix}; + const CONTINUED: &str = + "pip install --pre torch \\\n --index-url https://example.invalid/rocm6.4"; + let rendered = render_report( + "gpu-box", + &report(vec![Diagnosis { + id: "wheels".to_owned(), + title: "wrong wheels".to_owned(), + score: 90, + evidence: vec!["something".to_owned()], + fix: Some(Fix { + summary: "reinstall".to_owned(), + commands: vec![CONTINUED.to_owned()], + fix_id: "wheels".to_owned(), + ..Fix::default() + }), + }]), + 5, + ); + // The index URL must be inside the quoting, not stranded after it. + let quoted = rendered + .split_once("ssh gpu-box -- '") + .expect("the command is aimed at the remote") + .1; + let body = quoted + .split_once("'\n") + .map_or(quoted, |(body, _)| body) + .to_owned(); + assert!( + body.contains("--index-url https://example.invalid/rocm6.4"), + "the continuation was cut off, so the remote would install the wrong \ + wheels:\n{rendered}" + ); + } + + #[test] + fn the_report_names_the_machine_it_describes() { + // Without this the output is indistinguishable from a local report, and + // acting on it means fixing the wrong computer. + let rendered = render_report("gpu-box", &report(vec![]), 5); + assert!(rendered.starts_with("Health of gpu-box"), "{rendered}"); + } +} diff --git a/apps/rocm/src/remote/install.rs b/apps/rocm/src/remote/install.rs new file mode 100644 index 000000000..8e83c387a --- /dev/null +++ b/apps/rocm/src/remote/install.rs @@ -0,0 +1,417 @@ +// Copyright © Advanced Micro Devices, Inc., or its affiliates. +// +// SPDX-License-Identifier: MIT + +//! Installing ROCm on a machine nobody is watching. +//! +//! This is the riskiest thing `rocm remote` can do, and the parts that make it +//! safe are refusals rather than capability. The install itself already exists +//! and is already non-interactive; what is new here is deciding when it is +//! allowed to run. +//! +//! Three guards, all of them about a human not being present: +//! +//! The machine's state is scored against the failure-mode catalog *first*. That +//! catalog exists because ROCm installs go wrong in specific, recognisable +//! ways — a half-configured driver, a DKMS build against the wrong kernel, +//! Secure Boot refusing an unsigned module, repo files left over from a +//! previous attempt. Each of those needs a person to decide, and the wizard +//! that walks a human through them cannot do so over SSH. So anything the +//! catalog recognises stops the install and prints what it found. Only a +//! machine that is plainly missing ROCm, with nothing else wrong, proceeds. +//! +//! And a machine the catalog could not score at all is refused too, which is a +//! different refusal and not a subtlety. On a platform the catalog carries no +//! entries for it runs nothing and returns no findings — the same empty result +//! a healthy machine produces. "Nothing was checked" and "nothing is wrong" are +//! opposite facts wearing one shape, and only reading them apart keeps the +//! gating above from being satisfied by a machine nobody looked at. +//! +//! And privilege escalation is checked before it is needed. The driver install +//! runs commands through `sudo`, while the control channel refuses to answer +//! prompts by design — so a machine asking for a password does not fail, it +//! hangs. Establishing that sudo is passwordless first turns the most likely +//! real-world failure into a sentence instead of a stall. + +use anyhow::{Result, bail}; +use rocm_core::diagnose::DiagnoseReport; + +use super::doctor; +use super::transport::Transport; + +/// Why an unattended install was refused. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) enum Refusal { + /// The catalog has no entries for this platform, so it scored nothing. + /// Carries the catalog's own explanation. + NotEvaluated { reason: String }, + /// The catalog recognised something needing a person. + NeedsAHuman { findings: Vec }, + /// Privileged commands would block on a password prompt. + NeedsPasswordlessSudo, +} + +/// Decide whether ROCm may be installed on this machine unattended. +/// +/// Split from the doing so the judgement is testable on its own — the part +/// worth being sure about is what gets refused, not what gets run. +pub(crate) fn assess(report: &DiagnoseReport, passwordless_sudo: bool) -> Option { + // Asked first, and deliberately not by reading `matched`. When the catalog + // has no entries for a platform it does not run and returns an *empty* + // `matched` — the same shape a machine with nothing wrong produces. The two + // are opposite facts: one is a verdict, the other is the absence of one, and + // only `out_of_scope` tells them apart. Gating on `matched` alone therefore + // reads "nothing was checked" as "nothing else is wrong", which is the one + // way to reach an unattended, privileged, multi-minute install on a machine + // this module's premise says must have been looked at first. + if let Some(reason) = &report.out_of_scope { + return Some(Refusal::NotEvaluated { + reason: reason.clone(), + }); + } + let findings = recognised_problems(report); + if !findings.is_empty() { + return Some(Refusal::NeedsAHuman { findings }); + } + if !passwordless_sudo { + return Some(Refusal::NeedsPasswordlessSudo); + } + None +} + +/// Catalog findings strong enough to be treated as real. +/// +/// Weak signals are excluded deliberately: several checks open with a low score +/// for a situation that is merely *potentially* relevant, and treating those as +/// blockers would refuse installs on healthy machines until nobody trusted the +/// refusal. +fn recognised_problems(report: &DiagnoseReport) -> Vec { + report + .matched + .iter() + .filter(|diagnosis| diagnosis.score >= rocm_core::diagnose::MIN_SCORE_FOR_MATCH) + .map(|diagnosis| diagnosis.title.clone()) + .collect() +} + +/// Can privileged commands run without a prompt nobody will answer? +pub(crate) fn has_passwordless_sudo(transport: &dyn Transport) -> Result { + // `-n` makes sudo fail rather than prompt, which is the whole question. + Ok(transport.exec("sudo -n true")?.success) +} + +/// Install ROCm on the remote, having decided it is safe to. +pub(crate) fn install(transport: &dyn Transport, target: &str, remote_cli: &str) -> Result<()> { + println!("Installing ROCm on {target}. This can take several minutes."); + + let driver = transport.exec(&format!("{remote_cli} install driver --yes"))?; + if !driver.success { + bail!( + "the driver install failed on {target}: {}\n\ + Nothing further was attempted. Check the machine with \ + `rocm remote doctor {target}`.", + driver.stderr.trim() + ); + } + print_indented(&driver.stdout); + + let sdk = transport.exec(&format!("{remote_cli} install sdk --yes"))?; + if !sdk.success { + bail!( + "the driver installed on {target} but the ROCm SDK did not: {}\n\ + The machine is part-way through a setup; check it with \ + `rocm remote doctor {target}` before retrying.", + sdk.stderr.trim() + ); + } + print_indented(&sdk.stdout); + + // The driver install records that a reboot is needed rather than performing + // one. Saying so matters more here than locally: nobody is sitting at this + // machine to notice it behaving as though the install did not take. + if mentions_reboot(&driver.stdout) || mentions_reboot(&sdk.stdout) { + println!(); + println!("{target} needs a reboot before it can serve."); + println!(" reboot it, then run: rocm remote doctor {target}"); + } + Ok(()) +} + +fn mentions_reboot(output: &str) -> bool { + output.lines().any(|line| { + let line = line.to_ascii_lowercase(); + line.contains("reboot_required: true") || line.contains("reboot required") + }) +} + +fn print_indented(output: &str) { + for line in output.lines().filter(|line| !line.trim().is_empty()) { + println!(" {line}"); + } +} + +/// Explain a refusal, including what to run instead. +pub(crate) fn describe_refusal( + refusal: &Refusal, + target: &str, + report: &DiagnoseReport, + top: usize, +) -> String { + match refusal { + // Worded so it cannot be read as a clean bill of health, which is the + // whole hazard: an empty findings list looks identical to a healthy + // machine's. The catalog's own sentence is carried too, because it names + // the platform the remote reported and this one does not. + Refusal::NotEvaluated { reason } => format!( + "{target} runs a platform the ROCm failure catalog has no entries for, so nothing \ + was checked there and there is no verdict to install on.\n\n\ + What the check reported:\n {reason}\n\n\ + An unattended install is only allowed on a machine the catalog looked at and \ + found nothing wrong with. This is a machine it never looked at, which is a \ + different thing. Set it up yourself if you are sure it is supported:\n \ + ssh {target} -- rocm bootstrap setup" + ), + Refusal::NeedsAHuman { findings } => format!( + "{target} is not in a state that can be set up unattended.\n\n\ + What was found:\n{}\n\n\ + These are the situations the setup wizard exists to walk a person through, \ + and it cannot do that over a connection with nobody watching. Resolve them \ + first — the suggested fixes are below — then run this again.\n\n{}", + findings + .iter() + .map(|finding| format!(" - {finding}")) + .collect::>() + .join("\n"), + doctor::render_report(target, report, top) + ), + Refusal::NeedsPasswordlessSudo => format!( + "installing ROCm on {target} needs administrator rights, and this connection \ + cannot answer a password prompt — it would hang rather than fail.\n\n\ + Either set up passwordless sudo there, or install it yourself:\n \ + ssh {target} -- rocm bootstrap setup" + ), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::remote::transport::{ScriptedStep, ScriptedTransport}; + use rocm_core::diagnose::{Diagnosis, MIN_SCORE_FOR_MATCH}; + + fn report_with(diagnoses: Vec) -> DiagnoseReport { + let mut report = rocm_core::diagnose::diagnose( + &rocm_core::Examination::probe(rocm_core::FrameworkProbe::Skip), + "", + ); + report.matched = diagnoses; + report + } + + fn finding(title: &str, score: i32) -> Diagnosis { + Diagnosis { + id: "x".to_owned(), + title: title.to_owned(), + score, + ..Diagnosis::default() + } + } + + /// An examination of a machine the catalog has no entries for. + /// + /// `platform_family` reads `os_family` (and `is_wsl`), and the catalog + /// carries checkers for `linux`, `windows` and `wsl` only — so anything else + /// is a host on which nothing is ever run. + fn out_of_scope_examination() -> rocm_core::Examination { + let mut examination = rocm_core::Examination::probe(rocm_core::FrameworkProbe::Skip); + examination.is_wsl = false; + examination.os_family = "darwin".to_owned(); + examination + } + + /// A report for such a machine, produced by the real [`rocm_core::diagnose`] + /// rather than hand-built. + /// + /// What makes a platform out of scope is the catalog's own coverage rule. + /// Constructing a `DiagnoseReport { out_of_scope: Some(..), .. }` by hand + /// would assert against this test's idea of that rule instead of against the + /// catalog's, and would keep passing if the two ever diverged. + fn out_of_scope_report() -> DiagnoseReport { + let report = rocm_core::diagnose::diagnose(&out_of_scope_examination(), ""); + assert!( + report.out_of_scope.is_some(), + "the fixture must actually be a platform the catalog declines to score" + ); + assert!( + report.matched.is_empty(), + "an out-of-scope report carries no findings, which is the whole hazard" + ); + report + } + + #[test] + fn a_machine_the_catalog_never_scored_is_not_installed_on_unattended() { + // `matched` is empty here not because the machine is healthy but because + // no checker was ever run against it. Reading that emptiness as "nothing + // else wrong" is the one way to reach an unattended install with no + // catalog verdict at all — which is what this module says gates it. + let report = out_of_scope_report(); + assert!( + assess(&report, true).is_some(), + "a platform the catalog never evaluated must not be installed on unattended" + ); + } + + #[test] + fn an_unscored_platform_says_nothing_was_checked_rather_than_naming_a_finding() { + // The refusal a person reads has to say which of the two it is: the + // catalog found something, or the catalog was never run. Reporting an + // empty findings list would read as "nothing is wrong, but no". + let report = out_of_scope_report(); + let refusal = + assess(&report, true).expect("an unevaluated platform must produce a refusal"); + let message = describe_refusal(&refusal, "mac-mini", &report, 5); + assert!( + message.contains("nothing was checked"), + "the refusal must say the catalog never ran: {message}" + ); + assert!( + message.contains("mac-mini"), + "and name the machine it is about: {message}" + ); + } + + #[test] + fn an_unscored_platform_outranks_a_sudo_problem() { + // Same reasoning as a recognised finding outranking one: fixing sudo + // would not make this machine installable, so leading with sudo sends + // someone down a path that ends nowhere. + let report = out_of_scope_report(); + assert_ne!( + assess(&report, false), + Some(Refusal::NeedsPasswordlessSudo), + "an unevaluated platform is not a sudo problem" + ); + } + + #[test] + fn a_clean_machine_with_working_sudo_may_be_installed() { + assert_eq!(assess(&report_with(vec![]), true), None); + } + + #[test] + fn a_recognised_problem_stops_the_install_and_names_itself() { + // These are exactly the cases the interactive wizard exists for, and it + // cannot walk anybody through them over a connection nobody is watching. + let report = report_with(vec![finding("DKMS built against another kernel", 90)]); + assert_eq!( + assess(&report, true), + Some(Refusal::NeedsAHuman { + findings: vec!["DKMS built against another kernel".to_owned()] + }) + ); + } + + #[test] + fn a_weak_signal_does_not_block_an_otherwise_healthy_machine() { + // Several checks open with a low score for something merely potentially + // relevant. Treating those as blockers would refuse healthy machines + // until nobody believed the refusal. + let report = report_with(vec![finding( + "possibly in a container", + MIN_SCORE_FOR_MATCH - 1, + )]); + assert_eq!(assess(&report, true), None); + } + + #[test] + fn a_machine_that_would_prompt_for_a_password_is_refused_not_hung() { + // The control channel never answers prompts, so without this the install + // stalls instead of failing — the worst outcome of the three. + assert_eq!( + assess(&report_with(vec![]), false), + Some(Refusal::NeedsPasswordlessSudo) + ); + } + + #[test] + fn a_recognised_problem_outranks_a_sudo_problem() { + // Fixing sudo would not make this machine installable, so saying so + // first would send someone down the wrong path. + let report = report_with(vec![finding("Secure Boot is blocking the module", 80)]); + assert!(matches!( + assess(&report, false), + Some(Refusal::NeedsAHuman { .. }) + )); + } + + #[test] + fn passwordless_sudo_is_detected_without_triggering_a_prompt() { + let allowed = ScriptedTransport::new(vec![ScriptedStep::ok("sudo -n true", "")]); + assert!(has_passwordless_sudo(&allowed).unwrap()); + + let denied = ScriptedTransport::new(vec![ScriptedStep::fails( + "sudo -n true", + 1, + "a password is required", + )]); + assert!(!has_passwordless_sudo(&denied).unwrap()); + } + + #[test] + fn a_refusal_says_what_was_wrong_and_what_to_do() { + let report = report_with(vec![finding("stale repository files", 70)]); + let refusal = assess(&report, true).expect("refused"); + let message = describe_refusal(&refusal, "gpu-box", &report, 5); + assert!(message.contains("stale repository files"), "{message}"); + assert!(message.contains("Health of gpu-box"), "{message}"); + } + + #[test] + fn a_sudo_refusal_offers_the_manual_route() { + let message = describe_refusal( + &Refusal::NeedsPasswordlessSudo, + "gpu-box", + &report_with(vec![]), + 5, + ); + assert!(message.contains("would hang rather than fail"), "{message}"); + assert!( + message.contains("ssh gpu-box -- rocm bootstrap setup"), + "{message}" + ); + } + + #[test] + fn the_sdk_is_not_attempted_when_the_driver_fails() { + // Stacking a second install onto a failed one leaves a machine in a + // state neither step can describe. + let transport = ScriptedTransport::new(vec![ + ScriptedStep::fails("install driver --yes", 1, "no supported GPU"), + ScriptedStep::ok("install sdk --yes", ""), + ]); + let error = install(&transport, "gpu-box", "rocm") + .unwrap_err() + .to_string(); + assert!(error.contains("Nothing further was attempted"), "{error}"); + } + + #[test] + fn a_half_finished_install_says_so_rather_than_reporting_success() { + let transport = ScriptedTransport::new(vec![ + ScriptedStep::ok("install driver --yes", "status: completed"), + ScriptedStep::fails("install sdk --yes", 1, "no wheels for this platform"), + ]); + let error = install(&transport, "gpu-box", "rocm") + .unwrap_err() + .to_string(); + assert!(error.contains("part-way through"), "{error}"); + } + + #[test] + fn a_required_reboot_is_surfaced_because_nobody_is_watching_the_machine() { + assert!(mentions_reboot("execution:\n reboot_required: true\n")); + assert!(mentions_reboot("A reboot required before use")); + assert!(!mentions_reboot("status: completed")); + } +} diff --git a/apps/rocm/src/remote/mod.rs b/apps/rocm/src/remote/mod.rs new file mode 100644 index 000000000..3b53a6ec4 --- /dev/null +++ b/apps/rocm/src/remote/mod.rs @@ -0,0 +1,2235 @@ +// Copyright © Advanced Micro Devices, Inc., or its affiliates. +// +// SPDX-License-Identifier: MIT + +//! Driving a remote GPU host from this machine. +//! +//! The split this module is built around: **SSH is the control channel, not the +//! data path.** Everything that inspects or changes the remote — probing it, +//! starting a managed `rocm serve`, reading the service registry, tearing a +//! session down — goes over SSH via [`transport`]. The inference traffic does +//! not: the remote publishes its own loopback-bound service onto the tailnet +//! (see [`publish`]), so there is no local tunnel process and no local port to +//! keep alive. An endpoint therefore outlives the command that created it and +//! answers from any of the user's machines, not only this one. +//! +//! Two lifecycles, both on the remote and both able to fail alone: the model +//! server, and the publish pointing at it. `status` reports them as separate +//! facts rather than one health value, because the repair differs — a withdrawn +//! publish is re-declared, a dead server has to be started again. +//! +//! Keeping the control channel narrow is what makes this testable: it is one +//! trait with a scripted stand-in, and [`tailnet`]/[`publish`] parsing is pure, +//! so the flows below are unit tests with no network, no SSH server, and no +//! tailnet. + +use std::fmt::Write as _; + +use anyhow::{Context, Result, bail}; +use clap::Subcommand; +use rocm_core::{AppPaths, ManagedServiceRecord}; + +use session::RemoteSessionRecord; +use transport::{SshTransport, Transport}; + +pub(crate) mod bootstrap; +pub(crate) mod doctor; +pub(crate) mod install; +pub(crate) mod provision; +pub(crate) mod publish; +pub(crate) mod session; +pub(crate) mod tailnet; +pub(crate) mod transport; + +/// Default loopback port the model server binds on the remote. +const DEFAULT_REMOTE_PORT: u16 = 11434; +/// Default port the remote publishes to the tailnet. +const DEFAULT_TAILNET_PORT: u16 = 8000; +/// Release channel a remote installs from when nothing else is asked for. +const DEFAULT_CHANNEL: &str = "release"; + +#[derive(Subcommand, Debug)] +pub(crate) enum RemoteCommand { + /// List the machines on your tailnet that could host a model. + #[command(after_help = "EXAMPLES:\n \ +rocm remote targets\n \ +rocm remote targets --tag gpu")] + Targets { + /// Only show machines carrying this tailnet tag, such as `gpu`. + #[arg(long, value_name = "TAG")] + tag: Option, + }, + /// Serve a model on a remote machine and publish it to your tailnet. + #[command(after_help = "EXAMPLES:\n \ +rocm remote serve gpu-box qwen2.5-7b-instruct\n \ +rocm remote serve gpu-box qwen2.5-7b-instruct --tailnet-port 8080")] + Serve { + /// Machine to serve on: a tailnet name, or an SSH destination on it. + target: String, + /// Model name, alias, or a path on the remote machine. + model: String, + /// Engine to use on the remote. + #[arg(long)] + engine: Option, + /// GPU to serve on, as the remote sees it. + #[arg(long, value_name = "INDEX|auto")] + gpu: Option, + /// SSH port for the control channel. + #[arg(long, value_name = "PORT")] + ssh_port: Option, + /// Loopback port the model server binds on the remote. + #[arg(long, value_name = "PORT", default_value_t = DEFAULT_REMOTE_PORT)] + remote_port: u16, + /// Port the remote publishes to your tailnet. + #[arg(long, value_name = "PORT", default_value_t = DEFAULT_TAILNET_PORT)] + tailnet_port: u16, + /// Release channel to install from, if the remote needs the CLI. + /// + /// This CLI carries no record of the channel it was built from, so it + /// cannot match yours automatically. Name it if your machines should + /// track something other than release. + #[arg(long, default_value = DEFAULT_CHANNEL)] + channel: String, + /// Install ROCm on the machine if it does not have it. + /// + /// Off by default. Installing a GPU stack can run for minutes and may + /// need a reboot, which is a lot to start on a machine you are not + /// sitting at without saying so. Machines the failure catalog says need + /// a person are refused even with this set. + #[arg(long)] + install_rocm: bool, + }, + /// Check a remote machine's GPU and ROCm health. + #[command(after_help = "EXAMPLES:\n \ +rocm remote doctor gpu-box\n \ +rocm remote doctor gpu-box --symptom \"hip error 101\"")] + Doctor { + /// Machine to check: a tailnet name, or an SSH destination on it. + target: String, + /// Error text you saw, to sharpen the match. + #[arg(long)] + symptom: Option, + /// Show at most this many findings. + #[arg(long, default_value_t = 5)] + top: usize, + /// SSH port for the control channel. + #[arg(long, value_name = "PORT")] + ssh_port: Option, + }, + /// Show the remote sessions started from this machine. + Status { + /// Session id, or part of a machine name. Omit for all sessions. + session: Option, + }, + /// Re-publish a session's endpoint without restarting the model. + Attach { + /// Session id, or part of a machine name. + session: String, + }, + /// Stop a remote session: withdraw the endpoint and stop the model. + Stop { + /// Session id, or part of a machine name. + session: String, + /// Forget the session locally even if the machine cannot confirm it + /// stopped. + /// + /// For a machine that is gone for good. Everything the command could + /// not finish is listed, because the risk of forgetting a session is + /// that a live endpoint stops being anyone's problem. + #[arg(long)] + force: bool, + }, +} + +pub(crate) fn run(command: RemoteCommand) -> Result<()> { + match command { + RemoteCommand::Targets { tag } => targets(tag.as_deref()), + RemoteCommand::Serve { + target, + model, + engine, + gpu, + ssh_port, + remote_port, + tailnet_port, + channel, + install_rocm, + } => serve(&ServeRequest { + target, + model, + engine, + gpu, + ssh_port, + remote_port, + tailnet_port, + channel, + install_rocm, + }), + RemoteCommand::Doctor { + target, + symptom, + top, + ssh_port, + } => remote_doctor(&target, symptom.as_deref(), top, ssh_port), + RemoteCommand::Status { session } => status(session.as_deref()), + RemoteCommand::Attach { session } => attach(&session), + RemoteCommand::Stop { session, force } => stop(&session, force), + } +} + +/// Show candidate machines, or explain why we cannot see any. +/// +/// Discovery never fails the command for a missing or idle Tailscale. A user +/// asking "what can I reach" deserves an answer about their setup, not an error +/// exit — and `rocm remote targets` is precisely the command someone runs while +/// still setting Tailscale up. +fn targets(tag: Option<&str>) -> Result<()> { + match tailnet::local_status()? { + tailnet::TailnetAvailability::NotInstalled => { + println!( + "Tailscale is not installed on this machine, so there are no targets to list." + ); + println!(); + println!( + "`rocm remote` reaches GPU machines over a tailnet. Install Tailscale and run" + ); + println!("`tailscale up` on this machine and on the GPU machine, then try again."); + } + tailnet::TailnetAvailability::NotRunning { backend_state } => { + println!("Tailscale is installed but not connected (state: {backend_state})."); + println!(); + println!("Run `tailscale up` on this machine, then try again."); + } + tailnet::TailnetAvailability::Running(status) => { + print!("{}", tailnet::render_targets(&status, tag)); + } + } + Ok(()) +} + +/// Check a machine's health without starting anything on it. +fn remote_doctor( + target: &str, + symptom: Option<&str>, + top: usize, + ssh_port: Option, +) -> Result<()> { + // Resolved the same way `serve` resolves it, so a name that serves is a name + // that can be checked first — which is the order these are meant to be used + // in. + resolve_target(target)?; + let transport = SshTransport::new(target, ssh_port)?; + // Deliberately not `ensure_ready`: that provisions a missing CLI, and a + // health check has no business installing software on a machine it was only + // asked to look at. The daemon's read-only allowlist + // (`ensure_rocm_command_is_read_only` in `apps/rocmd/src/lib.rs`) admits + // `remote doctor` without the approval flow, which is only true while this + // stays a pure read. + let remote_cli = bootstrap::locate_cli(&transport, target)?; + let (_, report) = doctor::examine_remote(&transport, &remote_cli, symptom)?; + print!("{}", doctor::render_report(target, &report, top)); + Ok(()) +} + +pub(crate) struct ServeRequest { + pub(crate) target: String, + pub(crate) model: String, + pub(crate) engine: Option, + pub(crate) gpu: Option, + pub(crate) ssh_port: Option, + pub(crate) remote_port: u16, + pub(crate) tailnet_port: u16, + pub(crate) channel: String, + pub(crate) install_rocm: bool, +} + +fn serve(request: &ServeRequest) -> Result<()> { + let paths = AppPaths::discover()?; + let peer_host = resolve_target(&request.target)?; + let transport = SshTransport::new(&request.target, request.ssh_port)?; + serve_with_transport(&transport, &paths, &peer_host, request) +} + +/// The body of [`serve`], taking its transport and paths rather than building +/// them. +/// +/// This is the seam that makes the credential handoff below testable end to +/// end: a test can drive this with a [`transport::ScriptedTransport`] and an +/// isolated [`AppPaths`], and see the same stdin write the real command path +/// produces, instead of only checking [`remote_serve_command`] and +/// [`Transport::exec_with_stdin`] in isolation with nothing pairing them. +fn serve_with_transport( + transport: &dyn Transport, + paths: &AppPaths, + peer_host: &str, + request: &ServeRequest, +) -> Result<()> { + let session_id = RemoteSessionRecord::id_for(peer_host, request.remote_port); + // Refuse a name something already sits under, exactly as `publish` refuses a + // port that already forwards somewhere else. + // + // `id_for` is the machine and the port, with no nonce, so a second `serve` + // against the same box computes the same id. Without this the second run + // overwrites the first session's key when it mints its own, and — when it + // then fails to start, because the first session still holds the port — + // deletes the shared key outright. The user reads "failed to start the model" as "nothing + // happened", while the first session is still serving on a published tailnet + // endpoint that can no longer be called. + // + // Checked before the readiness probe, so a refusal costs no round trip and + // cannot provision a machine this command then declines to use. That is also + // the limit of what this check can do: it is not the claim. Provisioning sits + // between here and the write, so a run starting inside that window sees the + // same free name — `session::store_key` below is what actually takes it, in + // one indivisible step. This check exists because it is cheap and because it + // can see *what* is on disk, and a record and a stray credential need + // different remedies. + if session::exists(paths, &session_id) { + // Which remedy applies depends on what is actually on disk. A key with no + // record is not a session `stop` can reach: `load_all` enumerates `*.json` + // only, so `resolve` cannot see it and would answer "no remote sessions are + // recorded on this machine". Telling the user to stop it would be advice + // that provably fails, on the one state this guard exists to detect. + // + // The key-only case is reachable: `serve` mints the credential before it + // can know the remote service id, so any failure between the two — a + // registry it cannot parse, a publish that will not confirm — leaves the + // key behind with no record beside it. + if RemoteSessionRecord::path_in(paths, &session_id).exists() { + bail!( + "a session is already recorded for {} on port {}.\n\ + Serving again here would take over its credential and could leave it \ + running with no way to call it.\n\ + Stop it first: rocm remote stop {}\n\ + Or serve on another port with `--remote-port`.", + request.target, + request.remote_port, + session_id + ); + } + bail!( + "a credential from an earlier `rocm remote serve` on {} port {} is still on \ + this machine, with no session recorded beside it.\n\ + An earlier attempt got far enough to mint a key and not far enough to record \ + what it started, so a model may be running there untracked.\n\ + Check the machine: ssh {} -- rocm services list\n\ + Once you are sure nothing is using it, delete: {}\n\ + Or serve on another port with `--remote-port`.", + request.target, + request.remote_port, + request.target, + session::key_path(paths, &session_id).display() + ); + } + + println!("Preparing {} ...", request.target); + let remote_cli = bootstrap::ensure_ready_with( + transport, + &request.target, + &request.channel, + request.install_rocm, + )?; + + // Mint the credential before starting anything. A model that comes up + // unauthenticated and is then published is exposed for the window between + // the two, and the whole point of publishing is that the window is visible + // to every machine on the tailnet. + // + // Writing it is also what *claims* the session name, and that is the check + // that actually decides. The `session::exists` call above runs before + // `ensure_ready_with` so a refusal costs no round trip and cannot provision + // a machine this command then declines to use — but that placement is + // precisely what makes it unable to decide: provisioning takes minutes, and + // a second run started inside that window would see the same free name. + // `store_key` creates the file exclusively, so only one run can be here. + let api_key = rocm_core::generate_endpoint_api_key(); + if let Err(error) = session::store_key(paths, &session_id, &api_key) { + if error.downcast_ref::().is_some() { + // Nothing of ours is on disk — the create is what failed — so there + // is nothing to unwind, and in particular nothing to clear: the + // credential under this name belongs to the run that won, and the + // model it guards may already be published. + bail!( + "another `rocm remote serve` claimed {} on port {} while this one was \ + preparing the machine.\n\ + Both runs name the session after the machine and the port, so carrying on \ + would take over its credential and could leave it running with no way to \ + call it.\n\ + See what is there: rocm remote status\n\ + Or serve on another port with `--remote-port`.", + request.target, + request.remote_port + ); + } + return Err(error.context( + "refusing to publish a model endpoint whose API key could not be saved locally: \ + without it you would have no way to call the endpoint you are about to expose", + )); + } + + println!("Starting {} on {} ...", request.model, request.target); + let start = match transport + .exec_with_stdin(&remote_serve_command(&remote_cli, request), Some(&api_key)) + { + Ok(start) => start, + Err(error) => { + // The command may already have reached the remote — contact can be + // lost after the remote has begun starting the model — so its state + // is not knowable from here. Drop the key that now guards nothing, + // and say so rather than leaving the user to assume nothing + // happened. + // + // Note this also catches the case where ssh never reached the host + // at all, where nothing was started and the uncertainty is + // overstated. The inner error says "could not reach" plainly, so + // the user is not misled, but telling the two apart here would need + // the transport to report unreachability as something richer than a + // message. Left as is rather than grown a typed error for it. + session::clear_key(paths, &session_id); + return Err(error.context(format!( + "lost contact with {} while starting the model, so it may or may not be \ + running.\n\ + Check with: ssh {} -- {remote_cli} services list", + request.target, request.target + ))); + } + }; + if !start.success { + session::clear_key(paths, &session_id); + bail!( + "failed to start the model on {}: {}", + request.target, + start.stderr.trim() + ); + } + + // From here the model is running on someone's GPU. Every remaining failure + // has to leave the machine in a state the user can find and act on, so each + // one unwinds what has been done rather than returning and forgetting. + let remote_service_id = + match discover_started_service(transport, &remote_cli, request.remote_port) { + Ok(service_id) => service_id, + Err(error) => { + // The key stays. The model is very likely running and it was + // handed this credential, so deleting our only copy would leave + // a service the user can find but cannot call — or stop through + // its own API. Nothing can be stopped by name when the name is + // what could not be read, so say where to look and hand back the + // credential rather than implying it was all cleaned up. + return Err(error.context(format!( + "a model may now be running on {} port {} with nothing tracking it.\n\ + Check with: ssh {} -- {remote_cli} services list\n\ + Its API key was kept at {} — it is the only copy.", + request.target, + request.remote_port, + request.target, + session::key_path(paths, &session_id).display() + ))); + } + }; + + println!("Publishing to the tailnet ..."); + if let Err(error) = publish::publish(transport, request.tailnet_port, request.remote_port) { + // The model is up but unreachable. Stop it rather than leaving a GPU + // occupied by something nobody can call and nothing records. + let leftovers = unwind_partial_serve( + transport, + paths, + &session_id, + &remote_cli, + Some(&remote_service_id), + None, + ); + return Err(describe_leftovers(error, &request.target, &leftovers)); + } + + let base_url = base_url_for(peer_host, request.tailnet_port); + let record = RemoteSessionRecord { + session_id: session_id.clone(), + target: request.target.clone(), + peer_host: peer_host.to_owned(), + ssh_port: request.ssh_port, + model: request.model.clone(), + remote_service_id: remote_service_id.clone(), + remote_cli: remote_cli.clone(), + remote_port: request.remote_port, + tailnet_port: request.tailnet_port, + base_url, + created_at_unix_ms: RemoteSessionRecord::now(), + }; + if let Err(error) = record.write(paths) { + // The endpoint is live and published at this point. Without a record + // nothing on this machine knows it exists, so leaving it up would be + // exactly the untracked exposure the whole design tries to avoid. + let leftovers = unwind_partial_serve( + transport, + paths, + &session_id, + &remote_cli, + Some(&remote_service_id), + Some((request.tailnet_port, request.remote_port)), + ); + return Err(describe_leftovers( + error.context("could not record the session on this machine"), + &request.target, + &leftovers, + )); + } + + println!(); + println!("{}", render_started(paths, &record, &api_key)); + Ok(()) +} + +/// Undo as much of a half-finished `serve` as possible, returning whatever could +/// not be undone. +/// +/// Withdraw before stopping, for the same reason teardown does: an endpoint +/// still answering is worse than a process still running. A stopped model behind +/// a live publish refuses connections; a live model behind a forgotten publish is +/// a GPU endpoint on the tailnet that nothing is tracking. +/// +/// Every step's failure is collected rather than discarded. The caller needs to +/// tell the user what is still out there — silently swallowing a failed stop is +/// how a machine ends up with a model nobody remembers starting. +fn unwind_partial_serve( + transport: &dyn Transport, + paths: &AppPaths, + session_id: &str, + remote_cli: &str, + service_id: Option<&str>, + published: Option<(u16, u16)>, +) -> Vec { + let mut leftovers = Vec::new(); + + if let Some((tailnet_port, remote_port)) = published + && let Err(error) = publish::withdraw(transport, tailnet_port, remote_port) + { + leftovers.push(format!( + "the endpoint on port {tailnet_port} may still be published ({error})" + )); + } + + if let Some(service_id) = service_id { + match transport.exec(&format!( + "{remote_cli} services stop {} --yes", + shell_quote(service_id) + )) { + Ok(outcome) if outcome.success => {} + Ok(outcome) => leftovers.push(format!( + "the model ({service_id}) may still be running: {}", + outcome.stderr.trim() + )), + Err(error) => leftovers.push(format!( + "the model ({service_id}) may still be running: {error}" + )), + } + } + + // Only drop the credential when there is nothing left for it to guard. If any + // step above failed, the model may still be running and its endpoint may still + // be published, and the key is the only way to call it — deleting it here + // would leave a live tailnet endpoint that nobody can use and nothing tracks. + // + // This matches what the rest of the module already does: the + // `discover_started_service` failure path keeps the key and says where it is, + // and `stop_with_transport` keeps it whenever it cannot confirm both halves. + if leftovers.is_empty() { + session::clear_key(paths, session_id); + } else { + leftovers.push(format!( + "its API key was kept at {} — it is the only copy", + session::key_path(paths, session_id).display() + )); + } + leftovers +} + +/// Attach what could not be cleaned up to the error that caused it. +fn describe_leftovers(error: anyhow::Error, target: &str, leftovers: &[String]) -> anyhow::Error { + if leftovers.is_empty() { + return error; + } + error.context(format!( + "{target} was left with things this command could not undo:\n{}\n\ + Check it with: ssh {target} -- rocm services list", + leftovers + .iter() + .map(|leftover| format!(" - {leftover}")) + .collect::>() + .join("\n") + )) +} + +/// Resolve a user-supplied target to the tailnet name its endpoint is built on. +fn resolve_target(target: &str) -> Result { + match tailnet::local_status()? { + tailnet::TailnetAvailability::NotInstalled => bail!( + "Tailscale is not installed on this machine.\n\ + `rocm remote serve` publishes the model onto your tailnet, so both machines \ + need it. Install Tailscale and run `tailscale up`, then try again." + ), + tailnet::TailnetAvailability::NotRunning { backend_state } => bail!( + "Tailscale is installed but not connected (state: {backend_state}).\n\ + Run `tailscale up` on this machine, then try again." + ), + tailnet::TailnetAvailability::Running(status) => { + let Some(peer) = tailnet::resolve_peer(&status, target)? else { + bail!( + "`{target}` is not a machine on this tailnet.\n\ + Run `rocm remote targets` to see what is." + ); + }; + if !peer.online { + // Cheaper and far clearer than letting SSH time out. + bail!( + "`{target}` is on this tailnet but currently offline.\n\ + Start it, or run `rocm remote targets` to pick another machine." + ); + } + Ok(peer.endpoint_host().to_owned()) + } + } +} + +/// The remote command that starts the model. +/// +/// The API key arrives on stdin rather than in the command line: both the local +/// `ssh` invocation and the remote shell expose their arguments in the process +/// table, so an interpolated key would be readable by any other user on either +/// machine. `--require-api-key` is what makes the loopback bind authenticated +/// anyway, since the publish widens who can reach it. +fn remote_serve_command(remote_cli: &str, request: &ServeRequest) -> String { + let mut command = format!( + "IFS= read -r ROCM_SERVE_API_KEY; export ROCM_SERVE_API_KEY; \ + {remote_cli} serve {} --managed --require-api-key --host {} --port {}", + shell_quote(&request.model), + publish::LOOPBACK, + request.remote_port + ); + if let Some(engine) = &request.engine { + let _ = write!(command, " --engine {}", shell_quote(engine)); + } + if let Some(gpu) = &request.gpu { + let _ = write!(command, " --gpu {}", shell_quote(gpu)); + } + command +} + +/// Find the service the remote just started, by the port we asked it to bind. +fn discover_started_service( + transport: &dyn Transport, + remote_cli: &str, + remote_port: u16, +) -> Result { + let listing = transport + .run(&format!("{remote_cli} services list --json --all")) + .context("could not read the remote's service registry after starting the model")?; + let records: Vec = serde_json::from_str(&listing).context( + "could not understand the remote's service registry; the remote CLI may be a \ + different version than this one", + )?; + + records + .into_iter() + .filter(|record| record.port == remote_port) + // Several records can share a port over a machine's lifetime; the newest + // is the one we just started. + .max_by_key(|record| record.created_at_unix_ms) + .map(|record| record.service_id) + .with_context(|| { + format!("the remote started no service on port {remote_port}; nothing to publish") + }) +} + +fn base_url_for(peer_host: &str, tailnet_port: u16) -> String { + // `/v1` to match what the local serve path records, so a URL from either + // side can be pasted into the same client. + format!("http://{peer_host}:{tailnet_port}/v1") +} + +fn render_started(paths: &AppPaths, record: &RemoteSessionRecord, api_key: &str) -> String { + let mut output = String::new(); + let _ = writeln!(output, "Model serving on {}", record.target); + let _ = writeln!(output); + let _ = writeln!(output, " endpoint: {}", record.base_url); + let _ = writeln!(output, " api key: {api_key}"); + let _ = writeln!(output, " session: {}", record.session_id); + // Say where the key was kept. It is shown once here, and without this the + // only copy the user has is whatever their terminal still holds. + let _ = writeln!( + output, + " key file: {}", + session::key_path(paths, &record.session_id).display() + ); + let _ = writeln!(output); + // Say who can reach this. The loopback-only mental model from local serving + // does not carry over, and a user who assumes it does will not think to ask + // whether their tailnet ACLs are right. + let _ = writeln!( + output, + "This endpoint is reachable by every machine on your tailnet that your tailnet's" + ); + let _ = writeln!( + output, + "access rules allow. The API key above is what stops anyone else calling it." + ); + let _ = writeln!(output); + let _ = writeln!(output, " check: rocm remote status {}", record.session_id); + let _ = writeln!(output, " stop: rocm remote stop {}", record.session_id); + output +} + +/// How a session's model server is doing, as the remote reports it. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) enum ServerHealth { + Healthy, + Pending, + Failed, + /// The remote could not be reached at all. + Unreachable, + /// The remote answered, but not in a way we could read. + Error, + /// The remote reported a lifecycle state this CLI does not know. + /// + /// Distinct from `Failed`, which is a claim about the model. A word we do + /// not recognise usually means the remote runs a different version, and + /// calling that "failed" sends the user to restart something that may be + /// working perfectly. + Unrecognised { + raw: String, + }, + /// The remote has no record of this service any more. + Gone, +} + +impl ServerHealth { + const fn label(&self) -> &'static str { + match self { + Self::Healthy => "healthy", + Self::Pending => "starting", + Self::Failed => "failed", + Self::Unreachable => "unreachable", + Self::Error => "error", + Self::Gone => "gone", + Self::Unrecognised { .. } => "unrecognised", + } + } + + /// The label plus, where it helps, why we cannot say more. + fn describe(&self) -> String { + match self { + Self::Unrecognised { raw } => { + format!("unrecognised state `{raw}` (the machine may run a different CLI version)") + } + other => other.label().to_owned(), + } + } +} + +/// Map the remote registry's own lifecycle words onto the states we report. +fn health_from_status(raw: &str) -> ServerHealth { + match raw { + "ready" | "running" => ServerHealth::Healthy, + "starting" | "recovering" => ServerHealth::Pending, + "failed" | "stopped" => ServerHealth::Failed, + // Not folded into `Failed`. An unknown word almost always means version + // skew, and reporting it as a failure sends the user to restart a model + // that may be serving fine. + other => ServerHealth::Unrecognised { + raw: other.to_owned(), + }, + } +} + +/// Both halves of one session, as observed right now. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct SessionObservation { + pub(crate) server: ServerHealth, + /// The publishing state, or why it could not be read. Not an `Option`: + /// "could not read it" has two causes worth telling apart, and an `Option` + /// can only say that one of them happened. + pub(crate) publish: publish::PublishObservation, +} + +/// Probe one session over the control channel. +fn observe(transport: &dyn Transport, record: &RemoteSessionRecord) -> SessionObservation { + let listing = match transport.exec(&format!("{} services list --json --all", record.remote_cli)) + { + Ok(outcome) if outcome.success => outcome.stdout, + // Reached the machine but the command failed, versus could not reach it + // at all. Different problems, different fixes, so different words. + // The machine is reachable — only this one command failed — so the + // publishing state is still worth asking about, and the answer is a real + // observation rather than a guess. Reporting `Unreachable` here rendered + // as "the machine could not be asked", which is the wrong failure: it + // sends the user to check the network for a machine that just answered. + Ok(_) => { + return SessionObservation { + server: ServerHealth::Error, + publish: publish::observe(transport, record.tailnet_port, record.remote_port), + }; + } + Err(_) => { + return SessionObservation { + server: ServerHealth::Unreachable, + publish: publish::PublishObservation::Unreachable( + "the machine could not be reached over ssh".to_owned(), + ), + }; + } + }; + + let server = match serde_json::from_str::>(&listing) { + Ok(records) => records + .iter() + .find(|candidate| candidate.service_id == record.remote_service_id) + .map_or(ServerHealth::Gone, |found| { + health_from_status(&found.status) + }), + Err(_) => ServerHealth::Error, + }; + + SessionObservation { + server, + // `publish::observe`, not `publish_state(..).ok()`: the same two cases + // separated above for `services list` — reached-but-failed versus never + // reached — must stay separated here. `.ok()` collapsed them, so a + // remote that answered with a concrete reason was reported as one that + // was never asked. + publish: publish::observe(transport, record.tailnet_port, record.remote_port), + } +} + +fn status(session: Option<&str>) -> Result<()> { + let paths = AppPaths::discover()?; + let sessions = match session { + Some(needle) => vec![session::resolve(&paths, needle)?], + None => session::load_all(&paths)?, + }; + + if sessions.is_empty() { + println!("No remote sessions have been started from this machine."); + println!(); + println!("Start one with `rocm remote serve `."); + return Ok(()); + } + + let observations = sessions + .iter() + .map(|record| { + let observed = SshTransport::new(&record.target, record.ssh_port).map_or_else( + // A record naming a machine ssh cannot address is not a reason + // to hide every other session from the listing. + // + // `Unreachable`, not `Error`: nothing was ever sent, so this is + // the never-reached side of the split, not "answered in a way we + // could not read". And the constructor's own reason is carried + // rather than replaced with a generic line — it names what is + // wrong with the destination, which is the one thing that makes + // this fixable. + |error| SessionObservation { + server: ServerHealth::Unreachable, + publish: publish::PublishObservation::Unreachable(format!("{error:#}")), + }, + |transport| observe(&transport, record), + ); + (record.clone(), observed) + }) + .collect::>(); + print!("{}", render_status(&paths, &observations)); + Ok(()) +} + +/// Render sessions with their two lifecycles kept apart. +fn render_status( + paths: &AppPaths, + observations: &[(RemoteSessionRecord, SessionObservation)], +) -> String { + use publish::PublishObservation::{Failed, Known, Unreachable}; + + let mut output = String::new(); + let _ = writeln!(output, "Remote Sessions"); + let _ = writeln!(output); + + for (record, observed) in observations { + let _ = writeln!(output, "- {}", record.session_id); + let _ = writeln!(output, " machine: {}", record.target); + let _ = writeln!(output, " model: {}", record.model); + let _ = writeln!(output, " endpoint: {}", record.base_url); + // The path, not the credential: a listing lands in scrollback and CI + // logs, and the file beside it is owner-only for a reason. + let _ = writeln!( + output, + " key file: {}", + session::key_path(paths, &record.session_id).display() + ); + // Two facts, never collapsed into one: which of them is wrong decides + // whether the fix is `attach` or starting the model again. + let _ = writeln!(output, " model server: {}", observed.server.describe()); + let _ = writeln!( + output, + " endpoint published: {}", + match &observed.publish { + Known(publish::PublishState::Published) => "yes".to_owned(), + Known(publish::PublishState::Absent) => "no".to_owned(), + Known(publish::PublishState::Foreign { forwards_to }) => + format!("no — that port now forwards to {forwards_to}"), + // Not "no": Funnel is classified before anything is asked about + // our own forward and short-circuits, so this state says + // nothing either way about whether the endpoint is published. + // Leading with "no" answered a question it had not looked at, + // and buried the one thing here that needs acting on. + Known(publish::PublishState::FunnelAllowed) => format!( + "exposed — Tailscale Funnel is allowed on that port, which puts it on the \ + public internet; run `tailscale funnel --tcp={} off` on the remote, then \ + check again", + record.tailnet_port + ), + // The three below all mean "could not tell", and none may be + // read as "no": an endpoint that is still up must never render + // as one that is down, or the user stops looking for it. They + // stay distinct because the next step differs — fix the reply, + // fix the remote's tailscale, or fix the connection. + Known(publish::PublishState::Unreadable) => + "unknown — the machine's reply could not be read".to_owned(), + // Reached, and it told us why. Printing its own words beats + // "could not be asked", which describes a different failure and + // sends the user to check the network instead of the remote. + Failed(why) => format!("unknown — the machine answered: {why}"), + Unreachable(why) => format!("unknown — the machine could not be asked: {why}"), + } + ); + + if let Some(hint) = repair_hint(record, observed) { + let _ = writeln!(output, " {hint}"); + } + } + + // A publish outlives the machine's reboots, so a stale record is not merely + // untidy — it may be an endpoint still answering with nothing tracking it. + if observations + .iter() + .any(|(_, observed)| matches!(observed.server, ServerHealth::Gone)) + { + let _ = writeln!(output); + let _ = writeln!( + output, + "A session whose model server is gone may still be publishing its port." + ); + let _ = writeln!(output, "Run `rocm remote stop ` to clear it."); + } + output +} + +/// The one command that fixes what is wrong, when exactly one thing is. +fn repair_hint(record: &RemoteSessionRecord, observed: &SessionObservation) -> Option { + match (&observed.server, &observed.publish) { + ( + ServerHealth::Healthy, + publish::PublishObservation::Known(publish::PublishState::Absent), + ) => Some(format!( + "fix: rocm remote attach {} (the model is fine; only the endpoint is missing)", + record.session_id + )), + (ServerHealth::Failed | ServerHealth::Gone, _) => Some(format!( + "fix: rocm remote stop {} then serve again", + record.session_id + )), + _ => None, + } +} + +fn attach(needle: &str) -> Result<()> { + let paths = AppPaths::discover()?; + let record = session::resolve(&paths, needle)?; + let transport = SshTransport::new(&record.target, record.ssh_port)?; + attach_with_transport(&transport, &record) +} + +/// The body of [`attach`], taking its transport rather than building one. +/// +/// The same seam [`serve_with_transport`] exists for, and for the same reason: +/// without it the only testable part of `attach` is [`render_status`], and the +/// refusal below — the one thing standing between a dead model and an endpoint +/// that answers with connection refused — is reachable by no test at all. +fn attach_with_transport(transport: &dyn Transport, record: &RemoteSessionRecord) -> Result<()> { + // Re-declaring a publish is cheap, but doing it over a dead model server + // would produce an endpoint that answers with connection refused — worse + // than one that is honestly absent. + let observed = observe(transport, record); + match observed.server { + ServerHealth::Healthy | ServerHealth::Pending => {} + other => bail!( + "the model server for {} is {} on {}, so re-publishing would give you an \ + endpoint with nothing behind it.\n\ + Run `rocm remote stop {}` and serve again.", + record.session_id, + other.label(), + record.target, + record.session_id + ), + } + + publish::publish(transport, record.tailnet_port, record.remote_port)?; + println!("Endpoint re-published: {}", record.base_url); + println!("The model was not restarted."); + Ok(()) +} + +fn stop(needle: &str, force: bool) -> Result<()> { + let paths = AppPaths::discover()?; + let record = session::resolve(&paths, needle)?; + let transport = SshTransport::new(&record.target, record.ssh_port)?; + stop_with_transport(&transport, &paths, &record, force) +} + +/// The body of [`stop`], taking its transport and paths rather than building +/// them. +/// +/// Teardown is the path where getting it wrong is worst — forgetting a session +/// whose endpoint is still published leaves a GPU endpoint on the tailnet that +/// nothing tracks — and it was the one path no test could drive. +fn stop_with_transport( + transport: &dyn Transport, + paths: &AppPaths, + record: &RemoteSessionRecord, + force: bool, +) -> Result<()> { + // Withdraw before stopping the model. If only one of the two can be done, + // the endpoint being gone is the one that matters: a stopped model behind a + // live publish is a refused connection, but a live model behind a forgotten + // publish is an open GPU endpoint nobody is tracking. + let withdrawn = publish::withdraw(transport, record.tailnet_port, record.remote_port); + let stopped = transport.exec(&format!( + "{} services stop {} --yes", + record.remote_cli, + shell_quote(&record.remote_service_id) + )); + let stop_failure = describe_stop_failure(&stopped); + let model_stopped = stop_failure.is_none(); + + if !force { + if let Err(error) = withdrawn { + // Keep the record. Deleting it here would leave a published endpoint + // with nothing on this machine that can find it again. + bail!( + "could not confirm the endpoint for {} was withdrawn: {error}\n\ + The session is still listed so you can retry with `rocm remote stop {}`.\n\ + To clear it by hand: ssh {} -- tailscale serve --tcp={} off\n\ + To forget it locally anyway: rocm remote stop {} --force", + record.session_id, + record.session_id, + record.target, + record.tailnet_port, + record.session_id + ); + } + + if let Some(why) = &stop_failure { + // Same reasoning one step further in. The record is the only thing on + // this machine that knows the model's id and where it runs; dropping + // it while the model is still up leaves a GPU occupied by something + // the user can no longer name. + // + // The remote's own words are carried, as the withdraw branch above + // already does. Without them this arm reports only *that* the stop + // failed, which reads the same whether the machine refused, the + // service was already gone, or ssh never got there — three different + // next steps behind one sentence. + bail!( + "the endpoint for {} was withdrawn, but its model could not be stopped: {why}\n\ + The session is still listed so you can retry with `rocm remote stop {}`.\n\ + To check the machine: ssh {} -- {} services list\n\ + To forget it locally anyway: rocm remote stop {} --force", + record.session_id, + record.session_id, + record.target, + record.remote_cli, + record.session_id + ); + } + } + + session::clear_key(paths, &record.session_id); + record.remove(paths); + + print!( + "{}", + render_stopped(record, withdrawn.is_ok(), model_stopped, force) + ); + Ok(()) +} + +/// Why the remote could not stop the model, or `None` if it did. +/// +/// Keeps the two failure shapes apart the way the rest of this module does: +/// a machine that answered and refused carries its own stderr and exit status, +/// and one that was never reached carries the transport's reason. Collapsing +/// them into a bare bool is what left the teardown error with nothing to say. +fn describe_stop_failure(stopped: &Result) -> Option { + match stopped { + Ok(outcome) if outcome.success => None, + Ok(outcome) => { + let stderr = outcome.stderr.trim(); + let code = outcome + .code + .map_or_else(|| "signal".to_owned(), |code| code.to_string()); + Some(if stderr.is_empty() { + format!("the machine answered but the command failed (exit {code})") + } else { + format!("the machine answered (exit {code}): {stderr}") + }) + } + Err(error) => Some(format!("{error:#}")), + } +} + +/// Report a teardown, naming anything it could not finish. +/// +/// `--force` exists for a machine that is gone for good, and its whole risk is +/// that the user stops thinking about a session that may still be live. So a +/// forced stop is louder than a clean one, not quieter: it lists exactly what +/// may remain and the commands to deal with it once the machine is reachable. +fn render_stopped( + record: &RemoteSessionRecord, + withdrawn: bool, + model_stopped: bool, + forced: bool, +) -> String { + let mut output = String::new(); + let _ = writeln!(output, "Stopped {}.", record.session_id); + let _ = writeln!( + output, + " endpoint withdrawn: {}", + if withdrawn { "yes" } else { "NOT CONFIRMED" } + ); + let _ = writeln!( + output, + " model server stopped: {}", + if model_stopped { + "yes" + } else { + "NOT CONFIRMED" + } + ); + + if forced && !(withdrawn && model_stopped) { + let _ = writeln!(output); + let _ = writeln!( + output, + "The session was forgotten on this machine, but {} may still be running it.", + record.target + ); + if !withdrawn { + let _ = writeln!( + output, + " endpoint still reachable on the tailnet — clear it with:\n \ + ssh {} -- tailscale serve --tcp={} off", + record.target, record.tailnet_port + ); + } + if !model_stopped { + let _ = writeln!( + output, + " model may still hold the GPU — check with:\n \ + ssh {} -- {} services list", + record.target, record.remote_cli + ); + } + } + output +} + +/// Quote a value for a POSIX remote shell. +/// +/// Model names and engine flags are user-supplied values being placed into a +/// command line that a shell on another machine will interpret. Anything not +/// obviously inert gets single-quoted, with embedded single quotes closed and +/// re-opened, so no input can end the quoting and start a new command. +fn shell_quote(value: &str) -> String { + let inert = |character: char| { + character.is_ascii_alphanumeric() || matches!(character, '-' | '_' | '.' | '/' | ':' | '=') + }; + if !value.is_empty() && value.chars().all(inert) { + return value.to_owned(); + } + format!("'{}'", value.replace('\'', r"'\''")) +} + +#[cfg(test)] +mod tests { + use super::*; + use transport::{ScriptedStep, ScriptedTransport, TransportCall}; + + fn request() -> ServeRequest { + ServeRequest { + target: "gpu-box".to_owned(), + model: "qwen2.5-7b-instruct".to_owned(), + engine: None, + gpu: None, + ssh_port: None, + remote_port: 11434, + tailnet_port: 8000, + channel: DEFAULT_CHANNEL.to_owned(), + install_rocm: false, + } + } + + /// A serve-status document showing our own forward, for unwind tests. + const PUBLISHED_FIXTURE: &str = r#"{"TCP": {"8000": {"TCPForward": "127.0.0.1:11434"}}}"#; + + /// A rendering-only paths root: `render_status` needs one to name each + /// session's key file, and nothing here touches disk. + fn render_paths() -> AppPaths { + AppPaths { + config_dir: std::path::PathBuf::from("/tmp/rocm-render/config"), + data_dir: std::path::PathBuf::from("/tmp/rocm-render/data"), + cache_dir: std::path::PathBuf::from("/tmp/rocm-render/cache"), + } + } + + /// An isolated config/data root, so an unwind test never clears a real + /// endpoint key. + fn temp_paths(tag: &str) -> (std::path::PathBuf, AppPaths) { + let root = std::env::temp_dir().join(format!( + "rocm-remote-unwind-{tag}-{}-{}", + std::process::id(), + rocm_core::unix_time_millis() + )); + let _ = std::fs::remove_dir_all(&root); + std::fs::create_dir_all(&root).unwrap(); + ( + root.clone(), + AppPaths { + config_dir: root.join("config"), + data_dir: root.join("data"), + cache_dir: root.join("cache"), + }, + ) + } + + fn sample_record() -> RemoteSessionRecord { + RemoteSessionRecord { + session_id: "remote-gpu-box-11434".to_owned(), + target: "gpu-box".to_owned(), + peer_host: "gpu-box.example-tailnet.ts.net".to_owned(), + ssh_port: None, + model: "qwen".to_owned(), + remote_service_id: "svc-1".to_owned(), + remote_cli: "rocm".to_owned(), + remote_port: 11434, + tailnet_port: 8000, + base_url: "http://gpu-box.example-tailnet.ts.net:8000/v1".to_owned(), + created_at_unix_ms: 1, + } + } + + #[test] + fn the_remote_server_binds_loopback_but_demands_a_key() { + // The bind stays loopback — the publish is what widens reach — so the + // server would be credential-free without an explicit demand for a key. + let command = remote_serve_command("rocm", &request()); + assert!(command.contains("--host 127.0.0.1"), "{command}"); + assert!(command.contains("--require-api-key"), "{command}"); + assert!(command.contains("--managed"), "{command}"); + } + + #[test] + fn the_api_key_is_read_from_stdin_never_written_into_the_command() { + // Both machines expose command arguments in their process tables, so an + // interpolated key would be readable by any other user on either. + let command = remote_serve_command("rocm", &request()); + assert!( + command.starts_with("IFS= read -r ROCM_SERVE_API_KEY;"), + "{command}" + ); + assert!(command.contains("export ROCM_SERVE_API_KEY"), "{command}"); + } + + #[test] + fn optional_engine_and_gpu_are_threaded_through() { + let command = remote_serve_command( + "rocm", + &ServeRequest { + engine: Some("vllm".to_owned()), + gpu: Some("1".to_owned()), + ..request() + }, + ); + assert!(command.contains("--engine vllm"), "{command}"); + assert!(command.contains("--gpu 1"), "{command}"); + } + + #[cfg(unix)] + #[test] + fn hostile_values_survive_a_real_shell_as_one_literal_argument() { + // The property that matters is not the shape of the quoting but what a + // shell does with it. Ask one: each value must come back byte-identical, + // proving it was neither expanded nor split nor able to start a second + // command. + for value in [ + "x'; rm -rf ~; echo '", + "$(id)", + "`id`", + "a b", + "it's", + "*", + "--not-a-flag", + "qwen2.5-7b-instruct", + ] { + let output = std::process::Command::new("sh") + .arg("-c") + .arg(format!("printf %s {}", shell_quote(value))) + .output() + .expect("sh should run"); + assert_eq!( + String::from_utf8_lossy(&output.stdout), + value, + "shell mangled {value:?}" + ); + } + } + + #[test] + fn a_hostile_model_name_is_quoted_into_the_remote_command() { + let command = remote_serve_command( + "rocm", + &ServeRequest { + model: "x'; rm -rf ~; echo '".to_owned(), + ..request() + }, + ); + // Every embedded quote is closed and re-opened, so the payload cannot + // end the quoting and start a statement of its own. + assert!(command.contains(r"'\''"), "{command}"); + // And the flags we control still follow it as real flags. + assert!(command.contains("--require-api-key"), "{command}"); + } + + #[test] + fn shell_quoting_leaves_ordinary_values_alone_and_wraps_the_rest() { + for inert in ["qwen2.5-7b-instruct", "vllm", "/models/a.gguf", "auto"] { + assert_eq!(shell_quote(inert), inert); + } + assert_eq!(shell_quote("a b"), "'a b'"); + assert_eq!(shell_quote(""), "''"); + assert_eq!(shell_quote("it's"), r"'it'\''s'"); + } + + #[test] + fn the_newest_service_on_the_port_is_the_one_just_started() { + // A machine accumulates records on a port over its lifetime; picking an + // older one would publish a port pointing at a dead server. + let listing = r#"[ + {"service_id":"old","engine":"vllm","model_ref":"m","canonical_model_id":"m", + "host":"127.0.0.1","port":11434,"endpoint_url":"http://127.0.0.1:11434/v1", + "mode":"managed","status":"stopped","supervisor_pid":1, + "manifest_path":"/a","log_path":"/b","engine_state_path":"/c", + "created_at_unix_ms":100}, + {"service_id":"new","engine":"vllm","model_ref":"m","canonical_model_id":"m", + "host":"127.0.0.1","port":11434,"endpoint_url":"http://127.0.0.1:11434/v1", + "mode":"managed","status":"starting","supervisor_pid":2, + "manifest_path":"/a","log_path":"/b","engine_state_path":"/c", + "created_at_unix_ms":200}, + {"service_id":"other-port","engine":"vllm","model_ref":"m","canonical_model_id":"m", + "host":"127.0.0.1","port":9999,"endpoint_url":"http://127.0.0.1:9999/v1", + "mode":"managed","status":"ready","supervisor_pid":3, + "manifest_path":"/a","log_path":"/b","engine_state_path":"/c", + "created_at_unix_ms":300} + ]"#; + let transport = + ScriptedTransport::new(vec![ScriptedStep::ok("services list --json", listing)]); + assert_eq!( + discover_started_service(&transport, "rocm", 11434).unwrap(), + "new" + ); + } + + #[test] + fn a_registry_we_cannot_read_names_version_skew_as_the_likely_cause() { + let transport = ScriptedTransport::new(vec![ScriptedStep::ok( + "services list --json", + "not json at all", + )]); + let error = discover_started_service(&transport, "rocm", 11434) + .unwrap_err() + .to_string(); + assert!(error.contains("different version"), "{error}"); + } + + #[test] + fn the_two_lifecycles_are_reported_separately() { + // The whole reason for two columns: which one is broken decides whether + // the fix re-publishes or restarts. + let healthy_but_unpublished = vec![( + sample_record(), + SessionObservation { + server: ServerHealth::Healthy, + publish: publish::PublishObservation::Known(publish::PublishState::Absent), + }, + )]; + let rendered = render_status(&render_paths(), &healthy_but_unpublished); + assert!(rendered.contains("model server: healthy"), "{rendered}"); + assert!(rendered.contains("endpoint published: no"), "{rendered}"); + assert!( + rendered.contains("rocm remote attach"), + "a live model with no endpoint should point at attach, not a restart: {rendered}" + ); + } + + #[test] + fn a_funnel_exposed_port_is_reported_as_exposure_not_as_a_publish_answer() { + // Funnel is classified before anything is asked about our own forward + // and short-circuits, so this state says nothing either way about + // whether the endpoint is published. Leading the line with "no" + // answered a question it had not looked at, and buried the one thing + // on it that needs acting on. + // Port 443, not the 8000 the other fixtures use: Funnel only serves + // 443, 8443 and 10000, so a session that can reach this state at all + // is one started with `--tailnet-port`. Rendering the remedy for a + // port Funnel cannot listen on would pin a line no user could ever + // see. + let record = RemoteSessionRecord { + tailnet_port: 443, + ..sample_record() + }; + let rendered = render_status( + &render_paths(), + &[( + record.clone(), + SessionObservation { + server: ServerHealth::Healthy, + publish: publish::PublishObservation::Known( + publish::PublishState::FunnelAllowed, + ), + }, + )], + ); + assert!( + !rendered.contains("endpoint published: no"), + "Funnel exposure must not be rendered as an answer about publishing: {rendered}" + ); + assert!(rendered.contains("exposed"), "{rendered}"); + assert!( + rendered.contains("public internet"), + "the reason it matters must be on the line: {rendered}" + ); + // The remedy has to be copy-pasteable, so the real port belongs here + // rather than a `` placeholder the user has to substitute. + assert!( + rendered.contains(&format!( + "tailscale funnel --tcp={} off", + record.tailnet_port + )), + "{rendered}" + ); + } + + #[test] + fn a_dead_server_is_not_offered_a_republish() { + let rendered = render_status( + &render_paths(), + &[( + sample_record(), + SessionObservation { + server: ServerHealth::Failed, + publish: publish::PublishObservation::Known(publish::PublishState::Published), + }, + )], + ); + assert!(rendered.contains("model server: failed"), "{rendered}"); + assert!(!rendered.contains("attach"), "{rendered}"); + assert!(rendered.contains("stop"), "{rendered}"); + } + + #[test] + fn a_session_the_remote_has_forgotten_warns_about_a_stray_endpoint() { + // A publish survives reboots. A forgotten one is a GPU endpoint on the + // tailnet with nothing tracking it. + let rendered = render_status( + &render_paths(), + &[( + sample_record(), + SessionObservation { + server: ServerHealth::Gone, + publish: publish::PublishObservation::Known(publish::PublishState::Published), + }, + )], + ); + assert!(rendered.contains("may still be publishing"), "{rendered}"); + } + + #[test] + fn an_unreachable_machine_reads_differently_from_a_broken_command() { + let record = sample_record(); + + let unreachable = ScriptedTransport::new(vec![]); + assert_eq!( + observe(&unreachable, &record).server, + ServerHealth::Unreachable + ); + + let answered_badly = + ScriptedTransport::new(vec![ScriptedStep::fails("services list --json", 1, "boom")]); + assert_eq!( + observe(&answered_badly, &record).server, + ServerHealth::Error + ); + } + + #[test] + fn a_remote_that_answers_about_publishing_is_not_reported_as_one_that_was_never_asked() { + // The same distinction the test above pins for the model server, for + // the publish half. It used to be lost: `publish_state(..).ok()` + // mapped both causes to `None`, and the status line said "the machine + // could not be asked" for a machine that had answered with a concrete, + // actionable reason — sending the user to check the network instead of + // the remote's tailscale. + let record = sample_record(); + + // Reached: `services list` succeeds, then `serve status` fails with + // the remote's own words. + let answered = ScriptedTransport::new(vec![ + ScriptedStep::ok("services list --json", "[]"), + ScriptedStep::fails( + "tailscale serve status --json", + 127, + "tailscale: command not found", + ), + ]); + let observed = observe(&answered, &record); + let publish::PublishObservation::Failed(why) = &observed.publish else { + panic!("a remote that answered must not be reported as unreachable: {observed:?}"); + }; + assert!(why.contains("tailscale: command not found"), "{why}"); + + // Never reached at all: nothing scripted, so the transport errors. + let unreachable = ScriptedTransport::new(vec![]); + assert!( + matches!( + observe(&unreachable, &record).publish, + publish::PublishObservation::Unreachable(_) + ), + "a host that was never reached must say so" + ); + + // And the two must not render the same, which is the whole point. + let rendered_failed = render_status(&render_paths(), &[(record.clone(), observed)]); + let rendered_unreachable = render_status( + &render_paths(), + &[(record.clone(), observe(&unreachable, &record))], + ); + assert!( + rendered_failed.contains("tailscale: command not found"), + "the remote's own reason belongs on the status line: {rendered_failed}" + ); + assert_ne!( + rendered_failed, rendered_unreachable, + "answered-and-failed must not read identically to never-answered" + ); + } + + #[test] + fn a_service_missing_from_the_remote_registry_is_gone_not_failed() { + let record = sample_record(); + let transport = ScriptedTransport::new(vec![ + ScriptedStep::ok("services list --json", "[]"), + ScriptedStep::ok("tailscale serve status --json", "{}"), + ]); + assert_eq!(observe(&transport, &record).server, ServerHealth::Gone); + } + + #[test] + fn remote_lifecycle_words_map_onto_reported_health() { + for (raw, expected) in [ + ("ready", ServerHealth::Healthy), + ("running", ServerHealth::Healthy), + ("starting", ServerHealth::Pending), + ("recovering", ServerHealth::Pending), + ("failed", ServerHealth::Failed), + ("stopped", ServerHealth::Failed), + ] { + assert_eq!(health_from_status(raw), expected, "{raw}"); + } + } + + #[test] + fn a_state_this_version_does_not_know_is_not_called_a_failure() { + // Version skew is normal between the machine driving and the machine + // driven. Reporting an unknown word as "failed" sends the user to + // restart a model that may be serving perfectly. + let health = health_from_status("quiescing"); + assert_eq!( + health, + ServerHealth::Unrecognised { + raw: "quiescing".to_owned() + } + ); + assert!( + health.describe().contains("different CLI version"), + "{}", + health.describe() + ); + + let rendered = render_status( + &render_paths(), + &[( + sample_record(), + SessionObservation { + server: health, + publish: publish::PublishObservation::Known(publish::PublishState::Published), + }, + )], + ); + assert!( + rendered.contains("unrecognised state `quiescing`"), + "{rendered}" + ); + // And it must not be offered the dead-server repair. + assert!(!rendered.contains("then serve again"), "{rendered}"); + } + + #[test] + fn the_credential_is_recoverable_rather_than_shown_once_and_lost() { + // The key is printed once when serving. Storing it without ever naming + // where left the user's terminal scrollback as the only copy. + let started = render_started(&render_paths(), &sample_record(), "the-key"); + assert!(started.contains("key file:"), "{started}"); + + // Listings land in scrollback and CI logs, so they name the file rather + // than echoing what is in it. + let listed = render_status( + &render_paths(), + &[( + sample_record(), + SessionObservation { + server: ServerHealth::Healthy, + publish: publish::PublishObservation::Known(publish::PublishState::Published), + }, + )], + ); + assert!(listed.contains("key file:"), "{listed}"); + assert!( + !listed.contains("the-key"), + "a listing must not echo the credential:\n{listed}" + ); + // And it lives with the session, not in the local service registry. + assert!(listed.contains("remote-sessions"), "{listed}"); + } + + #[test] + fn the_started_message_states_who_can_reach_the_endpoint() { + // The loopback mental model from local serving does not carry over. A + // user who assumes it does will never check their tailnet access rules. + let rendered = render_started(&render_paths(), &sample_record(), "the-key"); + assert!( + rendered.contains("every machine on your tailnet"), + "{rendered}" + ); + assert!(rendered.contains("the-key"), "{rendered}"); + assert!( + rendered.contains("http://gpu-box.example-tailnet.ts.net:8000/v1"), + "{rendered}" + ); + } + + #[test] + fn the_endpoint_url_matches_what_local_serving_records() { + // Both sides record the OpenAI base including /v1, so a URL from either + // pastes into the same client unchanged. + assert_eq!( + base_url_for("gpu-box.example-tailnet.ts.net", 8000), + "http://gpu-box.example-tailnet.ts.net:8000/v1" + ); + } + + #[test] + fn a_teardown_that_leaves_the_model_running_keeps_the_record() { + // The record is the only thing on this machine holding the model's id + // and where it runs. Dropping it while the model is up leaves a GPU + // occupied by something the user can no longer name, let alone stop. + let rendered = render_stopped(&sample_record(), true, false, false); + assert!( + rendered.contains("model server stopped: NOT CONFIRMED"), + "{rendered}" + ); + } + + #[test] + fn a_forced_teardown_is_louder_than_a_clean_one() { + // Forgetting a session is only safe if the user is told exactly what may + // outlive it. A quiet --force is how a live endpoint stops being + // anyone's problem. + let forced = render_stopped(&sample_record(), false, false, true); + assert!(forced.contains("may still be running it"), "{forced}"); + assert!( + forced.contains("tailscale serve --tcp=8000 off"), + "a forced stop must say how to clear the endpoint: {forced}" + ); + assert!( + forced.contains("services list"), + "and how to find the model: {forced}" + ); + + // A clean stop stays quiet: there is nothing left to warn about. + let clean = render_stopped(&sample_record(), true, true, false); + assert!(!clean.contains("may still be running it"), "{clean}"); + assert!(clean.contains("endpoint withdrawn: yes"), "{clean}"); + } + + #[test] + fn a_clean_unwind_still_drops_the_key_it_minted() { + // The gate cuts both ways: when cleanup fully succeeds there is nothing + // left for the credential to guard, so leaving it behind would be the + // orphaned-secret half of the same mistake. + let (root, paths) = temp_paths("unwind-clean"); + session::store_key(&paths, "sess", "the-key").unwrap(); + let transport = ScriptedTransport::new(vec![ + ScriptedStep::ok("tailscale serve status --json", PUBLISHED_FIXTURE), + ScriptedStep::ok("tailscale serve --tcp=8000 off", ""), + ScriptedStep::ok("services stop", ""), + ]); + + let leftovers = + unwind_partial_serve(&transport, &paths, "sess", "rocm", Some("svc-1"), None); + + assert!(leftovers.is_empty(), "{leftovers:?}"); + assert!( + !session::key_path(&paths, "sess").exists(), + "a clean unwind must not leave the credential behind" + ); + let _ = std::fs::remove_dir_all(root); + } + + #[test] + fn serving_twice_on_one_machine_and_port_is_refused_rather_than_overwriting() { + // `id_for` is host + port with no nonce, so a second serve computes the + // same id. Without this refusal it overwrites the first session's key, + // then deletes it when the start fails on the port the first session + // still holds — leaving a live, published endpoint nobody can call. + let (root, paths) = temp_paths("serve-collision"); + let request = request(); + let peer_host = "gpu-box.example-tailnet.ts.net"; + let session_id = RemoteSessionRecord::id_for(peer_host, request.remote_port); + session::store_key(&paths, &session_id, "the-first-sessions-key").unwrap(); + + // No scripted steps: the refusal must land before anything is sent. + let transport = ScriptedTransport::new(Vec::new()); + let error = serve_with_transport(&transport, &paths, peer_host, &request) + .expect_err("a second serve on the same machine and port must be refused"); + let rendered = format!("{error:#}"); + + // A key with no record beside it is NOT a session `stop` can reach: + // `load_all` enumerates `*.json` only, so `resolve` would answer "no + // remote sessions are recorded on this machine". Pointing the user at + // `stop` here would be advice that provably fails. + assert!( + !rendered.contains("rocm remote stop"), + "a key-only leftover cannot be stopped, so the refusal must not say to: {rendered}" + ); + assert!( + rendered.contains(&session::key_path(&paths, &session_id).display().to_string()), + "the refusal must name the file the user has to deal with: {rendered}" + ); + + // The first session's credential is untouched. + assert_eq!( + std::fs::read_to_string(session::key_path(&paths, &session_id)).unwrap(), + "the-first-sessions-key" + ); + assert!( + transport.calls().is_empty(), + "nothing may be sent to the machine before the refusal: {:?}", + transport.calls() + ); + + // With a real record beside the key, `stop` *is* the remedy, and the + // refusal says so. + let mut record = sample_record(); + record.session_id = session_id.clone(); + record.peer_host = peer_host.to_owned(); + record.write(&paths).unwrap(); + let error = serve_with_transport(&transport, &paths, peer_host, &request) + .expect_err("a recorded session must still be refused"); + let rendered = format!("{error:#}"); + assert!(rendered.contains("already recorded"), "{rendered}"); + assert!( + rendered.contains(&format!("rocm remote stop {session_id}")), + "{rendered}" + ); + let _ = std::fs::remove_dir_all(root); + } + + #[test] + fn serves_racing_on_one_machine_and_port_leave_exactly_one_session() { + // The sequential refusal above is the outcome; this is the mechanism. + // `id_for` is host + port with no nonce, so every run here computes the + // same name, and the check that the name is free is separated from the + // write that takes it by the whole readiness probe — minutes, when that + // probe provisions a CLI. A check-then-write pair cannot survive two runs + // in that window: both see a free name, both mint a key under it, and the + // one that loses the port then deletes the other's credential on its way + // out, leaving a published endpoint nobody can call. + // + // Asserting on the end state would not see it. "One record, one key + // file" is true when both runs succeed too — they share a name. What is + // only true when the claim is atomic is that exactly one run *returns* + // having claimed it. + // + // The barrier is what makes the window real rather than hoped for: every + // thread reaches the free-name check before any of them has had time to + // write. + const RACERS: usize = 16; + let (root, paths) = temp_paths("serve-race"); + let peer_host = "gpu-box.example-tailnet.ts.net"; + let barrier = std::sync::Barrier::new(RACERS); + + let claimed = std::thread::scope(|scope| { + let racers = (0..RACERS) + .map(|_| { + scope.spawn(|| { + // One transport per thread: the scripted double records + // calls in a `RefCell` and is not shareable, and a real + // second `serve` would open its own connection anyway. + let transport = ScriptedTransport::new(full_serve_steps()); + barrier.wait(); + serve_with_transport(&transport, &paths, peer_host, &request()).is_ok() + }) + }) + .collect::>(); + racers + .into_iter() + .map(|racer| racer.join().expect("no racer may panic")) + .filter(|claimed| *claimed) + .count() + }); + + assert_eq!( + claimed, 1, + "exactly one concurrent serve may claim a machine and port; {claimed} did" + ); + + let session_id = RemoteSessionRecord::id_for(peer_host, request().remote_port); + let key = session::key_path(&paths, &session_id); + assert!( + key.exists(), + "the run that won must still hold its credential: a loser's cleanup may not \ + delete a key it did not mint" + ); + assert!( + !std::fs::read_to_string(&key).expect("read").is_empty(), + "the surviving credential must be the winner's, not an emptied file" + ); + assert_eq!( + session::load_all(&paths).expect("load").len(), + 1, + "one session, not one per racer" + ); + + let _ = std::fs::remove_dir_all(root); + } + + #[test] + fn a_failed_publish_stops_the_model_it_started() { + // Otherwise a GPU is held by something nobody can call and nothing + // records. + let (root, paths) = temp_paths("unwind-publish"); + let transport = ScriptedTransport::new(vec![ScriptedStep::ok("services stop", "")]); + + let leftovers = + unwind_partial_serve(&transport, &paths, "sess", "rocm", Some("svc-1"), None); + + assert!(leftovers.is_empty(), "{leftovers:?}"); + assert!( + transport.calls().iter().any(|call| matches!( + call, + TransportCall::Exec { command, .. } if command.contains("services stop svc-1") + )), + "the started model should have been stopped: {:?}", + transport.calls() + ); + let _ = std::fs::remove_dir_all(root); + } + + #[test] + fn what_the_unwind_could_not_undo_is_reported_not_swallowed() { + // A discarded cleanup failure is how a machine ends up running a model + // nobody remembers starting. + let (root, paths) = temp_paths("unwind-failed"); + session::store_key(&paths, "sess", "the-key").unwrap(); + let transport = ScriptedTransport::new(vec![ + ScriptedStep::ok("tailscale serve status --json", PUBLISHED_FIXTURE), + ScriptedStep::fails("tailscale serve --tcp=8000 off", 1, "daemon busy"), + ScriptedStep::fails("services stop", 1, "no such service"), + ]); + + let leftovers = unwind_partial_serve( + &transport, + &paths, + "sess", + "rocm", + Some("svc-1"), + Some((8000, 11434)), + ); + + assert_eq!(leftovers.len(), 3, "{leftovers:?}"); + assert!(leftovers[0].contains("still be published"), "{leftovers:?}"); + assert!(leftovers[1].contains("still be running"), "{leftovers:?}"); + + // The credential outlives a cleanup that could not finish. Both steps + // above failed, so the model may still be serving on a published + // endpoint, and this key is the only way to call it. + assert!(leftovers[2].contains("only copy"), "{leftovers:?}"); + assert!( + session::key_path(&paths, "sess").exists(), + "a failed unwind must not delete the key to a service that may still be running" + ); + + let described = + describe_leftovers(anyhow::anyhow!("publish failed"), "gpu-box", &leftovers) + .to_string(); + assert!(described.contains("could not undo"), "{described}"); + let _ = std::fs::remove_dir_all(root); + } + + /// Everything a fully successful `serve` touches on the remote, in call + /// order: the readiness probe, the serve command itself, the service + /// registry lookup that follows it, and the publish that follows that. + fn full_serve_steps() -> Vec { + let mut steps = ready_steps(); + steps.push(ScriptedStep::ok("read -r", "")); + steps.push(ScriptedStep::ok( + "services list --json --all", + r#"[{"service_id":"svc-1","engine":"vllm","model_ref":"m","canonical_model_id":"m", + "host":"127.0.0.1","port":11434,"endpoint_url":"http://127.0.0.1:11434/v1", + "mode":"managed","status":"ready","supervisor_pid":1, + "manifest_path":"/a","log_path":"/b","engine_state_path":"/c", + "created_at_unix_ms":1}]"#, + )); + // `publish` checks status once before writing and once after, to confirm + // the remote actually accepted the forward rather than trusting exit 0. + // Reporting it published both times keeps this fixture representing a + // machine with nothing else competing for the port. + steps.push(ScriptedStep::ok( + "tailscale serve status --json", + PUBLISHED_FIXTURE, + )); + steps.push(ScriptedStep::ok("tailscale serve --bg", "")); + steps + } + + /// The readiness probe steps `bootstrap::ensure_ready_with` needs to find a + /// machine with an existing CLI, ROCm, and Tailscale already present — the + /// same fixture `bootstrap`'s own tests use for a ready machine. + fn ready_steps() -> Vec { + vec![ + ScriptedStep::ok("rocm --version", "rocm 1.2.3"), + ScriptedStep::ok("command -v rocminfo", ""), + ScriptedStep::ok("command -v tailscale", ""), + ScriptedStep::ok("uname -s", "Linux\nx86_64\n"), + ] + } + + #[test] + fn serve_sends_the_key_over_stdin_when_it_starts_the_model() { + // Guards the pairing: the command reads stdin, and the caller actually + // supplies it. Either alone leaves the server without a credential. + // Driving this through `serve_with_transport` rather than calling + // `exec_with_stdin` directly is the point: a regression that stops the + // real call site from passing the key (say, reverting to `None`) has to + // fail this test, not just a check that never sees production code. + let (root, paths) = temp_paths("serve-stdin"); + let transport = ScriptedTransport::new(full_serve_steps()); + + serve_with_transport( + &transport, + &paths, + "gpu-box.example-tailnet.ts.net", + &request(), + ) + .expect("scripted serve"); + + let sent_key = transport.calls().iter().find_map(|call| match call { + TransportCall::Exec { + command, + stdin: Some(key), + } if command.contains("read -r ROCM_SERVE_API_KEY") => Some(key.clone()), + _ => None, + }); + assert!( + sent_key.is_some_and(|key| !key.is_empty()), + "the model-starting command must receive the key over stdin: {:?}", + transport.calls() + ); + let _ = std::fs::remove_dir_all(root); + } + + /// A `services list --json` reply holding one record on the session's port, + /// in the given lifecycle state. + fn listing_with_status(status: &str) -> String { + format!( + r#"[{{"service_id":"svc-1","engine":"vllm","model_ref":"m","canonical_model_id":"m", + "host":"127.0.0.1","port":11434,"endpoint_url":"http://127.0.0.1:11434/v1", + "mode":"managed","status":"{status}","supervisor_pid":1, + "manifest_path":"/a","log_path":"/b","engine_state_path":"/c", + "created_at_unix_ms":100}}]"# + ) + } + + /// Whether the transport was ever asked to declare a forward. + fn issued_a_publish(transport: &ScriptedTransport) -> bool { + transport.calls().iter().any(|call| { + matches!(call, TransportCall::Exec { command, .. } if command.contains("serve --bg")) + }) + } + + #[test] + fn a_machine_whose_registry_command_failed_is_still_asked_about_publishing() { + // Reached-but-failed is not never-reached. The endpoint question is + // still answerable, and reporting "the machine could not be asked" for a + // machine that just answered sends the user to check the network. + let transport = ScriptedTransport::new(vec![ + ScriptedStep::fails("services list --json", 127, "rocm: command not found"), + ScriptedStep::ok("tailscale serve status", PUBLISHED_FIXTURE), + ]); + + let observed = observe(&transport, &sample_record()); + assert_eq!(observed.server, ServerHealth::Error); + assert_eq!( + observed.publish, + publish::PublishObservation::Known(publish::PublishState::Published), + "a reachable machine's publishing state is a fact we can read" + ); + + let rendered = render_status(&render_paths(), &[(sample_record(), observed)]); + assert!(rendered.contains("endpoint published: yes"), "{rendered}"); + assert!( + !rendered.contains("could not be asked"), + "a machine that answered must not be reported as unreachable:\n{rendered}" + ); + } + + #[test] + fn attach_refuses_to_republish_over_a_model_that_is_not_running() { + // The refusal is the whole value of `attach`: re-declaring a forward is + // cheap, so without this guard a dead session is handed an endpoint that + // answers with connection refused — which reads as "up" to everything + // that only checks whether the port is published. + let transport = ScriptedTransport::new(vec![ + ScriptedStep::ok("services list --json", &listing_with_status("failed")), + ScriptedStep::ok("tailscale serve status", "{}"), + ]); + + let error = attach_with_transport(&transport, &sample_record()) + .expect_err("a dead model server must not be re-published over"); + let rendered = format!("{error:#}"); + assert!(rendered.contains("nothing behind it"), "{rendered}"); + assert!( + !issued_a_publish(&transport), + "a refused attach must not have declared a forward: {:?}", + transport.calls() + ); + } + + #[test] + fn attach_republishes_a_healthy_session_whose_endpoint_went_missing() { + // The other half: this is the case `status` sends the user here for, so + // it has to actually issue the publish rather than only decline to fail. + let transport = ScriptedTransport::new(vec![ + ScriptedStep::ok("services list --json", &listing_with_status("ready")), + ScriptedStep::ok("tailscale serve status", "{}"), + ScriptedStep::ok("serve --bg", ""), + ]); + // `publish` confirms by reading the state back, so the status step has to + // answer "absent" first and "ours" afterwards. One scripted reply cannot + // do both, so the confirmation is driven by a second transport below. + let error = attach_with_transport(&transport, &sample_record()) + .expect_err("an unconfirmed publish is an error"); + assert!( + format!("{error:#}").contains("does not report it as active"), + "{error:#}" + ); + assert!( + issued_a_publish(&transport), + "a healthy session must be re-published: {:?}", + transport.calls() + ); + + let confirming = ScriptedTransport::new(vec![ + ScriptedStep::ok("services list --json", &listing_with_status("ready")), + ScriptedStep::ok("tailscale serve status", PUBLISHED_FIXTURE), + ScriptedStep::ok("serve --bg", ""), + ]); + attach_with_transport(&confirming, &sample_record()).expect("a confirmed re-publish"); + } + + /// Put a session and its credential on disk, the way `serve` leaves them. + fn seeded_session(paths: &AppPaths) -> RemoteSessionRecord { + let record = sample_record(); + session::store_key(paths, &record.session_id, "the-key").unwrap(); + record.write(paths).unwrap(); + record + } + + #[test] + fn a_stop_that_cannot_stop_the_model_keeps_the_session_and_says_why() { + // Two properties in one flow, because they are the same decision: the + // record is the only thing on this machine that knows the model's id and + // where it runs, so dropping it while the model is still up leaves a GPU + // held by something the user can no longer name. And the error has to + // carry the remote's own words, or "could not be stopped" reads the same + // whether the machine refused, the service was gone, or ssh never landed. + // + // The endpoint is already absent here — withdrawn out of band, or never + // re-declared after a reboot — so `withdraw` succeeds early and the model + // stop is the only thing that can fail. That isolates the branch under + // test from the withdraw refusal above it. + let (root, paths) = temp_paths("stop-refused"); + let record = seeded_session(&paths); + let transport = ScriptedTransport::new(vec![ + ScriptedStep::ok("tailscale serve status", "{}"), + ScriptedStep::fails("services stop", 3, "engine still shutting down"), + ]); + + let error = stop_with_transport(&transport, &paths, &record, false) + .expect_err("an unconfirmed model stop must not drop the session"); + let rendered = format!("{error:#}"); + assert!( + rendered.contains("engine still shutting down"), + "{rendered}" + ); + assert!(rendered.contains("exit 3"), "{rendered}"); + assert!(rendered.contains("--force"), "{rendered}"); + + assert!( + session::resolve(&paths, &record.session_id).is_ok(), + "the session must still be listed so the user can retry" + ); + assert!( + session::key_path(&paths, &record.session_id).exists(), + "the credential must outlive a failed teardown: it is the only copy" + ); + let _ = std::fs::remove_dir_all(root); + } + + #[test] + fn a_forced_stop_forgets_a_session_the_machine_will_not_confirm() { + // `--force` is for a machine that is gone for good. It must drop the + // record even when neither half could be confirmed — that is its purpose + // — and the withdraw step here is the one that fails, which is the worse + // of the two to forget. + let (root, paths) = temp_paths("stop-forced"); + let record = seeded_session(&paths); + let transport = ScriptedTransport::new(vec![ + ScriptedStep::fails("tailscale serve status", 1, "tailscaled not running"), + ScriptedStep::fails("services stop", 1, "no such service"), + ]); + + stop_with_transport(&transport, &paths, &record, true).expect("a forced stop must succeed"); + + assert!( + session::resolve(&paths, &record.session_id).is_err(), + "a forced stop must forget the session locally" + ); + assert!( + !session::key_path(&paths, &record.session_id).exists(), + "a forgotten session must not leave its credential behind" + ); + let _ = std::fs::remove_dir_all(root); + } + + #[test] + fn a_clean_stop_removes_the_session_and_its_credential() { + let (root, paths) = temp_paths("stop-clean"); + let record = seeded_session(&paths); + let transport = ScriptedTransport::new(vec![ + ScriptedStep::ok("tailscale serve status", "{}"), + ScriptedStep::ok("services stop", ""), + ]); + + stop_with_transport(&transport, &paths, &record, false).expect("a confirmed teardown"); + + assert!( + session::resolve(&paths, &record.session_id).is_err(), + "a confirmed teardown must drop the session" + ); + assert!( + !session::key_path(&paths, &record.session_id).exists(), + "a dropped session must not leave its credential behind" + ); + let _ = std::fs::remove_dir_all(root); + } + + #[test] + fn a_withdrawal_the_remote_will_not_confirm_keeps_the_record() { + // `withdraw` reads the state back, and the machine still reports our + // forward. The record is the only thing on this machine that can find + // that endpoint again, so it has to survive — this is the failure the + // whole teardown path is shaped around. + let (root, paths) = temp_paths("stop-unconfirmed"); + let record = seeded_session(&paths); + let transport = ScriptedTransport::new(vec![ + ScriptedStep::ok("tailscale serve status", PUBLISHED_FIXTURE), + ScriptedStep::ok("tailscale serve --tcp=8000 off", ""), + ScriptedStep::ok("services stop", ""), + ]); + + let error = stop_with_transport(&transport, &paths, &record, false) + .expect_err("a withdrawal the remote does not confirm is not a teardown"); + assert!( + format!("{error:#}").contains("still reports port 8000 as published"), + "{error:#}" + ); + assert!( + session::resolve(&paths, &record.session_id).is_ok(), + "an unconfirmed withdrawal must keep the record: it is the only pointer to a live endpoint" + ); + assert!( + session::key_path(&paths, &record.session_id).exists(), + "and the credential with it — the endpoint may still be answering" + ); + let _ = std::fs::remove_dir_all(root); + } + + #[test] + fn a_stop_failure_names_the_machines_answer_or_the_reason_it_never_answered() { + // The two shapes this module keeps apart everywhere else. + let answered = describe_stop_failure(&Ok(transport::RemoteOutcome { + success: false, + code: Some(2), + stdout: String::new(), + stderr: " no such service ".to_owned(), + })) + .expect("a non-zero exit is a failure"); + assert!(answered.contains("exit 2"), "{answered}"); + assert!(answered.contains("no such service"), "{answered}"); + + let silent = describe_stop_failure(&Ok(transport::RemoteOutcome { + success: false, + code: None, + stdout: String::new(), + stderr: String::new(), + })) + .expect("a signalled exit is a failure"); + assert!(silent.contains("signal"), "{silent}"); + + let unreachable = describe_stop_failure(&Err(anyhow::anyhow!("could not reach gpu-box"))) + .expect("an unreachable machine is a failure"); + assert!( + unreachable.contains("could not reach gpu-box"), + "{unreachable}" + ); + + assert!( + describe_stop_failure(&Ok(transport::RemoteOutcome { + success: true, + code: Some(0), + stdout: String::new(), + stderr: String::new(), + })) + .is_none(), + "a successful stop is not a failure" + ); + } +} diff --git a/apps/rocm/src/remote/provision.rs b/apps/rocm/src/remote/provision.rs new file mode 100644 index 000000000..0a71b5a6c --- /dev/null +++ b/apps/rocm/src/remote/provision.rs @@ -0,0 +1,729 @@ +// Copyright © Advanced Micro Devices, Inc., or its affiliates. +// +// SPDX-License-Identifier: MIT + +//! Putting the ROCm CLI on a machine that does not have it. +//! +//! The obvious approach — copy the binary we are running — is wrong, and +//! quietly so. It only works when both machines share an OS and CPU +//! architecture, and when they do not the copy still lands, still runs as a +//! file, and fails with something unhelpful at the first invocation. +//! +//! So provisioning never copies this machine's binary. It asks the remote to +//! fetch its own build, using the project's own installer, which already knows +//! how to detect a platform and verify what it downloads — and does all of that +//! *on the remote*, for the remote. Only if the remote cannot reach the release +//! host does this machine fetch on its behalf, and then it fetches an artifact +//! built for the remote's platform, not for ours. +//! +//! Which of those two applies is not guessed in advance. "Does this machine +//! have internet" has no reliable signal from the outside, so the remote install +//! is simply attempted, and only *that command* failing selects the fallback. + +use std::path::PathBuf; + +use anyhow::{Context, Result, bail}; + +use super::bootstrap::{REMOTE_CLI_PATH, RemotePlatform}; +use super::transport::Transport; + +/// The installer, carried inside the binary rather than fetched. +/// +/// The fallback path needs an installer on a machine that by definition cannot +/// download one, and pushing the copy we were built with also guarantees the +/// installer and the CLI driving it agree about artifact naming and verification. +const INSTALLER: &str = include_str!("../../../../install.sh"); + +/// Public location of the same installer, for the common path where the remote +/// fetches it itself. +const INSTALLER_URL: &str = "https://raw.githubusercontent.com/ROCm/rocm-cli/main/install.sh"; + +/// Where pushed files land on the remote. A dedicated directory so a failed run +/// leaves something obvious to clean up rather than litter in /tmp. +const REMOTE_STAGING: &str = "$HOME/.rocm/provision"; + +/// How the CLI got onto the remote, for reporting. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) enum Provisioned { + /// The remote downloaded and installed its own build. + RemoteInstaller, + /// This machine fetched a build for the remote's platform and pushed it. + PushedMatchedArtifact { asset: String }, +} + +/// Install the CLI on the remote and return how to invoke it. +pub(crate) fn install_cli( + transport: &dyn Transport, + target: &str, + platform: Option<&RemotePlatform>, + channel: &str, +) -> Result<(String, Provisioned)> { + println!(" remote CLI: not found — installing it on {target} ..."); + + match run_remote_installer(transport, channel) { + Ok(()) => { + verify_remote_cli(transport, target)?; + println!(" remote CLI: installed by the remote itself"); + Ok((REMOTE_CLI_PATH.to_owned(), Provisioned::RemoteInstaller)) + } + Err(remote_error) => { + // Not a guess about connectivity: this specific command failed, so + // fall back to fetching on the remote's behalf. + println!(" remote CLI: the machine could not install it itself, fetching for it ..."); + let platform = platform.context( + "cannot fetch a build for the remote because its OS and CPU architecture \ + could not be determined, and installing a mismatched build would fail \ + in a way that is hard to diagnose", + )?; + let asset = push_matched_artifact(transport, platform, channel).with_context(|| { + format!( + "{target} could not install the CLI itself ({remote_error}), and \ + fetching a matching build for it failed too" + ) + })?; + verify_remote_cli(transport, target)?; + println!( + " remote CLI: installed from a build matching {}-{}", + platform.os, platform.arch + ); + Ok(( + REMOTE_CLI_PATH.to_owned(), + Provisioned::PushedMatchedArtifact { asset }, + )) + } + } +} + +/// Ask the remote to install its own build. +/// +/// Note what this deliberately does *not* forward: an alternate signing key set +/// in this machine's environment. On this path the remote fetches from the +/// release host and verifies with its own pinned trust roots, which is the +/// stricter outcome — a private-mirror key is only meaningful for the artifact +/// *this* machine fetched, which is the fallback path +/// ([`push_matched_artifact`]) and forwards it there. So the "the key we send +/// wins" guarantee is scoped to the fallback by design, and a remote that +/// cannot verify against the pinned roots fails here and falls through to it. +fn run_remote_installer(transport: &dyn Transport, channel: &str) -> Result<()> { + let outcome = transport.exec(&remote_installer_command(channel))?; + if !outcome.success { + bail!("{}", outcome.stderr.trim()); + } + Ok(()) +} + +fn remote_installer_command(channel: &str) -> String { + // Piped to `sh` on the remote so the remote's own platform detection, + // checksum and signature verification all run there, for it. + format!( + "curl -fsSL {INSTALLER_URL} | sh -s -- {}", + super::shell_quote(channel) + ) +} + +/// A staging directory that removes itself when it goes out of scope. +/// +/// Cleanup used to be a single call at the end of the happy path, so every `?` +/// and `bail!` above it left the tree behind — and `--install-rocm` is exactly +/// the command someone retries after a failure, so the leaks accumulated. A +/// `Drop` impl is what makes "on every path" true rather than aspirational. +struct StagingDir(PathBuf); + +impl StagingDir { + fn path(&self) -> &std::path::Path { + &self.0 + } +} + +impl Drop for StagingDir { + fn drop(&mut self) { + // Best-effort: a directory we cannot remove is not worth failing a + // provision that otherwise succeeded. + let _ = std::fs::remove_dir_all(&self.0); + } +} + +/// Fetch a build for the remote's platform on this machine, then push it. +fn push_matched_artifact( + transport: &dyn Transport, + platform: &RemotePlatform, + channel: &str, +) -> Result { + // Held for the whole function: dropping it is what deletes the tree, so it + // must outlive every use of `staging` below. + let staging_dir = StagingDir(tempdir_for_download()?); + let staging = staging_dir.path(); + let asset = download_for(platform, channel, staging)?; + + let remote_dir = REMOTE_STAGING; + transport + .run(&format!("mkdir -p {remote_dir}")) + .context("failed to make a staging directory on the remote")?; + + // The archive travels with its checksum and, when present, its signature, so + // the remote can repeat every check this machine made. Splitting the trust + // chain across two machines must not shorten it. + for (local, remote_name) in [ + (staging.join(&asset), asset.clone()), + ( + staging.join(format!("{asset}.sha256")), + format!("{asset}.sha256"), + ), + (staging.join(format!("{asset}.sig")), format!("{asset}.sig")), + ] { + if !local.exists() { + continue; + } + transport + .push_file(&local, &format!("{remote_dir}/{remote_name}")) + .with_context(|| format!("failed to copy {remote_name} to the remote"))?; + } + + let installer_path = staging.join("install.sh"); + std::fs::write(&installer_path, INSTALLER) + .context("failed to stage the installer for copying")?; + transport + .push_file(&installer_path, &format!("{remote_dir}/install.sh")) + .context("failed to copy the installer to the remote")?; + + // An alternate signing key set in this machine's environment is an escape + // hatch for private mirrors (see install.sh's resolve_public_keys). Forward + // it to the remote's own install.sh, or the remote falls back to the pinned + // production keys and rejects an archive this machine already trusted — + // shortening the trust chain the comment above insists on not shortening. + let signing_env = signing_env_fragment()?.unwrap_or_default(); + + let outcome = transport.exec(&install_from_archive_command( + &signing_env, + remote_dir, + &asset, + channel, + ))?; + if !outcome.success { + bail!( + "the remote rejected the build we fetched for it: {}", + outcome.stderr.trim() + ); + } + // No explicit cleanup: `staging_dir` removes the tree on drop, including + // on the error paths above. + Ok(asset) +} + +/// The remote command that installs from the archive we pushed. +/// +/// Split out so the *composition* is testable, not just the fragment builder. +/// Only the builder was covered before, so dropping `{signing_env}` from the +/// command — the one mistake that would silently take the forwarded trust root +/// out of play — passed the whole suite. +fn install_from_archive_command( + signing_env: &str, + remote_dir: &str, + asset: &str, + channel: &str, +) -> String { + // `remote_dir` stays outside the quoting. It is `$HOME/.rocm/provision`, left + // unexpanded on purpose so the *remote* shell resolves it, and single-quoting + // it suppresses that expansion — the far side then looks for a file whose + // name literally begins `$HOME`, and every push-provisioned install fails at + // `[ -f "${LOCAL_ARCHIVE}" ]`. Only the asset name, which comes from the + // installer's own report, needs quoting; an unquoted prefix concatenated with + // a quoted suffix is one word to the shell. + format!( + "{signing_env}ROCM_CLI_ARCHIVE={remote_dir}/{} sh {remote_dir}/install.sh {}", + super::shell_quote(asset), + super::shell_quote(channel) + ) +} + +/// Build the `NAME=value ` fragment (quoted, trailing space) that forwards an +/// alternate signing key to the remote's install.sh, or `None` if this +/// process's environment sets neither variable. +/// +/// Only `_PEM` ever carries a key across the wire. `_PATH` names a file on +/// *this* machine, and forwarding that path verbatim would tell the remote's +/// shell to open a file that is not there — the variable would be set but +/// useless. So a `_PATH` is read here and its *contents* are sent as `_PEM`. +/// +/// When a key *is* forwarded, `_PATH` goes with it — deliberately empty — so +/// that the forwarded key wins. `resolve_public_keys` consults `_PATH` first, +/// so a value exported on the far side (through `/etc/environment` and pam_env, +/// which apply to non-interactive sshd sessions) would otherwise beat the key +/// we just sent, and the two machines would verify against different trust +/// roots. An explicitly-empty export reads as unset to install.sh's +/// `[ -n ... ]`, which is what makes one token enough to close that. +/// +/// Note the scope: this guarantees *the key we send wins*, not that the +/// remote's own is always off. Forward nothing — the default, with neither +/// variable set here — and the fragment is empty, so a `_PATH` the remote +/// exports for itself still stands while this machine verifies against the +/// pinned keys. Blanking it unconditionally would close that too, at the cost +/// of overriding a remote operator's deliberate mirror-key config in the case +/// where we have no opinion at all. The divergence fails loud ("the remote +/// rejected the build we fetched for it") rather than silently trusting the +/// wrong root, so it is left as a decision rather than assumed. +/// +/// `_PATH` is therefore checked first, because that is the order install.sh's +/// own `resolve_public_keys` uses. These two must not disagree: what they are +/// choosing between is the trust root a signature is verified against, so if +/// they picked differently, a remote provision would accept a build that a +/// local install would reject, and the mismatch would surface as "the remote +/// rejected the build we fetched for it" — pointing at the artifact rather +/// than at the key. +/// +/// An empty value counts as unset, again matching install.sh, which tests +/// these with `[ -n ... ]`. Rust's `env::var` does not make that distinction +/// on its own: `FOO=` yields `Some("")`, which would otherwise forward an +/// empty key and silently discard the operator's real one. +fn signing_env_fragment() -> Result> { + fn set_and_non_empty(name: &str) -> Option { + std::env::var(name).ok().filter(|value| !value.is_empty()) + } + + signing_env_fragment_from( + set_and_non_empty("ROCM_CLI_SIGNING_PUBLIC_KEY_PEM"), + set_and_non_empty("ROCM_CLI_SIGNING_PUBLIC_KEY_PATH"), + |path: &str| std::fs::read_to_string(path), + ) +} + +fn signing_env_fragment_from( + pem_env: Option, + path_env: Option, + read_to_string: impl Fn(&str) -> std::io::Result, +) -> Result> { + // Defence in depth for callers that build these by hand rather than from + // the environment: the empty-is-unset rule belongs to the resolution, not + // to the one caller that happens to read env vars. + let pem_env = pem_env.filter(|value| !value.is_empty()); + let path_env = path_env.filter(|value| !value.is_empty()); + + let pem = match path_env { + Some(path) => Some(read_to_string(&path).with_context(|| { + format!( + "failed to read the signing key at {path} \ + (from ROCM_CLI_SIGNING_PUBLIC_KEY_PATH)" + ) + })?), + None => pem_env, + }; + Ok(pem.map(|pem| { + format!( + "ROCM_CLI_SIGNING_PUBLIC_KEY_PEM={} ROCM_CLI_SIGNING_PUBLIC_KEY_PATH= ", + super::shell_quote(&pem) + ) + })) +} + +/// Run the installer here in download-only mode, targeting the remote's +/// platform, and return the artifact's file name. +fn download_for( + platform: &RemotePlatform, + channel: &str, + into: &std::path::Path, +) -> Result { + let installer = into.join("install.sh"); + std::fs::create_dir_all(into) + .with_context(|| format!("failed to create {}", into.display()))?; + std::fs::write(&installer, INSTALLER).context("failed to stage the installer")?; + + let output = std::process::Command::new("sh") + .arg(&installer) + .arg(channel) + .env("ROCM_CLI_DOWNLOAD_ONLY", "1") + .env("ROCM_CLI_DOWNLOAD_DIR", into) + .env("ROCM_CLI_TARGET_OS", &platform.os) + .env("ROCM_CLI_TARGET_ARCH", &platform.arch) + .output() + .context("failed to run the installer to fetch a build for the remote")?; + + if !output.status.success() { + bail!( + "could not fetch a {}-{} build: {}", + platform.os, + platform.arch, + String::from_utf8_lossy(&output.stderr).trim() + ); + } + + parse_downloaded_asset(&String::from_utf8_lossy(&output.stdout)) + .context("the installer reported success but did not say which file it produced") +} + +/// Pull the artifact name out of the installer's `downloaded:` line. +fn parse_downloaded_asset(stdout: &str) -> Option { + stdout + .lines() + .rev() + .find_map(|line| line.trim().strip_prefix("downloaded:")) + .map(str::trim) + .and_then(|path| path.rsplit('/').next()) + .filter(|name| !name.is_empty()) + .map(ToOwned::to_owned) +} + +fn tempdir_for_download() -> Result { + // A PID-keyed path under the shared, world-writable system temp directory + // is predictable and PIDs get reused, so another local user could pre-stage + // (or symlink) that exact path ahead of us; the old code then either wrote + // the archive and signing material into whatever was already there, or had + // its `remove_dir_all` above follow a planted symlink somewhere unintended. + // Mixing in a nanosecond nonce makes the path unguessable, `create_dir` + // (not `_all`) refuses to silently adopt an existing entry, and 0700 keeps + // the contents unreadable to anyone else even if the name did leak. + let nonce = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map(|d| d.as_nanos()) + .unwrap_or_default(); + let directory = std::env::temp_dir().join(format!( + "rocm-remote-provision-{}-{nonce}", + std::process::id() + )); + create_restricted_dir(&directory) + .with_context(|| format!("failed to create {}", directory.display()))?; + Ok(directory) +} + +/// Create `path` accessible only by its owner, with the mode applied *at* +/// creation. +/// +/// What lands here is the release archive, its checksum and its signature — all +/// public material, so the point is not confidentiality. It is integrity, and +/// the window is narrower than it looks: `download_for` runs install.sh in +/// download-only mode, which fetches and verifies inside its *own* `mktemp -d` +/// and only then copies the proven artifacts out to here. So nothing in this +/// directory is waiting to be checked. +/// +/// What it is waiting for is the push. Between landing and the `scp` above, +/// another local user able to write here could swap the archive together with a +/// matching `.sha256`, and the remote would re-verify the pair it was given. A +/// signature stops that, and today one is always required: `resolve_public_keys` +/// falls back to the pinned release key with no channel argument, so +/// `public_keys` is non-empty on every channel and the `|| [ -n "${public_keys}" ]` +/// arm at `install.sh:431` fires regardless of the channel-gated `require_sig` +/// above it. Clearing both pinned key slots is the only way to reach an unsigned +/// install. 0700 is still what keeps a second local user out of the gap in the +/// meantime — and what keeps this sound if a channel ever ships without a key. +/// +/// (The signing key itself never lands here. It travels in the command prefix, +/// and install.sh writes it into its own `mktemp -d` on the far side.) +/// +/// Creating first and tightening afterwards leaves that window open for the +/// width of the umask. `DirBuilder::mode` closes it; this repo already uses the +/// same pattern in `dash.rs`'s `create_private_dir`, which documents why. +/// +/// Deliberately not recursive: the name carries a nonce, so an existing +/// directory means someone else got there first and must be an error rather +/// than something to adopt. +#[cfg(unix)] +fn create_restricted_dir(path: &std::path::Path) -> Result<()> { + use std::os::unix::fs::DirBuilderExt; + std::fs::DirBuilder::new().mode(0o700).create(path)?; + Ok(()) +} + +#[cfg(not(unix))] +fn create_restricted_dir(path: &std::path::Path) -> Result<()> { + std::fs::create_dir(path)?; + Ok(()) +} + +/// Confirm the freshly-installed CLI actually runs there. +/// +/// The check that catches a build which landed but cannot execute — the failure +/// mode copying our own binary produced silently, and the reason this module +/// exists. +fn verify_remote_cli(transport: &dyn Transport, target: &str) -> Result<()> { + let outcome = transport.exec(&format!("{REMOTE_CLI_PATH} --version"))?; + if !outcome.success { + bail!( + "the CLI was installed on {target} but does not run there: {}\n\ + This usually means the build does not match the machine's OS or CPU.", + outcome.stderr.trim() + ); + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::remote::transport::{ScriptedStep, ScriptedTransport}; + + #[test] + fn the_remote_installs_its_own_build_when_it_can() { + // The common path: the remote's own platform detection and verification + // run on the remote, so nothing here has to reason about its hardware. + let transport = ScriptedTransport::new(vec![ + ScriptedStep::ok("install.sh | sh", ""), + ScriptedStep::ok(".local/bin/rocm --version", "rocm 1.2.3"), + ]); + let (cli, how) = install_cli(&transport, "gpu-box", None, "release").expect("provisioned"); + assert_eq!(cli, REMOTE_CLI_PATH); + assert_eq!(how, Provisioned::RemoteInstaller); + } + + #[test] + fn the_signing_fragment_is_actually_prefixed_onto_the_remote_install_command() { + // The fragment builder was tested; its use was not. A refactor that + // dropped `{signing_env}` from the command would have left every + // signing test passing while the remote silently fell back to the + // pinned keys — the exact failure the forwarding exists to prevent. + let fragment = signing_env_fragment_from(Some("pem-content".to_owned()), None, |path| { + panic!("must not read {path}") + }) + .expect("no file read is attempted") + .expect("a _PEM was set"); + + let command = + install_from_archive_command(&fragment, "$HOME/.rocm/provision", "a.tgz", "release"); + assert!( + command.starts_with("ROCM_CLI_SIGNING_PUBLIC_KEY_PEM="), + "the key must lead the command, before ROCM_CLI_ARCHIVE: {command}" + ); + assert!( + command.contains("ROCM_CLI_SIGNING_PUBLIC_KEY_PATH= "), + "the empty _PATH must ride along, or the remote's own can win: {command}" + ); + assert!(command.contains("ROCM_CLI_ARCHIVE="), "{command}"); + assert!(command.contains("install.sh"), "{command}"); + + // And with nothing set locally, the command carries no signing vars at + // all rather than empty ones. + let bare = install_from_archive_command("", "$HOME/.rocm/provision", "a.tgz", "release"); + assert!(!bare.contains("SIGNING"), "{bare}"); + assert!(bare.starts_with("ROCM_CLI_ARCHIVE="), "{bare}"); + } + + #[test] + fn a_build_that_lands_but_cannot_run_is_caught_and_explained() { + // The exact failure that copying our own binary produced silently. + let transport = ScriptedTransport::new(vec![ + ScriptedStep::ok("install.sh | sh", ""), + ScriptedStep::fails(".local/bin/rocm --version", 126, "Exec format error"), + ]); + let error = install_cli(&transport, "gpu-box", None, "release") + .unwrap_err() + .to_string(); + assert!(error.contains("does not run there"), "{error}"); + assert!(error.contains("OS or CPU"), "{error}"); + } + + #[test] + fn an_unknown_remote_platform_refuses_rather_than_pushing_our_own_build() { + // Without knowing the target's platform the only thing left to send is + // this machine's binary, which is the mistake this module exists to + // avoid. Refuse instead. + let transport = ScriptedTransport::new(vec![ScriptedStep::fails( + "install.sh | sh", + 1, + "could not resolve host", + )]); + let error = install_cli(&transport, "gpu-box", None, "release") + .unwrap_err() + .to_string(); + assert!(error.contains("could not be determined"), "{error}"); + } + + #[test] + fn the_staging_path_still_expands_on_the_remote_shell() { + // `REMOTE_STAGING` is `$HOME/...`, left unexpanded so the far shell + // resolves it. Quoting it suppressed that, and the remote then looked + // for a file literally named `$HOME/...` — every push-provisioned + // install failed at `[ -f "${LOCAL_ARCHIVE}" ]`. + // + // Asked of a real shell rather than asserted about the string: the whole + // defect is what a shell does with the quoting, which is precisely what + // a `contains("ROCM_CLI_ARCHIVE=")` assertion cannot see. + let command = install_from_archive_command( + "", + REMOTE_STAGING, + "rocm-cli-release-linux-x86_64.tar.gz", + "release", + ); + let assignment = command + .split(" sh ") + .next() + .expect("the command opens with the env assignment"); + + let output = std::process::Command::new("sh") + .arg("-c") + .arg(format!( + "HOME=/home/tester; export HOME; \ + {assignment} sh -c 'printf %s \"$ROCM_CLI_ARCHIVE\"'" + )) + .output() + .expect("sh should run"); + + assert_eq!( + String::from_utf8_lossy(&output.stdout), + "/home/tester/.rocm/provision/rocm-cli-release-linux-x86_64.tar.gz", + "the staging path did not expand; the remote would look for a file \ + whose name begins with a literal `$HOME`: {command}" + ); + } + + #[test] + fn a_hostile_asset_name_is_still_quoted_into_the_install_command() { + // The prefix is unquoted so it can expand; the asset name is not, and + // must stay inert. Both halves have to hold at once — fixing the + // expansion by dropping the quoting entirely would trade a broken + // install for an injection. + let command = install_from_archive_command( + "", + REMOTE_STAGING, + "x'; rm -rf ~; echo '.tar.gz", + "release", + ); + assert!(command.contains(r"'\''"), "{command}"); + assert!(command.starts_with("ROCM_CLI_ARCHIVE=$HOME/"), "{command}"); + } + + #[test] + fn the_channel_reaches_the_remote_installer_quoted() { + assert!(remote_installer_command("nightly").contains("sh -s -- nightly")); + // A channel is user input landing in a remote shell like any other. + assert!(remote_installer_command("a; rm -rf /").contains(r"'a; rm -rf /'")); + } + + #[test] + fn the_downloaded_artifact_name_is_read_from_the_installers_own_report() { + let stdout = "rocm-cli installer\n channel: release\n\ + downloaded: /tmp/x/rocm-cli-linux-amd64.tar.gz\n"; + assert_eq!( + parse_downloaded_asset(stdout).as_deref(), + Some("rocm-cli-linux-amd64.tar.gz") + ); + assert_eq!(parse_downloaded_asset("no such line"), None); + } + + #[test] + fn the_installer_is_carried_in_the_binary_so_an_offline_remote_can_still_get_one() { + // A machine that cannot download an artifact cannot download an + // installer either, and shipping the one we were built with keeps the + // installer and this code agreeing about artifact names and checks. + assert!(INSTALLER.contains("rocm-cli installer")); + assert!( + INSTALLER.contains("ROCM_CLI_DOWNLOAD_ONLY"), + "the embedded installer must be the one supporting download-only mode" + ); + assert!(INSTALLER.contains("ROCM_CLI_ARCHIVE")); + } + + #[test] + fn neither_signing_var_set_forwards_nothing() { + let fragment = + signing_env_fragment_from(None, None, |path: &str| std::fs::read_to_string(path)) + .expect("no file read is attempted"); + assert_eq!(fragment, None); + } + + #[test] + fn an_explicit_pem_is_forwarded_as_is_when_it_is_the_only_one_set() { + // The trailing empty `_PATH` is load-bearing, so it is asserted as part + // of the whole fragment rather than trusted. Without it a key could be + // forwarded correctly and still lose: resolve_public_keys reads `_PATH` + // first, so a value the remote exports for itself — /etc/environment + // via pam_env reaches non-interactive sshd sessions — would beat the + // key we just sent. An explicitly-empty export is what install.sh's + // `[ -n ... ]` reads as unset. + let fragment = signing_env_fragment_from(Some("pem-content".to_owned()), None, |path| { + panic!("must not read {path}: no _PATH was set") + }) + .expect("no file read is attempted"); + assert_eq!( + fragment.as_deref(), + Some("ROCM_CLI_SIGNING_PUBLIC_KEY_PEM=pem-content ROCM_CLI_SIGNING_PUBLIC_KEY_PATH= ") + ); + } + + #[test] + fn a_path_wins_over_a_pem_because_that_is_what_install_sh_does() { + // install.sh's resolve_public_keys returns _PATH first and only falls + // through to _PEM. What the two are choosing between is the trust root + // a signature is checked against, so disagreeing here would let a + // remote provision verify against a different key than a local install + // — and the operator would see it as a rejected artifact, not a key + // mismatch. This test is the one that keeps the orders together. + let fragment = signing_env_fragment_from( + Some("pem-content".to_owned()), + Some("/etc/rocm-signing.pem".to_owned()), + |path| { + assert_eq!(path, "/etc/rocm-signing.pem"); + Ok("path-content".to_owned()) + }, + ) + .expect("the fake reader succeeds"); + assert_eq!( + fragment.as_deref(), + Some("ROCM_CLI_SIGNING_PUBLIC_KEY_PEM=path-content ROCM_CLI_SIGNING_PUBLIC_KEY_PATH= ") + ); + } + + #[test] + fn an_empty_value_counts_as_unset_the_way_install_sh_reads_it() { + // install.sh tests both with `[ -n ... ]`, so `FOO=` is unset to it. + // Rust's env::var disagrees — it yields Some(""). Without this, an + // empty _PEM alongside a real _PATH forwarded the empty one, throwing + // away the operator's chosen trust root with no diagnostic and quietly + // falling back to the pinned production keys. + let fragment = signing_env_fragment_from( + Some(String::new()), + Some("/etc/rocm-signing.pem".to_owned()), + |_| Ok("path-content".to_owned()), + ) + .expect("the fake reader succeeds"); + assert_eq!( + fragment.as_deref(), + Some("ROCM_CLI_SIGNING_PUBLIC_KEY_PEM=path-content ROCM_CLI_SIGNING_PUBLIC_KEY_PATH= ") + ); + + // The mirror case: an empty _PATH must not be opened, and must not + // suppress a real _PEM. + let fragment = signing_env_fragment_from( + Some("pem-content".to_owned()), + Some(String::new()), + |path| panic!("must not read {path:?}: an empty _PATH is unset"), + ) + .expect("no file read is attempted"); + assert_eq!( + fragment.as_deref(), + Some("ROCM_CLI_SIGNING_PUBLIC_KEY_PEM=pem-content ROCM_CLI_SIGNING_PUBLIC_KEY_PATH= ") + ); + + // Both empty is the same as neither set. + let fragment = + signing_env_fragment_from(Some(String::new()), Some(String::new()), |path| { + panic!("must not read {path:?}: an empty _PATH is unset") + }) + .expect("no file read is attempted"); + assert_eq!(fragment, None); + } + + #[test] + fn a_path_is_read_locally_and_its_content_is_forwarded_not_the_path() { + // The bug this guards against: forwarding _PATH verbatim names a file + // on this machine, which is meaningless to the remote shell that runs + // install.sh. Only file *content*, sent as _PEM, may cross the wire. + let fragment = + signing_env_fragment_from(None, Some("/etc/rocm-signing.pem".to_owned()), |path| { + assert_eq!(path, "/etc/rocm-signing.pem"); + Ok("-----BEGIN PUBLIC KEY-----\nabc\n-----END PUBLIC KEY-----\n".to_owned()) + }) + .expect("the fake reader succeeds"); + let fragment = fragment.expect("a _PATH was set"); + assert!(fragment.starts_with("ROCM_CLI_SIGNING_PUBLIC_KEY_PEM=")); + assert!(!fragment.contains("/etc/rocm-signing.pem")); + assert!(fragment.contains("BEGIN PUBLIC KEY")); + } + + #[test] + fn a_path_that_cannot_be_read_is_reported_rather_than_silently_dropped() { + let error = signing_env_fragment_from(None, Some("/no/such/file".to_owned()), |_| { + Err(std::io::Error::other("boom")) + }) + .unwrap_err() + .to_string(); + assert!(error.contains("/no/such/file"), "{error}"); + } +} diff --git a/apps/rocm/src/remote/publish.rs b/apps/rocm/src/remote/publish.rs new file mode 100644 index 000000000..87756e885 --- /dev/null +++ b/apps/rocm/src/remote/publish.rs @@ -0,0 +1,942 @@ +// Copyright © Advanced Micro Devices, Inc., or its affiliates. +// +// SPDX-License-Identifier: MIT + +//! Publishing a remote machine's loopback service onto the tailnet. +//! +//! This is the data path, and it is declared *by the remote*, not held open by +//! us. `rocm serve` binds `127.0.0.1` on the GPU machine as it always has; the +//! machine then tells its own Tailscale daemon to accept tailnet connections on +//! a port and forward them to that loopback address. Nothing on this end stays +//! running, which is why an endpoint survives the command that created it and +//! is reachable from the user's other machines rather than only this one. +//! +//! The cost of that, and the reason [`super`] insists on a credential: a +//! publish is visible to the whole tailnet, scoped only by its ACLs. That +//! assumes Tailscale Funnel is not already enabled for the port — Funnel is +//! what turns a tailnet-scoped forward into one reachable from the public +//! internet, and it is not something this module ever asks for. If the remote +//! already has it allowed for the target port, [`classify`] reports it as a +//! distinct state rather than folding it into "free", and both [`publish`] +//! and [`withdraw`] refuse to act until it is turned off by hand. Unlike a +//! point-to-point tunnel a publish also *outlives a reboot*, because it is +//! configuration rather than a process. A publish left behind is a GPU endpoint +//! nobody is tracking, so withdrawal is treated as a first-class operation that +//! reports failure loudly instead of being assumed to have worked. +//! +//! **Unverified against a live tailnet.** The command shapes below follow +//! Tailscale's documented surface, and the parsing follows the `ServeConfig` +//! struct definition, but neither has been run against a real daemon here. +//! Confirm both before relying on this. + +use anyhow::{Result, bail}; +use serde::Deserialize; +use std::collections::BTreeMap; + +use super::transport::Transport; + +/// Keys a real serve config carries. Seeing none of them in a non-empty +/// document means we were handed something else, whatever it parses as. +/// +/// Not all of these are actually inspected. `TCP`, `Foreground`, `Services`, +/// and `AllowFunnel` are parsed into [`RawServeConfig`] and drive +/// [`classify`], and a `TCP` entry's `HTTPS`/`HTTP`/`TerminateTLS` fields are +/// read too — not because this design ever asks Tailscale to terminate TLS, +/// but because someone else's handler still occupies the port we want. +/// `Web` is listed only as document-shape evidence: it confirms we are looking +/// at a real serve config, and a `Web` handler is reached through a `TCP` entry +/// we already inspect. Listing a key here without parsing it is exactly what let +/// `AllowFunnel` go unchecked for a release, so if a key stays evidence-only, +/// say so here rather than leaving a reader to assume otherwise. +const SERVE_CONFIG_KEYS: &[&str] = &["TCP", "Web", "Services", "AllowFunnel", "Foreground"]; + +/// Loopback address a published port forwards to. The model server binds here +/// and nowhere else; the publish is the only thing that widens its reach. +pub(crate) const LOOPBACK: &str = "127.0.0.1"; + +/// `tailscale serve status --json`, as much of it as we read. +/// +/// Only the TCP forwards matter: this design never asks Tailscale to terminate +/// TLS or serve HTTP on our behalf, because the model server already speaks the +/// protocol the caller wants and putting a proxy in between would only add a +/// place for the two to disagree. +#[derive(Debug, Default, Deserialize)] +struct RawServeConfig { + /// Keyed by port. Go renders integer map keys as strings, so these arrive + /// as `"8000"` rather than `8000`. + #[serde(rename = "TCP", default)] + tcp: BTreeMap, + /// Per-session configuration, used when a serve was started in the + /// foreground. We always publish in the background, so anything here + /// belongs to someone else — but a forward is a forward, and missing one + /// would report a live endpoint as absent. + #[serde(rename = "Foreground", default)] + foreground: BTreeMap, + /// Tailscale "Services" (VIP services), keyed by service name. Nests its + /// own `TCP` map the same way `Foreground` does — a forward declared here + /// is still a forward, and missing it has the same failure mode as + /// missing a foreground one: a live foreign endpoint reads as absent, and + /// the `Foreign` ownership guard never fires. + #[serde(rename = "Services", default)] + services: BTreeMap, + /// Keyed `host:port`. `true` means the port is exposed to the public + /// internet via Tailscale Funnel, not just the tailnet — something this + /// module never asks for. We only ever check this for our own target + /// port, so we do not track which host set it. + #[serde(rename = "AllowFunnel", default)] + allow_funnel: BTreeMap, +} + +/// Shape shared by `Foreground` sessions and `Services` entries: both nest a +/// serve config under their own key, so both can carry a `TCP` map *and* their +/// own `AllowFunnel`. +/// +/// Parsing `AllowFunnel` here too is not symmetry for its own sake. A nested +/// entry was already trusted to carry a forward — that is why `tcp` is read — +/// and reading the forward while ignoring the Funnel permission beside it is +/// the same asymmetry that let `AllowFunnel` go unchecked for a release at the +/// top level. `tailscale funnel ` in its foreground form is a common way +/// for the permission to be set, so the nesting the forward can hide in is the +/// nesting the permission can hide in. +#[derive(Debug, Default, Deserialize)] +struct RawNestedConfig { + #[serde(rename = "TCP", default)] + tcp: BTreeMap, + #[serde(rename = "AllowFunnel", default)] + allow_funnel: BTreeMap, +} + +#[derive(Debug, Default, Deserialize, Clone)] +struct RawTcpHandler { + /// Destination, as `host:port`. Absent for a TLS-terminating handler, which + /// is not something we create — but is something we must recognise, because + /// it still occupies the port. + #[serde(rename = "TCPForward", default)] + tcp_forward: Option, + /// Tailscale terminates TLS and serves web content on this port. + #[serde(rename = "HTTPS", default)] + https: bool, + /// As `HTTPS`, without TLS. + #[serde(rename = "HTTP", default)] + http: bool, + /// Terminates TLS for the named host and hands the plaintext on. + #[serde(rename = "TerminateTLS", default)] + terminate_tls: Option, +} + +impl RawTcpHandler { + /// What to tell the user is sitting on the port, when it is not a forward. + /// + /// Named rather than described as "something", because the refusal it feeds + /// is the user's only clue about what they would have destroyed. + fn describe_holder(&self) -> String { + if self.https { + "an existing Tailscale HTTPS handler".to_owned() + } else if self.http { + "an existing Tailscale HTTP handler".to_owned() + } else if self.terminate_tls.is_some() { + "an existing Tailscale TLS-terminating handler".to_owned() + } else { + // A handler we cannot name is still a handler. Refusing on it is the + // safe direction: a newer Tailscale adding a kind must not read as + // an empty port. + "an existing Tailscale handler this CLI does not recognise".to_owned() + } + } +} + +/// What the remote's Tailscale says about one port. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) enum PublishState { + /// The port forwards to the loopback address we expect. + Published, + /// The port is not forwarded at all. + Absent, + /// The port forwards somewhere else. Not ours to withdraw, and a warning + /// that two things are competing for it. + Foreign { forwards_to: String }, + /// Tailscale Funnel is allowed for this port, regardless of what (if + /// anything) forwards to it. Funnel is what exposes a port to the public + /// internet rather than just the tailnet, and this module never turns it + /// on. Deliberately not folded into `Absent` or `Published`: publishing + /// over it would complete a public exposure nobody asked this command + /// for, and withdrawing our forward while it stays on would not close + /// anything. + FunnelAllowed, + /// The remote answered with something we could not read. + /// + /// Deliberately not folded into `Absent`. "I looked and there is no + /// forward" and "I could not tell" differ exactly where it matters: the + /// first confirms a withdrawal, the second must not, or a malformed reply + /// becomes a report that an endpoint is gone while it is still published. + Unreadable, +} + +/// Ask the remote which of its ports are forwarded. +pub(crate) fn publish_state( + transport: &dyn Transport, + tailnet_port: u16, + remote_port: u16, +) -> Result { + match observe(transport, tailnet_port, remote_port) { + PublishObservation::Known(state) => Ok(state), + PublishObservation::Failed(why) | PublishObservation::Unreachable(why) => bail!("{why}"), + } +} + +/// Why an attempt to read the publishing state did not produce one. +/// +/// Separate from [`PublishState`] because these are not states of the *port* — +/// they are states of our attempt to look at it. Folding them in would let +/// "I could not ask" render as a fact about the endpoint. +/// +/// The split between the two failure arms matters for the same reason +/// [`super::ServerHealth`] splits `Error` from `Unreachable`: reached-but-failed +/// carries the remote's own words about what is wrong, and never-reached does +/// not. Different problems, different fixes, so different words. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) enum PublishObservation { + /// The remote answered and the reply was classified. + Known(PublishState), + /// Reached the machine, but `tailscale serve status` exited non-zero. + /// Carries the exit status and the remote's stderr. + Failed(String), + /// Could not reach the machine at all. Carries the transport's reason. + Unreachable(String), +} + +/// Read the remote's publishing state, reporting *how* the attempt failed +/// rather than only that it did. +/// +/// Mirrors the three-way match `super::observe` already does for +/// `services list`: answered-and-succeeded, answered-and-failed, never +/// answered. +pub(crate) fn observe( + transport: &dyn Transport, + tailnet_port: u16, + remote_port: u16, +) -> PublishObservation { + match transport.exec("tailscale serve status --json") { + Ok(outcome) if outcome.success => { + PublishObservation::Known(classify(&outcome.stdout, tailnet_port, remote_port)) + } + Ok(outcome) => PublishObservation::Failed(format!( + "could not read the remote's tailnet publishing state (exit {}): {}", + outcome + .code + .map_or_else(|| "signal".to_owned(), |code| code.to_string()), + outcome.stderr.trim() + )), + Err(error) => PublishObservation::Unreachable(format!("{error:#}")), + } +} + +/// Decide what a serve-status document says about one port. Pure, so the whole +/// classification is testable against fixtures. +fn classify(status_json: &str, tailnet_port: u16, remote_port: u16) -> PublishState { + let trimmed = status_json.trim(); + // A machine publishing nothing prints an empty document; some versions + // print literal `null` for an unset config. Neither is an error. + if trimmed.is_empty() || trimmed == "null" { + return PublishState::Absent; + } + // Parse loosely first. `RawServeConfig` defaults every field, so a document + // that is valid JSON but not a serve config — an error object, a newer shape + // we do not know — would deserialize to an empty config and read as "nothing + // published". After a withdrawal that is indistinguishable from success, + // which is the failure this whole module exists to avoid. + let Ok(document) = serde_json::from_str::(trimmed) else { + return PublishState::Unreadable; + }; + let Some(fields) = document.as_object() else { + return PublishState::Unreadable; + }; + // An empty object is a real, legitimate answer: nothing is published. + if !fields.is_empty() + && !fields + .keys() + .any(|key| SERVE_CONFIG_KEYS.contains(&key.as_str())) + { + return PublishState::Unreadable; + } + let Ok(config) = serde_json::from_value::(document) else { + return PublishState::Unreadable; + }; + + let key = tailnet_port.to_string(); + + // Checked before we ask whether anything forwards to the port at all: a + // Funnel-enabled port is a distinct hazard independent of what (if + // anything) currently forwards to it, and must never be read as merely + // "free" or folded into a normal `Published`/`Absent` result. + // + // Every nesting a forward can hide in is searched, for the same reason the + // forward lookup below searches them: a permission we do not look at reads + // as one that is not set. + let funnel_anywhere = funnel_allowed_for_port(&config.allow_funnel, tailnet_port) + || config + .foreground + .values() + .chain(config.services.values()) + .any(|nested| funnel_allowed_for_port(&nested.allow_funnel, tailnet_port)); + if funnel_anywhere { + return PublishState::FunnelAllowed; + } + + let handler = config + .tcp + .get(&key) + .or_else(|| { + config + .foreground + .values() + .find_map(|session| session.tcp.get(&key)) + }) + .or_else(|| config.services.values().find_map(|svc| svc.tcp.get(&key))); + + // An entry at this port means something holds it, whether or not that + // something is a forward. A TLS-terminating handler (`HTTPS`, `HTTP`, + // `TerminateTLS`) carries no `TCPForward` at all, so reading only that field + // and mapping its absence to `Absent` reports an occupied port as free — and + // `publish` takes a free port, destroying whatever was serving there. + // + // The question this answers is "is anything holding this port", not "is this + // our forward". Only the first decides whether it is safe to write. + let Some(handler) = handler else { + return PublishState::Absent; + }; + let Some(forward) = handler.tcp_forward.clone() else { + return PublishState::Foreign { + forwards_to: handler.describe_holder(), + }; + }; + + let expected = forward_target(remote_port); + if forward == expected { + PublishState::Published + } else { + PublishState::Foreign { + forwards_to: forward, + } + } +} + +/// Where a published port should point: the model server's loopback bind. +fn forward_target(remote_port: u16) -> String { + format!("{LOOPBACK}:{remote_port}") +} + +/// True if Funnel is allowed for `port` in an `AllowFunnel` map. +/// +/// `AllowFunnel` is keyed `host:port`, where the host is a tailnet DNS name we +/// do not otherwise track. We only care whether *our* port is exposed, so we +/// match on the port suffix rather than requiring an exact key. +fn funnel_allowed_for_port(allow_funnel: &BTreeMap, port: u16) -> bool { + let port = port.to_string(); + allow_funnel + .iter() + .any(|(host_port, allowed)| *allowed && host_port.rsplit(':').next() == Some(port.as_str())) +} + +/// Command that declares the forward on the remote. +/// +/// `--bg` is what makes it outlive the SSH command that issued it; without it +/// the publish would die with our connection and the endpoint would vanish the +/// moment `serve` returned. +fn publish_command(tailnet_port: u16, remote_port: u16) -> String { + format!( + "tailscale serve --bg --tcp={tailnet_port} tcp://{}", + forward_target(remote_port) + ) +} + +/// Command that removes the forward. +fn withdraw_command(tailnet_port: u16) -> String { + format!("tailscale serve --tcp={tailnet_port} off") +} + +/// Claim the port, declare the forward, then confirm the remote agrees. +/// +/// Ownership is established *before* writing, not after. `tailscale serve` +/// overwrites whatever holds a port without complaint, so checking afterwards +/// is too late — by then the other forward is already gone and the state we +/// read back is our own, which reads as success. A second session reusing a +/// port would silently take the first one's endpoint away. +/// +/// The confirmation afterwards is still needed, and is not ceremony: +/// `tailscale serve` can exit zero while the tailnet's policy declines to +/// publish, and trusting the exit code hands the user a URL that never answers. +pub(crate) fn publish( + transport: &dyn Transport, + tailnet_port: u16, + remote_port: u16, +) -> Result<()> { + match publish_state(transport, tailnet_port, remote_port)? { + // Free, or already pointing where we want it. Re-declaring our own is + // harmless and keeps `attach` idempotent. + PublishState::Absent | PublishState::Published => {} + PublishState::Foreign { forwards_to } => bail!( + "port {tailnet_port} on the remote already forwards to {forwards_to}.\n\ + Refusing to take it over — publishing here would silently break whatever \ + is using it. Choose another port with `--tailnet-port`." + ), + PublishState::FunnelAllowed => bail!( + "port {tailnet_port} on the remote has Tailscale Funnel allowed, which exposes it \ + to the public internet rather than just the tailnet.\n\ + Refusing to publish over it — turn Funnel off first: \ + `tailscale funnel --tcp={tailnet_port} off`." + ), + PublishState::Unreadable => bail!( + "port {tailnet_port} could not be checked before publishing, so there is no way \ + to tell whether something else is already using it.\n\ + Check it by hand: tailscale serve status" + ), + } + + let outcome = transport.exec(&publish_command(tailnet_port, remote_port))?; + if !outcome.success { + bail!( + "the remote refused to publish port {tailnet_port} on the tailnet: {}\n\ + This is usually the tailnet's own policy. Check that the machine is allowed \ + to serve, then try again.", + outcome.stderr.trim() + ); + } + + match publish_state(transport, tailnet_port, remote_port)? { + PublishState::Published => Ok(()), + PublishState::Absent => bail!( + "the remote accepted the publish for port {tailnet_port} but does not report it \ + as active, so the endpoint would not answer" + ), + // Something took the port between our check and our write. + PublishState::Foreign { forwards_to } => bail!( + "port {tailnet_port} on the remote now forwards to {forwards_to} rather than to \ + this model server; something else claimed it. Choose another port with \ + `--tailnet-port`." + ), + // Funnel was turned on between our check and our write. Reported the + // same way as the pre-check case: publishing must not be allowed to + // complete a public-internet exposure nobody asked for. + PublishState::FunnelAllowed => bail!( + "port {tailnet_port} on the remote now has Tailscale Funnel allowed, exposing it to \ + the public internet rather than just the tailnet.\n\ + Turn Funnel off: `tailscale funnel --tcp={tailnet_port} off`, then try again." + ), + PublishState::Unreadable => bail!( + "the remote accepted the publish for port {tailnet_port} but its reply could \ + not be read, so there is no way to confirm the endpoint answers.\n\ + Check it with `rocm remote status`." + ), + } +} + +/// Remove the forward, and confirm it is gone. +/// +/// Returns an error when withdrawal cannot be confirmed. Callers must not treat +/// that as cosmetic: because a publish is configuration rather than a process, +/// an unwithdrawn one survives reboots and keeps a GPU endpoint on the tailnet +/// with nothing tracking it. +pub(crate) fn withdraw( + transport: &dyn Transport, + tailnet_port: u16, + remote_port: u16, +) -> Result<()> { + // Establish it is ours before turning it off. `tailscale serve … off` takes a + // port, not a forward, so it would happily tear down whatever is on that + // port — including something another tool or another person put there after + // our session was recorded. + match publish_state(transport, tailnet_port, remote_port)? { + PublishState::Published => {} + // Already gone. Nothing to do, and nothing to complain about: teardown + // has to be safe to retry after a partial one. + PublishState::Absent => return Ok(()), + PublishState::Foreign { forwards_to } => bail!( + "port {tailnet_port} on the remote now forwards to {forwards_to}, not to this \ + session's model server.\n\ + Refusing to turn it off — it belongs to something else." + ), + PublishState::FunnelAllowed => bail!( + "port {tailnet_port} on the remote has Tailscale Funnel allowed. Withdrawing our \ + forward would not close the public-internet exposure, so this needs a human \ + decision, not a silent teardown.\n\ + Turn Funnel off first: `tailscale funnel --tcp={tailnet_port} off`." + ), + PublishState::Unreadable => bail!( + "port {tailnet_port} could not be checked before withdrawing it, so there is no \ + way to tell whether it is still this session's endpoint.\n\ + Check it by hand: tailscale serve status" + ), + } + + let outcome = transport.exec(&withdraw_command(tailnet_port))?; + if !outcome.success { + bail!( + "failed to withdraw port {tailnet_port} on the remote: {}", + outcome.stderr.trim() + ); + } + match publish_state(transport, tailnet_port, remote_port)? { + PublishState::Absent | PublishState::Foreign { .. } => Ok(()), + PublishState::Published => { + bail!("the remote still reports port {tailnet_port} as published after withdrawing it") + } + // Our forward is gone, but Funnel is still allowed for the port. The + // port may still be reachable from the public internet, so this is + // not the clean close the caller asked for. + PublishState::FunnelAllowed => bail!( + "port {tailnet_port} was turned off, but Tailscale Funnel is still allowed for it, \ + so it may still be reachable from the public internet.\n\ + Turn it off: `tailscale funnel --tcp={tailnet_port} off`." + ), + // An unreadable reply is not a withdrawal. Accepting it here would be the + // exact failure this function exists to prevent: reporting an endpoint + // gone while it is still published, on a machine nobody is watching. + PublishState::Unreadable => bail!( + "port {tailnet_port} was asked to stop publishing, but the remote's reply could \ + not be read, so it cannot be confirmed withdrawn.\n\ + Check it by hand: tailscale serve status" + ), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::remote::transport::{ScriptedStep, ScriptedTransport, TransportCall}; + + const PUBLISHED: &str = r#"{ + "TCP": { "8000": { "TCPForward": "127.0.0.1:11434" } } + }"#; + + #[test] + fn output_we_cannot_read_is_never_reported_as_nothing_published() { + // The distinction the module turns on. Folded into `Absent`, a garbled + // or unexpected reply after a withdrawal reads as "gone" and the + // endpoint stays up with nobody tracking it. + // + // Two shapes matter. Not JSON at all: + for document in ["not json", "{\"TCP\": ", "an error page"] { + assert_eq!( + classify(document, 8000, 11434), + PublishState::Unreadable, + "{document:?}" + ); + } + // And valid JSON that is not a serve config. Every field defaults, so + // without the key check these deserialize to an empty config and look + // like a machine publishing nothing. + for document in [ + r#"{"error": "not logged in"}"#, + r#"{"SomeFutureShape": {"TCP": {}}}"#, + "[]", + r#""a string""#, + "42", + ] { + assert_eq!( + classify(document, 8000, 11434), + PublishState::Unreadable, + "{document}" + ); + } + + // But an empty object is a real answer, and so is a config carrying a + // key we know even when our port is absent from it. + assert_eq!(classify("{}", 8000, 11434), PublishState::Absent); + assert_eq!( + classify(r#"{"Web": {}}"#, 8000, 11434), + PublishState::Absent + ); + } + + #[test] + fn neither_publish_nor_withdraw_accepts_an_unreadable_reply() { + // Withdrawal is the dangerous half: accepting this reports an endpoint + // torn down while it is still live. + let garbled = ScriptedTransport::new(vec![ + ScriptedStep::ok("tailscale serve --tcp=8000 off", ""), + ScriptedStep::ok("tailscale serve status --json", r#"{"unexpected": true}"#), + ]); + let error = withdraw(&garbled, 8000, 11434).unwrap_err().to_string(); + assert!(error.contains("could not be checked"), "{error}"); + + let publishing = ScriptedTransport::new(vec![ + ScriptedStep::ok("tailscale serve --bg", ""), + ScriptedStep::ok("tailscale serve status --json", r#"{"unexpected": true}"#), + ]); + let error = publish(&publishing, 8000, 11434).unwrap_err().to_string(); + assert!(error.contains("could not be checked"), "{error}"); + } + + #[test] + fn a_matching_forward_is_recognised_as_ours() { + assert_eq!(classify(PUBLISHED, 8000, 11434), PublishState::Published); + } + + #[test] + fn an_empty_or_null_config_means_nothing_is_published() { + // A machine that has never published prints an empty or null document. + // Treating that as a parse failure would turn a normal state into noise. + for document in ["", " ", "null", "{}"] { + assert_eq!( + classify(document, 8000, 11434), + PublishState::Absent, + "{document:?}" + ); + } + } + + #[test] + fn a_port_pointing_elsewhere_is_not_treated_as_ours() { + // Withdrawing this would tear down whatever else is using the port. + let other = r#"{"TCP": {"8000": {"TCPForward": "127.0.0.1:9999"}}}"#; + assert_eq!( + classify(other, 8000, 11434), + PublishState::Foreign { + forwards_to: "127.0.0.1:9999".to_owned() + } + ); + } + + #[test] + fn an_https_handler_holds_the_port_even_though_it_is_not_a_forward() { + // A handler with no TCPForward terminates TLS and serves web content. It + // is not a passthrough to our model server — and it is also not nothing. + // Reading it as `Absent` told `publish` the port was free, and publishing + // over it destroys whatever was being served there, which is exactly what + // `Foreign` exists to refuse. + for (document, expected) in [ + ( + r#"{"TCP": {"8000": {"HTTPS": true}}}"#, + "an existing Tailscale HTTPS handler", + ), + ( + r#"{"TCP": {"8000": {"HTTP": true}}}"#, + "an existing Tailscale HTTP handler", + ), + ( + r#"{"TCP": {"8000": {"TerminateTLS": "box.tail1234.ts.net"}}}"#, + "an existing Tailscale TLS-terminating handler", + ), + // A kind this CLI has never heard of still holds the port. Refusing + // is the safe direction when a newer Tailscale adds one. + ( + r#"{"TCP": {"8000": {"SomeFutureHandler": true}}}"#, + "an existing Tailscale handler this CLI does not recognise", + ), + ] { + assert_eq!( + classify(document, 8000, 11434), + PublishState::Foreign { + forwards_to: expected.to_owned() + }, + "{document}" + ); + } + + // The port genuinely being absent is still absent — the guard must not + // have become "refuse whenever the document mentions the port". + assert_eq!( + classify(r#"{"TCP": {"9999": {"HTTPS": true}}}"#, 8000, 11434), + PublishState::Absent + ); + } + + #[test] + fn a_held_port_is_refused_rather_than_taken_over() { + // The consequence that makes the classification matter: `publish` treats + // `Absent` as free to claim, and `tailscale serve` overwrites without + // complaint. + let transport = ScriptedTransport::new(vec![ScriptedStep::ok( + "tailscale serve status", + r#"{"TCP": {"8000": {"HTTPS": true}}}"#, + )]); + let error = publish(&transport, 8000, 11434) + .expect_err("a port held by someone else must not be taken over"); + let rendered = format!("{error:#}"); + assert!(rendered.contains("HTTPS handler"), "{rendered}"); + assert!(rendered.contains("--tailnet-port"), "{rendered}"); + assert!( + !transport + .calls() + .iter() + .any(|call| matches!(call, TransportCall::Exec { command, .. } if command.contains("serve --bg"))), + "nothing may be written to a port we refused: {:?}", + transport.calls() + ); + } + + #[test] + fn a_foreground_publish_is_still_a_publish() { + // We always publish in the background, but a forward someone else + // started in the foreground still answers on that port. Missing it would + // report a live endpoint as absent. + let foreground = r#"{ + "Foreground": { "sess-1": { "TCP": { "8000": { "TCPForward": "127.0.0.1:11434" } } } } + }"#; + assert_eq!(classify(foreground, 8000, 11434), PublishState::Published); + } + + #[test] + fn ports_are_matched_as_string_keys_not_numbers() { + // Go renders integer map keys as strings; looking for a numeric key + // would silently never match and report every endpoint as absent. + assert_eq!(classify(PUBLISHED, 8001, 11434), PublishState::Absent); + } + + #[test] + fn a_service_forward_is_still_a_forward() { + // Same blind spot as `Foreground`: a forward declared under a named + // Service still answers on that port. Missing it defeats the + // `Foreign` ownership guard for whatever it points to. + let service = r#"{ + "Services": { "svc:my-app": { "TCP": { "8000": { "TCPForward": "127.0.0.1:9999" } } } } + }"#; + assert_eq!( + classify(service, 8000, 11434), + PublishState::Foreign { + forwards_to: "127.0.0.1:9999".to_owned() + } + ); + } + + #[test] + fn our_own_forward_under_a_service_is_recognised() { + let service = r#"{ + "Services": { "svc:my-app": { "TCP": { "8000": { "TCPForward": "127.0.0.1:11434" } } } } + }"#; + assert_eq!(classify(service, 8000, 11434), PublishState::Published); + } + + #[test] + fn funnel_allowed_on_our_port_is_never_read_as_free() { + // No matching TCPForward at all: without the AllowFunnel check this + // reads as `Absent`, and `publish()` would happily complete the + // exposure it was never asked to create. + let funnel_only = r#"{"AllowFunnel": {"my-machine.tail1234.ts.net:443": true}}"#; + assert_eq!( + classify(funnel_only, 443, 11434), + PublishState::FunnelAllowed + ); + } + + #[test] + fn funnel_allowed_overrides_a_matching_forward() { + // Even when our own forward is in place, Funnel being allowed means + // the port is reachable from the public internet, not just the + // tailnet. That must not be reported as an ordinary `Published`. + let both = r#"{ + "TCP": { "443": { "TCPForward": "127.0.0.1:11434" } }, + "AllowFunnel": { "my-machine.tail1234.ts.net:443": true } + }"#; + assert_eq!(classify(both, 443, 11434), PublishState::FunnelAllowed); + } + + #[test] + fn funnel_allowed_on_another_port_does_not_affect_ours() { + let other_port = r#"{"AllowFunnel": {"my-machine.tail1234.ts.net:8443": true}}"#; + assert_eq!(classify(other_port, 443, 11434), PublishState::Absent); + } + + #[test] + fn funnel_allowed_inside_a_foreground_session_is_still_funnel() { + // `tailscale funnel ` run in the foreground records its permission + // under that session rather than at the top level. The forward lookup + // already searches this nesting; a guard that does not search it reports + // a publicly-exposed port as free. + let nested = r#"{"Foreground": {"sess-1": { + "AllowFunnel": {"my-machine.tail1234.ts.net:443": true} + }}}"#; + assert_eq!(classify(nested, 443, 11434), PublishState::FunnelAllowed); + } + + #[test] + fn funnel_allowed_inside_a_service_entry_is_still_funnel() { + let nested = r#"{"Services": {"svc:web": { + "AllowFunnel": {"my-machine.tail1234.ts.net:8443": true} + }}}"#; + assert_eq!(classify(nested, 8443, 11434), PublishState::FunnelAllowed); + } + + #[test] + fn a_nested_funnel_on_another_port_leaves_ours_alone() { + // The guard must not become "any Funnel anywhere refuses everything": + // exposure is per-port, so a Funnel on 8443 says nothing about 443. + let nested = r#"{"Foreground": {"sess-1": { + "AllowFunnel": {"my-machine.tail1234.ts.net:8443": true}, + "TCP": {"443": {"TCPForward": "127.0.0.1:11434"}} + }}}"#; + assert_eq!(classify(nested, 443, 11434), PublishState::Published); + } + + #[test] + fn funnel_disabled_entry_does_not_trip_the_guard() { + // The map can carry `false` entries for a port Funnel was allowed for + // and then turned off. Only `true` matters. + let disabled = r#"{"AllowFunnel": {"my-machine.tail1234.ts.net:443": false}}"#; + assert_eq!(classify(disabled, 443, 11434), PublishState::Absent); + } + + #[test] + fn publish_refuses_a_funnel_enabled_port() { + let transport = ScriptedTransport::new(vec![ScriptedStep::ok( + "tailscale serve status --json", + r#"{"AllowFunnel": {"my-machine.tail1234.ts.net:443": true}}"#, + )]); + let error = publish(&transport, 443, 11434).unwrap_err().to_string(); + assert!(error.contains("tailscale funnel --tcp=443 off"), "{error}"); + // Nothing should have been written. + assert!( + !transport.calls().iter().any(|call| matches!( + call, + crate::remote::transport::TransportCall::Exec { command, .. } + if command.contains("--bg") + )), + "{:?}", + transport.calls() + ); + } + + #[test] + fn withdraw_refuses_a_funnel_enabled_port() { + let transport = ScriptedTransport::new(vec![ScriptedStep::ok( + "tailscale serve status --json", + r#"{"AllowFunnel": {"my-machine.tail1234.ts.net:443": true}}"#, + )]); + let error = withdraw(&transport, 443, 11434).unwrap_err().to_string(); + assert!(error.contains("tailscale funnel --tcp=443 off"), "{error}"); + assert!( + !transport.calls().iter().any(|call| matches!( + call, + crate::remote::transport::TransportCall::Exec { command, .. } + if command.contains(" off") + )), + "{:?}", + transport.calls() + ); + } + + #[test] + fn publishing_runs_in_the_background_and_targets_loopback() { + // Without --bg the forward dies with the SSH command that made it, and + // the endpoint vanishes the moment serve returns. + let command = publish_command(8000, 11434); + assert!(command.contains("--bg"), "{command}"); + assert!(command.contains("--tcp=8000"), "{command}"); + assert!(command.contains("tcp://127.0.0.1:11434"), "{command}"); + } + + #[test] + fn a_publish_the_remote_does_not_confirm_is_an_error() { + // `tailscale serve` can exit zero while tailnet policy declines to + // publish. Trusting the exit code hands out a URL that never answers. + // Free before, still nothing after: the daemon accepted and did nothing. + let transport = ScriptedTransport::new(vec![ + ScriptedStep::ok("tailscale serve --bg", ""), + ScriptedStep::ok("tailscale serve status --json", "{}"), + ]); + let error = publish(&transport, 8000, 11434).unwrap_err().to_string(); + assert!(error.contains("would not answer"), "{error}"); + } + + #[test] + fn a_confirmed_publish_succeeds() { + let transport = ScriptedTransport::new(vec![ + ScriptedStep::ok("tailscale serve --bg", ""), + ScriptedStep::ok("tailscale serve status --json", PUBLISHED), + ]); + publish(&transport, 8000, 11434).expect("publish confirmed"); + } + + #[test] + fn a_port_someone_else_is_using_is_not_taken_over() { + // `tailscale serve` overwrites a port without complaint, so a check + // after the write is too late: the other forward is already gone and + // what we read back is our own. A second session on the same port would + // silently take the first one's endpoint away. + let occupied = ScriptedTransport::new(vec![ScriptedStep::ok( + "tailscale serve status --json", + r#"{"TCP": {"8000": {"TCPForward": "127.0.0.1:9999"}}}"#, + )]); + + let error = publish(&occupied, 8000, 11434).unwrap_err().to_string(); + assert!(error.contains("Refusing to take it over"), "{error}"); + assert!( + !occupied.calls().iter().any(|call| matches!( + call, + crate::remote::transport::TransportCall::Exec { command, .. } + if command.contains("--bg") + )), + "nothing should have been written: {:?}", + occupied.calls() + ); + } + + #[test] + fn re_publishing_our_own_forward_is_allowed() { + // `attach` re-declares an endpoint that is already ours; that has to + // stay idempotent rather than tripping the ownership guard. + let ours = ScriptedTransport::new(vec![ + ScriptedStep::ok("tailscale serve status --json", PUBLISHED), + ScriptedStep::ok("tailscale serve --bg", ""), + ]); + publish(&ours, 8000, 11434).expect("re-publishing our own forward"); + } + + #[test] + fn a_refused_publish_points_at_tailnet_policy() { + // The port is free; the daemon refuses the write itself. + let transport = ScriptedTransport::new(vec![ + ScriptedStep::ok("tailscale serve status --json", "{}"), + ScriptedStep::fails("tailscale serve --bg", 1, "serve not allowed"), + ]); + let error = publish(&transport, 8000, 11434).unwrap_err().to_string(); + assert!(error.contains("policy"), "{error}"); + } + + #[test] + fn a_port_that_now_belongs_to_something_else_is_not_turned_off() { + // `serve … off` takes a port, not a forward, so without this check a + // teardown tears down whatever happens to hold the port — possibly + // another tool's, or another person's, put there after our session was + // recorded. + let hijacked = ScriptedTransport::new(vec![ScriptedStep::ok( + "tailscale serve status --json", + r#"{"TCP": {"8000": {"TCPForward": "127.0.0.1:9999"}}}"#, + )]); + let error = withdraw(&hijacked, 8000, 11434).unwrap_err().to_string(); + assert!(error.contains("belongs to something else"), "{error}"); + // And nothing was turned off. + assert!( + !hijacked + .calls() + .iter() + .any(|call| matches!(call, crate::remote::transport::TransportCall::Exec { command, .. } if command.contains("off"))), + "a foreign forward must not be touched" + ); + } + + #[test] + fn withdrawing_an_already_absent_endpoint_is_not_an_error() { + // Teardown has to be safe to retry after a partial one. + let gone = ScriptedTransport::new(vec![ScriptedStep::ok( + "tailscale serve status --json", + "{}", + )]); + withdraw(&gone, 8000, 11434).expect("already gone is success"); + } + + #[test] + fn withdrawal_is_confirmed_not_assumed() { + // A publish outlives a reboot, so an unwithdrawn one is a GPU endpoint + // left on the tailnet with nothing tracking it. + // Ownership is probed first, so a stubborn remote answers PUBLISHED both + // before and after the `off` — which is exactly the state that must fail. + let stubborn = ScriptedTransport::new(vec![ + ScriptedStep::ok("tailscale serve --tcp=8000 off", ""), + ScriptedStep::ok("tailscale serve status --json", PUBLISHED), + ]); + let error = withdraw(&stubborn, 8000, 11434).unwrap_err().to_string(); + assert!(error.contains("still reports"), "{error}"); + } +} diff --git a/apps/rocm/src/remote/session.rs b/apps/rocm/src/remote/session.rs new file mode 100644 index 000000000..a0a852c24 --- /dev/null +++ b/apps/rocm/src/remote/session.rs @@ -0,0 +1,718 @@ +// Copyright © Advanced Micro Devices, Inc., or its affiliates. +// +// SPDX-License-Identifier: MIT + +//! Local records of work running on other machines. +//! +//! A session is the pairing of two things on the remote host: a managed model +//! server, and a tailnet publish pointing at it. Neither lives here — this is +//! only the note-to-self that lets a later `rocm remote status` or +//! `rocm remote stop`, run from a different shell or after a reboot, find them +//! again. +//! +//! One JSON file per session, mirroring the managed-service registry's shape. +//! The record is a point-in-time snapshot, never a source of truth about +//! liveness: `status` re-probes the remote rather than trusting what was +//! written here, because both halves can disappear without anyone updating a +//! local file. + +use std::fs; +use std::path::{Path, PathBuf}; + +use anyhow::{Context, Result, bail}; +use rocm_core::{AppPaths, unix_time_millis}; +use serde::{Deserialize, Serialize}; + +/// A model served on a remote machine and published to the tailnet. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub(crate) struct RemoteSessionRecord { + /// Stable identifier, also the record's filename. + pub(crate) session_id: String, + /// The target exactly as the user typed it, so messages echo their words + /// back rather than a resolved name they never used. + pub(crate) target: String, + /// The peer's MagicDNS name — what the endpoint URL is built from. + pub(crate) peer_host: String, + /// Explicit SSH port for the control channel, when one was given. + #[serde(default)] + pub(crate) ssh_port: Option, + pub(crate) model: String, + /// Id of the managed service on the remote's own registry. + pub(crate) remote_service_id: String, + /// How to invoke the CLI on the remote (may be a path, not just `rocm`). + pub(crate) remote_cli: String, + /// Loopback port the model server is bound to *on the remote*. + pub(crate) remote_port: u16, + /// Port the remote publishes to the tailnet. + pub(crate) tailnet_port: u16, + /// The address a user calls, including the OpenAI-compatible path. + pub(crate) base_url: String, + pub(crate) created_at_unix_ms: u128, +} + +impl RemoteSessionRecord { + /// Identifier for a session, derived from the peer and the port it serves on + /// rather than randomly. + /// + /// Deterministic on purpose: re-running `serve` against the same machine and + /// port resolves to the same record instead of accumulating a new one per + /// invocation, which is what turns a repeated command into an idempotent + /// action rather than a leak. + pub(crate) fn id_for(peer_host: &str, remote_port: u16) -> String { + let host = peer_host + .chars() + .map(|character| { + if character.is_ascii_alphanumeric() || character == '-' { + character.to_ascii_lowercase() + } else { + // Anything else becomes a dash: the id is used as a filename, + // and a MagicDNS name carries dots while a raw IPv6 address + // carries colons, which Windows rejects outright. + '-' + } + }) + .collect::(); + let host = host.trim_matches('-'); + format!("remote-{host}-{remote_port}") + } + + /// Check an id against the rules the rest of the CLI uses for anything that + /// becomes a path segment. + /// + /// [`Self::id_for`] only ever produces safe ids, but a record is also read + /// back from disk, and deserializing does not re-run that construction. An + /// id carrying `..` or a separator would make [`Self::path`] resolve outside + /// the sessions directory — and `stop` deletes that path. + /// + /// Deliberately delegating to [`rocm_core::ServiceId`] rather than spelling + /// the rules again here. It is the type the endpoint-key path builder + /// documents as its contract, this record's id is passed straight to that + /// builder, and a second hand-written rule set is exactly how the two drift + /// into disagreeing about what is safe. + fn validate_id(session_id: &str) -> Result<()> { + rocm_core::ServiceId::new(session_id) + .map(|_| ()) + .with_context(|| format!("unusable remote session id `{session_id}`")) + } + + pub(crate) fn path_in(paths: &AppPaths, session_id: &str) -> PathBuf { + paths + .remote_sessions_dir() + .join(format!("{session_id}.json")) + } + + pub(crate) fn path(&self, paths: &AppPaths) -> PathBuf { + Self::path_in(paths, &self.session_id) + } + + pub(crate) fn write(&self, paths: &AppPaths) -> Result<()> { + let directory = paths.remote_sessions_dir(); + fs::create_dir_all(&directory) + .with_context(|| format!("failed to create {}", directory.display()))?; + let path = self.path(paths); + let bytes = + serde_json::to_vec_pretty(self).context("failed to serialize the remote session")?; + + // Write beside the record and rename over it. A plain write is not + // atomic, so two serves racing on the same machine and port could leave + // a half-written file — which `load_all` then skips, hiding a session + // whose model is running and whose endpoint is published. Rename is + // atomic on both platforms we target, so a reader sees either the old + // record or the new one. + // Unique per writer. A single shared `.json.tmp` just moves the + // race: two serves for the same machine and port would write the same + // staging file and rename each other's half-written bytes into place. + let staging = staging_path_for(&path); + fs::write(&staging, &bytes) + .with_context(|| format!("failed to write {}", staging.display()))?; + fs::rename(&staging, &path).with_context(|| { + let _ = fs::remove_file(&staging); + format!("failed to replace {}", path.display()) + }) + } + + /// Forget this session locally. Best-effort and idempotent, so a teardown + /// that has already removed it can call this without special-casing. + pub(crate) fn remove(&self, paths: &AppPaths) { + let _ = fs::remove_file(self.path(paths)); + } + + /// Timestamp helper so callers do not each reach for the clock. + pub(crate) fn now() -> u128 { + unix_time_millis() + } +} + +/// Whether anything is already filed under this session id. +/// +/// The id is derived from the machine and the port, not minted per attempt, so +/// two `serve` runs against the same machine and port compute the same one. That +/// makes the id a *name*, not a claim. +/// +/// This answers a question; it does not take the name. Only [`store_key`] does +/// that, and it has to, because the gap between asking and writing is where two +/// runs both get "free" for an answer — a gap that spans the whole readiness +/// probe on the serve path, which is minutes when it provisions a CLI. So this +/// is for *diagnosis*: it runs first because it is cheap and because it can tell +/// a recorded session from a stray credential, and each of those needs a +/// different remedy. The claim itself comes later and is what actually decides. +/// +/// The key is checked as well as the record, because they are written at +/// different moments: a session that failed between minting its key and writing +/// its record leaves the key alone on disk. +pub(crate) fn exists(paths: &AppPaths, session_id: &str) -> bool { + RemoteSessionRecord::path_in(paths, session_id).exists() || key_path(paths, session_id).exists() +} + +/// A staging path no other writer will pick, for the write-then-rename above. +/// +/// The pid separates processes and the timestamp separates most writes, but two +/// threads in one process can reach the same millisecond — and a shared staging +/// file means each renames the other's half-written bytes into place, which is +/// the very race the rename exists to prevent. The counter closes that: it is +/// per-process and monotonic, so no two calls here can agree. +fn staging_path_for(path: &std::path::Path) -> PathBuf { + use std::sync::atomic::{AtomicU64, Ordering}; + static NEXT: AtomicU64 = AtomicU64::new(0); + path.with_extension(format!( + "{}.{}.{}.tmp", + std::process::id(), + unix_time_millis(), + NEXT.fetch_add(1, Ordering::Relaxed) + )) +} + +/// Where a session's endpoint credential is kept. +/// +/// Beside the session record, not in the managed-service registry. The endpoint +/// key helpers used for local serving build their path under `services_dir()` +/// from a *service* id, and a remote session is not a local service — it is a +/// different registry with its own ids. Borrowing that directory made two +/// namespaces share one folder and made a remote session's id masquerade as a +/// service id, which is a collision waiting to happen and reads wrong to anyone +/// inspecting either registry. +pub(crate) fn key_path(paths: &AppPaths, session_id: &str) -> PathBuf { + paths + .remote_sessions_dir() + .join(format!("{session_id}.endpoint-key")) +} + +/// Returned by [`store_key`] when the session name is already held. +/// +/// A distinct type rather than a message, because the caller has to tell it +/// apart from a full disk or a permission problem: one means somebody else owns +/// this session and we must leave every file under it alone, the other means our +/// own write failed and there is nothing of ours on disk to protect. +#[derive(Debug)] +pub(crate) struct NameAlreadyHeld { + pub(crate) path: PathBuf, +} + +impl std::fmt::Display for NameAlreadyHeld { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + formatter, + "the session name is already held by {}", + self.path.display() + ) + } +} + +impl std::error::Error for NameAlreadyHeld {} + +/// Claim the session name by creating the credential that guards its endpoint, +/// owner-only, failing if the name is already held. +/// +/// The creation is the claim, and it is one syscall. `id_for` is the machine and +/// the port with no nonce, so two `serve` runs against one machine compute the +/// same name; asking [`exists`] and then writing lets both get past the question +/// before either answers it, and on the serve path the two are separated by the +/// entire readiness probe. The loser of that race then overwrites the holder's +/// key, fails to start a model on a port the holder still has, and clears the +/// key on its way out — leaving the holder serving a published tailnet endpoint +/// that nobody, including its owner, can call. +/// +/// `create_new` is what closes it: `O_EXCL` on unix, `CREATE_NEW` on Windows, +/// and on both the existence check and the creation are one indivisible step. It +/// is the same discipline [`super::publish`] applies to a port — establish +/// ownership before writing, never after — with the difference that here the +/// establishing and the writing can be the same act. +/// +/// The id is validated first: it becomes a path segment, and this is the second +/// place that matters after the record itself. +pub(crate) fn store_key(paths: &AppPaths, session_id: &str, key: &str) -> Result<()> { + RemoteSessionRecord::validate_id(session_id)?; + let directory = paths.remote_sessions_dir(); + fs::create_dir_all(&directory) + .with_context(|| format!("failed to create {}", directory.display()))?; + let path = key_path(paths, session_id); + + let mut options = fs::OpenOptions::new(); + options.write(true).create_new(true); + // 0600 at creation, not tightened afterwards: a credential must never exist, + // even briefly, at whatever the umask would have given it. + #[cfg(unix)] + { + use std::os::unix::fs::OpenOptionsExt as _; + options.mode(0o600); + } + + let mut file = match options.open(&path) { + Ok(file) => file, + Err(error) if error.kind() == std::io::ErrorKind::AlreadyExists => { + return Err(anyhow::Error::new(NameAlreadyHeld { path })); + } + Err(error) => { + return Err(anyhow::Error::new(error)) + .with_context(|| format!("failed to write {}", path.display())); + } + }; + std::io::Write::write_all(&mut file, key.as_bytes()) + .with_context(|| format!("failed to write {}", path.display())) +} + +/// Forget a session's credential. Idempotent, so teardown can call it without +/// checking. +pub(crate) fn clear_key(paths: &AppPaths, session_id: &str) { + if RemoteSessionRecord::validate_id(session_id).is_ok() { + let _ = fs::remove_file(key_path(paths, session_id)); + } +} + +/// Every session recorded on this machine, ordered by id so listings are stable. +pub(crate) fn load_all(paths: &AppPaths) -> Result> { + let directory = paths.remote_sessions_dir(); + if !directory.is_dir() { + return Ok(Vec::new()); + } + + let mut records = Vec::new(); + for entry in fs::read_dir(&directory) + .with_context(|| format!("failed to read {}", directory.display()))? + { + let path = entry?.path(); + // A staging file from an interrupted write is not a record; its + // extension is `tmp`, so this already skips it, but say why. + if path.extension().is_none_or(|extension| extension != "json") { + continue; + } + match read_one(&path) { + Ok(record) => records.push(record), + // One unreadable record must not hide the rest. A half-written file + // from an interrupted run would otherwise make every session + // invisible, including the ones still publishing an endpoint. + Err(error) => { + eprintln!( + "warning: ignoring unreadable remote session {}: {error}", + path.display() + ); + } + } + } + records.sort_by(|left, right| left.session_id.cmp(&right.session_id)); + Ok(records) +} + +fn read_one(path: &Path) -> Result { + let bytes = fs::read(path).with_context(|| format!("failed to read {}", path.display()))?; + let record: RemoteSessionRecord = serde_json::from_slice(&bytes) + .with_context(|| format!("failed to parse {}", path.display()))?; + RemoteSessionRecord::validate_id(&record.session_id) + .with_context(|| format!("refusing to act on {}", path.display()))?; + Ok(record) +} + +/// Find the one session a user meant. +/// +/// Accepts an exact session id, or any unambiguous fragment of the target or +/// peer name. Ambiguity is refused rather than resolved by picking: `stop` on +/// the wrong session tears down someone's running model. +pub(crate) fn resolve(paths: &AppPaths, needle: &str) -> Result { + let sessions = load_all(paths)?; + if sessions.is_empty() { + bail!("no remote sessions are recorded on this machine"); + } + + if let Some(found) = sessions + .iter() + .find(|session| session.session_id.eq_ignore_ascii_case(needle)) + { + return Ok(found.clone()); + } + + let lowered = needle.to_ascii_lowercase(); + let matches = sessions + .iter() + .filter(|session| { + session.target.to_ascii_lowercase().contains(&lowered) + || session.peer_host.to_ascii_lowercase().contains(&lowered) + || session.session_id.to_ascii_lowercase().contains(&lowered) + }) + .collect::>(); + + match matches.as_slice() { + [] => bail!( + "no remote session matches `{needle}`. Recorded sessions: {}", + sessions + .iter() + .map(|session| session.session_id.as_str()) + .collect::>() + .join(", ") + ), + [only] => Ok((*only).clone()), + several => bail!( + "`{needle}` matches more than one remote session: {}\n\ + Name one of them exactly.", + several + .iter() + .map(|session| session.session_id.as_str()) + .collect::>() + .join(", ") + ), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn temp_paths(tag: &str) -> (PathBuf, AppPaths) { + let root = std::env::temp_dir().join(format!( + "rocm-remote-session-{tag}-{}-{}", + std::process::id(), + unix_time_millis() + )); + let _ = fs::remove_dir_all(&root); + fs::create_dir_all(&root).unwrap(); + ( + root.clone(), + AppPaths { + config_dir: root.join("config"), + data_dir: root.join("data"), + cache_dir: root.join("cache"), + }, + ) + } + + fn sample(peer_host: &str, remote_port: u16) -> RemoteSessionRecord { + RemoteSessionRecord { + session_id: RemoteSessionRecord::id_for(peer_host, remote_port), + target: peer_host.split('.').next().unwrap_or(peer_host).to_owned(), + peer_host: peer_host.to_owned(), + ssh_port: None, + model: "qwen".to_owned(), + remote_service_id: "svc-1".to_owned(), + remote_cli: "rocm".to_owned(), + remote_port, + tailnet_port: 8000, + base_url: format!("http://{peer_host}:8000/v1"), + created_at_unix_ms: 1, + } + } + + #[test] + fn a_session_id_is_stable_and_safe_as_a_filename() { + // Re-serving the same machine and port must land on the same record + // rather than accumulating one per invocation. + let first = RemoteSessionRecord::id_for("gpu-box-1.example-tailnet.ts.net", 11434); + let second = RemoteSessionRecord::id_for("gpu-box-1.example-tailnet.ts.net", 11434); + assert_eq!(first, second); + assert_ne!(first, RemoteSessionRecord::id_for("gpu-box-1", 11435)); + + // Dots and colons are not filename-safe everywhere; Windows rejects the + // colons an IPv6 address carries outright. + for id in [ + RemoteSessionRecord::id_for("gpu-box-1.example-tailnet.ts.net", 11434), + RemoteSessionRecord::id_for("fd7a:115c:a1e0::3", 11434), + ] { + assert!( + id.chars().all(|c| c.is_ascii_alphanumeric() || c == '-'), + "not filename-safe: {id}" + ); + } + } + + #[test] + fn a_session_round_trips_through_disk() -> Result<()> { + let (root, paths) = temp_paths("roundtrip"); + let record = sample("gpu-box-1.example-tailnet.ts.net", 11434); + record.write(&paths)?; + + let loaded = load_all(&paths)?; + assert_eq!(loaded, vec![record.clone()]); + + record.remove(&paths); + assert!(load_all(&paths)?.is_empty()); + // Removing twice must be safe: teardown calls it without checking. + record.remove(&paths); + + let _ = fs::remove_dir_all(root); + Ok(()) + } + + #[test] + fn a_record_is_replaced_atomically_so_a_reader_never_sees_half_of_one() { + // A torn record is skipped by `load_all`, which hides a session whose + // model is running and whose endpoint is published — the one thing that + // must never happen quietly. + let (root, paths) = temp_paths("atomic-write"); + let first = sample("gpu-box.example-tailnet.ts.net", 11434); + first.write(&paths).unwrap(); + + let mut second = first; + second.model = "a-different-model".to_owned(); + second.write(&paths).unwrap(); + + let loaded = load_all(&paths).unwrap(); + assert_eq!(loaded.len(), 1, "a rewrite must replace, not accumulate"); + assert_eq!(loaded[0].model, "a-different-model"); + + // No staging file is left behind for `load_all` to trip over. + let leftovers: Vec<_> = fs::read_dir(paths.remote_sessions_dir()) + .unwrap() + .filter_map(Result::ok) + .map(|entry| entry.file_name().to_string_lossy().into_owned()) + .filter(|name| name.contains("tmp")) + .collect(); + assert!(leftovers.is_empty(), "staging files left: {leftovers:?}"); + + let _ = fs::remove_dir_all(root); + } + + /// The assertions above describe the *outcome* of an atomic replace, and an + /// in-place `fs::write` produces the same outcome — so they pass with the + /// staging file and the rename deleted outright, certifying a guarantee + /// nothing exercises. This observes the mechanism instead. + /// + /// A rename replaces the directory entry, so the destination gets a new + /// inode. An in-place write keeps the inode and truncates, which is the + /// window where a reader sees half a record. Comparing inodes across a + /// rewrite therefore distinguishes the two without needing to catch the race. + #[cfg(unix)] + #[test] + fn replacing_a_record_swaps_the_file_rather_than_rewriting_it_in_place() { + use std::os::unix::fs::MetadataExt as _; + + let (root, paths) = temp_paths("atomic-swap"); + let first = sample("gpu-box.example-tailnet.ts.net", 11434); + first.write(&paths).unwrap(); + let before = fs::metadata(first.path(&paths)).unwrap().ino(); + + let mut second = first; + second.model = "a-different-model".to_owned(); + second.write(&paths).unwrap(); + let after = fs::metadata(second.path(&paths)).unwrap().ino(); + + assert_ne!( + before, after, + "the record was rewritten in place, so a reader can observe it truncated; \ + it must be staged and renamed over" + ); + + let _ = fs::remove_dir_all(root); + } + + #[test] + fn two_writers_in_one_process_do_not_share_a_staging_path() { + // The staging name carries the pid and a millisecond timestamp. Two + // threads writing the same session id inside one millisecond would + // otherwise pick the same staging file and rename each other's + // half-written bytes into place — the exact race the unique name exists + // to prevent. + let (root, paths) = temp_paths("atomic-staging-unique"); + let record = sample("gpu-box.example-tailnet.ts.net", 11434); + let names: std::collections::HashSet = (0..64) + .map(|_| { + staging_path_for(&record.path(&paths)) + .file_name() + .unwrap() + .to_string_lossy() + .into_owned() + }) + .collect(); + assert_eq!( + names.len(), + 64, + "staging names collided within a millisecond: {names:?}" + ); + + let _ = fs::remove_dir_all(root); + } + + #[test] + fn an_unreadable_record_does_not_hide_the_others() -> Result<()> { + // A half-written file from an interrupted run must not make every + // session invisible — including ones still publishing an endpoint that + // the user now has no listed way to find and stop. + let (root, paths) = temp_paths("corrupt"); + sample("gpu-box-1.example-tailnet.ts.net", 11434).write(&paths)?; + fs::write( + paths.remote_sessions_dir().join("truncated.json"), + b"{\"session_id\": ", + )?; + + let loaded = load_all(&paths)?; + assert_eq!(loaded.len(), 1); + assert_eq!(loaded[0].peer_host, "gpu-box-1.example-tailnet.ts.net"); + + let _ = fs::remove_dir_all(root); + Ok(()) + } + + #[test] + fn a_credential_is_kept_beside_its_session_and_is_recoverable() { + // Stored owner-only, next to the record rather than in the local + // service registry — a remote session is not a local service, and + // sharing that folder made one registry's ids masquerade as another's. + let (root, paths) = temp_paths("session-key"); + store_key(&paths, "remote-gpu-box-11434", "s3cret").expect("store"); + + // Asserting the relationship, and saying so in the message rather than + // dumping the path: a failure here is about *where* the credential + // landed, which the value alone does not explain. + let path = key_path(&paths, "remote-gpu-box-11434"); + assert!( + path.starts_with(paths.remote_sessions_dir()), + "a session's credential must live beside its record" + ); + assert!( + !path.starts_with(paths.services_dir()), + "a remote session must not write into the local service registry" + ); + // The user sees the key once when serving; without a readable file they + // could never recover it. + assert_eq!(fs::read_to_string(&path).expect("read"), "s3cret"); + + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt as _; + let mode = fs::metadata(&path).unwrap().permissions().mode(); + assert_eq!(mode & 0o777, 0o600, "a credential must be owner-only"); + } + + clear_key(&paths, "remote-gpu-box-11434"); + assert!(!path.exists()); + // Teardown calls this unconditionally, so a second clear must be safe. + clear_key(&paths, "remote-gpu-box-11434"); + let _ = fs::remove_dir_all(root); + } + + #[test] + fn a_session_credential_is_claimed_rather_than_overwritten() { + // Writing the key is what claims the session name, so it has to be the + // thing that fails when the name is already taken. Checking first and + // writing second cannot do that: the id is the machine and the port with + // no nonce, so two `serve` runs compute the same name, and both can pass + // a check before either writes. + // + // This observes the exclusivity itself rather than the end state. "Only + // one key file exists afterwards" is true either way — an overwrite + // leaves exactly one file too. + let (root, paths) = temp_paths("claim-exclusive"); + store_key(&paths, "remote-gpu-box-11434", "the-first-sessions-key").expect("first claim"); + + let second = store_key(&paths, "remote-gpu-box-11434", "a-second-runs-key"); + assert!( + second.is_err(), + "a second claim on a name already held must fail rather than take it over" + ); + assert_eq!( + fs::read_to_string(key_path(&paths, "remote-gpu-box-11434")).expect("read"), + "the-first-sessions-key", + "the holder's credential must survive another run's attempt to claim the name" + ); + + // Released, the name can be claimed again — teardown and a later serve + // on the same machine and port have to keep working. + clear_key(&paths, "remote-gpu-box-11434"); + store_key(&paths, "remote-gpu-box-11434", "a-later-sessions-key") + .expect("a released name is claimable again"); + + let _ = fs::remove_dir_all(root); + } + + #[test] + fn an_unusable_id_cannot_place_a_credential_outside_the_sessions_directory() { + let (root, paths) = temp_paths("session-key-escape"); + assert!(store_key(&paths, "../../escaped", "s3cret").is_err()); + // And clearing one is a no-op rather than a delete somewhere else. + clear_key(&paths, "../../escaped"); + let _ = fs::remove_dir_all(root); + } + + #[test] + fn a_record_with_an_unusable_id_is_refused_rather_than_followed() { + // Ids are constructed safe, but records are also read back from disk and + // deserializing does not re-run that construction. An id carrying a + // traversal would make the record's path resolve outside the sessions + // directory — and `stop` deletes that path. + let (root, paths) = temp_paths("unsafe-id"); + let mut record = sample("gpu-box.example-tailnet.ts.net", 11434); + record.session_id = "../../escaped".to_owned(); + fs::create_dir_all(paths.remote_sessions_dir()).unwrap(); + fs::write( + paths.remote_sessions_dir().join("planted.json"), + serde_json::to_vec(&record).unwrap(), + ) + .unwrap(); + + // Skipped like any other unreadable record: the rest stay visible. + sample("other-box.example-tailnet.ts.net", 11434) + .write(&paths) + .unwrap(); + let loaded = load_all(&paths).unwrap(); + assert_eq!(loaded.len(), 1); + assert_eq!(loaded[0].peer_host, "other-box.example-tailnet.ts.net"); + + // The rules come from the shared service-id type, so a traversal, a + // separator or an empty id is refused with a message naming the rule. + for bad in ["../../escaped", "a/b", "a\\b", ""] { + assert!( + RemoteSessionRecord::validate_id(bad).is_err(), + "should refuse {bad:?}" + ); + } + assert!(RemoteSessionRecord::validate_id("remote-gpu-box-11434").is_ok()); + + let _ = fs::remove_dir_all(root); + } + + #[test] + fn a_session_resolves_by_id_or_by_an_unambiguous_fragment() -> Result<()> { + let (root, paths) = temp_paths("resolve"); + let first = sample("gpu-box-1.example-tailnet.ts.net", 11434); + let second = sample("other-box.example-tailnet.ts.net", 11434); + first.write(&paths)?; + second.write(&paths)?; + + assert_eq!(resolve(&paths, &first.session_id)?, first); + assert_eq!(resolve(&paths, "gpu-box-1")?, first); + assert_eq!(resolve(&paths, "other")?, second); + + let _ = fs::remove_dir_all(root); + Ok(()) + } + + #[test] + fn an_ambiguous_session_is_refused_rather_than_guessed() -> Result<()> { + // Stopping the wrong session tears down a model someone is using. + let (root, paths) = temp_paths("ambiguous"); + sample("gpu-box-1.example-tailnet.ts.net", 11434).write(&paths)?; + sample("gpu-box-2.example-tailnet.ts.net", 11434).write(&paths)?; + + let error = resolve(&paths, "gpu-box").unwrap_err().to_string(); + assert!(error.contains("gpu-box-1"), "{error}"); + assert!(error.contains("gpu-box-2"), "{error}"); + + let unknown = resolve(&paths, "nothing-like-this") + .unwrap_err() + .to_string(); + assert!( + unknown.contains("Recorded sessions"), + "an unmatched name should list what does exist: {unknown}" + ); + + let _ = fs::remove_dir_all(root); + Ok(()) + } +} diff --git a/apps/rocm/src/remote/tailnet.rs b/apps/rocm/src/remote/tailnet.rs new file mode 100644 index 000000000..5b01103e4 --- /dev/null +++ b/apps/rocm/src/remote/tailnet.rs @@ -0,0 +1,714 @@ +// Copyright © Advanced Micro Devices, Inc., or its affiliates. +// +// SPDX-License-Identifier: MIT + +//! Reading the local tailnet: which machines exist, and how to address one. +//! +//! Everything here talks to the *local* Tailscale daemon and nothing else. No +//! SSH, no session, no call to a peer — listing candidate machines must not +//! require being able to reach them, or discovery would only ever show you what +//! you already knew how to contact. +//! +//! Two consumers: `rocm remote targets`, which renders the list, and the serve +//! path, which resolves one user-supplied name to a peer. They share a resolver +//! so a name that lists is a name that serves. +//! +//! Authorization is deliberately absent. Which peers a user may reach is the +//! tailnet admin's ACL policy, enforced by Tailscale; this module reports what +//! `tailscale status` already says and passes that tool's failures through +//! unchanged rather than reinterpreting them. + +use std::collections::BTreeMap; +use std::fmt::Write as _; +use std::io::ErrorKind; +use std::process::{Command, Stdio}; + +use anyhow::{Context, Result, bail}; +use serde::Deserialize; + +/// Backend state the Tailscale daemon reports when it is actually usable. +const BACKEND_RUNNING: &str = "Running"; + +/// What the local machine's Tailscale looks like right now. +/// +/// Three outcomes rather than a `Result`, because callers treat them +/// differently: `targets` reports the first two calmly and exits successfully, +/// while serving over the tailnet cannot proceed without the third. +#[derive(Debug, Clone)] +pub(crate) enum TailnetAvailability { + /// No `tailscale` on `PATH`. + NotInstalled, + /// Installed, but the daemon is not in a usable state — most often the user + /// has not run `tailscale up` yet. + NotRunning { + backend_state: String, + }, + Running(TailnetStatus), +} + +/// The local view of the tailnet. +#[derive(Debug, Clone)] +pub(crate) struct TailnetStatus { + /// This machine, when the daemon reports it. + pub(crate) this_machine: Option, + /// Every other machine on the tailnet, ordered by host name so output is + /// stable between runs rather than following a hash map's iteration order. + pub(crate) peers: Vec, + /// True when Tailscale is running without a kernel network device. + /// + /// There is no routable local address in that mode, so `ssh` cannot dial a + /// tailnet IP directly and has to be routed through Tailscale's own relay + /// instead. Taken from the daemon's `TUN` flag, whose documented meaning is + /// exactly this — worth preferring over inferring the mode from a failed + /// connection after the fact. + pub(crate) userspace_networking: bool, +} + +/// One machine on the tailnet. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct TailnetPeer { + /// Short host name. Not guaranteed unique across a tailnet. + pub(crate) host: String, + /// Fully-qualified MagicDNS name, trailing dot removed. + pub(crate) dns_name: String, + pub(crate) addresses: Vec, + pub(crate) os: String, + pub(crate) tags: Vec, + pub(crate) online: bool, +} + +impl TailnetPeer { + /// The name to hand to `ssh` and to build an endpoint URL from. + /// + /// Prefers the MagicDNS name over a raw address: it survives the peer being + /// reassigned an address, and it is what a user recognises in a URL. + pub(crate) fn endpoint_host(&self) -> &str { + if self.dns_name.is_empty() { + self.addresses.first().map_or("", String::as_str) + } else { + &self.dns_name + } + } + + fn has_tag(&self, tag: &str) -> bool { + // Accept `gpu` for `tag:gpu`: the `tag:` prefix is Tailscale's wire + // format, and making users type it adds nothing. + let wanted = tag.strip_prefix("tag:").unwrap_or(tag); + self.tags + .iter() + .any(|owned| owned.strip_prefix("tag:").unwrap_or(owned) == wanted) + } + + /// Whether `needle` names this peer exactly, by DNS name, host name, or + /// address. + fn matches_exactly(&self, needle: &str) -> bool { + let needle = needle.trim_end_matches('.'); + self.dns_name.eq_ignore_ascii_case(needle) + || self.host.eq_ignore_ascii_case(needle) + || self.addresses.iter().any(|address| address == needle) + } + + /// Whether `needle` appears in this peer's names — the loose fallback used + /// only when nothing matched exactly. + fn matches_loosely(&self, needle: &str) -> bool { + let needle = needle.to_ascii_lowercase(); + self.host.to_ascii_lowercase().contains(&needle) + || self.dns_name.to_ascii_lowercase().contains(&needle) + } +} + +/// `tailscale status --json`, as much of it as we read. +/// +/// Unknown fields are ignored by default, so a newer Tailscale adding output +/// cannot break us. `backend_state` is the one field with no default: every +/// real status document has it, so its absence means we were handed something +/// that is not a status document at all, and failing there gives a far better +/// error than a confidently empty peer list. +#[derive(Debug, Deserialize)] +struct RawStatus { + #[serde(rename = "BackendState")] + backend_state: String, + #[serde(rename = "TUN", default)] + tun: bool, + #[serde(rename = "Self")] + this_machine: Option, + #[serde(rename = "Peer", default)] + peers: BTreeMap, +} + +#[derive(Debug, Deserialize)] +struct RawPeer { + #[serde(rename = "HostName", default)] + host_name: String, + #[serde(rename = "DNSName", default)] + dns_name: String, + #[serde(rename = "OS", default)] + os: String, + #[serde(rename = "TailscaleIPs", default)] + tailscale_ips: Vec, + #[serde(rename = "Tags", default)] + tags: Vec, + #[serde(rename = "Online", default)] + online: bool, +} + +impl RawPeer { + fn into_peer(self) -> TailnetPeer { + TailnetPeer { + host: self.host_name, + // MagicDNS names arrive as absolute FQDNs with a trailing dot. + // Leaving it on produces `http://box.tail.ts.net.:8000` in printed + // URLs, which is technically valid and looks like a typo. + dns_name: self.dns_name.trim_end_matches('.').to_owned(), + addresses: self.tailscale_ips, + os: self.os, + tags: self.tags, + online: self.online, + } + } +} + +/// Ask the local Tailscale daemon what it can see. +pub(crate) fn local_status() -> Result { + let output = match Command::new("tailscale") + .args(["status", "--json"]) + .stdin(Stdio::null()) + .output() + { + Ok(output) => output, + Err(error) if error.kind() == ErrorKind::NotFound => { + return Ok(TailnetAvailability::NotInstalled); + } + Err(error) => { + return Err(error).context("failed to run `tailscale status --json`"); + } + }; + + if !output.status.success() { + // Pass Tailscale's own words through. It knows why it is unhappy — + // logged out, daemon not running, permission denied — and paraphrasing + // that into our own vocabulary only loses detail. + bail!( + "`tailscale status --json` failed: {}", + String::from_utf8_lossy(&output.stderr).trim() + ); + } + + let status = parse_status(&String::from_utf8_lossy(&output.stdout))?; + if status_is_running(&status.0) { + Ok(TailnetAvailability::Running(status.1)) + } else { + Ok(TailnetAvailability::NotRunning { + backend_state: status.0, + }) + } +} + +fn status_is_running(backend_state: &str) -> bool { + backend_state == BACKEND_RUNNING +} + +/// Parse `tailscale status --json`, returning the backend state alongside the +/// tailnet view. Pure, so the whole surface is testable against fixtures. +fn parse_status(json: &str) -> Result<(String, TailnetStatus)> { + let raw: RawStatus = serde_json::from_str(json).context( + "failed to parse `tailscale status --json`; the output was not a status document", + )?; + + let mut peers = raw + .peers + .into_values() + .map(RawPeer::into_peer) + .collect::>(); + // A map gives no useful order, so impose one. Sorting by the name a user + // would type keeps repeated runs comparable. + peers.sort_by(|left, right| { + left.host + .to_ascii_lowercase() + .cmp(&right.host.to_ascii_lowercase()) + .then_with(|| left.dns_name.cmp(&right.dns_name)) + }); + + Ok(( + raw.backend_state, + TailnetStatus { + this_machine: raw.this_machine.map(RawPeer::into_peer), + peers, + userspace_networking: !raw.tun, + }, + )) +} + +/// Find the peer a user meant by `needle`. +/// +/// `Ok(None)` means no peer matched — a fact, not a failure, so the caller can +/// decide whether that is fatal. An ambiguous name *is* an error: silently +/// picking one of several GPU boxes is the kind of guess that gets a model +/// served on someone else's machine. +pub(crate) fn resolve_peer<'a>( + status: &'a TailnetStatus, + needle: &str, +) -> Result> { + // An SSH destination may carry a user prefix; the tailnet knows nothing + // about that, so compare on the host part. + let needle = needle.rsplit('@').next().unwrap_or(needle).trim(); + if needle.is_empty() { + return Ok(None); + } + + // Collected rather than `find`-ed. `matches_exactly` compares on `host`, + // which is documented above as not unique across a tailnet — two machines + // sharing an OS hostname is ordinary, since MagicDNS dedupes `DNSName` and + // not `HostName`. Taking the first would be exactly the silent guess the doc + // comment above forbids, and the loose path below already refuses it. + let exact = status + .peers + .iter() + .filter(|peer| peer.matches_exactly(needle)) + .collect::>(); + match exact.as_slice() { + [] => {} + [only] => return Ok(Some(only)), + several => return Err(ambiguous(needle, several)), + } + + let loose = status + .peers + .iter() + .filter(|peer| peer.matches_loosely(needle)) + .collect::>(); + match loose.as_slice() { + [] => Ok(None), + [only] => Ok(Some(only)), + several => Err(ambiguous(needle, several)), + } +} + +/// The refusal both match passes share: name the candidates rather than pick one. +fn ambiguous(needle: &str, candidates: &[&TailnetPeer]) -> anyhow::Error { + anyhow::anyhow!( + "`{needle}` matches more than one machine on the tailnet: {}\n\ + Name one of them exactly.", + candidates + .iter() + .map(|peer| peer.dns_name.as_str()) + .collect::>() + .join(", ") + ) +} + +/// Render the candidate machines. +/// +/// Purely informational, and says so: being listed here means the tailnet can +/// see the machine, not that it has a GPU, ROCm, or the CLI. Each entry points +/// at the command that actually answers that question. +pub(crate) fn render_targets(status: &TailnetStatus, tag: Option<&str>) -> String { + let peers = status + .peers + .iter() + .filter(|peer| tag.is_none_or(|tag| peer.has_tag(tag))) + .collect::>(); + + let mut output = String::new(); + let _ = writeln!(output, "Remote Targets"); + let _ = writeln!(output); + // Name the machine we are looking *from*. A user on more than one tailnet + // otherwise has no way to tell which one this list describes. + if let Some(this_machine) = &status.this_machine { + let _ = writeln!(output, "This machine: {}", this_machine.host); + } + let online = peers.iter().filter(|peer| peer.online).count(); + let _ = writeln!( + output, + "Status: {online} online, {} offline", + peers.len() - online + ); + let _ = writeln!(output); + + if peers.is_empty() { + let _ = match tag { + Some(tag) => writeln!(output, "No tailnet machines are tagged `{tag}`."), + None => writeln!(output, "No other machines are on this tailnet."), + }; + return output; + } + + for peer in peers { + let _ = writeln!(output, "- {}", peer.host); + let _ = writeln!(output, " address: {}", peer.endpoint_host()); + if !peer.addresses.is_empty() { + let _ = writeln!(output, " ip: {}", peer.addresses.join(", ")); + } + if !peer.os.is_empty() { + let _ = writeln!(output, " os: {}", peer.os); + } + if !peer.tags.is_empty() { + let _ = writeln!(output, " tags: {}", peer.tags.join(", ")); + } + let _ = writeln!( + output, + " online: {}", + if peer.online { "yes" } else { "no" } + ); + // Quoted, because this line is offered to be copied and run and the name + // in it is not ours. `HostName` comes straight from the peer's own + // `tailscale status` entry — any device on the tailnet chooses its own — + // so an unquoted one renders a command that does something other than + // what it reads as, on the machine of whoever pastes it. + let _ = writeln!( + output, + " check: rocm remote doctor {}", + super::shell_quote(&peer.host) + ); + } + + let _ = writeln!(output); + if status.userspace_networking { + // Worth saying out loud: in this mode traffic cannot take a direct path + // and is relayed instead, which caps throughput. Someone about to serve + // a model should know that before they blame the GPU. + let _ = writeln!( + output, + "Note: Tailscale is running here without a network device, so traffic to these" + ); + let _ = writeln!( + output, + "machines is relayed rather than sent directly, which limits throughput." + ); + let _ = writeln!(output); + } + let _ = writeln!( + output, + "Listed machines are reachable on the tailnet. That does not mean they have a GPU," + ); + let _ = writeln!( + output, + "ROCm, or the ROCm CLI — run `rocm remote doctor ` to find out." + ); + if tag.is_none() { + let _ = writeln!( + output, + "Narrow this list with `--tag ` once your tailnet tags its GPU machines." + ); + } + output +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Shaped after a real `tailscale status --json`: absolute MagicDNS names, + /// a peer map keyed by node public key, tags only on some machines, and one + /// offline host. + const FIXTURE: &str = r#"{ + "Version": "1.999.0-t01b2c3", + "TUN": true, + "BackendState": "Running", + "MagicDNSSuffix": "example-tailnet.ts.net", + "Self": { + "ID": "nSELF", + "HostName": "laptop", + "DNSName": "laptop.example-tailnet.ts.net.", + "OS": "linux", + "TailscaleIPs": ["100.88.0.9"], + "Online": true + }, + "Peer": { + "nodekey:bbb": { + "ID": "nBBB", + "HostName": "gpu-box-2", + "DNSName": "gpu-box-2.example-tailnet.ts.net.", + "OS": "linux", + "TailscaleIPs": ["100.88.14.37"], + "Tags": ["tag:gpu"], + "Online": false, + "LastSeen": "2026-08-30T11:02:41Z" + }, + "nodekey:aaa": { + "ID": "nAAA", + "HostName": "gpu-box-1", + "DNSName": "gpu-box-1.example-tailnet.ts.net.", + "OS": "linux", + "TailscaleIPs": ["100.88.14.21", "fd7a:115c:a1e0::3"], + "Tags": ["tag:gpu", "tag:prod"], + "Online": true + }, + "nodekey:ccc": { + "ID": "nCCC", + "HostName": "phone", + "DNSName": "phone.example-tailnet.ts.net.", + "OS": "iOS", + "TailscaleIPs": ["100.88.51.6"], + "Online": true + } + } + }"#; + + fn fixture() -> TailnetStatus { + parse_status(FIXTURE).expect("fixture parses").1 + } + + #[test] + fn parsing_orders_peers_and_strips_the_magicdns_trailing_dot() { + let status = fixture(); + + // A peer map has no inherent order; unsorted output would shuffle + // between runs and make the listing unreadable. + assert_eq!( + status + .peers + .iter() + .map(|peer| peer.host.as_str()) + .collect::>(), + vec!["gpu-box-1", "gpu-box-2", "phone"] + ); + // Left on, the trailing dot shows up in printed URLs looking like a typo. + assert_eq!(status.peers[0].dns_name, "gpu-box-1.example-tailnet.ts.net"); + assert_eq!(status.this_machine.unwrap().host, "laptop"); + } + + #[test] + fn parsing_tolerates_a_newer_tailscale_but_rejects_a_non_status_document() { + // Tailscale adds output over time; new keys must not break us. + let with_extras = FIXTURE.replace( + "\"BackendState\": \"Running\",", + "\"BackendState\": \"Running\", \"SomethingAddedLater\": {\"a\": 1},", + ); + assert!(parse_status(&with_extras).is_ok()); + + // But a document with no backend state is not a status document, and + // treating it as one would report an empty tailnet with full confidence. + let not_a_status = r#"{"Peer": {}}"#; + let error = parse_status(not_a_status).unwrap_err().to_string(); + assert!(error.contains("not a status document"), "{error}"); + } + + #[test] + fn userspace_networking_is_read_from_the_daemon_not_guessed() { + // With a kernel device, ssh can dial a tailnet address directly. + assert!(!fixture().userspace_networking); + + // Without one there is no routable address, and the connection has to be + // relayed. Tailscale reports this directly, so there is no need to infer + // it from a connection that already failed. + let userspace = FIXTURE.replace("\"TUN\": true", "\"TUN\": false"); + assert!(parse_status(&userspace).unwrap().1.userspace_networking); + } + + #[test] + fn a_daemon_that_is_not_up_is_reported_as_such_not_as_an_empty_tailnet() { + let logged_out = FIXTURE.replace( + "\"BackendState\": \"Running\"", + "\"BackendState\": \"NeedsLogin\"", + ); + let (state, _) = parse_status(&logged_out).expect("still parses"); + assert!(!status_is_running(&state)); + assert_eq!(state, "NeedsLogin"); + } + + #[test] + fn a_peer_resolves_by_host_name_dns_name_or_address() { + let status = fixture(); + for needle in [ + "gpu-box-1", + "GPU-BOX-1", + "gpu-box-1.example-tailnet.ts.net", + "gpu-box-1.example-tailnet.ts.net.", + "100.88.14.21", + "user@gpu-box-1", + ] { + let peer = resolve_peer(&status, needle) + .unwrap_or_else(|error| panic!("{needle}: {error}")) + .unwrap_or_else(|| panic!("{needle} should resolve")); + assert_eq!(peer.host, "gpu-box-1", "resolving {needle}"); + } + } + + #[test] + fn an_ambiguous_name_is_refused_and_names_the_candidates() { + // Quietly picking one of two GPU boxes would serve a model on a machine + // the user did not choose. + let status = fixture(); + let error = resolve_peer(&status, "gpu-box").unwrap_err().to_string(); + assert!(error.contains("gpu-box-1"), "{error}"); + assert!(error.contains("gpu-box-2"), "{error}"); + } + + #[test] + fn two_machines_sharing_a_host_name_are_refused_not_silently_picked() { + // `HostName` is not unique on a tailnet — MagicDNS dedupes `DNSName`, + // not the OS hostname — so two machines both called `gpu-box` is + // ordinary, not contrived. Both match *exactly*, which is the path that + // used to take the first and serve a model on whichever peer happened to + // sort first. + const DUPLICATE_HOSTS: &str = r#"{ + "Version": "1.999.0", + "TUN": true, + "BackendState": "Running", + "MagicDNSSuffix": "example-tailnet.ts.net", + "Self": {"ID": "nS", "HostName": "laptop", + "DNSName": "laptop.example-tailnet.ts.net.", "OS": "linux", + "TailscaleIPs": ["100.88.0.9"], "Online": true}, + "Peer": { + "nodekey:aaa": {"ID": "nA", "HostName": "gpu-box", + "DNSName": "gpu-box.example-tailnet.ts.net.", "OS": "linux", + "TailscaleIPs": ["100.88.14.21"], "Online": true}, + "nodekey:bbb": {"ID": "nB", "HostName": "gpu-box", + "DNSName": "gpu-box-2.example-tailnet.ts.net.", "OS": "linux", + "TailscaleIPs": ["100.88.14.37"], "Online": true} + } + }"#; + let status = parse_status(DUPLICATE_HOSTS).expect("fixture parses").1; + + let error = resolve_peer(&status, "gpu-box") + .expect_err("an exactly-ambiguous name must be refused, not guessed") + .to_string(); + // Named by DNS name, not host name: printing the colliding host twice + // would tell the user nothing they could act on. + assert!(error.contains("gpu-box.example-tailnet.ts.net"), "{error}"); + assert!( + error.contains("gpu-box-2.example-tailnet.ts.net"), + "{error}" + ); + + // And a name that singles one of them out still resolves. + let peer = resolve_peer(&status, "gpu-box-2.example-tailnet.ts.net") + .expect("an unambiguous exact name resolves") + .expect("the peer is present"); + assert_eq!(peer.dns_name, "gpu-box-2.example-tailnet.ts.net"); + } + + #[test] + fn an_unknown_name_is_absent_rather_than_an_error() { + // Not finding a peer is a fact the caller judges, not a failure here. + assert!( + resolve_peer(&fixture(), "not-on-this-tailnet") + .expect("no match is not an error") + .is_none() + ); + } + + #[test] + fn an_exact_name_wins_over_a_longer_one_that_contains_it() { + let status = parse_status(&FIXTURE.replace( + "\"HostName\": \"phone\"", + "\"HostName\": \"gpu-box-1-spare\"", + )) + .unwrap() + .1; + let peer = resolve_peer(&status, "gpu-box-1") + .expect("exact match must not be ambiguous") + .expect("resolves"); + assert_eq!(peer.host, "gpu-box-1"); + } + + #[test] + fn targets_can_be_narrowed_to_tagged_machines() { + let status = fixture(); + + let all = render_targets(&status, None); + assert!(all.contains("- gpu-box-1")); + assert!(all.contains("- phone")); + + // The `tag:` prefix is Tailscale's wire format; users should not have to + // type it, so both spellings filter the same way. + for spelling in ["gpu", "tag:gpu"] { + let tagged = render_targets(&status, Some(spelling)); + assert!(tagged.contains("- gpu-box-1"), "{spelling}"); + assert!(!tagged.contains("- phone"), "{spelling}"); + assert!(tagged.contains("Status: 1 online, 1 offline"), "{spelling}"); + } + } + + #[test] + fn targets_states_that_listing_is_not_readiness() { + // The whole risk of this command is being read as "these machines can + // serve". It has to say plainly that it means no such thing. + let rendered = render_targets(&fixture(), None); + assert!( + rendered.contains("does not mean they have a GPU"), + "{rendered}" + ); + assert!( + rendered.contains("rocm remote doctor gpu-box-1"), + "{rendered}" + ); + assert!( + rendered.contains("online: no"), + "an offline peer is shown, not hidden" + ); + } + + #[test] + fn targets_names_the_machine_the_list_is_seen_from() { + // Someone on more than one tailnet cannot otherwise tell which one they + // are looking at. + assert!(render_targets(&fixture(), None).contains("This machine: laptop")); + } + + #[test] + fn a_peer_name_carrying_shell_syntax_is_quoted_before_it_is_offered_to_run() { + // Every device on a tailnet chooses its own `HostName`, so this string is + // not ours. It is rendered into a line the user is invited to copy, which + // makes an unquoted one a command that does something other than what it + // reads as — on the machine of whoever pastes it. + const HOSTILE: &str = r#"{ + "Version": "1.999.0", "TUN": true, "BackendState": "Running", + "MagicDNSSuffix": "example-tailnet.ts.net", + "Self": {"ID": "nS", "HostName": "laptop", + "DNSName": "laptop.example-tailnet.ts.net.", "OS": "linux", + "TailscaleIPs": ["100.88.0.9"], "Online": true}, + "Peer": { + "nodekey:aaa": {"ID": "nA", "HostName": "box; rm -rf ~", + "DNSName": "box.example-tailnet.ts.net.", "OS": "linux", + "TailscaleIPs": ["100.88.14.21"], "Online": true} + } + }"#; + let status = parse_status(HOSTILE).expect("fixture parses").1; + let rendered = render_targets(&status, None); + + assert!( + !rendered.contains("doctor box; rm -rf ~"), + "a peer name must not reach the advice line unquoted:\n{rendered}" + ); + // Ask a real shell what the rendered line would actually pass along. + let line = rendered + .lines() + .map(str::trim) + .find(|line| line.starts_with("check: ")) + .expect("every peer carries a check line") + .trim_start_matches("check: ") + .to_owned(); + let output = std::process::Command::new("sh") + .arg("-c") + .arg(format!("rocm() {{ printf %s \"$3\"; }}\n{line}")) + .output() + .expect("sh should run"); + assert_eq!( + String::from_utf8_lossy(&output.stdout), + "box; rm -rf ~", + "the name must arrive as one literal argument" + ); + } + + #[test] + fn targets_warns_when_traffic_to_peers_will_be_relayed() { + assert!(!render_targets(&fixture(), None).contains("relayed")); + + let userspace = parse_status(&FIXTURE.replace("\"TUN\": true", "\"TUN\": false")) + .unwrap() + .1; + let rendered = render_targets(&userspace, None); + assert!( + rendered.contains("relayed"), + "a throughput cap should be stated before someone blames the GPU: {rendered}" + ); + } + + #[test] + fn an_empty_tag_filter_explains_itself_rather_than_printing_nothing() { + let rendered = render_targets(&fixture(), Some("nonexistent")); + assert!(rendered.contains("No tailnet machines are tagged `nonexistent`.")); + } +} diff --git a/apps/rocm/src/remote/transport.rs b/apps/rocm/src/remote/transport.rs new file mode 100644 index 000000000..3497f165c --- /dev/null +++ b/apps/rocm/src/remote/transport.rs @@ -0,0 +1,918 @@ +// Copyright © Advanced Micro Devices, Inc., or its affiliates. +// +// SPDX-License-Identifier: MIT + +//! The control channel to a remote GPU host. +//! +//! `rocm remote` needs exactly two things from a remote machine: run a command +//! and read its output, and copy a file over. [`Transport`] captures those and +//! nothing else, so the orchestration above it stays agnostic to how the host is +//! reached. [`SshTransport`] is the one real implementation. +//! +//! Why shell out to the system `ssh`/`scp` instead of a Rust SSH crate: the +//! user's existing keys, agent, `~/.ssh/config`, `ProxyJump` hosts, and host +//! aliases all keep working with no extra configuration, and there is no new +//! dependency surface for something as security-sensitive as an SSH client. It +//! also keeps this code synchronous, matching the command handlers around it. +//! +//! Note what is deliberately *absent*: there is no port-forwarding method. An +//! earlier prototype opened a detached `ssh -L` tunnel and tracked its PID; this +//! design instead has the remote publish its own service onto the tailnet, so +//! the data path never runs through the control channel and no local process +//! outlives the command. + +#[cfg(test)] +use std::cell::RefCell; +use std::io::Write as _; +use std::path::Path; +use std::process::{Command, Stdio}; + +use anyhow::{Context, Result, bail}; + +/// Captured result of running a command on the remote host. +/// +/// The `success` flag distinguishes a clean-but-non-zero remote exit — say +/// `rocm --version` on a host that has no CLI — from a transport failure where +/// the host could not be reached at all. Every readiness probe depends on that +/// distinction: "answered, and said no" and "never answered" call for different +/// errors, and collapsing them produces the classic misleading +/// "ROCm is not installed" on a host that is merely offline. +#[derive(Debug, Clone)] +pub(crate) struct RemoteOutcome { + pub(crate) success: bool, + pub(crate) code: Option, + pub(crate) stdout: String, + pub(crate) stderr: String, +} + +impl RemoteOutcome { + /// How the command exited, for error messages: an exit status, or `signal` + /// when it was killed and carries no code. + fn exit_label(&self) -> String { + self.code + .map_or_else(|| "signal".to_owned(), |code| code.to_string()) + } +} + +/// Ways to reach a remote host. See the module docs for the design rationale. +pub(crate) trait Transport { + /// Run `command` on the remote host, optionally feeding it `stdin`, and + /// capture the outcome. + /// + /// Returns `Err` only when the command could not be launched or the host is + /// unreachable; a non-zero remote exit is reported through + /// [`RemoteOutcome::success`]. + /// + /// `stdin` exists so a secret never has to travel in a command line. Both + /// the local `ssh` invocation and the remote shell expose their arguments in + /// the process table, so an API key interpolated into `command` would be + /// readable by any other user on either machine; piped in, it is not. + fn exec_with_stdin(&self, command: &str, stdin: Option<&str>) -> Result; + + /// Run `command` on the remote host and capture its outcome. + fn exec(&self, command: &str) -> Result { + self.exec_with_stdin(command, None) + } + + /// 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.exit_label(), + outcome.stderr.trim(), + ); + } + Ok(outcome.stdout) + } + + /// Copy a local file to `remote_path` on the host. + fn push_file(&self, local_path: &Path, remote_path: &str) -> Result<()>; +} + +/// Options that let repeated commands share one connection. +/// +/// All three are needed or none are. OpenSSH defaults `ControlPath` to `none`, +/// and with no socket path it ignores `ControlMaster` entirely — so emitting the +/// master and persist options alone is a claim the client does not honour, and +/// every status poll silently pays a fresh handshake. When no private socket +/// directory can be established we emit nothing rather than options that look +/// like multiplexing and are not. +/// +/// `%C` is OpenSSH's hash of the connection's identity, so one socket per +/// destination without building a filename out of user-supplied host strings. +fn multiplex_args(control_path: Option<&str>) -> Vec { + let Some(control_path) = control_path else { + return Vec::new(); + }; + vec![ + "-o".to_owned(), + "ControlMaster=auto".to_owned(), + "-o".to_owned(), + format!("ControlPath={control_path}"), + // Long enough that a burst of polls reuses one connection, short enough + // that nothing lingers after a command finishes. + "-o".to_owned(), + "ControlPersist=60s".to_owned(), + ] +} + +/// A private directory to keep control sockets in, or `None` if we cannot get +/// one — in which case multiplexing is skipped rather than half-configured. +#[cfg(unix)] +fn control_socket_path() -> Option { + use std::os::unix::fs::PermissionsExt as _; + use std::sync::OnceLock; + + static PATH: OnceLock> = OnceLock::new(); + PATH.get_or_init(|| { + let directory = rocm_core::AppPaths::discover().ok()?.data_dir.join("ssh"); + std::fs::create_dir_all(&directory).ok()?; + // Owner-only: a control socket is an authenticated channel to the remote, + // so anything that can reach it can run commands there as this user. + std::fs::set_permissions(&directory, std::fs::Permissions::from_mode(0o700)).ok()?; + + let candidate = directory.join("cm-%C").to_string_lossy().into_owned(); + // Unix socket paths are capped near 104 bytes; `%C` expands to a 40-char + // hash. Past the cap ssh fails outright, so a deep home directory should + // cost multiplexing, not every remote command. + (candidate.len() + 40 < 100).then_some(candidate) + }) + .clone() +} + +/// Windows has no connection multiplexing to configure. +/// +/// Win32-OpenSSH does not implement `ControlMaster`; setting it there produces +/// warnings at best. Returning `None` keeps the argument builder honest instead +/// of emitting options the platform ignores. +#[cfg(not(unix))] +fn control_socket_path() -> Option { + None +} + +/// Exit status `ssh` uses for its own failures rather than the remote command's. +/// +/// Documented in ssh(1): "ssh exits with the exit status of the remote command +/// or with 255 if an error occurred." +const SSH_TRANSPORT_FAILURE: i32 = 255; + +/// Environment variable naming an alternative SSH configuration file. +/// +/// `ssh` resolves `~/.ssh/config` from the account database rather than from +/// `HOME`, so there is otherwise no way to point this at a different one — +/// awkward for anyone keeping a per-project or per-tenant ssh config, and the +/// reason the end-to-end harness could not drive the real binary at all. +const SSH_CONFIG_ENV: &str = "ROCM_REMOTE_SSH_CONFIG"; + +/// `-F ` when an alternative config was named, nothing otherwise. +fn config_args() -> Vec { + std::env::var(SSH_CONFIG_ENV) + .ok() + .filter(|path| !path.trim().is_empty()) + .map_or_else(Vec::new, |path| vec!["-F".to_owned(), path]) +} + +/// Reject a destination that `ssh` would read as an option rather than a host. +/// +/// The destination has to precede the `--` guard — that is where `ssh` expects +/// it — so unlike the remote command it is not protected by it. A value +/// beginning with `-` is therefore parsed as a local option, and +/// `-oProxyCommand=…` would run a command on *this* machine. Hostnames cannot +/// begin with a hyphen anyway (RFC 1123), so nothing legitimate is lost by +/// refusing outright rather than trying to escape it. +fn validate_destination(destination: &str) -> Result<()> { + if destination.trim().is_empty() { + bail!("a remote machine must be named"); + } + // Check the host half too: `user@-h` puts the hyphen where ssh still sees it. + let host = destination.rsplit('@').next().unwrap_or(destination); + if destination.starts_with('-') || host.starts_with('-') { + bail!( + "`{destination}` is not a usable machine name: a name starting with `-` would be \ + read as an option by ssh rather than as a host" + ); + } + Ok(()) +} + +/// Reject an `scp` path argument that starts with `-`. +/// +/// The two sides are not at equal risk, and it is worth being exact about +/// which is which. The **local** path is pushed as its own argument, so a +/// leading `-` genuinely is read as an option rather than a file — the same +/// risk [`validate_destination`] guards against for the ssh destination. The +/// **remote** path is not: it is built as `{destination}:{remote_path}`, so +/// the argument always begins with the destination and scp can never take it +/// for an option. Guarding it anyway is a shape check, not a safety one — a +/// remote path in that shape means this process built the argument wrongly, +/// and saying so here beats chasing an odd path on the far machine. +/// +/// Neither is expected to come from outside this process today, but an +/// artifact name or staging path is not a hostname either, and both cross +/// this boundary as plain strings rather than as a type that already rules a +/// leading `-` out. +fn validate_scp_path(path: &str, which: &str) -> Result<()> { + if path.starts_with('-') { + bail!( + "`{path}` is not a usable {which} path for scp: a path here never starts \ + with `-`, and one that does is either read as an option or a sign the \ + argument was built wrongly" + ); + } + Ok(()) +} + +/// SSH-backed control channel that shells out to the system `ssh`/`scp`. +#[derive(Debug, Clone)] +pub(crate) struct SshTransport { + /// SSH destination as accepted by `ssh` (e.g. `user@host`, or a + /// `~/.ssh/config` host alias). + destination: String, + /// Optional explicit port; `None` uses the ssh default or whatever + /// `~/.ssh/config` specifies for this destination. + port: Option, +} + +impl SshTransport { + pub(crate) fn new(destination: impl Into, port: Option) -> Result { + let destination = destination.into(); + validate_destination(&destination)?; + Ok(Self { + destination, + port: port.filter(|port| *port != 0), + }) + } + + /// Options applied to every invocation. + /// + /// `BatchMode=yes` is the load-bearing one: without it a host that wants a + /// password or key passphrase blocks forever behind a prompt nobody is + /// watching, which for a status poll or a scripted run means a hang rather + /// than an error. `ConnectTimeout` bounds an unreachable host the same way. + fn base_ssh_args(&self, control_path: Option<&str>) -> Vec { + let mut args = config_args(); + args.extend([ + "-o".to_owned(), + "BatchMode=yes".to_owned(), + "-o".to_owned(), + "ConnectTimeout=10".to_owned(), + ]); + args.extend(multiplex_args(control_path)); + 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, excluding the `ssh` + /// program name. Split out so argument construction is unit-testable without + /// a network or an SSH server. + fn exec_argv(&self, command: &str) -> Vec { + self.exec_argv_with(command, control_socket_path().as_deref()) + } + + /// The argument builder proper, with the socket location passed in. + /// + /// Separated so tests exercise both the multiplexing and no-multiplexing + /// forms without reaching for a real socket directory — discovering one + /// creates a directory under the user's home, which a unit test has no + /// business doing. + fn exec_argv_with(&self, command: &str, control_path: Option<&str>) -> Vec { + let mut args = self.base_ssh_args(control_path); + args.push(self.destination.clone()); + // `--` stops ssh from parsing anything in the remote command as its own + // option, so a model name or flag that happens to start with `-` reaches + // the remote intact instead of being swallowed locally. + args.push("--".to_owned()); + args.push(command.to_owned()); + args + } + + /// Argument vector for `scp`, excluding the `scp` program name. + fn scp_argv(&self, local_path: &str, remote_path: &str) -> Result> { + validate_scp_path(local_path, "local")?; + validate_scp_path(remote_path, "remote")?; + let mut args = config_args(); + args.extend([ + "-o".to_owned(), + "BatchMode=yes".to_owned(), + "-o".to_owned(), + "ConnectTimeout=10".to_owned(), + ]); + // scp spells the port with a capital -P, unlike ssh's lowercase -p. A + // perennial bug when argument-building code is copy-pasted between the + // two, so it has its own test. + 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)); + Ok(args) + } +} + +impl Transport for SshTransport { + fn exec_with_stdin(&self, command: &str, stdin: Option<&str>) -> Result { + let mut ssh = Command::new("ssh"); + ssh.args(self.exec_argv(command)); + run_with_piped_io(ssh, stdin, &self.destination) + } + + fn push_file(&self, local_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() { + bail!( + "failed to copy {local} to {}:{remote_path}: {}", + self.destination, + String::from_utf8_lossy(&output.stderr).trim(), + ); + } + Ok(()) + } +} + +/// Spawn `command` with all three stdio streams piped, optionally feeding it +/// `stdin`, and collect the outcome. +/// +/// Split out of [`SshTransport::exec_with_stdin`] so the pipe handling can be +/// driven by any child process. What it gets right is concurrency-sensitive and +/// only observable against a program that both consumes a large stdin and +/// produces output — and `ssh` is not required to be that program. A shell is, +/// which is what makes the property testable without a reachable host. +/// +/// `destination` appears only in error messages. +fn run_with_piped_io( + mut command: Command, + stdin: Option<&str>, + destination: &str, +) -> Result { + let mut child = command + .stdin(if stdin.is_some() { + Stdio::piped() + } else { + Stdio::null() + }) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .spawn() + .with_context(|| format!("failed to launch ssh to run a command on {destination}"))?; + + // Writing the payload and reading stdout/stderr must not be sequenced + // one after the other: ssh does not have to drain stdin before it + // starts producing output, and with both pipes bounded by the OS + // (commonly ~64KiB), a payload larger than that paired with any + // remote output can deadlock — us blocked in `write_all` waiting for + // ssh to read more of stdin, ssh blocked writing stdout/stderr + // waiting for us to read it, and `wait_with_output` (which would + // drain both) never even reached. Writing on its own thread lets the + // main thread reach `wait_with_output` immediately, which drains + // stdout/stderr concurrently with the write. + let writer = if let Some(payload) = stdin { + let mut handle = child + .stdin + .take() + .context("ssh stdin was not available to write to")?; + let payload = payload.to_owned(); + // `std::io::Result`, not `anyhow::Result`: the join site has to ask + // whether this was a broken pipe, and an `ErrorKind` it can match + // on directly is a far sturdier way to answer that than downcasting + // back out of a context chain. + Some(std::thread::spawn(move || -> std::io::Result<()> { + handle.write_all(payload.as_bytes())?; + // Dropping closes the pipe, which is what tells the remote + // reader the input has ended. Without it a remote `read` + // waits forever. + drop(handle); + Ok(()) + })) + } else { + None + }; + + let waited = child.wait_with_output(); + // Joined before anything can return, including the error path below: an + // un-joined handle detaches the thread, leaving it blocked writing into + // a pipe with no reader for as long as the process lives. + let written = writer.map(std::thread::JoinHandle::join); + + let output = waited + .with_context(|| format!("failed to read the result of a command on {destination}"))?; + + // 255 is ssh's own: it could not connect, could not authenticate, or the + // connection broke. Reporting it as a remote answer is what turns an + // offline or key-less host into "no ROCm installation was found", which + // sends the user to install something on a machine they never reached. + // A remote command can in principle exit 255 itself; the cost of reading + // that rare case as unreachable is far smaller than the cost of the + // confident wrong diagnosis it replaces. + // + // Checked before the write result on purpose. ssh exiting without + // reading stdin breaks the pipe, so an unreachable host produces *two* + // errors describing one event; `output` is the ground truth and the + // write error is the echo. Consulting the echo first is what reported + // "failed to send input" for a host that was never reached, which + // `serve_with_transport` then wrapped as "may or may not be running". + if output.status.code() == Some(SSH_TRANSPORT_FAILURE) { + bail!( + "could not reach {destination} over ssh: {}", + String::from_utf8_lossy(&output.stderr).trim() + ); + } + + match written { + // A panic here is a bug on this side of the connection, so it is + // reported rather than folded into the outcome. Note it is only + // reached when the two returns above did not fire: an unreachable + // host or an unreadable outcome is the more actionable answer, and + // a panic in *this* closure — `write_all`, `drop`, `Ok(())` — is + // not something that happens in practice. Grow the closure and + // that ordering is worth revisiting. + Some(Err(_)) => bail!("the stdin writer thread for {destination} panicked"), + Some(Ok(Err(error))) => { + // A broken pipe means the remote stopped reading. When the + // command also failed, its own exit code and stderr say why, + // and they say it better than "broken pipe" does — so the write + // error is demoted to a symptom. When the command *succeeded*, + // nothing explains the unread payload and it stays an error: + // the sole caller's payload is an API key, and a model started + // without the key that was meant to guard it is not a success. + // + // What keeps a *truncated* key from slipping through the first + // branch is a property of the caller, not of this function: + // `remote_serve_command` puts `IFS= read -r` first in the + // remote command, so a broken pipe means the read never + // finished, which means the command cannot have exited 0. + // Reorder that command and this demotion needs rechecking. + let explained_by_the_command = + error.kind() == std::io::ErrorKind::BrokenPipe && !output.status.success(); + if !explained_by_the_command { + return Err(error) + .with_context(|| format!("failed to send input to {destination}")); + } + } + Some(Ok(Ok(()))) | None => {} + } + + 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(), + }) +} + +/// One programmed reply in a [`ScriptedTransport`]. +#[cfg(test)] +#[derive(Debug, Clone)] +pub(crate) struct ScriptedStep { + /// Substring the remote command must contain for this step to apply. + /// + /// A substring rather than an exact string on purpose: a test that pins + /// whole command lines re-breaks every time an unrelated flag is added, + /// which trains people to update expectations without reading them. + matches: String, + outcome: RemoteOutcome, +} + +#[cfg(test)] +impl ScriptedStep { + /// A step whose command succeeds, returning `stdout`. + pub(crate) fn ok(matches: &str, stdout: &str) -> Self { + Self { + matches: matches.to_owned(), + outcome: RemoteOutcome { + success: true, + code: Some(0), + stdout: stdout.to_owned(), + stderr: String::new(), + }, + } + } + + /// A step whose command runs but exits non-zero — the remote answering "no", + /// as distinct from being unreachable. + pub(crate) fn fails(matches: &str, code: i32, stderr: &str) -> Self { + Self { + matches: matches.to_owned(), + outcome: RemoteOutcome { + success: false, + code: Some(code), + stdout: String::new(), + stderr: stderr.to_owned(), + }, + } + } +} + +/// What a [`ScriptedTransport`] was asked to do, in order. +#[cfg(test)] +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) enum TransportCall { + Exec { + command: String, + stdin: Option, + }, + PushFile { + local_path: String, + remote_path: String, + }, +} + +/// A [`Transport`] stand-in that replays programmed outcomes and records every +/// call, so bootstrap and session flows can be tested without a network, an SSH +/// server, or a tailnet. +/// +/// Unmatched commands are a hard error rather than a benign default. A silently +/// tolerant double lets a flow test keep passing after the flow stops issuing a +/// command it is supposed to issue, which is exactly the regression these tests +/// exist to catch. +#[cfg(test)] +pub(crate) struct ScriptedTransport { + steps: Vec, + calls: RefCell>, + push_result: RefCell>, +} + +#[cfg(test)] +impl ScriptedTransport { + pub(crate) fn new(steps: Vec) -> Self { + Self { + steps, + calls: RefCell::new(Vec::new()), + push_result: RefCell::new(Ok(())), + } + } + + /// Make the next and all subsequent file copies fail, for testing the + /// provisioning fallback paths. + pub(crate) fn failing_push(self, message: &str) -> Self { + *self.push_result.borrow_mut() = Err(message.to_owned()); + self + } + + /// Every call made so far, in order. + pub(crate) fn calls(&self) -> Vec { + self.calls.borrow().clone() + } +} + +#[cfg(test)] +impl Transport for ScriptedTransport { + fn exec_with_stdin(&self, command: &str, stdin: Option<&str>) -> Result { + self.calls.borrow_mut().push(TransportCall::Exec { + command: command.to_owned(), + stdin: stdin.map(ToOwned::to_owned), + }); + self.steps + .iter() + .filter(|step| command.contains(&step.matches)) + // Most specific wins. One remote command is often a suffix of + // another — `rocm --version` and `$HOME/.local/bin/rocm --version` + // differ only by a prefix — so first-match would answer the second + // with the first's reply and silently test the wrong branch. + .max_by_key(|step| step.matches.len()) + .map(|step| step.outcome.clone()) + .with_context(|| { + format!( + "the scripted transport has no reply for: {command}\n programmed: {:?}", + self.steps + .iter() + .map(|step| step.matches.as_str()) + .collect::>() + ) + }) + } + + fn push_file(&self, local_path: &Path, remote_path: &str) -> Result<()> { + self.calls.borrow_mut().push(TransportCall::PushFile { + local_path: local_path.to_string_lossy().into_owned(), + remote_path: remote_path.to_owned(), + }); + match &*self.push_result.borrow() { + Ok(()) => Ok(()), + Err(message) => bail!("{message}"), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn exec_argv_never_prompts_and_guards_the_remote_command() { + let transport = SshTransport::new("user@gpubox", None).unwrap(); + let argv = transport.exec_argv_with("rocm --version", None); + + // Without BatchMode a host wanting a passphrase hangs instead of failing. + assert!(argv.windows(2).any(|pair| pair == ["-o", "BatchMode=yes"])); + assert_eq!(argv.last().unwrap(), "rocm --version"); + // The remote command sits behind `--`, immediately after the destination. + let guard = argv.iter().position(|arg| arg == "--").unwrap(); + assert_eq!(argv[guard - 1], "user@gpubox"); + } + + #[test] + fn a_name_ssh_would_read_as_an_option_is_refused() { + // The destination sits before the `--` guard, so unlike the remote + // command it is not shielded by it. `-oProxyCommand=…` as a "host" runs + // a command on this machine. + for hostile in ["-oProxyCommand=touch /tmp/pwned", "--fake", "user@-oX", ""] { + let error = SshTransport::new(hostile, None) + .expect_err(&format!("{hostile:?} should be refused")) + .to_string(); + assert!(!error.is_empty()); + } + + // Ordinary destinations still work, including a user prefix and an alias. + for fine in [ + "gpu-box", + "user@gpu-box", + "gpu-box.example-tailnet.ts.net", + "100.88.14.21", + ] { + SshTransport::new(fine, None).unwrap_or_else(|error| panic!("{fine}: {error}")); + } + } + + #[test] + fn connection_reuse_is_either_fully_configured_or_not_claimed() { + // OpenSSH ignores ControlMaster when ControlPath is unset, so the two + // options without the third are multiplexing that silently never + // happens — every status poll paying a fresh handshake while the code + // claims otherwise. + let configured = multiplex_args(Some("/tmp/rocm/cm-%C")); + assert!( + configured + .windows(2) + .any(|pair| pair == ["-o", "ControlMaster=auto"]), + "{configured:?}" + ); + assert!( + configured + .iter() + .any(|arg| arg == "ControlPath=/tmp/rocm/cm-%C"), + "{configured:?}" + ); + + // With nowhere private to put the socket, emit nothing at all. + assert!(multiplex_args(None).is_empty()); + } + + #[test] + fn no_multiplexing_option_ever_appears_without_a_socket_path() { + // The invariant, stated over the real argument vector rather than the + // helper: a reader scanning for ControlMaster should never find it + // orphaned, in either configuration. + let transport = SshTransport::new("gpubox", None).unwrap(); + for control_path in [None, Some("/tmp/rocm-ssh/cm-%C")] { + let argv = transport.exec_argv_with("true", control_path); + let has_master = argv.iter().any(|arg| arg.starts_with("ControlMaster")); + let has_path = argv.iter().any(|arg| arg.starts_with("ControlPath")); + assert_eq!( + has_master, has_path, + "ControlMaster and ControlPath must appear together: {argv:?}" + ); + assert_eq!(has_master, control_path.is_some(), "{argv:?}"); + } + } + + #[test] + fn exec_argv_threads_an_explicit_port_lowercase() { + let transport = SshTransport::new("user@gpubox", Some(2222)).unwrap(); + let argv = transport.exec_argv_with("echo hi", None); + assert!(argv.windows(2).any(|pair| pair == ["-p", "2222"])); + } + + #[test] + fn exec_argv_omits_an_unset_port_so_ssh_config_decides() { + // Passing no port must leave the choice to ~/.ssh/config rather than + // hardcoding 22, or a configured non-standard port is silently ignored. + let transport = SshTransport::new("gpubox", None).unwrap(); + let argv = transport.exec_argv_with("echo hi", None); + assert!(!argv.iter().any(|arg| arg == "-p")); + } + + #[test] + fn scp_argv_uses_uppercase_port_and_a_remote_colon_path() { + let transport = SshTransport::new("user@gpubox", Some(2222)).unwrap(); + let argv = transport.scp_argv("/tmp/rocm", "/tmp/rocm").unwrap(); + // Capital -P: scp's port flag differs from ssh's, and getting it wrong + // silently copies to the default port instead. + assert!(argv.windows(2).any(|pair| pair == ["-P", "2222"])); + assert_eq!(argv.last().unwrap(), "user@gpubox:/tmp/rocm"); + } + + #[test] + fn scp_argv_refuses_a_local_or_remote_path_starting_with_a_dash() { + // Both are refused, for different reasons. The local path is its own + // argument, so scp really would read a leading `-` as an option — + // that half mirrors validate_destination's guard. The remote path is + // prefixed with `{destination}:` before it is passed, so it can never + // be mistaken for an option; refusing it is a check on our own + // argument building rather than a safety guard. + let transport = SshTransport::new("user@gpubox", None).unwrap(); + let local = transport + .scp_argv("-oProxyCommand=evil", "/tmp/rocm") + .unwrap_err() + .to_string(); + assert!(local.contains("local path"), "{local}"); + + let remote = transport + .scp_argv("/tmp/rocm", "-oProxyCommand=evil") + .unwrap_err() + .to_string(); + assert!(remote.contains("remote path"), "{remote}"); + } + + #[test] + fn sshs_own_failure_code_is_not_read_as_a_remote_answer() { + // Exercised through a real `ssh` against a port nothing is listening on, + // because the property under test is what the client does, not what a + // constant says. Asserting the constant against itself proved nothing. + let unreachable = SshTransport::new("127.0.0.1", Some(1)).expect("valid destination"); + let error = unreachable + .exec("true") + .expect_err("an unreachable host must be a transport failure, not a remote answer"); + + let message = format!("{error:#}"); + assert!( + message.contains("could not reach"), + "an unreachable host must say so rather than look like a command that failed: \ + {message}" + ); + } + + #[cfg(unix)] + #[test] + fn a_large_payload_and_a_flood_of_output_do_not_deadlock() { + // The bug the writer thread exists to prevent, and — until this test — + // the one thing in this file that no test failed on if the fix were + // reverted. Both pipes are bounded (commonly ~64KiB): write the payload + // inline and the child fills stdout while we are still blocked in + // `write_all`, so neither side can move and `wait_with_output` is never + // reached. + // + // `sh` rather than `ssh`, because the property is about our pipe + // handling, not about ssh — and a reachable host would make this a + // container test instead of a unit one. + // + // The ordering is the whole test: the child floods stdout *first* and + // only then reads stdin. Run the two concurrently instead and the child + // keeps draining stdin throughout, so an inline write finishes and + // nothing hangs — a version of this test that backgrounded the writer + // passed with the fix reverted, which is to say it tested nothing. + let mut command = Command::new("sh"); + command.args([ + "-c", + "yes deadlock-canary | head -c 2000000; cat >/dev/null", + ]); + + let payload = "x".repeat(4 << 20); + let (tx, rx) = std::sync::mpsc::channel(); + std::thread::spawn(move || { + let _ = tx.send(run_with_piped_io(command, Some(&payload), "local-shell")); + }); + + // A deadlock shows up as a timeout, not a failed assertion, so the + // bound is the assertion. Ten seconds is far past the ~0.1s this takes + // when it works, and far short of hanging the suite. + let outcome = rx + .recv_timeout(std::time::Duration::from_secs(10)) + .expect("writing a large payload while the child floods stdout must not deadlock") + .expect("the shell ran"); + + assert!(outcome.success, "{outcome:?}"); + assert!( + outcome.stdout.len() > 64 * 1024, + "the child must have produced more than one pipe buffer of output, or the \ + test would pass without exercising the hazard: {} bytes", + outcome.stdout.len() + ); + } + + #[test] + fn an_unreachable_host_says_so_even_when_a_payload_was_being_written() { + // Same unreachable host as above, but with stdin to deliver. ssh exits + // 255 without ever reading it, so the writer thread's `write_all` ends + // in EPIPE. Two errors then describe the same event, and only one is + // the truth the user needs: consulting the write error first reports + // "failed to send input", which `serve_with_transport` wraps as "may + // or may not be running" — telling the user the model's state is + // unknown when in fact nothing was ever started. The process outcome + // is ground truth and must win. + // + // The payload is deliberately larger than a pipe buffer (commonly + // ~64KiB). A short one lands in the buffer and returns success with no + // reader on the other end, so the race decides whether the bug appears + // and the test would pass on a machine that simply scheduled the + // writer first. Overflowing the buffer makes `write_all` block until + // ssh exits, which makes the EPIPE certain rather than likely — and + // exercises the very overflow the writer thread exists to survive. + let unreachable = SshTransport::new("127.0.0.1", Some(1)).expect("valid destination"); + let payload = "x".repeat(1 << 20); + let error = unreachable + .exec_with_stdin("true", Some(&payload)) + .expect_err("an unreachable host must be a transport failure, not a remote answer"); + + let message = format!("{error:#}"); + assert!( + message.contains("could not reach"), + "the broken pipe from the payload write must not displace ssh's own 255: {message}" + ); + } + + #[test] + fn a_transport_failure_is_not_something_a_probe_can_mistake_for_absence() { + // The bug this guards: a probe asks "is ROCm here?", ssh answers 255 + // because it never connected, and the caller reads a normal unsuccessful + // outcome and tells the user to install ROCm on a machine it never + // reached. `run` must surface an Err, not an outcome. + let unreachable = SshTransport::new("127.0.0.1", Some(1)).expect("valid destination"); + assert!(unreachable.run("command -v rocminfo").is_err()); + } + + #[test] + fn scripted_transport_separates_a_refusal_from_being_unreachable() { + // A remote that answers "no CLI here" (exit 127) is not the same as a + // remote we could not reach; the first is an outcome, the second an Err. + let transport = ScriptedTransport::new(vec![ScriptedStep::fails( + "rocm --version", + 127, + "command not found", + )]); + + let outcome = transport.exec("rocm --version").expect("host answered"); + assert!(!outcome.success); + assert_eq!(outcome.code, Some(127)); + + let unreachable = transport.exec("tailscale status"); + assert!( + unreachable.is_err(), + "an unscripted command must fail loudly, not return a benign default" + ); + } + + #[test] + fn run_surfaces_the_failing_command_and_its_stderr() { + let transport = ScriptedTransport::new(vec![ScriptedStep::fails( + "rocm serve", + 1, + "no GPU available", + )]); + + let error = transport + .run("rocm serve tiny-model") + .unwrap_err() + .to_string(); + assert!(error.contains("no GPU available"), "{error}"); + assert!(error.contains("rocm serve tiny-model"), "{error}"); + } + + #[test] + fn scripted_transport_records_stdin_so_secrets_stay_out_of_argv() { + // The recording is what lets a later test assert an API key was piped in + // rather than interpolated into the command line. + let transport = ScriptedTransport::new(vec![ScriptedStep::ok("read -r", "")]); + transport + .exec_with_stdin("read -r KEY; exec rocm serve", Some("s3cret")) + .expect("scripted"); + + assert_eq!( + transport.calls(), + vec![TransportCall::Exec { + command: "read -r KEY; exec rocm serve".to_owned(), + stdin: Some("s3cret".to_owned()), + }] + ); + } + + #[test] + fn scripted_transport_reports_a_failing_file_copy() { + let transport = ScriptedTransport::new(vec![]).failing_push("no space left on device"); + + let error = transport + .push_file(Path::new("/tmp/rocm"), "~/.local/bin/rocm") + .unwrap_err() + .to_string(); + assert!(error.contains("no space left on device"), "{error}"); + assert_eq!( + transport.calls(), + vec![TransportCall::PushFile { + local_path: "/tmp/rocm".to_owned(), + remote_path: "~/.local/bin/rocm".to_owned(), + }] + ); + } +} diff --git a/apps/rocmd/src/lib.rs b/apps/rocmd/src/lib.rs index dfc287f45..3762212e3 100644 --- a/apps/rocmd/src/lib.rs +++ b/apps/rocmd/src/lib.rs @@ -2526,6 +2526,14 @@ fn ensure_rocm_command_is_read_only(args: &[String]) -> Result<()> { // itself reopen onboarding). Mirrors the bin's rocm_command classifier so // the read-only allowlist is consistent across binaries. Some("setup") => second.as_deref().is_none_or(|value| value == "status"), + // `remote targets` reads the local tailnet, `doctor` fetches another + // machine's state and scores it here, `status` probes sessions that + // already exist. None of them change anything on either machine. + // `serve`, `attach` and `stop` start, publish or tear down, so they stay + // off the list and go through the approval UI like any other mutation. + Some("remote") => second + .as_deref() + .is_some_and(|value| matches!(value, "targets" | "doctor" | "status")), _ => false, }; if read_only { @@ -3253,6 +3261,32 @@ fn supervise_service( ); record.gpu_indices = gpu_indices; record.engine_recipe_json = engine_recipe_json.clone(); + // Carried over from whatever is on disk. `ManagedServiceRecord::new` starts + // this false, so rebuilding a record here without restoring it would not + // just skip the check now — it would write the weakened record back and + // disarm every later `rocm services restart` as well. + // + // Propagated, not defaulted. This read arms the guard below, so it is not + // best-effort the way an identical-looking call feeding a printed warning + // would be. `load_managed_services` already *skips* unparseable records, so + // an `Err` here is a real I/O failure — and a missing directory is `Ok` + // anyway. Swallowing it would say "no service ever required a key", the + // key-file fallback is false precisely when a service has been stopped, and + // the weakened record would then be written back at the bottom of this + // function. That is the outcome the comment above says must not happen. + let previously_required = load_managed_services(paths) + .context( + "could not read the service registry to check whether this service requires an \ + endpoint API key; refusing to recover it rather than assume it does not", + )? + .iter() + .any(|existing| existing.service_id == record.service_id && existing.requires_api_key); + // Only what the registry recorded. The `|| key-file-is-present` clause that + // used to be here re-derived the flag the same way `spawn_managed_engine_child` + // did, and was wrong for the same reason: a public bind always has a key file + // whether or not auth was ever demanded, so recovery re-armed this on services + // that never asked for it and refused them with the wrong remediation. + record.requires_api_key = previously_required; // Refuse a keyless public respawn before the manifest write, so a refused // attempt leaves the recorded restart_count and timestamps intact instead of // clobbering them with a record no live process will ever back. The spawn @@ -3262,6 +3296,7 @@ fn supervise_service( rocm_engine_protocol::endpoint_key_file_if_present(paths, &record.service_id) .and_then(|path| rocm_engine_protocol::endpoint_api_key_file_if_valid(&path)) .is_some(), + record.requires_api_key, )?; record.write()?; @@ -3300,7 +3335,11 @@ fn supervise_service( // no auth, so fail closed instead — an unreachable service is recoverable, // an anonymous public one is not. let endpoint_key_applied = apply_endpoint_key_env(&mut command, paths, &record.service_id); - ensure_public_service_has_endpoint_key(&record.host, endpoint_key_applied)?; + ensure_public_service_has_endpoint_key( + &record.host, + endpoint_key_applied, + record.requires_api_key, + )?; let mut child = command .spawn() .with_context(|| format!("failed to spawn engine supervisor child for {engine}"))?; @@ -4822,6 +4861,7 @@ fn handle_server_recover_event_with_record( rocm_engine_protocol::endpoint_key_file_if_present(paths, &record.service_id) .and_then(|path| rocm_engine_protocol::endpoint_api_key_file_if_valid(&path)) .is_some(), + record.requires_api_key, ) { return record_event( paths, @@ -5226,6 +5266,261 @@ mod tests { use rocm_core::ModelRecipeArtifactSourcePolicyRecord; use std::path::PathBuf; + #[test] + fn recovery_refuses_a_service_that_lost_a_key_it_was_launched_with() { + // The daemon keeps its own copy of this guard, and it only knew about + // public binds. A loopback service that something republishes — a + // tailnet publish outlives this daemon, let alone the process — would + // be recovered without authentication, and the rebuilt record would + // then disarm `rocm services restart` too. + for host in ["127.0.0.1", "localhost", "::1"] { + let error = super::ensure_public_service_has_endpoint_key(host, false, true) + .expect_err("a service launched with a key must not be recovered without one"); + assert!( + format!("{error:#}").contains("without authentication"), + "{error:#}" + ); + } + // With the key still present, recovery proceeds. + super::ensure_public_service_has_endpoint_key("127.0.0.1", true, true).unwrap(); + // And a service that never had one is untouched. + super::ensure_public_service_has_endpoint_key("127.0.0.1", false, false).unwrap(); + } + + #[test] + fn remote_read_only_verbs_are_allowed_and_mutating_ones_are_not() { + let allow = |args: &[&str]| { + let owned = args.iter().map(|a| (*a).to_owned()).collect::>(); + super::ensure_rocm_command_is_read_only(&owned) + }; + + // These read: the local tailnet, another machine's state, sessions that + // already exist. Rejecting them made the whole family unusable here even + // though none of them change anything. + for args in [ + &["remote", "targets"][..], + &["remote", "targets", "--tag", "gpu"][..], + &["remote", "doctor", "gpu-box"][..], + &["remote", "status"][..], + ] { + allow(args).unwrap_or_else(|error| panic!("{args:?} should be read-only: {error:#}")); + } + + // These start, publish or tear down, so they go through approval. + for args in [ + &["remote", "serve", "gpu-box", "a-model"][..], + &["remote", "attach", "sess"][..], + &["remote", "stop", "sess"][..], + &["remote"][..], + ] { + assert!(allow(args).is_err(), "{args:?} must not be read-only"); + } + } + + /// Drive `supervise_service` far enough to reach the key guard, and return + /// what it did. + /// + /// The guard sits before the manifest write and well before any spawn, so a + /// refusal returns without starting a process — which is what makes the real + /// call site testable at all. The arguments below are the shape a recovery + /// re-exec passes: a loopback bind, no GPU, no recipe. + /// + /// This exists because testing `ensure_public_service_has_endpoint_key` + /// directly with literal arguments cannot catch the defect that actually + /// happened twice in this crate's history — the guard being *wired up* with + /// the wrong value at its call site. + /// + /// Bounded, and the bound is the assertion. A guard that fails to refuse + /// does not return an error — it falls through to the engine spawn and + /// supervises a child that never exits, so an unbounded call would hang the + /// suite instead of failing it. Both callers below are regression tests for + /// a fail-*open*, which is exactly the shape that turns into a hang. + fn supervise_at_the_guard(paths: &AppPaths, service_id: &str) -> Result<()> { + let paths = paths.clone(); + let service_id = service_id.to_owned(); + let (sender, receiver) = std::sync::mpsc::channel(); + std::thread::spawn(move || { + let outcome = supervise_service( + &paths, + service_id, + "llamacpp".to_owned(), + "a-model".to_owned(), + "a-model".to_owned(), + None, + None, + "127.0.0.1".to_owned(), + 11434, + "gpu_required".to_owned(), + None, + None, + ); + let _ = sender.send(outcome.map_err(|error| format!("{error:#}"))); + }); + receiver + .recv_timeout(std::time::Duration::from_secs(30)) + .unwrap_or_else(|_| { + panic!( + "supervise_service did not return within 30s: the key guard let the call \ + through and it reached the engine spawn, which is the fail-open this test \ + exists to catch" + ) + }) + .map_err(anyhow::Error::msg) + } + + /// Write a service record into the registry the way a live service would + /// have left it behind. + fn seed_registry(paths: &AppPaths, service_id: &str, requires_api_key: bool) { + fs::create_dir_all(paths.services_dir()).unwrap(); + let mut record = ManagedServiceRecord::new( + paths, + service_id.to_owned(), + "llamacpp".to_owned(), + "a-model".to_owned(), + "a-model".to_owned(), + "127.0.0.1".to_owned(), + 11434, + "managed", + std::process::id(), + None, + None, + Some("gpu_required".to_owned()), + ); + record.requires_api_key = requires_api_key; + record.write().unwrap(); + } + + /// Drive `supervise_service` and report whether the key guard let it past. + /// + /// Decided on what the call *returns*, not on any file. The obvious + /// observable — the manifest appearing — is useless here, because + /// `seed_registry` has already written one, so polling for it passes + /// whatever the guard does. That mistake was made first and caught by + /// mutating the code the test claims to protect. + /// + /// A call the guard admits does not return: it carries on to the engine + /// spawn. So the guard's refusal is the only thing that comes back quickly, + /// and it is identified by its message rather than by the mere fact of an + /// error — a later, unrelated failure must not read as a refusal. + fn guard_admits(paths: &AppPaths, service_id: &str) -> bool { + let owned_paths = paths.clone(); + let owned_id = service_id.to_owned(); + let (sender, receiver) = std::sync::mpsc::channel(); + std::thread::spawn(move || { + let outcome = supervise_service( + &owned_paths, + owned_id, + "llamacpp".to_owned(), + "a-model".to_owned(), + "a-model".to_owned(), + None, + None, + "127.0.0.1".to_owned(), + 11434, + "gpu_required".to_owned(), + None, + None, + ); + let _ = sender.send(outcome.map_err(|error| format!("{error:#}"))); + }); + + match receiver.recv_timeout(std::time::Duration::from_secs(10)) { + // The key guard refused, by its own words. + Ok(Err(rendered)) if rendered.contains("without authentication") => false, + // Anything else means it got past the guard: it either finished, or + // failed later for a reason that is not this guard, or is still + // running because it reached the spawn. + _ => true, + } + } + + #[test] + fn a_service_that_never_required_a_key_is_not_refused_for_lacking_one() { + // The other direction of the guard, and the one no test covered. + // Hardcoding `record.requires_api_key = true` at the restore site passes + // every other test in this crate, because they all seed a service that + // *does* require a key. This is the case that catches it. + // + // Two records are seeded, not one: with a single record the + // `existing.service_id == record.service_id` half of the lookup does + // nothing, so dropping that comparison would go unnoticed and one + // service's requirement would leak onto another's. + let (root, paths) = temp_app_paths("supervise-no-key-needed"); + seed_registry(&paths, "svc-needs-key", true); + seed_registry(&paths, "svc-plain", false); + + assert!( + guard_admits(&paths, "svc-plain"), + "a loopback service that never asked for a key must not be refused for lacking one" + ); + + let _ = fs::remove_dir_all(&root); + } + + #[test] + fn supervising_a_service_that_required_a_key_refuses_when_the_key_is_gone() { + // The real call site, not the guard in isolation. `supervise_service` + // rebuilds the record with `ManagedServiceRecord::new`, which starts + // `requires_api_key` false, and restores it from the registry. Passing + // the wrong value here — a literal, or the freshly-built field before it + // is restored — is exactly the miswiring that shipped twice in this + // crate and that a literal-argument unit test cannot see. + let (root, paths) = temp_app_paths("supervise-requires-key"); + seed_registry(&paths, "svc-needs-key", true); + + let error = supervise_at_the_guard(&paths, "svc-needs-key") + .expect_err("a service that required a key must not be recovered without one"); + assert!( + format!("{error:#}").contains("without authentication"), + "{error:#}" + ); + + // And the refusal must not have weakened what is on disk. The guard runs + // before `record.write()` precisely so a refused attempt leaves the + // recorded requirement armed for the next attempt. + let stored = load_managed_services(&paths).unwrap(); + let stored = stored + .iter() + .find(|candidate| candidate.service_id == "svc-needs-key") + .expect("the seeded record must survive a refused recovery"); + assert!( + stored.requires_api_key, + "a refused recovery must not disarm the requirement" + ); + + let _ = fs::remove_dir_all(&root); + } + + #[test] + fn a_registry_that_cannot_be_read_refuses_recovery_rather_than_assuming_no_key() { + // `load_managed_services` already skips records it cannot parse, so an + // `Err` from it is a real I/O failure — and a missing directory is `Ok`. + // Defaulting it away therefore says "no service ever required a key", + // which is fail-open on an auth gate and, worse, gets written back. + // + // The failure is provoked portably: a directory named like a record + // makes the `fs::read` inside the loop fail rather than the read_dir. + let (root, paths) = temp_app_paths("supervise-unreadable-registry"); + fs::create_dir_all(paths.services_dir().join("not-a-record.json")).unwrap(); + + let error = supervise_at_the_guard(&paths, "svc-unknown") + .expect_err("an unreadable registry must refuse, not assume no key was required"); + let rendered = format!("{error:#}"); + assert!( + rendered.contains("could not read the service registry"), + "{rendered}" + ); + + // Nothing was written: a registry we could not read is not a registry we + // may add a weakened record to. + assert!( + !paths.service_manifest_path("svc-unknown").exists(), + "a refused recovery must not persist a record" + ); + + let _ = fs::remove_dir_all(&root); + } + /// Regression: a failed write must not leave a `.tmp-*` scratch file /// behind. The name is unique per attempt, so before this an orphan /// accumulated per retry — and when the failure is a full disk, those @@ -7592,15 +7887,15 @@ mod tests { // Daemon recovery re-execs `rocmd supervise` for the recorded host. With // the key gone the child would listen on that public host anonymously, // so the spawn must be refused instead. - let error = ensure_public_service_has_endpoint_key("0.0.0.0", false).unwrap_err(); + let error = ensure_public_service_has_endpoint_key("0.0.0.0", false, false).unwrap_err(); assert!( error.to_string().contains("without authentication"), "{error:#}" ); - ensure_public_service_has_endpoint_key("0.0.0.0", true).unwrap(); + ensure_public_service_has_endpoint_key("0.0.0.0", true, false).unwrap(); for host in ["127.0.0.1", "localhost", "::1"] { - ensure_public_service_has_endpoint_key(host, false) + ensure_public_service_has_endpoint_key(host, false, false) .unwrap_or_else(|error| panic!("{host} must not require a key: {error:#}")); } } @@ -9562,7 +9857,24 @@ fn apply_endpoint_key_env( /// Mirrors the guard of the same name in `rocm`; the shared /// [`rocm_engine_protocol::is_public_bind_host`] keeps the two classifications /// identical for a given `ManagedServiceRecord::host`. -fn ensure_public_service_has_endpoint_key(host: &str, key_present: bool) -> Result<()> { +fn ensure_public_service_has_endpoint_key( + host: &str, + key_present: bool, + requires_api_key: bool, +) -> Result<()> { + // The bind address is not the whole story. A service bound to loopback is + // only private until something republishes the port, and a tailnet publish + // outlives both the process and this daemon. The requirement is recorded on + // the service precisely so recovery can honour it without re-deriving it + // from an address that no longer answers the question. + if requires_api_key && !key_present { + bail!( + "refusing to recover a service that was launched with an endpoint API key but no \ + longer has one: it would come back up without authentication, and something \ + outside this machine may still be publishing its port. Relaunch it with \ + `rocm serve --require-api-key` to issue a new key." + ); + } if rocm_engine_protocol::is_public_bind_host(host) && !key_present { bail!( "refusing to respawn a service bound to the public host `{host}` without an endpoint \ diff --git a/crates/e2e-report/src/lib.rs b/crates/e2e-report/src/lib.rs index 5d0d3e157..d8f801fe0 100644 --- a/crates/e2e-report/src/lib.rs +++ b/crates/e2e-report/src/lib.rs @@ -1665,6 +1665,12 @@ struct CommandKey { /// coverage % reflects the real surface (a deliberate, reviewable denominator /// beats silently drifting). const KNOWN_COMMAND_SURFACE: &[&str] = &[ + "rocm remote targets", + "rocm remote serve", + "rocm remote doctor", + "rocm remote status", + "rocm remote attach", + "rocm remote stop", "rocm examine", "rocm diagnose", "rocm fix", diff --git a/crates/rocm-core/src/lib.rs b/crates/rocm-core/src/lib.rs index 77bd6bbd0..19ee9b6a7 100644 --- a/crates/rocm-core/src/lib.rs +++ b/crates/rocm-core/src/lib.rs @@ -1842,6 +1842,17 @@ impl AppPaths { self.services_dir().join("launch.lock") } + /// Where `rocm remote` records the sessions it started on other machines. + /// + /// Kept beside [`Self::services_dir`] and following the same file-per-record + /// shape, but deliberately separate: these describe work running on a + /// *different* machine, and anything walking the local service registry + /// (status rendering, the daemon's recovery supervisor) must not mistake a + /// remote session for a local server it can supervise. + pub fn remote_sessions_dir(&self) -> PathBuf { + self.data_dir.join("remote-sessions") + } + pub fn audit_dir(&self) -> PathBuf { self.data_dir.join("audit") } @@ -7548,6 +7559,15 @@ pub struct ManagedServiceRecord { /// identity token. `None` until the engine state records one. #[serde(default)] pub engine_start_ticks: Option, + /// Whether this service must never come back up without an endpoint key. + /// + /// The bind address alone cannot answer that. A service bound to loopback is + /// unreachable from elsewhere *until something republishes the port* — a + /// tailnet publish, a proxy, a container port map — and the publish outlives + /// the process. So the requirement has to be recorded next to the service and + /// survive a restart, exactly as the key itself does. + #[serde(default)] + pub requires_api_key: bool, #[serde(default)] pub runtime_id: Option, #[serde(default)] @@ -7631,6 +7651,9 @@ impl ManagedServiceRecord { engine_pid: None, supervisor_start_ticks: None, engine_start_ticks: None, + // Off unless a caller says otherwise: local loopback serving stays + // credential-free, which is the unchanged default. + requires_api_key: false, runtime_id, env_id, device_policy, diff --git a/install.sh b/install.sh index 5466bfe08..1bf9d3184 100755 --- a/install.sh +++ b/install.sh @@ -306,11 +306,15 @@ ensure_installer_process_path() { esac } -os="$(uname -s)" -arch="$(uname -m)" +# Which machine the artifact is *for*. Normally this machine, but the target can +# be named explicitly so an artifact can be fetched here on behalf of a different +# one -- see ROCM_CLI_DOWNLOAD_ONLY below. Detection is the only thing that +# changes; naming a target does not relax any verification. +os="${ROCM_CLI_TARGET_OS:-$(uname -s)}" +arch="${ROCM_CLI_TARGET_ARCH:-$(uname -m)}" case "${os}" in - Linux) platform_os="linux" ;; + Linux|linux) platform_os="linux" ;; *) fail "unsupported OS: ${os} (installer currently supports Linux x86_64 only)" ;; @@ -323,6 +327,29 @@ case "${arch}" in ;; esac +# Fetch and verify an artifact, then stop without installing it. Used to obtain a +# build for a machine that cannot reach the release host itself: this machine +# downloads it, verifies it, and hands the verified files over for delivery. +# +# The artifact is emitted with its checksum and signature sidecars so that +# whatever installs it later can repeat every check performed here. Nothing is +# skipped on either side -- this splits the trust chain across two machines, it +# does not shorten it. +DOWNLOAD_ONLY=0 +if truthy "${ROCM_CLI_DOWNLOAD_ONLY:-0}"; then + DOWNLOAD_ONLY=1 +fi +DOWNLOAD_DIR="${ROCM_CLI_DOWNLOAD_DIR:-.}" + +# Install from an artifact already on disk instead of downloading one. The +# receiving half of the split above. Its .sha256 sidecar is required, not +# optional: an artifact that arrived over the network is exactly the one whose +# integrity still has to be proven. +LOCAL_ARCHIVE="${ROCM_CLI_ARCHIVE:-}" +if [ -n "${LOCAL_ARCHIVE}" ] && [ "${DOWNLOAD_ONLY}" -eq 1 ]; then + fail "ROCM_CLI_ARCHIVE and ROCM_CLI_DOWNLOAD_ONLY are mutually exclusive" +fi + case "${CHANNEL}" in nightly) asset_base="rocm-cli-nightly-${platform_os}-${platform_arch}.tar.gz" @@ -358,11 +385,29 @@ sig_path="${archive_path}.sig" echo "rocm-cli installer" echo " repo: ${REPO}" echo " channel: ${CHANNEL}" -echo " install_dir: ${INSTALL_DIR}" -echo " download: ${archive_url}" +echo " platform: ${platform_os}-${platform_arch}" +if [ -n "${LOCAL_ARCHIVE}" ]; then + echo " archive: ${LOCAL_ARCHIVE}" +elif [ "${DOWNLOAD_ONLY}" -eq 1 ]; then + echo " download_only: ${DOWNLOAD_DIR}" + echo " download: ${archive_url}" +else + echo " install_dir: ${INSTALL_DIR}" + echo " download: ${archive_url}" +fi -fetch "${archive_url}" "${archive_path}" -fetch "${sha_url}" "${sha_path}" +if [ -n "${LOCAL_ARCHIVE}" ]; then + [ -f "${LOCAL_ARCHIVE}" ] || fail "archive not found: ${LOCAL_ARCHIVE}" + [ -f "${LOCAL_ARCHIVE}.sha256" ] || fail "archive checksum not found: ${LOCAL_ARCHIVE}.sha256" + cp "${LOCAL_ARCHIVE}" "${archive_path}" + cp "${LOCAL_ARCHIVE}.sha256" "${sha_path}" + if [ -f "${LOCAL_ARCHIVE}.sig" ]; then + cp "${LOCAL_ARCHIVE}.sig" "${sig_path}" + fi +else + fetch "${archive_url}" "${archive_path}" + fetch "${sha_url}" "${sha_path}" +fi expected="$(awk '{print $1}' "${sha_path}" | head -n1)" [ -n "${expected}" ] || fail "checksum file did not contain a sha256 digest" @@ -385,11 +430,31 @@ fi if [ "${require_sig}" -eq 1 ] || [ -n "${public_keys}" ]; then [ -n "${public_keys}" ] || fail "signature verification requires ROCM_CLI_SIGNING_PUBLIC_KEY_PATH or ROCM_CLI_SIGNING_PUBLIC_KEY_PEM" - fetch "${sig_url}" "${sig_path}" "required signature sidecar is missing or unavailable: ${sig_url}" + if [ -n "${LOCAL_ARCHIVE}" ]; then + # A locally-supplied archive has no URL to fall back on: the signature had to + # travel with it. Refusing here is the point -- an artifact delivered out of + # band is precisely the one whose provenance cannot be assumed. + [ -f "${sig_path}" ] || fail "required signature sidecar is missing: ${LOCAL_ARCHIVE}.sig" + else + fetch "${sig_url}" "${sig_path}" "required signature sidecar is missing or unavailable: ${sig_url}" + fi verify_signature "${archive_path}" "${sig_path}" "${public_keys}" echo "signature verified" fi +# Everything above ran unchanged. Only now, with the artifact proven, is it +# handed over rather than installed. +if [ "${DOWNLOAD_ONLY}" -eq 1 ]; then + mkdir -p "${DOWNLOAD_DIR}" + install -m 0644 "${archive_path}" "${DOWNLOAD_DIR}/${asset_base}" + install -m 0644 "${sha_path}" "${DOWNLOAD_DIR}/${asset_base}.sha256" + if [ -f "${sig_path}" ]; then + install -m 0644 "${sig_path}" "${DOWNLOAD_DIR}/${asset_base}.sig" + fi + echo "downloaded: ${DOWNLOAD_DIR}/${asset_base}" + exit 0 +fi + extract_dir="${tmp_dir}/extract" mkdir -p "${extract_dir}" tar -xzf "${archive_path}" -C "${extract_dir}" diff --git a/tests/e2e-cucumber/src/expectation.rs b/tests/e2e-cucumber/src/expectation.rs index 1c850f591..76d3fbf9f 100644 --- a/tests/e2e-cucumber/src/expectation.rs +++ b/tests/e2e-cucumber/src/expectation.rs @@ -884,13 +884,13 @@ serve_timeout_secs = 90 ]); // MI300X has eight devices → the premise holds → the scenario runs. assert_eq!( - resolve(&d, &cap("mi300x"), &m, false, false, false), + resolve(&d, &cap("mi300x"), &m, Included::default()), Expectation::ExpectPass ); // Strix Halo has exactly one. `@requires-gpu` is satisfied there, which // is why the scenario used to run and fail on its premise. assert!(matches!( - resolve(&d, &cap("strix-ubuntu"), &m, false, false, false), + resolve(&d, &cap("strix-ubuntu"), &m, Included::default()), Expectation::Skip { .. } )); // No GPU at all, and a host whose count could not be probed (WSL), skip @@ -898,7 +898,7 @@ serve_timeout_secs = 90 for host in ["mock", "wsl2"] { assert!( matches!( - resolve(&d, &cap(host), &m, false, false, false), + resolve(&d, &cap(host), &m, Included::default()), Expectation::Skip { .. } ), "{host} must not run a multi-GPU scenario" @@ -944,11 +944,11 @@ serve_timeout_secs = 90 )); assert!(rocr.requires_multi_gpu, "serve-20 must carry the gate"); assert!(matches!( - resolve(&rocr, &cap("strix-ubuntu"), &m, false, false, false), + resolve(&rocr, &cap("strix-ubuntu"), &m, Included::default()), Expectation::Skip { .. } )); assert_eq!( - resolve(&rocr, &cap("mi300x"), &m, false, false, false), + resolve(&rocr, &cap("mi300x"), &m, Included::default()), Expectation::ExpectPass ); @@ -960,7 +960,7 @@ serve_timeout_secs = 90 ); for host in ["strix-ubuntu", "mi300x"] { assert_eq!( - resolve(&masked, &cap(host), &m, false, false, false), + resolve(&masked, &cap(host), &m, Included::default()), Expectation::ExpectPass, "{host} must still run serve-19" ); @@ -975,7 +975,7 @@ serve_timeout_secs = 90 let m = Expectations::default(); let d = decl(&["id:x", "requires-gpu", "requires-multi-gpu"]); let Expectation::Skip { reason } = - resolve(&d, &cap("strix-ubuntu"), &m, false, false, false) + resolve(&d, &cap("strix-ubuntu"), &m, Included::default()) else { panic!("a single-GPU host must skip a multi-GPU scenario"); }; From 6f470307931233259c43f96d2e88d29753dfc190 Mon Sep 17 00:00:00 2001 From: Eugene Volen Date: Mon, 14 Sep 2026 08:57:46 +0000 Subject: [PATCH 2/5] test(remote): cover the remote control channel against a real SSH server MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The unit tests drive a scripted stand-in, which proves the control flow but assumes the real tools behave a certain way. These check that assumption against a real OpenSSH server in a container — no GPU, no ROCm, no tailnet, since the remote's `rocm` and `tailscale` are stand-ins. - `tests/remote-ssh/run.sh` checks the tool contracts: argument handling, exit-code propagation, a credential delivered on stdin and absent from the command line, file copy, batch-mode refusal, the shape Tailscale Funnel takes in the serve config, and that withdrawing a published endpoint actually removes it. - `tests/remote-ssh/run-e2e.sh` drives the built binary through the whole flow: discover, probe, serve, publish, reconcile status, re-publish after an out-of-band withdrawal, tear down, and refuse to publish over a Funnel-exposed port. - 14 cucumber scenarios in `features/remote.feature`; the six needing a host on the other end of a real SSH connection carry `@requires-docker` and skip with a reason where no container runtime exists. - Both scripts run on a new `remote control channel (containerised)` CI lane, gated on the `heavy` path filter. The Funnel fixtures use 443, not the default tailnet port: Funnel serves only 443, 8443 and 10000, so an AllowFunnel entry on any other port is a document the daemon cannot produce and a test against it proves nothing. `resolve` takes an `Included` struct rather than a row of same-typed bools, so a mis-ordered argument cannot silently change which set runs. Signed-off-by: Eugene Volen --- .github/workflows/ci.yml | 60 ++ tests/e2e-cucumber/Cargo.toml | 6 + tests/e2e-cucumber/features/remote.feature | 111 +++ tests/e2e-cucumber/src/bin/fake-tailscale.rs | 44 ++ tests/e2e-cucumber/src/expectation.rs | 549 +++++++++++++-- tests/e2e-cucumber/tests/e2e.rs | 25 +- tests/e2e-cucumber/tests/e2e/remote_steps.rs | 702 +++++++++++++++++++ tests/e2e-cucumber/tests/feature_naming.rs | 1 + tests/remote-ssh/Dockerfile | 66 ++ tests/remote-ssh/examination.json | 53 ++ tests/remote-ssh/fake-rocm.sh | 91 +++ tests/remote-ssh/fake-tailscale.sh | 100 +++ tests/remote-ssh/local-tailscale.sh | 24 + tests/remote-ssh/run-e2e.sh | 250 +++++++ tests/remote-ssh/run.sh | 297 ++++++++ 15 files changed, 2320 insertions(+), 59 deletions(-) create mode 100644 tests/e2e-cucumber/features/remote.feature create mode 100644 tests/e2e-cucumber/src/bin/fake-tailscale.rs create mode 100644 tests/e2e-cucumber/tests/e2e/remote_steps.rs create mode 100644 tests/remote-ssh/Dockerfile create mode 100644 tests/remote-ssh/examination.json create mode 100755 tests/remote-ssh/fake-rocm.sh create mode 100755 tests/remote-ssh/fake-tailscale.sh create mode 100755 tests/remote-ssh/local-tailscale.sh create mode 100755 tests/remote-ssh/run-e2e.sh create mode 100755 tests/remote-ssh/run.sh diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 2a1489b78..a80f861ad 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -133,6 +133,9 @@ jobs: - '**/*.ps1' - '**/*.feature' - 'tests/e2e-cucumber/**' + # The containerised control-channel lane; its Dockerfile matches + # none of the extension globs above. + - 'tests/remote-ssh/**' - 'install*' # Pinned-key consistency compares docs/keys/* against the installers, # so a canonical-key-only change must trigger the heavy job that runs @@ -858,6 +861,55 @@ jobs: prek-version: 0.4.5 extra-args: --all-files --no-group local-tools + remote-ssh: + # `rocm remote` drives ssh, scp, and the remote's service/serve tooling. Its + # unit tests use a scripted stand-in, which proves the control flow but + # assumes the real tools behave a certain way. This job checks that + # assumption against a real OpenSSH server in a container: argument + # handling, exit-code propagation, stdin delivery of a credential, file + # copy, batch-mode refusal, the shape Tailscale Funnel takes in the serve + # config (including that the fake refuses the ports the real daemon + # refuses), and — the one that matters most — that withdrawing a published + # endpoint actually removes it. A publish is configuration rather than a + # process, so it survives reboots, and a withdrawal that silently does + # nothing leaves a GPU endpoint exposed with nothing tracking it. + # + # No GPU, no ROCm, no tailnet: the remote's `rocm` and `tailscale` are + # stand-ins that answer in the shapes the real tools do. + name: remote control channel (containerised) + runs-on: ubuntu-latest + timeout-minutes: 30 + needs: changes + if: needs.changes.outputs.heavy == 'true' + steps: + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + + # Docker and the OpenSSH client are both present on ubuntu-latest; jq is + # the only extra, and only for asserting on JSON contracts. + - name: Install jq + run: sudo apt-get update && sudo apt-get install -y jq + + - name: Check the control-channel contracts + run: tests/remote-ssh/run.sh + + # The pinned toolchain, and the rust-cache that comes with it, as every + # other cargo job in this workflow uses. Without it this lane builds the + # whole workspace cold against a different compiler from the rest of CI, + # inside a 30-minute timeout. + - uses: actions-rust-lang/setup-rust-toolchain@166cdcfd11aee3cb47222f9ddb555ce30ddb9659 # v1.17.0 + + # The same container, driven by the real binary this time: discover a + # target, probe it, serve, publish, reconcile status, re-publish after an + # out-of-band withdrawal, tear down, and refuse to publish over a + # Funnel-exposed port. Catches the orchestration mistakes a scripted + # stand-in cannot — an argument the CLI builds wrongly, a session record + # it fails to clean up. + - name: Build the CLI for the end-to-end run + run: cargo build -p rocm + + - name: Drive rocm remote end to end + run: tests/remote-ssh/run-e2e.sh + clippy: runs-on: ubuntu-latest timeout-minutes: 20 @@ -1271,6 +1323,14 @@ jobs: - name: Run E2E tests if: steps.gate.outputs.run == 'true' + # E2E_INCLUDE_DOCKER turns on the `rocm remote` scenarios that need a + # second machine, stood up as a container. This is the only lane that + # opts in: GitHub-hosted runners can build the fixture image, while the + # self-hosted GPU runners have a daemon but no route to the package + # mirror it installs from, so they skip rather than fail on an image + # they could never build. + env: + E2E_INCLUDE_DOCKER: "1" run: cargo xtask e2e - name: Upload E2E report diff --git a/tests/e2e-cucumber/Cargo.toml b/tests/e2e-cucumber/Cargo.toml index e7c3e8a8c..88e74611d 100644 --- a/tests/e2e-cucumber/Cargo.toml +++ b/tests/e2e-cucumber/Cargo.toml @@ -16,6 +16,12 @@ workspace = true name = "rocm-demo-env" path = "src/bin/rocm-demo-env.rs" +# Stand-in for the Tailscale CLI, put on PATH by the `rocm remote` scenarios so +# they do not depend on the developer's own machine being on a tailnet. +[[bin]] +name = "fake-tailscale" +path = "src/bin/fake-tailscale.rs" + [dependencies] axum.workspace = true cucumber = { version = "0.23", features = ["output-json", "output-junit"] } diff --git a/tests/e2e-cucumber/features/remote.feature b/tests/e2e-cucumber/features/remote.feature new file mode 100644 index 000000000..c502f80fb --- /dev/null +++ b/tests/e2e-cucumber/features/remote.feature @@ -0,0 +1,111 @@ +Feature: Working with GPU machines over a private network + + # These cover the parts of `rocm remote` that need neither a real tailnet nor a + # second machine: discovery, refusals, and the local session list. Serving, + # publishing and teardown need a reachable remote and are covered end to end by + # tests/remote-ssh/run-e2e.sh, which drives this same binary against a + # containerised stand-in. + + @id:remote-targets-lists-machines + Scenario: remote-01 - The user sees which machines could host a model + Given a private network with a GPU machine and a phone + When the user asks which remote targets exist + Then both machines are listed + And the listing says it is not a readiness check + + @id:remote-targets-tag-filter + Scenario: remote-02 - The list can be narrowed to machines marked as GPU nodes + Given a private network with a GPU machine and a phone + When the user asks for remote targets tagged as GPU machines + Then only the GPU machine is listed + + @id:remote-targets-offline-shown + Scenario: remote-03 - A machine that is offline is still listed, marked offline + Given a private network whose GPU machine is offline + When the user asks which remote targets exist + Then the GPU machine is listed as offline + + @id:remote-targets-not-connected + Scenario: remote-04 - Discovery explains itself when the private network is not connected + Given the private network client is installed but not connected + When the user asks which remote targets exist + Then the user is told it is not connected and how to connect + And the command still succeeds + + @id:remote-serve-unknown-machine + Scenario: remote-05 - Serving refuses a machine that is not on the private network + Given a private network with a GPU machine and a phone + When the user asks to serve a model on a machine that is not there + Then the user is told it is not on the network + And they are pointed at the list of machines that are + + @id:remote-serve-offline-machine + Scenario: remote-06 - Serving refuses an offline machine instead of waiting for it + Given a private network whose GPU machine is offline + When the user asks to serve a model on the GPU machine + Then the user is told the machine is offline + + @id:remote-status-no-sessions + Scenario: remote-07 - The user is told when they have no remote sessions + When the user asks about their remote sessions + Then the user is told there are none and how to start one + + @id:remote-stop-unknown-session + Scenario: remote-08 - Stopping a session that does not exist is refused + When the user asks to stop a remote session that does not exist + Then the user is told no such session is recorded + + @id:remote-doctor-refuses-to-install @requires-docker + Scenario: remote-09 - Checking a machine's health never installs anything on it + Given a reachable GPU machine with no ROCm CLI on the private network + When the user checks the health of a machine that has no ROCm CLI + Then the user is told nothing was installed and how to install it + And they are pointed at serving as the command that provisions + + # The successful paths need a second machine. `rocm remote` drives a real SSH + # connection, so no amount of local stubbing produces one — these stand a + # container up instead. The GPU and the private network are still stand-ins; + # what is real is the connection, the commands the CLI builds, and the records + # it keeps. + # + # Opt-in via E2E_INCLUDE_DOCKER=1, set on the GitHub-hosted `E2E tests` lane. + # A working daemon is not enough on its own: the self-hosted GPU runners have + # one but cannot reach the package mirror the fixture image builds from, so + # they skip these rather than failing on an image they could never build. + + @id:remote-serve-publishes-endpoint @requires-docker + Scenario: remote-10 - Serving a model on another machine gives back a usable endpoint + Given a reachable GPU machine on the private network + When the user serves a model on that machine + Then the user is given an endpoint and a credential + And the user is told the endpoint is reachable by the whole network + And the machine is publishing that endpoint + + @id:remote-status-reports-both-halves @requires-docker + Scenario: remote-11 - Status reports the model and the endpoint separately + Given a model serving on a reachable GPU machine + When the user asks about their remote sessions + Then the model and the endpoint are both reported healthy + + @id:remote-attach-restores-endpoint @requires-docker + Scenario: remote-12 - An endpoint withdrawn behind the user's back is repairable + Given a model serving on a reachable GPU machine + When the endpoint is withdrawn on the machine itself + And the user asks about their remote sessions + Then the model is still healthy but the endpoint is reported gone + When the user re-publishes the endpoint + Then the endpoint is restored without restarting the model + + @id:remote-stop-clears-both @requires-docker + Scenario: remote-13 - Stopping a session leaves nothing running and nothing exposed + Given a model serving on a reachable GPU machine + When the user stops the session + Then the endpoint and the model are both reported stopped + And the machine is publishing nothing + And the session is no longer listed + + @id:remote-doctor-reports-remote-health @requires-docker + Scenario: remote-14 - Checking a remote machine reports that machine's health + Given a reachable GPU machine on the private network + When the user checks that machine's health + Then the report names that machine diff --git a/tests/e2e-cucumber/src/bin/fake-tailscale.rs b/tests/e2e-cucumber/src/bin/fake-tailscale.rs new file mode 100644 index 000000000..1f07be0f0 --- /dev/null +++ b/tests/e2e-cucumber/src/bin/fake-tailscale.rs @@ -0,0 +1,44 @@ +// Copyright © Advanced Micro Devices, Inc., or its affiliates. +// +// SPDX-License-Identifier: MIT + +//! A stand-in for the Tailscale CLI, for scenarios that need a tailnet. +//! +//! `rocm remote` shells out to `tailscale status --json` to find candidate +//! machines and to check one is online before dialling it. A scenario that has +//! to be deterministic cannot depend on whether the developer's machine happens +//! to have Tailscale installed, connected, or peered with anything — so it puts +//! this on `PATH` instead and points it at a status document it wrote. +//! +//! A Rust binary rather than a shell script so the scenarios run on Windows too. +//! +//! `FAKE_TAILSCALE_STATUS` names the document to serve. Unset, it reports a +//! daemon that is installed but not logged in, which is its own scenario. + +use std::io::Write as _; + +fn main() { + let args = std::env::args().skip(1).collect::>(); + let mut stdout = std::io::stdout(); + + match args.first().map(String::as_str) { + Some("status") => { + let document = std::env::var("FAKE_TAILSCALE_STATUS") + .ok() + .and_then(|path| std::fs::read_to_string(path).ok()) + .unwrap_or_else(|| r#"{"BackendState":"NeedsLogin"}"#.to_owned()); + let _ = writeln!(stdout, "{document}"); + } + other => { + // Anything else is a command a scenario did not set up. Failing + // loudly beats a silent empty answer that the CLI would read as a + // legitimate "nothing here". + let _ = writeln!( + std::io::stderr(), + "fake tailscale: unsupported command: {}", + other.unwrap_or("(none)") + ); + std::process::exit(2); + } + } +} diff --git a/tests/e2e-cucumber/src/expectation.rs b/tests/e2e-cucumber/src/expectation.rs index 76d3fbf9f..bf382a487 100644 --- a/tests/e2e-cucumber/src/expectation.rs +++ b/tests/e2e-cucumber/src/expectation.rs @@ -24,6 +24,7 @@ use crate::capability::HostCapability; const ID_PREFIX: &str = "id:"; const REQUIRES_ENGINE_PREFIX: &str = "requires-engine:"; const REQUIRES_OS_PREFIX: &str = "requires-os:"; +const REQUIRES_DOCKER_TAG: &str = "requires-docker"; const REQUIRES_GPU_TAG: &str = "requires-gpu"; const REQUIRES_MULTI_GPU_TAG: &str = "requires-multi-gpu"; const REQUIRES_GFX_TARGET_TAG: &str = "requires-gfx-target"; @@ -52,6 +53,24 @@ pub enum Expectation { Skip { reason: String }, } +/// Whether a container runtime is usable here. +/// +/// Probed rather than assumed: a developer machine may have the client without +/// a running daemon. +/// +/// Necessary but not sufficient — see the `requires_docker` gate, which also +/// wants an explicit opt-in. A daemon answering does not mean the fixture image +/// can be built: a self-hosted runner behind a restricted network has both a +/// working daemon and no route to the package mirror the image installs from. +fn docker_available() -> bool { + std::process::Command::new("docker") + .arg("info") + .stdout(std::process::Stdio::null()) + .stderr(std::process::Stdio::null()) + .status() + .is_ok_and(|status| status.success()) +} + /// Facts extracted from a scenario's tags. #[derive(Debug, Clone)] pub struct ScenarioDecl { @@ -112,6 +131,12 @@ pub struct ScenarioDecl { /// xfail `serve_timeout_secs` in expectations.toml (which shortens a known-bug /// serve to fail fast); this lengthens a genuinely-slow expected-pass serve. pub serve_timeout_secs: Option, + /// `@requires-docker`: the scenario's premise is a second machine, stood up + /// as a container. `rocm remote` drives a real SSH connection to a real + /// host, and no amount of local stubbing produces one — so the successful + /// serve/status/attach/stop paths can only be exercised where a container + /// runtime exists. Skipped elsewhere rather than silently uncovered. + pub requires_docker: bool, /// `@nightly`: an expensive scenario (e.g. a large-model serve) that is skipped /// on ordinary per-PR / on-demand runs to keep them fast, and only runs when /// the nightly workflow opts in via `E2E_INCLUDE_NIGHTLY`. @@ -134,6 +159,7 @@ impl ScenarioDecl { /// leading `@` shape cucumber-rs supplies at runtime and bare unit fixtures. pub fn from_tags>(tags: &[S]) -> Self { let mut id = None; + let mut requires_docker = false; let mut requires_gpu = false; let mut requires_multi_gpu = false; let mut requires_gfx_target = false; @@ -159,6 +185,8 @@ impl ScenarioDecl { requires_os = Some(rest.to_ascii_lowercase()); } else if let Some(rest) = tag.strip_prefix(SERVE_TIMEOUT_PREFIX) { serve_timeout_secs = rest.parse::().ok(); + } else if tag == REQUIRES_DOCKER_TAG { + requires_docker = true; } else if tag == REQUIRES_GPU_TAG { requires_gpu = true; } else if tag == REQUIRES_MULTI_GPU_TAG { @@ -190,6 +218,7 @@ impl ScenarioDecl { requires_engine, requires_os, serve_timeout_secs, + requires_docker, nightly, lifecycle, merge_queue, @@ -405,33 +434,64 @@ pub struct PlatformManifest<'a> { pub expectations: Vec, } +/// Which opt-in scenario sets a run includes. +/// +/// A struct rather than a row of bools: they are all the same type, so a +/// mis-ordered argument silently changes which set runs and the compiler cannot +/// help. Naming them at the call site is what makes that visible. +/// +/// Every field defaults to `false`, because each one guards a set that is +/// expensive, slow, or needs something the host may not have. A run opts in; +/// nothing opts in on its behalf. +/// +/// - `nightly` — set by the nightly workflow (via `E2E_INCLUDE_NIGHTLY`). +/// Ordinary per-PR and on-demand runs leave it false so expensive `@nightly` +/// scenarios stay out of the fast path. +/// - `lifecycle` — set (via `E2E_INCLUDE_LIFECYCLE`) only when the caller opts +/// into the expensive, OS-mutating release-lifecycle scenarios. The default +/// fast suite keeps them out. +/// - `docker` — set (via `E2E_INCLUDE_DOCKER`) by the GitHub-hosted lane, for +/// `@requires-docker` scenarios that stand a second machine up in a container. +/// Necessary but not sufficient: `docker_available` must also answer, since +/// a self-hosted GPU runner has a daemon but no route to the package mirror +/// the fixture image builds from. +/// - `merge_queue` — set only in the merge queue (via `E2E_MERGE_QUEUE`). Per-PR +/// runs leave it false so heavy `@merge-queue` serves stay off the PR path (a +/// cheaper per-engine canary covers them) and run once before the change lands. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct Included { + pub nightly: bool, + pub lifecycle: bool, + pub docker: bool, + pub merge_queue: bool, +} + /// Resolve a scenario's expectation on this host. /// /// 1. Not-applicable → `Skip`: a `@nightly` scenario when nightly isn't included, -/// a `@merge-queue` scenario outside the merge queue, a `@requires-gpu` -/// scenario on a host with no AMD GPU, a `@requires-multi-gpu` scenario on a -/// host that does not have more than one, a `@requires-bare-metal` scenario on -/// WSL2, a `@requires-os:` scenario on a different OS, or a scenario whose +/// a `@merge-queue` scenario outside the merge queue, a `@requires-docker` +/// scenario without a usable container runtime, a `@requires-gpu` scenario on +/// a host with no AMD GPU, a `@requires-multi-gpu` scenario on a host that +/// does not have more than one, a `@requires-bare-metal` scenario on WSL2, a +/// `@requires-os:` scenario on a different OS, or a scenario whose /// effective engine can't start. /// 2. First matching `expectations.toml` condition → `ExpectXfail`. /// 3. Otherwise → `ExpectPass`. /// -/// `include_nightly` is set by the nightly workflow (via `E2E_INCLUDE_NIGHTLY`); -/// ordinary per-PR / on-demand runs pass `false` so expensive `@nightly` -/// scenarios stay out of the fast path. `include_lifecycle` is set (via -/// `E2E_INCLUDE_LIFECYCLE`) only when the caller opts into the expensive, -/// OS-mutating release-lifecycle scenarios; the default fast suite keeps them out. -/// `include_merge_queue` is set only in the merge queue (via `E2E_MERGE_QUEUE`); -/// per-PR runs pass `false` so heavy `@merge-queue` serves stay off the PR path (a -/// cheaper per-engine canary covers them) and run once before the change lands. +/// Which opt-in sets are in play is carried by [`Included`], which documents +/// what each one gates and who sets it. pub fn resolve( decl: &ScenarioDecl, cap: &HostCapability, matrix: &Expectations, - include_nightly: bool, - include_lifecycle: bool, - include_merge_queue: bool, + included: Included, ) -> Expectation { + let Included { + nightly: include_nightly, + lifecycle: include_lifecycle, + docker: include_docker, + merge_queue: include_merge_queue, + } = included; // (1) Applicability / skip. if decl.nightly && !include_nightly { return Expectation::Skip { @@ -448,6 +508,13 @@ pub fn resolve( reason: "merge-queue-only scenario; set E2E_MERGE_QUEUE to run".to_owned(), }; } + if decl.requires_docker && !(include_docker && docker_available()) { + return Expectation::Skip { + reason: "needs a second machine in a container; set E2E_INCLUDE_DOCKER=1 on a \ + runner that can build the fixture image" + .to_owned(), + }; + } if decl.requires_gpu && !cap.has_amd_gpu { return Expectation::Skip { reason: "requires an AMD GPU; none detected on this host".to_owned(), @@ -710,14 +777,12 @@ serve_timeout_secs = 90 &scenario, &cap("wsl-no-passthrough"), &matrix, - false, - false, - false, + Included::default(), ), Expectation::ExpectPass ); assert!(matches!( - resolve(&scenario, &cap("mock"), &matrix, false, false, false,), + resolve(&scenario, &cap("mock"), &matrix, Included::default()), Expectation::Skip { .. } )); } @@ -739,17 +804,47 @@ serve_timeout_secs = 90 let d = decl(&["id:big", "requires-gpu", "nightly"]); assert!(d.nightly); assert!(matches!( - resolve(&d, &cap("mi300x"), &m, false, false, false), + resolve( + &d, + &cap("mi300x"), + &m, + Included { + nightly: false, + lifecycle: false, + docker: false, + merge_queue: false + } + ), Expectation::Skip { .. } )); assert_eq!( - resolve(&d, &cap("mi300x"), &m, true, false, false), + resolve( + &d, + &cap("mi300x"), + &m, + Included { + nightly: true, + lifecycle: false, + docker: false, + merge_queue: false + } + ), Expectation::ExpectPass ); // The nightly gate is cheapest-first: a @nightly scenario that ALSO can't // run here (no GPU) still skips regardless of the include flag. assert!(matches!( - resolve(&d, &cap("mock"), &m, true, false, false), + resolve( + &d, + &cap("mock"), + &m, + Included { + nightly: true, + lifecycle: false, + docker: false, + merge_queue: false + } + ), Expectation::Skip { .. } )); } @@ -766,16 +861,78 @@ serve_timeout_secs = 90 ]); assert!(d.lifecycle); assert!(matches!( - resolve(&d, &cap("strix-ubuntu"), &m, false, false, false), + resolve( + &d, + &cap("strix-ubuntu"), + &m, + Included { + nightly: false, + lifecycle: false, + docker: false, + merge_queue: false + } + ), Expectation::Skip { .. } )); assert_eq!( - resolve(&d, &cap("strix-ubuntu"), &m, false, true, false), + resolve( + &d, + &cap("strix-ubuntu"), + &m, + Included { + nightly: false, + lifecycle: true, + docker: false, + merge_queue: false + } + ), Expectation::ExpectPass ); // Even when included, an inapplicable OS still skips (os gate is checked). assert!(matches!( - resolve(&d, &cap("strix-windows"), &m, false, true, false), + resolve( + &d, + &cap("strix-windows"), + &m, + Included { + nightly: false, + lifecycle: true, + docker: false, + merge_queue: false + } + ), + Expectation::Skip { .. } + )); + } + + #[test] + fn docker_scenario_skips_unless_included() { + // A container-backed scenario is opt-in, not merely "docker is here". + // A runner can have a working daemon and still be unable to build the + // fixture image — the self-hosted GPU boxes have exactly that shape, and + // failing there on an image they could never build told us nothing. + let d = ScenarioDecl { + requires_docker: true, + ..decl(&["@id:x"]) + }; + let m = Expectations::default(); + assert!(matches!( + resolve(&d, &cap("mi300x"), &m, Included::default()), + Expectation::Skip { .. } + )); + // And it is not opted in by any of the other sets. + assert!(matches!( + resolve( + &d, + &cap("mi300x"), + &m, + Included { + nightly: true, + lifecycle: true, + merge_queue: true, + docker: false, + } + ), Expectation::Skip { .. } )); } @@ -792,22 +949,62 @@ serve_timeout_secs = 90 ]); assert!(d.merge_queue); assert!(matches!( - resolve(&d, &cap("mi300x"), &m, false, false, false), + resolve( + &d, + &cap("mi300x"), + &m, + Included { + nightly: false, + lifecycle: false, + docker: false, + merge_queue: false + } + ), Expectation::Skip { .. } )); assert_eq!( - resolve(&d, &cap("mi300x"), &m, false, false, true), + resolve( + &d, + &cap("mi300x"), + &m, + Included { + nightly: false, + lifecycle: false, + docker: false, + merge_queue: true + } + ), Expectation::ExpectPass ); // Independent of the nightly axis: a merge-queue scenario is not opted in // by E2E_INCLUDE_NIGHTLY. assert!(matches!( - resolve(&d, &cap("mi300x"), &m, true, false, false), + resolve( + &d, + &cap("mi300x"), + &m, + Included { + nightly: true, + lifecycle: false, + docker: false, + merge_queue: false + } + ), Expectation::Skip { .. } )); // Cheapest-first: still skips where it can't run at all (no GPU). assert!(matches!( - resolve(&d, &cap("mock"), &m, false, false, true), + resolve( + &d, + &cap("mock"), + &m, + Included { + nightly: false, + lifecycle: false, + docker: false, + merge_queue: true + } + ), Expectation::Skip { .. } )); } @@ -832,17 +1029,47 @@ serve_timeout_secs = 90 // MI300X: default engine vLLM → xfail. assert!(matches!( - resolve(&d, &cap("mi300x"), &m, false, false, false), + resolve( + &d, + &cap("mi300x"), + &m, + Included { + nightly: false, + lifecycle: false, + docker: false, + merge_queue: false + } + ), Expectation::ExpectXfail { .. } )); // Strix Ubuntu: gfx1151 → lemonade default → NOT vLLM → expect-pass. assert_eq!( - resolve(&d, &cap("strix-ubuntu"), &m, false, false, false), + resolve( + &d, + &cap("strix-ubuntu"), + &m, + Included { + nightly: false, + lifecycle: false, + docker: false, + merge_queue: false + } + ), Expectation::ExpectPass ); // Strix Windows: lemonade default → expect-pass (this is the XPASS fix). assert_eq!( - resolve(&d, &cap("strix-windows"), &m, false, false, false), + resolve( + &d, + &cap("strix-windows"), + &m, + Included { + nightly: false, + lifecycle: false, + docker: false, + merge_queue: false + } + ), Expectation::ExpectPass ); } @@ -852,7 +1079,17 @@ serve_timeout_secs = 90 let m = eai7333_matrix(); let d = decl(&["id:serve-default-engine-inference", "requires-gpu"]); assert!(matches!( - resolve(&d, &cap("mock"), &m, false, false, false), + resolve( + &d, + &cap("mock"), + &m, + Included { + nightly: false, + lifecycle: false, + docker: false, + merge_queue: false + } + ), Expectation::Skip { .. } )); } @@ -991,16 +1228,46 @@ serve_timeout_secs = 90 let d = decl(&["id:serve-no-gpu-fails-fast", "requires-no-gpu"]); // The mock host has no AMD GPU → the no-GPU premise applies → runs. assert_eq!( - resolve(&d, &cap("mock"), &m, false, false, false), + resolve( + &d, + &cap("mock"), + &m, + Included { + nightly: false, + lifecycle: false, + docker: false, + merge_queue: false + } + ), Expectation::ExpectPass ); // Every GPU host skips it — the premise can't hold there. assert!(matches!( - resolve(&d, &cap("mi300x"), &m, false, false, false), + resolve( + &d, + &cap("mi300x"), + &m, + Included { + nightly: false, + lifecycle: false, + docker: false, + merge_queue: false + } + ), Expectation::Skip { .. } )); assert!(matches!( - resolve(&d, &cap("strix-ubuntu"), &m, false, false, false), + resolve( + &d, + &cap("strix-ubuntu"), + &m, + Included { + nightly: false, + lifecycle: false, + docker: false, + merge_queue: false + } + ), Expectation::Skip { .. } )); } @@ -1011,7 +1278,17 @@ serve_timeout_secs = 90 let d = decl(&["id:diagnose-matches-known-symptom", "requires-bare-metal"]); for wsl in ["wsl2", "wsl"] { assert!(matches!( - resolve(&d, &cap(wsl), &m, false, false, false), + resolve( + &d, + &cap(wsl), + &m, + Included { + nightly: false, + lifecycle: false, + docker: false, + merge_queue: false + } + ), Expectation::Skip { .. } )); } @@ -1019,7 +1296,17 @@ serve_timeout_secs = 90 // which is where these scenarios earn their keep as a required check. for host in ["mock", "mi300x", "strix-ubuntu", "strix-windows"] { assert_eq!( - resolve(&d, &cap(host), &m, false, false, false), + resolve( + &d, + &cap(host), + &m, + Included { + nightly: false, + lifecycle: false, + docker: false, + merge_queue: false + } + ), Expectation::ExpectPass, "{host} is bare metal and must still run the scenario" ); @@ -1043,21 +1330,21 @@ serve_timeout_secs = 90 for wsl in ["wsl", "wsl2"] { assert!( matches!( - resolve(&d, &cap(wsl), &m, false, false, false), + resolve(&d, &cap(wsl), &m, Included::default()), Expectation::Skip { .. } ), "{wsl} is linux but not bare metal, so the scenario has no premise" ); } assert!(matches!( - resolve(&d, &cap("strix-windows"), &m, false, false, false), + resolve(&d, &cap("strix-windows"), &m, Included::default()), Expectation::Skip { .. } )); // `mock` is deliberately absent: the fixture models it as os_family // "other", so it cannot stand for the real mock lane here. for host in ["mi300x", "strix-ubuntu"] { assert_eq!( - resolve(&d, &cap(host), &m, false, false, false), + resolve(&d, &cap(host), &m, Included::default()), Expectation::ExpectPass, "{host} is native Linux and must still run the scenario" ); @@ -1073,13 +1360,33 @@ serve_timeout_secs = 90 assert!(d.requires_wsl); for wsl in ["wsl", "wsl2"] { assert_eq!( - resolve(&d, &cap(wsl), &m, false, false, false), + resolve( + &d, + &cap(wsl), + &m, + Included { + nightly: false, + lifecycle: false, + docker: false, + merge_queue: false + } + ), Expectation::ExpectPass ); } for platform in ["mock", "mi300x", "strix-ubuntu", "strix-windows"] { assert!(matches!( - resolve(&d, &cap(platform), &m, false, false, false), + resolve( + &d, + &cap(platform), + &m, + Included { + nightly: false, + lifecycle: false, + docker: false, + merge_queue: false + } + ), Expectation::Skip { .. } )); } @@ -1093,7 +1400,17 @@ serve_timeout_secs = 90 let m = Expectations::default(); let d = decl(&["id:some-linux-scenario", "requires-os:linux"]); assert_eq!( - resolve(&d, &cap("wsl2"), &m, false, false, false), + resolve( + &d, + &cap("wsl2"), + &m, + Included { + nightly: false, + lifecycle: false, + docker: false, + merge_queue: false + } + ), Expectation::ExpectPass ); } @@ -1114,7 +1431,17 @@ reason = "unrelated open bug" .unwrap(); let d = decl(&["id:diagnose-matches-known-symptom", "requires-bare-metal"]); assert!(matches!( - resolve(&d, &cap("wsl2"), &m, false, false, false), + resolve( + &d, + &cap("wsl2"), + &m, + Included { + nightly: false, + lifecycle: false, + docker: false, + merge_queue: false + } + ), Expectation::Skip { .. } )); } @@ -1130,12 +1457,32 @@ reason = "unrelated open bug" ]); // MI300X: vLLM available → not skipped (expect-pass here, no matrix entry). assert_eq!( - resolve(&d, &cap("mi300x"), &m, false, false, false), + resolve( + &d, + &cap("mi300x"), + &m, + Included { + nightly: false, + lifecycle: false, + docker: false, + merge_queue: false + } + ), Expectation::ExpectPass ); // Strix Windows: vLLM can't start → skip (N/A). assert!(matches!( - resolve(&d, &cap("strix-windows"), &m, false, false, false), + resolve( + &d, + &cap("strix-windows"), + &m, + Included { + nightly: false, + lifecycle: false, + docker: false, + merge_queue: false + } + ), Expectation::Skip { .. } )); } @@ -1149,15 +1496,45 @@ reason = "unrelated open bug" // Runs on a Linux GPU host; skips where os_family != linux (windows, and // the "other" fixture host). assert_eq!( - resolve(&d, &cap("strix-ubuntu"), &m, false, false, false), + resolve( + &d, + &cap("strix-ubuntu"), + &m, + Included { + nightly: false, + lifecycle: false, + docker: false, + merge_queue: false + } + ), Expectation::ExpectPass ); assert!(matches!( - resolve(&d, &cap("strix-windows"), &m, false, false, false), + resolve( + &d, + &cap("strix-windows"), + &m, + Included { + nightly: false, + lifecycle: false, + docker: false, + merge_queue: false + } + ), Expectation::Skip { .. } )); assert!(matches!( - resolve(&d, &cap("mock"), &m, false, false, false), + resolve( + &d, + &cap("mock"), + &m, + Included { + nightly: false, + lifecycle: false, + docker: false, + merge_queue: false + } + ), Expectation::Skip { .. } )); } @@ -1167,11 +1544,31 @@ reason = "unrelated open bug" let m = Expectations::default(); let d = decl(&["id:examine-version"]); assert_eq!( - resolve(&d, &cap("mock"), &m, false, false, false), + resolve( + &d, + &cap("mock"), + &m, + Included { + nightly: false, + lifecycle: false, + docker: false, + merge_queue: false + } + ), Expectation::ExpectPass ); assert_eq!( - resolve(&d, &cap("mi300x"), &m, false, false, false), + resolve( + &d, + &cap("mi300x"), + &m, + Included { + nightly: false, + lifecycle: false, + docker: false, + merge_queue: false + } + ), Expectation::ExpectPass ); } @@ -1190,11 +1587,31 @@ reason = "short-name not surfaced" let d = decl(&["id:serve-short-name-expansion"]); // No requires-gpu → runs everywhere, always xfail. assert!(matches!( - resolve(&d, &cap("mock"), &m, false, false, false), + resolve( + &d, + &cap("mock"), + &m, + Included { + nightly: false, + lifecycle: false, + docker: false, + merge_queue: false + } + ), Expectation::ExpectXfail { .. } )); assert!(matches!( - resolve(&d, &cap("mi300x"), &m, false, false, false), + resolve( + &d, + &cap("mi300x"), + &m, + Included { + nightly: false, + lifecycle: false, + docker: false, + merge_queue: false + } + ), Expectation::ExpectXfail { .. } )); } @@ -1217,12 +1634,32 @@ reason = "lemonade vulkan fallback" ]); // Strix Ubuntu (linux, lemonade) → xfail. assert!(matches!( - resolve(&d, &cap("strix-ubuntu"), &m, false, false, false), + resolve( + &d, + &cap("strix-ubuntu"), + &m, + Included { + nightly: false, + lifecycle: false, + docker: false, + merge_queue: false + } + ), Expectation::ExpectXfail { .. } )); // Strix Windows (windows, lemonade) → os mismatch → expect-pass. assert_eq!( - resolve(&d, &cap("strix-windows"), &m, false, false, false), + resolve( + &d, + &cap("strix-windows"), + &m, + Included { + nightly: false, + lifecycle: false, + docker: false, + merge_queue: false + } + ), Expectation::ExpectPass ); } diff --git a/tests/e2e-cucumber/tests/e2e.rs b/tests/e2e-cucumber/tests/e2e.rs index 5adeccf3a..815154a62 100644 --- a/tests/e2e-cucumber/tests/e2e.rs +++ b/tests/e2e-cucumber/tests/e2e.rs @@ -30,6 +30,7 @@ mod e2e { pub mod examine_steps; pub mod lifecycle_steps; pub mod logs_steps; + pub mod remote_steps; pub mod runtime_lifecycle_steps; pub mod runtime_steps; pub mod service_cleanup_steps; @@ -56,6 +57,14 @@ pub struct E2eWorld { pub cli_outputs: Option>, pub cli_stderr: Option, pub cli_rc: Option, + /// Extra environment for `rocm remote` scenarios: a `PATH` carrying the + /// tailscale stand-in, and the status document it should serve. Set by a + /// Given step so the When steps stay about what the user does. + pub remote_env: Vec<(String, String)>, + /// Container standing in for a second machine, for the `@requires-docker` + /// scenarios. Held on the World so it lives for the scenario and is torn + /// down when the World drops, even if a step panics. + pub remote_machine: Option, /// Name of the scenario currently executing, set by the `before` hook. Used /// to tie each recorded `rocm` invocation to its scenario so the coverage /// report can join commands to pass/fail results. @@ -218,6 +227,8 @@ impl Default for E2eWorld { cli_outputs: None, cli_stderr: None, cli_rc: None, + remote_env: Vec::new(), + remote_machine: None, current_scenario: None, isolated_root: Some(root), legacy_rocm_path: None, @@ -1122,6 +1133,11 @@ async fn main() { // install/uninstall) are skipped unless the caller opts in via // `E2E_INCLUDE_LIFECYCLE`, so the default `cargo xtask e2e` stays fast. let include_lifecycle = std::env::var_os("E2E_INCLUDE_LIFECYCLE").is_some_and(|v| v == "1"); + // The container-backed remote scenarios need a runner that can *build* the + // fixture image, which a working daemon alone does not guarantee — a + // restricted network gives you one without the other. Opt in explicitly on + // the lanes where it holds. + let include_docker = std::env::var_os("E2E_INCLUDE_DOCKER").is_some_and(|v| v == "1"); // CI runs just the lifecycle set after opting in. Keep this selection inside // our custom filter instead of cucumber's `--tags`/`-n`: cucumber 0.23 uses // either its CLI filter OR this closure, so CLI selection would bypass OS, @@ -1224,9 +1240,12 @@ async fn main() { &decl, cap, matrix, - include_nightly, - include_lifecycle, - include_merge_queue, + e2e_cucumber::expectation::Included { + nightly: include_nightly, + lifecycle: include_lifecycle, + docker: include_docker, + merge_queue: include_merge_queue, + }, ); let run = (!only_lifecycle || decl.lifecycle) && !matches!(expectation, Expectation::Skip { .. }); diff --git a/tests/e2e-cucumber/tests/e2e/remote_steps.rs b/tests/e2e-cucumber/tests/e2e/remote_steps.rs new file mode 100644 index 000000000..096f054dc --- /dev/null +++ b/tests/e2e-cucumber/tests/e2e/remote_steps.rs @@ -0,0 +1,702 @@ +// Copyright © Advanced Micro Devices, Inc., or its affiliates. +// +// SPDX-License-Identifier: MIT + +//! Steps for `rocm remote`. +//! +//! Every scenario here plants a stand-in `tailscale` on the CLI's `PATH` and +//! points it at a status document the step wrote. Without that these would pass +//! or fail depending on whether the machine running them happens to have +//! Tailscale installed and peered — and "no targets" would look the same as +//! "the feature is broken". + +use std::path::PathBuf; + +use cucumber::{given, then, when}; + +use crate::E2eWorld; + +/// A tailnet with one GPU machine and one machine that is not a GPU machine, so +/// tag filtering has something to exclude. +fn tailnet(gpu_online: bool) -> String { + format!( + r#"{{ + "BackendState": "Running", + "TUN": true, + "MagicDNSSuffix": "example-tailnet.ts.net", + "Self": {{ "HostName": "laptop", "DNSName": "laptop.example-tailnet.ts.net.", "Online": true }}, + "Peer": {{ + "nodekey:aaa": {{ + "HostName": "gpu-box", + "DNSName": "gpu-box.example-tailnet.ts.net.", + "OS": "linux", + "TailscaleIPs": ["100.88.14.21"], + "Tags": ["tag:gpu"], + "Online": {gpu_online} + }}, + "nodekey:bbb": {{ + "HostName": "phone", + "DNSName": "phone.example-tailnet.ts.net.", + "OS": "iOS", + "TailscaleIPs": ["100.88.51.6"], + "Online": true + }} + }} + }}"# + ) +} + +/// Put the stand-in on `PATH` and, when given, the status document it serves. +fn plant_tailscale(world: &E2eWorld, status: Option<&str>) -> Vec<(String, String)> { + let root = world + .isolated_root + .as_ref() + .expect("scenario root") + .path() + .to_path_buf(); + let bin_dir = root.join("fake-bin"); + std::fs::create_dir_all(&bin_dir).expect("create fake bin dir"); + + let built = fake_tailscale_binary(); + let installed = bin_dir.join(if cfg!(windows) { + "tailscale.exe" + } else { + "tailscale" + }); + std::fs::copy(&built, &installed).unwrap_or_else(|error| { + panic!( + "failed to install the tailscale stand-in from {}: {error}", + built.display() + ) + }); + + let mut env = vec![( + "PATH".to_owned(), + format!( + "{}{}{}", + bin_dir.display(), + if cfg!(windows) { ";" } else { ":" }, + std::env::var("PATH").unwrap_or_default() + ), + )]; + if let Some(status) = status { + let path = root.join("tailscale-status.json"); + std::fs::write(&path, status).expect("write status document"); + env.push(( + "FAKE_TAILSCALE_STATUS".to_owned(), + path.display().to_string(), + )); + } + env +} + +/// The stand-in is a bin target of this crate, so cargo has already built it +/// next to the test binary. +fn fake_tailscale_binary() -> PathBuf { + let mut dir = std::env::current_exe().expect("test binary path"); + dir.pop(); + if dir.ends_with("deps") { + dir.pop(); + } + let candidate = dir.join(if cfg!(windows) { + "fake-tailscale.exe" + } else { + "fake-tailscale" + }); + assert!( + candidate.exists(), + "the tailscale stand-in was not built at {}; it is a [[bin]] of this crate", + candidate.display() + ); + candidate +} + +fn run_remote(world: &mut E2eWorld, args: &[&str], env: &[(String, String)]) { + let borrowed = env + .iter() + .map(|(key, value)| (key.as_str(), value.as_str())) + .collect::>(); + let (stdout, stderr, rc) = crate::run_rocm_with_stdin(world, args, "", &borrowed); + world.cli_output = Some(stdout); + world.cli_stderr = Some(stderr); + world.cli_rc = Some(rc); +} + +/// Everything the command said, however it said it. A refusal is an error, and +/// a scenario asserting on what the user was told should not care which stream +/// carried it. +fn said(world: &E2eWorld) -> String { + format!( + "{}\n{}", + world.cli_output.clone().unwrap_or_default(), + world.cli_stderr.clone().unwrap_or_default() + ) +} + +#[given("a private network with a GPU machine and a phone")] +async fn given_tailnet(world: &mut E2eWorld) { + world.remote_env = plant_tailscale(world, Some(&tailnet(true))); +} + +#[given("a private network whose GPU machine is offline")] +async fn given_offline_tailnet(world: &mut E2eWorld) { + world.remote_env = plant_tailscale(world, Some(&tailnet(false))); +} + +#[given("the private network client is installed but not connected")] +async fn given_not_connected(world: &mut E2eWorld) { + // No status document: the stand-in then reports a daemon that is installed + // but not logged in, which is the state this scenario is about. + world.remote_env = plant_tailscale(world, None); +} + +#[when("the user asks which remote targets exist")] +async fn when_targets(world: &mut E2eWorld) { + let env = world.remote_env.clone(); + run_remote(world, &["remote", "targets"], &env); +} + +#[when("the user asks for remote targets tagged as GPU machines")] +async fn when_targets_tagged(world: &mut E2eWorld) { + let env = world.remote_env.clone(); + run_remote(world, &["remote", "targets", "--tag", "gpu"], &env); +} + +#[when("the user asks to serve a model on a machine that is not there")] +async fn when_serve_unknown(world: &mut E2eWorld) { + let env = world.remote_env.clone(); + run_remote( + world, + &["remote", "serve", "not-a-machine", "some-model"], + &env, + ); +} + +#[when("the user asks to serve a model on the GPU machine")] +async fn when_serve_gpu_box(world: &mut E2eWorld) { + let env = world.remote_env.clone(); + run_remote(world, &["remote", "serve", "gpu-box", "some-model"], &env); +} + +#[when("the user asks about their remote sessions")] +async fn when_status(world: &mut E2eWorld) { + let env = world.remote_env.clone(); + run_remote(world, &["remote", "status"], &env); +} + +#[when("the user asks to stop a remote session that does not exist")] +async fn when_stop_unknown(world: &mut E2eWorld) { + let env = world.remote_env.clone(); + run_remote(world, &["remote", "stop", "no-such-session"], &env); +} + +#[when("the user checks the health of a machine that has no ROCm CLI")] +async fn when_doctor_without_cli(world: &mut E2eWorld) { + // Whatever `Given` ran set world.remote_env: either no container (the + // machine is unreachable over ssh) or the no-CLI container (reachable, but + // genuinely missing the binary). Both must refuse rather than install. + let env = world.remote_env.clone(); + run_remote(world, &["remote", "doctor", "gpu-box"], &env); +} + +#[then("both machines are listed")] +async fn then_both_listed(world: &mut E2eWorld) { + let said = said(world); + assert!(said.contains("- gpu-box"), "{said}"); + assert!(said.contains("- phone"), "{said}"); +} + +#[then("the listing says it is not a readiness check")] +async fn then_not_readiness(world: &mut E2eWorld) { + let said = said(world); + assert!( + said.contains("does not mean they have a GPU"), + "a list of reachable machines must not read as a list of usable ones:\n{said}" + ); +} + +#[then("only the GPU machine is listed")] +async fn then_only_gpu(world: &mut E2eWorld) { + let said = said(world); + assert!(said.contains("- gpu-box"), "{said}"); + assert!(!said.contains("- phone"), "{said}"); +} + +#[then("the GPU machine is listed as offline")] +async fn then_offline_listed(world: &mut E2eWorld) { + let said = said(world); + assert!(said.contains("- gpu-box"), "{said}"); + assert!(said.contains("online: no"), "{said}"); +} + +#[then("the user is told it is not connected and how to connect")] +async fn then_not_connected(world: &mut E2eWorld) { + let said = said(world); + assert!(said.contains("not connected"), "{said}"); + assert!(said.contains("tailscale up"), "{said}"); +} + +#[then("the command still succeeds")] +async fn then_succeeds(world: &mut E2eWorld) { + assert_eq!( + world.cli_rc, + Some(0), + "asking what exists should answer, not fail:\n{}", + said(world) + ); +} + +#[then("the user is told it is not on the network")] +async fn then_not_on_network(world: &mut E2eWorld) { + let said = said(world); + assert!(said.contains("not a machine on this tailnet"), "{said}"); + assert_ne!(world.cli_rc, Some(0), "{said}"); +} + +#[then("they are pointed at the list of machines that are")] +async fn then_pointed_at_targets(world: &mut E2eWorld) { + let said = said(world); + assert!(said.contains("rocm remote targets"), "{said}"); +} + +#[then("the user is told the machine is offline")] +async fn then_told_offline(world: &mut E2eWorld) { + let said = said(world); + assert!(said.contains("offline"), "{said}"); + assert_ne!(world.cli_rc, Some(0), "{said}"); +} + +#[then("the user is told there are none and how to start one")] +async fn then_no_sessions(world: &mut E2eWorld) { + let said = said(world); + assert!(said.contains("No remote sessions"), "{said}"); + assert!(said.contains("rocm remote serve"), "{said}"); +} + +#[then("the user is told no such session is recorded")] +async fn then_no_such_session(world: &mut E2eWorld) { + let said = said(world); + assert!(said.contains("no remote sessions are recorded"), "{said}"); + assert_ne!(world.cli_rc, Some(0), "{said}"); +} + +#[then("the user is told nothing was installed and how to install it")] +async fn then_doctor_installed_nothing(world: &mut E2eWorld) { + let said = said(world); + assert_ne!( + world.cli_rc, + Some(0), + "a health check with no CLI must fail:\n{said}" + ); + // No "or it never got there" alternative: this scenario's Given starts a + // real container, so the machine is reachable by construction. Accepting + // the unreachable wording too would let the scenario pass having checked + // only that the command failed somehow — and it would keep passing if the + // container stopped coming up. + assert!( + said.contains("only reads"), + "a read-only command must say it changed nothing:\n{said}" + ); +} + +#[then("they are pointed at serving as the command that provisions")] +async fn then_doctor_points_at_serve(world: &mut E2eWorld) { + let said = said(world); + // Asserted unconditionally, for the same reason as above: guarding this on + // the output already containing "only reads" made the step vacuous on every + // other path, so an unexpected one passed silently instead of failing. + assert!( + said.contains("rocm remote serve"), + "the health check must name serving as what provisions the machine:\n{said}" + ); +} + +// ── Scenarios needing a second machine ───────────────────────────── +// +// `rocm remote` opens a real SSH connection, so the successful paths cannot be +// covered by stubbing alone — there has to be a host on the other end. These +// stand one up as a container built from tests/remote-ssh, whose `rocm` and +// `tailscale` are stand-ins. What is real: the connection, the arguments the +// CLI builds, and the session records it writes. + +/// A container standing in for a GPU machine, torn down with the World. +#[derive(Debug)] +pub struct RemoteMachine { + container: String, + /// Env the CLI needs to reach it: an ssh config, the tailnet stand-in, PATH. + pub env: Vec<(String, String)>, +} + +impl Drop for RemoteMachine { + fn drop(&mut self) { + let _ = std::process::Command::new("docker") + .args(["rm", "-f", &self.container]) + .stdout(std::process::Stdio::null()) + .stderr(std::process::Stdio::null()) + .status(); + } +} + +impl RemoteMachine { + /// Run a command on the machine itself, to set up or inspect state the CLI + /// is not responsible for. + /// A failed command is a failed test, not an empty string. The callers of + /// this assert on what the container reports — including + /// `then_publishing_nothing`, which asserts a port is *absent* from the serve + /// config. That assertion passes trivially on empty output, so a broken + /// container or a missing binary would read as "the endpoint was withdrawn" + /// and quietly retire the one check that does not take the CLI's word for it. + pub fn exec(&self, args: &[&str]) -> String { + let output = std::process::Command::new("docker") + .arg("exec") + .arg(&self.container) + .args(args) + .output() + .expect("docker exec"); + assert!( + output.status.success(), + "`docker exec {} {args:?}` failed ({}): {}", + self.container, + output.status, + String::from_utf8_lossy(&output.stderr).trim() + ); + String::from_utf8_lossy(&output.stdout).into_owned() + } +} + +fn repo_root() -> PathBuf { + PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .parent() + .and_then(std::path::Path::parent) + .expect("repo root") + .to_path_buf() +} + +fn docker(args: &[&str]) -> std::process::Output { + std::process::Command::new("docker") + .args(args) + .output() + .expect("docker") +} + +/// A port that was free a moment ago. +/// +/// Only a candidate: the listener is closed before the caller uses the port, so +/// something else can take it in between. [`start_on_a_free_port`] is what makes +/// that survivable. +fn free_port() -> u16 { + std::net::TcpListener::bind("127.0.0.1:0") + .expect("bind") + .local_addr() + .expect("addr") + .port() +} + +/// Start the stand-in machine on a loopback port, retrying if the port was +/// taken between picking it and publishing it. +/// +/// Returns the container name and the port it is reachable on. +fn start_on_a_free_port(tag: &str) -> (String, u16) { + let mut last_error = String::new(); + for _ in 0..10 { + let port = free_port(); + let container = format!("rocm-e2e-remote-{}-{port}", std::process::id()); + let run = docker(&[ + "run", + "-d", + "--name", + &container, + "-p", + &format!("127.0.0.1:{port}:22"), + tag, + ]); + if run.status.success() { + return (container, port); + } + last_error = String::from_utf8_lossy(&run.stderr).into_owned(); + // A failed `run` can still leave the named container behind, which + // would make every retry fail on the name rather than the port. + let _ = docker(&["rm", "-f", &container]); + } + panic!("failed to start the stand-in machine after 10 attempts: {last_error}"); +} + +/// Build the image, start it, and set up everything the CLI needs to reach it. +fn start_remote_machine(world: &E2eWorld) -> RemoteMachine { + start_remote_machine_variant(world, true) +} + +/// Same stand-in, but built without a `rocm` binary at all. Used for the one +/// scenario that has to prove a health check refuses to install onto a +/// reachable machine, rather than merely failing to connect to one. +fn start_remote_machine_without_cli(world: &E2eWorld) -> RemoteMachine { + start_remote_machine_variant(world, false) +} + +fn start_remote_machine_variant(world: &E2eWorld, include_rocm_cli: bool) -> RemoteMachine { + let fixtures = repo_root().join("tests").join("remote-ssh"); + // A distinct tag per variant: the two must not clobber each other's image + // when scenarios run concurrently or interleaved. + let tag = if include_rocm_cli { + "rocm-remote-ssh-test" + } else { + "rocm-remote-ssh-test-no-cli" + }; + let build = docker(&[ + "build", + "-q", + "--build-arg", + &format!( + "APK_REPO_FLAGS={}", + std::env::var("ROCM_TEST_APK_REPOS").unwrap_or_default() + ), + "--build-arg", + &format!("INCLUDE_ROCM_CLI={}", i32::from(include_rocm_cli)), + "-t", + tag, + fixtures.to_str().expect("fixtures path"), + ]); + assert!( + build.status.success(), + "failed to build the stand-in machine: {}", + String::from_utf8_lossy(&build.stderr) + ); + + let root = world + .isolated_root + .as_ref() + .expect("scenario root") + .path() + .to_path_buf(); + // Retried rather than assumed: `free_port` closes its listener before + // docker binds, and this lane runs scenarios 64-way concurrent, so the + // port can be taken in between. Letting docker pick (`-p 127.0.0.1:0:22`) + // would close the gap outright, but `tests/remote-ssh/run.sh` documents + // that some Docker setups do not forward ephemeral publications to the + // host at all — so the explicit port stays and the race is handled by + // trying again with a different one. + let (container, port) = start_on_a_free_port(tag); + let mut machine = RemoteMachine { + container: container.clone(), + env: Vec::new(), + }; + + // A key for this scenario only. + let key = root.join("id"); + let keygen = std::process::Command::new("ssh-keygen") + .args(["-q", "-t", "ed25519", "-N", "", "-f"]) + .arg(&key) + .status() + .expect("ssh-keygen"); + assert!(keygen.success(), "ssh-keygen failed"); + let public = std::fs::read_to_string(key.with_extension("pub")).expect("read public key"); + let install = std::process::Command::new("docker") + .args([ + "exec", + &container, + "sh", + "-c", + &format!( + "printf '%s\\n' '{}' >> /root/.ssh/authorized_keys", + public.trim() + ), + ]) + .status() + .expect("install key"); + assert!(install.success(), "failed to authorize the scenario key"); + + // ssh resolves ~/.ssh/config from the account database rather than from + // HOME, so the CLI has to be told where this scenario's config is. + let ssh_config = root.join("ssh-config"); + std::fs::write( + &ssh_config, + format!( + "Host gpu-box\n HostName 127.0.0.1\n Port {port}\n User root\n \ + IdentityFile {}\n IdentitiesOnly yes\n StrictHostKeyChecking no\n \ + UserKnownHostsFile /dev/null\n LogLevel ERROR\n", + key.display() + ), + ) + .expect("write ssh config"); + + let status = root.join("tailnet.json"); + std::fs::write(&status, tailnet(true)).expect("write tailnet status"); + let mut env = plant_tailscale(world, Some(&tailnet(true))); + env.push(( + "FAKE_TAILSCALE_STATUS".to_owned(), + status.display().to_string(), + )); + env.push(( + "ROCM_REMOTE_SSH_CONFIG".to_owned(), + ssh_config.display().to_string(), + )); + + // Wait for sshd rather than sleeping a fixed amount. + let mut ready = false; + for _ in 0..80 { + let probe = std::process::Command::new("ssh") + .args(["-o", "BatchMode=yes", "-o", "ConnectTimeout=5", "-F"]) + .arg(&ssh_config) + .args(["gpu-box", "true"]) + .status(); + if probe.is_ok_and(|status| status.success()) { + ready = true; + break; + } + std::thread::sleep(std::time::Duration::from_millis(250)); + } + assert!(ready, "the stand-in machine never accepted a connection"); + + machine.env = env; + machine +} + +const fn machine(world: &E2eWorld) -> &RemoteMachine { + world.remote_machine.as_ref().expect("no machine started") +} + +#[given("a reachable GPU machine on the private network")] +async fn given_reachable_machine(world: &mut E2eWorld) { + let started = start_remote_machine(world); + world.remote_env = started.env.clone(); + world.remote_machine = Some(started); +} + +#[given("a reachable GPU machine with no ROCm CLI on the private network")] +async fn given_reachable_machine_without_cli(world: &mut E2eWorld) { + let started = start_remote_machine_without_cli(world); + world.remote_env = started.env.clone(); + world.remote_machine = Some(started); +} + +#[given("a model serving on a reachable GPU machine")] +async fn given_serving_machine(world: &mut E2eWorld) { + given_reachable_machine(world).await; + let env = world.remote_env.clone(); + run_remote(world, &["remote", "serve", "gpu-box", "test-model"], &env); + assert_eq!(world.cli_rc, Some(0), "serve failed:\n{}", said(world)); +} + +#[when("the user serves a model on that machine")] +async fn when_serve_on_machine(world: &mut E2eWorld) { + let env = world.remote_env.clone(); + run_remote(world, &["remote", "serve", "gpu-box", "test-model"], &env); +} + +#[when("the user checks that machine's health")] +async fn when_check_health(world: &mut E2eWorld) { + let env = world.remote_env.clone(); + run_remote(world, &["remote", "doctor", "gpu-box"], &env); +} + +#[when("the endpoint is withdrawn on the machine itself")] +async fn when_withdrawn_remotely(world: &mut E2eWorld) { + machine(world).exec(&["tailscale", "serve", "--tcp=8000", "off"]); +} + +#[when("the user re-publishes the endpoint")] +async fn when_reattach(world: &mut E2eWorld) { + let env = world.remote_env.clone(); + let session = current_session(world); + run_remote(world, &["remote", "attach", &session], &env); +} + +#[when("the user stops the session")] +async fn when_stop_session(world: &mut E2eWorld) { + let env = world.remote_env.clone(); + let session = current_session(world); + run_remote(world, &["remote", "stop", &session], &env); +} + +/// The one session this scenario started, read back from the CLI's own listing +/// rather than reconstructed — if the id were derived here the test could pass +/// against a listing that shows something else. +fn current_session(world: &mut E2eWorld) -> String { + let env = world.remote_env.clone(); + run_remote(world, &["remote", "status"], &env); + said(world) + .lines() + .find_map(|line| line.trim().strip_prefix("- ").map(str::to_owned)) + .expect("no session in the listing") +} + +#[then("the user is given an endpoint and a credential")] +async fn then_endpoint_and_credential(world: &mut E2eWorld) { + let said = said(world); + assert!(said.contains("endpoint: http://gpu-box."), "{said}"); + assert!(said.contains("api key:"), "{said}"); +} + +#[then("the user is told the endpoint is reachable by the whole network")] +async fn then_reach_stated(world: &mut E2eWorld) { + let said = said(world); + assert!( + said.contains("every machine on your tailnet"), + "the loopback intuition from local serving does not carry over and must be \ + corrected explicitly:\n{said}" + ); +} + +#[then("the machine is publishing that endpoint")] +async fn then_machine_publishing(world: &mut E2eWorld) { + let config = machine(world).exec(&["tailscale", "serve", "status", "--json"]); + assert!( + config.contains("\"8000\""), + "machine is not publishing:\n{config}" + ); +} + +#[then("the model and the endpoint are both reported healthy")] +async fn then_both_healthy(world: &mut E2eWorld) { + let said = said(world); + assert!(said.contains("model server: healthy"), "{said}"); + assert!(said.contains("endpoint published: yes"), "{said}"); +} + +#[then("the model is still healthy but the endpoint is reported gone")] +async fn then_model_healthy_endpoint_gone(world: &mut E2eWorld) { + let said = said(world); + assert!(said.contains("model server: healthy"), "{said}"); + assert!(said.contains("endpoint published: no"), "{said}"); +} + +#[then("the endpoint is restored without restarting the model")] +async fn then_endpoint_restored(world: &mut E2eWorld) { + let said = said(world); + assert!(said.contains("Endpoint re-published"), "{said}"); + assert!(said.contains("not restarted"), "{said}"); +} + +#[then("the endpoint and the model are both reported stopped")] +async fn then_both_stopped(world: &mut E2eWorld) { + let said = said(world); + assert!(said.contains("endpoint withdrawn: yes"), "{said}"); + assert!(said.contains("model server stopped: yes"), "{said}"); +} + +#[then("the machine is publishing nothing")] +async fn then_publishing_nothing(world: &mut E2eWorld) { + let config = machine(world).exec(&["tailscale", "serve", "status", "--json"]); + assert!( + !config.contains("\"8000\""), + "an endpoint was left published:\n{config}" + ); +} + +#[then("the session is no longer listed")] +async fn then_session_gone(world: &mut E2eWorld) { + let env = world.remote_env.clone(); + run_remote(world, &["remote", "status"], &env); + let said = said(world); + assert!(said.contains("No remote sessions"), "{said}"); +} + +#[then("the report names that machine")] +async fn then_report_names_machine(world: &mut E2eWorld) { + let said = said(world); + assert!( + said.contains("Health of gpu-box"), + "a remote report indistinguishable from a local one gets acted on against the \ + wrong computer:\n{said}" + ); +} diff --git a/tests/e2e-cucumber/tests/feature_naming.rs b/tests/e2e-cucumber/tests/feature_naming.rs index 318d55fb8..38859b430 100644 --- a/tests/e2e-cucumber/tests/feature_naming.rs +++ b/tests/e2e-cucumber/tests/feature_naming.rs @@ -36,6 +36,7 @@ const FEATURE_KEYS: &[(&str, &str)] = &[ ("logs.feature", "logs"), ("model_serving.feature", "serve"), ("networking.feature", "networking"), + ("remote.feature", "remote"), // Not `runtime`: `runtime_setup.feature` owns that key, and two files // sharing one key would collide on every index (`runtime-01` in both). ("runtime_lifecycle.feature", "runtime-lifecycle"), diff --git a/tests/remote-ssh/Dockerfile b/tests/remote-ssh/Dockerfile new file mode 100644 index 000000000..5ce3acda3 --- /dev/null +++ b/tests/remote-ssh/Dockerfile @@ -0,0 +1,66 @@ +# Copyright © Advanced Micro Devices, Inc., or its affiliates. +# +# SPDX-License-Identifier: MIT + +# A stand-in for a remote GPU machine: real OpenSSH, fake everything else. +# +# The point is to exercise the parts that only a real SSH server can prove — +# argument handling, stdin delivery, file copy, batch-mode refusal — without +# needing a GPU, ROCm, or a tailnet. The `rocm` and `tailscale` here are shell +# scripts that answer in the shapes the real tools do. +FROM alpine:3.20 + +# Empty by default, so the image builds against the distribution's normal +# HTTPS repositories. Some corporate networks intercept TLS and break that; +# those can pass plain-HTTP repository flags here rather than the Dockerfile +# shipping a weaker default for everyone. Package integrity does not rest on +# the transport either way — apk verifies signatures on what it fetches. +ARG APK_REPO_FLAGS="" + +# shellcheck disable=SC2086 # APK_REPO_FLAGS is deliberately word-split. +RUN apk add --no-cache ${APK_REPO_FLAGS} openssh-server openssh-sftp-server jq \ + && ssh-keygen -A \ + && mkdir -p /root/.ssh /var/lib/fake-rocm \ + && chmod 700 /root/.ssh + +# Key-only for the account the control channel uses. +RUN printf '%s\n' \ + 'PermitRootLogin prohibit-password' \ + 'PasswordAuthentication no' \ + 'PubkeyAuthentication yes' \ + 'AcceptEnv ROCM_*' \ + > /etc/ssh/sshd_config.d/rocm-test.conf + +# ...but one account that *does* ask for a password, so the batch-mode check has +# something real to prove. With password auth disabled everywhere, a connection +# fails for lack of a key whether or not batch mode is set, and the test passes +# without demonstrating anything. This account makes the two behave differently: +# refused immediately with batch mode, waiting for input without it. +RUN adduser -D prompt-only \ + && echo 'prompt-only:correct-horse' | chpasswd \ + && printf '%s\n' \ + 'Match User prompt-only' \ + ' PasswordAuthentication yes' \ + ' PubkeyAuthentication no' \ + >> /etc/ssh/sshd_config.d/rocm-test.conf + +# Set to 0 to build a machine that is reachable over SSH but has no `rocm` +# binary at all — the shape `remote doctor` needs to prove it refuses to +# install rather than failing earlier for lack of a connection. +ARG INCLUDE_ROCM_CLI=1 + +COPY fake-rocm.sh /tmp/fake-rocm.sh +COPY fake-tailscale.sh /usr/local/bin/tailscale +COPY examination.json /var/lib/fake-rocm/examination.json +RUN chmod +x /tmp/fake-rocm.sh /usr/local/bin/tailscale \ + && if [ "${INCLUDE_ROCM_CLI}" = "1" ]; then cp /tmp/fake-rocm.sh /usr/local/bin/rocm; fi \ + && rm -f /tmp/fake-rocm.sh \ + && echo '[]' > /var/lib/fake-rocm/services.json \ + && echo '{}' > /var/lib/fake-rocm/serve.json \ + # The readiness probe looks for a ROCm install independently of the CLI, so + # the stand-in needs the marker it looks for or it is refused as a GPU-less + # machine before anything else is exercised. + && mkdir -p /opt/rocm + +EXPOSE 22 +CMD ["/usr/sbin/sshd", "-D", "-e"] diff --git a/tests/remote-ssh/examination.json b/tests/remote-ssh/examination.json new file mode 100644 index 000000000..053d2ab19 --- /dev/null +++ b/tests/remote-ssh/examination.json @@ -0,0 +1,53 @@ +{ + "os_family": "linux", + "os_version": "6.8.0", + "distro_id": "ubuntu", + "distro_version": "24.04", + "kernel_release": "6.8.0-generic", + "kernel_cmdline": "", + "is_wsl": false, + "cpu_vendor": "AuthenticAMD", + "cpu_model": "", + "gpus": [], + "has_amd_gpu": false, + "has_nvidia_gpu": false, + "has_apu": false, + "has_discrete_amd": false, + "amdgpu_loaded": null, + "amdgpu_blacklisted_in": [], + "amdkfd_loaded": null, + "secure_boot": "unknown", + "iommu_kernel_param": "", + "kfd": null, + "render_devices": [], + "user_name": "", + "user_groups": [], + "in_render_group": null, + "in_video_group": null, + "rocm_version": "", + "rocm_install_method": "", + "rocm_path": "", + "rocminfo_present": false, + "rocminfo_status": "", + "hip_libs_on_ld_path": null, + "rocm_repos_seen": [], + "hip_sdk_path": "", + "hip_sdk_version": "", + "hipinfo_present": false, + "hipinfo_status": "", + "adrenalin_version": "", + "msvc_redist_present": null, + "framework": "unknown", + "framework_version": "", + "framework_rocm_version": "", + "framework_arch_list": [], + "framework_notes": [], + "framework_source": "", + "env": {}, + "in_container": false, + "container_kind": "", + "dmesg_amdgpu_tail": [], + "notes": [], + "probe_failures": [], + "status": "ok" +} diff --git a/tests/remote-ssh/fake-rocm.sh b/tests/remote-ssh/fake-rocm.sh new file mode 100755 index 000000000..42cb2c275 --- /dev/null +++ b/tests/remote-ssh/fake-rocm.sh @@ -0,0 +1,91 @@ +#!/bin/sh +# Copyright © Advanced Micro Devices, Inc., or its affiliates. +# +# SPDX-License-Identifier: MIT + +# A stand-in for the ROCm CLI on the test remote. It answers in the shapes the +# real one does for the handful of commands the control channel issues, and +# keeps its service registry in a file so a `serve` is visible to a later +# `services list`. +# +# It also records the API key it was handed on stdin, so a test can prove the +# credential arrived by that route and not through the command line. + +set -eu + +STATE_DIR="${FAKE_ROCM_STATE:-/var/lib/fake-rocm}" +SERVICES="${STATE_DIR}/services.json" +KEY_RECORD="${STATE_DIR}/last-api-key" +ARGV_RECORD="${STATE_DIR}/last-argv" + +mkdir -p "${STATE_DIR}" +printf '%s\n' "$*" > "${ARGV_RECORD}" + +case "${1:-}" in + --version) + echo "rocm 0.0.0-fake" + ;; + + serve) + shift + port=11434 + while [ "$#" -gt 0 ]; do + case "$1" in + --port) port="$2"; shift 2 ;; + *) shift ;; + esac + done + # The key reaches us as an exported environment variable, put there by the + # remote shell reading it off stdin. Recording it is how a test tells that + # apart from a key pasted into the command line. + printf '%s' "${ROCM_SERVE_API_KEY:-}" > "${KEY_RECORD}" + + service_id="svc-fake-${port}" + jq --arg id "${service_id}" --argjson port "${port}" \ + --argjson now "$(date +%s)000" ' + . + [{ + service_id: $id, engine: "fake", model_ref: "m", canonical_model_id: "m", + host: "127.0.0.1", port: $port, + endpoint_url: ("http://127.0.0.1:" + ($port|tostring) + "/v1"), + mode: "managed", status: "ready", supervisor_pid: 1, + manifest_path: "/tmp/m", log_path: "/tmp/l", engine_state_path: "/tmp/e", + created_at_unix_ms: $now + }]' "${SERVICES}" > "${SERVICES}.tmp" + mv "${SERVICES}.tmp" "${SERVICES}" + echo "started ${service_id}" + ;; + + services) + case "${2:-}" in + list) + # Only the JSON form is exercised; that is the machine-readable contract + # the remote orchestration actually reads. + cat "${SERVICES}" + ;; + stop) + service_id="${3:-}" + jq --arg id "${service_id}" 'map(select(.service_id != $id))' \ + "${SERVICES}" > "${SERVICES}.tmp" + mv "${SERVICES}.tmp" "${SERVICES}" + echo "stopped ${service_id}" + ;; + *) + echo "fake rocm: unsupported services subcommand: ${2:-}" >&2 + exit 2 + ;; + esac + ;; + + examine) + # A complete examination, not a plausible-looking fragment. The scoring runs + # on the calling machine and deserializes this into a fixed struct with no + # optional fields, so a partial document is not "close enough" — it fails to + # parse, which is exactly the mistake an earlier version of this stub made. + cat "${STATE_DIR}/examination.json" + ;; + + *) + echo "fake rocm: unsupported command: ${1:-}" >&2 + exit 2 + ;; +esac diff --git a/tests/remote-ssh/fake-tailscale.sh b/tests/remote-ssh/fake-tailscale.sh new file mode 100755 index 000000000..f7da29f0b --- /dev/null +++ b/tests/remote-ssh/fake-tailscale.sh @@ -0,0 +1,100 @@ +#!/bin/sh +# Copyright © Advanced Micro Devices, Inc., or its affiliates. +# +# SPDX-License-Identifier: MIT + +# A stand-in for Tailscale on the test remote, covering `serve` and `funnel`. +# +# It keeps its forwards in a file shaped like the real ServeConfig, so publish, +# inspect and withdraw are genuinely stateful: withdrawing has to actually +# remove the entry for a test to see it gone. That is the property worth +# checking here — a publish outlives reboots, so a withdrawal that silently does +# nothing leaves an endpoint exposed with nothing tracking it. +# +# `funnel` exists here only as a way to put the remote into the state the CLI +# must refuse to publish over. The CLI never runs it — it reads `AllowFunnel` +# out of the serve config and bails — but without a fake that can *write* that +# key, the classifier's exposure branch is only ever exercised against +# hand-written fixtures, never against something daemon-shaped. + +set -eu + +STATE_DIR="${FAKE_ROCM_STATE:-/var/lib/fake-rocm}" +SERVE="${STATE_DIR}/serve.json" +# The real daemon keys AllowFunnel by `host:port` using the node's tailnet DNS +# name. Any plausible name will do — the CLI matches on the port suffix because +# it does not track the host. +FUNNEL_HOST="${FAKE_TAILNET_HOST:-fake-gpu.tail1234.ts.net}" +mkdir -p "${STATE_DIR}" +[ -f "${SERVE}" ] || echo '{}' > "${SERVE}" + +command="${1:-}" +if [ "${command}" != "serve" ] && [ "${command}" != "funnel" ]; then + echo "fake tailscale: unsupported command: ${command}" >&2 + exit 2 +fi +shift + +if [ "${1:-}" = "status" ]; then + cat "${SERVE}" + exit 0 +fi + +port="" +target="" +off=0 +for arg in "$@"; do + case "${arg}" in + --bg) ;; + --tcp=*) port="${arg#--tcp=}" ;; + tcp://*) target="${arg#tcp://}" ;; + off) off=1 ;; + *) ;; + esac +done + +[ -n "${port}" ] || { echo "fake tailscale: no --tcp port given" >&2; exit 2; } + +# `funnel --tcp=N on|off` toggles AllowFunnel and touches nothing else, which is +# what makes it a usable fixture: a port can be Funnel-exposed with no forward +# behind it, or with one, and the CLI has to refuse in both cases. +if [ "${command}" = "funnel" ]; then + # The real daemon serves Funnel on 443, 8443 and 10000 only, and rejects + # anything else. Refusing the same ports here keeps the fixture from + # encoding an AllowFunnel entry that tailscaled could never emit — a test + # that passes against an impossible document proves nothing about the + # states the CLI will actually meet. + case "${port}" in + 443|8443|10000) ;; + *) + echo "fake tailscale: funnel is only supported on ports 443, 8443 and 10000" >&2 + exit 1 + ;; + esac + if [ "${off}" -eq 1 ]; then + jq --arg key "${FUNNEL_HOST}:${port}" \ + 'if .AllowFunnel then .AllowFunnel |= del(.[$key]) else . end' \ + "${SERVE}" > "${SERVE}.tmp" + else + jq --arg key "${FUNNEL_HOST}:${port}" \ + '.AllowFunnel = ((.AllowFunnel // {}) + {($key): true})' \ + "${SERVE}" > "${SERVE}.tmp" + fi + mv "${SERVE}.tmp" "${SERVE}" + exit 0 +fi + +if [ "${off}" -eq 1 ]; then + jq --arg port "${port}" 'if .TCP then .TCP |= del(.[$port]) else . end' \ + "${SERVE}" > "${SERVE}.tmp" + mv "${SERVE}.tmp" "${SERVE}" + exit 0 +fi + +[ -n "${target}" ] || { echo "fake tailscale: no forward target given" >&2; exit 2; } +# Integer map keys serialize as strings, which is what the real daemon emits and +# what the parser has to match. +jq --arg port "${port}" --arg target "${target}" \ + '.TCP = ((.TCP // {}) + {($port): {TCPForward: $target}})' \ + "${SERVE}" > "${SERVE}.tmp" +mv "${SERVE}.tmp" "${SERVE}" diff --git a/tests/remote-ssh/local-tailscale.sh b/tests/remote-ssh/local-tailscale.sh new file mode 100755 index 000000000..357bc235b --- /dev/null +++ b/tests/remote-ssh/local-tailscale.sh @@ -0,0 +1,24 @@ +#!/bin/sh +# Copyright © Advanced Micro Devices, Inc., or its affiliates. +# +# SPDX-License-Identifier: MIT + +# A stand-in for Tailscale on the *calling* machine, for the end-to-end harness. +# +# `rocm remote` shells out locally to discover peers and to check one is online +# before dialling it. This reports a single peer — the container the harness +# started — as an online, GPU-tagged machine whose address is loopback, which is +# where that container's SSH port is published. +# +# It answers `status` only. The remote half of the tailnet story (`serve`) is a +# different stub that runs inside the container. + +set -eu + +if [ "${1:-}" = "status" ]; then + cat "${FAKE_LOCAL_TAILSCALE_STATUS:?FAKE_LOCAL_TAILSCALE_STATUS must point at a status document}" + exit 0 +fi + +echo "fake local tailscale: unsupported command: ${1:-}" >&2 +exit 2 diff --git a/tests/remote-ssh/run-e2e.sh b/tests/remote-ssh/run-e2e.sh new file mode 100755 index 000000000..955b93afe --- /dev/null +++ b/tests/remote-ssh/run-e2e.sh @@ -0,0 +1,250 @@ +#!/usr/bin/env bash +# Copyright © Advanced Micro Devices, Inc., or its affiliates. +# +# SPDX-License-Identifier: MIT + +# Drive the real `rocm remote` end to end against a stand-in GPU machine. +# +# Unlike run.sh, which checks that the tools behave as assumed, this runs the +# actual binary through the whole orchestration: discover a target, probe it, +# start a model, publish an endpoint, reconcile status, re-publish after the +# endpoint is withdrawn out of band, and tear down. Real `ssh`, real argument +# building, real session records on disk. +# +# What is faked, and therefore what this does NOT prove: there is no GPU, no +# model, and no tailnet. `tailscale` is a stub on both sides, so the published +# endpoint does not carry traffic — reachability is the one thing here that +# still needs a real two-node tailnet to confirm. +# +# Usage: tests/remote-ssh/run-e2e.sh +# Requires: docker, ssh, jq, and a built `rocm` (cargo build -p rocm). + +set -euo pipefail + +HERE="$(cd "$(dirname "$0")" && pwd)" +REPO="$(cd "${HERE}/../.." && pwd)" +IMAGE="rocm-remote-ssh-test" +CONTAINER="rocm-remote-e2e-$$" +WORK="$(mktemp -d)" +FAILURES=0 +PORT="" + +cleanup() { + docker rm -f "${CONTAINER}" >/dev/null 2>&1 || true + rm -rf "${WORK}" +} +trap cleanup EXIT INT TERM + +pass() { echo " ok $1"; } +fail() { echo " FAIL $1${2:+ — $2}"; FAILURES=$((FAILURES + 1)); } + +expect_contains() { + local name="$1" needle="$2" haystack="$3" + if [[ "${haystack}" == *"${needle}"* ]]; then + pass "${name}" + else + fail "${name}" "expected to find '${needle}' in:"$'\n'"${haystack}" + fi +} + +expect_absent() { + local name="$1" needle="$2" haystack="$3" + if [[ "${haystack}" != *"${needle}"* ]]; then + pass "${name}" + else + fail "${name}" "did not expect '${needle}' in:"$'\n'"${haystack}" + fi +} + +for tool in docker ssh jq ssh-keygen; do + command -v "${tool}" >/dev/null 2>&1 || { echo "missing required tool: ${tool}" >&2; exit 1; } +done + +ROCM_BIN="${ROCM_BIN:-${REPO}/target/debug/rocm}" +[[ -x "${ROCM_BIN}" ]] || { echo "no rocm binary at ${ROCM_BIN}; run: cargo build -p rocm" >&2; exit 1; } + +find_free_port() { + local candidate + for _ in $(seq 1 50); do + candidate=$(( 20000 + RANDOM % 20000 )) + if ! (exec 3<>"/dev/tcp/127.0.0.1/${candidate}") 2>/dev/null; then + printf '%s' "${candidate}" + return 0 + fi + done + echo "could not find a free port" >&2 + return 1 +} + +echo "building the stand-in remote" +docker build -q --build-arg "APK_REPO_FLAGS=${ROCM_TEST_APK_REPOS:-}" -t "${IMAGE}" "${HERE}" >/dev/null + +PORT="$(find_free_port)" +echo "starting it on port ${PORT}" +docker run -d --name "${CONTAINER}" -p "127.0.0.1:${PORT}:22" "${IMAGE}" >/dev/null + +ssh-keygen -q -t ed25519 -N '' -f "${WORK}/id" -C rocm-remote-e2e +docker exec -i "${CONTAINER}" sh -c 'cat >> /root/.ssh/authorized_keys' < "${WORK}/id.pub" +docker exec "${CONTAINER}" chmod 600 /root/.ssh/authorized_keys + +# An isolated HOME: the CLI reads ~/.ssh/config through the system ssh, and puts +# its own session records under ~/.rocm. Both should be this run's, not the +# developer's. +export HOME="${WORK}/home" +mkdir -p "${HOME}/.ssh" +chmod 700 "${HOME}/.ssh" +cat > "${HOME}/.ssh/config" < "${WORK}/tailscale-status.json" <<'EOF' +{ + "BackendState": "Running", + "TUN": true, + "MagicDNSSuffix": "example-tailnet.ts.net", + "Self": { "HostName": "laptop", "DNSName": "laptop.example-tailnet.ts.net.", "Online": true }, + "Peer": { + "nodekey:aaa": { + "HostName": "gpu-box", + "DNSName": "gpu-box.example-tailnet.ts.net.", + "OS": "linux", + "TailscaleIPs": ["100.88.14.21"], + "Tags": ["tag:gpu"], + "Online": true + } + } +} +EOF +mkdir -p "${WORK}/bin" +cp "${HERE}/local-tailscale.sh" "${WORK}/bin/tailscale" +chmod +x "${WORK}/bin/tailscale" +export FAKE_LOCAL_TAILSCALE_STATUS="${WORK}/tailscale-status.json" +export PATH="${WORK}/bin:${PATH}" + +rocm() { "${ROCM_BIN}" "$@" 2>&1; } +in_container() { docker exec "${CONTAINER}" "$@"; } + +ready=0 +for _ in $(seq 1 60); do + if ssh -o BatchMode=yes -o ConnectTimeout=5 -F "${ROCM_REMOTE_SSH_CONFIG}" gpu-box true >/dev/null 2>&1; then + ready=1 + break + fi + sleep 0.25 +done +[[ "${ready}" -eq 1 ]] || { echo "the stand-in remote never accepted a connection" >&2; exit 1; } + +echo +echo "discovery" +out="$(rocm remote targets)" +expect_contains "the machine is discovered" "- gpu-box" "${out}" +expect_contains "and reported online" "online: yes" "${out}" +out="$(rocm remote targets --tag gpu)" +expect_contains "and can be narrowed by tag" "- gpu-box" "${out}" + +echo +echo "health" +out="$(rocm remote doctor gpu-box || true)" +expect_contains "the remote's own state is fetched and scored here" "Health of gpu-box" "${out}" + +echo +echo "serve" +out="$(rocm remote serve gpu-box test-model)" +expect_contains "a model is served" "Model serving on gpu-box" "${out}" +expect_contains "and an endpoint is printed" "http://gpu-box.example-tailnet.ts.net:8000/v1" "${out}" +expect_contains "with a credential" "api key:" "${out}" +expect_contains "and who can reach it is stated" "every machine on your tailnet" "${out}" + +serve_config="$(in_container tailscale serve status --json)" +expect_contains "the remote published the endpoint" '"8000"' "${serve_config}" + +# The credential must have arrived by stdin, not on a command line either +# machine exposes in its process table. +recorded_argv="$(in_container cat /var/lib/fake-rocm/last-argv)" +recorded_key="$(in_container cat /var/lib/fake-rocm/last-api-key)" +if [[ -n "${recorded_key}" && "${recorded_argv}" != *"${recorded_key}"* ]]; then + pass "the credential reached the model without appearing in its arguments" +else + fail "the credential reached the model without appearing in its arguments" \ + "key='${recorded_key}' argv='${recorded_argv}'" +fi + +echo +echo "status" +out="$(rocm remote status)" +expect_contains "the model is reported healthy" "model server: healthy" "${out}" +expect_contains "and the endpoint published" "endpoint published: yes" "${out}" + +echo +echo "the endpoint is withdrawn behind our back" +in_container tailscale serve --tcp=8000 off +out="$(rocm remote status)" +expect_contains "the model is still healthy" "model server: healthy" "${out}" +expect_contains "but the endpoint is reported gone" "endpoint published: no" "${out}" +expect_contains "and re-publishing is the suggested fix" "rocm remote attach" "${out}" + +session_id="$(rocm remote status | sed -n 's/^- \(remote-.*\)$/\1/p' | head -n1)" +out="$(rocm remote attach "${session_id}")" +expect_contains "attaching restores the endpoint" "Endpoint re-published" "${out}" +expect_contains "without restarting the model" "not restarted" "${out}" + +serve_config="$(in_container tailscale serve status --json)" +expect_contains "the remote published the endpoint" '"8000"' "${serve_config}" + +echo +echo "teardown" +out="$(rocm remote stop "${session_id}")" +expect_contains "the endpoint is withdrawn" "endpoint withdrawn: yes" "${out}" +expect_contains "and the model stopped" "model server stopped: yes" "${out}" + +serve_config="$(in_container tailscale serve status --json)" +expect_absent "nothing is left published on the machine" '"8000"' "${serve_config}" +out="$(rocm remote status)" +expect_contains "and the session is gone from this machine" "No remote sessions" "${out}" + +echo +echo "a Funnel-exposed port is refused" +# Run last, on a torn-down machine, so the port is free and the only thing +# standing between the CLI and publishing is the exposure itself. Funnel puts a +# port on the public internet rather than just the tailnet, so publishing a +# model endpoint over one turns "everyone on your tailnet" into "everyone". +# +# The unit tests cover this classification against hand-written fixtures. What +# they cannot show is that the CLI reads the key the daemon actually writes, +# which is the whole reason this runs here. +# +# On 443, because that is one of the three ports Funnel can actually serve — +# the CLI has to be driven at `--tailnet-port 443` to meet the state at all. +# That is worth knowing on its own: at the default tailnet port of 8000 this +# guard can never fire, since Funnel cannot listen there. +in_container tailscale funnel --tcp=443 on +out="$(rocm remote serve gpu-box test-model --tailnet-port 443 || true)" +expect_contains "publishing over it is refused" "Tailscale Funnel allowed" "${out}" +expect_contains "and the way out is named" "tailscale funnel --tcp=443 off" "${out}" + +serve_config="$(in_container tailscale serve status --json)" +expect_absent "and nothing was published in spite of the refusal" '"TCPForward"' "${serve_config}" + +in_container tailscale funnel --tcp=443 off + +echo +if [[ "${FAILURES}" -eq 0 ]]; then + echo "all checks passed" +else + echo "${FAILURES} check(s) failed" + exit 1 +fi diff --git a/tests/remote-ssh/run.sh b/tests/remote-ssh/run.sh new file mode 100755 index 000000000..03ee44ed3 --- /dev/null +++ b/tests/remote-ssh/run.sh @@ -0,0 +1,297 @@ +#!/usr/bin/env bash +# Copyright © Advanced Micro Devices, Inc., or its affiliates. +# +# SPDX-License-Identifier: MIT + +# Check the assumptions `rocm remote` makes about the tools it drives, against a +# real OpenSSH server rather than a stand-in. +# +# Scope, stated plainly: this exercises the *tools and their contracts* — ssh, +# scp, the service listing, the serve config — not the Rust call paths, which +# are unit-tested against a scripted transport. The transport is private to a +# binary crate, so no integration test can reach it. What is proven here is that +# the behaviour those unit tests assume is the behaviour the real tools have. +# +# Usage: tests/remote-ssh/run.sh +# Requires: docker, ssh, scp, jq. No GPU, no ROCm, no tailnet. + +set -euo pipefail + +HERE="$(cd "$(dirname "$0")" && pwd)" +IMAGE="rocm-remote-ssh-test" +CONTAINER="rocm-remote-ssh-test-$$" +WORK="$(mktemp -d)" +FAILURES=0 +PORT="" + +cleanup() { + docker rm -f "${CONTAINER}" >/dev/null 2>&1 || true + rm -rf "${WORK}" +} +trap cleanup EXIT INT TERM + +pass() { echo " ok $1"; } +fail() { echo " FAIL $1${2:+ — $2}"; FAILURES=$((FAILURES + 1)); } + +# Compare an actual value with an expected one. +expect() { + local name="$1" expected="$2" actual="$3" + if [[ "${actual}" == "${expected}" ]]; then + pass "${name}" + else + fail "${name}" "expected '${expected}', got '${actual}'" + fi +} + +# Assert a value contains a substring. For output whose exact wording is not +# the point but whose identity is — distinguishing a tool's own refusal from +# ssh failing for reasons of its own, say. +expect_contains() { + local name="$1" needle="$2" haystack="$3" + if [[ "${haystack}" == *"${needle}"* ]]; then + pass "${name}" + else + fail "${name}" "expected to find '${needle}' in:"$'\n'"${haystack}" + fi +} + +# Assert a jq filter holds over a JSON document. +expect_json() { + local name="$1" filter="$2" document="$3" + if printf '%s' "${document}" | jq -e "${filter}" >/dev/null 2>&1; then + pass "${name}" + else + fail "${name}" "filter '${filter}' did not hold" + fi +} + +for tool in docker ssh scp jq ssh-keygen; do + command -v "${tool}" >/dev/null 2>&1 || { echo "missing required tool: ${tool}" >&2; exit 1; } +done + +# An explicit host port, not an ephemeral one: some Docker setups do not forward +# ephemeral or loopback-bound publications to the host at all. +find_free_port() { + local candidate + for _ in $(seq 1 50); do + candidate=$(( 20000 + RANDOM % 20000 )) + if ! (exec 3<>"/dev/tcp/127.0.0.1/${candidate}") 2>/dev/null; then + printf '%s' "${candidate}" + return 0 + fi + done + echo "could not find a free port" >&2 + return 1 +} + +echo "building the stand-in remote" +# ROCM_TEST_APK_REPOS lets a network that intercepts TLS point the package +# manager at plain-HTTP mirrors. Unset everywhere else, which is the default. +docker build -q \ + --build-arg "APK_REPO_FLAGS=${ROCM_TEST_APK_REPOS:-}" \ + -t "${IMAGE}" "${HERE}" >/dev/null + +PORT="$(find_free_port)" +echo "starting it on port ${PORT}" +# Loopback-only, no fallback: this image ships a password account +# (Dockerfile's `prompt-only`) so the batch-mode check has something real to +# refuse. Falling back to an all-interfaces publish on a failed loopback bind +# would put that guessable-password shell on the LAN instead of just failing. +docker run -d --name "${CONTAINER}" -p "127.0.0.1:${PORT}:22" "${IMAGE}" >/dev/null + +ssh-keygen -q -t ed25519 -N '' -f "${WORK}/id" -C rocm-remote-test +docker exec -i "${CONTAINER}" sh -c 'cat >> /root/.ssh/authorized_keys' < "${WORK}/id.pub" +docker exec "${CONTAINER}" chmod 600 /root/.ssh/authorized_keys + +# Mirrors what the control channel builds: never prompt, fail fast, reuse +# briefly. Host-key checking is off only because this container is recreated +# with a new key every run. +SSH_COMMON=( + -o BatchMode=yes + -o ConnectTimeout=10 + -o StrictHostKeyChecking=no + -o UserKnownHostsFile=/dev/null + -o LogLevel=ERROR + -i "${WORK}/id" +) +remote() { + ssh "${SSH_COMMON[@]}" -o ControlMaster=auto -o ControlPersist=60s \ + -p "${PORT}" root@127.0.0.1 -- "$@" +} + +# Wait for sshd rather than sleeping a fixed amount. +ready=0 +for _ in $(seq 1 60); do + if remote true >/dev/null 2>&1; then ready=1; break; fi + sleep 0.25 +done +[[ "${ready}" -eq 1 ]] || { echo "the stand-in remote never accepted a connection" >&2; exit 1; } + +echo +echo "control channel" + +expect "runs a command and returns its output" "hello" "$(remote echo hello 2>/dev/null || true)" + +# A non-zero remote exit must be a readable outcome, not a transport failure. +# Quoted as one string: ssh joins its arguments into a single remote command, +# so `remote sh -c 'exit 7'` would reach the remote as `sh -c exit 7` and run +# `exit` with an ignored argument. +set +e +remote 'exit 7' >/dev/null 2>&1 +code=$? +set -e +expect "a non-zero remote exit arrives as an exit code" "7" "${code}" + +# The `--` guard: an argument starting with a dash must reach the remote intact +# rather than being eaten by the local ssh. +expect "an argument starting with a dash reaches the remote" "--managed" \ + "$(remote 'echo --managed' 2>/dev/null || true)" + +# The property behind passing a credential on stdin instead of in the command. +piped="$(printf 'sekrit' | ssh "${SSH_COMMON[@]}" -p "${PORT}" root@127.0.0.1 -- \ + 'IFS= read -r K; printf %s "$K"' 2>/dev/null || true)" +expect "a value piped in arrives on the remote's stdin" "sekrit" "${piped}" + +echo "hello-from-here" > "${WORK}/pushed.txt" +scp "${SSH_COMMON[@]}" -P "${PORT}" "${WORK}/pushed.txt" root@127.0.0.1:/tmp/pushed.txt >/dev/null 2>&1 +expect "scp copies a file, taking its port as -P" "hello-from-here" \ + "$(remote cat /tmp/pushed.txt 2>/dev/null || true)" + +# Batch mode is what stops a status poll hanging forever on a host that wants a +# password. The container has one account that genuinely asks for one, so the +# two settings behave differently — against a key-only account both fail for +# lack of a key and the check would prove nothing. +# +# Batch mode: refused immediately, without reading the password waiting on stdin. +start=$(date +%s) +set +e +printf 'correct-horse\n' | timeout 20 ssh "${SSH_COMMON[@]}" -p "${PORT}" \ + prompt-only@127.0.0.1 -- true >/dev/null 2>&1 +denied=$? +set -e +elapsed=$(( $(date +%s) - start )) +if [[ "${denied}" -ne 0 && "${denied}" -ne 124 ]]; then + pass "an account that wants a password is refused, not prompted" +else + fail "an account that wants a password is refused, not prompted" "rc=${denied}" +fi +if [[ "${elapsed}" -lt 15 ]]; then + pass "and refused promptly" +else + fail "and refused promptly" "took ${elapsed}s" +fi + +# The counterfactual, which is what makes the check above mean anything: the +# same account, the same server, without batch mode. ssh reads a password from +# the terminal rather than from stdin, so this needs a pty to reproduce what a +# user would hit — and there it sits waiting for input nobody will type, which +# is precisely the hang batch mode exists to prevent. Timing out here is the +# expected result. +if command -v script >/dev/null 2>&1; then + set +e + timeout 15 script -qec "ssh -o BatchMode=no -o StrictHostKeyChecking=no \ + -o UserKnownHostsFile=/dev/null -o LogLevel=ERROR -o NumberOfPasswordPrompts=1 \ + -o ConnectTimeout=10 -p ${PORT} prompt-only@127.0.0.1 -- true" /dev/null \ + /dev/null 2>&1 + hung=$? + set -e + if [[ "${hung}" -eq 124 ]]; then + pass "and without batch mode the same host waits for input instead" + else + # Not a product failure: some ssh builds give up on a pty with no reader. + # Say so rather than asserting a behaviour this environment does not show. + pass "and without batch mode the same host does not refuse cleanly (rc=${hung})" + fi +else + echo " skip batch-mode counterfactual (no 'script' to allocate a pty)" +fi + +echo +echo "service registry contract" + +remote rocm serve m --managed --host 127.0.0.1 --port 11434 >/dev/null +listing="$(remote rocm services list --json)" +expect_json "the listing is JSON" "." "${listing}" +expect_json "the started service appears on its port" \ + '.[] | select(.port == 11434)' "${listing}" + +# The credential must have travelled by stdin, not in the command line — both +# machines expose command arguments in their process tables. +expect "the key reached the remote process" "sekrit-key" \ + "$(printf 'sekrit-key' | ssh "${SSH_COMMON[@]}" -p "${PORT}" root@127.0.0.1 -- \ + 'IFS= read -r K; export ROCM_SERVE_API_KEY="$K"; rocm serve m --port 11500 >/dev/null; cat /var/lib/fake-rocm/last-api-key' 2>/dev/null || true)" +# Read without `|| true`: an unreadable file yields an empty string, which +# trivially does not contain the key, so the absence check below would pass +# having verified nothing. This is the check backing the claim that a credential +# never reaches a process table, so it has to fail loudly when it cannot look. +recorded_argv="$(remote cat /var/lib/fake-rocm/last-argv)" +if [[ -z "${recorded_argv}" ]]; then + fail "and never appeared in the command line" \ + "could not read /var/lib/fake-rocm/last-argv, so nothing was actually checked" +elif printf '%s' "${recorded_argv}" | grep -q 'sekrit-key'; then + fail "and never appeared in the command line" "found in: ${recorded_argv}" +else + pass "and never appeared in the command line" +fi + +echo +echo "endpoint publishing" + +remote tailscale serve --bg --tcp=8000 tcp://127.0.0.1:11434 +serve_config="$(remote tailscale serve status --json)" +# Integer map keys serialize as strings; a parser looking for a number would +# match nothing and report every endpoint as absent. +expect_json "the published port is keyed as a string" '.TCP."8000"' "${serve_config}" +expect_json "it forwards to the model server's loopback port" \ + '.TCP."8000".TCPForward == "127.0.0.1:11434"' "${serve_config}" + +# The one that matters most. A publish is configuration, not a process: it +# outlives reboots, so a withdrawal that quietly does nothing leaves a GPU +# endpoint exposed with nothing tracking it. +remote tailscale serve --tcp=8000 off +expect_json "withdrawing actually removes the forward" '.TCP."8000" == null' \ + "$(remote tailscale serve status --json)" + +# Funnel is the hazard the publish path refuses over, and it lives in a +# different key than the forwards do. Asserting its shape here is what lets the +# e2e check below mean something: without it, a fake that wrote AllowFunnel in +# the wrong place would make the CLI look correctly cautious while actually +# reading nothing. +# +# 443 rather than 8000: Funnel only serves 443, 8443 and 10000, so an +# AllowFunnel entry for any other port is a document the daemon cannot produce. +remote tailscale funnel --tcp=443 on +funnel_config="$(remote tailscale serve status --json)" +expect_json "allowing Funnel is keyed host:port, not by port alone" \ + '.AllowFunnel | keys | length == 1 and (.[0] | endswith(":443"))' "${funnel_config}" +expect_json "and records it as allowed rather than merely present" \ + '.AllowFunnel | to_entries[0].value == true' "${funnel_config}" +expect_json "and says nothing about a forward, which is a separate question" \ + '.TCP."443" == null' "${funnel_config}" + +# Asserting on the message, not merely on a non-zero exit: `remote` is an ssh +# wrapper, and ssh fails non-zero for its own reasons too — a dropped control +# socket, a missing fixture. Taking any failure as proof would let a broken +# harness report this as passing. +refusal="$(remote tailscale funnel --tcp=8000 on 2>&1 || true)" +expect_contains "a port Funnel cannot serve is refused" \ + "funnel is only supported on ports 443, 8443 and 10000" "${refusal}" + +remote tailscale funnel --tcp=443 off +expect_json "turning Funnel off clears the exposure" \ + '(.AllowFunnel // {}) | length == 0' "$(remote tailscale serve status --json)" + +echo +echo "teardown" +remote rocm services stop svc-fake-11434 --yes >/dev/null +expect_json "the stopped service leaves the registry" \ + 'map(select(.service_id == "svc-fake-11434")) | length == 0' \ + "$(remote rocm services list --json)" + +echo +if [[ "${FAILURES}" -eq 0 ]]; then + echo "all checks passed" +else + echo "${FAILURES} check(s) failed" + exit 1 +fi From 58501f3772a62fd46dc0aecbdd7f04b371066ce6 Mon Sep 17 00:00:00 2001 From: Eugene Volen Date: Mon, 14 Sep 2026 08:57:54 +0000 Subject: [PATCH 3/5] docs(remote): document the remote commands and how to exercise them MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Covers the `rocm remote` surface, and in docs/testing.md how to run the two container-backed lanes — including the ROCM_TEST_APK_REPOS escape hatch a network that intercepts TLS needs to build the fixture image. Signed-off-by: Eugene Volen --- README.md | 43 ++++++++++++++++++++++++++++++++++++++++++- docs/testing.md | 42 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 84 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index cbad76471..ee9d41fab 100644 --- a/README.md +++ b/README.md @@ -231,6 +231,7 @@ form works depends on the engine your GPU selects. | `rocm install sdk` | Install TheRock ROCm wheels into a managed Python environment | | `rocm install driver` | Install the AMD kernel driver on Linux | | `rocm serve ` | Start a local OpenAI-compatible model server | +| `rocm remote serve ` | Serve a model on another GPU machine over your private network | | `rocm dash` | Open the full-screen telemetry dashboard | | `rocm bench load --endpoint ` | Load-test a local OpenAI-compatible endpoint | | `rocm setup status` | Show first-time setup state | @@ -498,7 +499,7 @@ model across multiple GPUs is not supported. Because selection uses the Manage background servers started with `--managed`: ``` -rocm services list [--all] +rocm services list [--all] [--json] rocm services logs rocm services stop [--yes] rocm services restart [--yes] @@ -519,6 +520,46 @@ however recent — that is the flag `prune` names in its own summary when it reports how many records it kept for being too recent. The two cannot be combined. +`--json` prints the service records verbatim, for scripting and for the remote +orchestration below. + +### Remote machines (preview) + +Run a model on a different GPU machine and reach it from your own. Both machines +join a [Tailscale](https://tailscale.com) network; the GPU machine serves the +model on its own loopback address and publishes that port onto the network, so +the endpoint keeps working after the command exits and answers from any of your +machines rather than only the one that started it. + +```console +rocm remote targets [--tag ] +rocm remote doctor [--symptom ] +rocm remote serve [--engine ] [--gpu ] + [--tailnet-port ] [--install-rocm] +rocm remote status [] +rocm remote attach +rocm remote stop [--force] +``` + +`targets` lists machines on your network — it does not check whether they can +actually serve, which is what `doctor` is for. `serve` prepares the machine +(installing the CLI if it is missing), starts the model, publishes the endpoint +and prints the address together with an API key. + +**The endpoint is reachable by every machine on your network that your network's +access rules allow**, not just yours. That API key is what stops anyone else +using it, so `rocm remote` always sets one — unlike local serving, which is +credential-free because only your own machine can reach it. + +`status` reports the model and the endpoint separately, because either can fail +alone: a healthy model with no endpoint needs `attach`, not a restart. `stop` +withdraws the endpoint and stops the model, and keeps the session listed if it +cannot confirm both — use `--force` to forget one whose machine is gone. + +Communication with the machine uses your existing `ssh` setup. Set +`ROCM_REMOTE_SSH_CONFIG` to point at a configuration file other than the +default. + ### Dashboard ``` diff --git a/docs/testing.md b/docs/testing.md index 8d2fdaebc..a60b3eff0 100644 --- a/docs/testing.md +++ b/docs/testing.md @@ -34,6 +34,48 @@ If the workspace is already built: python scripts/smoke_local.py --skip-build ``` +## Remote control-channel checks + +`rocm remote` drives `ssh`, `scp`, and the remote machine's own tooling. Its unit +tests use a scripted stand-in, which proves control flow but assumes those tools +behave a certain way. Two scripts check that assumption against a real OpenSSH +server in a container — no GPU, no ROCm, no tailnet, since the remote's `rocm` +and `tailscale` are stand-ins: + +```bash +tests/remote-ssh/run.sh # the tools and their contracts +tests/remote-ssh/run-e2e.sh # the real binary, end to end +``` + +`run.sh` covers argument handling, exit-code propagation, delivering a +credential on stdin, file copy, batch-mode refusal, that withdrawing a +published endpoint actually removes it, and the shape Tailscale Funnel takes in +the serve config — including that a port Funnel cannot serve is refused. + +`run-e2e.sh` runs the built `rocm` through the whole flow — discover, probe, +serve, publish, reconcile status, re-publish after an out-of-band withdrawal, +tear down, and refuse to publish over a Funnel-exposed port — and needs +`cargo build -p rocm` first. It does not prove the endpoint carries traffic; +that needs a real two-node tailnet. + +The `rocm remote` cucumber scenarios that need a second machine carry +`@requires-docker` and run only when `E2E_INCLUDE_DOCKER=1`, which the +GitHub-hosted `E2E tests` lane sets. A working daemon alone is not enough — the +self-hosted GPU runners have one but cannot reach the package mirror the fixture +image builds from, so they skip those scenarios. To run them locally: + +```bash +E2E_INCLUDE_DOCKER=1 cargo xtask e2e +``` + +On a network that intercepts TLS, point the container's package manager at +plain-HTTP mirrors: + +```bash +export ROCM_TEST_APK_REPOS="--repository http://dl-cdn.alpinelinux.org/alpine/v3.20/main \ + --repository http://dl-cdn.alpinelinux.org/alpine/v3.20/community" +``` + ## CI test selection On pull requests, CI does not test the whole workspace. The `test` job runs From 15025a527003e8d840e88fd9f69e5d315c22b25c Mon Sep 17 00:00:00 2001 From: Eugene Volen Date: Fri, 25 Sep 2026 12:21:21 +0000 Subject: [PATCH 4/5] fix(remote): let the shell enforce the read-before-serve invariant MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The remote serve command joined the API-key `read`, the export, and `rocm serve` with `;`, so the compound's exit status was whatever `serve` returned. `run_with_piped_io` demotes a broken-pipe write error to a symptom whenever the command also failed, and that is only sound if a failed read cannot be followed by a successful serve. The ordering held by prose, not by the shell: reorder the command and a truncated key would classify as success on the one path guarding the model. `&&` makes the shell hold it instead. Tightening the chain also means the key has to arrive as a complete line. `IFS= read -r` returns non-zero at EOF without a terminator even though it did assign the variable, so the bare key the caller used to send now short-circuits the chain and nothing starts. Send it with a trailing newline, which is also what gives the short circuit meaning: a write truncated part-way delivers no newline, so the read fails and no model comes up holding half a key. Both halves are pinned against a real shell — the command's exit status for the chain, and the recorded stdin at the real call site for the terminator. A string assertion can see neither. Also create the remote staging directory 0700, matching the local `create_restricted_dir`. Its test asserts the parent's mode as well as the leaf's, because `chmod` names the leaf and would set it either way: checking the leaf alone passed with `umask 077` deleted, leaving the half the doc comment calls load-bearing unable to fail. A parent that `mkdir -p` created is covered by the umask and by nothing else. The archive, checksum, and signature sit there before any of them is verified, so a permissive umask on a shared box let another local user read or race them. `umask 077` closes the window at creation and the chmod repairs a directory an earlier run left open. Gate the remote-ssh CI job on clippy and prek and skip it on manual dispatch, as every other heavy job in the workflow already does, so a fast lint failure short-circuits before the containerised build. Record the teardown check-then-act window in `withdraw` as a known limitation: tailscale offers no compare-and-swap on a port, so only a daemon holding the claim could close it, and re-reading state before the exec would narrow the window while implying it was shut. Signed-off-by: Eugene Volen --- .github/workflows/ci.yml | 7 +- apps/rocm/src/remote/mod.rs | 117 ++++++++++++++++++++++++++++-- apps/rocm/src/remote/provision.rs | 107 ++++++++++++++++++++++++++- apps/rocm/src/remote/publish.rs | 8 ++ apps/rocm/src/remote/transport.rs | 10 ++- 5 files changed, 237 insertions(+), 12 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index a80f861ad..65f570109 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -879,8 +879,11 @@ jobs: name: remote control channel (containerised) runs-on: ubuntu-latest timeout-minutes: 30 - needs: changes - if: needs.changes.outputs.heavy == 'true' + # Gated like every other heavy job here: a fast lint failure should + # short-circuit before this spends 30 minutes on a container and a cold + # cargo build, and a manual dispatch should stay a fast loop. + needs: [changes, clippy, prek] + if: github.event_name != 'workflow_dispatch' && needs.changes.outputs.heavy == 'true' steps: - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 diff --git a/apps/rocm/src/remote/mod.rs b/apps/rocm/src/remote/mod.rs index 3b53a6ec4..3c366035a 100644 --- a/apps/rocm/src/remote/mod.rs +++ b/apps/rocm/src/remote/mod.rs @@ -368,9 +368,18 @@ fn serve_with_transport( } println!("Starting {} on {} ...", request.model, request.target); - let start = match transport - .exec_with_stdin(&remote_serve_command(&remote_cli, request), Some(&api_key)) - { + // The trailing newline is what makes the `&&` chain in + // `remote_serve_command` work, not cosmetics: `IFS= read -r` returns + // non-zero when it hits EOF without one, even though it did assign the + // variable. Send the bare key and the remote reads it, reports failure, and + // the `&&` stops the model from ever starting. It is also what makes the + // short circuit *mean* something — a write truncated part-way delivers no + // newline, so the read fails and nothing serves with half a key. + let key_payload = format!("{api_key}\n"); + let start = match transport.exec_with_stdin( + &remote_serve_command(&remote_cli, request), + Some(&key_payload), + ) { Ok(start) => start, Err(error) => { // The command may already have reached the remote — contact can be @@ -599,9 +608,16 @@ fn resolve_target(target: &str) -> Result { /// table, so an interpolated key would be readable by any other user on either /// machine. `--require-api-key` is what makes the loopback bind authenticated /// anyway, since the publish widens who can reach it. +/// +/// Joined with `&&`, not `;`, and that is load-bearing rather than stylistic. +/// [`transport::run_with_piped_io`] treats a broken pipe on the stdin writer as a +/// mere symptom whenever the command itself failed, which is only sound if a +/// failed `read` cannot be followed by a successful `serve`. `&&` is what makes +/// the shell enforce that; under `;` the compound's status is whatever `serve` +/// returned, so a truncated key could report success on the credential path. fn remote_serve_command(remote_cli: &str, request: &ServeRequest) -> String { let mut command = format!( - "IFS= read -r ROCM_SERVE_API_KEY; export ROCM_SERVE_API_KEY; \ + "IFS= read -r ROCM_SERVE_API_KEY && export ROCM_SERVE_API_KEY && \ {remote_cli} serve {} --managed --require-api-key --host {} --port {}", shell_quote(&request.model), publish::LOOPBACK, @@ -1241,12 +1257,85 @@ mod tests { // interpolated key would be readable by any other user on either. let command = remote_serve_command("rocm", &request()); assert!( - command.starts_with("IFS= read -r ROCM_SERVE_API_KEY;"), + command.starts_with("IFS= read -r ROCM_SERVE_API_KEY &&"), "{command}" ); assert!(command.contains("export ROCM_SERVE_API_KEY"), "{command}"); } + /// Run the generated serve command under a real shell with `stdin_payload` + /// on its stdin, and report whether it succeeded and whether it reached the + /// serve step. + /// + /// `echo REACHED_SERVE` stands in for the remote CLI: reaching it means the + /// `&&` chain did not short-circuit, and the marker says so in the failure. + fn serve_command_under_a_shell(stdin_payload: &str) -> (bool, bool) { + use std::io::Write as _; + + let command = remote_serve_command("echo REACHED_SERVE", &request()); + let mut child = std::process::Command::new("sh") + .arg("-c") + .arg(&command) + .stdin(std::process::Stdio::piped()) + .stdout(std::process::Stdio::piped()) + .stderr(std::process::Stdio::null()) + .spawn() + .expect("failed to run the generated command under a shell"); + child + .stdin + .take() + .expect("stdin was piped") + .write_all(stdin_payload.as_bytes()) + .expect("the payload is far smaller than a pipe buffer"); + + let output = child.wait_with_output().expect("the shell should finish"); + ( + output.status.success(), + String::from_utf8_lossy(&output.stdout).contains("REACHED_SERVE"), + ) + } + + #[test] + fn only_a_complete_key_lets_the_command_reach_serve() { + // The invariant `transport::run_with_piped_io` leans on: it demotes a + // broken-pipe write error to a symptom whenever the command also failed, + // which is only sound if a failed `read` cannot be followed by a + // successful `serve`. Under `;` it could be — the compound's status + // would be whatever `serve` returned — and a truncated API key would + // then classify as success on the one path guarding the model. + // + // Asked of a real shell, and asked for the *exit status*: the whole + // property is what `&&` does to a compound command, which no assertion + // about the string can see. The same reasoning as + // `the_staging_path_still_expands_on_the_remote_shell` in `provision`. + + // Nothing at all — the far side of a pipe that broke before any byte. + assert_eq!( + serve_command_under_a_shell(""), + (false, false), + "no key arrived, yet the command served or reported success" + ); + + // A key with no terminating newline. This is the case that matters and + // the one a closed-stdin test cannot see: `read` assigns the variable + // but still returns non-zero at EOF, so this is indistinguishable from + // a truncated write — and must not serve. It is also exactly what the + // caller used to send, which made the `&&` chain refuse every real + // start until the caller began terminating the payload. + assert_eq!( + serve_command_under_a_shell("abc123"), + (false, false), + "an unterminated key is a truncated key; it must not reach serve" + ); + + // A complete line: the shape `serve_with_transport` actually sends. + assert_eq!( + serve_command_under_a_shell("abc123\n"), + (true, true), + "a complete key must reach serve, or no remote model ever starts" + ); + } + #[test] fn optional_engine_and_gpu_are_threaded_through() { let command = remote_serve_command( @@ -1957,10 +2046,22 @@ mod tests { } if command.contains("read -r ROCM_SERVE_API_KEY") => Some(key.clone()), _ => None, }); + let sent_key = sent_key.unwrap_or_else(|| { + panic!( + "the model-starting command must receive the key over stdin: {:?}", + transport.calls() + ) + }); + assert!(!sent_key.trim().is_empty(), "the key sent was blank"); + // The terminator is as load-bearing as the key. `IFS= read -r` returns + // non-zero at EOF without one, so an unterminated payload short-circuits + // the `&&` chain and the model never starts — which is what the + // container lane caught after the chain was tightened, and what no + // assertion about the command string could see. assert!( - sent_key.is_some_and(|key| !key.is_empty()), - "the model-starting command must receive the key over stdin: {:?}", - transport.calls() + sent_key.ends_with('\n'), + "the key must arrive as a complete line or `read` fails and nothing \ + serves; got {sent_key:?}" ); let _ = std::fs::remove_dir_all(root); } diff --git a/apps/rocm/src/remote/provision.rs b/apps/rocm/src/remote/provision.rs index 0a71b5a6c..239253dcd 100644 --- a/apps/rocm/src/remote/provision.rs +++ b/apps/rocm/src/remote/provision.rs @@ -143,6 +143,28 @@ impl Drop for StagingDir { } } +/// Make the remote's staging directory, restricted to its owner. +/// +/// 0700 on the far side for the same reason [`create_restricted_dir`] insists on +/// it here: the archive, its checksum, and its signature sit in this directory +/// *before* they are verified, so another user on a shared remote box could +/// otherwise read them, or swap one out between the push and the check. The +/// signature gate still has to be defeated for that to become an install, but +/// "only one gate left" is not the posture the local side settles for. +/// +/// `umask 077` covers the creation itself — a mode applied afterwards leaves a +/// window at whatever the remote's umask happens to be — and the `chmod` then +/// repairs a directory a previous run already left too permissive. `mkdir -m` +/// does neither: it ignores an existing directory and skips the parent +/// directories `-p` creates. +/// +/// `remote_dir` is interpolated unquoted so the far shell still expands the +/// `$HOME` in [`REMOTE_STAGING`]; see +/// `the_staging_path_still_expands_on_the_remote_shell`. +fn staging_dir_command(remote_dir: &str) -> String { + format!("(umask 077 && mkdir -p {remote_dir}) && chmod 700 {remote_dir}") +} + /// Fetch a build for the remote's platform on this machine, then push it. fn push_matched_artifact( transport: &dyn Transport, @@ -157,7 +179,7 @@ fn push_matched_artifact( let remote_dir = REMOTE_STAGING; transport - .run(&format!("mkdir -p {remote_dir}")) + .run(&staging_dir_command(remote_dir)) .context("failed to make a staging directory on the remote")?; // The archive travels with its checksum and, when present, its signature, so @@ -563,6 +585,89 @@ mod tests { ); } + #[test] + #[cfg(unix)] + fn the_remote_staging_directory_is_restricted_to_its_owner() { + // The remote counterpart of `create_restricted_dir`'s 0700: the archive, + // checksum, and signature land here before any of them is verified. + // + // Asked of a real shell for the resulting *mode*, not of the string: a + // `contains("umask 077")` assertion would keep passing through the two + // ways this actually goes wrong — `mkdir -m` semantics, and a directory + // a previous run left permissive. + use std::os::unix::fs::PermissionsExt; + + let root = std::env::temp_dir().join(format!( + "rocm-remote-staging-mode-{}-{}", + std::process::id(), + rocm_core::unix_time_millis() + )); + let _ = std::fs::remove_dir_all(&root); + let fresh = root.join("fresh/provision"); + let stale = root.join("stale/provision"); + + // The second case has to start wrong to prove the `chmod` repairs it — + // 0755 is what a default umask would have left behind. + std::fs::create_dir_all(&stale).expect("the stale dir"); + std::fs::set_permissions(&stale, std::fs::Permissions::from_mode(0o755)) + .expect("loosen the stale dir"); + + // `parent_is_ours` marks the case where `mkdir -p` creates the parent + // itself. In the stale case this test pre-created it, so its mode + // reflects this process's umask rather than anything the command did. + for (directory, parent_is_ours) in [(&fresh, true), (&stale, false)] { + let status = std::process::Command::new("sh") + .arg("-c") + // A permissive umask on the invoking side, so a pass here means + // the command set the mode rather than inheriting a lucky one. + .arg(format!( + "umask 022 && {}", + staging_dir_command(&directory.to_string_lossy()) + )) + .status() + .expect("sh should run"); + assert!( + status.success(), + "the staging command failed for {directory:?}" + ); + + let mode = std::fs::metadata(directory) + .expect("the staging dir exists") + .permissions() + .mode() + & 0o777; + assert_eq!( + mode, 0o700, + "{directory:?} is mode {mode:o}; another user on the remote could \ + read or race the artifacts staged there before they are verified" + ); + + // The parent, and this is the half that makes the `umask` testable. + // `chmod` names the leaf only, so it sets the leaf whether or not + // the `umask` is there — asserting the leaf alone passes with + // `umask 077` deleted, which is precisely the branch the doc comment + // calls load-bearing. A parent that `mkdir -p` created is covered by + // the umask and by nothing else, so under the harness's `umask 022` + // it is 0700 with the umask and 0755 without. + if parent_is_ours { + let parent = directory.parent().expect("the staging dir has a parent"); + let parent_mode = std::fs::metadata(parent) + .expect("the parent exists") + .permissions() + .mode() + & 0o777; + assert_eq!( + parent_mode, 0o700, + "{parent:?} is mode {parent_mode:o}; the umask is not covering \ + directory creation, so the staged artifacts are readable for \ + the window between `mkdir` and `chmod`" + ); + } + } + + let _ = std::fs::remove_dir_all(&root); + } + #[test] fn a_hostile_asset_name_is_still_quoted_into_the_install_command() { // The prefix is unquoted so it can expand; the asset name is not, and diff --git a/apps/rocm/src/remote/publish.rs b/apps/rocm/src/remote/publish.rs index 87756e885..9610ab293 100644 --- a/apps/rocm/src/remote/publish.rs +++ b/apps/rocm/src/remote/publish.rs @@ -446,6 +446,14 @@ pub(crate) fn withdraw( // port, not a forward, so it would happily tear down whatever is on that // port — including something another tool or another person put there after // our session was recorded. + // + // Known and accepted limitation: this is a check, and the `off` below is the + // act, so a third party republishing onto this port in between still gets + // torn down and reported `Ok`. Re-reading the state immediately before the + // exec would only narrow that window, not close it — `tailscale` offers no + // compare-and-swap on a port, so nothing short of a daemon holding the claim + // makes check-and-act atomic. Narrowing it is not worth the extra round trip + // and the false impression of safety it would give the next reader. match publish_state(transport, tailnet_port, remote_port)? { PublishState::Published => {} // Already gone. Nothing to do, and nothing to complain about: teardown diff --git a/apps/rocm/src/remote/transport.rs b/apps/rocm/src/remote/transport.rs index 3497f165c..4d2a00368 100644 --- a/apps/rocm/src/remote/transport.rs +++ b/apps/rocm/src/remote/transport.rs @@ -453,7 +453,15 @@ fn run_with_piped_io( // `remote_serve_command` puts `IFS= read -r` first in the // remote command, so a broken pipe means the read never // finished, which means the command cannot have exited 0. - // Reorder that command and this demotion needs rechecking. + // + // That ordering is enforced by the shell rather than by prose: + // the caller joins the read, the export, and `serve` with `&&`, + // so a failed read short-circuits and the compound's status + // reflects it. Under `;` the status would be whatever `serve` + // returned and this demotion would be unsound. A reorder that + // breaks the coupling fails + // `only_a_complete_key_lets_the_command_reach_serve` in + // `remote::mod`, which asks a real shell for the exit status. let explained_by_the_command = error.kind() == std::io::ErrorKind::BrokenPipe && !output.status.success(); if !explained_by_the_command { From cc276ef800f760eb07123c22980747c7065a97da Mon Sep 17 00:00:00 2001 From: Eugene Volen Date: Fri, 25 Sep 2026 12:21:31 +0000 Subject: [PATCH 5/5] fix(serve): refuse to reuse an unauthenticated service for a keyed serve `spawn_managed_engine_child` treats a live managed service for the same engine and model as satisfying the request. That early return happens before `record.requires_api_key` is set and before the endpoint key is checked, so `--require-api-key` was accepted and discarded: the caller got a server that never had a key, and no error. `rocm remote serve` is what makes that dangerous. It publishes the port onto the tailnet and prints the key it just minted under "the API key above is what stops anyone else calling it". Reusing an unauthenticated service made that sentence false about an endpoint every machine on the tailnet can reach. Refuse instead of upgrading the record. The engine reads its key once, at launch, so a server already running without one keeps serving anonymously whatever is written afterwards; flipping the flag and writing a key file would leave the record claiming auth that nothing enforces, which is worse than the original bug. The refusal names the service to stop, and is the same shape as the recipe-mismatch refusal beside it. Reachability was narrow: the remote path defaults to a different port than local serving, and publish discovers by exact port, so an all-defaults collision already failed closed. It needed a service already serving that model on the requested port. Signed-off-by: Eugene Volen --- apps/rocm/src/main.rs | 113 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 113 insertions(+) diff --git a/apps/rocm/src/main.rs b/apps/rocm/src/main.rs index 6c6f25439..bdef760e4 100644 --- a/apps/rocm/src/main.rs +++ b/apps/rocm/src/main.rs @@ -6302,6 +6302,32 @@ fn spawn_managed_engine_child( resolve.canonical_model_id ); } + // Reuse cannot satisfy a demand for auth the running server never got. + // The engine reads its key once, at launch, from the environment this + // function builds below — so a server started without one keeps serving + // anonymously no matter what is written afterwards. Upgrading the record + // here would be worse than doing nothing: the record would claim auth + // that the live process does not enforce, and + // `ensure_public_service_has_endpoint_key` would pass on the strength of + // a key file nothing reads. + // + // This is what `rocm remote serve` relies on. It publishes a loopback + // port onto the tailnet and prints "the API key above is what stops + // anyone else calling it". Reusing an unauthenticated service silently + // would make that sentence false about an endpoint the whole tailnet can + // reach. Refusing is the only answer that fails closed, and it is the + // same shape as the recipe mismatch above. + if require_api_key && !existing.requires_api_key { + bail!( + "managed service `{}` is already running for engine `{engine}` and model `{}` \ + without authentication, and a running server cannot be given a key it did not \ + start with; stop it with `rocm services stop {}` and run the command again to \ + serve it with `--require-api-key`", + existing.service_id, + resolve.canonical_model_id, + existing.service_id + ); + } record_cli_audit_event( paths, "service", @@ -28433,6 +28459,93 @@ install therock"; Ok(()) } + #[test] + fn spawn_managed_engine_child_refuses_to_reuse_an_unauthenticated_service() -> Result<()> { + // `--require-api-key` used to be accepted and dropped on this path. The + // reuse branch returns before the flag is recorded and before the + // endpoint key is checked, so a caller demanding auth got a server that + // never had it, with no error. `rocm remote serve` then published that + // endpoint onto the tailnet and printed a freshly minted key under "the + // API key above is what stops anyone else calling it" — a false + // assurance about an endpoint the whole tailnet can reach. + // + // Asserted through `spawn_managed_engine_child` rather than against the + // guard's own arguments: the defect was the early return, so only the + // real call site can fail for it. + let (root, paths) = test_paths("dup-managed-unauthenticated-reuse"); + paths.ensure()?; + let mut existing = ManagedServiceRecord::new( + &paths, + "lemonade-qwen-3000", + "lemonade", + "qwen", + "qwen-canonical", + "127.0.0.1", + 11520, + "managed", + std::process::id(), + None, + None, + None, + ); + existing.status = "ready".to_owned(); + existing.engine_pid = Some(std::process::id()); + // The state that matters: live, matching, and serving without auth. + existing.requires_api_key = false; + existing.write()?; + + let resolve = ResolveModelResponse { + canonical_model_id: "qwen-canonical".to_owned(), + task: "chat".to_owned(), + source: "hf".to_owned(), + revision: "main".to_owned(), + loader: "llama.cpp".to_owned(), + trust_remote_code: false, + chat_template_mode: "auto".to_owned(), + dtype: "auto".to_owned(), + device_policy: DevicePolicy::GpuPreferred, + estimated_memory: "unknown".to_owned(), + launch_defaults: serde_json::json!({}), + engine_recipe: None, + warnings: Vec::new(), + }; + + let result = spawn_managed_engine_child( + &paths, + "lemonade", + "lemonade-qwen-3001", + "qwen", + &resolve, + "127.0.0.1", + 11520, + &resolve.device_policy, + &[], + None, + None, + None, + // The demand that used to be silently discarded. + true, + ); + let _ = fs::remove_dir_all(root); + + let Err(error) = result else { + panic!( + "reusing an unauthenticated service must not satisfy `--require-api-key`; a \ + satisfied reuse leaves the endpoint open while the caller is told it is not" + ) + }; + let message = error.to_string(); + assert!( + message.contains("without authentication"), + "the refusal must say why it refused: {message}" + ); + assert!( + message.contains("rocm services stop lemonade-qwen-3000"), + "the refusal must name the way out, with the service to stop: {message}" + ); + Ok(()) + } + #[test] fn a_managed_spawn_refuses_an_invalid_key_file_on_a_service_that_requires_one() -> Result<()> { // Drives the real call site, not the guard's own arguments. The service