diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 2a1489b78..65f570109 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,58 @@ 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 + # 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 + + # 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 +1326,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/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/apps/rocm/src/main.rs b/apps/rocm/src/main.rs index 8b9e399bd..bdef760e4 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())?; @@ -6238,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", @@ -6276,6 +6366,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 +6414,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 +6500,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 +6517,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 +6815,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 +6832,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 +7083,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 +13713,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 +17162,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 +17551,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 +19821,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 +21459,7 @@ fn treat_as_natural_language(args: &[String]) -> bool { "comfyui", "comfy", "services", + "remote", "automations", "config", "logs", @@ -25261,6 +25476,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 +25497,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 +27199,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 +28440,7 @@ install therock"; None, None, Some(&requested_recipe), + false, ); let _ = fs::remove_dir_all(root); @@ -28077,6 +28459,163 @@ 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 + // 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 +28727,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 +28959,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 +29002,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 +29011,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 +29025,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 +29068,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 +29162,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..3c366035a --- /dev/null +++ b/apps/rocm/src/remote/mod.rs @@ -0,0 +1,2336 @@ +// 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); + // 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 + // 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. +/// +/// 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 && \ + {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}"); + } + + /// 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( + "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, + }); + 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.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); + } + + /// 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..239253dcd --- /dev/null +++ b/apps/rocm/src/remote/provision.rs @@ -0,0 +1,834 @@ +// 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); + } +} + +/// 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, + 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(&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 + // 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] + #[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 + // 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..9610ab293 --- /dev/null +++ b/apps/rocm/src/remote/publish.rs @@ -0,0 +1,950 @@ +// 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. + // + // 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 + // 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..4d2a00368 --- /dev/null +++ b/apps/rocm/src/remote/transport.rs @@ -0,0 +1,926 @@ +// 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. + // + // 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 { + 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/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 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/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 1c850f591..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 { .. } )); } @@ -884,13 +1121,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 +1135,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 +1181,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 +1197,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 +1212,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"); }; @@ -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