From 05a256c76f2ae794f2a694e6ce59de6f8a8e44ad Mon Sep 17 00:00:00 2001 From: Juho Vainio Date: Fri, 16 Oct 2026 14:34:19 +0300 Subject: [PATCH 01/22] ROCMAI-554: add engines/aim Docker-based AIM serving engine Adds a new opt-in rocm-cli serving engine (--engine aim) that runs AMD Enterprise AI AIM pre-built Docker containers, alongside the existing native vLLM and Lemonade engines. Unlike vLLM/Lemonade, AIM tracks a Docker container name as its service identity rather than a supervised child PID: `docker run -d --name rocm-aim-` detaches immediately, so there is no long-lived client process to signal. Stop uses `docker stop`/`docker kill` daemon-side, sidestepping the SIGKILL-orphans-the-container risk found during the ROCMAI-552 spike. GPU passthrough uses raw device flags (--device=/dev/kfd --device=/dev/dri --group-add --ipc=host) confirmed on gfx1151 hardware; no AMD Container Toolkit needed. Registry metadata (accelerator family/type, supported gfx targets, model labels) is fetched via two plain Docker Hub v2 API calls without pulling the image, per the ROCMAI-553 spike. Also wires "aim" into the small set of hardcoded cross-crate registration points: CLI engine dispatch and inventory listings in apps/rocm, the dash daemon's per-engine metrics parser selection, the dash TUI's OpenAI-compatible engine list, generation-default flag translation in serve_cmd, and the xtask crate-dependency layering allowlist. Signed-off-by: Juho Vainio --- Cargo.lock | 14 + Cargo.toml | 1 + apps/rocm/Cargo.toml | 2 + apps/rocm/src/engines_cmd.rs | 2 +- apps/rocm/src/main.rs | 26 +- apps/rocm/src/serve_cmd.rs | 3 +- .../src/engine_registry.rs | 15 +- crates/rocm-dash-tui/src/app/chat.rs | 2 +- engines/aim/Cargo.toml | 27 ++ engines/aim/src/docker.rs | 205 +++++++++++ engines/aim/src/install.rs | 38 ++ engines/aim/src/lib.rs | 347 ++++++++++++++++++ engines/aim/src/main.rs | 7 + engines/aim/src/process.rs | 347 ++++++++++++++++++ engines/aim/src/registry.rs | 185 ++++++++++ engines/aim/src/state.rs | 248 +++++++++++++ xtask/src/crate_edges.rs | 3 + 17 files changed, 1463 insertions(+), 9 deletions(-) create mode 100644 engines/aim/Cargo.toml create mode 100644 engines/aim/src/docker.rs create mode 100644 engines/aim/src/install.rs create mode 100644 engines/aim/src/lib.rs create mode 100644 engines/aim/src/main.rs create mode 100644 engines/aim/src/process.rs create mode 100644 engines/aim/src/registry.rs create mode 100644 engines/aim/src/state.rs diff --git a/Cargo.lock b/Cargo.lock index a46a744a0..59882919b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3660,6 +3660,7 @@ dependencies = [ "rocm-dash-daemon", "rocm-dash-tui", "rocm-deps", + "rocm-engine-aim", "rocm-engine-lemonade", "rocm-engine-protocol", "rocm-engine-vllm", @@ -3796,6 +3797,19 @@ dependencies = [ "toml", ] +[[package]] +name = "rocm-engine-aim" +version = "0.1.0" +dependencies = [ + "anyhow", + "clap", + "rocm-core", + "rocm-engine-protocol", + "serde", + "serde_json", + "ureq", +] + [[package]] name = "rocm-engine-lemonade" version = "0.1.0" diff --git a/Cargo.toml b/Cargo.toml index 4583977ea..a8a3d7705 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -12,6 +12,7 @@ members = [ "crates/rocm-dash-daemon", "crates/rocm-dash-tui", "crates/e2e-report", + "engines/aim", "engines/lemonade", "engines/vllm", "xtask", diff --git a/apps/rocm/Cargo.toml b/apps/rocm/Cargo.toml index a319b11c6..5db1f4e8d 100644 --- a/apps/rocm/Cargo.toml +++ b/apps/rocm/Cargo.toml @@ -15,6 +15,7 @@ e2e-test-hooks = [ "rocm-core/e2e-test-hooks", "rocm-engine-lemonade/e2e-test-hooks", "rocm-engine-vllm/e2e-test-hooks", + "rocm-engine-aim/e2e-test-hooks", ] [dependencies] @@ -39,6 +40,7 @@ rocm-engine-lemonade = { path = "../../engines/lemonade" } rocm-engine-protocol = { path = "../../crates/rocm-engine-protocol" } rocmd = { path = "../rocmd" } rocm-engine-vllm = { path = "../../engines/vllm" } +rocm-engine-aim = { path = "../../engines/aim" } rpassword.workspace = true serde.workspace = true serde_json.workspace = true diff --git a/apps/rocm/src/engines_cmd.rs b/apps/rocm/src/engines_cmd.rs index 44ae5c9d6..5a0c6981c 100644 --- a/apps/rocm/src/engines_cmd.rs +++ b/apps/rocm/src/engines_cmd.rs @@ -178,7 +178,7 @@ fn resolve_engine_install_runtime_id( } pub(crate) fn engine_manages_own_runtime(engine: &str) -> bool { - engine == "lemonade" + engine == "lemonade" || engine == "aim" } fn env_root_for_runtime( diff --git a/apps/rocm/src/main.rs b/apps/rocm/src/main.rs index 37d4c6e42..50a8a6ef6 100644 --- a/apps/rocm/src/main.rs +++ b/apps/rocm/src/main.rs @@ -131,7 +131,7 @@ static BUILTIN_ENGINE_ENV_LOCK: OnceLock> = OnceLock::new(); /// User-selectable serving engines, in the order shown in `--help` and shell /// completions. Used as the clap `value_parser` for every user-facing `--engine` /// argument so the possible values stay in sync across the CLI surface. -const SUPPORTED_ENGINES: [&str; 2] = ["lemonade", "vllm"]; +const SUPPORTED_ENGINES: [&str; 3] = ["lemonade", "vllm", "aim"]; #[derive(Parser, Debug)] #[command( @@ -2764,6 +2764,10 @@ const fn builtin_engine_inventory() -> &'static [(&'static str, &'static str)] { "vllm", "Linux/WSL ROCm GPU serving engine through external vLLM", ), + ( + "aim", + "AMD Enterprise AI pre-built Docker serving containers", + ), ] } @@ -9965,6 +9969,10 @@ pub(crate) const fn engine_inventory() -> &'static [(&'static str, &'static str) "vllm", "Linux/WSL ROCm GPU serving engine through external vLLM", ), + ( + "aim", + "AMD Enterprise AI pre-built Docker serving containers", + ), ] } @@ -10258,6 +10266,7 @@ fn builtin_engine_request( envelope.clone(), )), "vllm" => Some(rocm_engine_vllm::builtin_handle_envelope(envelope.clone())), + "aim" => Some(rocm_engine_aim::builtin_handle_envelope(envelope.clone())), _ => None, } } @@ -10276,7 +10285,7 @@ fn run_builtin_engine_stdio(engine: &str) -> Result<()> { } fn builtin_engine_available(engine: &str) -> bool { - matches!(engine, "lemonade" | "vllm") + matches!(engine, "lemonade" | "vllm" | "aim") } #[allow(clippy::too_many_arguments)] @@ -10322,6 +10331,19 @@ fn run_builtin_engine_serve_http( log_path, engine_recipe, ), + "aim" => rocm_engine_aim::builtin_serve_http( + service_id, + model_ref, + host, + port, + parsed_policy, + gpu_indices, + runtime_id, + env_id, + state_path, + log_path, + engine_recipe, + ), other => bail!("engine `{other}` is not built into this rocm binary"), } } diff --git a/apps/rocm/src/serve_cmd.rs b/apps/rocm/src/serve_cmd.rs index 275115f2a..b2d823008 100644 --- a/apps/rocm/src/serve_cmd.rs +++ b/apps/rocm/src/serve_cmd.rs @@ -275,7 +275,8 @@ fn engine_recipe_with_generation_defaults( engine: engine.to_owned(), ..EngineRecipeHint::default() }); - if engine.eq_ignore_ascii_case("vllm") { + if engine.eq_ignore_ascii_case("vllm") || engine.eq_ignore_ascii_case("aim") { + // AIM images bundle vLLM, so its `--override-generation-config` flag applies unchanged. let mut overrides = serde_json::Map::new(); if let Some(temperature) = defaults.temperature { overrides.insert("temperature".to_owned(), serde_json::json!(temperature)); diff --git a/crates/rocm-dash-collectors/src/engine_registry.rs b/crates/rocm-dash-collectors/src/engine_registry.rs index 208e75af3..4b50e7d52 100644 --- a/crates/rocm-dash-collectors/src/engine_registry.rs +++ b/crates/rocm-dash-collectors/src/engine_registry.rs @@ -19,6 +19,10 @@ pub enum EngineKind { Vllm, /// Lemonade Server — JSON `/api/v1/stats`. Lemonade, + /// AMD Enterprise AI AIM — bundles vLLM, so its `/metrics` surface is + /// presumed identical; unconfirmed against a real container until the + /// Strix Halo validation pass. + Aim, } impl EngineKind { @@ -27,6 +31,7 @@ impl EngineKind { match self { Self::Vllm => "vllm", Self::Lemonade => "lemonade", + Self::Aim => "aim", } } @@ -36,7 +41,7 @@ impl EngineKind { /// unmanaged/external discovery where no registry record exists. pub const fn default_port(self) -> u16 { match self { - Self::Vllm => 8000, + Self::Vllm | Self::Aim => 8000, Self::Lemonade => crate::lemonade::LEMONADE_PORT, // 13305 } } @@ -46,16 +51,18 @@ impl EngineKind { match label.trim().to_ascii_lowercase().as_str() { "vllm" => Some(Self::Vllm), "lemonade" => Some(Self::Lemonade), + "aim" => Some(Self::Aim), _ => None, } } /// Parse this engine's raw scrape body into an [`InstanceSample`] using the /// engine-appropriate parser. vLLM dispatches to the **unchanged** - /// `vllm_prom::parse`; Lemonade to `lemonade::parse_stats`. + /// `vllm_prom::parse`; Lemonade to `lemonade::parse_stats`; AIM reuses the + /// vLLM parser since AIM images bundle vLLM. pub fn parse_sample(self, body: &str) -> InstanceSample { match self { - Self::Vllm => crate::vllm_prom::parse(body), + Self::Vllm | Self::Aim => crate::vllm_prom::parse(body), Self::Lemonade => crate::lemonade::parse_stats(body), } } @@ -67,7 +74,7 @@ mod tests { #[test] fn engine_kind_labels_and_ports_and_from_label_roundtrip() { - for kind in [EngineKind::Vllm, EngineKind::Lemonade] { + for kind in [EngineKind::Vllm, EngineKind::Lemonade, EngineKind::Aim] { assert_eq!(EngineKind::from_label(kind.label()), Some(kind)); } assert_eq!(EngineKind::Lemonade.default_port(), 13305); diff --git a/crates/rocm-dash-tui/src/app/chat.rs b/crates/rocm-dash-tui/src/app/chat.rs index fc0224b07..c3d01593c 100644 --- a/crates/rocm-dash-tui/src/app/chat.rs +++ b/crates/rocm-dash-tui/src/app/chat.rs @@ -127,7 +127,7 @@ pub(super) fn build_local_agent( /// Local engines that expose an OpenAI-compatible `/v1` surface the dash chat /// can talk to directly. A managed service running one of these is a valid /// auto-detected chat endpoint regardless of which port it bound. -const OPENAI_COMPATIBLE_ENGINES: &[&str] = &["vllm", "lemonade"]; +const OPENAI_COMPATIBLE_ENGINES: &[&str] = &["vllm", "lemonade", "aim"]; /// A managed-service endpoint the dash chat can route to, picked from the /// read-only `services` tool payload. diff --git a/engines/aim/Cargo.toml b/engines/aim/Cargo.toml new file mode 100644 index 000000000..39a87bdd3 --- /dev/null +++ b/engines/aim/Cargo.toml @@ -0,0 +1,27 @@ +[package] +name = "rocm-engine-aim" +version.workspace = true +edition.workspace = true +license.workspace = true +repository.workspace = true +rust-version.workspace = true +publish.workspace = true + +[lints] +workspace = true + +[features] +e2e-test-hooks = ["rocm-core/e2e-test-hooks"] + +[[bin]] +name = "rocm-engine-aim" +path = "src/main.rs" + +[dependencies] +anyhow.workspace = true +clap.workspace = true +rocm-core = { path = "../../crates/rocm-core" } +rocm-engine-protocol = { path = "../../crates/rocm-engine-protocol" } +serde.workspace = true +serde_json.workspace = true +ureq = { version = "2.12", features = ["json", "native-certs"] } diff --git a/engines/aim/src/docker.rs b/engines/aim/src/docker.rs new file mode 100644 index 000000000..679b46368 --- /dev/null +++ b/engines/aim/src/docker.rs @@ -0,0 +1,205 @@ +// Copyright © Advanced Micro Devices, Inc., or its affiliates. +// +// SPDX-License-Identifier: MIT + +//! Thin wrapper around the `docker` CLI binary. No `bollard`/Docker API +//! dependency: shelling out mirrors the pattern `engines/lemonade/src/process.rs` +//! already uses for its own CLI (`run_lemonade_pull`/`run_lemonade_backend_install`). + +use anyhow::{Context, Result, bail}; +use std::process::{Command, Stdio}; +use std::time::Duration; + +pub(crate) const DOCKER_BIN: &str = "docker"; + +#[allow(clippy::too_many_arguments)] +pub(crate) fn run( + name: &str, + image: &str, + host: &str, + port: u16, + container_port: u16, + gpu_indices: &[u32], + render_gid: Option, +) -> Result<()> { + let mut command = Command::new(DOCKER_BIN); + command + .arg("run") + .arg("-d") + .arg("--name") + .arg(name) + .arg("-p") + .arg(format!("{host}:{port}:{container_port}")); + + if !gpu_indices.is_empty() { + if let Some(csv) = rocm_engine_protocol::gpu_indices_to_csv(gpu_indices) { + command.arg("-e").arg(format!("HIP_VISIBLE_DEVICES={csv}")); + } + command + .arg("--device=/dev/kfd") + .arg("--device=/dev/dri") + .arg("--ipc=host"); + if let Some(gid) = render_gid { + command.arg("--group-add").arg(gid.to_string()); + } + } + + command.arg(image); + command.stdin(Stdio::null()); + + let output = command + .output() + .with_context(|| format!("failed to run `{DOCKER_BIN} run` for container `{name}`"))?; + if !output.status.success() { + bail!( + "`docker run --name {name}` failed: {}", + rocm_core::command_failure_detail(&output) + ); + } + Ok(()) +} + +/// Resolves the host's `render` group gid dynamically (host-specific, so it +/// cannot be hardcoded). Returns `None` if the group is absent (e.g. non-Linux). +pub(crate) fn render_group_gid() -> Option { + let output = Command::new("getent") + .arg("group") + .arg("render") + .output() + .ok()?; + if !output.status.success() { + return None; + } + let text = String::from_utf8_lossy(&output.stdout); + // `getent group render` prints `render:x::`. + text.trim().split(':').nth(2)?.parse().ok() +} + +pub(crate) fn stop(name: &str, force: bool, grace: Duration) -> Result { + let mut command = Command::new(DOCKER_BIN); + if force { + command.arg("kill").arg(name); + } else { + command + .arg("stop") + .arg("-t") + .arg(grace.as_secs().to_string()) + .arg(name); + } + let output = command.stdin(Stdio::null()).output().with_context(|| { + format!( + "failed to run `docker {}` for container `{name}`", + if force { "kill" } else { "stop" } + ) + })?; + if !output.status.success() { + // A container that's already gone is not a stop failure. + let detail = rocm_core::command_failure_detail(&output); + if detail.contains("No such container") { + return Ok(false); + } + bail!("failed to stop container `{name}`: {detail}"); + } + Ok(true) +} + +pub(crate) fn inspect_running(name: &str) -> Result> { + inspect_format(name, "{{.State.Running}}").map(|value| value.map(|text| text == "true")) +} + +pub(crate) fn inspect_exit_code(name: &str) -> Result> { + inspect_format(name, "{{.State.ExitCode}}") + .map(|value| value.and_then(|text| text.parse().ok())) +} + +fn inspect_format(name: &str, format: &str) -> Result> { + let output = Command::new(DOCKER_BIN) + .arg("inspect") + .arg("--format") + .arg(format) + .arg(name) + .stdin(Stdio::null()) + .output() + .with_context(|| format!("failed to run `docker inspect` for container `{name}`"))?; + if !output.status.success() { + return Ok(None); + } + Ok(Some( + String::from_utf8_lossy(&output.stdout).trim().to_owned(), + )) +} + +pub(crate) fn logs(name: &str, tail: usize) -> Result> { + let output = Command::new(DOCKER_BIN) + .arg("logs") + .arg("--tail") + .arg(tail.to_string()) + .arg(name) + .stdin(Stdio::null()) + .output() + .with_context(|| format!("failed to run `docker logs` for container `{name}`"))?; + if !output.status.success() { + bail!( + "`docker logs {name}` failed: {}", + rocm_core::command_failure_detail(&output) + ); + } + let mut lines: Vec = String::from_utf8_lossy(&output.stdout) + .lines() + .map(ToOwned::to_owned) + .collect(); + lines.extend( + String::from_utf8_lossy(&output.stderr) + .lines() + .map(ToOwned::to_owned), + ); + Ok(lines) +} + +pub(crate) fn image_present(image: &str) -> Result { + let status = Command::new(DOCKER_BIN) + .arg("image") + .arg("inspect") + .arg(image) + .stdin(Stdio::null()) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .status() + .with_context(|| format!("failed to run `docker image inspect {image}`"))?; + Ok(status.success()) +} + +pub(crate) fn pull(image: &str) -> Result<()> { + let status = Command::new(DOCKER_BIN) + .arg("pull") + .arg(image) + .stdin(Stdio::null()) + .stdout(Stdio::inherit()) + .stderr(Stdio::inherit()) + .status() + .with_context(|| format!("failed to run `docker pull {image}`"))?; + if !status.success() { + bail!("`docker pull {image}` failed with status {status}"); + } + Ok(()) +} + +pub(crate) fn daemon_reachable() -> bool { + Command::new(DOCKER_BIN) + .arg("info") + .stdin(Stdio::null()) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .status() + .is_ok_and(|status| status.success()) +} + +#[cfg(test)] +mod tests { + #[test] + fn render_group_gid_parses_getent_format() { + let line = "render:x:993:"; + let gid: Option = line.trim().split(':').nth(2).and_then(|s| s.parse().ok()); + assert_eq!(gid, Some(993)); + } +} diff --git a/engines/aim/src/install.rs b/engines/aim/src/install.rs new file mode 100644 index 000000000..3d6b797ed --- /dev/null +++ b/engines/aim/src/install.rs @@ -0,0 +1,38 @@ +// Copyright © Advanced Micro Devices, Inc., or its affiliates. +// +// SPDX-License-Identifier: MIT + +//! Much smaller than vLLM's/Lemonade's `install.rs`: there is no venv, wheel, +//! or torch-alignment logic. AIM images are self-contained containers, so +//! "install" for this engine means "Docker is present and reachable" — the +//! actual model image is pulled lazily per-launch (see `process::launch_service`), +//! since each `model_ref` names its own image. + +use anyhow::{Result, bail}; +use rocm_engine_protocol::InstallRequest; +use rocm_engine_protocol::InstallResponse; + +pub(crate) fn install_response(request: InstallRequest) -> Result { + rocm_core::require_nonempty(&request.runtime_id, "runtime_id")?; + + if !crate::docker::daemon_reachable() { + bail!( + "docker daemon is not reachable; install Docker and ensure it is running before using the AIM engine" + ); + } + + Ok(InstallResponse { + env_id: request.runtime_id, + env_path: "docker".to_owned(), + python_executable: "docker".to_owned(), + runtime_kind: Some("docker".to_owned()), + runtime_executable: Some("docker".to_owned()), + managed_env: Some(false), + installed_packages: Vec::new(), + capabilities: crate::capabilities(), + lock_hash: "docker".to_owned(), + warnings: vec![ + "AIM model images are pulled on first launch, not during install".to_owned(), + ], + }) +} diff --git a/engines/aim/src/lib.rs b/engines/aim/src/lib.rs new file mode 100644 index 000000000..29a128490 --- /dev/null +++ b/engines/aim/src/lib.rs @@ -0,0 +1,347 @@ +// Copyright © Advanced Micro Devices, Inc., or its affiliates. +// +// SPDX-License-Identifier: MIT + +use anyhow::{Context, Result}; +use clap::{Parser, Subcommand}; +use rocm_core::AppPaths; +use rocm_engine_protocol::{ + DetectRequest, DetectResponse, DevicePolicy, EngineCapabilities, EngineDeviceAvailability, + EngineMethod, EngineRecipeHint, EngineRequestEnvelope, EngineResponseEnvelope, InstallRequest, + LaunchRequest, ResolveModelRequest, +}; +use serde::Serialize; +use serde::de::DeserializeOwned; +use serde_json::Value; +use std::io::{Read, Write}; +use std::path::PathBuf; + +mod docker; +mod install; +mod process; +mod registry; +mod state; + +pub(crate) const ENGINE_NAME: &str = "aim"; +pub(crate) const DEFAULT_HOST: &str = "127.0.0.1"; +/// AIM images bundle vLLM, which defaults to port 8000 inside the container. +pub(crate) const DEFAULT_AIM_PORT: u16 = 8000; + +#[derive(Parser)] +struct Cli { + #[command(subcommand)] + command: CommandKind, +} + +#[derive(Subcommand)] +enum CommandKind { + Detect, + Capabilities, + Install { + #[arg(long, default_value = "aim")] + runtime_id: String, + #[arg(long)] + reinstall: bool, + }, + ResolveModel { + #[arg(long)] + model_ref: String, + #[arg(long)] + device_policy: Option, + }, + Launch { + #[arg(long)] + service_id: String, + #[arg(long)] + model_ref: String, + #[arg(long, default_value = DEFAULT_HOST)] + host: String, + #[arg(long, default_value_t = DEFAULT_AIM_PORT)] + port: u16, + #[arg(long)] + device_policy: Option, + #[arg(long)] + gpu: Option, + #[arg(long)] + runtime_id: Option, + #[arg(long)] + env_id: Option, + }, + Stdio, + ServeHttp { + #[arg(long)] + service_id: String, + #[arg(long)] + model_ref: String, + #[arg(long, default_value = DEFAULT_HOST)] + host: String, + #[arg(long, default_value_t = DEFAULT_AIM_PORT)] + port: u16, + #[arg(long, default_value = "gpu_required")] + device_policy: String, + #[arg(long)] + gpu: Option, + #[arg(long)] + runtime_id: Option, + #[arg(long)] + env_id: Option, + #[arg(long)] + state_path: PathBuf, + #[arg(long)] + engine_recipe_json: Option, + }, +} + +pub(crate) struct ServeHttpRequest { + pub service_id: String, + pub model_ref: String, + pub host: String, + pub port: u16, + pub device_policy: DevicePolicy, + pub gpu_indices: Vec, + #[allow(dead_code)] // ponytail: accepted for protocol-shape parity; AIM has no runtime concept + pub runtime_id: Option, + #[allow(dead_code)] + // ponytail: accepted for protocol-shape parity; AIM has no managed-env concept + pub env_id: Option, + pub state_path: PathBuf, + #[allow(dead_code)] + // ponytail: no host-side log file in `docker run -d` mode; logs come from `docker logs` + pub log_path: Option, + pub engine_recipe: Option, +} + +pub fn run_cli() -> Result<()> { + let cli = Cli::parse(); + match cli.command { + CommandKind::Detect => print_json(&detect_response()), + CommandKind::Capabilities => print_json(&capabilities()), + CommandKind::Install { + runtime_id, + reinstall, + } => print_json(&crate::install::install_response(InstallRequest { + runtime_id, + python_version: None, + reinstall, + env_root: None, + })?), + CommandKind::ResolveModel { + model_ref, + device_policy, + } => print_json(&crate::process::resolve_model_response( + ResolveModelRequest { + model_ref, + runtime_id: None, + device_policy: device_policy + .as_deref() + .map(crate::process::parse_device_policy_arg) + .transpose()?, + recipe_override: None, + engine_recipe: None, + }, + )?), + CommandKind::Launch { + service_id, + model_ref, + host, + port, + device_policy, + gpu, + runtime_id, + env_id, + } => print_json(&crate::process::launch_service(LaunchRequest { + service_id, + env_id, + runtime_id, + model_ref, + host, + port, + device_policy: device_policy + .as_deref() + .map(crate::process::parse_device_policy_arg) + .transpose()?, + endpoint_mode: None, + engine_recipe: None, + gpu_selection: crate::process::parse_gpu_selection_arg(gpu.as_deref())?, + })?), + CommandKind::Stdio => { + let envelope = read_request()?; + let response = handle_envelope(envelope); + print_json(&response) + } + CommandKind::ServeHttp { + service_id, + model_ref, + host, + port, + device_policy, + gpu, + runtime_id, + env_id, + state_path, + engine_recipe_json, + } => { + let gpu_indices = crate::process::parse_gpu_indices_arg(gpu.as_deref())?; + let log_path = Some( + AppPaths::discover()? + .engine_logs_dir(ENGINE_NAME) + .join(format!("{service_id}.log")), + ); + crate::process::serve_http(ServeHttpRequest { + service_id, + model_ref, + host, + port, + device_policy: crate::process::parse_device_policy_arg(&device_policy)?, + gpu_indices, + runtime_id, + env_id, + state_path, + log_path, + engine_recipe: crate::process::parse_engine_recipe_json(engine_recipe_json)?, + }) + } + } +} + +pub fn builtin_handle_envelope(envelope: EngineRequestEnvelope) -> EngineResponseEnvelope { + handle_envelope(envelope) +} + +#[allow(clippy::too_many_arguments)] +pub fn builtin_serve_http( + service_id: String, + model_ref: String, + host: String, + port: u16, + device_policy: DevicePolicy, + gpu_indices: Vec, + runtime_id: Option, + env_id: Option, + state_path: PathBuf, + log_path: Option, + engine_recipe: Option, +) -> Result<()> { + crate::process::serve_http(ServeHttpRequest { + service_id, + model_ref, + host, + port, + device_policy, + gpu_indices, + runtime_id, + env_id, + state_path, + log_path, + engine_recipe, + }) +} + +fn handle_envelope(envelope: EngineRequestEnvelope) -> EngineResponseEnvelope { + match envelope.method { + EngineMethod::Detect => { + deserialize_and_respond::(envelope.payload, |_request| { + Ok(detect_response()) + }) + } + EngineMethod::Capabilities => EngineResponseEnvelope::success(capabilities()), + EngineMethod::Install => { + deserialize_and_respond(envelope.payload, crate::install::install_response) + } + EngineMethod::ResolveModel => { + deserialize_and_respond(envelope.payload, crate::process::resolve_model_response) + } + EngineMethod::Launch => { + deserialize_and_respond(envelope.payload, crate::process::launch_service) + } + EngineMethod::Healthcheck => { + deserialize_and_respond(envelope.payload, crate::state::healthcheck_service) + } + EngineMethod::Endpoint => { + deserialize_and_respond(envelope.payload, crate::state::endpoint_response) + } + EngineMethod::Stop => { + deserialize_and_respond(envelope.payload, crate::process::stop_service) + } + EngineMethod::Logs => { + deserialize_and_respond(envelope.payload, crate::state::logs_response) + } + } +} + +fn deserialize_and_respond(payload: Value, handler: F) -> EngineResponseEnvelope +where + T: DeserializeOwned, + F: FnOnce(T) -> Result, + U: Serialize, +{ + let request: T = match serde_json::from_value(payload) { + Ok(request) => request, + Err(error) => { + return EngineResponseEnvelope::failure("invalid_payload", error.to_string()); + } + }; + match handler(request) { + Ok(response) => EngineResponseEnvelope::success(response), + Err(error) => EngineResponseEnvelope::failure("request_failed", error.to_string()), + } +} + +fn detect_response() -> DetectResponse { + let reachable = crate::docker::daemon_reachable(); + let mut notes = Vec::new(); + if !reachable { + notes.push( + "docker daemon is not reachable; install Docker and ensure it is running to use the AIM engine" + .to_owned(), + ); + } + DetectResponse { + installed: reachable, + env_id: None, + runtime_kind: Some("docker".to_owned()), + runtime_executable: Some("docker".to_owned()), + managed_env: Some(false), + python_version: None, + torch_version: None, + transformers_version: None, + available_devices: vec![EngineDeviceAvailability { + kind: "rocm_gpu".to_owned(), + available: reachable, + reason: if reachable { + None + } else { + Some("docker daemon unreachable".to_owned()) + }, + }], + capabilities: capabilities(), + notes, + } +} + +pub(crate) fn capabilities() -> EngineCapabilities { + EngineCapabilities { + cpu: true, + rocm_gpu: true, + openai_compatible: true, + tool_calling: false, + quantized_models: "image-reported".to_owned(), + reasoning_parser: false, + } +} + +fn read_request() -> Result { + let mut input = String::new(); + std::io::stdin() + .read_to_string(&mut input) + .context("failed to read engine stdio request")?; + serde_json::from_str(&input).context("failed to parse engine stdio request") +} + +fn print_json(value: &T) -> Result<()> { + let stdout = std::io::stdout(); + let mut handle = stdout.lock(); + serde_json::to_writer_pretty(&mut handle, value)?; + writeln!(handle)?; + Ok(()) +} diff --git a/engines/aim/src/main.rs b/engines/aim/src/main.rs new file mode 100644 index 000000000..90c18b91f --- /dev/null +++ b/engines/aim/src/main.rs @@ -0,0 +1,7 @@ +// Copyright © Advanced Micro Devices, Inc., or its affiliates. +// +// SPDX-License-Identifier: MIT + +fn main() -> anyhow::Result<()> { + rocm_engine_aim::run_cli() +} diff --git a/engines/aim/src/process.rs b/engines/aim/src/process.rs new file mode 100644 index 000000000..891cffb89 --- /dev/null +++ b/engines/aim/src/process.rs @@ -0,0 +1,347 @@ +// Copyright © Advanced Micro Devices, Inc., or its affiliates. +// +// SPDX-License-Identifier: MIT + +use anyhow::{Context, Result, bail}; +use rocm_core::{AppPaths, require_nonempty}; +use rocm_engine_protocol::{ + DevicePolicy, ENGINE_RECIPE_CONTRACT_VERSION, EngineRecipeHint, GpuSelection, LaunchRequest, + LaunchResponse, ResolveModelRequest, ResolveModelResponse, StopRequest, StopResponse, +}; +use serde_json::json; +use std::time::{Duration, Instant}; + +use crate::ServeHttpRequest; + +const STOP_GRACE: Duration = Duration::from_secs(10); +const READY_POLL_INTERVAL: Duration = Duration::from_millis(500); +const DEFAULT_READY_TIMEOUT: Duration = Duration::from_mins(5); + +/// Deterministic, Docker-name-sanitized container identity for a service. +/// This *is* the service's identity for the AIM engine — no PID/start-ticks +/// tracking, since `docker run -d` leaves no long-lived client process to +/// supervise (see plan: SIGKILL to a foreground `docker run` client orphans +/// the container, so this engine never signals a client at all). +pub(crate) fn container_name(service_id: &str) -> String { + let sanitized: String = service_id + .chars() + .map(|c| { + if c.is_ascii_alphanumeric() || c == '_' || c == '.' || c == '-' { + c + } else { + '-' + } + }) + .collect(); + format!("rocm-aim-{sanitized}") +} + +pub(crate) fn resolve_model_response(request: ResolveModelRequest) -> Result { + require_nonempty(&request.model_ref, "model_ref")?; + let engine_recipe = accepted_engine_recipe(request.engine_recipe)?; + + let metadata = crate::registry::fetch_image_metadata(&request.model_ref)?; + let mut warnings = Vec::new(); + + let device_policy = match metadata.accelerator_type.as_deref() { + Some("gpu") => DevicePolicy::GpuRequired, + Some("cpu") => DevicePolicy::CpuOnly, + _ => { + warnings.push( + "AIM image did not report AIM_ACCELERATOR_TYPE; assuming GPU-preferred".to_owned(), + ); + DevicePolicy::GpuPreferred + } + }; + + if device_policy != DevicePolicy::CpuOnly + && let Some(host_gfx) = rocm_core::detect_host_gfx_target() + && !metadata.pytorch_rocm_arch.is_empty() + && !metadata + .pytorch_rocm_arch + .iter() + .any(|arch| arch == &host_gfx) + { + warnings.push(format!( + "AIM image targets {:?} but this host reports {host_gfx}; it may not run", + metadata.pytorch_rocm_arch + )); + } + + Ok(ResolveModelResponse { + canonical_model_id: metadata + .canonical_name + .unwrap_or_else(|| request.model_ref.clone()), + task: "text-generation".to_owned(), + source: "docker_registry".to_owned(), + revision: "main".to_owned(), + loader: "aim".to_owned(), + trust_remote_code: false, + chat_template_mode: "engine_default".to_owned(), + dtype: "auto".to_owned(), + device_policy, + estimated_memory: "engine-reported".to_owned(), + launch_defaults: json!({ + "endpoint_mode": "openai", + "host": crate::DEFAULT_HOST, + "port": crate::DEFAULT_AIM_PORT, + }), + engine_recipe, + warnings, + }) +} + +fn accepted_engine_recipe( + engine_recipe: Option, +) -> Result> { + let Some(hint) = engine_recipe else { + return Ok(None); + }; + if hint.engine != crate::ENGINE_NAME || hint.contract_version != ENGINE_RECIPE_CONTRACT_VERSION + { + bail!( + "engine recipe hint is not valid for `{}` (contract {ENGINE_RECIPE_CONTRACT_VERSION})", + crate::ENGINE_NAME + ); + } + Ok(Some(hint)) +} + +pub(crate) fn parse_engine_recipe_json(value: Option) -> Result> { + let Some(raw) = value else { + return Ok(None); + }; + let hint: EngineRecipeHint = + serde_json::from_str(&raw).context("failed to parse engine recipe JSON")?; + accepted_engine_recipe(Some(hint)) +} + +pub(crate) fn launch_service(request: LaunchRequest) -> Result { + require_nonempty(&request.service_id, "service_id")?; + require_nonempty(&request.model_ref, "model_ref")?; + let engine_recipe = accepted_engine_recipe(request.engine_recipe)?; + let gpu_indices = rocm_engine_protocol::launch_gpu_indices(request.gpu_selection.as_ref()); + let device_policy = request.device_policy.unwrap_or(DevicePolicy::GpuPreferred); + + let paths = AppPaths::discover()?; + let state_path = paths + .engine_state_dir(crate::ENGINE_NAME) + .join(format!("{}.json", request.service_id)); + + run_container( + &request.service_id, + &request.model_ref, + &request.host, + request.port, + &gpu_indices, + )?; + + let endpoint_url = endpoint_url(&request.host, request.port); + crate::state::write_running_state( + &request.service_id, + &request.model_ref, + &request.host, + request.port, + device_policy, + engine_recipe.as_ref(), + )?; + + let log_path = format!("docker://{}", container_name(&request.service_id)); + Ok(LaunchResponse { + service_id: request.service_id, + // ponytail: `docker run -d` backgrounds the container itself — there is + // no long-lived client process to report a real PID for. + pid: 0, + endpoint_url, + log_path, + state_path: state_path.display().to_string(), + }) +} + +/// Starts the named service's container if it isn't already running. +/// Container names are globally unique in Docker (unlike host PIDs), so +/// unconditionally calling `docker run --name X` twice for the same service +/// fails with a name conflict — `launch_service` and `serve_http` must agree +/// on "already running? skip `docker run`" instead of vLLM's host-process +/// model, where two independent spawns are safe because each gets its own PID. +fn run_container( + service_id: &str, + model_ref: &str, + host: &str, + port: u16, + gpu_indices: &[u32], +) -> Result<()> { + let name = container_name(service_id); + if crate::docker::inspect_running(&name)?.unwrap_or(false) { + return Ok(()); + } + + if !crate::docker::image_present(model_ref)? { + crate::docker::pull(model_ref)?; + } + + let metadata = crate::registry::fetch_image_metadata(model_ref).unwrap_or_default(); + let use_gpu = metadata.accelerator_type.as_deref() != Some("cpu") && !gpu_indices.is_empty(); + let render_gid = if use_gpu { + crate::docker::render_group_gid() + } else { + None + }; + + crate::docker::run( + &name, + model_ref, + host, + port, + crate::DEFAULT_AIM_PORT, + if use_gpu { gpu_indices } else { &[] }, + render_gid, + ) +} + +pub(crate) fn serve_http(request: ServeHttpRequest) -> Result<()> { + run_container( + &request.service_id, + &request.model_ref, + &request.host, + request.port, + &request.gpu_indices, + )?; + + crate::state::write_running_state( + &request.service_id, + &request.model_ref, + &request.host, + request.port, + request.device_policy, + request.engine_recipe.as_ref(), + )?; + + let name = container_name(&request.service_id); + let timeout = DEFAULT_READY_TIMEOUT; + match wait_for_container_ready( + &name, + &request.host, + request.port, + &request.model_ref, + timeout, + ) { + Ok(()) => Ok(()), + Err(error) => { + crate::state::write_terminal_state(&request.state_path, "failed")?; + Err(error) + } + } +} + +/// Mirrors `engines/vllm/src/process.rs`'s `wait_for_vllm_ready` architecture: +/// early-exit check, monotonic timeout, listed-then-inference-verified gate. +/// `docker::inspect_running` replaces `child.try_wait()` since there is no +/// `std::process::Child` for a detached container. +fn wait_for_container_ready( + name: &str, + host: &str, + port: u16, + model_ref: &str, + timeout: Duration, +) -> Result<()> { + let start = Instant::now(); + let endpoint = endpoint_url(host, port); + loop { + if crate::docker::inspect_running(name)?.is_some_and(|running| !running) { + let exit_code = crate::docker::inspect_exit_code(name)?.unwrap_or(-1); + let tail = crate::docker::logs(name, 80).unwrap_or_default().join("\n"); + bail!("AIM container `{name}` exited early (exit code {exit_code}):\n{tail}"); + } + if start.elapsed() > timeout { + let tail = crate::docker::logs(name, 80).unwrap_or_default().join("\n"); + bail!( + "AIM container `{name}` did not become ready within {}s:\n{tail}", + timeout.as_secs() + ); + } + + let listed = + crate::state::query_loaded_model_endpoint(&endpoint, Some(model_ref)).unwrap_or(false); + if listed + && crate::state::query_inference_probe_endpoint(&endpoint, model_ref).unwrap_or(false) + { + return Ok(()); + } + + std::thread::sleep(READY_POLL_INTERVAL); + } +} + +pub(crate) fn stop_service(request: StopRequest) -> Result { + require_nonempty(&request.service_id, "service_id")?; + let name = container_name(&request.service_id); + let stopped = crate::docker::stop(&name, request.force, STOP_GRACE)?; + if stopped { + let state_path = AppPaths::discover()? + .engine_state_dir(crate::ENGINE_NAME) + .join(format!("{}.json", request.service_id)); + crate::state::write_terminal_state(&state_path, "stopped")?; + } + Ok(StopResponse { + stopped, + graceful: stopped && !request.force, + }) +} + +pub(crate) fn endpoint_url(host: &str, port: u16) -> String { + format!("{}/v1", rocm_core::format_http_base_url(host, port)) +} + +pub(crate) fn parse_device_policy_arg(policy: &str) -> Result { + match policy { + "gpu" | "gpu_required" => Ok(DevicePolicy::GpuRequired), + "gpu_preferred" => Ok(DevicePolicy::GpuPreferred), + "cpu" | "cpu_only" => Ok(DevicePolicy::CpuOnly), + other => bail!("unsupported device policy: {other}"), + } +} + +pub(crate) fn parse_gpu_selection_arg(value: Option<&str>) -> Result> { + value + .map(|raw| GpuSelection::parse_cli_value(raw).map_err(anyhow::Error::msg)) + .transpose() +} + +pub(crate) fn parse_gpu_indices_arg(value: Option<&str>) -> Result> { + Ok(rocm_engine_protocol::launch_gpu_indices( + parse_gpu_selection_arg(value)?.as_ref(), + )) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn container_name_sanitizes_unsafe_characters() { + assert_eq!(container_name("svc 1/a"), "rocm-aim-svc-1-a"); + } + + #[test] + fn container_name_preserves_safe_characters() { + assert_eq!(container_name("svc_1.a-b"), "rocm-aim-svc_1.a-b"); + } + + #[test] + fn parse_device_policy_arg_accepts_known_aliases() { + assert_eq!( + parse_device_policy_arg("gpu").unwrap(), + DevicePolicy::GpuRequired + ); + assert_eq!( + parse_device_policy_arg("gpu_preferred").unwrap(), + DevicePolicy::GpuPreferred + ); + assert_eq!( + parse_device_policy_arg("cpu_only").unwrap(), + DevicePolicy::CpuOnly + ); + assert!(parse_device_policy_arg("bogus").is_err()); + } +} diff --git a/engines/aim/src/registry.rs b/engines/aim/src/registry.rs new file mode 100644 index 000000000..aef439ff4 --- /dev/null +++ b/engines/aim/src/registry.rs @@ -0,0 +1,185 @@ +// Copyright © Advanced Micro Devices, Inc., or its affiliates. +// +// SPDX-License-Identifier: MIT + +//! Docker Hub v2 registry client for AIM image metadata (ROCMAI-553 finding): +//! the `amdenterpriseai` org's image config blob carries structured +//! `Env`/`Labels` fetchable via two plain HTTP calls, without pulling the image. + +use anyhow::{Context, Result}; +use serde_json::Value; +use std::time::Duration; + +const REGISTRY_HOST: &str = "registry-1.docker.io"; +const AUTH_HOST: &str = "auth.docker.io"; +const DEFAULT_TIMEOUT: Duration = Duration::from_secs(10); +pub(crate) const DEFAULT_ORG: &str = "amdenterpriseai"; + +#[derive(Debug, Default, Clone)] +pub(crate) struct AimImageMetadata { + pub accelerator_family: Option, + pub accelerator_type: Option, + pub pytorch_rocm_arch: Vec, + pub canonical_name: Option, +} + +/// Splits `org/repo:tag` (or bare `repo:tag`, defaulting to [`DEFAULT_ORG`]) +/// into `(repo, tag)` where `repo` includes the org prefix Docker Hub expects. +fn split_image_ref(image_ref: &str) -> (String, String) { + let (name, tag) = image_ref.split_once(':').unwrap_or((image_ref, "latest")); + let repo = if name.contains('/') { + name.to_owned() + } else { + format!("{DEFAULT_ORG}/{name}") + }; + (repo, tag.to_owned()) +} + +fn agent() -> ureq::Agent { + ureq::AgentBuilder::new() + .timeout_connect(DEFAULT_TIMEOUT) + .timeout(DEFAULT_TIMEOUT) + .build() +} + +fn anonymous_pull_token(repo: &str) -> Result { + let url = format!( + "https://{AUTH_HOST}/token?service=registry.docker.io&scope=repository:{repo}:pull" + ); + let body: Value = agent() + .get(&url) + .set("User-Agent", "rocm-cli") + .call() + .with_context(|| format!("failed to fetch registry auth token for {repo}"))? + .into_json() + .with_context(|| format!("failed to parse registry auth token response for {repo}"))?; + body.get("token") + .and_then(Value::as_str) + .map(ToOwned::to_owned) + .with_context(|| format!("registry auth response for {repo} had no token")) +} + +pub(crate) fn fetch_image_metadata(image_ref: &str) -> Result { + let (repo, tag) = split_image_ref(image_ref); + let token = anonymous_pull_token(&repo)?; + + let manifest_url = format!("https://{REGISTRY_HOST}/v2/{repo}/manifests/{tag}"); + let manifest: Value = agent() + .get(&manifest_url) + .set("User-Agent", "rocm-cli") + .set("Authorization", &format!("Bearer {token}")) + .set( + "Accept", + "application/vnd.docker.distribution.manifest.v2+json, application/vnd.oci.image.manifest.v1+json", + ) + .call() + .with_context(|| format!("failed to fetch manifest for {image_ref}"))? + .into_json() + .with_context(|| format!("failed to parse manifest for {image_ref}"))?; + + let config_digest = manifest + .get("config") + .and_then(|config| config.get("digest")) + .and_then(Value::as_str) + .with_context(|| format!("manifest for {image_ref} had no config digest"))?; + + let blob_url = format!("https://{REGISTRY_HOST}/v2/{repo}/blobs/{config_digest}"); + let config: Value = agent() + .get(&blob_url) + .set("User-Agent", "rocm-cli") + .set("Authorization", &format!("Bearer {token}")) + .call() + .with_context(|| format!("failed to fetch config blob for {image_ref}"))? + .into_json() + .with_context(|| format!("failed to parse config blob for {image_ref}"))?; + + Ok(parse_config_blob(&config)) +} + +fn parse_config_blob(config: &Value) -> AimImageMetadata { + let mut metadata = AimImageMetadata::default(); + + let env_entries = config + .pointer("/config/Env") + .and_then(Value::as_array) + .map(|entries| entries.iter().filter_map(Value::as_str).collect::>()) + .unwrap_or_default(); + for entry in env_entries { + if let Some((key, value)) = entry.split_once('=') { + match key { + "AIM_ACCELERATOR_FAMILY" => metadata.accelerator_family = Some(value.to_owned()), + "AIM_ACCELERATOR_TYPE" => metadata.accelerator_type = Some(value.to_owned()), + "PYTORCH_ROCM_ARCH" => { + metadata.pytorch_rocm_arch = value.split(';').map(ToOwned::to_owned).collect(); + } + _ => {} + } + } + } + + if let Some(labels) = config.pointer("/config/Labels").and_then(Value::as_object) { + metadata.canonical_name = labels + .get("com.amd.aim.model.canonicalName") + .and_then(Value::as_str) + .map(ToOwned::to_owned); + } + + metadata +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn split_image_ref_defaults_org_for_bare_repo() { + assert_eq!( + split_image_ref("aim-meta-llama-llama-3-2-1b-instruct:0.11.1"), + ( + "amdenterpriseai/aim-meta-llama-llama-3-2-1b-instruct".to_owned(), + "0.11.1".to_owned() + ) + ); + } + + #[test] + fn split_image_ref_preserves_explicit_org() { + assert_eq!( + split_image_ref("amdenterpriseai/aim-foo:latest"), + ("amdenterpriseai/aim-foo".to_owned(), "latest".to_owned()) + ); + } + + #[test] + fn split_image_ref_defaults_tag_to_latest() { + assert_eq!( + split_image_ref("amdenterpriseai/aim-foo"), + ("amdenterpriseai/aim-foo".to_owned(), "latest".to_owned()) + ); + } + + #[test] + fn parse_config_blob_extracts_known_fields() { + let config = serde_json::json!({ + "config": { + "Env": [ + "PATH=/usr/bin", + "AIM_ACCELERATOR_FAMILY=radeon", + "AIM_ACCELERATOR_TYPE=gpu", + "PYTORCH_ROCM_ARCH=gfx1100;gfx1151", + ], + "Labels": { + "com.amd.aim.model.canonicalName": "Llama-3.2-1B-Instruct", + }, + }, + }); + let metadata = parse_config_blob(&config); + assert_eq!(metadata.accelerator_family.as_deref(), Some("radeon")); + assert_eq!(metadata.accelerator_type.as_deref(), Some("gpu")); + assert_eq!(metadata.pytorch_rocm_arch, vec!["gfx1100", "gfx1151"]); + assert_eq!( + metadata.canonical_name.as_deref(), + Some("Llama-3.2-1B-Instruct") + ); + } +} diff --git a/engines/aim/src/state.rs b/engines/aim/src/state.rs new file mode 100644 index 000000000..9d8fd4796 --- /dev/null +++ b/engines/aim/src/state.rs @@ -0,0 +1,248 @@ +// Copyright © Advanced Micro Devices, Inc., or its affiliates. +// +// SPDX-License-Identifier: MIT + +use anyhow::{Context, Result}; +use rocm_core::{AppPaths, openai_models_endpoint_has_model, require_nonempty}; +use rocm_engine_protocol::{ + DEFAULT_LOG_TAIL_LINES, DevicePolicy, EndpointRequest, EndpointResponse, EngineRecipeHint, + HealthcheckRequest, HealthcheckResponse, LogsRequest, LogsResponse, +}; +use serde_json::{Value, json}; +use std::fs; +use std::path::{Path, PathBuf}; +use std::time::Duration; + +const HEALTHCHECK_TIMEOUT_MS: u64 = 700; + +pub(crate) fn healthcheck_service(request: HealthcheckRequest) -> Result { + require_nonempty(&request.service_id, "service_id")?; + let state_path = state_path(&request.service_id); + let state = read_service_state(&state_path).ok(); + + let endpoint_url = state + .as_ref() + .and_then(endpoint_url_from_state) + .unwrap_or_default(); + let model_ref = state.as_ref().and_then(|s| value_string(s, "model_ref")); + let state_status = state + .as_ref() + .and_then(|s| value_string(s, "status")) + .unwrap_or_else(|| "unknown".to_owned()); + let device = if state.is_some() { + "rocm_gpu" + } else { + "unknown" + }; + + let listed = !endpoint_url.is_empty() + && query_loaded_model_endpoint(&endpoint_url, model_ref.as_deref()).unwrap_or(false); + let ready = listed + && model_ref.as_deref().is_some_and(|model_ref| { + inference_verified(&state_path, state.as_ref(), &endpoint_url, model_ref) + }); + + Ok(HealthcheckResponse::for_readiness( + listed, + ready, + &state_status, + device, + )) +} + +pub(crate) fn endpoint_response(request: EndpointRequest) -> Result { + require_nonempty(&request.service_id, "service_id")?; + let state = read_service_state(&state_path(&request.service_id)) + .with_context(|| format!("service state not found for `{}`", request.service_id))?; + let endpoint_url = endpoint_url_from_state(&state) + .with_context(|| format!("service `{}` has no endpoint URL", request.service_id))?; + Ok(EndpointResponse { + endpoint_url, + api_style: "openai".to_owned(), + supported_routes: vec![ + "/health".to_owned(), + "/v1/models".to_owned(), + "/v1/chat/completions".to_owned(), + "/v1/completions".to_owned(), + ], + }) +} + +pub(crate) fn logs_response(request: LogsRequest) -> Result { + require_nonempty(&request.service_id, "service_id")?; + let limit = request.tail_lines.unwrap_or(DEFAULT_LOG_TAIL_LINES); + let name = crate::process::container_name(&request.service_id); + let recent_lines = crate::docker::logs(&name, limit).unwrap_or_default(); + Ok(LogsResponse { + log_path: format!("docker://{name}"), + recent_lines, + }) +} + +pub(crate) fn state_path(service_id: &str) -> PathBuf { + AppPaths::discover().map_or_else( + |_| PathBuf::from(format!("{service_id}.json")), + |paths| { + paths + .engine_state_dir(crate::ENGINE_NAME) + .join(format!("{service_id}.json")) + }, + ) +} + +pub(crate) fn write_running_state( + service_id: &str, + model_ref: &str, + host: &str, + port: u16, + device_policy: DevicePolicy, + engine_recipe: Option<&EngineRecipeHint>, +) -> Result<()> { + let endpoint = crate::process::endpoint_url(host, port); + write_state( + &state_path(service_id), + &json!({ + "service_id": service_id, + "engine": crate::ENGINE_NAME, + "status": "running", + "container_name": crate::process::container_name(service_id), + "model_ref": model_ref, + "host": host, + "port": port, + "endpoint_url": endpoint, + "device_policy": device_policy, + "engine_recipe": engine_recipe, + "started_at_unix_ms": current_unix_millis(), + }), + ) +} + +pub(crate) fn write_terminal_state(state_path: &Path, status: &str) -> Result<()> { + let mut state = read_service_state(state_path).unwrap_or_else(|_| json!({})); + if let Some(object) = state.as_object_mut() { + object.insert("status".to_owned(), json!(status)); + object.insert( + "stopped_at_unix_ms".to_owned(), + json!(current_unix_millis()), + ); + } + write_state(state_path, &state) +} + +pub(crate) fn read_service_state(path: &Path) -> Result { + let text = + fs::read_to_string(path).with_context(|| format!("failed to read {}", path.display()))?; + serde_json::from_str(&text).with_context(|| format!("failed to parse {}", path.display())) +} + +fn write_state(path: &Path, value: &Value) -> Result<()> { + if let Some(parent) = path.parent() { + fs::create_dir_all(parent) + .with_context(|| format!("failed to create {}", parent.display()))?; + } + fs::write(path, serde_json::to_vec_pretty(value)?) + .with_context(|| format!("failed to write {}", path.display())) +} + +fn endpoint_url_from_state(state: &Value) -> Option { + if let Some(endpoint) = value_string(state, "endpoint_url") { + return Some(endpoint); + } + let host = value_string(state, "host")?; + let port = state.get("port")?.as_u64()? as u16; + Some(crate::process::endpoint_url(&host, port)) +} + +pub(crate) fn query_loaded_model_endpoint( + endpoint_url: &str, + model_ref: Option<&str>, +) -> Result { + let endpoint_api_key = rocm_engine_protocol::resolve_endpoint_api_key(); + openai_models_endpoint_has_model( + endpoint_url, + model_ref, + endpoint_api_key.as_deref(), + Duration::from_millis(HEALTHCHECK_TIMEOUT_MS), + ) +} + +pub(crate) fn query_inference_probe_endpoint(endpoint_url: &str, model_ref: &str) -> Result { + if model_ref.trim().is_empty() { + return Ok(false); + } + let endpoint_api_key = rocm_engine_protocol::resolve_endpoint_api_key(); + rocm_core::openai_chat_completion_probe( + endpoint_url, + model_ref, + endpoint_api_key.as_deref(), + rocm_core::INFERENCE_PROBE_TIMEOUT, + ) +} + +fn inference_verified( + state_path: &Path, + state: Option<&Value>, + endpoint_url: &str, + model_ref: &str, +) -> bool { + let endpoint_api_key = rocm_engine_protocol::resolve_endpoint_api_key(); + rocm_core::engine_state_inference_verified( + state_path, + state, + endpoint_url, + model_ref, + endpoint_api_key.as_deref(), + ) +} + +fn value_string(value: &Value, key: &str) -> Option { + value + .get(key) + .and_then(Value::as_str) + .map(ToOwned::to_owned) +} + +fn current_unix_millis() -> u128 { + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map_or(0, |d| d.as_millis()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn endpoint_url_falls_back_to_host_and_port() { + let state = json!({"host": "127.0.0.1", "port": 8000}); + assert_eq!( + endpoint_url_from_state(&state).as_deref(), + Some("http://127.0.0.1:8000/v1") + ); + } + + #[test] + fn endpoint_url_prefers_explicit_value() { + let state = json!({"endpoint_url": "http://example.test/v1", "host": "ignored", "port": 1}); + assert_eq!( + endpoint_url_from_state(&state).as_deref(), + Some("http://example.test/v1") + ); + } + + #[test] + fn write_and_read_terminal_state_round_trips() { + let dir = std::env::temp_dir().join(format!( + "rocm-aim-state-test-{}-{}", + std::process::id(), + current_unix_millis() + )); + let path = dir.join("svc.json"); + write_state(&path, &json!({"status": "running"})).unwrap(); + write_terminal_state(&path, "stopped").unwrap(); + let state = read_service_state(&path).unwrap(); + assert_eq!(value_string(&state, "status").as_deref(), Some("stopped")); + assert!(state.get("stopped_at_unix_ms").is_some()); + let _ = fs::remove_dir_all(&dir); + } +} diff --git a/xtask/src/crate_edges.rs b/xtask/src/crate_edges.rs index 8c931ad66..d13c978e1 100644 --- a/xtask/src/crate_edges.rs +++ b/xtask/src/crate_edges.rs @@ -47,6 +47,7 @@ const ALLOWLIST: &[(&str, &str)] = &[ ("rocm", "rocm-dash-core"), ("rocm", "rocm-dash-daemon"), ("rocm", "rocm-dash-tui"), + ("rocm", "rocm-engine-aim"), ("rocm", "rocm-engine-lemonade"), ("rocm", "rocm-engine-protocol"), ("rocm", "rocmd"), @@ -59,6 +60,8 @@ const ALLOWLIST: &[(&str, &str)] = &[ ("rocm-dash-daemon", "rocm-dash-collectors"), ("rocm-dash-tui", "rocm-dash-core"), ("rocm-dash-tui", "rocm-deps"), + ("rocm-engine-aim", "rocm-core"), + ("rocm-engine-aim", "rocm-engine-protocol"), ("rocm-engine-lemonade", "rocm-core"), ("rocm-engine-lemonade", "rocm-deps"), ("rocm-engine-lemonade", "rocm-engine-protocol"), From de3b9bb91fad6e8f9c78531827e78e10a1f213b1 Mon Sep 17 00:00:00 2001 From: Juho Vainio Date: Fri, 16 Oct 2026 14:52:40 +0300 Subject: [PATCH 02/22] ROCMAI-554: add e2e-cucumber coverage for the aim engine - HostCapability::engine_available gets a dedicated "aim" arm: the engine's presence in `rocm engines list` is a static inventory entry, not a liveness probe, so without this an @requires-gpu @requires-engine:aim scenario would run (and fail) rather than skip on any host without a reachable Docker daemon. Relocates the existing docker_available() probe out of expectation.rs (private, used only by the unrelated @requires-docker/`rocm remote` gate) into capability.rs so both call sites share it. - Adds model_serving.feature scenario serve-25 and its `a model is being served on aim` step, serving the raw AIM image reference directly (no model catalog in this PR's scope) and asserting inference works end to end. Signed-off-by: Juho Vainio --- .../features/model_serving.feature | 14 +++++++++++ tests/e2e-cucumber/src/capability.rs | 25 +++++++++++++++++++ tests/e2e-cucumber/src/expectation.rs | 20 +-------------- tests/e2e-cucumber/tests/e2e/serving_steps.rs | 24 ++++++++++++++++++ 4 files changed, 64 insertions(+), 19 deletions(-) diff --git a/tests/e2e-cucumber/features/model_serving.feature b/tests/e2e-cucumber/features/model_serving.feature index 90caa1aee..8c9bac328 100644 --- a/tests/e2e-cucumber/features/model_serving.feature +++ b/tests/e2e-cucumber/features/model_serving.feature @@ -341,3 +341,17 @@ Feature: Model serving Scenario: serve-25 - The recommended model list prints ASCII-only text When the user lists recommended models Then the model list output contains only ASCII characters + + # AIM serve + inference canary (ROCMAI-554). AIM has no model catalog in this + # PR's scope (ROCMAI-555) — the model argument IS the raw AIM Docker image + # reference, not a short recipe name. `@requires-engine:aim` only skips here + # when Docker itself is unreachable (see `HostCapability::engine_available`'s + # "aim" arm); a Strix Halo self-hosted runner with Docker is where this + # actually runs. + @id:serve-aim-inference @requires-gpu @requires-engine:aim + Scenario: serve-26 - A model served on AIM responds to inference requests + Given a managed runtime is active + And a model is being served on aim + When the user sends a chat completion request + Then the response contains a model reply + And the response identifies the correct model diff --git a/tests/e2e-cucumber/src/capability.rs b/tests/e2e-cucumber/src/capability.rs index 548b4700a..010935016 100644 --- a/tests/e2e-cucumber/src/capability.rs +++ b/tests/e2e-cucumber/src/capability.rs @@ -34,6 +34,30 @@ fn rocm_binary() -> String { std::env::var("ROCM_CLI_BINARY").unwrap_or_else(|_| "rocm".to_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 for `@requires-docker` — see the +/// `requires_docker` gate in `expectation.rs`, 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. +/// +/// Also used by [`HostCapability::engine_available`] for `"aim"`: unlike +/// lemonade/vllm, the AIM adapter is always present in the binary's inventory +/// (a static description, not a liveness probe — see `rocm engines list`), so +/// "can this engine start" must check for a reachable Docker daemon itself. +pub(crate) 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()) +} + /// What `rocm serve ` would pick with no `--engine`, from GPU family + OS. /// /// Mirrors the product precedence in `preferred_serve_engine_for_host_gpu_summary` @@ -140,6 +164,7 @@ impl HostCapability { !self.os_family.eq_ignore_ascii_case("windows") && family_prefers_vllm(self.gfx_target.as_deref()) } + "aim" => self.available_engines.iter().any(|e| e == engine) && docker_available(), _ => self.available_engines.iter().any(|e| e == engine), } } diff --git a/tests/e2e-cucumber/src/expectation.rs b/tests/e2e-cucumber/src/expectation.rs index 2ffe4e1ed..578a699c4 100644 --- a/tests/e2e-cucumber/src/expectation.rs +++ b/tests/e2e-cucumber/src/expectation.rs @@ -16,7 +16,7 @@ use std::collections::BTreeMap; use serde::Deserialize; -use crate::capability::HostCapability; +use crate::capability::{HostCapability, docker_available}; /// Tag prefixes in their normalized bare form. cucumber-rs currently supplies /// scenario tags with a leading `@`; [`ScenarioDecl::from_tags`] accepts either @@ -61,24 +61,6 @@ 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 { diff --git a/tests/e2e-cucumber/tests/e2e/serving_steps.rs b/tests/e2e-cucumber/tests/e2e/serving_steps.rs index 7ce18b326..2c85fa355 100644 --- a/tests/e2e-cucumber/tests/e2e/serving_steps.rs +++ b/tests/e2e-cucumber/tests/e2e/serving_steps.rs @@ -633,6 +633,30 @@ async fn setup_lemonade_hf_checkpoint_model(world: &mut E2eWorld) { .await; } +#[given("a model is being served on aim")] +async fn setup_aim_model(world: &mut E2eWorld) { + // AIM has no model catalog in this PR's scope (ROCMAI-555): the model + // argument IS the raw AIM Docker image reference, not a short recipe name + // like the vLLM/Lemonade scenarios above. This is the smallest model image + // the ROCMAI-553 registry spike found — `instinct`/`gpu`, and its + // `PYTORCH_ROCM_ARCH` lists `gfx1151` (Strix Halo, the only runner this + // scenario actually executes on). + let image = "amdenterpriseai/aim-meta-llama-llama-3-2-1b-instruct:0.11.1"; + // The container bundles vLLM, which is expected to echo the model's baked-in + // HF id (not the Docker image reference) as `model` — mirroring + // `engine_registry.rs`'s identical "presumed, unconfirmed until Strix Halo + // validation pass" note. Confirm this id against the real container during + // that pass and adjust if it reports something else. + let model = "meta-llama/Llama-3.2-1B-Instruct"; + serve_and_wait( + world, + &["serve", image, "--engine", "aim", "--managed"], + model, + "Llama-3.2-1B", + ) + .await; +} + #[given("a large model is being served on GPU")] async fn setup_large_gpu_model(world: &mut E2eWorld) { // Large-model nightly coverage follows the host's serving path. vLLM hosts From ed9fe91b4c1cdffe0f182eee7e76d44db889e3ca Mon Sep 17 00:00:00 2001 From: Juho Vainio Date: Fri, 16 Oct 2026 15:27:39 +0300 Subject: [PATCH 03/22] ROCMAI-554: fix AIM model-id handling for docker pull and readiness match MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit canonical_model_id was set to the registry's HF-style canonicalName label, but apps/rocm re-threads it back into the engine subprocess as the model_ref used for docker pull/run — an HF-style id (mixed case, not a valid Docker repo path) made docker pull fail outright. Keep canonical_model_id as the image reference instead. That in turn meant the readiness poll was matching /v1/models against the image ref, which the container never reports (it serves under the canonical HF-style name) — pass the registry's canonical_name through for that check instead, falling back to the image ref if absent. Signed-off-by: Juho Vainio --- engines/aim/src/process.rs | 32 ++++++++++++++++++++++---------- 1 file changed, 22 insertions(+), 10 deletions(-) diff --git a/engines/aim/src/process.rs b/engines/aim/src/process.rs index 891cffb89..2dfc1d78c 100644 --- a/engines/aim/src/process.rs +++ b/engines/aim/src/process.rs @@ -69,9 +69,15 @@ pub(crate) fn resolve_model_response(request: ResolveModelRequest) -> Result Result { request.port, &gpu_indices, )?; - let endpoint_url = endpoint_url(&request.host, request.port); crate::state::write_running_state( &request.service_id, @@ -170,17 +175,17 @@ fn run_container( host: &str, port: u16, gpu_indices: &[u32], -) -> Result<()> { +) -> Result> { let name = container_name(service_id); + let metadata = crate::registry::fetch_image_metadata(model_ref).unwrap_or_default(); if crate::docker::inspect_running(&name)?.unwrap_or(false) { - return Ok(()); + return Ok(metadata.canonical_name); } if !crate::docker::image_present(model_ref)? { crate::docker::pull(model_ref)?; } - let metadata = crate::registry::fetch_image_metadata(model_ref).unwrap_or_default(); let use_gpu = metadata.accelerator_type.as_deref() != Some("cpu") && !gpu_indices.is_empty(); let render_gid = if use_gpu { crate::docker::render_group_gid() @@ -196,11 +201,12 @@ fn run_container( crate::DEFAULT_AIM_PORT, if use_gpu { gpu_indices } else { &[] }, render_gid, - ) + )?; + Ok(metadata.canonical_name) } pub(crate) fn serve_http(request: ServeHttpRequest) -> Result<()> { - run_container( + let canonical_name = run_container( &request.service_id, &request.model_ref, &request.host, @@ -219,11 +225,17 @@ pub(crate) fn serve_http(request: ServeHttpRequest) -> Result<()> { let name = container_name(&request.service_id); let timeout = DEFAULT_READY_TIMEOUT; + // The container self-reports under the registry's HF-style canonical name + // (e.g. `meta-llama/Llama-3.2-1B-Instruct`), not the Docker image reference + // used for `docker pull`/`run` — matching `/v1/models` against the image ref + // would never find it. Fall back to the image ref only when the registry + // didn't advertise a canonical name. + let expected_model_id = canonical_name.as_deref().unwrap_or(&request.model_ref); match wait_for_container_ready( &name, &request.host, request.port, - &request.model_ref, + expected_model_id, timeout, ) { Ok(()) => Ok(()), From 2db4bacfc56bf4e4abb56719264c8de52921683e Mon Sep 17 00:00:00 2001 From: Juho Vainio Date: Fri, 16 Oct 2026 16:49:53 +0300 Subject: [PATCH 04/22] ROCMAI-554: add generic --engine-env passthrough for engine recipes Adds EngineRecipeHint::extra_env (BTreeMap) and a repeatable `--engine-env KEY=VALUE` flag on `rocm serve`, consumed only by engines/aim's docker run to emit `-e KEY=VALUE` args (e.g. to set AIM_PROFILE_ID). Generic by design so any engine can opt in later. Signed-off-by: Juho Vainio --- apps/rocm/src/main.rs | 8 ++++++ apps/rocm/src/serve_cmd.rs | 37 ++++++++++++++++++++++++++ crates/rocm-engine-protocol/src/lib.rs | 11 ++++++++ engines/aim/src/docker.rs | 29 ++++++++++++++++++++ engines/aim/src/process.rs | 5 ++++ engines/vllm/src/process.rs | 1 + engines/vllm/src/state.rs | 1 + 7 files changed, 92 insertions(+) diff --git a/apps/rocm/src/main.rs b/apps/rocm/src/main.rs index 50a8a6ef6..8c23f083c 100644 --- a/apps/rocm/src/main.rs +++ b/apps/rocm/src/main.rs @@ -564,6 +564,12 @@ rocm serve qwen --verbose --device gpu_required")] /// authenticated too. #[arg(long)] api_key: Option, + /// Extra `KEY=VALUE` environment variable to pass through to the engine's + /// own process/container (repeatable). Engine-specific; currently only + /// `--engine aim` reads these (e.g. `AIM_PROFILE_ID`). Ignored by engines + /// that don't consume engine-env overrides. + #[arg(long = "engine-env", value_name = "KEY=VALUE")] + engine_env: Vec, }, /// [preview] Install, start, stop, or inspect ComfyUI. #[command(alias = "comfy")] @@ -2510,6 +2516,7 @@ fn dispatch(cli: Cli) -> Result<()> { top_p, max_tokens, api_key, + engine_env, }) => serve(ServeArgs { model, engine, @@ -2531,6 +2538,7 @@ fn dispatch(cli: Cli) -> Result<()> { top_p, max_tokens, api_key, + engine_env, }), Some(Command::Comfyui { command }) => comfyui(command), Some(Command::Services { command }) => services(command), diff --git a/apps/rocm/src/serve_cmd.rs b/apps/rocm/src/serve_cmd.rs index b2d823008..0e19f0c4f 100644 --- a/apps/rocm/src/serve_cmd.rs +++ b/apps/rocm/src/serve_cmd.rs @@ -21,6 +21,7 @@ //! to `service_launch.rs` in ROCMAI-91 (Phase 6b), along with its test module; //! this file's `use` list below points there now instead of at the crate root. +use std::collections::BTreeMap; use std::fmt::Write as _; use anyhow::{Result, bail}; @@ -180,6 +181,7 @@ fn protocol_engine_recipe_hint( engine: engine_recipe.engine.clone(), required_flags: engine_recipe.required_flags.clone(), parser_settings: engine_recipe.parser_settings.clone(), + extra_env: BTreeMap::new(), preferred_endpoint: engine_recipe.preferred_endpoint.as_ref().map(|endpoint| { EngineRecipeEndpointHint { endpoint_mode: endpoint.endpoint_mode.clone(), @@ -483,6 +485,38 @@ fn parse_gpu_memory_utilization(value: Option<&str>) -> Result> { Ok(Some(parsed)) } +/// Applies `--engine-env KEY=VALUE` overrides to the engine recipe's `extra_env` +/// map. Generic by design (not gated to a specific engine): the mechanism exists +/// so any engine that reads `EngineRecipeHint::extra_env` (currently only +/// `engines/aim`, for overrides like `AIM_PROFILE_ID`) can accept host-specific +/// env tweaks without a dedicated flag per variable. A minimal hint is +/// synthesized when none exists. With no `--engine-env` the hint passes through +/// unchanged. +fn engine_recipe_with_extra_env( + engine: &str, + hint: Option, + engine_env: &[String], +) -> Result> { + if engine_env.is_empty() { + return Ok(hint); + } + let mut hint = hint.unwrap_or_else(|| EngineRecipeHint { + contract_version: ENGINE_RECIPE_CONTRACT_VERSION.to_owned(), + engine: engine.to_owned(), + ..EngineRecipeHint::default() + }); + for entry in engine_env { + let (key, value) = entry + .split_once('=') + .ok_or_else(|| anyhow::anyhow!("--engine-env expects `KEY=VALUE`; got `{entry}`"))?; + if key.is_empty() { + bail!("--engine-env expects `KEY=VALUE` with a non-empty key; got `{entry}`"); + } + hint.extra_env.insert(key.to_owned(), value.to_owned()); + } + Ok(Some(hint)) +} + /// Whether the resolved engine recipe launches vLLM with tool calling enabled. fn engine_recipe_enables_tool_choice(hint: Option<&EngineRecipeHint>) -> bool { hint.is_some_and(|hint| { @@ -515,6 +549,7 @@ pub(crate) struct ServeArgs { pub(crate) top_p: Option, pub(crate) max_tokens: Option, pub(crate) api_key: Option, + pub(crate) engine_env: Vec, } pub(crate) fn serve(args: ServeArgs) -> Result<()> { @@ -539,6 +574,7 @@ pub(crate) fn serve(args: ServeArgs) -> Result<()> { top_p, max_tokens, api_key, + engine_env, } = args; let _ = managed; // background is now the default; --managed is accepted as an explicit synonym. validate_bind_host(&host, allow_public_bind)?; @@ -627,6 +663,7 @@ pub(crate) fn serve(args: ServeArgs) -> Result<()> { engine_recipe, generation_defaults, )?; + let engine_recipe = engine_recipe_with_extra_env(&selected_engine, engine_recipe, &engine_env)?; let tool_call_note = if tool_call_parser.is_some() && !engine_serves_vllm { Some(format!( "note: --tool-call-parser applies only to vLLM; ignored for engine '{selected_engine}'" diff --git a/crates/rocm-engine-protocol/src/lib.rs b/crates/rocm-engine-protocol/src/lib.rs index fbca3104f..eff367bf9 100644 --- a/crates/rocm-engine-protocol/src/lib.rs +++ b/crates/rocm-engine-protocol/src/lib.rs @@ -511,6 +511,11 @@ pub struct EngineRecipeHint { pub required_flags: Vec, #[serde(default)] pub parser_settings: BTreeMap, + /// Engine-specific environment variables to pass through verbatim to the + /// engine's own process/container (e.g. AIM's `AIM_PROFILE_ID` override). + /// Generic by design — only `engines/aim` currently reads this. + #[serde(default)] + pub extra_env: BTreeMap, #[serde(default, skip_serializing_if = "Option::is_none")] pub preferred_endpoint: Option, #[serde(default)] @@ -888,6 +893,7 @@ mod tests { "reasoning_parser".to_owned(), "qwen3".to_owned(), )]), + extra_env: BTreeMap::from([("AIM_PROFILE_ID".to_owned(), "foo".to_owned())]), preferred_endpoint: Some(EngineRecipeEndpointHint { endpoint_mode: "openai".to_owned(), settings: BTreeMap::from([("streaming".to_owned(), "true".to_owned())]), @@ -921,6 +927,10 @@ mod tests { ); assert_eq!(hint.unsupported_combinations.len(), 1); assert_eq!(hint.notes, vec!["signed recipe metadata"]); + assert_eq!( + hint.extra_env.get("AIM_PROFILE_ID").map(String::as_str), + Some("foo") + ); } #[test] @@ -939,6 +949,7 @@ mod tests { engine: "vllm".to_owned(), required_flags: vec!["--enable-auto-tool-choice".to_owned()], parser_settings: BTreeMap::new(), + extra_env: BTreeMap::new(), preferred_endpoint: None, unsupported_combinations: Vec::new(), notes: Vec::new(), diff --git a/engines/aim/src/docker.rs b/engines/aim/src/docker.rs index 679b46368..deb291f37 100644 --- a/engines/aim/src/docker.rs +++ b/engines/aim/src/docker.rs @@ -7,11 +7,21 @@ //! already uses for its own CLI (`run_lemonade_pull`/`run_lemonade_backend_install`). use anyhow::{Context, Result, bail}; +use std::collections::BTreeMap; use std::process::{Command, Stdio}; use std::time::Duration; pub(crate) const DOCKER_BIN: &str = "docker"; +/// Builds `-e KEY=VALUE` args for engine-specific env overrides (e.g. AIM's +/// `AIM_PROFILE_ID`). Sorted (`BTreeMap`) for deterministic argv in tests/logs. +fn extra_env_args(extra_env: &BTreeMap) -> Vec { + extra_env + .iter() + .flat_map(|(key, value)| ["-e".to_owned(), format!("{key}={value}")]) + .collect() +} + #[allow(clippy::too_many_arguments)] pub(crate) fn run( name: &str, @@ -21,6 +31,7 @@ pub(crate) fn run( container_port: u16, gpu_indices: &[u32], render_gid: Option, + extra_env: &BTreeMap, ) -> Result<()> { let mut command = Command::new(DOCKER_BIN); command @@ -44,6 +55,8 @@ pub(crate) fn run( } } + command.args(extra_env_args(extra_env)); + command.arg(image); command.stdin(Stdio::null()); @@ -196,10 +209,26 @@ pub(crate) fn daemon_reachable() -> bool { #[cfg(test)] mod tests { + use super::extra_env_args; + use std::collections::BTreeMap; + #[test] fn render_group_gid_parses_getent_format() { let line = "render:x:993:"; let gid: Option = line.trim().split(':').nth(2).and_then(|s| s.parse().ok()); assert_eq!(gid, Some(993)); } + + #[test] + fn extra_env_args_builds_sorted_flags() { + assert!(extra_env_args(&BTreeMap::new()).is_empty()); + + let mut extra_env = BTreeMap::new(); + extra_env.insert("AIM_PROFILE_ID".to_owned(), "vllm-foo".to_owned()); + extra_env.insert("ANOTHER".to_owned(), "x".to_owned()); + assert_eq!( + extra_env_args(&extra_env), + vec!["-e", "AIM_PROFILE_ID=vllm-foo", "-e", "ANOTHER=x"] + ); + } } diff --git a/engines/aim/src/process.rs b/engines/aim/src/process.rs index 2dfc1d78c..d52033f08 100644 --- a/engines/aim/src/process.rs +++ b/engines/aim/src/process.rs @@ -140,6 +140,7 @@ pub(crate) fn launch_service(request: LaunchRequest) -> Result { &request.host, request.port, &gpu_indices, + engine_recipe.as_ref().map(|hint| &hint.extra_env), )?; let endpoint_url = endpoint_url(&request.host, request.port); crate::state::write_running_state( @@ -175,6 +176,7 @@ fn run_container( host: &str, port: u16, gpu_indices: &[u32], + extra_env: Option<&std::collections::BTreeMap>, ) -> Result> { let name = container_name(service_id); let metadata = crate::registry::fetch_image_metadata(model_ref).unwrap_or_default(); @@ -193,6 +195,7 @@ fn run_container( None }; + let empty_env = std::collections::BTreeMap::default(); crate::docker::run( &name, model_ref, @@ -201,6 +204,7 @@ fn run_container( crate::DEFAULT_AIM_PORT, if use_gpu { gpu_indices } else { &[] }, render_gid, + extra_env.unwrap_or(&empty_env), )?; Ok(metadata.canonical_name) } @@ -212,6 +216,7 @@ pub(crate) fn serve_http(request: ServeHttpRequest) -> Result<()> { &request.host, request.port, &request.gpu_indices, + request.engine_recipe.as_ref().map(|hint| &hint.extra_env), )?; crate::state::write_running_state( diff --git a/engines/vllm/src/process.rs b/engines/vllm/src/process.rs index 86281468e..17889eb9c 100644 --- a/engines/vllm/src/process.rs +++ b/engines/vllm/src/process.rs @@ -866,6 +866,7 @@ mod tests { engine: engine.to_owned(), required_flags: vec!["--enable-auto-tool-choice".to_owned()], parser_settings: std::collections::BTreeMap::default(), + extra_env: std::collections::BTreeMap::default(), preferred_endpoint: None, unsupported_combinations: Vec::new(), notes: vec!["test recipe".to_owned()], diff --git a/engines/vllm/src/state.rs b/engines/vllm/src/state.rs index 6ad87f937..bca142849 100644 --- a/engines/vllm/src/state.rs +++ b/engines/vllm/src/state.rs @@ -652,6 +652,7 @@ mod tests { engine: crate::ENGINE_NAME.to_owned(), required_flags: vec!["--enable-auto-tool-choice".to_owned()], parser_settings: std::collections::BTreeMap::default(), + extra_env: std::collections::BTreeMap::default(), preferred_endpoint: None, unsupported_combinations: Vec::new(), notes: Vec::new(), From 1ae2386a3d414e2bd0d3884ea04019d99ed1c25c Mon Sep 17 00:00:00 2001 From: Juho Vainio Date: Fri, 16 Oct 2026 18:10:21 +0300 Subject: [PATCH 05/22] ROCMAI-554: make AIM container ready-timeout configurable MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Mirrors engines/vllm's ROCM_CLI_VLLM_READY_TIMEOUT_SECS pattern. The previous hardcoded 5-minute DEFAULT_READY_TIMEOUT caused a live Strix Halo validation run to misreport a genuinely-serving AIM container as status: failed — a cold-start HF weight download for a 9B model legitimately exceeded 5 minutes, so the supervisor timed out and exited before the container actually became ready, with nothing left alive to correct the recorded status afterward. Signed-off-by: Juho Vainio --- engines/aim/src/process.rs | 44 +++++++++++++++++++++++++++++++++++++- 1 file changed, 43 insertions(+), 1 deletion(-) diff --git a/engines/aim/src/process.rs b/engines/aim/src/process.rs index d52033f08..c933f727a 100644 --- a/engines/aim/src/process.rs +++ b/engines/aim/src/process.rs @@ -17,6 +17,23 @@ const STOP_GRACE: Duration = Duration::from_secs(10); const READY_POLL_INTERVAL: Duration = Duration::from_millis(500); const DEFAULT_READY_TIMEOUT: Duration = Duration::from_mins(5); +/// Time to wait for the AIM container to become ready before reporting failure. +/// +/// Defaults to [`DEFAULT_READY_TIMEOUT`]. A cold start (image's first HF weight +/// download, first-decode compile) can exceed the default, so the timeout is +/// configurable via `ROCM_CLI_AIM_READY_TIMEOUT_SECS` (a positive integer number +/// of seconds) — mirrors `engines/vllm`'s `ROCM_CLI_VLLM_READY_TIMEOUT_SECS`. +fn aim_ready_timeout() -> Duration { + resolve_aim_ready_timeout(std::env::var("ROCM_CLI_AIM_READY_TIMEOUT_SECS").ok()) +} + +fn resolve_aim_ready_timeout(override_value: Option) -> Duration { + override_value + .and_then(|value| value.trim().parse::().ok()) + .filter(|secs| *secs > 0) + .map_or(DEFAULT_READY_TIMEOUT, Duration::from_secs) +} + /// Deterministic, Docker-name-sanitized container identity for a service. /// This *is* the service's identity for the AIM engine — no PID/start-ticks /// tracking, since `docker run -d` leaves no long-lived client process to @@ -229,7 +246,7 @@ pub(crate) fn serve_http(request: ServeHttpRequest) -> Result<()> { )?; let name = container_name(&request.service_id); - let timeout = DEFAULT_READY_TIMEOUT; + let timeout = aim_ready_timeout(); // The container self-reports under the registry's HF-style canonical name // (e.g. `meta-llama/Llama-3.2-1B-Instruct`), not the Docker image reference // used for `docker pull`/`run` — matching `/v1/models` against the image ref @@ -335,6 +352,31 @@ pub(crate) fn parse_gpu_indices_arg(value: Option<&str>) -> Result> { mod tests { use super::*; + #[test] + fn aim_ready_timeout_uses_default_without_override() { + assert_eq!(resolve_aim_ready_timeout(None), DEFAULT_READY_TIMEOUT); + } + + #[test] + fn aim_ready_timeout_honors_positive_override() { + assert_eq!( + resolve_aim_ready_timeout(Some(" 900 ".to_owned())), + Duration::from_mins(15) + ); + } + + #[test] + fn aim_ready_timeout_ignores_invalid_or_zero_override() { + assert_eq!( + resolve_aim_ready_timeout(Some("0".to_owned())), + DEFAULT_READY_TIMEOUT + ); + assert_eq!( + resolve_aim_ready_timeout(Some("not-a-number".to_owned())), + DEFAULT_READY_TIMEOUT + ); + } + #[test] fn container_name_sanitizes_unsafe_characters() { assert_eq!(container_name("svc 1/a"), "rocm-aim-svc-1-a"); From b69971be3bb11f8161288ad3df34e6fd9b2dcbf7 Mon Sep 17 00:00:00 2001 From: Juho Vainio Date: Fri, 16 Oct 2026 18:26:07 +0300 Subject: [PATCH 06/22] ROCMAI-554: fix e2e-cucumber AIM scenario to use a Strix-Halo-compatible image The serve-25 AIM scenario hardcoded aim-meta-llama-llama-3-2-1b-instruct, whose 48 bundled profiles are all Instinct-series-only (MI300X/MI250X/ MI355X/MI325X/MI350X) -- confirmed via live container logs to report "Found 0 compatible profiles" on Strix Halo/gfx1151, contradicting the step's own comment. Swap in aim-radeon-qwen-qwen3-5-9b, hardware-verified serving real inference on Strix Halo, with an explicit --engine-env AIM_PROFILE_ID=vllm-w7900-bf16-tp1-latency since Radeon profile auto-selection isn't reliable yet. Also widen the scenario's readiness window via @serve-timeout:1200 for the 9B model's cold-start HF download, and propagate ROCM_CLI_AIM_READY_TIMEOUT_SECS alongside the existing vLLM timeout var in isolate_cmd so the CLI's own AIM readiness cap (just made configurable in e77c0e85) widens to match. Signed-off-by: Juho Vainio --- .../features/model_serving.feature | 7 +++- tests/e2e-cucumber/tests/e2e.rs | 16 +++++--- tests/e2e-cucumber/tests/e2e/serving_steps.rs | 38 ++++++++++++------- 3 files changed, 40 insertions(+), 21 deletions(-) diff --git a/tests/e2e-cucumber/features/model_serving.feature b/tests/e2e-cucumber/features/model_serving.feature index 8c9bac328..70cb04546 100644 --- a/tests/e2e-cucumber/features/model_serving.feature +++ b/tests/e2e-cucumber/features/model_serving.feature @@ -347,8 +347,11 @@ Feature: Model serving # reference, not a short recipe name. `@requires-engine:aim` only skips here # when Docker itself is unreachable (see `HostCapability::engine_available`'s # "aim" arm); a Strix Halo self-hosted runner with Docker is where this - # actually runs. - @id:serve-aim-inference @requires-gpu @requires-engine:aim + # actually runs. `@serve-timeout:1200` widens both the harness's readiness + # poll and the CLI's own ROCM_CLI_AIM_READY_TIMEOUT_SECS (see isolate_cmd) — + # a cold-start HF weight download for this 9B model legitimately exceeds the + # engine's default 5-minute readiness cap. + @id:serve-aim-inference @requires-gpu @requires-engine:aim @serve-timeout:1200 Scenario: serve-26 - A model served on AIM responds to inference requests Given a managed runtime is active And a model is being served on aim diff --git a/tests/e2e-cucumber/tests/e2e.rs b/tests/e2e-cucumber/tests/e2e.rs index 3ce355c58..40df0241c 100644 --- a/tests/e2e-cucumber/tests/e2e.rs +++ b/tests/e2e-cucumber/tests/e2e.rs @@ -397,14 +397,18 @@ impl E2eWorld { } // When a scenario declares a longer serve-readiness window (via a // `@serve-timeout:` tag → serve_timeout_override), also raise the - // CLI's OWN vLLM readiness cap to match. `rocm serve --managed` otherwise - // SIGTERM-kills a vLLM that isn't ready within its default (5 min, - // EAI-7393), which a large model's cold load legitimately exceeds — so - // extending only the harness's poll (`model_is_ready`) isn't enough; the CLI - // would kill the server first. Keeping the two in lockstep makes the - // big-model serve actually reach ready (verified on MI300X with Qwen3.6-27B). + // CLI's OWN per-engine readiness cap(s) to match. `rocm serve --managed` + // otherwise kills/fails a server that isn't ready within its default (5 + // min, EAI-7393), which a large model's cold load legitimately exceeds — + // so extending only the harness's poll (`model_is_ready`) isn't enough; + // the CLI would terminate the server first. Keeping them in lockstep + // makes the big-model serve actually reach ready (vLLM verified on + // MI300X with Qwen3.6-27B; AIM verified on Strix Halo with a 9B model's + // cold-start HF download). Setting both vars unconditionally is harmless + // — the unused one is a no-op for whichever engine isn't in play. if let Some(secs) = self.serve_timeout_override { env.push(("ROCM_CLI_VLLM_READY_TIMEOUT_SECS", secs.to_string().into())); + env.push(("ROCM_CLI_AIM_READY_TIMEOUT_SECS", secs.to_string().into())); } env } diff --git a/tests/e2e-cucumber/tests/e2e/serving_steps.rs b/tests/e2e-cucumber/tests/e2e/serving_steps.rs index 2c85fa355..c14a33cc5 100644 --- a/tests/e2e-cucumber/tests/e2e/serving_steps.rs +++ b/tests/e2e-cucumber/tests/e2e/serving_steps.rs @@ -637,22 +637,34 @@ async fn setup_lemonade_hf_checkpoint_model(world: &mut E2eWorld) { async fn setup_aim_model(world: &mut E2eWorld) { // AIM has no model catalog in this PR's scope (ROCMAI-555): the model // argument IS the raw AIM Docker image reference, not a short recipe name - // like the vLLM/Lemonade scenarios above. This is the smallest model image - // the ROCMAI-553 registry spike found — `instinct`/`gpu`, and its - // `PYTORCH_ROCM_ARCH` lists `gfx1151` (Strix Halo, the only runner this - // scenario actually executes on). - let image = "amdenterpriseai/aim-meta-llama-llama-3-2-1b-instruct:0.11.1"; - // The container bundles vLLM, which is expected to echo the model's baked-in - // HF id (not the Docker image reference) as `model` — mirroring - // `engine_registry.rs`'s identical "presumed, unconfirmed until Strix Halo - // validation pass" note. Confirm this id against the real container during - // that pass and adjust if it reports something else. - let model = "meta-llama/Llama-3.2-1B-Instruct"; + // like the vLLM/Lemonade scenarios above. + // + // The `aim-meta-llama-llama-3-2-1b-instruct` image originally used here + // bundles only Instinct-series (MI300X/MI250X/MI355X/MI325X/MI350X) + // profiles — confirmed via its container logs to report "Found 0 + // compatible profiles" on Strix Halo/gfx1151, the only runner this + // scenario executes on. `aim-radeon-qwen-qwen3-5-9b` is a Radeon-targeted + // image, hardware-verified serving real inference on Strix Halo; its + // profile still needs an explicit pick via `--engine-env` because + // auto-selection isn't reliable for Radeon targets yet. + let image = "amdenterpriseai/aim-radeon-qwen-qwen3-5-9b:0.12.0-preview"; + // The container bundles vLLM, which echoes the model's baked-in HF id + // (not the Docker image reference) as `model` — confirmed live on Strix + // Halo. + let model = "Qwen/Qwen3.5-9B"; serve_and_wait( world, - &["serve", image, "--engine", "aim", "--managed"], + &[ + "serve", + image, + "--engine", + "aim", + "--engine-env", + "AIM_PROFILE_ID=vllm-w7900-bf16-tp1-latency", + "--managed", + ], model, - "Llama-3.2-1B", + "Qwen3.5-9B", ) .await; } From cf3d3b27bb9a6dc2cc3154ef2aeaeb8456f3fd19 Mon Sep 17 00:00:00 2001 From: Juho Vainio Date: Mon, 19 Oct 2026 08:14:37 +0300 Subject: [PATCH 07/22] ROCMAI-554: rule the aim engine out when Docker is unreachable The host-fit gate only accounted for vLLM-on-Windows. aim needs a Docker daemon to serve at all, and nothing checked for one before serve_cmd tried to launch a container -- the /model adapter-availability note and rocm serve's readiness verdict would both silently assume aim was fine. Extend engine_ruled_out_by_platform (the single source of truth both callers already share, per its own doc comment) to also rule aim out when rocm_engine_aim::builtin_docker_daemon_reachable() is false, and branch the two message-producing functions on engine name so the "ruled out" text names the actual reason (Docker vs native Windows) instead of always blaming Windows. Also wire the remaining small aim dispatch gaps: friendly_engine_label, infer_engine_from_request, and a builtin_docker_daemon_reachable() wrapper on engines/aim mirroring the existing builtin_handle_envelope pattern. Signed-off-by: Juho Vainio --- apps/rocm/src/main.rs | 42 +++++++++++++++++++++++++++++++++--------- engines/aim/src/lib.rs | 8 ++++++++ 2 files changed, 41 insertions(+), 9 deletions(-) diff --git a/apps/rocm/src/main.rs b/apps/rocm/src/main.rs index 8c23f083c..cc8cf545a 100644 --- a/apps/rocm/src/main.rs +++ b/apps/rocm/src/main.rs @@ -3183,8 +3183,10 @@ fn assess_model_for_host( /// /// One place, so the `/model` adapter-availability note and the readiness /// verdict cannot answer differently about the same host and engine. -const fn engine_ruled_out_by_platform(engine: &str) -> bool { - rocm_core::runtime_is_windows() && engine.eq_ignore_ascii_case("vllm") +fn engine_ruled_out_by_platform(engine: &str) -> bool { + (rocm_core::runtime_is_windows() && engine.eq_ignore_ascii_case("vllm")) + || (engine.eq_ignore_ascii_case("aim") + && !rocm_engine_aim::builtin_docker_daemon_reachable()) } /// The `unsupported_here` message for an engine, given whether the platform @@ -3197,8 +3199,13 @@ const fn engine_ruled_out_by_platform(engine: &str) -> bool { /// exercise the "ruled out" branch on Linux CI and the "allowed" branch on a /// Windows runner without either one being unreachable. fn unsupported_here_for(engine: &str, ruled_out: bool) -> Option { - ruled_out - .then(|| format!("{engine} has no adapter on native Windows; serve it from WSL or Linux")) + ruled_out.then(|| { + if engine.eq_ignore_ascii_case("aim") { + format!("{engine} needs a reachable Docker daemon to serve; start Docker or choose a different engine") + } else { + format!("{engine} has no adapter on native Windows; serve it from WSL or Linux") + } + }) } /// The catalog entry a report should name, and whether a fix was offered for it. @@ -6066,13 +6073,18 @@ fn append_model_engine_support_lines( } #[allow(dead_code)] -const fn model_registry_adapter_availability_note(engine: &str) -> Option<&'static str> { - if engine_ruled_out_by_platform(engine) { +fn model_registry_adapter_availability_note(engine: &str) -> Option<&'static str> { + if !engine_ruled_out_by_platform(engine) { + return None; + } + if engine.eq_ignore_ascii_case("aim") { Some( - "runtime_status=unsupported_native_windows reason=native Windows skipped; use WSL/Linux vLLM ROCm; gpu_execution_required=true; run /engine for adapter details", + "runtime_status=unsupported_docker_unavailable reason=Docker daemon not reachable; start Docker or choose a different engine; run /engine for adapter details", ) } else { - None + Some( + "runtime_status=unsupported_native_windows reason=native Windows skipped; use WSL/Linux vLLM ROCm; gpu_execution_required=true; run /engine for adapter details", + ) } } @@ -8177,6 +8189,7 @@ fn friendly_engine_label(engine: &str) -> &str { match engine { "lemonade" => "Lemonade", "vllm" => "vLLM", + "aim" => "AIM", other => other, } } @@ -9985,7 +9998,7 @@ pub(crate) const fn engine_inventory() -> &'static [(&'static str, &'static str) } fn infer_engine_from_request(lower: &str) -> Option<&'static str> { - ["lemonade", "vllm"] + ["lemonade", "vllm", "aim"] .into_iter() .find(|engine| lower.contains(*engine)) } @@ -21069,6 +21082,17 @@ pub(crate) mod tests { assert_eq!(unsupported_here_for("llamacpp", false), None); } + /// `aim` is ruled out for a different reason than `vllm` on Windows + /// (Docker, not the OS), so the message has to say Docker, not WSL. + #[test] + fn unsupported_here_for_names_docker_for_aim() { + let message = unsupported_here_for("aim", true).expect("a ruled-out engine gets a message"); + assert!( + message.to_lowercase().contains("docker"), + "the message has to name the actual reason, Docker: {message:?}" + ); + } + /// Regression test for the engine-consistency bug: `assess_model_on_this_host` /// used to call `AppPaths::discover()` twice -- once (implicitly `None`) for /// the GPU summary, once for the config lookup a few lines later -- so the diff --git a/engines/aim/src/lib.rs b/engines/aim/src/lib.rs index 29a128490..6f621435b 100644 --- a/engines/aim/src/lib.rs +++ b/engines/aim/src/lib.rs @@ -208,6 +208,14 @@ pub fn builtin_handle_envelope(envelope: EngineRequestEnvelope) -> EngineRespons handle_envelope(envelope) } +/// Whether the Docker daemon is reachable on this host. +/// +/// Exposed so `apps/rocm`'s host-fit gate can rule the `aim` engine out +/// before a model is served, not just report it after the fact via `Detect`. +pub fn builtin_docker_daemon_reachable() -> bool { + crate::docker::daemon_reachable() +} + #[allow(clippy::too_many_arguments)] pub fn builtin_serve_http( service_id: String, From 18e99696eca9f3f08dfe7d8625f5ad386c9cf0de Mon Sep 17 00:00:00 2001 From: Juho Vainio Date: Mon, 19 Oct 2026 08:15:34 +0300 Subject: [PATCH 08/22] ROCMAI-555: surface a hardware-verified AIM recipe in the model catalog rocm model/rocm models only lists a recipe under a platform heading when that platform's engines array contains the recipe's first preferred engine -- no platform listed "aim", so any AIM-only recipe would have been silently invisible regardless of how many were added. Add "aim" to the existing Radeon platform (the hardware-verified image below targets gfx110X-all/gfx120X-all, the same family already listed there) and add a recipe for Qwen/Qwen3.5-9B backed by amdenterpriseai/aim-radeon-qwen-qwen3-5-9b:0.12.0-preview, the image confirmed serving real inference on Strix Halo in the e2e-cucumber AIM scenario. The catalog has no structural field for an engine's required env vars (engine_recipes has no extra_env, unlike the protocol layer's EngineRecipeHint, which serve_cmd always populates empty from a catalog recipe) -- so the mandatory AIM_PROFILE_ID is documented via the existing engine_recipes notes field instead of adding new schema, and the image reference goes in model_id_override since that's what serve_model_ref_for_engine already forwards as the engine's model_ref. Signed-off-by: Juho Vainio --- crates/rocm-core/src/model_catalog.json | 43 ++++++++++++++++++++++++- 1 file changed, 42 insertions(+), 1 deletion(-) diff --git a/crates/rocm-core/src/model_catalog.json b/crates/rocm-core/src/model_catalog.json index f970bbbbf..98727ed74 100644 --- a/crates/rocm-core/src/model_catalog.json +++ b/crates/rocm-core/src/model_catalog.json @@ -10,7 +10,7 @@ }, { "label": "AMD Radeon - RX 7000/9000, Radeon PRO (Lemonade / llama.cpp)", - "engines": ["lemonade", "llama.cpp", "llamacpp"], + "engines": ["lemonade", "llama.cpp", "llamacpp", "aim"], "gfx_families": ["gfx110X-all", "gfx120X-all"] }, { @@ -508,6 +508,47 @@ "warnings": [ "multimodal model; text chat is the default serve path" ] + }, + { + "canonical_model_id": "Qwen/Qwen3.5-9B", + "aliases": [ + "qwen3.5-9b", + "aim-qwen3.5-9b" + ], + "task": "chat", + "source": "recipe_index", + "revision": "main", + "loader": "aim", + "trust_remote_code": false, + "dtype": "auto", + "device_policy": "gpu_required", + "min_gpu_mem_gb": 24, + "recommended_system_ram_gb": 32, + "quantization": "none; served as-is by the AIM image's bundled vLLM", + "artifact_hint": "AMD Enterprise AI (AIM) pre-built Docker image; pulled and run by the aim engine, not by Hugging Face resolution", + "artifacts": [], + "engine_recipes": [ + { + "engine": "aim", + "required_flags": [], + "parser_settings": {}, + "preferred_endpoint": null, + "unsupported_combinations": [], + "notes": [ + "requires --engine-env AIM_PROFILE_ID=vllm-w7900-bf16-tp1-latency; Radeon profile auto-selection isn't reliable yet", + "requires a reachable Docker daemon; the image is pulled on first serve" + ], + "model_id_override": "amdenterpriseai/aim-radeon-qwen-qwen3-5-9b:0.12.0-preview" + } + ], + "manual_alternatives": [], + "chat_template_mode": "auto", + "preferred_engines": [ + "aim" + ], + "warnings": [ + "Radeon-targeted AIM image; requires Docker and an explicit AIM_PROFILE_ID" + ] } ] } From bed4a2f6a6ef2c3ea0be37dd8b17e028cb893f7e Mon Sep 17 00:00:00 2001 From: Juho Vainio Date: Mon, 19 Oct 2026 08:45:26 +0300 Subject: [PATCH 09/22] ROCMAI-555: correct stale AIM comments after catalog entry landed Both comments claimed AIM had no model catalog entry "in this PR's scope" -- no longer true since the catalog recipe was added. Reworded to explain why this canary still uses the raw image reference (hardware-verified pin, not catalog discovery) instead. Signed-off-by: Juho Vainio --- .../features/model_serving.feature | 20 ++++++++++--------- tests/e2e-cucumber/tests/e2e/serving_steps.rs | 6 +++--- 2 files changed, 14 insertions(+), 12 deletions(-) diff --git a/tests/e2e-cucumber/features/model_serving.feature b/tests/e2e-cucumber/features/model_serving.feature index 70cb04546..73903e2d3 100644 --- a/tests/e2e-cucumber/features/model_serving.feature +++ b/tests/e2e-cucumber/features/model_serving.feature @@ -342,15 +342,17 @@ Feature: Model serving When the user lists recommended models Then the model list output contains only ASCII characters - # AIM serve + inference canary (ROCMAI-554). AIM has no model catalog in this - # PR's scope (ROCMAI-555) — the model argument IS the raw AIM Docker image - # reference, not a short recipe name. `@requires-engine:aim` only skips here - # when Docker itself is unreachable (see `HostCapability::engine_available`'s - # "aim" arm); a Strix Halo self-hosted runner with Docker is where this - # actually runs. `@serve-timeout:1200` widens both the harness's readiness - # poll and the CLI's own ROCM_CLI_AIM_READY_TIMEOUT_SECS (see isolate_cmd) — - # a cold-start HF weight download for this 9B model legitimately exceeds the - # engine's default 5-minute readiness cap. + # AIM serve + inference canary (ROCMAI-554). This scenario deliberately still + # serves via the raw AIM Docker image reference rather than the catalog + # recipe added in ROCMAI-555 — it's a hardware-verified canary pinned to one + # exact image/profile, not a catalog-discovery test. `@requires-engine:aim` + # only skips here when Docker itself is unreachable (see + # `HostCapability::engine_available`'s "aim" arm); a Strix Halo self-hosted + # runner with Docker is where this actually runs. `@serve-timeout:1200` + # widens both the harness's readiness poll and the CLI's own + # ROCM_CLI_AIM_READY_TIMEOUT_SECS (see isolate_cmd) — a cold-start HF weight + # download for this 9B model legitimately exceeds the engine's default + # 5-minute readiness cap. @id:serve-aim-inference @requires-gpu @requires-engine:aim @serve-timeout:1200 Scenario: serve-26 - A model served on AIM responds to inference requests Given a managed runtime is active diff --git a/tests/e2e-cucumber/tests/e2e/serving_steps.rs b/tests/e2e-cucumber/tests/e2e/serving_steps.rs index c14a33cc5..0d7481bd1 100644 --- a/tests/e2e-cucumber/tests/e2e/serving_steps.rs +++ b/tests/e2e-cucumber/tests/e2e/serving_steps.rs @@ -635,9 +635,9 @@ async fn setup_lemonade_hf_checkpoint_model(world: &mut E2eWorld) { #[given("a model is being served on aim")] async fn setup_aim_model(world: &mut E2eWorld) { - // AIM has no model catalog in this PR's scope (ROCMAI-555): the model - // argument IS the raw AIM Docker image reference, not a short recipe name - // like the vLLM/Lemonade scenarios above. + // Serves via the raw AIM Docker image reference rather than the catalog + // recipe added in ROCMAI-555 — this is a hardware-verified serve+inference + // canary pinned to one exact image/profile, not a catalog-discovery test. // // The `aim-meta-llama-llama-3-2-1b-instruct` image originally used here // bundles only Instinct-series (MI300X/MI250X/MI355X/MI325X/MI350X) From 96d8532173c5ddbcfebdba7e9233ea6d7824b3c3 Mon Sep 17 00:00:00 2001 From: Juho Vainio Date: Mon, 19 Oct 2026 09:09:23 +0300 Subject: [PATCH 10/22] ROCMAI-555: add e2e coverage for AIM catalog discovery and availability Covers two of the epic's three catalog Gherkin scenarios as automated e2e-cucumber tests: AIM images listed in `rocm model` output, and a listed AIM image shown as unavailable with a reason when Docker can't be reached. Both exercise `model --verbose`, since the AIM recipe is curated out of the plain view. The third scenario (stale catalog entries) is left out: no staleness-threshold mechanism exists yet anywhere in the codebase -- that belongs to ROCMAI-556, not this ticket. Signed-off-by: Juho Vainio --- .../features/model_catalog.feature | 23 ++++++++++ tests/e2e-cucumber/tests/e2e.rs | 1 + .../tests/e2e/model_catalog_steps.rs | 46 +++++++++++++++++++ 3 files changed, 70 insertions(+) create mode 100644 tests/e2e-cucumber/features/model_catalog.feature create mode 100644 tests/e2e-cucumber/tests/e2e/model_catalog_steps.rs diff --git a/tests/e2e-cucumber/features/model_catalog.feature b/tests/e2e-cucumber/features/model_catalog.feature new file mode 100644 index 000000000..1f25a592d --- /dev/null +++ b/tests/e2e-cucumber/features/model_catalog.feature @@ -0,0 +1,23 @@ +Feature: Model catalog listing + + # ROCMAI-555: the AIM engine recipe (Qwen/Qwen3.5-9B) is curated out of the + # default `rocm model` view (it is not `featured`), so only `--verbose` + # actually lists it. `model` and its `models` alias share one parser branch, + # so exercising `model --verbose` covers the epic's "rocm model or rocm + # models" wording either way. + @id:model-catalog-lists-aim-recipe + Scenario: Discovering available AIM images via rocm model + When the user asks for the model catalog in verbose detail + Then the catalog lists an AIM engine recipe + + # The epic names two reasons a listed AIM image can be unusable: an + # unsupported GPU family, or no Docker. Only the Docker check runs at the + # catalog layer today (see `engine_ruled_out_by_platform` in main.rs); + # GPU-family fit is a `rocm serve`/`rocm diagnose --model` readiness check, + # not a catalog property. This cross-checks against the harness's own Docker + # probe rather than hardcoding "unavailable", so it stays honest on a host + # (e.g. the Strix Halo AIM runner) where Docker really is reachable. + @id:model-catalog-flags-aim-unavailable-without-docker + Scenario: Host cannot run a listed AIM image + When the user asks for the model catalog in verbose detail + Then the catalog's AIM entry matches this host's real Docker availability diff --git a/tests/e2e-cucumber/tests/e2e.rs b/tests/e2e-cucumber/tests/e2e.rs index 40df0241c..c1e937995 100644 --- a/tests/e2e-cucumber/tests/e2e.rs +++ b/tests/e2e-cucumber/tests/e2e.rs @@ -33,6 +33,7 @@ mod e2e { pub mod examine_steps; pub mod lifecycle_steps; pub mod logs_steps; + pub mod model_catalog_steps; pub mod remote_steps; pub mod runtime_lifecycle_steps; pub mod runtime_steps; diff --git a/tests/e2e-cucumber/tests/e2e/model_catalog_steps.rs b/tests/e2e-cucumber/tests/e2e/model_catalog_steps.rs new file mode 100644 index 000000000..2c509f16d --- /dev/null +++ b/tests/e2e-cucumber/tests/e2e/model_catalog_steps.rs @@ -0,0 +1,46 @@ +// Copyright © Advanced Micro Devices, Inc., or its affiliates. +// +// SPDX-License-Identifier: MIT + +use cucumber::{then, when}; + +use crate::E2eWorld; + +#[when("the user asks for the model catalog in verbose detail")] +async fn user_asks_for_verbose_catalog(world: &mut E2eWorld) { + let (stdout, _, rc) = crate::run_rocm(world, &["model", "--verbose"]); + world.cli_output = Some(stdout); + world.cli_rc = Some(rc); +} + +#[then("the catalog lists an AIM engine recipe")] +async fn catalog_lists_aim_recipe(world: &mut E2eWorld) { + let output = world.cli_output.as_ref().expect("no model catalog output"); + assert_eq!( + world.cli_rc, + Some(0), + "listing the catalog is not a failure:\n{output}" + ); + assert!( + output.contains("Qwen/Qwen3.5-9B") && output.contains("engines=[aim]"), + "expected the AIM recipe listed by canonical id and engine:\n{output}" + ); +} + +#[then("the catalog's AIM entry matches this host's real Docker availability")] +async fn catalog_aim_entry_matches_docker_availability(world: &mut E2eWorld) { + let output = world.cli_output.as_ref().expect("no model catalog output"); + let aim_available = e2e_cucumber::capability::host_capability().engine_available("aim"); + if aim_available { + assert!( + !output.contains("unsupported_docker_unavailable"), + "Docker is reachable on this host, so the AIM entry should not be flagged unavailable:\n{output}" + ); + } else { + assert!( + output.contains("unsupported_docker_unavailable") + && output.contains("Docker daemon not reachable"), + "Docker is unreachable on this host, so the AIM entry should name the reason:\n{output}" + ); + } +} From be35c7e0b7aa802c7445fa05210cc13e41e6e03d Mon Sep 17 00:00:00 2001 From: Juho Vainio Date: Mon, 19 Oct 2026 13:39:31 +0300 Subject: [PATCH 11/22] ROCMAI-554: add WSL2 GPU passthrough args to AIM's docker run Signed-off-by: Juho Vainio --- engines/aim/src/docker.rs | 59 +++++++++++++++++++++++++++++++++------ 1 file changed, 51 insertions(+), 8 deletions(-) diff --git a/engines/aim/src/docker.rs b/engines/aim/src/docker.rs index deb291f37..3ade442bf 100644 --- a/engines/aim/src/docker.rs +++ b/engines/aim/src/docker.rs @@ -46,13 +46,7 @@ pub(crate) fn run( if let Some(csv) = rocm_engine_protocol::gpu_indices_to_csv(gpu_indices) { command.arg("-e").arg(format!("HIP_VISIBLE_DEVICES={csv}")); } - command - .arg("--device=/dev/kfd") - .arg("--device=/dev/dri") - .arg("--ipc=host"); - if let Some(gid) = render_gid { - command.arg("--group-add").arg(gid.to_string()); - } + command.args(gpu_device_args(rocm_core::is_wsl_host(), render_gid)); } command.args(extra_env_args(extra_env)); @@ -72,6 +66,28 @@ pub(crate) fn run( Ok(()) } +/// GPU device/mount args for `docker run`, branched by host. WSL2 exposes the +/// GPU through the dxcore device model (`/dev/dxg`, with the host's dxcore +/// libs bind-mounted in since the container image doesn't ship them) — it has +/// no `/dev/kfd`, `/dev/dri`, or render-group gid, which are native-Linux +/// kernel concepts that don't apply there. +fn gpu_device_args(wsl: bool, render_gid: Option) -> Vec { + let mut args = vec!["--ipc=host".to_owned()]; + if wsl { + args.push("--device=/dev/dxg".to_owned()); + args.push("-v".to_owned()); + args.push("/usr/lib/wsl/lib:/usr/lib/wsl/lib".to_owned()); + } else { + args.push("--device=/dev/kfd".to_owned()); + args.push("--device=/dev/dri".to_owned()); + if let Some(gid) = render_gid { + args.push("--group-add".to_owned()); + args.push(gid.to_string()); + } + } + args +} + /// Resolves the host's `render` group gid dynamically (host-specific, so it /// cannot be hardcoded). Returns `None` if the group is absent (e.g. non-Linux). pub(crate) fn render_group_gid() -> Option { @@ -209,9 +225,36 @@ pub(crate) fn daemon_reachable() -> bool { #[cfg(test)] mod tests { - use super::extra_env_args; + use super::{extra_env_args, gpu_device_args}; use std::collections::BTreeMap; + #[test] + fn gpu_device_args_native_linux_uses_kfd_dri_and_render_gid() { + assert_eq!( + gpu_device_args(false, Some(993)), + vec![ + "--ipc=host", + "--device=/dev/kfd", + "--device=/dev/dri", + "--group-add", + "993" + ] + ); + } + + #[test] + fn gpu_device_args_wsl2_uses_dxg_device_and_wsl_lib_mount() { + assert_eq!( + gpu_device_args(true, Some(993)), + vec![ + "--ipc=host", + "--device=/dev/dxg", + "-v", + "/usr/lib/wsl/lib:/usr/lib/wsl/lib" + ] + ); + } + #[test] fn render_group_gid_parses_getent_format() { let line = "render:x:993:"; From 272916080bd9dd717869b4fa5fd7a88b616a4c1e Mon Sep 17 00:00:00 2001 From: Juho Vainio Date: Mon, 19 Oct 2026 13:40:31 +0300 Subject: [PATCH 12/22] ROCMAI-554: add unit test for AIM install's Docker-unreachable bail-out Signed-off-by: Juho Vainio --- engines/aim/src/install.rs | 71 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 71 insertions(+) diff --git a/engines/aim/src/install.rs b/engines/aim/src/install.rs index 3d6b797ed..d15f1a7cb 100644 --- a/engines/aim/src/install.rs +++ b/engines/aim/src/install.rs @@ -36,3 +36,74 @@ pub(crate) fn install_response(request: InstallRequest) -> Result Date: Mon, 19 Oct 2026 13:41:04 +0300 Subject: [PATCH 13/22] ROCMAI-555: add optional image_digest field to ModelRecipeEngineRecord Signed-off-by: Juho Vainio --- apps/rocm/src/main.rs | 1 + apps/rocm/src/serve_cmd.rs | 2 ++ crates/rocm-core/src/model_catalog.rs | 7 +++++++ 3 files changed, 10 insertions(+) diff --git a/apps/rocm/src/main.rs b/apps/rocm/src/main.rs index cc8cf545a..661062c15 100644 --- a/apps/rocm/src/main.rs +++ b/apps/rocm/src/main.rs @@ -19897,6 +19897,7 @@ pub(crate) mod tests { }], notes: vec!["metadata only; not applied to launches yet".to_owned()], model_id_override: None, + image_digest: None, }]; let mut output = String::new(); diff --git a/apps/rocm/src/serve_cmd.rs b/apps/rocm/src/serve_cmd.rs index 0e19f0c4f..30610a898 100644 --- a/apps/rocm/src/serve_cmd.rs +++ b/apps/rocm/src/serve_cmd.rs @@ -1173,6 +1173,7 @@ mod tests { ], notes: vec!["adapter hint".to_owned()], model_id_override: None, + image_digest: None, }, rocm_core::ModelRecipeEngineRecord { engine: "lemonade".to_owned(), @@ -1182,6 +1183,7 @@ mod tests { unsupported_combinations: Vec::new(), notes: Vec::new(), model_id_override: None, + image_digest: None, }, ]; diff --git a/crates/rocm-core/src/model_catalog.rs b/crates/rocm-core/src/model_catalog.rs index 5bd43f3eb..8e0a8481c 100644 --- a/crates/rocm-core/src/model_catalog.rs +++ b/crates/rocm-core/src/model_catalog.rs @@ -80,6 +80,12 @@ pub struct ModelRecipeEngineRecord { /// id for Lemonade versus a Hugging Face repo id for vLLM). #[serde(default)] pub model_id_override: Option, + /// Pins the container image this engine serves the model from to a specific + /// digest (`sha256:...`), rather than a floating tag. Currently informational + /// only (no verification against it yet); used by container-based engines + /// such as AIM. + #[serde(default)] + pub image_digest: Option, } #[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq)] @@ -945,6 +951,7 @@ mod tests { }], notes: vec!["metadata only; adapter protocol does not consume this yet".to_owned()], model_id_override: None, + image_digest: None, }); let index = ModelRecipeIndexDocument { schema_version: 1, From a786a3d7076105681c7b8e53e1a2006512733abb Mon Sep 17 00:00:00 2001 From: Juho Vainio Date: Mon, 19 Oct 2026 13:41:36 +0300 Subject: [PATCH 14/22] ROCMAI-554: add AIM lifecycle-parity e2e scenario (stop, then restart) Signed-off-by: Juho Vainio --- .../features/model_serving.feature | 17 +++++++++++ tests/e2e-cucumber/tests/e2e/serving_steps.rs | 28 +++++++++++++++++++ 2 files changed, 45 insertions(+) diff --git a/tests/e2e-cucumber/features/model_serving.feature b/tests/e2e-cucumber/features/model_serving.feature index 73903e2d3..045bb9d19 100644 --- a/tests/e2e-cucumber/features/model_serving.feature +++ b/tests/e2e-cucumber/features/model_serving.feature @@ -360,3 +360,20 @@ Feature: Model serving When the user sends a chat completion request Then the response contains a model reply And the response identifies the correct model + + # Lifecycle parity for AIM (ROCMAI-554). AIM's engine integration is new + # enough that nothing had yet proven a container-backed service can be + # stopped and relaunched like any other managed engine's — this reuses the + # same stop path the serve-retry logic and scenario teardown already rely + # on (`stop_managed_services`, via the `stop_aim_service` step), and the + # same serve+inference fixture `serve-26` uses, rather than adding a new + # stop/restart mechanism just for AIM. + @id:serve-aim-lifecycle-parity @requires-gpu @requires-engine:aim @serve-timeout:1200 + Scenario: serve-27 - A stopped AIM service can be restarted and serves inference again + Given a managed runtime is active + And a model is being served on aim + When the user stops the AIM service + And the user serves that same model on aim again + And the user sends a chat completion request + Then the response contains a model reply + And the response identifies the correct model diff --git a/tests/e2e-cucumber/tests/e2e/serving_steps.rs b/tests/e2e-cucumber/tests/e2e/serving_steps.rs index 0d7481bd1..8cb551853 100644 --- a/tests/e2e-cucumber/tests/e2e/serving_steps.rs +++ b/tests/e2e-cucumber/tests/e2e/serving_steps.rs @@ -635,6 +635,13 @@ async fn setup_lemonade_hf_checkpoint_model(world: &mut E2eWorld) { #[given("a model is being served on aim")] async fn setup_aim_model(world: &mut E2eWorld) { + serve_aim_canary(world).await; +} + +/// Serves the same hardware-verified AIM canary [`setup_aim_model`] does. +/// Factored out so the lifecycle-parity scenario can call it a second time, +/// after stopping the service, without duplicating the image/profile pin. +async fn serve_aim_canary(world: &mut E2eWorld) { // Serves via the raw AIM Docker image reference rather than the catalog // recipe added in ROCMAI-555 — this is a hardware-verified serve+inference // canary pinned to one exact image/profile, not a catalog-discovery test. @@ -669,6 +676,27 @@ async fn setup_aim_model(world: &mut E2eWorld) { .await; } +/// Stops the AIM service this scenario started and asserts the stop actually +/// took: reuses [`stop_scenario_services`] (itself `crate::stop_managed_services`, +/// the same real-record stop the serve-retry path and scenario teardown use), +/// rather than a hardcoded service id — the AIM canary's id comes from a real +/// `--managed` launch, not a planted record, so there is no fixed id to match. +#[when("the user stops the AIM service")] +async fn stop_aim_service(world: &mut E2eWorld) { + let status = stop_scenario_services(world); + assert!( + status.contains(": stopped"), + "expected the AIM service to stop cleanly, got: {status}" + ); +} + +/// Serves the same AIM canary again, proving the engine can be launched a +/// second time after a stop — the restart half of lifecycle parity. +#[when("the user serves that same model on aim again")] +async fn reserve_aim_model(world: &mut E2eWorld) { + serve_aim_canary(world).await; +} + #[given("a large model is being served on GPU")] async fn setup_large_gpu_model(world: &mut E2eWorld) { // Large-model nightly coverage follows the host's serving path. vLLM hosts From 893a36330ac2351d88eadf54b268c2166e622839 Mon Sep 17 00:00:00 2001 From: Juho Vainio Date: Thu, 22 Oct 2026 16:47:17 +0300 Subject: [PATCH 15/22] ROCMAI-554: fix CI failures from the AIM engine addition - tests/e2e-cucumber/src/capability.rs: parse_engines_list() had a "lemonade" | "vllm" allowlist that never recognized "aim", so HostCapability::engine_available("aim") always returned false even when the real `rocm engines list` output includes it and Docker is reachable. This contradicted the production catalog's own (correct) Docker-reachability check and failed model-catalog-flags-aim-unavailable-without-docker on every runner where Docker actually works (main E2E job and both Strix Halo lanes). - tests/e2e-cucumber/features/model_catalog.feature, tests/e2e-cucumber/tests/feature_naming.rs: the new model_catalog.feature scenarios were missing from FEATURE_KEYS and their names lacked the required "model-catalog-NN - " index prefix, tripping the feature-naming drift guard. - coverage-floors.toml: add the missing rocm-engine-aim floor (cargo xtask coverage --bless), which CI requires for every gated crate. Signed-off-by: Juho Vainio --- coverage-floors.toml | 1 + tests/e2e-cucumber/features/model_catalog.feature | 4 ++-- tests/e2e-cucumber/src/capability.rs | 2 +- tests/e2e-cucumber/tests/feature_naming.rs | 1 + 4 files changed, 5 insertions(+), 3 deletions(-) diff --git a/coverage-floors.toml b/coverage-floors.toml index b1b9d7a95..545139a74 100644 --- a/coverage-floors.toml +++ b/coverage-floors.toml @@ -19,6 +19,7 @@ "rocm-dash-daemon" = 82.62 # 2066 lines "rocm-dash-tui" = 89.33 # 24535 lines "rocm-deps" = 100.00 # 33 lines +"rocm-engine-aim" = 33.11 # 1090 lines "rocm-engine-lemonade" = 63.82 # 5554 lines "rocm-engine-protocol" = 98.37 # 491 lines "rocm-engine-vllm" = 66.76 # 3030 lines diff --git a/tests/e2e-cucumber/features/model_catalog.feature b/tests/e2e-cucumber/features/model_catalog.feature index 1f25a592d..fb69719d4 100644 --- a/tests/e2e-cucumber/features/model_catalog.feature +++ b/tests/e2e-cucumber/features/model_catalog.feature @@ -6,7 +6,7 @@ Feature: Model catalog listing # so exercising `model --verbose` covers the epic's "rocm model or rocm # models" wording either way. @id:model-catalog-lists-aim-recipe - Scenario: Discovering available AIM images via rocm model + Scenario: model-catalog-01 - Discovering available AIM images via rocm model When the user asks for the model catalog in verbose detail Then the catalog lists an AIM engine recipe @@ -18,6 +18,6 @@ Feature: Model catalog listing # probe rather than hardcoding "unavailable", so it stays honest on a host # (e.g. the Strix Halo AIM runner) where Docker really is reachable. @id:model-catalog-flags-aim-unavailable-without-docker - Scenario: Host cannot run a listed AIM image + Scenario: model-catalog-02 - Host cannot run a listed AIM image When the user asks for the model catalog in verbose detail Then the catalog's AIM entry matches this host's real Docker availability diff --git a/tests/e2e-cucumber/src/capability.rs b/tests/e2e-cucumber/src/capability.rs index 010935016..47b1b6146 100644 --- a/tests/e2e-cucumber/src/capability.rs +++ b/tests/e2e-cucumber/src/capability.rs @@ -672,7 +672,7 @@ fn parse_engines_list(text: &str) -> Vec { continue; } let first = trimmed.split_whitespace().next().unwrap_or(""); - if matches!(first, "lemonade" | "vllm") && !engines.iter().any(|e| e == first) { + if matches!(first, "lemonade" | "vllm" | "aim") && !engines.iter().any(|e| e == first) { engines.push(first.to_owned()); } } diff --git a/tests/e2e-cucumber/tests/feature_naming.rs b/tests/e2e-cucumber/tests/feature_naming.rs index c40fc7ec0..8070140dc 100644 --- a/tests/e2e-cucumber/tests/feature_naming.rs +++ b/tests/e2e-cucumber/tests/feature_naming.rs @@ -40,6 +40,7 @@ const FEATURE_KEYS: &[(&str, &str)] = &[ // about the ones a managed `rocm serve` leaves behind. ("local_server_records.feature", "server-records"), ("logs.feature", "logs"), + ("model_catalog.feature", "model-catalog"), ("model_serving.feature", "serve"), ("networking.feature", "networking"), ("remote.feature", "remote"), From 07c33feb16edd91e426ecd1b71bc1a6cc3de1cb8 Mon Sep 17 00:00:00 2001 From: Juho Vainio Date: Thu, 22 Oct 2026 19:35:40 +0300 Subject: [PATCH 16/22] Lower rocm-engine-aim coverage floor to match CI's measurement CI measures 32.75% line coverage for rocm-engine-aim, which falls outside the 0.2pp tolerance of the locally-blessed 33.11% floor. Rebless to CI's actual measured value. Signed-off-by: Juho Vainio --- coverage-floors.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/coverage-floors.toml b/coverage-floors.toml index 545139a74..aa2ece4ed 100644 --- a/coverage-floors.toml +++ b/coverage-floors.toml @@ -19,7 +19,7 @@ "rocm-dash-daemon" = 82.62 # 2066 lines "rocm-dash-tui" = 89.33 # 24535 lines "rocm-deps" = 100.00 # 33 lines -"rocm-engine-aim" = 33.11 # 1090 lines +"rocm-engine-aim" = 32.75 # 1090 lines "rocm-engine-lemonade" = 63.82 # 5554 lines "rocm-engine-protocol" = 98.37 # 491 lines "rocm-engine-vllm" = 66.76 # 3030 lines From f5e75426a3fabdb57b536775c6f3c7b53c3a9fbb Mon Sep 17 00:00:00 2001 From: Juho Vainio Date: Mon, 26 Oct 2026 11:30:17 +0200 Subject: [PATCH 17/22] Address PR #708 review feedback (CHANGES_REQUESTED) Fixes 10 blocking findings and 6 non-blocking findings from the automated review (rominf, review 5625108576): Blocking: - Validate --engine-env keys before launching AIM containers to reject empty/malformed names instead of silently passing them to Docker. - Reject AIM model refs that escape the configured cache root via `..` or absolute paths before mounting them into the container. - Surface Docker daemon unreachable errors during AIM install as a clear bail-out instead of a generic failure. - Stop AIM `stop`/`restart` from reporting success when the underlying `docker stop` call fails; propagate the real error. - Fix AIM lifecycle parity: `restart` now goes through the same readiness/health checks as `serve`, with an e2e scenario covering stop-then-restart. - Add unit test coverage for AIM install's Docker-unreachable bail-out. - Add e2e coverage for AIM image-digest pinning via the new optional `image_digest` field on `ModelRecipeEngineRecord`. - Correct the PR description's scope (was mis-scoped to ROCMAI-555; actual work spans ROCMAI-554/555). - Lower the rocm-engine-aim coverage floor to match CI's actual measurement instead of an aspirational number that was failing CI. - Add an e2e-cucumber fake-docker test harness so AIM subprocess behavior (success, failure, timeout) can be exercised without a real Docker daemon in CI. Non-blocking: - Warn users when --engine-env is passed to an engine that doesn't consume it, via a new engine_env_ignored_note, with unit coverage for AIM (consumes it, no note), empty input (no note), and non-consuming engines (note shown). - Redact extra_env secret values (e.g. HF_TOKEN) before persisting AIM running-state to disk, keeping keys for visibility; covered by two new unit tests confirming secrets never land on disk while keys do. - (4 further small fixes folded into the blocking commits above where they shared the same touched code.) Left as judgment calls per the review's own "Decisions for the author" section (replied on the PR, not changed in code): - extra_env as a generic passthrough vs. an engine-specific typed channel: kept generic; AIM is the only consumer today and a typed channel would be speculative for a single engine. - shelling out to the `docker` CLI vs. the Docker API: kept the CLI; it matches this repo's existing engine-process conventions and avoids a new dependency for one engine. - model weights in the container's writable layer vs. a cache volume: left as-is; moving to a volume is a real option but changes AIM's on-disk layout and warrants its own follow-up rather than riding in on this review-response commit. Also not code-changed (recorded for the reviewer, not oversights): - Running serve-26/27 once on real hardware: flagged as a manual verification step, not something a unit/e2e test can substitute for. - Success-path coverage of wait_for_container_ready/serve_http beyond what the new fake-docker harness exercises: the harness covers the failure/timeout paths that mattered most; a full fake OpenAI-compatible HTTP server for the success path is a larger addition saved for a follow-up if needed. - tests/e2e-cucumber/tests/e2e/serving_steps.rs:1527: confirmed not a bug, deliberately left unchanged. Verified: cargo build --workspace --all-targets, cargo test --workspace (one pre-existing flaky comfyui test unrelated to this change), and cargo clippy --workspace --all-targets -- -D warnings all clean. Signed-off-by: Juho Vainio --- AGENTS.md | 2 +- README.md | 11 +- apps/rocm/src/chat.rs | 6 +- apps/rocm/src/engines_cmd.rs | 13 +- apps/rocm/src/main.rs | 188 +++++-- apps/rocm/src/serve_cmd.rs | 175 ++++++- apps/rocmd/src/service.rs | 19 + crates/rocm-core/src/model_catalog.json | 18 +- crates/rocm-core/src/model_catalog.rs | 18 +- crates/rocm-dash-tui/src/ui/serve_wizard.rs | 2 +- docs/architecture.md | 12 +- docs/engine-plugins.md | 12 +- docs/llm-tool-use.md | 2 +- docs/manual-testing.md | 22 + docs/testing.md | 1 + engines/aim/src/docker.rs | 261 +++++++++- engines/aim/src/lib.rs | 108 +++- engines/aim/src/process.rs | 489 +++++++++++++++++- engines/aim/src/registry.rs | 87 +++- engines/aim/src/state.rs | 80 +++ skills/rocm-cli-assistant/SKILL.md | 2 +- tests/e2e-cucumber/Cargo.toml | 7 + .../features/model_catalog.feature | 13 +- .../features/model_serving.feature | 8 +- tests/e2e-cucumber/src/bin/fake-docker.rs | 19 + tests/e2e-cucumber/src/capability.rs | 11 +- .../tests/e2e/model_catalog_steps.rs | 40 +- 27 files changed, 1495 insertions(+), 131 deletions(-) create mode 100644 tests/e2e-cucumber/src/bin/fake-docker.rs diff --git a/AGENTS.md b/AGENTS.md index 0557ab19a..8c6b3cf9b 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -205,7 +205,7 @@ Current workspace members: - apps: `apps/rocm`, `apps/rocmd` - shared crates: `crates/rocm-core`, `crates/rocm-engine-protocol` -- engine crates: `engines/lemonade`, `engines/vllm` +- engine crates: `engines/lemonade`, `engines/vllm`, `engines/aim` Shared UI components — reuse rather than hand-rolling new ones: `apps/rocm/src/cli_progress.rs`'s `Spinner` for a caller-driven indicator that diff --git a/README.md b/README.md index 1071c097e..4992ddce7 100644 --- a/README.md +++ b/README.md @@ -618,7 +618,7 @@ rocm engines install [--runtime-id KEY] [--python-version X.Y] [--reins rocm engines shell [--runtime-id KEY | --env-id ID] [--shell PATH] ``` -Supported engines: `lemonade`, `vllm`. +Supported engines: `lemonade`, `vllm`, `aim`. ### Will a model run here? @@ -646,7 +646,7 @@ outside the catalog, this just cannot tell you in advance how it will go. Start a local OpenAI-compatible model server: ``` -rocm serve [--engine lemonade|vllm] +rocm serve [--engine lemonade|vllm|aim] [--device gpu_required|gpu_preferred] [--gpu auto|] [--runtime-id KEY | --env-id ID] @@ -655,8 +655,15 @@ rocm serve [--engine lemonade|vllm] [--no-smoke-test] [--allow-public-bind] [--temperature TEMP] [--top-p PROB] [--max-tokens N] + [--engine-env KEY=VALUE]... ``` +`--engine aim` runs the model in a Docker container instead of a host +process; it requires a reachable Docker daemon and is Linux/WSL2 only. +`--engine-env` passes an extra `KEY=VALUE` environment variable through to +the engine's own process/container (repeatable); currently only `aim` +consumes these (e.g. `AIM_PROFILE_ID`). + `--temperature` (>= 0.0), `--top-p` (0.0-1.0), and `--max-tokens` (> 0) set server-wide sampling defaults for the launched engine. They apply only to `vllm` and `lemonade`; other engines reject them. For vLLM they are folded diff --git a/apps/rocm/src/chat.rs b/apps/rocm/src/chat.rs index fb496fe33..99a277c33 100644 --- a/apps/rocm/src/chat.rs +++ b/apps/rocm/src/chat.rs @@ -323,7 +323,7 @@ fn parse_bracketed_csv(line: &str, start: &str, end: &str) -> Vec { fn is_engine_status_line(line: &str) -> bool { matches!( line.split_once(':').map(|(engine, _)| engine), - Some("lemonade" | "vllm") + Some("lemonade" | "vllm" | "aim") ) } @@ -434,7 +434,7 @@ pub(crate) fn parse_deterministic_engine_rows(tool_text: &str) -> Vec Option<&'static str> { - ["lemonade", "vllm"].into_iter().find(|engine| { + ["lemonade", "vllm", "aim"].into_iter().find(|engine| { line == *engine || line .strip_prefix(*engine) @@ -1414,7 +1414,7 @@ fn fallback_config_tool_call_for_prompt(normalized: &str) -> Option Result> { if engine_manages_own_runtime(engine) { - return env_root_for_self_managed_engine(paths, config); + return env_root_for_self_managed_engine(paths, config, engine); } env_root_for_runtime(paths, engine, runtime_id) } +/// Only `lemonade` installs into a therock-managed Python runtime; `aim` runs +/// entirely in Docker and has no env root to match against. Answering `None` +/// for it here (rather than returning whatever unrelated runtime happens to +/// be active) keeps `detect_runtime_matches_env_root` from failing aim's +/// Detect response for a reason that has nothing to do with aim. fn env_root_for_self_managed_engine( paths: &AppPaths, config: &RocmCliConfig, + engine: &str, ) -> Result> { + if engine != "lemonade" { + return Ok(None); + } recover_setup_runtime_registration(paths, config)?; let manifests = therock::load_runtime_manifests(paths)?; for selector in [ @@ -334,7 +343,7 @@ pub(crate) fn ensure_self_managed_engine_ready( return Ok(()); } let runtime_id = managed_engine_runtime_id(engine); - let env_root = env_root_for_self_managed_engine(paths, config)?; + let env_root = env_root_for_self_managed_engine(paths, config, engine)?; let detect = engine_request::<_, DetectResponse>( Some(paths), engine, diff --git a/apps/rocm/src/main.rs b/apps/rocm/src/main.rs index 661062c15..ebcc12ac3 100644 --- a/apps/rocm/src/main.rs +++ b/apps/rocm/src/main.rs @@ -3179,33 +3179,51 @@ fn assess_model_for_host( rocm_core::model_readiness::assess_model_readiness(model_ref, catalog, host, &engine_for) } -/// Whether the platform gate rules this engine out on this host. +/// Why the platform gate rules an engine out on this host, if it does. /// -/// One place, so the `/model` adapter-availability note and the readiness -/// verdict cannot answer differently about the same host and engine. -fn engine_ruled_out_by_platform(engine: &str) -> bool { - (rocm_core::runtime_is_windows() && engine.eq_ignore_ascii_case("vllm")) - || (engine.eq_ignore_ascii_case("aim") - && !rocm_engine_aim::builtin_docker_daemon_reachable()) +/// Returned by [`engine_ruled_out_by_platform`], the one place that decides +/// *why*, so the `/model` adapter-availability note and the readiness verdict +/// -- which only decide how to word it -- cannot disagree about the same host +/// and engine. +#[derive(Clone, Copy, PartialEq, Eq, Debug)] +enum PlatformRuleOut { + Allowed, + NativeWindowsVllm, + AimDockerUnavailable, +} + +/// Whether, and why, the platform gate rules this engine out on this host. +fn engine_ruled_out_by_platform(engine: &str) -> PlatformRuleOut { + if rocm_core::runtime_is_windows() && engine.eq_ignore_ascii_case("vllm") { + PlatformRuleOut::NativeWindowsVllm + } else if engine.eq_ignore_ascii_case("aim") + && !rocm_engine_aim::builtin_docker_daemon_reachable() + { + PlatformRuleOut::AimDockerUnavailable + } else { + PlatformRuleOut::Allowed + } } -/// The `unsupported_here` message for an engine, given whether the platform -/// gate already ruled it out. +/// The `unsupported_here` message for an engine, given why the platform gate +/// already ruled it out. /// /// Split out from the `engine_for` closure in [`assess_model_for_host`] so it /// can be unit-tested on both branches directly, instead of only indirectly -/// through whichever platform the test happens to run on. `ruled_out` is a -/// plain `bool` rather than re-deriving it from `engine` here, so a test can +/// through whichever platform the test happens to run on. `ruled_out` is +/// passed in rather than re-derived from `engine` here, so a test can /// exercise the "ruled out" branch on Linux CI and the "allowed" branch on a /// Windows runner without either one being unreachable. -fn unsupported_here_for(engine: &str, ruled_out: bool) -> Option { - ruled_out.then(|| { - if engine.eq_ignore_ascii_case("aim") { - format!("{engine} needs a reachable Docker daemon to serve; start Docker or choose a different engine") - } else { - format!("{engine} has no adapter on native Windows; serve it from WSL or Linux") - } - }) +fn unsupported_here_for(engine: &str, ruled_out: PlatformRuleOut) -> Option { + match ruled_out { + PlatformRuleOut::Allowed => None, + PlatformRuleOut::AimDockerUnavailable => Some(format!( + "{engine} needs a reachable Docker daemon to serve; start Docker or choose a different engine" + )), + PlatformRuleOut::NativeWindowsVllm => Some(format!( + "{engine} has no adapter on native Windows; serve it from WSL or Linux" + )), + } } /// The catalog entry a report should name, and whether a fix was offered for it. @@ -6074,17 +6092,14 @@ fn append_model_engine_support_lines( #[allow(dead_code)] fn model_registry_adapter_availability_note(engine: &str) -> Option<&'static str> { - if !engine_ruled_out_by_platform(engine) { - return None; - } - if engine.eq_ignore_ascii_case("aim") { - Some( + match engine_ruled_out_by_platform(engine) { + PlatformRuleOut::Allowed => None, + PlatformRuleOut::AimDockerUnavailable => Some( "runtime_status=unsupported_docker_unavailable reason=Docker daemon not reachable; start Docker or choose a different engine; run /engine for adapter details", - ) - } else { - Some( + ), + PlatformRuleOut::NativeWindowsVllm => Some( "runtime_status=unsupported_native_windows reason=native Windows skipped; use WSL/Linux vLLM ROCm; gpu_execution_required=true; run /engine for adapter details", - ) + ), } } @@ -9998,9 +10013,16 @@ pub(crate) const fn engine_inventory() -> &'static [(&'static str, &'static str) } fn infer_engine_from_request(lower: &str) -> Option<&'static str> { + // Word-boundary match, not substring: "aim" as a bare `.contains` also + // matches "aiming"/"claim"/"maintain", silently steering an unrelated + // chat request onto the Docker-backed AIM engine. + let words: std::collections::HashSet<&str> = lower + .split(|c: char| !c.is_ascii_alphanumeric()) + .filter(|word| !word.is_empty()) + .collect(); ["lemonade", "vllm", "aim"] .into_iter() - .find(|engine| lower.contains(*engine)) + .find(|engine| words.contains(engine)) } fn infer_model_from_request(request: &str) -> Option<&str> { @@ -12029,6 +12051,24 @@ pub(crate) mod tests { use crate::runtimes_cmd::runtime_manifest_path; use crate::test_support::write_test_pip_runtime; + /// "aim" must match as its own word, not as a substring of an unrelated + /// word like "aiming"/"claim"/"maintain" in a free-text chat request. + #[test] + fn infer_engine_from_request_does_not_match_aim_as_a_substring() { + assert_eq!( + infer_engine_from_request("serve a small model, aiming for low latency"), + None + ); + assert_eq!( + infer_engine_from_request("please explain how to maintain a claim"), + None + ); + assert_eq!( + infer_engine_from_request("serve qwen3.5-9b with aim"), + Some("aim") + ); + } + /// Serializes and isolates tests that read or mutate process-global env vars. /// /// `std::env::set_var`/`remove_var` are `unsafe` in edition 2024 because @@ -19757,7 +19797,7 @@ pub(crate) mod tests { assert!(rendered.contains("Recommended models - run one with `rocm serve `")); // Featured current models are grouped under their hardware target, each // with the quant that fits a single GPU. - assert!(rendered.contains("AMD Ryzen AI - Strix Halo (Lemonade / llama.cpp)")); + assert!(rendered.contains("AMD Ryzen AI - Strix Halo (Lemonade / llama.cpp / AIM)")); assert!(rendered.contains("AMD Instinct - MI300X, MI350X, MI355X (vLLM)")); // Strix Halo GGUF entries carry the servable owner/repo:variant ref. assert!(rendered.contains("unsloth/Qwen3.6-35B-A3B-GGUF:Q4_K_M")); @@ -19898,6 +19938,7 @@ pub(crate) mod tests { notes: vec!["metadata only; not applied to launches yet".to_owned()], model_id_override: None, image_digest: None, + extra_env: BTreeMap::new(), }]; let mut output = String::new(); @@ -21062,8 +21103,8 @@ pub(crate) mod tests { /// points at WSL/Linux -- the actual escape hatch, not a guess. #[test] fn unsupported_here_for_names_the_engine_when_ruled_out() { - let message = - unsupported_here_for("vllm", true).expect("a ruled-out engine gets a message"); + let message = unsupported_here_for("vllm", PlatformRuleOut::NativeWindowsVllm) + .expect("a ruled-out engine gets a message"); assert!( message.contains("vllm"), "the message has to name the engine it is talking about: {message:?}" @@ -21079,21 +21120,98 @@ pub(crate) mod tests { /// OS this test happens to run on. #[test] fn unsupported_here_for_is_none_when_allowed() { - assert_eq!(unsupported_here_for("vllm", false), None); - assert_eq!(unsupported_here_for("llamacpp", false), None); + assert_eq!(unsupported_here_for("vllm", PlatformRuleOut::Allowed), None); + assert_eq!( + unsupported_here_for("llamacpp", PlatformRuleOut::Allowed), + None + ); } /// `aim` is ruled out for a different reason than `vllm` on Windows /// (Docker, not the OS), so the message has to say Docker, not WSL. #[test] fn unsupported_here_for_names_docker_for_aim() { - let message = unsupported_here_for("aim", true).expect("a ruled-out engine gets a message"); + let message = unsupported_here_for("aim", PlatformRuleOut::AimDockerUnavailable) + .expect("a ruled-out engine gets a message"); assert!( message.to_lowercase().contains("docker"), "the message has to name the actual reason, Docker: {message:?}" ); } + /// Child half of [`engine_ruled_out_by_platform_gates_aim_on_docker_reachability`]. + /// + /// Runs in its own process so the hostile `PATH` (a fake `docker` that + /// always fails) can be set with `Command::env` -- mutating `PATH` + /// in-process would need `unsafe`, which the workspace denies. + #[test] + fn docker_unreachable_subprocess_helper() { + if std::env::var_os("ROCM_AIM_DOCKER_UNREACHABLE_TEST_DIR").is_none() { + return; + } + assert_eq!( + engine_ruled_out_by_platform("aim"), + PlatformRuleOut::AimDockerUnavailable, + "docker is unreachable, so the platform gate must rule aim out" + ); + let message = + unsupported_here_for("aim", engine_ruled_out_by_platform("aim")).expect("ruled out"); + assert!( + message.to_lowercase().contains("docker"), + "the gate and its message must agree: {message:?}" + ); + } + + /// Fix for the review finding that `unsupported_here_for_names_docker_for_aim` + /// only pinned the message's wording, never the real + /// `engine_ruled_out_by_platform("aim")` gate it is paired with in + /// `assess_model_for_host`. This drives the gate itself against a `docker` + /// that always fails, rather than hand-feeding it `ruled_out: true`. + #[cfg(unix)] + #[test] + fn engine_ruled_out_by_platform_gates_aim_on_docker_reachability() { + let dir = std::env::temp_dir().join(format!( + "rocm-aim-docker-unreachable-{}-{}", + std::process::id(), + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos() + )); + let fake_bin = dir.join("bin"); + std::fs::create_dir_all(&fake_bin).unwrap(); + rocm_core::fake_executable::write_fake_executable( + &fake_bin.join("docker"), + "#!/bin/sh\nexit 1\n", + ) + .unwrap(); + + let mut entries = vec![fake_bin]; + entries.extend(std::env::split_paths( + &std::env::var_os("PATH").unwrap_or_default(), + )); + + let status = std::process::Command::new(std::env::current_exe().unwrap()) + .args([ + "--exact", + "tests::docker_unreachable_subprocess_helper", + "--nocapture", + ]) + .env( + "ROCM_AIM_DOCKER_UNREACHABLE_TEST_DIR", + std::path::PathBuf::from("1"), + ) + .env("PATH", std::env::join_paths(entries).unwrap()) + .stdin(std::process::Stdio::null()) + .stdout(std::process::Stdio::null()) + .stderr(std::process::Stdio::null()) + .status() + .unwrap(); + assert!(status.success(), "helper subprocess failed"); + + std::fs::remove_dir_all(&dir).ok(); + } + /// Regression test for the engine-consistency bug: `assess_model_on_this_host` /// used to call `AppPaths::discover()` twice -- once (implicitly `None`) for /// the GPU summary, once for the config lookup a few lines later -- so the diff --git a/apps/rocm/src/serve_cmd.rs b/apps/rocm/src/serve_cmd.rs index 30610a898..fdb68adc1 100644 --- a/apps/rocm/src/serve_cmd.rs +++ b/apps/rocm/src/serve_cmd.rs @@ -21,7 +21,6 @@ //! to `service_launch.rs` in ROCMAI-91 (Phase 6b), along with its test module; //! this file's `use` list below points there now instead of at the crate root. -use std::collections::BTreeMap; use std::fmt::Write as _; use anyhow::{Result, bail}; @@ -130,16 +129,21 @@ fn serve_model_ref_for_engine( else { return model.to_owned(); }; - if let Some(override_id) = recipe + let engine_recipe = recipe .engine_recipes .iter() - .find(|engine_recipe| engine_recipe.engine.eq_ignore_ascii_case(selected_engine)) + .find(|engine_recipe| engine_recipe.engine.eq_ignore_ascii_case(selected_engine)); + let base = engine_recipe .and_then(|engine_recipe| engine_recipe.model_id_override.as_deref()) .filter(|value| !value.trim().is_empty()) - { - return override_id.to_owned(); + .map_or_else(|| recipe.canonical_model_id.clone(), ToOwned::to_owned); + // Pins a container-based engine (e.g. AIM) to the catalog's `image_digest` + // rather than a floating tag, so a `docker pull`/`run` keyed on this ref + // actually enforces it instead of the field being purely informational. + match engine_recipe.and_then(|engine_recipe| engine_recipe.image_digest.as_deref()) { + Some(digest) if !digest.trim().is_empty() => format!("{base}@{digest}"), + _ => base, } - recipe.canonical_model_id.clone() } fn serve_engine_selection_line(selection: &ServeEngineSelection) -> String { @@ -181,7 +185,7 @@ fn protocol_engine_recipe_hint( engine: engine_recipe.engine.clone(), required_flags: engine_recipe.required_flags.clone(), parser_settings: engine_recipe.parser_settings.clone(), - extra_env: BTreeMap::new(), + extra_env: engine_recipe.extra_env.clone(), preferred_endpoint: engine_recipe.preferred_endpoint.as_ref().map(|endpoint| { EngineRecipeEndpointHint { endpoint_mode: endpoint.endpoint_mode.clone(), @@ -277,8 +281,7 @@ fn engine_recipe_with_generation_defaults( engine: engine.to_owned(), ..EngineRecipeHint::default() }); - if engine.eq_ignore_ascii_case("vllm") || engine.eq_ignore_ascii_case("aim") { - // AIM images bundle vLLM, so its `--override-generation-config` flag applies unchanged. + if engine.eq_ignore_ascii_case("vllm") { let mut overrides = serde_json::Map::new(); if let Some(temperature) = defaults.temperature { overrides.insert("temperature".to_owned(), serde_json::json!(temperature)); @@ -517,6 +520,16 @@ fn engine_recipe_with_extra_env( Ok(Some(hint)) } +/// `--engine-env` note shown when the flag was passed but the selected engine +/// never consumes `EngineRecipeHint::extra_env` (only `aim` does; see +/// `engine_recipe_with_extra_env`), so the override would otherwise be silently +/// dropped with no indication to the user. +fn engine_env_ignored_note(engine: &str, engine_env: &[String]) -> Option { + (!engine_env.is_empty() && !engine.eq_ignore_ascii_case("aim")).then(|| { + format!("--engine-env applies only to engine 'aim'; ignored for engine '{engine}'") + }) +} + /// Whether the resolved engine recipe launches vLLM with tool calling enabled. fn engine_recipe_enables_tool_choice(hint: Option<&EngineRecipeHint>) -> bool { hint.is_some_and(|hint| { @@ -664,6 +677,7 @@ pub(crate) fn serve(args: ServeArgs) -> Result<()> { generation_defaults, )?; let engine_recipe = engine_recipe_with_extra_env(&selected_engine, engine_recipe, &engine_env)?; + let engine_env_note = engine_env_ignored_note(&selected_engine, &engine_env); let tool_call_note = if tool_call_parser.is_some() && !engine_serves_vllm { Some(format!( "note: --tool-call-parser applies only to vLLM; ignored for engine '{selected_engine}'" @@ -879,6 +893,9 @@ pub(crate) fn serve(args: ServeArgs) -> Result<()> { if let Some(note) = &gpu_memory_utilization_note { println!(" note: {note}"); } + if let Some(note) = &engine_env_note { + println!(" note: {note}"); + } } let managed_runtime_id = resolved_selection.runtime_id.clone(); @@ -948,7 +965,7 @@ pub(crate) fn serve(args: ServeArgs) -> Result<()> { }; spinner.clear(); - let notes = collect_serve_notes( + let mut notes = collect_serve_notes( cpu_only, &gpu_selection, rocr_visible_devices_set, @@ -958,6 +975,9 @@ pub(crate) fn serve(args: ServeArgs) -> Result<()> { host_gpu_summary.as_ref(), engine_serves_vllm, ); + if let Some(note) = &engine_env_note { + notes.push(note.clone()); + } let summary = serve_summary::DeploymentSummary { engine: selected_engine.clone(), requested_model: model, @@ -1110,6 +1130,22 @@ mod tests { ); } + #[test] + fn serve_model_ref_for_engine_pins_image_digest_when_set() { + let mut recipe = resolve_builtin_model_recipe("qwen32b").expect("qwen32b recipe"); + let vllm_recipe = recipe + .engine_recipes + .iter_mut() + .find(|engine_recipe| engine_recipe.engine == "vllm") + .expect("qwen32b has a vllm engine recipe"); + vllm_recipe.image_digest = Some("sha256:abcd1234".to_owned()); + + assert_eq!( + serve_model_ref_for_engine("qwen32b", Some(&recipe), "vllm"), + "Qwen/Qwen3-32B-FP8@sha256:abcd1234" + ); + } + #[test] fn explicit_engine_override_keeps_alias_when_shared_recipe_is_for_another_engine() { // `qwen-smoke` is a Lemonade-only GGUF recipe (no vLLM engine recipe). @@ -1127,6 +1163,22 @@ mod tests { ); } + #[test] + fn serve_model_ref_for_engine_ignores_blank_image_digest() { + let mut recipe = resolve_builtin_model_recipe("qwen32b").expect("qwen32b recipe"); + recipe + .engine_recipes + .iter_mut() + .find(|engine_recipe| engine_recipe.engine == "vllm") + .expect("vllm engine recipe") + .image_digest = Some(" ".to_owned()); + + assert_eq!( + serve_model_ref_for_engine("qwen32b", Some(&recipe), "vllm"), + "Qwen/Qwen3-32B-FP8" + ); + } + #[test] fn serve_engine_selection_respects_explicit_and_configured_engines() { let recipe = resolve_builtin_model_recipe("qwen32b").expect("qwen32b recipe"); @@ -1174,6 +1226,7 @@ mod tests { notes: vec!["adapter hint".to_owned()], model_id_override: None, image_digest: None, + extra_env: BTreeMap::new(), }, rocm_core::ModelRecipeEngineRecord { engine: "lemonade".to_owned(), @@ -1184,6 +1237,7 @@ mod tests { notes: Vec::new(), model_id_override: None, image_digest: None, + extra_env: BTreeMap::new(), }, ]; @@ -1548,6 +1602,26 @@ mod tests { ); } + #[test] + fn generation_defaults_reject_aim_instead_of_silently_dropping() { + // AIM's engine code only ever reads `extra_env`, never `required_flags` — a + // container launch has no argv to carry `--override-generation-config` + // into, unlike vLLM's own host process. Writing it into `required_flags` + // anyway would silently drop `--temperature`/`--top-p`/`--max-tokens` for + // `rocm serve --engine aim`, so this must fail fast instead. + let error = engine_recipe_with_generation_defaults( + "aim", + None, + ServeGenerationDefaults { + temperature: Some(0.5), + top_p: None, + max_tokens: None, + }, + ) + .unwrap_err(); + assert!(error.to_string().contains("aim")); + } + #[test] fn generation_defaults_translate_to_lemonade_llama_server_flags() { let hint = engine_recipe_with_generation_defaults( @@ -1602,6 +1676,87 @@ mod tests { ); } + #[test] + fn extra_env_absent_passes_hint_through_unchanged() { + assert!( + engine_recipe_with_extra_env("aim", None, &[]) + .unwrap() + .is_none() + ); + + let authored = EngineRecipeHint { + contract_version: ENGINE_RECIPE_CONTRACT_VERSION.to_owned(), + engine: "aim".to_owned(), + ..EngineRecipeHint::default() + }; + let hint = engine_recipe_with_extra_env("aim", Some(authored.clone()), &[]).unwrap(); + assert_eq!(hint, Some(authored)); + } + + #[test] + fn extra_env_synthesizes_hint_and_merges_into_authored_one() { + let hint = engine_recipe_with_extra_env("aim", None, &["AIM_PROFILE_ID=foo".to_owned()]) + .unwrap() + .expect("--engine-env should synthesize a hint"); + assert_eq!(hint.engine, "aim"); + assert_eq!( + hint.extra_env.get("AIM_PROFILE_ID"), + Some(&"foo".to_owned()) + ); + + let authored = EngineRecipeHint { + contract_version: ENGINE_RECIPE_CONTRACT_VERSION.to_owned(), + engine: "aim".to_owned(), + extra_env: BTreeMap::from([("EXISTING".to_owned(), "kept".to_owned())]), + ..EngineRecipeHint::default() + }; + let hint = + engine_recipe_with_extra_env("aim", Some(authored), &["AIM_PROFILE_ID=foo".to_owned()]) + .unwrap() + .unwrap(); + assert_eq!(hint.extra_env.get("EXISTING"), Some(&"kept".to_owned())); + assert_eq!( + hint.extra_env.get("AIM_PROFILE_ID"), + Some(&"foo".to_owned()) + ); + } + + #[test] + fn extra_env_rejects_entry_with_no_equals_sign() { + let error = + engine_recipe_with_extra_env("aim", None, &["AIM_PROFILE_ID".to_owned()]).unwrap_err(); + assert!(error.to_string().contains("KEY=VALUE")); + } + + #[test] + fn extra_env_rejects_entry_with_empty_key() { + let error = + engine_recipe_with_extra_env("aim", None, &["=some-value".to_owned()]).unwrap_err(); + assert!(error.to_string().contains("non-empty key")); + } + + #[test] + fn engine_env_ignored_note_is_none_for_aim() { + assert_eq!( + engine_env_ignored_note("aim", &["AIM_PROFILE_ID=foo".to_owned()]), + None + ); + } + + #[test] + fn engine_env_ignored_note_is_none_when_engine_env_is_empty() { + assert_eq!(engine_env_ignored_note("vllm", &[]), None); + } + + #[test] + fn engine_env_ignored_note_warns_for_engines_that_do_not_consume_it() { + let note = engine_env_ignored_note("vllm", &["FOO=bar".to_owned()]) + .expect("vllm does not consume extra_env"); + assert!(note.contains("--engine-env")); + assert!(note.contains("'aim'")); + assert!(note.contains("'vllm'")); + } + #[test] fn engine_recipe_enables_tool_choice_reflects_flags() { assert!(!engine_recipe_enables_tool_choice(None)); diff --git a/apps/rocmd/src/service.rs b/apps/rocmd/src/service.rs index 1a8c3335e..8745468ff 100644 --- a/apps/rocmd/src/service.rs +++ b/apps/rocmd/src/service.rs @@ -7,6 +7,7 @@ use rocm_core::{ AppPaths, AutomationRuntimeState, ManagedServiceRecord, RocmCliConfig, WatcherRuntimeSnapshot, builtin_watchers, daemon_binary_path, unix_time_millis, }; +use rocm_engine_protocol::{EngineMethod, StopRequest, StopResponse}; use serde_json::Value; use serde_json::json; use std::fs; @@ -48,6 +49,24 @@ pub(crate) fn stop_managed_service_with( let requested_at = rocm_core::unix_time_millis(); stopping.stop_requested_unix_ms = Some(requested_at); stopping.write()?; + // Ask the engine to stop before signalling recorded PIDs. For a + // PID-supervised engine (vLLM, Lemonade) this redundantly re-signals the + // same PID `terminate_recorded_service_pids` below already handles — + // harmless, same as `rocm services stop` (`stop_internal_managed_service`) + // already does unconditionally. For a container-backed engine like `aim`, + // this is the only thing that actually stops the container: its recorded + // PIDs belong to a short-lived launcher, not the container itself, so PID + // signalling alone leaks it running. + let _ = crate::common::engine_request::<_, StopResponse>( + paths, + &stopping.engine, + service_id, + EngineMethod::Stop, + &StopRequest { + service_id: stopping.service_id.clone(), + force: true, + }, + ); let stops = terminate_recorded_service_pids(&stopping, terminate); let pids_where = |keep: fn(&RecordedPidStop) -> bool| -> Vec { stops diff --git a/crates/rocm-core/src/model_catalog.json b/crates/rocm-core/src/model_catalog.json index 98727ed74..54b65e486 100644 --- a/crates/rocm-core/src/model_catalog.json +++ b/crates/rocm-core/src/model_catalog.json @@ -4,12 +4,12 @@ "generated_at_unix_ms": null, "platforms": [ { - "label": "AMD Ryzen AI - Strix Halo (Lemonade / llama.cpp)", - "engines": ["lemonade", "llama.cpp", "llamacpp"], + "label": "AMD Ryzen AI - Strix Halo (Lemonade / llama.cpp / AIM)", + "engines": ["lemonade", "llama.cpp", "llamacpp", "aim"], "gfx_families": ["gfx1151"] }, { - "label": "AMD Radeon - RX 7000/9000, Radeon PRO (Lemonade / llama.cpp)", + "label": "AMD Radeon - RX 7000/9000, Radeon PRO (Lemonade / llama.cpp / AIM)", "engines": ["lemonade", "llama.cpp", "llamacpp", "aim"], "gfx_families": ["gfx110X-all", "gfx120X-all"] }, @@ -510,10 +510,9 @@ ] }, { - "canonical_model_id": "Qwen/Qwen3.5-9B", + "canonical_model_id": "aim-qwen3.5-9b", "aliases": [ - "qwen3.5-9b", - "aim-qwen3.5-9b" + "amdenterpriseai/aim-radeon-qwen-qwen3-5-9b" ], "task": "chat", "source": "recipe_index", @@ -535,10 +534,13 @@ "preferred_endpoint": null, "unsupported_combinations": [], "notes": [ - "requires --engine-env AIM_PROFILE_ID=vllm-w7900-bf16-tp1-latency; Radeon profile auto-selection isn't reliable yet", + "pins AIM_PROFILE_ID=vllm-w7900-bf16-tp1-latency via extra_env; Radeon profile auto-selection isn't reliable yet, override with --engine-env if needed", "requires a reachable Docker daemon; the image is pulled on first serve" ], - "model_id_override": "amdenterpriseai/aim-radeon-qwen-qwen3-5-9b:0.12.0-preview" + "model_id_override": "amdenterpriseai/aim-radeon-qwen-qwen3-5-9b:0.12.0-preview", + "extra_env": { + "AIM_PROFILE_ID": "vllm-w7900-bf16-tp1-latency" + } } ], "manual_alternatives": [], diff --git a/crates/rocm-core/src/model_catalog.rs b/crates/rocm-core/src/model_catalog.rs index 8e0a8481c..8febf7e58 100644 --- a/crates/rocm-core/src/model_catalog.rs +++ b/crates/rocm-core/src/model_catalog.rs @@ -69,6 +69,11 @@ pub struct ModelRecipeEngineRecord { pub required_flags: Vec, #[serde(default)] pub parser_settings: BTreeMap, + /// Environment variables this engine recipe requires at launch (for example + /// AIM's `AIM_PROFILE_ID`), carried into the launched [`EngineRecipeHint`]'s + /// `extra_env` alongside any `--engine-env` CLI overrides. + #[serde(default)] + pub extra_env: BTreeMap, #[serde(default)] pub preferred_endpoint: Option, #[serde(default)] @@ -81,9 +86,9 @@ pub struct ModelRecipeEngineRecord { #[serde(default)] pub model_id_override: Option, /// Pins the container image this engine serves the model from to a specific - /// digest (`sha256:...`), rather than a floating tag. Currently informational - /// only (no verification against it yet); used by container-based engines - /// such as AIM. + /// digest (`sha256:...`), rather than a floating tag. Enforced by appending + /// `@{digest}` to the resolved model ref in `serve_model_ref_for_engine`; + /// used by container-based engines such as AIM. #[serde(default)] pub image_digest: Option, } @@ -659,6 +664,10 @@ fn validate_model_recipe_engine_record( require_nonempty(key, "engine parser setting key")?; require_nonempty(value, "engine parser setting value")?; } + for (key, value) in &engine_recipe.extra_env { + require_nonempty(key, "engine extra_env key")?; + require_nonempty(value, "engine extra_env value")?; + } if let Some(endpoint) = engine_recipe.preferred_endpoint.as_ref() { require_nonempty(&endpoint.endpoint_mode, "engine preferred endpoint mode")?; for (key, value) in &endpoint.settings { @@ -843,7 +852,7 @@ mod tests { let strix = resolve_builtin_model_recipe("qwen").expect("qwen assistant"); assert_eq!( model_recipe_target_platform_label(&strix, &platforms), - "AMD Ryzen AI - Strix Halo (Lemonade / llama.cpp)" + "AMD Ryzen AI - Strix Halo (Lemonade / llama.cpp / AIM)" ); let mi300x = resolve_builtin_model_recipe("qwen3.6-27b").expect("qwen3.6-27b"); assert_eq!( @@ -941,6 +950,7 @@ mod tests { engine: "vllm".to_owned(), required_flags: vec!["--enable-auto-tool-choice".to_owned()], parser_settings: BTreeMap::from([("reasoning_parser".to_owned(), "qwen3".to_owned())]), + extra_env: BTreeMap::new(), preferred_endpoint: Some(ModelRecipeEndpointRecord { endpoint_mode: "openai".to_owned(), settings: BTreeMap::from([("streaming".to_owned(), "true".to_owned())]), diff --git a/crates/rocm-dash-tui/src/ui/serve_wizard.rs b/crates/rocm-dash-tui/src/ui/serve_wizard.rs index daff4b3f6..03449f71a 100644 --- a/crates/rocm-dash-tui/src/ui/serve_wizard.rs +++ b/crates/rocm-dash-tui/src/ui/serve_wizard.rs @@ -38,7 +38,7 @@ use crate::ui::theme::Theme; /// Engine inventory — names mirror `apps/rocm` `engine_inventory()`. Kept /// TUI-local (a stable, small list) so this layer needs no `rocm-core` dep. -pub const ENGINES: &[&str] = &["lemonade", "vllm"]; +pub const ENGINES: &[&str] = &["lemonade", "vllm", "aim"]; /// Device-policy choices. Index 0 omits `--device` entirely (engine default); /// the rest mirror `rocm-core`'s validated `gpu_required|gpu_preferred|cpu_only`. diff --git a/docs/architecture.md b/docs/architecture.md index e2976b59d..6a889eac4 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -138,7 +138,7 @@ Modularized (EAI-8032, Phase 1 of EAI-7768's sequencing): These modules are private `mod` declarations — [`lib.rs`](../crates/e2e-report/src/lib.rs) re-exports only the selected public API surface (`XfailReport`, `evaluate_xfail`, `scenario_results_by_id`, `generate`, `RunMeta`, `PlatformReports`, `load_platform_reports`, `generate_consolidated`, `consolidated_summary_markdown`) via `pub use`, so consumers reach it through the crate root rather than through module-qualified paths like `e2e_report::parse::...`. `generate_consolidated` and `consolidated_summary_markdown` take the opaque `PlatformReports` bundle [`load_platform_reports`](../crates/e2e-report/src/consolidated.rs) returns (ROCMAI-484 / ROCm/rocm-cli#640) — bundling the parsed reports with the `inputs` they came from means a caller producing both outputs can no longer pass the two out of sync with each other, and `PlatformReport` itself went back to `pub(crate)` once the bundle was the only thing outside callers needed. This is a deliberate encapsulation choice tighter than the full-domain-extraction convention's `pub mod x;` + `pub use x::{...};` default described above — nothing outside this crate needs the module paths themselves, only the re-exported items. Unlike [`components.rs`](../crates/e2e-report/src/components.rs), which, as noted above, keeps the module graph acyclic by depending only on [`parse.rs`](../crates/e2e-report/src/parse.rs), [`command_coverage.rs`](../crates/e2e-report/src/command_coverage.rs) and [`consolidated.rs`](../crates/e2e-report/src/consolidated.rs) depend on each other — [`consolidated.rs`](../crates/e2e-report/src/consolidated.rs) calls `command_coverage::command_coverage_markdown`, and [`command_coverage.rs`](../crates/e2e-report/src/command_coverage.rs) consumes `consolidated::PlatformReport` (via its `commands`, `label` and `scenario_pass_map()`) — a back-edge the split introduces without eliminating. -### `engines/lemonade`, `engines/vllm` — inference engine adapters +### `engines/lemonade`, `engines/vllm`, `engines/aim` — inference engine adapters `engines/vllm` follows full domain extraction: @@ -159,3 +159,13 @@ These modules are private `mod` declarations — [`lib.rs`](../crates/e2e-report - [`state.rs`](../engines/lemonade/src/state.rs) — service state file I/O, device-policy/GPU-selection parsing, and port/endpoint helpers. [`lib.rs`](../engines/lemonade/src/lib.rs) keeps only CLI/envelope dispatch and the service-lifecycle orchestration that calls across those modules. + +`engines/aim` runs models as Docker containers instead of a host process: + +- [`docker.rs`](../engines/aim/src/docker.rs) — thin wrapper around the `docker` CLI binary; no `bollard`/Docker API dependency. +- [`registry.rs`](../engines/aim/src/registry.rs) — Docker Hub v2 registry client for AIM image metadata (image ref normalization, tag/digest parsing). +- [`process.rs`](../engines/aim/src/process.rs) — container run/readiness/stop lifecycle. +- [`state.rs`](../engines/aim/src/state.rs) — service state file I/O. +- [`install.rs`](../engines/aim/src/install.rs) — much smaller than vLLM's/Lemonade's `install.rs`: there is no venv, wheel, or runtime archive to manage, only a Docker-reachability check. + +[`lib.rs`](../engines/aim/src/lib.rs) keeps only CLI/envelope dispatch. diff --git a/docs/engine-plugins.md b/docs/engine-plugins.md index 9f15fd09c..497293c2c 100644 --- a/docs/engine-plugins.md +++ b/docs/engine-plugins.md @@ -18,15 +18,15 @@ The first directory is the preferred location for external adapters. Use a binary name in the form `rocm-engine-` on Linux/WSL and `rocm-engine-.exe` on Windows. -Packaged first-party adapters are `lemonade` and `vllm`. Linux/WSL-only ROCm GPU -adapters (such as `vllm`) fail explicitly on native Windows instead of selecting -a CPU fallback. +Packaged first-party adapters are `lemonade`, `vllm`, and `aim`. Linux/WSL-only +ROCm GPU adapters (such as `vllm` and `aim`) fail explicitly on native Windows +instead of selecting a CPU fallback. The engine-selecting commands (`rocm serve --engine`, `rocm engines install`/`shell`, `rocm config set-engine`/`set-default-engine`) currently accept -only the built-in `lemonade` and `vllm` engines. Discovery still lists external -plugins under the search directories above, but selecting one by name from the -CLI is not supported while the engine set is limited to the two built-ins. +only the built-in `lemonade`, `vllm`, and `aim` engines. Discovery still lists +external plugins under the search directories above, but selecting one by name +from the CLI is not supported while the engine set is limited to the built-ins. The `lemonade` adapter uses Lemonade embeddable and prefers Lemonade's `llamacpp:rocm` backend, falling back to `llamacpp:vulkan` when ROCm is diff --git a/docs/llm-tool-use.md b/docs/llm-tool-use.md index 7f3525d2d..e040c3fb8 100644 --- a/docs/llm-tool-use.md +++ b/docs/llm-tool-use.md @@ -82,7 +82,7 @@ before running it: ``` The assistant can request installing a serving engine through the existing -engine surface. The supported engines are `lemonade` and `vllm`: +engine surface. The supported engines are `lemonade`, `vllm`, and `aim`: ```json {"name":"rocm_command","arguments":{"args":["engines","install","vllm"],"reason":"Install the vLLM serving engine."}} diff --git a/docs/manual-testing.md b/docs/manual-testing.md index d1604e2e2..c7afeb6d0 100644 --- a/docs/manual-testing.md +++ b/docs/manual-testing.md @@ -721,3 +721,25 @@ Expected result: `PATH change left in place` and names it, and the `PATH` change stays. - With `rocm uninstall --keep-binaries`, the review says the `PATH` change stays too, and it does. + +## 9. AIM (Docker) Engine Verification + +Linux/WSL2 only, and requires a reachable Docker daemon (`docker info` must +succeed). + +```bash +rocm serve Qwen/Qwen3.5-9B --engine aim --managed +rocm services list --all +rocm services stop qwen3.5-9b +``` + +Expected result: + +- The command pulls/starts an AIM Docker container (first run takes longer + while the image is pulled) and the model is reachable on the printed port + once the service shows `ready`. +- `rocm services list --all` shows the service's engine as `aim`. +- Stopping the service stops and removes the container; `docker ps -a` no + longer lists it. +- On native Windows, the same command fails explicitly instead of falling + back to another engine. diff --git a/docs/testing.md b/docs/testing.md index d4432a028..ee4298d11 100644 --- a/docs/testing.md +++ b/docs/testing.md @@ -742,6 +742,7 @@ cargo test -p rocm --bin rocm update_report_policy cargo test -p rocm-engine-lemonade stdio_protocol_routes_all_methods_without_side_effects cargo test -p rocm-engine-vllm stdio_protocol_routes_all_methods_without_side_effects cargo test -p rocm-engine-vllm +cargo test -p rocm-engine-aim cargo test -p rocmd event_collector cargo test -p rocmd event_dispatcher ``` diff --git a/engines/aim/src/docker.rs b/engines/aim/src/docker.rs index 3ade442bf..8d9ea76a0 100644 --- a/engines/aim/src/docker.rs +++ b/engines/aim/src/docker.rs @@ -99,8 +99,11 @@ pub(crate) fn render_group_gid() -> Option { if !output.status.success() { return None; } - let text = String::from_utf8_lossy(&output.stdout); - // `getent group render` prints `render:x::`. + parse_getent_gid(&String::from_utf8_lossy(&output.stdout)) +} + +// `getent group render` prints `render:x::`. +fn parse_getent_gid(text: &str) -> Option { text.trim().split(':').nth(2)?.parse().ok() } @@ -213,19 +216,21 @@ pub(crate) fn pull(image: &str) -> Result<()> { Ok(()) } +/// Bounds `docker info` so a wedged daemon cannot hang this probe (and every +/// caller that gates on it, like `rocm model --verbose`/`rocm diagnose +/// --model`) for the OS's own TCP/socket-connect timeout. +const DAEMON_PROBE_TIMEOUT: Duration = Duration::from_secs(2); + pub(crate) fn daemon_reachable() -> bool { - Command::new(DOCKER_BIN) - .arg("info") - .stdin(Stdio::null()) - .stdout(Stdio::null()) - .stderr(Stdio::null()) - .status() - .is_ok_and(|status| status.success()) + let mut command = Command::new(DOCKER_BIN); + command.arg("info"); + rocm_core::process::run_with_timeout(command, DAEMON_PROBE_TIMEOUT, "docker info") + .is_ok_and(|output| output.status.success()) } #[cfg(test)] mod tests { - use super::{extra_env_args, gpu_device_args}; + use super::{extra_env_args, gpu_device_args, parse_getent_gid}; use std::collections::BTreeMap; #[test] @@ -257,9 +262,9 @@ mod tests { #[test] fn render_group_gid_parses_getent_format() { - let line = "render:x:993:"; - let gid: Option = line.trim().split(':').nth(2).and_then(|s| s.parse().ok()); - assert_eq!(gid, Some(993)); + assert_eq!(parse_getent_gid("render:x:993:"), Some(993)); + assert_eq!(parse_getent_gid("render:x:not-a-number:"), None); + assert_eq!(parse_getent_gid(""), None); } #[test] @@ -275,3 +280,233 @@ mod tests { ); } } + +/// Exercises the `Command`-invoking functions (`run`/`stop`/`inspect_*`/`pull`/ +/// `image_present`) against a fake `docker` on `PATH`, mirroring `install.rs`'s +/// subprocess pattern: mutating `PATH` for the *current* process needs +/// `unsafe` (denied workspace-wide), so each scenario re-spawns this test +/// binary filtered to just itself, with a hostile `PATH` set via the safe +/// `Command::env` on that child. +#[cfg(all(test, unix))] +mod subprocess_tests { + use super::{image_present, inspect_exit_code, inspect_running, pull, run, stop}; + use rocm_core::fake_executable::write_fake_executable; + use std::collections::BTreeMap; + use std::process::{Command, Stdio}; + use std::time::Duration; + + const SENTINEL: &str = "ROCM_AIM_DOCKER_SUBPROCESS_TEST"; + + /// One parameterized fake `docker`, driven entirely by inherited env vars, + /// reused by every scenario below instead of planting N bespoke scripts. + const FAKE_DOCKER_SCRIPT: &str = r#"#!/bin/sh +case "$1" in + run) + [ -n "$FAKE_DOCKER_RUN_STDERR" ] && echo "$FAKE_DOCKER_RUN_STDERR" >&2 + exit "${FAKE_DOCKER_RUN_EXIT:-0}" + ;; + kill|stop) + if [ -n "$FAKE_DOCKER_STOP_NO_SUCH_CONTAINER" ]; then + echo "Error response from daemon: No such container: fake" >&2 + exit 1 + fi + [ -n "$FAKE_DOCKER_STOP_STDERR" ] && echo "$FAKE_DOCKER_STOP_STDERR" >&2 + exit "${FAKE_DOCKER_STOP_EXIT:-0}" + ;; + inspect) + case "$3" in + *Running*) echo "${FAKE_DOCKER_RUNNING:-false}" ;; + *ExitCode*) echo "${FAKE_DOCKER_EXIT_CODE:-0}" ;; + esac + exit 0 + ;; + image) + exit "${FAKE_DOCKER_IMAGE_PRESENT_EXIT:-1}" + ;; + pull) + exit "${FAKE_DOCKER_PULL_EXIT:-0}" + ;; + logs) + [ -n "$FAKE_DOCKER_LOGS" ] && echo "$FAKE_DOCKER_LOGS" + exit 0 + ;; + info) + exit "${FAKE_DOCKER_INFO_EXIT:-0}" + ;; +esac +exit 0 +"#; + + fn in_subprocess() -> bool { + std::env::var_os(SENTINEL).is_some() + } + + /// Plants the fake `docker`, re-execs this test binary with just + /// `test_path` selected and `extra_env` set, and asserts it succeeded. + fn run_with_fake_docker(test_path: &str, extra_env: &[(&str, &str)]) { + let dir = std::env::temp_dir().join(format!( + "aim-docker-test-{}-{}", + std::process::id(), + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos() + )); + let fake_bin = dir.join("bin"); + std::fs::create_dir_all(&fake_bin).unwrap(); + write_fake_executable(&fake_bin.join("docker"), FAKE_DOCKER_SCRIPT).unwrap(); + + let mut entries = vec![fake_bin]; + entries.extend(std::env::split_paths( + &std::env::var_os("PATH").unwrap_or_default(), + )); + + let mut command = Command::new(std::env::current_exe().unwrap()); + command + .args(["--exact", test_path, "--nocapture"]) + .env(SENTINEL, "1") + .env("PATH", std::env::join_paths(entries).unwrap()) + .stdin(Stdio::null()); + for (key, value) in extra_env { + command.env(key, value); + } + let status = command.status().unwrap(); + assert!(status.success(), "helper subprocess failed for {test_path}"); + + std::fs::remove_dir_all(&dir).ok(); + } + + #[test] + fn run_succeeds_against_fake_docker() { + if in_subprocess() { + run( + "fake-container", + "fake-image", + "127.0.0.1", + 8080, + 8000, + &[], + None, + &BTreeMap::new(), + ) + .expect("fake docker run should succeed"); + return; + } + run_with_fake_docker( + "docker::subprocess_tests::run_succeeds_against_fake_docker", + &[], + ); + } + + #[test] + fn run_bails_with_container_name_and_detail_on_failure() { + if in_subprocess() { + let error = run( + "fake-container", + "fake-image", + "127.0.0.1", + 8080, + 8000, + &[], + None, + &BTreeMap::new(), + ) + .expect_err("fake docker run should fail"); + let message = error.to_string(); + assert!(message.contains("fake-container"), "{message}"); + assert!(message.contains("boom"), "{message}"); + return; + } + run_with_fake_docker( + "docker::subprocess_tests::run_bails_with_container_name_and_detail_on_failure", + &[ + ("FAKE_DOCKER_RUN_EXIT", "1"), + ("FAKE_DOCKER_RUN_STDERR", "boom"), + ], + ); + } + + #[test] + fn stop_force_kill_and_graceful_stop_both_succeed() { + if in_subprocess() { + assert!(stop("fake-container", true, Duration::from_secs(1)).unwrap()); + assert!(stop("fake-container", false, Duration::from_secs(1)).unwrap()); + return; + } + run_with_fake_docker( + "docker::subprocess_tests::stop_force_kill_and_graceful_stop_both_succeed", + &[], + ); + } + + #[test] + fn stop_missing_container_returns_false_not_error() { + if in_subprocess() { + let stopped = stop("fake-container", true, Duration::from_secs(1)).unwrap(); + assert!(!stopped); + return; + } + run_with_fake_docker( + "docker::subprocess_tests::stop_missing_container_returns_false_not_error", + &[("FAKE_DOCKER_STOP_NO_SUCH_CONTAINER", "1")], + ); + } + + #[test] + fn stop_generic_failure_bails_with_detail() { + if in_subprocess() { + let error = + stop("fake-container", true, Duration::from_secs(1)).expect_err("should fail"); + assert!(error.to_string().contains("disk full"), "{error}"); + return; + } + run_with_fake_docker( + "docker::subprocess_tests::stop_generic_failure_bails_with_detail", + &[ + ("FAKE_DOCKER_STOP_EXIT", "1"), + ("FAKE_DOCKER_STOP_STDERR", "disk full"), + ], + ); + } + + #[test] + fn inspect_running_and_exit_code_parse_fake_docker_output() { + if in_subprocess() { + assert_eq!(inspect_running("fake-container").unwrap(), Some(true)); + assert_eq!(inspect_exit_code("fake-container").unwrap(), Some(137)); + return; + } + run_with_fake_docker( + "docker::subprocess_tests::inspect_running_and_exit_code_parse_fake_docker_output", + &[ + ("FAKE_DOCKER_RUNNING", "true"), + ("FAKE_DOCKER_EXIT_CODE", "137"), + ], + ); + } + + #[test] + fn pull_bails_on_failure() { + if in_subprocess() { + let error = pull("fake-image").expect_err("fake docker pull should fail"); + assert!(error.to_string().contains("fake-image"), "{error}"); + return; + } + run_with_fake_docker( + "docker::subprocess_tests::pull_bails_on_failure", + &[("FAKE_DOCKER_PULL_EXIT", "1")], + ); + } + + #[test] + fn image_present_reports_true_when_fake_docker_reports_it() { + if in_subprocess() { + assert!(image_present("fake-image").unwrap()); + return; + } + run_with_fake_docker( + "docker::subprocess_tests::image_present_reports_true_when_fake_docker_reports_it", + &[("FAKE_DOCKER_IMAGE_PRESENT_EXIT", "0")], + ); + } +} diff --git a/engines/aim/src/lib.rs b/engines/aim/src/lib.rs index 6f621435b..273a73bc8 100644 --- a/engines/aim/src/lib.rs +++ b/engines/aim/src/lib.rs @@ -114,7 +114,7 @@ pub(crate) struct ServeHttpRequest { pub fn run_cli() -> Result<()> { let cli = Cli::parse(); match cli.command { - CommandKind::Detect => print_json(&detect_response()), + CommandKind::Detect => print_json(&detect_response(None)), CommandKind::Capabilities => print_json(&capabilities()), CommandKind::Install { runtime_id, @@ -248,8 +248,8 @@ pub fn builtin_serve_http( fn handle_envelope(envelope: EngineRequestEnvelope) -> EngineResponseEnvelope { match envelope.method { EngineMethod::Detect => { - deserialize_and_respond::(envelope.payload, |_request| { - Ok(detect_response()) + deserialize_and_respond::(envelope.payload, |request| { + Ok(detect_response(request.runtime_id)) }) } EngineMethod::Capabilities => EngineResponseEnvelope::success(capabilities()), @@ -295,7 +295,11 @@ where } } -fn detect_response() -> DetectResponse { +/// `env_id` echoes back `runtime_id` when Docker is reachable, matching what +/// `install_response` reports for the same `runtime_id` -- aim has no actual +/// install step beyond "Docker is present", so Detect and Install have to +/// agree on readiness or every serve re-runs Install for nothing. +fn detect_response(runtime_id: Option) -> DetectResponse { let reachable = crate::docker::daemon_reachable(); let mut notes = Vec::new(); if !reachable { @@ -306,7 +310,7 @@ fn detect_response() -> DetectResponse { } DetectResponse { installed: reachable, - env_id: None, + env_id: reachable.then_some(runtime_id).flatten(), runtime_kind: Some("docker".to_owned()), runtime_executable: Some("docker".to_owned()), managed_env: Some(false), @@ -315,8 +319,10 @@ fn detect_response() -> DetectResponse { transformers_version: None, available_devices: vec![EngineDeviceAvailability { kind: "rocm_gpu".to_owned(), - available: reachable, - reason: if reachable { + available: reachable && !cfg!(windows), + reason: if cfg!(windows) { + Some(windows_gpu_unsupported_message().to_owned()) + } else if reachable { None } else { Some("docker daemon unreachable".to_owned()) @@ -327,10 +333,91 @@ fn detect_response() -> DetectResponse { } } +#[cfg(all(test, unix))] +mod detect_tests { + use super::detect_response; + use rocm_core::fake_executable::write_fake_executable; + use std::path::PathBuf; + use std::process::{Command, Stdio}; + + const DOCKER_REACHABLE_TEST_DIR: &str = "ROCM_AIM_DOCKER_REACHABLE_TEST_DIR"; + + /// Child half of [`detect_response_echoes_runtime_id_as_env_id_when_docker_is_reachable`]. + /// + /// Runs in its own process so the `PATH` override (a fake `docker` that + /// always succeeds) can be set with `Command::env` -- mutating `PATH` + /// in-process would need `unsafe`, which the workspace denies. + #[test] + fn docker_reachable_subprocess_helper() { + if std::env::var_os(DOCKER_REACHABLE_TEST_DIR).is_none() { + return; + } + let response = detect_response(Some("aim".to_owned())); + assert_eq!( + response.env_id, + Some("aim".to_owned()), + "env_id must echo the requested runtime_id once Detect reports installed, or \ + `ensure_self_managed_engine_ready` will treat aim as never installed" + ); + } + + /// Fix for the review finding that `ensure_self_managed_engine_ready` + /// re-runs Install on every aim serve: Detect always reported `env_id: + /// None`, which could never equal the `runtime_id` the caller compares it + /// against. + #[test] + fn detect_response_echoes_runtime_id_as_env_id_when_docker_is_reachable() { + let dir = std::env::temp_dir().join(format!( + "aim-detect-docker-reachable-{}-{}", + std::process::id(), + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos() + )); + let fake_bin = dir.join("bin"); + std::fs::create_dir_all(&fake_bin).unwrap(); + write_fake_executable(&fake_bin.join("docker"), "#!/bin/sh\nexit 0\n").unwrap(); + + let mut entries = vec![fake_bin]; + entries.extend(std::env::split_paths( + &std::env::var_os("PATH").unwrap_or_default(), + )); + + let status = Command::new(std::env::current_exe().unwrap()) + .args([ + "--exact", + "detect_tests::docker_reachable_subprocess_helper", + "--nocapture", + ]) + .env(DOCKER_REACHABLE_TEST_DIR, PathBuf::from("1")) + .env("PATH", std::env::join_paths(entries).unwrap()) + .stdin(Stdio::null()) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .status() + .unwrap(); + assert!(status.success(), "helper subprocess failed"); + + std::fs::remove_dir_all(&dir).ok(); + } + + #[test] + fn detect_response_has_no_env_id_when_runtime_id_is_absent() { + // Standalone CLI mode (`rocm-engine-aim detect`) passes no runtime_id; + // there is nothing to echo back regardless of Docker reachability. + assert_eq!(detect_response(None).env_id, None); + } +} + pub(crate) fn capabilities() -> EngineCapabilities { EngineCapabilities { cpu: true, - rocm_gpu: true, + // `crate::docker::gpu_device_args` only knows WSL2 (`/dev/dxg`) and + // assumed-native-Linux (`/dev/kfd`/`/dev/dri`) device shapes; native + // Windows has neither, so GPU passthrough is unimplemented there + // (mirrors `engines/vllm`'s `rocm_gpu: !cfg!(windows)`). + rocm_gpu: !cfg!(windows), openai_compatible: true, tool_calling: false, quantized_models: "image-reported".to_owned(), @@ -338,6 +425,11 @@ pub(crate) fn capabilities() -> EngineCapabilities { } } +pub(crate) const fn windows_gpu_unsupported_message() -> &'static str { + "AIM GPU passthrough is supported by rocm-cli only on Linux/WSL; native Windows AIM GPU \ +serving is skipped." +} + fn read_request() -> Result { let mut input = String::new(); std::io::stdin() diff --git a/engines/aim/src/process.rs b/engines/aim/src/process.rs index c933f727a..360ecf782 100644 --- a/engines/aim/src/process.rs +++ b/engines/aim/src/process.rs @@ -37,8 +37,11 @@ fn resolve_aim_ready_timeout(override_value: Option) -> Duration { /// Deterministic, Docker-name-sanitized container identity for a service. /// This *is* the service's identity for the AIM engine — no PID/start-ticks /// tracking, since `docker run -d` leaves no long-lived client process to -/// supervise (see plan: SIGKILL to a foreground `docker run` client orphans -/// the container, so this engine never signals a client at all). +/// supervise. Unlike vLLM/Lemonade, this engine never signals a client +/// process at all: the container is stopped by name (`docker stop`/`kill`), +/// which sidesteps the SIGKILL-orphan hazard those engines have to guard +/// against (a killed `docker run` client would otherwise leave the container +/// running, with nothing left to track it by). pub(crate) fn container_name(service_id: &str) -> String { let sanitized: String = service_id .chars() @@ -156,6 +159,7 @@ pub(crate) fn launch_service(request: LaunchRequest) -> Result { &request.model_ref, &request.host, request.port, + &device_policy, &gpu_indices, engine_recipe.as_ref().map(|hint| &hint.extra_env), )?; @@ -192,46 +196,105 @@ fn run_container( model_ref: &str, host: &str, port: u16, + device_policy: &DevicePolicy, gpu_indices: &[u32], extra_env: Option<&std::collections::BTreeMap>, ) -> Result> { let name = container_name(service_id); - let metadata = crate::registry::fetch_image_metadata(model_ref).unwrap_or_default(); + // Normalize once and reuse everywhere below: metadata and the actual + // `docker pull`/`run` must resolve the exact same image, or the fetched + // accelerator/GPU-arch metadata silently describes a different image + // than the one that runs (ROCMAI-554 review finding). + let image_ref = crate::registry::normalize_image_ref(model_ref); + let metadata = crate::registry::fetch_image_metadata(&image_ref).unwrap_or_else(|error| { + eprintln!( + "Warning: could not fetch AIM image metadata for `{image_ref}`: {error:#}; \ + assuming GPU-preferred with no accelerator-arch check." + ); + crate::registry::AimImageMetadata::default() + }); if crate::docker::inspect_running(&name)?.unwrap_or(false) { return Ok(metadata.canonical_name); } - if !crate::docker::image_present(model_ref)? { - crate::docker::pull(model_ref)?; + if !crate::docker::image_present(&image_ref)? { + crate::docker::pull(&image_ref)?; } - let use_gpu = metadata.accelerator_type.as_deref() != Some("cpu") && !gpu_indices.is_empty(); + // `crate::docker::gpu_device_args` has no native-Windows device shape, and + // an image reporting AIM_ACCELERATOR_TYPE=cpu has nothing to pass GPU args + // to, so neither counts as "a GPU is usable" regardless of what the caller + // selected. + let gpu_usable = !cfg!(windows) && !gpu_indices.is_empty(); + let image_wants_cpu = metadata.accelerator_type.as_deref() == Some("cpu"); + let use_gpu = resolve_use_gpu( + service_id, + &image_ref, + device_policy, + gpu_usable, + image_wants_cpu, + )?; let render_gid = if use_gpu { crate::docker::render_group_gid() } else { None }; - let empty_env = std::collections::BTreeMap::default(); + let mut env = extra_env.cloned().unwrap_or_default(); + // AIM images bundle vLLM's OpenAI server, so it is `VLLM_API_KEY` — the + // same variable `engines/vllm` sets on its own child process — that this + // engine's actual server enforces. The container is the only process this + // engine ever runs, so without this the require-api-key setting a public + // `rocm serve --engine aim` was given is silently never enforced. + if let Some(api_key) = rocm_engine_protocol::resolve_endpoint_api_key() { + env.insert("VLLM_API_KEY".to_owned(), api_key); + } crate::docker::run( &name, - model_ref, + &image_ref, host, port, crate::DEFAULT_AIM_PORT, if use_gpu { gpu_indices } else { &[] }, render_gid, - extra_env.unwrap_or(&empty_env), + &env, )?; Ok(metadata.canonical_name) } +/// Enforces `device_policy` instead of treating GPU use as a soft heuristic: +/// `CpuOnly` never requests a GPU, and `GpuRequired` fails fast instead of +/// silently falling back to a CPU launch when no GPU turns out to be usable. +fn resolve_use_gpu( + service_id: &str, + model_ref: &str, + device_policy: &DevicePolicy, + gpu_usable: bool, + image_wants_cpu: bool, +) -> Result { + Ok(match device_policy { + DevicePolicy::CpuOnly => false, + DevicePolicy::GpuRequired => { + if !gpu_usable || image_wants_cpu { + bail!( + "AIM service `{service_id}` has device_policy=gpu_required but no GPU is \ + usable for `{model_ref}` (native Windows has no GPU passthrough, no GPU \ + indices were selected, or the image reports AIM_ACCELERATOR_TYPE=cpu)" + ); + } + true + } + DevicePolicy::GpuPreferred => gpu_usable && !image_wants_cpu, + }) +} + pub(crate) fn serve_http(request: ServeHttpRequest) -> Result<()> { let canonical_name = run_container( &request.service_id, &request.model_ref, &request.host, request.port, + &request.device_policy, &request.gpu_indices, request.engine_recipe.as_ref().map(|hint| &hint.extra_env), )?; @@ -260,14 +323,35 @@ pub(crate) fn serve_http(request: ServeHttpRequest) -> Result<()> { expected_model_id, timeout, ) { - Ok(()) => Ok(()), + Ok(()) => block_until_container_stops(&name, &request.state_path), Err(error) => { + // A readiness timeout (unlike an early container exit) leaves the + // container running; stop it so a failed launch doesn't leak an + // orphaned container nothing will ever signal again (the service + // record is about to be marked `failed`, not `running`, so + // `apps/rocmd`'s own stop path will never find it). + let _ = crate::docker::stop(&name, true, STOP_GRACE); crate::state::write_terminal_state(&request.state_path, "failed")?; Err(error) } } } +/// Blocks for the container's remaining lifetime, mirroring `engines/vllm` +/// and `engines/lemonade`'s blocking `child.wait()`: `apps/rocmd`'s generic +/// supervisor (`supervise_service`) treats this process's exit as the +/// service's exit, so returning right after readiness (as this used to) +/// made a still-running, actively-serving container read `stopped` the +/// instant it became ready. A transient `docker inspect` failure is treated +/// as "still running" rather than ending the loop, matching this file's +/// other `unwrap_or`-style tolerance of CLI hiccups elsewhere. +fn block_until_container_stops(name: &str, state_path: &std::path::Path) -> Result<()> { + while crate::docker::inspect_running(name).unwrap_or(Some(true)) != Some(false) { + std::thread::sleep(READY_POLL_INTERVAL); + } + crate::state::write_terminal_state(state_path, "stopped") +} + /// Mirrors `engines/vllm/src/process.rs`'s `wait_for_vllm_ready` architecture: /// early-exit check, monotonic timeout, listed-then-inference-verified gate. /// `docker::inspect_running` replaces `child.try_wait()` since there is no @@ -403,4 +487,389 @@ mod tests { ); assert!(parse_device_policy_arg("bogus").is_err()); } + + #[test] + fn resolve_use_gpu_cpu_only_never_uses_gpu() { + assert!(!resolve_use_gpu("svc", "model", &DevicePolicy::CpuOnly, true, false).unwrap()); + } + + #[test] + fn resolve_use_gpu_gpu_required_uses_gpu_when_usable() { + assert!(resolve_use_gpu("svc", "model", &DevicePolicy::GpuRequired, true, false).unwrap()); + } + + #[test] + fn resolve_use_gpu_gpu_required_fails_fast_when_no_gpu_usable() { + assert!(resolve_use_gpu("svc", "model", &DevicePolicy::GpuRequired, false, false).is_err()); + } + + #[test] + fn resolve_use_gpu_gpu_required_fails_fast_when_image_wants_cpu() { + assert!(resolve_use_gpu("svc", "model", &DevicePolicy::GpuRequired, true, true).is_err()); + } + + #[test] + fn resolve_use_gpu_gpu_preferred_falls_back_to_cpu_without_failing() { + assert!( + !resolve_use_gpu("svc", "model", &DevicePolicy::GpuPreferred, false, false).unwrap() + ); + assert!(!resolve_use_gpu("svc", "model", &DevicePolicy::GpuPreferred, true, true).unwrap()); + assert!(resolve_use_gpu("svc", "model", &DevicePolicy::GpuPreferred, true, false).unwrap()); + } +} + +/// `run_container`/`wait_for_container_ready`/`serve_http`/`stop_service` all +/// shell out to `docker`, so (same reasoning as `docker.rs`'s +/// `subprocess_tests`) each scenario below re-spawns this test binary with a +/// fake `docker` on `PATH`. `fetch_image_metadata` (called unconditionally by +/// `run_container`) hits the real Docker Hub registry for a deliberately +/// bogus image and is expected to 404; that failure is already handled by a +/// production fallback (`unwrap_or_else` to `AimImageMetadata::default()`), +/// so no mock/seam is needed, just the real ~1s network round trip. +#[cfg(all(test, unix))] +mod subprocess_tests { + use super::*; + use rocm_core::fake_executable::write_fake_executable; + use std::process::{Command, Stdio}; + + const SENTINEL: &str = "ROCM_AIM_PROCESS_SUBPROCESS_TEST"; + + const FAKE_DOCKER_SCRIPT: &str = r#"#!/bin/sh +case "$1" in + run) + [ -n "$FAKE_DOCKER_RUN_STDERR" ] && echo "$FAKE_DOCKER_RUN_STDERR" >&2 + exit "${FAKE_DOCKER_RUN_EXIT:-0}" + ;; + kill|stop) + if [ -n "$FAKE_DOCKER_STOP_NO_SUCH_CONTAINER" ]; then + echo "Error response from daemon: No such container: fake" >&2 + exit 1 + fi + exit "${FAKE_DOCKER_STOP_EXIT:-0}" + ;; + inspect) + case "$3" in + *Running*) echo "${FAKE_DOCKER_RUNNING:-false}" ;; + *ExitCode*) echo "${FAKE_DOCKER_EXIT_CODE:-0}" ;; + esac + exit 0 + ;; + image) + exit "${FAKE_DOCKER_IMAGE_PRESENT_EXIT:-1}" + ;; + pull) + exit "${FAKE_DOCKER_PULL_EXIT:-0}" + ;; + logs) + [ -n "$FAKE_DOCKER_LOGS" ] && echo "$FAKE_DOCKER_LOGS" + exit 0 + ;; +esac +exit 0 +"#; + + fn in_subprocess() -> bool { + std::env::var_os(SENTINEL).is_some() + } + + fn run_with_fake_docker(test_path: &str, extra_env: &[(&str, &str)]) { + let dir = std::env::temp_dir().join(format!( + "aim-process-test-{}-{}", + std::process::id(), + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos() + )); + let fake_bin = dir.join("bin"); + std::fs::create_dir_all(&fake_bin).unwrap(); + write_fake_executable(&fake_bin.join("docker"), FAKE_DOCKER_SCRIPT).unwrap(); + + let mut entries = vec![fake_bin]; + entries.extend(std::env::split_paths( + &std::env::var_os("PATH").unwrap_or_default(), + )); + + let mut command = Command::new(std::env::current_exe().unwrap()); + command + .args(["--exact", test_path, "--nocapture"]) + .env(SENTINEL, "1") + .env("PATH", std::env::join_paths(entries).unwrap()) + .stdin(Stdio::null()); + for (key, value) in extra_env { + command.env(key, value); + } + let status = command.status().unwrap(); + assert!(status.success(), "helper subprocess failed for {test_path}"); + + std::fs::remove_dir_all(&dir).ok(); + } + + const BOGUS_IMAGE: &str = "rocm-cli-test/does-not-exist-on-docker-hub-aim-fixture"; + + #[test] + fn run_container_skips_docker_run_when_already_running() { + if in_subprocess() { + let result = run_container( + "svc-already-running", + BOGUS_IMAGE, + "127.0.0.1", + 18080, + &DevicePolicy::GpuPreferred, + &[], + None, + ); + assert!(result.is_ok(), "{result:?}"); + return; + } + run_with_fake_docker( + "process::subprocess_tests::run_container_skips_docker_run_when_already_running", + &[ + ("FAKE_DOCKER_RUNNING", "true"), + ("FAKE_DOCKER_PULL_EXIT", "1"), + ("FAKE_DOCKER_RUN_EXIT", "1"), + ], + ); + } + + #[test] + fn run_container_pulls_when_image_absent_then_runs() { + if in_subprocess() { + let result = run_container( + "svc-pull-then-run", + BOGUS_IMAGE, + "127.0.0.1", + 18081, + &DevicePolicy::GpuPreferred, + &[], + None, + ); + assert!(result.is_ok(), "{result:?}"); + return; + } + run_with_fake_docker( + "process::subprocess_tests::run_container_pulls_when_image_absent_then_runs", + &[ + ("FAKE_DOCKER_IMAGE_PRESENT_EXIT", "1"), + ("FAKE_DOCKER_PULL_EXIT", "0"), + ("FAKE_DOCKER_RUN_EXIT", "0"), + ], + ); + } + + #[test] + fn run_container_bails_when_pull_fails() { + if in_subprocess() { + let error = run_container( + "svc-pull-fails", + BOGUS_IMAGE, + "127.0.0.1", + 18082, + &DevicePolicy::GpuPreferred, + &[], + None, + ) + .expect_err("pull failure should bail"); + assert!(error.to_string().contains("pull"), "{error}"); + return; + } + run_with_fake_docker( + "process::subprocess_tests::run_container_bails_when_pull_fails", + &[ + ("FAKE_DOCKER_IMAGE_PRESENT_EXIT", "1"), + ("FAKE_DOCKER_PULL_EXIT", "1"), + ], + ); + } + + #[test] + fn run_container_bails_when_docker_run_fails() { + if in_subprocess() { + let error = run_container( + "svc-run-fails", + BOGUS_IMAGE, + "127.0.0.1", + 18083, + &DevicePolicy::GpuPreferred, + &[], + None, + ) + .expect_err("docker run failure should bail"); + let message = error.to_string(); + assert!(message.contains("rocm-aim-svc-run-fails"), "{message}"); + assert!(message.contains("boom"), "{message}"); + return; + } + run_with_fake_docker( + "process::subprocess_tests::run_container_bails_when_docker_run_fails", + &[ + ("FAKE_DOCKER_IMAGE_PRESENT_EXIT", "0"), + ("FAKE_DOCKER_RUN_EXIT", "1"), + ("FAKE_DOCKER_RUN_STDERR", "boom"), + ], + ); + } + + #[test] + fn wait_for_container_ready_bails_when_container_exits_early() { + if in_subprocess() { + let error = wait_for_container_ready( + "fake-container", + "127.0.0.1", + 19999, + "model-x", + Duration::from_secs(5), + ) + .expect_err("early container exit should bail"); + let message = error.to_string(); + assert!(message.contains("exited early"), "{message}"); + assert!(message.contains('7'), "{message}"); + assert!(message.contains("widget crashed"), "{message}"); + return; + } + run_with_fake_docker( + "process::subprocess_tests::wait_for_container_ready_bails_when_container_exits_early", + &[ + ("FAKE_DOCKER_RUNNING", "false"), + ("FAKE_DOCKER_EXIT_CODE", "7"), + ("FAKE_DOCKER_LOGS", "widget crashed"), + ], + ); + } + + #[test] + fn wait_for_container_ready_bails_on_timeout() { + if in_subprocess() { + // Unbound port: the readiness HTTP probe fails fast (connection + // refused), so this only exercises the timeout bookkeeping, not a + // real OpenAI-compatible server (left as a judgment call — see the + // PR summary). + let error = wait_for_container_ready( + "fake-container", + "127.0.0.1", + 59123, + "model-x", + Duration::from_millis(10), + ) + .expect_err("timeout should bail"); + assert!( + error.to_string().contains("did not become ready"), + "{error}" + ); + return; + } + run_with_fake_docker( + "process::subprocess_tests::wait_for_container_ready_bails_on_timeout", + &[("FAKE_DOCKER_RUNNING", "true")], + ); + } + + #[test] + fn serve_http_stops_container_and_marks_state_failed_on_readiness_failure() { + if in_subprocess() { + let dir = std::env::temp_dir().join(format!( + "aim-serve-http-test-{}", + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos() + )); + let state_path = dir.join("svc.json"); + + let request = ServeHttpRequest { + service_id: "svc-serve-http-fails".to_owned(), + model_ref: BOGUS_IMAGE.to_owned(), + host: "127.0.0.1".to_owned(), + port: 18084, + device_policy: DevicePolicy::GpuPreferred, + gpu_indices: vec![], + runtime_id: None, + env_id: None, + state_path: state_path.clone(), + log_path: None, + engine_recipe: None, + }; + serve_http(request).expect_err("readiness failure should surface as an error"); + + let state = crate::state::read_service_state(&state_path).unwrap(); + assert_eq!(state.get("status").and_then(|v| v.as_str()), Some("failed")); + + std::fs::remove_dir_all(&dir).ok(); + return; + } + run_with_fake_docker( + "process::subprocess_tests::serve_http_stops_container_and_marks_state_failed_on_readiness_failure", + &[ + ("FAKE_DOCKER_RUNNING", "false"), + ("FAKE_DOCKER_EXIT_CODE", "1"), + ], + ); + } + + #[test] + fn stop_service_force_stop_succeeds_and_writes_stopped_state() { + if in_subprocess() { + let response = stop_service(StopRequest { + service_id: "svc-stop-succeeds".to_owned(), + force: true, + }) + .unwrap(); + assert!(response.stopped); + assert!(!response.graceful); + + let state_path = crate::state::state_path("svc-stop-succeeds"); + let state = crate::state::read_service_state(&state_path).unwrap(); + assert_eq!( + state.get("status").and_then(|v| v.as_str()), + Some("stopped") + ); + return; + } + let dir = std::env::temp_dir().join(format!( + "aim-stop-service-test-{}", + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos() + )); + std::fs::create_dir_all(&dir).unwrap(); + run_with_fake_docker( + "process::subprocess_tests::stop_service_force_stop_succeeds_and_writes_stopped_state", + &[("ROCM_CLI_DATA_DIR", dir.to_str().unwrap())], + ); + std::fs::remove_dir_all(&dir).ok(); + } + + #[test] + fn stop_service_missing_container_reports_not_stopped_without_writing_state() { + if in_subprocess() { + let response = stop_service(StopRequest { + service_id: "svc-stop-missing".to_owned(), + force: true, + }) + .unwrap(); + assert!(!response.stopped); + assert!(!response.graceful); + + let state_path = crate::state::state_path("svc-stop-missing"); + assert!(!state_path.exists()); + return; + } + let dir = std::env::temp_dir().join(format!( + "aim-stop-service-missing-test-{}", + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos() + )); + std::fs::create_dir_all(&dir).unwrap(); + run_with_fake_docker( + "process::subprocess_tests::stop_service_missing_container_reports_not_stopped_without_writing_state", + &[ + ("ROCM_CLI_DATA_DIR", dir.to_str().unwrap()), + ("FAKE_DOCKER_STOP_NO_SUCH_CONTAINER", "1"), + ], + ); + std::fs::remove_dir_all(&dir).ok(); + } } diff --git a/engines/aim/src/registry.rs b/engines/aim/src/registry.rs index aef439ff4..47b74d9ce 100644 --- a/engines/aim/src/registry.rs +++ b/engines/aim/src/registry.rs @@ -23,16 +23,34 @@ pub(crate) struct AimImageMetadata { pub canonical_name: Option, } -/// Splits `org/repo:tag` (or bare `repo:tag`, defaulting to [`DEFAULT_ORG`]) -/// into `(repo, tag)` where `repo` includes the org prefix Docker Hub expects. -fn split_image_ref(image_ref: &str) -> (String, String) { - let (name, tag) = image_ref.split_once(':').unwrap_or((image_ref, "latest")); - let repo = if name.contains('/') { - name.to_owned() +/// Normalizes a user- or catalog-supplied image reference so the registry +/// metadata lookup and the real `docker pull`/`run` always resolve the exact +/// same image. A bare (org-less) name defaults to [`DEFAULT_ORG`], matching +/// this engine's own catalog convention — Docker's own bare-name default +/// (`library/`) would silently point at an unrelated official image instead. +pub(crate) fn normalize_image_ref(image_ref: &str) -> String { + if image_ref.contains('/') { + image_ref.to_owned() } else { - format!("{DEFAULT_ORG}/{name}") - }; - (repo, tag.to_owned()) + format!("{DEFAULT_ORG}/{image_ref}") + } +} + +/// Splits a normalized (see [`normalize_image_ref`]) `org/repo[:tag|@digest]` +/// ref into `(repo, reference)`. Treats `@sha256:...` as a digest suffix, and +/// only looks for a tag's `:` after the last `/`, so a `host:port/repo` or +/// `repo@sha256:...` ref isn't mis-split by a naive first-`:` search. +fn split_image_ref(image_ref: &str) -> (String, String) { + if let Some((name, digest)) = image_ref.split_once('@') { + return (name.to_owned(), digest.to_owned()); + } + let tag_sep = image_ref + .rfind(':') + .filter(|&i| i > image_ref.rfind('/').unwrap_or(0)); + match tag_sep { + Some(i) => (image_ref[..i].to_owned(), image_ref[i + 1..].to_owned()), + None => (image_ref.to_owned(), "latest".to_owned()), + } } fn agent() -> ureq::Agent { @@ -60,10 +78,11 @@ fn anonymous_pull_token(repo: &str) -> Result { } pub(crate) fn fetch_image_metadata(image_ref: &str) -> Result { - let (repo, tag) = split_image_ref(image_ref); + let normalized = normalize_image_ref(image_ref); + let (repo, reference) = split_image_ref(&normalized); let token = anonymous_pull_token(&repo)?; - let manifest_url = format!("https://{REGISTRY_HOST}/v2/{repo}/manifests/{tag}"); + let manifest_url = format!("https://{REGISTRY_HOST}/v2/{repo}/manifests/{reference}"); let manifest: Value = agent() .get(&manifest_url) .set("User-Agent", "rocm-cli") @@ -132,13 +151,22 @@ mod tests { use super::*; #[test] - fn split_image_ref_defaults_org_for_bare_repo() { + fn normalize_image_ref_defaults_org_for_bare_name() { assert_eq!( - split_image_ref("aim-meta-llama-llama-3-2-1b-instruct:0.11.1"), - ( - "amdenterpriseai/aim-meta-llama-llama-3-2-1b-instruct".to_owned(), - "0.11.1".to_owned() - ) + normalize_image_ref("aim-meta-llama-llama-3-2-1b-instruct:0.11.1"), + "amdenterpriseai/aim-meta-llama-llama-3-2-1b-instruct:0.11.1" + ); + } + + #[test] + fn normalize_image_ref_preserves_explicit_org_or_registry() { + assert_eq!( + normalize_image_ref("amdenterpriseai/aim-foo:latest"), + "amdenterpriseai/aim-foo:latest" + ); + assert_eq!( + normalize_image_ref("myregistry.example.com:5000/aim-foo:latest"), + "myregistry.example.com:5000/aim-foo:latest" ); } @@ -158,6 +186,31 @@ mod tests { ); } + #[test] + fn split_image_ref_handles_digest_refs() { + assert_eq!( + split_image_ref("amdenterpriseai/aim-foo@sha256:abcd1234"), + ( + "amdenterpriseai/aim-foo".to_owned(), + "sha256:abcd1234".to_owned() + ) + ); + } + + #[test] + fn split_image_ref_handles_registry_host_with_port() { + // The first `:` here belongs to the registry host's port, not a tag + // separator — a naive `split_once(':')` would mis-parse this as + // repo="myregistry.example.com", tag="5000/aim-foo:v1". + assert_eq!( + split_image_ref("myregistry.example.com:5000/aim-foo:v1"), + ( + "myregistry.example.com:5000/aim-foo".to_owned(), + "v1".to_owned() + ) + ); + } + #[test] fn parse_config_blob_extracts_known_fields() { let config = serde_json::json!({ diff --git a/engines/aim/src/state.rs b/engines/aim/src/state.rs index 9d8fd4796..31adc2453 100644 --- a/engines/aim/src/state.rs +++ b/engines/aim/src/state.rs @@ -90,6 +90,15 @@ pub(crate) fn state_path(service_id: &str) -> PathBuf { ) } +/// Replaces every `extra_env` value with a redaction marker, keeping the keys. +fn redact_engine_recipe_extra_env(hint: &EngineRecipeHint) -> EngineRecipeHint { + let mut hint = hint.clone(); + for value in hint.extra_env.values_mut() { + *value = "[redacted]".to_owned(); + } + hint +} + pub(crate) fn write_running_state( service_id: &str, model_ref: &str, @@ -99,6 +108,11 @@ pub(crate) fn write_running_state( engine_recipe: Option<&EngineRecipeHint>, ) -> Result<()> { let endpoint = crate::process::endpoint_url(host, port); + // `extra_env` can carry secrets (e.g. an HF token passed via `--engine-env`); + // this state file is status/introspection only (never read back to relaunch + // the container, see `crate::process::launch_service`), so redact values + // before writing it to disk while keeping the keys for visibility. + let engine_recipe = engine_recipe.map(redact_engine_recipe_extra_env); write_state( &state_path(service_id), &json!({ @@ -230,6 +244,72 @@ mod tests { ); } + #[test] + fn healthcheck_reports_unknown_for_a_service_with_no_state_file() { + let response = healthcheck_service(HealthcheckRequest { + service_id: format!("svc-never-tracked-{}", current_unix_millis()), + }) + .unwrap(); + assert!(!response.model_loaded); + assert_eq!(response.status, "unknown"); + assert_eq!(response.device, "unknown"); + } + + #[test] + fn redact_engine_recipe_extra_env_masks_values_but_keeps_keys() { + let hint = EngineRecipeHint { + contract_version: "0.1.0".to_owned(), + engine: "aim".to_owned(), + extra_env: std::collections::BTreeMap::from([( + "HF_TOKEN".to_owned(), + "hf_super_secret".to_owned(), + )]), + ..EngineRecipeHint::default() + }; + let redacted = redact_engine_recipe_extra_env(&hint); + assert_eq!( + redacted.extra_env.get("HF_TOKEN").map(String::as_str), + Some("[redacted]") + ); + } + + #[test] + fn write_running_state_never_persists_extra_env_secrets_to_disk() { + let dir = std::env::temp_dir().join(format!( + "rocm-aim-state-secret-test-{}-{}", + std::process::id(), + current_unix_millis() + )); + let service_id = format!("svc-secret-{}", current_unix_millis()); + let hint = EngineRecipeHint { + contract_version: "0.1.0".to_owned(), + engine: "aim".to_owned(), + extra_env: std::collections::BTreeMap::from([( + "HF_TOKEN".to_owned(), + "hf_super_secret".to_owned(), + )]), + ..EngineRecipeHint::default() + }; + let path = state_path(&service_id); + // `write_running_state` always writes under `state_path(service_id)`, which + // lives outside `dir`; only `dir` is cleaned up below, matching the other + // tests in this module that leave the real state dir alone. + write_running_state( + &service_id, + "some/model", + "127.0.0.1", + 8000, + DevicePolicy::GpuPreferred, + Some(&hint), + ) + .unwrap(); + let on_disk = fs::read_to_string(&path).unwrap(); + assert!(!on_disk.contains("hf_super_secret")); + assert!(on_disk.contains("HF_TOKEN")); + let _ = fs::remove_file(&path); + let _ = fs::remove_dir_all(&dir); + } + #[test] fn write_and_read_terminal_state_round_trips() { let dir = std::env::temp_dir().join(format!( diff --git a/skills/rocm-cli-assistant/SKILL.md b/skills/rocm-cli-assistant/SKILL.md index 0b54a2208..66af24a85 100644 --- a/skills/rocm-cli-assistant/SKILL.md +++ b/skills/rocm-cli-assistant/SKILL.md @@ -24,7 +24,7 @@ Use this skill when answering ROCm CLI local assistant questions. ## Engines And Assistant -- vLLM and Lemonade are serving engines. +- vLLM, Lemonade, and AIM (Docker-backed) are serving engines. - The built-in assistant is fixed to qwen served by Lemonade with GPU required. Do not switch the built-in assistant to vLLM. - Installing an engine and running a model server are different states. Answer each separately when the user asks. - `rocm serve` accepts `--gpu auto|` to pick the AMD GPU. `auto` (default) prefers a GPU that looks idle from VRAM telemetry (`amd-smi`, or the amdgpu DRM sysfs counters when `amd-smi` is absent) and is not already used by another rocm-cli server (managed or foreground), falling back to the GPU with the most free memory; a single index pins one GPU. Serving one model across multiple GPUs is not supported. Do not suggest CPU fallback when a GPU is busy or out of range. diff --git a/tests/e2e-cucumber/Cargo.toml b/tests/e2e-cucumber/Cargo.toml index fe2265a92..09259b2a4 100644 --- a/tests/e2e-cucumber/Cargo.toml +++ b/tests/e2e-cucumber/Cargo.toml @@ -33,6 +33,13 @@ path = "src/bin/fake-uv.rs" name = "fake-python" path = "src/bin/fake-python.rs" +# Stand-in for `docker` that always fails, planted by the model-catalog +# scenario proving the AIM engine's Docker-unavailable gate deterministically +# (see `src/bin/fake-docker.rs`). +[[bin]] +name = "fake-docker" +path = "src/bin/fake-docker.rs" + [dependencies] axum.workspace = true cucumber = { version = "0.23", features = ["output-json", "output-junit"] } diff --git a/tests/e2e-cucumber/features/model_catalog.feature b/tests/e2e-cucumber/features/model_catalog.feature index fb69719d4..5831995ac 100644 --- a/tests/e2e-cucumber/features/model_catalog.feature +++ b/tests/e2e-cucumber/features/model_catalog.feature @@ -1,6 +1,6 @@ Feature: Model catalog listing - # ROCMAI-555: the AIM engine recipe (Qwen/Qwen3.5-9B) is curated out of the + # ROCMAI-555: the AIM engine recipe (aim-qwen3.5-9b) is curated out of the # default `rocm model` view (it is not `featured`), so only `--verbose` # actually lists it. `model` and its `models` alias share one parser branch, # so exercising `model --verbose` covers the epic's "rocm model or rocm @@ -21,3 +21,14 @@ Feature: Model catalog listing Scenario: model-catalog-02 - Host cannot run a listed AIM image When the user asks for the model catalog in verbose detail Then the catalog's AIM entry matches this host's real Docker availability + + # model-catalog-02 cross-checks the harness's own `docker info` probe, which + # only proves the two probes agree, not that the gating logic itself is + # correct — it stays green even if the "no Docker" branch were deleted + # entirely on a Docker-reachable host. This plants a `docker` that always + # fails so the unavailable branch is exercised regardless of host state. + @id:model-catalog-flags-aim-unavailable-with-no-docker-daemon + Scenario: model-catalog-03 - AIM image is flagged unavailable with no Docker daemon + Given no Docker daemon is reachable + When the user asks for the model catalog in verbose detail + Then the catalog's AIM entry is flagged unavailable for lack of Docker diff --git a/tests/e2e-cucumber/features/model_serving.feature b/tests/e2e-cucumber/features/model_serving.feature index 045bb9d19..25fed56b7 100644 --- a/tests/e2e-cucumber/features/model_serving.feature +++ b/tests/e2e-cucumber/features/model_serving.feature @@ -346,9 +346,11 @@ Feature: Model serving # serves via the raw AIM Docker image reference rather than the catalog # recipe added in ROCMAI-555 — it's a hardware-verified canary pinned to one # exact image/profile, not a catalog-discovery test. `@requires-engine:aim` - # only skips here when Docker itself is unreachable (see - # `HostCapability::engine_available`'s "aim" arm); a Strix Halo self-hosted - # runner with Docker is where this actually runs. `@serve-timeout:1200` + # skips here when Docker is unreachable OR the gfx family isn't one aim is + # verified on (see `HostCapability::engine_available`'s "aim" arm — Radeon + # gfx110X/120X and Strix Halo gfx1151, not Instinct); a Strix Halo + # self-hosted runner with Docker is where this actually runs. + # `@serve-timeout:1200` # widens both the harness's readiness poll and the CLI's own # ROCM_CLI_AIM_READY_TIMEOUT_SECS (see isolate_cmd) — a cold-start HF weight # download for this 9B model legitimately exceeds the engine's default diff --git a/tests/e2e-cucumber/src/bin/fake-docker.rs b/tests/e2e-cucumber/src/bin/fake-docker.rs new file mode 100644 index 000000000..d6e55cf18 --- /dev/null +++ b/tests/e2e-cucumber/src/bin/fake-docker.rs @@ -0,0 +1,19 @@ +// Copyright © Advanced Micro Devices, Inc., or its affiliates. +// +// SPDX-License-Identifier: MIT + +//! A stand-in for `docker` that always fails, for scenarios proving the AIM +//! engine's Docker-unavailable gate. +//! +//! `daemon_reachable()` (engines/aim) and the harness's own Docker probe +//! (`capability::docker_available`) both just run `docker info` and check the +//! exit status, so a `docker` that always exits non-zero deterministically +//! simulates "no reachable daemon" regardless of whether the host running the +//! suite actually has Docker installed. +//! +//! A Rust binary rather than a shell script so the scenarios run on Windows +//! too. + +fn main() { + std::process::exit(1); +} diff --git a/tests/e2e-cucumber/src/capability.rs b/tests/e2e-cucumber/src/capability.rs index 47b1b6146..3f4a03727 100644 --- a/tests/e2e-cucumber/src/capability.rs +++ b/tests/e2e-cucumber/src/capability.rs @@ -164,7 +164,16 @@ impl HostCapability { !self.os_family.eq_ignore_ascii_case("windows") && family_prefers_vllm(self.gfx_target.as_deref()) } - "aim" => self.available_engines.iter().any(|e| e == engine) && docker_available(), + "aim" => { + self.available_engines.iter().any(|e| e == engine) + && docker_available() + // aim's verified hardware is Radeon (gfx110X/120X) and + // Strix Halo (gfx1151) — exactly the families + // `family_prefers_vllm` excludes. Without this, a + // Docker-equipped Instinct lane would run the Radeon aim + // image with a W7900 profile. + && !family_prefers_vllm(self.gfx_target.as_deref()) + } _ => self.available_engines.iter().any(|e| e == engine), } } diff --git a/tests/e2e-cucumber/tests/e2e/model_catalog_steps.rs b/tests/e2e-cucumber/tests/e2e/model_catalog_steps.rs index 2c509f16d..d05baf84d 100644 --- a/tests/e2e-cucumber/tests/e2e/model_catalog_steps.rs +++ b/tests/e2e-cucumber/tests/e2e/model_catalog_steps.rs @@ -2,13 +2,16 @@ // // SPDX-License-Identifier: MIT -use cucumber::{then, when}; +use cucumber::{given, then, when}; +use e2e_cucumber::fixture_binary::fixture_target_name; use crate::E2eWorld; #[when("the user asks for the model catalog in verbose detail")] async fn user_asks_for_verbose_catalog(world: &mut E2eWorld) { - let (stdout, _, rc) = crate::run_rocm(world, &["model", "--verbose"]); + // `run_rocm_with_scenario_env` so a `Given` step that armed `command_env` + // (e.g. a fake-docker PATH) takes effect; it is a no-op otherwise. + let (stdout, _, rc) = crate::run_rocm_with_scenario_env(world, &["model", "--verbose"]); world.cli_output = Some(stdout); world.cli_rc = Some(rc); } @@ -22,7 +25,7 @@ async fn catalog_lists_aim_recipe(world: &mut E2eWorld) { "listing the catalog is not a failure:\n{output}" ); assert!( - output.contains("Qwen/Qwen3.5-9B") && output.contains("engines=[aim]"), + output.contains("aim-qwen3.5-9b") && output.contains("engines=[aim]"), "expected the AIM recipe listed by canonical id and engine:\n{output}" ); } @@ -44,3 +47,34 @@ async fn catalog_aim_entry_matches_docker_availability(world: &mut E2eWorld) { ); } } + +/// Plant a `docker` on `PATH` that always fails, so the catalog's +/// Docker-unavailable gate is exercised deterministically regardless of +/// whether this host actually has Docker. +#[given("no Docker daemon is reachable")] +async fn no_docker_daemon_reachable(world: &mut E2eWorld) { + let root = world + .isolated_root + .as_ref() + .expect("scenario root") + .path() + .to_path_buf(); + let bin_dir = root.join("fake-bin"); + let installed = bin_dir.join(fixture_target_name("docker")); + crate::install_fixture_binary("fake-docker", &installed); + + let mut path = std::ffi::OsString::from(&bin_dir); + path.push(if cfg!(windows) { ";" } else { ":" }); + path.push(std::env::var_os("PATH").unwrap_or_default()); + world.command_env.push(("PATH", path)); +} + +#[then("the catalog's AIM entry is flagged unavailable for lack of Docker")] +async fn catalog_aim_entry_flagged_unavailable(world: &mut E2eWorld) { + let output = world.cli_output.as_ref().expect("no model catalog output"); + assert!( + output.contains("unsupported_docker_unavailable") + && output.contains("Docker daemon not reachable"), + "expected the AIM entry to name the Docker-unavailable reason:\n{output}" + ); +} From 2f614d542a748b2d2f8e692994b87d6bf6a286aa Mon Sep 17 00:00:00 2001 From: Juho Vainio Date: Mon, 26 Oct 2026 14:09:31 +0200 Subject: [PATCH 18/22] Address PR #708 second review: AIM container lifecycle, secrets, served name Containers are now removed on stop, on a failed start and before a re-run, so recovery and restart can reuse the name. serve_http follows the exit code (a crash or vanished container is failed, a deliberate stop is stopped) and tolerates only a bounded run of docker inspect errors. Env values reach docker through a 0600 env file instead of argv; a bare --engine-env KEY forwards from the environment and is never stored. Image metadata is read from the pulled image, with Docker Hub as an optional hint, and the model name the container serves is reported as served_model_id and used for readiness, smoke test, summary and liveness. GPU planning honours --device and resolves an empty selection like vLLM; native Windows rules AIM out. rocmd's engine Stop result is reported and decides confirmation for container engines. image_digest is validated, the shipped AIM entry is pinned, and AIM is listed in the remaining engine inventories. Adds an argv-recording fake docker and serve_http lifecycle tests, and removes the host-dependent model-catalog-02 scenario. Signed-off-by: Juho Vainio --- Cargo.lock | 1 + apps/rocm/src/chat.rs | 47 +- apps/rocm/src/chat_host_facts.rs | 13 +- apps/rocm/src/main.rs | 36 +- apps/rocm/src/serve_cmd.rs | 133 +- apps/rocm/src/service_launch.rs | 6 +- apps/rocmd/src/common.rs | 4 + apps/rocmd/src/service.rs | 200 ++- crates/rocm-core/src/model_catalog.json | 1 + crates/rocm-core/src/model_catalog.rs | 71 +- crates/rocm-core/src/net.rs | 16 +- crates/rocm-core/src/service_record.rs | 52 + .../src/engine_registry.rs | 15 +- crates/rocm-engine-protocol/src/lib.rs | 18 + docs/architecture.md | 6 +- docs/manual-testing.md | 35 +- docs/testing.md | 1 + engines/aim/Cargo.toml | 1 + engines/aim/src/docker.rs | 660 ++++++++-- engines/aim/src/lib.rs | 33 +- engines/aim/src/process.rs | 1171 ++++++++++++----- engines/aim/src/registry.rs | 127 +- engines/aim/src/state.rs | 101 +- engines/lemonade/src/lib.rs | 1 + engines/vllm/src/process.rs | 1 + .../features/model_catalog.feature | 19 +- tests/e2e-cucumber/tests/e2e/examine_steps.rs | 2 +- .../tests/e2e/model_catalog_steps.rs | 18 - 28 files changed, 2144 insertions(+), 645 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 59882919b..d80da7bfd 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3807,6 +3807,7 @@ dependencies = [ "rocm-engine-protocol", "serde", "serde_json", + "tempfile", "ureq", ] diff --git a/apps/rocm/src/chat.rs b/apps/rocm/src/chat.rs index 99a277c33..73678e436 100644 --- a/apps/rocm/src/chat.rs +++ b/apps/rocm/src/chat.rs @@ -1414,17 +1414,16 @@ fn fallback_config_tool_call_for_prompt(normalized: &str) -> Option, - /// Extra `KEY=VALUE` environment variable to pass through to the engine's - /// own process/container (repeatable). Engine-specific; currently only - /// `--engine aim` reads these (e.g. `AIM_PROFILE_ID`). Ignored by engines - /// that don't consume engine-env overrides. - #[arg(long = "engine-env", value_name = "KEY=VALUE")] + /// Extra environment variable for the engine's container (repeatable). + /// Engine-specific; currently only `--engine aim` reads these (e.g. + /// `AIM_PROFILE_ID=...`) and other engines ignore it. `KEY=VALUE` is + /// recorded with the service, so for a secret such as `HF_TOKEN` pass a + /// bare `KEY` instead: it is forwarded from your environment and never + /// stored. + #[arg(long = "engine-env", value_name = "KEY[=VALUE]")] engine_env: Vec, }, /// [preview] Install, start, stop, or inspect ComfyUI. @@ -3189,6 +3191,8 @@ fn assess_model_for_host( enum PlatformRuleOut { Allowed, NativeWindowsVllm, + /// AIM's Docker GPU passthrough exists only for Linux and WSL2. + NativeWindowsAim, AimDockerUnavailable, } @@ -3196,6 +3200,8 @@ enum PlatformRuleOut { fn engine_ruled_out_by_platform(engine: &str) -> PlatformRuleOut { if rocm_core::runtime_is_windows() && engine.eq_ignore_ascii_case("vllm") { PlatformRuleOut::NativeWindowsVllm + } else if rocm_core::runtime_is_windows() && engine.eq_ignore_ascii_case("aim") { + PlatformRuleOut::NativeWindowsAim } else if engine.eq_ignore_ascii_case("aim") && !rocm_engine_aim::builtin_docker_daemon_reachable() { @@ -3220,7 +3226,7 @@ fn unsupported_here_for(engine: &str, ruled_out: PlatformRuleOut) -> Option Some(format!( "{engine} needs a reachable Docker daemon to serve; start Docker or choose a different engine" )), - PlatformRuleOut::NativeWindowsVllm => Some(format!( + PlatformRuleOut::NativeWindowsVllm | PlatformRuleOut::NativeWindowsAim => Some(format!( "{engine} has no adapter on native Windows; serve it from WSL or Linux" )), } @@ -6100,6 +6106,9 @@ fn model_registry_adapter_availability_note(engine: &str) -> Option<&'static str PlatformRuleOut::NativeWindowsVllm => Some( "runtime_status=unsupported_native_windows reason=native Windows skipped; use WSL/Linux vLLM ROCm; gpu_execution_required=true; run /engine for adapter details", ), + PlatformRuleOut::NativeWindowsAim => Some( + "runtime_status=unsupported_native_windows reason=native Windows skipped; AIM GPU serving is Linux/WSL2 only; run /engine for adapter details", + ), } } @@ -7377,7 +7386,7 @@ fn restart_internal_managed_service_with( &started.engine, &started.host, started.port, - &started.canonical_model_id, + started.served_model_name(), endpoint_api_key.as_deref(), Duration::from_secs(45), ); @@ -21118,6 +21127,19 @@ pub(crate) mod tests { /// When the platform gate allows an engine, there is nothing to report -- /// asserted directly against the pure helper, not by relying on whichever /// OS this test happens to run on. + /// AIM's Docker GPU passthrough only exists on Linux/WSL2, so native + /// Windows must be ruled out (the Docker-reachability gate alone would pass + /// on Docker Desktop and launch on CPU). + #[test] + fn unsupported_here_for_names_wsl_for_aim_on_native_windows() { + let message = unsupported_here_for("aim", PlatformRuleOut::NativeWindowsAim) + .expect("a ruled-out engine gets a message"); + assert!( + message.contains("aim") && message.contains("WSL"), + "{message:?}" + ); + } + #[test] fn unsupported_here_for_is_none_when_allowed() { assert_eq!(unsupported_here_for("vllm", PlatformRuleOut::Allowed), None); diff --git a/apps/rocm/src/serve_cmd.rs b/apps/rocm/src/serve_cmd.rs index fdb68adc1..65f0502f6 100644 --- a/apps/rocm/src/serve_cmd.rs +++ b/apps/rocm/src/serve_cmd.rs @@ -488,19 +488,22 @@ fn parse_gpu_memory_utilization(value: Option<&str>) -> Result> { Ok(Some(parsed)) } -/// Applies `--engine-env KEY=VALUE` overrides to the engine recipe's `extra_env` -/// map. Generic by design (not gated to a specific engine): the mechanism exists -/// so any engine that reads `EngineRecipeHint::extra_env` (currently only -/// `engines/aim`, for overrides like `AIM_PROFILE_ID`) can accept host-specific -/// env tweaks without a dedicated flag per variable. A minimal hint is -/// synthesized when none exists. With no `--engine-env` the hint passes through -/// unchanged. +/// Env vars the engine sets itself, which `--engine-env` may not override. +const RESERVED_ENGINE_ENV_KEYS: [&str; 2] = ["HIP_VISIBLE_DEVICES", "VLLM_API_KEY"]; + +/// Applies `--engine-env` overrides to the engine recipe's `extra_env` map, for +/// the one engine that consumes it (`aim`, e.g. `AIM_PROFILE_ID`): `KEY=VALUE` +/// sets a value, and a bare `KEY` forwards `KEY` from the engine process's own +/// environment so a secret such as `HF_TOKEN` never lands in argv or the +/// service record (recorded as an empty value, meaning "forward"). Other engines +/// ignore the flag (see [`engine_env_ignored_note`]) and get their hint back +/// unchanged. A minimal hint is synthesized when none exists. fn engine_recipe_with_extra_env( engine: &str, hint: Option, engine_env: &[String], ) -> Result> { - if engine_env.is_empty() { + if engine_env.is_empty() || !engine.eq_ignore_ascii_case("aim") { return Ok(hint); } let mut hint = hint.unwrap_or_else(|| EngineRecipeHint { @@ -509,11 +512,18 @@ fn engine_recipe_with_extra_env( ..EngineRecipeHint::default() }); for entry in engine_env { - let (key, value) = entry - .split_once('=') - .ok_or_else(|| anyhow::anyhow!("--engine-env expects `KEY=VALUE`; got `{entry}`"))?; - if key.is_empty() { - bail!("--engine-env expects `KEY=VALUE` with a non-empty key; got `{entry}`"); + let (key, value) = entry.split_once('=').unwrap_or((entry.as_str(), "")); + if key.is_empty() || !key.chars().all(|c| c.is_ascii_alphanumeric() || c == '_') { + bail!( + "--engine-env expects `KEY=VALUE` (or a bare `KEY` to forward it from the \ + environment) with a non-empty name of letters, digits and `_`; got `{entry}`" + ); + } + if RESERVED_ENGINE_ENV_KEYS.contains(&key) { + bail!("--engine-env cannot set `{key}`; rocm-cli manages it for the engine"); + } + if value.contains(['\n', '\r']) { + bail!("--engine-env value for `{key}` cannot contain newlines"); } hint.extra_env.insert(key.to_owned(), value.to_owned()); } @@ -959,7 +969,7 @@ pub(crate) fn serve(args: ServeArgs) -> Result<()> { // 0600 key file by service id, so the smoke test authenticates // against a protected public endpoint without threading the secret // through here. - serve_summary::run_smoke_test(&paths, &resolve.canonical_model_id) + serve_summary::run_smoke_test(&paths, resolve.served_model_name()) } else { serve_summary::SmokeMetrics::default() }; @@ -981,7 +991,7 @@ pub(crate) fn serve(args: ServeArgs) -> Result<()> { let summary = serve_summary::DeploymentSummary { engine: selected_engine.clone(), requested_model: model, - api_model: resolve.canonical_model_id, + api_model: resolve.served_model_name().to_owned(), chat_endpoint: format!("{}/chat/completions", report.endpoint_url), service_id: report.service_id.clone(), status: report.status.clone(), @@ -1131,19 +1141,56 @@ mod tests { } #[test] - fn serve_model_ref_for_engine_pins_image_digest_when_set() { - let mut recipe = resolve_builtin_model_recipe("qwen32b").expect("qwen32b recipe"); - let vllm_recipe = recipe - .engine_recipes - .iter_mut() - .find(|engine_recipe| engine_recipe.engine == "vllm") - .expect("qwen32b has a vllm engine recipe"); - vllm_recipe.image_digest = Some("sha256:abcd1234".to_owned()); + fn serve_model_ref_for_engine_pins_the_shipped_aim_image_to_its_digest() { + let recipe = resolve_builtin_model_recipe("aim-qwen3.5-9b").expect("aim catalog entry"); + let pinned = serve_model_ref_for_engine("aim-qwen3.5-9b", Some(&recipe), "aim"); + let (image, digest) = pinned + .split_once('@') + .expect("the AIM entry must be pinned"); + assert_eq!( + image, + "amdenterpriseai/aim-radeon-qwen-qwen3-5-9b:0.12.0-preview" + ); + assert!( + digest.starts_with("sha256:") && digest.len() == "sha256:".len() + 64, + "{digest}" + ); + // Other engines are unaffected by the AIM entry's pin. + let qwen32b = resolve_builtin_model_recipe("qwen32b").expect("qwen32b recipe"); + assert_eq!( + serve_model_ref_for_engine("qwen32b", Some(&qwen32b), "vllm"), + "Qwen/Qwen3-32B-FP8" + ); + } + #[test] + fn catalog_extra_env_reaches_the_aim_hint_and_engine_env_overrides_it() { + let recipe = resolve_builtin_model_recipe("aim-qwen3.5-9b").expect("aim catalog entry"); + let hint = protocol_engine_recipe_hint(&recipe, "aim").expect("aim hint"); assert_eq!( - serve_model_ref_for_engine("qwen32b", Some(&recipe), "vllm"), - "Qwen/Qwen3-32B-FP8@sha256:abcd1234" + hint.extra_env.get("AIM_PROFILE_ID").map(String::as_str), + Some("vllm-w7900-bf16-tp1-latency") + ); + + let overridden = engine_recipe_with_extra_env( + "aim", + Some(hint), + &[ + "AIM_PROFILE_ID=vllm-other".to_owned(), + "HF_TOKEN".to_owned(), + ], + ) + .unwrap() + .unwrap(); + assert_eq!( + overridden + .extra_env + .get("AIM_PROFILE_ID") + .map(String::as_str), + Some("vllm-other"), + "`--engine-env` must beat the catalog's value" ); + assert_eq!(overridden.extra_env.get("HF_TOKEN"), Some(&String::new())); } #[test] @@ -1722,17 +1769,37 @@ mod tests { } #[test] - fn extra_env_rejects_entry_with_no_equals_sign() { - let error = - engine_recipe_with_extra_env("aim", None, &["AIM_PROFILE_ID".to_owned()]).unwrap_err(); - assert!(error.to_string().contains("KEY=VALUE")); + fn extra_env_bare_key_forwards_from_the_environment_without_recording_a_value() { + let hint = engine_recipe_with_extra_env("aim", None, &["HF_TOKEN".to_owned()]) + .unwrap() + .unwrap(); + assert_eq!(hint.extra_env.get("HF_TOKEN"), Some(&String::new())); } #[test] - fn extra_env_rejects_entry_with_empty_key() { - let error = - engine_recipe_with_extra_env("aim", None, &["=some-value".to_owned()]).unwrap_err(); - assert!(error.to_string().contains("non-empty key")); + fn extra_env_rejects_entry_with_empty_or_malformed_key() { + for entry in ["=some-value", "", "BAD KEY=1", "A-B=1"] { + let error = engine_recipe_with_extra_env("aim", None, &[entry.to_owned()]).unwrap_err(); + assert!(error.to_string().contains("KEY=VALUE"), "{entry}: {error}"); + } + } + + #[test] + fn extra_env_rejects_reserved_keys_and_newlines() { + for entry in ["VLLM_API_KEY=x", "HIP_VISIBLE_DEVICES=0", "VLLM_API_KEY"] { + let error = engine_recipe_with_extra_env("aim", None, &[entry.to_owned()]).unwrap_err(); + assert!(error.to_string().contains("manages it"), "{entry}: {error}"); + } + assert!(engine_recipe_with_extra_env("aim", None, &["A=b\nc".to_owned()]).is_err()); + } + + #[test] + fn extra_env_leaves_the_hint_alone_for_engines_that_ignore_it() { + assert_eq!( + engine_recipe_with_extra_env("vllm", None, &["FOO=bar".to_owned()]).unwrap(), + None, + "a flag reported as ignored must not change what the engine is sent" + ); } #[test] diff --git a/apps/rocm/src/service_launch.rs b/apps/rocm/src/service_launch.rs index 4e71cf58b..400e2168d 100644 --- a/apps/rocm/src/service_launch.rs +++ b/apps/rocm/src/service_launch.rs @@ -414,6 +414,7 @@ fn spawn_managed_engine_child( Some(device_policy_name(device_policy).to_owned()), ); record.gpu_indices = gpu_indices.to_vec(); + record.served_model_id = resolve.served_model_id.clone(); record.engine_recipe_json = requested_recipe_json; // The flag the user actually passed, carried through rather than re-derived. // @@ -586,7 +587,7 @@ pub(crate) fn start_managed_service( engine, host, port, - &resolve.canonical_model_id, + resolve.served_model_name(), endpoint_api_key, Duration::from_secs(45), on_wait_tick, @@ -1217,6 +1218,7 @@ mod tests { let resolve = ResolveModelResponse { canonical_model_id: "qwen-canonical".to_owned(), + served_model_id: None, task: "chat".to_owned(), source: "hf".to_owned(), revision: "main".to_owned(), @@ -1306,6 +1308,7 @@ mod tests { let resolve = ResolveModelResponse { canonical_model_id: "qwen-canonical".to_owned(), + served_model_id: None, task: "chat".to_owned(), source: "hf".to_owned(), revision: "main".to_owned(), @@ -1390,6 +1393,7 @@ mod tests { let resolve = ResolveModelResponse { canonical_model_id: "qwen-canonical".to_owned(), + served_model_id: None, task: "chat".to_owned(), source: "hf".to_owned(), revision: "main".to_owned(), diff --git a/apps/rocmd/src/common.rs b/apps/rocmd/src/common.rs index 7f1df1746..87dfa11eb 100644 --- a/apps/rocmd/src/common.rs +++ b/apps/rocmd/src/common.rs @@ -189,6 +189,10 @@ const fn rocmd_engine_inventory() -> &'static [(&'static str, &'static str)] { "vllm", "Linux/WSL ROCm GPU serving engine through external vLLM", ), + ( + "aim", + "AMD Enterprise AI pre-built Docker serving containers", + ), ] } diff --git a/apps/rocmd/src/service.rs b/apps/rocmd/src/service.rs index 8745468ff..934a75771 100644 --- a/apps/rocmd/src/service.rs +++ b/apps/rocmd/src/service.rs @@ -38,6 +38,39 @@ pub(crate) fn stop_managed_service_with( paths: &AppPaths, service_id: &str, terminate: impl Fn(&rocm_core::ProcessIdentity) -> rocm_core::TerminationOutcome, +) -> Result { + stop_managed_service_via(paths, service_id, terminate, request_engine_stop) +} + +fn request_engine_stop(paths: &AppPaths, record: &ManagedServiceRecord) -> Result { + crate::common::engine_request::<_, StopResponse>( + paths, + &record.engine, + &record.service_id, + EngineMethod::Stop, + &StopRequest { + service_id: record.service_id.clone(), + force: true, + }, + ) +} + +/// Engines whose service is a container rather than a recorded process: the +/// engine `Stop` is the only thing that stops one, so its verdict, not the PID +/// signalling, decides whether the stop is confirmed. For a PID-supervised +/// engine the PIDs are authoritative and the engine `Stop` is reported only. +fn engine_stop_decides_confirmation(engine: &str) -> bool { + engine == "aim" +} + +/// [`stop_managed_service_with`], with the engine `Stop` request supplied by the +/// caller as well, so a test can see what was asked of the engine and make it +/// fail. +pub(crate) fn stop_managed_service_via( + paths: &AppPaths, + service_id: &str, + terminate: impl Fn(&rocm_core::ProcessIdentity) -> rocm_core::TerminationOutcome, + engine_stop: impl Fn(&AppPaths, &ManagedServiceRecord) -> Result, ) -> Result { let mut stopping = find_managed_service(paths, service_id)?; // Record the request before the stop starts, not only after it ends. A stop @@ -49,24 +82,12 @@ pub(crate) fn stop_managed_service_with( let requested_at = rocm_core::unix_time_millis(); stopping.stop_requested_unix_ms = Some(requested_at); stopping.write()?; - // Ask the engine to stop before signalling recorded PIDs. For a - // PID-supervised engine (vLLM, Lemonade) this redundantly re-signals the - // same PID `terminate_recorded_service_pids` below already handles — - // harmless, same as `rocm services stop` (`stop_internal_managed_service`) - // already does unconditionally. For a container-backed engine like `aim`, - // this is the only thing that actually stops the container: its recorded - // PIDs belong to a short-lived launcher, not the container itself, so PID - // signalling alone leaks it running. - let _ = crate::common::engine_request::<_, StopResponse>( - paths, - &stopping.engine, - service_id, - EngineMethod::Stop, - &StopRequest { - service_id: stopping.service_id.clone(), - force: true, - }, - ); + // Ask the engine to stop before signalling recorded PIDs. Lemonade is + // skipped: `rocm services stop` unloads its model over HTTP rather than + // sending it an engine `Stop`, and this path has never touched it. For a + // container-backed engine (`aim`) this is the only thing that stops the + // container, since its recorded PIDs are only the supervisor that follows it. + let engine_stop = (stopping.engine != "lemonade").then(|| engine_stop(paths, &stopping)); let stops = terminate_recorded_service_pids(&stopping, terminate); let pids_where = |keep: fn(&RecordedPidStop) -> bool| -> Vec { stops @@ -105,8 +126,11 @@ pub(crate) fn stop_managed_service_with( // stop was running — those are processes it never touched, and claiming the // service stopped, clearing their PIDs or dropping the key they serve with // would each be false. - let all_stopped = - stops.iter().all(|stop| stop.stopped) && stopping.names_same_processes(&record); + let engine_stop_confirms = !engine_stop_decides_confirmation(&stopping.engine) + || engine_stop.as_ref().is_some_and(Result::is_ok); + let all_stopped = stops.iter().all(|stop| stop.stopped) + && stopping.names_same_processes(&record) + && engine_stop_confirms; // A PID whose identity cannot be read at all — `/proc//stat` // unreadable while `kill(pid, 0)` still answers: another uid, `hidepid`, an // LSM — comes back `Unverified`, which never counts as stopped. Nor can it @@ -169,10 +193,28 @@ pub(crate) fn stop_managed_service_with( "force_signaled_pids": force_signaled_pids, "skipped_pids": skipped_pids, "pid_outcomes": pid_outcomes, + "engine_stop": engine_stop_report(engine_stop), "stopped": all_stopped, })) } +/// The engine `Stop` outcome as it appears in a stop report; `null` when no +/// engine `Stop` was sent. +fn engine_stop_report(engine_stop: Option>) -> Value { + match engine_stop { + None => Value::Null, + Some(Ok(response)) => json!({ + "attempted": true, + "stopped": response.stopped, + "graceful": response.graceful, + }), + Some(Err(error)) => json!({ + "attempted": true, + "error": error.to_string(), + }), + } +} + /// Whether a [`stop_managed_service`] report confirms that every recorded /// process is gone. A report without the verdict counts as unconfirmed, so a /// caller can never mistake a missing field for a completed stop. @@ -1864,6 +1906,124 @@ mod tests { Ok(()) } + /// Seeds a keyed service and relabels it as belonging to `engine`. + fn seed_service_for_engine( + paths: &AppPaths, + service_id: &str, + port: u16, + engine: &str, + ) -> Result<()> { + seed_keyed_service(paths, service_id, port)?; + let mut record = load_service_record(paths, service_id)?; + record.engine = engine.to_owned(); + record.write() + } + + /// A container-backed service (`aim`) has no process whose exit proves it + /// stopped, so a stop must send the engine a `Stop` and report its result. + #[test] + fn stopping_an_aim_service_sends_the_engine_a_forced_stop() -> Result<()> { + let (root, paths) = temp_app_paths("stop-aim-sends-engine-stop"); + paths.ensure()?; + let service_id = "svc-stop-aim-sends-engine-stop"; + seed_service_for_engine(&paths, service_id, 11460, "aim")?; + let requested = std::cell::RefCell::new(Vec::new()); + + let result = stop_managed_service_via( + &paths, + service_id, + |_| rocm_core::TerminationOutcome::Graceful, + |_, record| { + requested + .borrow_mut() + .push((record.engine.clone(), record.service_id.clone())); + Ok(StopResponse { + stopped: true, + graceful: false, + }) + }, + ); + fs::remove_dir_all(root).ok(); + + let report = result?; + assert_eq!( + requested.into_inner(), + vec![("aim".to_owned(), service_id.to_owned())] + ); + assert_eq!(report["stopped"], true); + assert_eq!(report["engine_stop"]["attempted"], true); + assert_eq!(report["engine_stop"]["stopped"], true); + Ok(()) + } + + /// If the container cannot be stopped the service is still running, however + /// cleanly its supervisor process went away. + #[test] + fn a_failed_engine_stop_leaves_an_aim_service_unconfirmed() -> Result<()> { + let (root, paths) = temp_app_paths("stop-aim-engine-stop-fails"); + paths.ensure()?; + let service_id = "svc-stop-aim-engine-stop-fails"; + seed_service_for_engine(&paths, service_id, 11461, "aim")?; + + let result = stop_managed_service_via( + &paths, + service_id, + |_| rocm_core::TerminationOutcome::Graceful, + |_, _| anyhow::bail!("docker kill failed"), + ); + let reloaded = load_service_record(&paths, service_id); + fs::remove_dir_all(root).ok(); + + let report = result?; + assert_eq!(report["stopped"], false); + assert!( + report["engine_stop"]["error"] + .as_str() + .is_some_and(|error| error.contains("docker kill failed")), + "{report}" + ); + let reloaded = reloaded?; + assert_ne!(reloaded.status, "stopped"); + assert!(reloaded.stop_requested_unix_ms.is_some()); + Ok(()) + } + + /// Lemonade is unloaded over HTTP by `rocm services stop`, never sent an + /// engine `Stop`; this path must not newly start doing so. A PID-supervised + /// engine gets the `Stop` but its PIDs stay authoritative for the verdict. + #[test] + fn engine_stop_is_skipped_for_lemonade_and_advisory_for_pid_engines() -> Result<()> { + let (root, paths) = temp_app_paths("stop-engine-stop-policy"); + paths.ensure()?; + seed_service_for_engine(&paths, "svc-stop-policy-lemonade", 11462, "lemonade")?; + seed_service_for_engine(&paths, "svc-stop-policy-vllm", 11463, "vllm")?; + + let lemonade = stop_managed_service_via( + &paths, + "svc-stop-policy-lemonade", + |_| rocm_core::TerminationOutcome::Graceful, + |_, _| panic!("lemonade must not be sent an engine Stop"), + ); + let vllm = stop_managed_service_via( + &paths, + "svc-stop-policy-vllm", + |_| rocm_core::TerminationOutcome::Graceful, + |_, _| anyhow::bail!("engine unreachable"), + ); + fs::remove_dir_all(root).ok(); + + let lemonade = lemonade?; + assert_eq!(lemonade["stopped"], true); + assert!(lemonade["engine_stop"].is_null()); + let vllm = vllm?; + assert_eq!( + vllm["stopped"], true, + "the PIDs decide for a process engine" + ); + assert!(vllm["engine_stop"]["error"].is_string()); + Ok(()) + } + /// The stop applies its result to the record as it is *after* the stop, not /// to the copy it loaded before: anything another writer persisted in the /// meantime survives. diff --git a/crates/rocm-core/src/model_catalog.json b/crates/rocm-core/src/model_catalog.json index 54b65e486..c107637f3 100644 --- a/crates/rocm-core/src/model_catalog.json +++ b/crates/rocm-core/src/model_catalog.json @@ -538,6 +538,7 @@ "requires a reachable Docker daemon; the image is pulled on first serve" ], "model_id_override": "amdenterpriseai/aim-radeon-qwen-qwen3-5-9b:0.12.0-preview", + "image_digest": "sha256:ee1201261f358be881318de5930ce9505d17309486f747a83147d7d042fc95ce", "extra_env": { "AIM_PROFILE_ID": "vllm-w7900-bf16-tp1-latency" } diff --git a/crates/rocm-core/src/model_catalog.rs b/crates/rocm-core/src/model_catalog.rs index 8febf7e58..4be915292 100644 --- a/crates/rocm-core/src/model_catalog.rs +++ b/crates/rocm-core/src/model_catalog.rs @@ -86,9 +86,10 @@ pub struct ModelRecipeEngineRecord { #[serde(default)] pub model_id_override: Option, /// Pins the container image this engine serves the model from to a specific - /// digest (`sha256:...`), rather than a floating tag. Enforced by appending - /// `@{digest}` to the resolved model ref in `serve_model_ref_for_engine`; - /// used by container-based engines such as AIM. + /// digest (`sha256:` plus 64 hex characters), rather than a floating tag. + /// Enforced by appending `@{digest}` to the resolved model ref in + /// `serve_model_ref_for_engine`. Only valid on a container engine's record + /// (currently `aim`); catalog validation rejects it elsewhere. #[serde(default)] pub image_digest: Option, } @@ -652,6 +653,12 @@ fn recipe_artifact_url_host(uri: &str) -> Option { (!host.is_empty()).then_some(host) } +fn is_sha256_digest(value: &str) -> bool { + value + .strip_prefix("sha256:") + .is_some_and(|hex| hex.len() == 64 && hex.chars().all(|c| c.is_ascii_hexdigit())) +} + fn validate_model_recipe_engine_record( engine_recipe: &ModelRecipeEngineRecord, canonical_model_id: &str, @@ -685,6 +692,23 @@ fn validate_model_recipe_engine_record( if let Some(model_id_override) = engine_recipe.model_id_override.as_deref() { require_nonempty(model_id_override, "engine model id override")?; } + if let Some(digest) = engine_recipe.image_digest.as_deref() { + // The digest is appended to the image reference by the serve path, which + // only means something for an engine that pulls an image. + if !engine_recipe.engine.eq_ignore_ascii_case("aim") { + bail!( + "engine recipe for `{}` on `{canonical_model_id}` sets image_digest, which only \ + applies to container engines (aim)", + engine_recipe.engine + ); + } + if !is_sha256_digest(digest) { + bail!( + "engine recipe image_digest `{digest}` on `{canonical_model_id}` must be \ + `sha256:` followed by 64 hex characters" + ); + } + } if engine_recipe.required_flags.is_empty() && engine_recipe.parser_settings.is_empty() && engine_recipe.preferred_endpoint.is_none() @@ -936,6 +960,47 @@ mod tests { ); } + #[test] + fn model_recipe_index_validates_image_digest_shape_and_engine() { + let build = |engine: &str, digest: &str| { + let mut recipe = sample_recipe_with_artifact("Qwen/Test-1B", &["test-qwen"]); + recipe.engine_recipes.push(ModelRecipeEngineRecord { + engine: engine.to_owned(), + required_flags: vec!["--flag".to_owned()], + parser_settings: BTreeMap::new(), + extra_env: BTreeMap::new(), + preferred_endpoint: None, + unsupported_combinations: Vec::new(), + notes: Vec::new(), + model_id_override: None, + image_digest: Some(digest.to_owned()), + }); + ModelRecipeIndexDocument { + schema_version: 1, + source: Some("fixture".to_owned()), + generated_at_unix_ms: Some(123), + platforms: Vec::new(), + recipes: vec![recipe], + } + }; + let good = format!("sha256:{}", "a".repeat(64)); + build("aim", &good) + .validate() + .expect("a valid aim digest passes"); + assert!( + build("aim", "sha256:abcd1234").validate().is_err(), + "a truncated digest must be rejected" + ); + assert!( + build("aim", &"a".repeat(64)).validate().is_err(), + "a digest without the sha256: prefix must be rejected" + ); + assert!( + build("vllm", &good).validate().is_err(), + "image_digest only applies to container engines" + ); + } + #[test] fn model_recipe_index_validates_artifact_metadata() -> Result<()> { let mut recipe = sample_recipe_with_artifact("Qwen/Test-1B", &["test-qwen"]); diff --git a/crates/rocm-core/src/net.rs b/crates/rocm-core/src/net.rs index 03b6616cb..dbcb31378 100644 --- a/crates/rocm-core/src/net.rs +++ b/crates/rocm-core/src/net.rs @@ -800,13 +800,7 @@ pub fn managed_service_endpoint_model_ready( if record.endpoint_url.trim().is_empty() { return Ok(false); } - let expected = if !record.canonical_model_id.trim().is_empty() { - Some(record.canonical_model_id.as_str()) - } else if !record.model_ref.trim().is_empty() { - Some(record.model_ref.as_str()) - } else { - None - }; + let expected = Some(record.served_model_name()).filter(|name| !name.trim().is_empty()); openai_models_endpoint_has_model(&record.endpoint_url, expected, endpoint_api_key, timeout) } @@ -880,15 +874,11 @@ pub fn managed_service_endpoint_readiness( { return outcome(EndpointReadiness::Listing, false); } - let model_ref = if record.canonical_model_id.trim().is_empty() { - record.model_ref.as_str() - } else { - record.canonical_model_id.as_str() - }; + let model_ref = record.served_model_name().to_owned(); record.inference_probe_attempted_at_unix_ms = Some(now); if !openai_chat_completion_probe( &record.endpoint_url, - model_ref, + &model_ref, endpoint_api_key, probe_timeout, ) diff --git a/crates/rocm-core/src/service_record.rs b/crates/rocm-core/src/service_record.rs index ce0459001..65cf009f6 100644 --- a/crates/rocm-core/src/service_record.rs +++ b/crates/rocm-core/src/service_record.rs @@ -31,6 +31,12 @@ pub struct ManagedServiceRecord { pub engine: String, pub model_ref: String, pub canonical_model_id: String, + /// The name the running server lists in `/v1/models`, when it differs from + /// `canonical_model_id` (a container engine's `canonical_model_id` is the + /// image to pull). Adopted from the engine state file; see + /// [`Self::served_model_name`]. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub served_model_id: Option, pub host: String, pub port: u16, pub endpoint_url: String, @@ -134,6 +140,7 @@ impl ManagedServiceRecord { engine, model_ref: model_ref.into(), canonical_model_id: canonical_model_id.into(), + served_model_id: None, host, port, mode: mode.into(), @@ -163,6 +170,21 @@ impl ManagedServiceRecord { } } + /// The model name to look for in the endpoint's `/v1/models` and to send in + /// requests: the served name when the engine reported one, else the + /// canonical id, else the model ref. + #[must_use] + pub fn served_model_name(&self) -> &str { + [ + self.served_model_id.as_deref().unwrap_or_default(), + self.canonical_model_id.as_str(), + self.model_ref.as_str(), + ] + .into_iter() + .find(|name| !name.trim().is_empty()) + .unwrap_or_default() + } + /// Drop the per-run state that a restart invalidates, and count the restart. /// /// A restart reuses this manifest but spawns a different server with an @@ -277,6 +299,13 @@ impl ManagedServiceRecord { { self.env_id = Some(env_id.to_owned()); } + if let Some(served) = state + .get("served_model_id") + .and_then(serde_json::Value::as_str) + .filter(|value| !value.trim().is_empty()) + { + self.served_model_id = Some(served.to_owned()); + } // Adopt the engine server PID together with ITS OWN start-time token so // a later stop verifies the server process, not the launcher. The ticks // key must match the PID key: `server_pid`↔`server_start_ticks`, @@ -400,6 +429,29 @@ mod tests { record } + /// A container engine serves under a different name than the image it pulls; + /// the record adopts that name from the engine state, and every readiness + /// probe asks for it rather than for the image reference. + #[test] + fn engine_refresh_adopts_the_served_model_name_for_readiness_probes() { + let mut record = probe_test_record(11438); + record.canonical_model_id = "org/aim-image:1".to_owned(); + assert_eq!(record.served_model_name(), "org/aim-image:1"); + + let record = refresh_from_state( + "served-name", + &serde_json::json!({ "status": "running", "served_model_id": "Org/Served-Model" }), + ); + assert_eq!(record.served_model_id.as_deref(), Some("Org/Served-Model")); + assert_eq!(record.served_model_name(), "Org/Served-Model"); + + let record = refresh_from_state( + "blank-served-name", + &serde_json::json!({ "status": "running", "served_model_id": " " }), + ); + assert_eq!(record.served_model_id, None); + } + /// vLLM records one process as both `pid` and `server_pid` with only /// `start_ticks`. The refresh prefers `server_pid`, so it has to find that /// process's token under `start_ticks`, or the PID is adopted with none. diff --git a/crates/rocm-dash-collectors/src/engine_registry.rs b/crates/rocm-dash-collectors/src/engine_registry.rs index 4b50e7d52..208e75af3 100644 --- a/crates/rocm-dash-collectors/src/engine_registry.rs +++ b/crates/rocm-dash-collectors/src/engine_registry.rs @@ -19,10 +19,6 @@ pub enum EngineKind { Vllm, /// Lemonade Server — JSON `/api/v1/stats`. Lemonade, - /// AMD Enterprise AI AIM — bundles vLLM, so its `/metrics` surface is - /// presumed identical; unconfirmed against a real container until the - /// Strix Halo validation pass. - Aim, } impl EngineKind { @@ -31,7 +27,6 @@ impl EngineKind { match self { Self::Vllm => "vllm", Self::Lemonade => "lemonade", - Self::Aim => "aim", } } @@ -41,7 +36,7 @@ impl EngineKind { /// unmanaged/external discovery where no registry record exists. pub const fn default_port(self) -> u16 { match self { - Self::Vllm | Self::Aim => 8000, + Self::Vllm => 8000, Self::Lemonade => crate::lemonade::LEMONADE_PORT, // 13305 } } @@ -51,18 +46,16 @@ impl EngineKind { match label.trim().to_ascii_lowercase().as_str() { "vllm" => Some(Self::Vllm), "lemonade" => Some(Self::Lemonade), - "aim" => Some(Self::Aim), _ => None, } } /// Parse this engine's raw scrape body into an [`InstanceSample`] using the /// engine-appropriate parser. vLLM dispatches to the **unchanged** - /// `vllm_prom::parse`; Lemonade to `lemonade::parse_stats`; AIM reuses the - /// vLLM parser since AIM images bundle vLLM. + /// `vllm_prom::parse`; Lemonade to `lemonade::parse_stats`. pub fn parse_sample(self, body: &str) -> InstanceSample { match self { - Self::Vllm | Self::Aim => crate::vllm_prom::parse(body), + Self::Vllm => crate::vllm_prom::parse(body), Self::Lemonade => crate::lemonade::parse_stats(body), } } @@ -74,7 +67,7 @@ mod tests { #[test] fn engine_kind_labels_and_ports_and_from_label_roundtrip() { - for kind in [EngineKind::Vllm, EngineKind::Lemonade, EngineKind::Aim] { + for kind in [EngineKind::Vllm, EngineKind::Lemonade] { assert_eq!(EngineKind::from_label(kind.label()), Some(kind)); } assert_eq!(EngineKind::Lemonade.default_port(), 13305); diff --git a/crates/rocm-engine-protocol/src/lib.rs b/crates/rocm-engine-protocol/src/lib.rs index eff367bf9..6053b39d1 100644 --- a/crates/rocm-engine-protocol/src/lib.rs +++ b/crates/rocm-engine-protocol/src/lib.rs @@ -488,6 +488,12 @@ pub struct InstallResponse { #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ResolveModelResponse { pub canonical_model_id: String, + /// The model name the running server advertises in `/v1/models` and accepts + /// in requests, when it differs from `canonical_model_id` (a container + /// engine's `canonical_model_id` is the image to pull, not the served name). + /// `None` means they are the same. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub served_model_id: Option, pub task: String, pub source: String, pub revision: String, @@ -503,6 +509,18 @@ pub struct ResolveModelResponse { pub warnings: Vec, } +impl ResolveModelResponse { + /// The model name the running server answers to: `served_model_id` when the + /// engine reported one, else `canonical_model_id`. + #[must_use] + pub fn served_model_name(&self) -> &str { + self.served_model_id + .as_deref() + .filter(|name| !name.trim().is_empty()) + .unwrap_or(&self.canonical_model_id) + } +} + #[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)] pub struct EngineRecipeHint { pub contract_version: String, diff --git a/docs/architecture.md b/docs/architecture.md index 6a889eac4..8b89dc787 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -163,9 +163,9 @@ These modules are private `mod` declarations — [`lib.rs`](../crates/e2e-report `engines/aim` runs models as Docker containers instead of a host process: - [`docker.rs`](../engines/aim/src/docker.rs) — thin wrapper around the `docker` CLI binary; no `bollard`/Docker API dependency. -- [`registry.rs`](../engines/aim/src/registry.rs) — Docker Hub v2 registry client for AIM image metadata (image ref normalization, tag/digest parsing). -- [`process.rs`](../engines/aim/src/process.rs) — container run/readiness/stop lifecycle. -- [`state.rs`](../engines/aim/src/state.rs) — service state file I/O. +- [`registry.rs`](../engines/aim/src/registry.rs) — AIM image metadata (`Env`/`Labels`): read from the local daemon once the image is pulled, which works for any registry, with a best-effort Docker Hub v2 lookup as a pre-pull hint (image ref normalization, tag/digest parsing). The hint is optional; its failure never blocks a launch. +- [`process.rs`](../engines/aim/src/process.rs) — container run/readiness/stop lifecycle. `serve_http` is the process rocmd supervises: it blocks for the container's lifetime and records how it ended (`failed` and a non-zero exit for a crash or a vanished container, so recovery runs). A stale container is removed before a re-run, and a stop removes the container so its name and downloaded weights are released. +- [`state.rs`](../engines/aim/src/state.rs) — service state file I/O, including the model name the container actually serves (`served_model_id`), which differs from the image reference. - [`install.rs`](../engines/aim/src/install.rs) — much smaller than vLLM's/Lemonade's `install.rs`: there is no venv, wheel, or runtime archive to manage, only a Docker-reachability check. [`lib.rs`](../engines/aim/src/lib.rs) keeps only CLI/envelope dispatch. diff --git a/docs/manual-testing.md b/docs/manual-testing.md index c7afeb6d0..788414721 100644 --- a/docs/manual-testing.md +++ b/docs/manual-testing.md @@ -728,18 +728,41 @@ Linux/WSL2 only, and requires a reachable Docker daemon (`docker info` must succeed). ```bash -rocm serve Qwen/Qwen3.5-9B --engine aim --managed -rocm services list --all -rocm services stop qwen3.5-9b +rocm serve aim-qwen3.5-9b --engine aim --managed +rocm services list --all # note the service id, e.g. aim-- +rocm services stop +docker ps -a --filter name=rocm-aim- +rocm services restart +``` + +`aim-qwen3.5-9b` is the catalog entry (`rocm model --verbose`). It supplies the +pinned image and `AIM_PROFILE_ID`; to serve a raw image reference instead, pass +the profile yourself: + +```bash +rocm serve amdenterpriseai/aim-radeon-qwen-qwen3-5-9b:0.12.0-preview --engine aim \ + --managed --engine-env AIM_PROFILE_ID=vllm-w7900-bf16-tp1-latency ``` +`--engine-env KEY=VALUE` is stored with the service record, so use it for +non-secret values only. For a secret such as a gated model's `HF_TOKEN`, export +it and pass the bare name (`export HF_TOKEN=...`, then `--engine-env HF_TOKEN`): +it is forwarded from your environment and never stored. `--engine-env` is +ignored by the other engines (the command prints a note saying so). + Expected result: - The command pulls/starts an AIM Docker container (first run takes longer while the image is pulled) and the model is reachable on the printed port - once the service shows `ready`. -- `rocm services list --all` shows the service's engine as `aim`. + once the service shows `ready`. The model name to send is the one printed as + the API model, which is the name the container serves, not the image + reference. +- `rocm services list --all` shows the service's engine as `aim`, and the + service stays `ready` while the container runs. - Stopping the service stops and removes the container; `docker ps -a` no - longer lists it. + longer lists it, and `rocm services restart` brings the same service id back + up. +- Killing the container directly (`docker kill rocm-aim-`) makes the + service read `failed` rather than `stopped`, so rocmd's recovery restarts it. - On native Windows, the same command fails explicitly instead of falling back to another engine. diff --git a/docs/testing.md b/docs/testing.md index ee4298d11..4021ea44e 100644 --- a/docs/testing.md +++ b/docs/testing.md @@ -743,6 +743,7 @@ cargo test -p rocm-engine-lemonade stdio_protocol_routes_all_methods_without_sid cargo test -p rocm-engine-vllm stdio_protocol_routes_all_methods_without_side_effects cargo test -p rocm-engine-vllm cargo test -p rocm-engine-aim +cargo test -p rocm --bin rocm extra_env cargo test -p rocmd event_collector cargo test -p rocmd event_dispatcher ``` diff --git a/engines/aim/Cargo.toml b/engines/aim/Cargo.toml index 39a87bdd3..97b7d0265 100644 --- a/engines/aim/Cargo.toml +++ b/engines/aim/Cargo.toml @@ -24,4 +24,5 @@ rocm-core = { path = "../../crates/rocm-core" } rocm-engine-protocol = { path = "../../crates/rocm-engine-protocol" } serde.workspace = true serde_json.workspace = true +tempfile = "3" ureq = { version = "2.12", features = ["json", "native-certs"] } diff --git a/engines/aim/src/docker.rs b/engines/aim/src/docker.rs index 8d9ea76a0..211f4ef53 100644 --- a/engines/aim/src/docker.rs +++ b/engines/aim/src/docker.rs @@ -8,58 +8,135 @@ use anyhow::{Context, Result, bail}; use std::collections::BTreeMap; -use std::process::{Command, Stdio}; +use std::io::Write; +use std::path::Path; +use std::process::{Command, Output, Stdio}; use std::time::Duration; pub(crate) const DOCKER_BIN: &str = "docker"; -/// Builds `-e KEY=VALUE` args for engine-specific env overrides (e.g. AIM's -/// `AIM_PROFILE_ID`). Sorted (`BTreeMap`) for deterministic argv in tests/logs. -fn extra_env_args(extra_env: &BTreeMap) -> Vec { - extra_env - .iter() - .flat_map(|(key, value)| ["-e".to_owned(), format!("{key}={value}")]) - .collect() +/// Bound for the short-lived `docker` calls (`inspect`, `logs`, `rm`, `kill`), +/// so a wedged daemon cannot hang the supervisor or a stop request. +const CLI_TIMEOUT: Duration = Duration::from_secs(30); + +fn docker_output(command: Command, timeout: Duration, label: &str) -> Result { + rocm_core::process::run_with_timeout(command, timeout, label) } -#[allow(clippy::too_many_arguments)] -pub(crate) fn run( - name: &str, - image: &str, - host: &str, - port: u16, - container_port: u16, - gpu_indices: &[u32], - render_gid: Option, - extra_env: &BTreeMap, -) -> Result<()> { - let mut command = Command::new(DOCKER_BIN); - command - .arg("run") - .arg("-d") - .arg("--name") - .arg(name) - .arg("-p") - .arg(format!("{host}:{port}:{container_port}")); - - if !gpu_indices.is_empty() { - if let Some(csv) = rocm_engine_protocol::gpu_indices_to_csv(gpu_indices) { - command.arg("-e").arg(format!("HIP_VISIBLE_DEVICES={csv}")); +pub(crate) struct RunSpec<'a> { + pub name: &'a str, + pub image: &'a str, + pub host: &'a str, + pub port: u16, + pub container_port: u16, + /// Whether to pass GPU devices through at all. + pub use_gpu: bool, + /// Pins `HIP_VISIBLE_DEVICES`; empty leaves the container's default visibility. + pub gpu_indices: &'a [u32], + pub render_gid: Option, + /// `KEY=VALUE` entries go through a 0600 `--env-file`, never argv. An empty + /// value means "forward `KEY` from this process's environment" (`-e KEY`). + pub env: &'a BTreeMap, +} + +/// Splits `env` into entries destined for the `--env-file` and bare names to +/// forward from the engine process's own environment. +fn split_env(env: &BTreeMap) -> (Vec<(&str, &str)>, Vec<&str>) { + let mut file_entries = Vec::new(); + let mut forwarded = Vec::new(); + for (key, value) in env { + if value.is_empty() { + forwarded.push(key.as_str()); + } else { + file_entries.push((key.as_str(), value.as_str())); } - command.args(gpu_device_args(rocm_core::is_wsl_host(), render_gid)); } + (file_entries, forwarded) +} + +fn run_args( + spec: &RunSpec<'_>, + wsl: bool, + env_file: Option<&Path>, + forwarded: &[&str], +) -> Vec { + let mut args = vec![ + "run".to_owned(), + "-d".to_owned(), + "--name".to_owned(), + spec.name.to_owned(), + "-p".to_owned(), + format!("{}:{}:{}", spec.host, spec.port, spec.container_port), + ]; + if spec.use_gpu { + if let Some(csv) = rocm_engine_protocol::gpu_indices_to_csv(spec.gpu_indices) { + args.push("-e".to_owned()); + args.push(format!("HIP_VISIBLE_DEVICES={csv}")); + } + args.extend(gpu_device_args(wsl, spec.render_gid)); + } + if let Some(path) = env_file { + args.push("--env-file".to_owned()); + args.push(path.display().to_string()); + } + for key in forwarded { + args.push("-e".to_owned()); + args.push((*key).to_owned()); + } + args.push(spec.image.to_owned()); + args +} - command.args(extra_env_args(extra_env)); +fn write_env_file(entries: &[(&str, &str)]) -> Result { + // `NamedTempFile` is created 0600 on unix and removed on drop. + let mut file = tempfile::NamedTempFile::new().context("failed to create docker env file")?; + for (key, value) in entries { + if key.contains(['=', '\n']) || value.contains('\n') { + bail!("engine env `{key}` cannot contain newlines or `=` in its name"); + } + writeln!(file, "{key}={value}").context("failed to write docker env file")?; + } + file.flush().context("failed to write docker env file")?; + Ok(file) +} - command.arg(image); - command.stdin(Stdio::null()); +pub(crate) fn run(spec: &RunSpec<'_>) -> Result<()> { + let (file_entries, forwarded) = split_env(spec.env); + for key in &forwarded { + if std::env::var_os(key).is_none() { + bail!( + "engine env `{key}` was passed without a value, but `{key}` is not set in this \ + process's environment; export it before serving or pass `{key}=VALUE`" + ); + } + } + let env_file = if file_entries.is_empty() { + None + } else { + Some(write_env_file(&file_entries)?) + }; + let args = run_args( + spec, + rocm_core::is_wsl_host(), + env_file.as_ref().map(tempfile::NamedTempFile::path), + &forwarded, + ); - let output = command - .output() - .with_context(|| format!("failed to run `{DOCKER_BIN} run` for container `{name}`"))?; + let mut command = Command::new(DOCKER_BIN); + command.args(&args).stdin(Stdio::null()); + // `docker run -d` returns once the container is created, but a first run + // can take a while on a busy daemon; bound it generously rather than not at all. + let output = + docker_output(command, Duration::from_mins(2), "docker run").with_context(|| { + format!( + "failed to run `{DOCKER_BIN} run` for container `{}`", + spec.name + ) + })?; if !output.status.success() { bail!( - "`docker run --name {name}` failed: {}", + "`docker run --name {}` failed: {}", + spec.name, rocm_core::command_failure_detail(&output) ); } @@ -107,6 +184,8 @@ fn parse_getent_gid(text: &str) -> Option { text.trim().split(':').nth(2)?.parse().ok() } +/// Returns `Ok(false)` when there was nothing running to stop (no such +/// container, or it already exited). pub(crate) fn stop(name: &str, force: bool, grace: Duration) -> Result { let mut command = Command::new(DOCKER_BIN); if force { @@ -118,16 +197,12 @@ pub(crate) fn stop(name: &str, force: bool, grace: Duration) -> Result { .arg(grace.as_secs().to_string()) .arg(name); } - let output = command.stdin(Stdio::null()).output().with_context(|| { - format!( - "failed to run `docker {}` for container `{name}`", - if force { "kill" } else { "stop" } - ) - })?; + let label = if force { "docker kill" } else { "docker stop" }; + let output = docker_output(command, grace + CLI_TIMEOUT, label) + .with_context(|| format!("failed to run `{label}` for container `{name}`"))?; if !output.status.success() { - // A container that's already gone is not a stop failure. let detail = rocm_core::command_failure_detail(&output); - if detail.contains("No such container") { + if is_nothing_to_stop(&detail) { return Ok(false); } bail!("failed to stop container `{name}`: {detail}"); @@ -135,40 +210,75 @@ pub(crate) fn stop(name: &str, force: bool, grace: Duration) -> Result { Ok(true) } -pub(crate) fn inspect_running(name: &str) -> Result> { - inspect_format(name, "{{.State.Running}}").map(|value| value.map(|text| text == "true")) +fn is_nothing_to_stop(detail: &str) -> bool { + detail.contains("No such container") || detail.contains("is not running") +} + +/// Force-removes the container (stopping it first if needed) so its name and +/// writable layer, including any downloaded weights, are released. A container +/// that is already gone is not an error. +pub(crate) fn remove(name: &str) -> Result<()> { + let mut command = Command::new(DOCKER_BIN); + command.arg("rm").arg("-f").arg(name); + let output = docker_output(command, CLI_TIMEOUT, "docker rm") + .with_context(|| format!("failed to run `docker rm` for container `{name}`"))?; + if !output.status.success() { + let detail = rocm_core::command_failure_detail(&output); + if detail.contains("No such container") { + return Ok(()); + } + bail!("failed to remove container `{name}`: {detail}"); + } + Ok(()) } -pub(crate) fn inspect_exit_code(name: &str) -> Result> { - inspect_format(name, "{{.State.ExitCode}}") - .map(|value| value.and_then(|text| text.parse().ok())) +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum ContainerState { + Running, + /// Exited (or created but never started) with this exit code. + Exited(i32), + /// No container by this name exists. + Missing, } -fn inspect_format(name: &str, format: &str) -> Result> { - let output = Command::new(DOCKER_BIN) +/// A daemon that is unreachable or erroring is an `Err`, distinct from a +/// container that is definitively gone. +pub(crate) fn container_state(name: &str) -> Result { + let mut command = Command::new(DOCKER_BIN); + command .arg("inspect") .arg("--format") - .arg(format) - .arg(name) - .stdin(Stdio::null()) - .output() + .arg("{{.State.Running}} {{.State.ExitCode}}") + .arg(name); + let output = docker_output(command, CLI_TIMEOUT, "docker inspect") .with_context(|| format!("failed to run `docker inspect` for container `{name}`"))?; if !output.status.success() { - return Ok(None); + let detail = rocm_core::command_failure_detail(&output); + if detail.contains("No such") { + return Ok(ContainerState::Missing); + } + bail!("`docker inspect {name}` failed: {detail}"); + } + parse_container_state(&String::from_utf8_lossy(&output.stdout)) +} + +fn parse_container_state(text: &str) -> Result { + let mut parts = text.split_whitespace(); + match (parts.next(), parts.next().map(str::parse::)) { + (Some("true"), _) => Ok(ContainerState::Running), + (Some("false"), Some(Ok(code))) => Ok(ContainerState::Exited(code)), + _ => bail!("unexpected `docker inspect` output: {:?}", text.trim()), } - Ok(Some( - String::from_utf8_lossy(&output.stdout).trim().to_owned(), - )) } pub(crate) fn logs(name: &str, tail: usize) -> Result> { - let output = Command::new(DOCKER_BIN) + let mut command = Command::new(DOCKER_BIN); + command .arg("logs") .arg("--tail") .arg(tail.to_string()) - .arg(name) - .stdin(Stdio::null()) - .output() + .arg(name); + let output = docker_output(command, CLI_TIMEOUT, "docker logs") .with_context(|| format!("failed to run `docker logs` for container `{name}`"))?; if !output.status.success() { bail!( @@ -188,6 +298,24 @@ pub(crate) fn logs(name: &str, tail: usize) -> Result> { Ok(lines) } +/// The image's `Config` block (`Env`, `Labels`) from the local daemon, once the +/// image has been pulled. Works for any registry, unlike the Docker Hub lookup. +pub(crate) fn image_config(image: &str) -> Result> { + let mut command = Command::new(DOCKER_BIN); + command + .arg("image") + .arg("inspect") + .arg("--format") + .arg("{{json .Config}}") + .arg(image); + let output = docker_output(command, CLI_TIMEOUT, "docker image inspect") + .with_context(|| format!("failed to run `docker image inspect {image}`"))?; + if !output.status.success() { + return Ok(None); + } + Ok(serde_json::from_slice(&output.stdout).ok()) +} + pub(crate) fn image_present(image: &str) -> Result { let status = Command::new(DOCKER_BIN) .arg("image") @@ -230,8 +358,26 @@ pub(crate) fn daemon_reachable() -> bool { #[cfg(test)] mod tests { - use super::{extra_env_args, gpu_device_args, parse_getent_gid}; + use super::{ + ContainerState, RunSpec, gpu_device_args, is_nothing_to_stop, parse_container_state, + parse_getent_gid, run_args, split_env, + }; use std::collections::BTreeMap; + use std::path::Path; + + fn spec<'a>(gpu_indices: &'a [u32], env: &'a BTreeMap) -> RunSpec<'a> { + RunSpec { + name: "rocm-aim-svc", + image: "org/aim-foo:1", + host: "127.0.0.1", + port: 8080, + container_port: 8000, + gpu_indices, + use_gpu: !gpu_indices.is_empty(), + render_gid: Some(993), + env, + } + } #[test] fn gpu_device_args_native_linux_uses_kfd_dri_and_render_gid() { @@ -268,30 +414,110 @@ mod tests { } #[test] - fn extra_env_args_builds_sorted_flags() { - assert!(extra_env_args(&BTreeMap::new()).is_empty()); + fn run_args_without_gpu_or_env_is_just_name_port_and_image() { + let env = BTreeMap::new(); + assert_eq!( + run_args(&spec(&[], &env), false, None, &[]), + vec![ + "run", + "-d", + "--name", + "rocm-aim-svc", + "-p", + "127.0.0.1:8080:8000", + "org/aim-foo:1" + ] + ); + } + + #[test] + fn run_args_with_gpu_sets_visible_devices_and_device_flags_before_the_image() { + let env = BTreeMap::new(); + let args = run_args(&spec(&[1, 3], &env), false, None, &[]); + assert_eq!(args.last().map(String::as_str), Some("org/aim-foo:1")); + assert!( + args.windows(2) + .any(|pair| pair == ["-e", "HIP_VISIBLE_DEVICES=1,3"]) + ); + assert!(args.contains(&"--device=/dev/kfd".to_owned())); + assert!(args.windows(2).any(|pair| pair == ["--group-add", "993"])); + } + + #[test] + fn run_args_with_gpu_and_no_pinned_indices_omits_visible_devices() { + let env = BTreeMap::new(); + let mut gpu_spec = spec(&[], &env); + gpu_spec.use_gpu = true; + let args = run_args(&gpu_spec, false, None, &[]); + assert!( + !args + .iter() + .any(|arg| arg.starts_with("HIP_VISIBLE_DEVICES")) + ); + assert!(args.contains(&"--device=/dev/kfd".to_owned())); + } + + #[test] + fn run_args_never_puts_env_values_on_the_command_line() { + let env = BTreeMap::from([ + ("AIM_PROFILE_ID".to_owned(), "secret-profile".to_owned()), + ("HF_TOKEN".to_owned(), String::new()), + ]); + let (file_entries, forwarded) = split_env(&env); + assert_eq!(file_entries, vec![("AIM_PROFILE_ID", "secret-profile")]); + assert_eq!(forwarded, vec!["HF_TOKEN"]); + + let args = run_args( + &spec(&[], &env), + false, + Some(Path::new("/tmp/env-file")), + &forwarded, + ); + assert!( + args.windows(2) + .any(|pair| pair == ["--env-file", "/tmp/env-file"]) + ); + assert!(args.windows(2).any(|pair| pair == ["-e", "HF_TOKEN"])); + assert!(!args.iter().any(|arg| arg.contains("secret-profile"))); + assert_eq!(args.last().map(String::as_str), Some("org/aim-foo:1")); + } - let mut extra_env = BTreeMap::new(); - extra_env.insert("AIM_PROFILE_ID".to_owned(), "vllm-foo".to_owned()); - extra_env.insert("ANOTHER".to_owned(), "x".to_owned()); + #[test] + fn parse_container_state_reads_running_and_exit_code() { + assert_eq!( + parse_container_state("true 0\n").unwrap(), + ContainerState::Running + ); assert_eq!( - extra_env_args(&extra_env), - vec!["-e", "AIM_PROFILE_ID=vllm-foo", "-e", "ANOTHER=x"] + parse_container_state("false 137").unwrap(), + ContainerState::Exited(137) ); + assert!(parse_container_state("").is_err()); + assert!(parse_container_state("false oops").is_err()); + } + + #[test] + fn is_nothing_to_stop_covers_missing_and_already_exited_containers() { + assert!(is_nothing_to_stop("Error: No such container: x")); + assert!(is_nothing_to_stop( + "Error response from daemon: cannot kill container: x: container x is not running" + )); + assert!(!is_nothing_to_stop("permission denied")); } } -/// Exercises the `Command`-invoking functions (`run`/`stop`/`inspect_*`/`pull`/ -/// `image_present`) against a fake `docker` on `PATH`, mirroring `install.rs`'s -/// subprocess pattern: mutating `PATH` for the *current* process needs -/// `unsafe` (denied workspace-wide), so each scenario re-spawns this test -/// binary filtered to just itself, with a hostile `PATH` set via the safe -/// `Command::env` on that child. +/// Exercises the `Command`-invoking functions against a fake `docker` on +/// `PATH`, mirroring `install.rs`'s subprocess pattern: mutating `PATH` for the +/// *current* process needs `unsafe` (denied workspace-wide), so each scenario +/// re-spawns this test binary filtered to just itself, with a hostile `PATH` +/// set via the safe `Command::env` on that child. The fake records every argv +/// (and a copy of any `--env-file`) so the parent can assert on them. #[cfg(all(test, unix))] -mod subprocess_tests { - use super::{image_present, inspect_exit_code, inspect_running, pull, run, stop}; +pub(crate) mod subprocess_tests { + use super::{ContainerState, RunSpec, container_state, pull, remove, run, stop}; use rocm_core::fake_executable::write_fake_executable; use std::collections::BTreeMap; + use std::path::PathBuf; use std::process::{Command, Stdio}; use std::time::Duration; @@ -299,9 +525,17 @@ mod subprocess_tests { /// One parameterized fake `docker`, driven entirely by inherited env vars, /// reused by every scenario below instead of planting N bespoke scripts. - const FAKE_DOCKER_SCRIPT: &str = r#"#!/bin/sh + pub(crate) const FAKE_DOCKER_SCRIPT: &str = r#"#!/bin/sh +[ -n "$FAKE_DOCKER_ARGV_LOG" ] && echo "$*" >> "$FAKE_DOCKER_ARGV_LOG" case "$1" in run) + prev="" + for arg in "$@"; do + if [ "$prev" = "--env-file" ] && [ -n "$FAKE_DOCKER_ENV_COPY" ]; then + cp "$arg" "$FAKE_DOCKER_ENV_COPY" + fi + prev="$arg" + done [ -n "$FAKE_DOCKER_RUN_STDERR" ] && echo "$FAKE_DOCKER_RUN_STDERR" >&2 exit "${FAKE_DOCKER_RUN_EXIT:-0}" ;; @@ -310,17 +544,49 @@ case "$1" in echo "Error response from daemon: No such container: fake" >&2 exit 1 fi + if [ -n "$FAKE_DOCKER_STOP_NOT_RUNNING" ]; then + echo "Error response from daemon: container fake is not running" >&2 + exit 1 + fi [ -n "$FAKE_DOCKER_STOP_STDERR" ] && echo "$FAKE_DOCKER_STOP_STDERR" >&2 exit "${FAKE_DOCKER_STOP_EXIT:-0}" ;; + rm) + if [ -n "$FAKE_DOCKER_RM_NO_SUCH_CONTAINER" ]; then + echo "Error response from daemon: No such container: fake" >&2 + exit 1 + fi + exit "${FAKE_DOCKER_RM_EXIT:-0}" + ;; inspect) - case "$3" in - *Running*) echo "${FAKE_DOCKER_RUNNING:-false}" ;; - *ExitCode*) echo "${FAKE_DOCKER_EXIT_CODE:-0}" ;; - esac + if [ -n "$FAKE_DOCKER_INSPECT_SEQUENCE" ]; then + echo x >> "$FAKE_DOCKER_DIR/inspect.count" + n=$(wc -l < "$FAKE_DOCKER_DIR/inspect.count" | tr -d ' ') + entry=$(echo "$FAKE_DOCKER_INSPECT_SEQUENCE" | tr ',' '\n' | sed -n "${n}p") + [ -z "$entry" ] && entry=$(echo "$FAKE_DOCKER_INSPECT_SEQUENCE" | tr ',' '\n' | tail -n 1) + if [ "$entry" = "missing" ]; then + echo "Error: No such object: fake" >&2 + exit 1 + fi + echo "$entry" + exit 0 + fi + if [ -n "$FAKE_DOCKER_INSPECT_MISSING" ]; then + echo "Error: No such object: fake" >&2 + exit 1 + fi + if [ -n "$FAKE_DOCKER_INSPECT_ERROR" ]; then + echo "Cannot connect to the Docker daemon" >&2 + exit 1 + fi + echo "${FAKE_DOCKER_RUNNING:-false} ${FAKE_DOCKER_EXIT_CODE:-0}" exit 0 ;; image) + if [ -n "$FAKE_DOCKER_IMAGE_CONFIG" ]; then + echo "$FAKE_DOCKER_IMAGE_CONFIG" + exit 0 + fi exit "${FAKE_DOCKER_IMAGE_PRESENT_EXIT:-1}" ;; pull) @@ -337,13 +603,39 @@ esac exit 0 "#; - fn in_subprocess() -> bool { + pub(crate) fn in_subprocess() -> bool { std::env::var_os(SENTINEL).is_some() } + /// Scratch dir holding the fake `docker` plus its argv/env-file recordings; + /// removed on drop. + pub(crate) struct Fixture { + dir: PathBuf, + } + + impl Fixture { + pub(crate) fn argv(&self) -> Vec { + std::fs::read_to_string(self.dir.join("argv.log")) + .unwrap_or_default() + .lines() + .map(ToOwned::to_owned) + .collect() + } + + pub(crate) fn env_copy(&self) -> Option { + std::fs::read_to_string(self.dir.join("env.copy")).ok() + } + } + + impl Drop for Fixture { + fn drop(&mut self) { + std::fs::remove_dir_all(&self.dir).ok(); + } + } + /// Plants the fake `docker`, re-execs this test binary with just /// `test_path` selected and `extra_env` set, and asserts it succeeded. - fn run_with_fake_docker(test_path: &str, extra_env: &[(&str, &str)]) { + pub(crate) fn run_with_fake_docker(test_path: &str, extra_env: &[(&str, &str)]) -> Fixture { let dir = std::env::temp_dir().join(format!( "aim-docker-test-{}-{}", std::process::id(), @@ -355,6 +647,7 @@ exit 0 let fake_bin = dir.join("bin"); std::fs::create_dir_all(&fake_bin).unwrap(); write_fake_executable(&fake_bin.join("docker"), FAKE_DOCKER_SCRIPT).unwrap(); + let fixture = Fixture { dir: dir.clone() }; let mut entries = vec![fake_bin]; entries.extend(std::env::split_paths( @@ -366,52 +659,95 @@ exit 0 .args(["--exact", test_path, "--nocapture"]) .env(SENTINEL, "1") .env("PATH", std::env::join_paths(entries).unwrap()) + .env("FAKE_DOCKER_ARGV_LOG", dir.join("argv.log")) + .env("FAKE_DOCKER_ENV_COPY", dir.join("env.copy")) + .env("FAKE_DOCKER_DIR", &dir) + // Keep any state a child writes out of the user's real data dir. + .env("ROCM_CLI_DATA_DIR", dir.join("data")) .stdin(Stdio::null()); for (key, value) in extra_env { command.env(key, value); } let status = command.status().unwrap(); assert!(status.success(), "helper subprocess failed for {test_path}"); + fixture + } - std::fs::remove_dir_all(&dir).ok(); + const fn plain_spec(env: &BTreeMap) -> RunSpec<'_> { + RunSpec { + name: "fake-container", + image: "fake-image", + host: "127.0.0.1", + port: 8080, + container_port: 8000, + use_gpu: false, + gpu_indices: &[], + render_gid: None, + env, + } } #[test] fn run_succeeds_against_fake_docker() { if in_subprocess() { - run( - "fake-container", - "fake-image", - "127.0.0.1", - 8080, - 8000, - &[], - None, - &BTreeMap::new(), - ) - .expect("fake docker run should succeed"); + run(&plain_spec(&BTreeMap::new())).expect("fake docker run should succeed"); return; } - run_with_fake_docker( + let fixture = run_with_fake_docker( "docker::subprocess_tests::run_succeeds_against_fake_docker", &[], ); + assert_eq!( + fixture.argv(), + vec!["run -d --name fake-container -p 127.0.0.1:8080:8000 fake-image"] + ); + } + + #[test] + fn run_keeps_env_values_off_argv_and_in_a_private_env_file() { + if in_subprocess() { + let env = BTreeMap::from([ + ("AIM_PROFILE_ID".to_owned(), "profile-123".to_owned()), + ("HF_TOKEN".to_owned(), String::new()), + ]); + run(&plain_spec(&env)).expect("fake docker run should succeed"); + return; + } + let fixture = run_with_fake_docker( + "docker::subprocess_tests::run_keeps_env_values_off_argv_and_in_a_private_env_file", + &[("HF_TOKEN", "hf_super_secret")], + ); + let argv = fixture.argv().join("\n"); + assert!(argv.contains("--env-file"), "{argv}"); + assert!(argv.contains("-e HF_TOKEN"), "{argv}"); + assert!(!argv.contains("profile-123"), "{argv}"); + assert!(!argv.contains("hf_super_secret"), "{argv}"); + assert_eq!( + fixture.env_copy().as_deref(), + Some("AIM_PROFILE_ID=profile-123\n") + ); + } + + #[test] + fn run_refuses_a_forwarded_env_var_that_is_not_set() { + if in_subprocess() { + let env = BTreeMap::from([("ROCM_AIM_TEST_UNSET_VAR".to_owned(), String::new())]); + let error = run(&plain_spec(&env)).expect_err("unset forwarded var must fail"); + assert!(error.to_string().contains("ROCM_AIM_TEST_UNSET_VAR")); + return; + } + let fixture = run_with_fake_docker( + "docker::subprocess_tests::run_refuses_a_forwarded_env_var_that_is_not_set", + &[], + ); + assert!(fixture.argv().is_empty(), "docker must not have been run"); } #[test] fn run_bails_with_container_name_and_detail_on_failure() { if in_subprocess() { - let error = run( - "fake-container", - "fake-image", - "127.0.0.1", - 8080, - 8000, - &[], - None, - &BTreeMap::new(), - ) - .expect_err("fake docker run should fail"); + let error = + run(&plain_spec(&BTreeMap::new())).expect_err("fake docker run should fail"); let message = error.to_string(); assert!(message.contains("fake-container"), "{message}"); assert!(message.contains("boom"), "{message}"); @@ -433,23 +769,31 @@ exit 0 assert!(stop("fake-container", false, Duration::from_secs(1)).unwrap()); return; } - run_with_fake_docker( + let fixture = run_with_fake_docker( "docker::subprocess_tests::stop_force_kill_and_graceful_stop_both_succeed", &[], ); + assert_eq!( + fixture.argv(), + vec!["kill fake-container", "stop -t 1 fake-container"] + ); } #[test] - fn stop_missing_container_returns_false_not_error() { + fn stop_missing_or_already_exited_container_returns_false_not_error() { if in_subprocess() { let stopped = stop("fake-container", true, Duration::from_secs(1)).unwrap(); assert!(!stopped); return; } run_with_fake_docker( - "docker::subprocess_tests::stop_missing_container_returns_false_not_error", + "docker::subprocess_tests::stop_missing_or_already_exited_container_returns_false_not_error", &[("FAKE_DOCKER_STOP_NO_SUCH_CONTAINER", "1")], ); + run_with_fake_docker( + "docker::subprocess_tests::stop_missing_or_already_exited_container_returns_false_not_error", + &[("FAKE_DOCKER_STOP_NOT_RUNNING", "1")], + ); } #[test] @@ -470,19 +814,77 @@ exit 0 } #[test] - fn inspect_running_and_exit_code_parse_fake_docker_output() { + fn remove_force_removes_and_tolerates_a_missing_container() { + if in_subprocess() { + remove("fake-container").expect("rm should succeed or tolerate a missing container"); + return; + } + let fixture = run_with_fake_docker( + "docker::subprocess_tests::remove_force_removes_and_tolerates_a_missing_container", + &[], + ); + assert_eq!(fixture.argv(), vec!["rm -f fake-container"]); + run_with_fake_docker( + "docker::subprocess_tests::remove_force_removes_and_tolerates_a_missing_container", + &[("FAKE_DOCKER_RM_NO_SUCH_CONTAINER", "1")], + ); + } + + #[test] + fn remove_surfaces_other_failures() { if in_subprocess() { - assert_eq!(inspect_running("fake-container").unwrap(), Some(true)); - assert_eq!(inspect_exit_code("fake-container").unwrap(), Some(137)); + let error = remove("fake-container").expect_err("should fail"); + assert!(error.to_string().contains("fake-container"), "{error}"); return; } run_with_fake_docker( - "docker::subprocess_tests::inspect_running_and_exit_code_parse_fake_docker_output", + "docker::subprocess_tests::remove_surfaces_other_failures", + &[("FAKE_DOCKER_RM_EXIT", "1")], + ); + } + + #[test] + fn container_state_distinguishes_running_exited_missing_and_daemon_errors() { + if in_subprocess() { + let state = container_state("fake-container"); + match std::env::var("FAKE_DOCKER_EXPECT").as_deref() { + Ok("running") => assert_eq!(state.unwrap(), ContainerState::Running), + Ok("exited") => assert_eq!(state.unwrap(), ContainerState::Exited(137)), + Ok("missing") => assert_eq!(state.unwrap(), ContainerState::Missing), + Ok("error") => assert!(state.is_err()), + other => panic!("unexpected FAKE_DOCKER_EXPECT {other:?}"), + } + return; + } + let test = "docker::subprocess_tests::container_state_distinguishes_running_exited_missing_and_daemon_errors"; + run_with_fake_docker( + test, &[ + ("FAKE_DOCKER_EXPECT", "running"), ("FAKE_DOCKER_RUNNING", "true"), + ], + ); + run_with_fake_docker( + test, + &[ + ("FAKE_DOCKER_EXPECT", "exited"), ("FAKE_DOCKER_EXIT_CODE", "137"), ], ); + run_with_fake_docker( + test, + &[ + ("FAKE_DOCKER_EXPECT", "missing"), + ("FAKE_DOCKER_INSPECT_MISSING", "1"), + ], + ); + run_with_fake_docker( + test, + &[ + ("FAKE_DOCKER_EXPECT", "error"), + ("FAKE_DOCKER_INSPECT_ERROR", "1"), + ], + ); } #[test] @@ -501,7 +903,7 @@ exit 0 #[test] fn image_present_reports_true_when_fake_docker_reports_it() { if in_subprocess() { - assert!(image_present("fake-image").unwrap()); + assert!(super::image_present("fake-image").unwrap()); return; } run_with_fake_docker( @@ -509,4 +911,20 @@ exit 0 &[("FAKE_DOCKER_IMAGE_PRESENT_EXIT", "0")], ); } + + #[test] + fn image_config_parses_the_local_images_config_block() { + if in_subprocess() { + let config = super::image_config("fake-image").unwrap().expect("config"); + assert_eq!(config["Labels"]["k"], "v"); + return; + } + run_with_fake_docker( + "docker::subprocess_tests::image_config_parses_the_local_images_config_block", + &[( + "FAKE_DOCKER_IMAGE_CONFIG", + r#"{"Env":["A=b"],"Labels":{"k":"v"}}"#, + )], + ); + } } diff --git a/engines/aim/src/lib.rs b/engines/aim/src/lib.rs index 273a73bc8..3bab19210 100644 --- a/engines/aim/src/lib.rs +++ b/engines/aim/src/lib.rs @@ -4,7 +4,6 @@ use anyhow::{Context, Result}; use clap::{Parser, Subcommand}; -use rocm_core::AppPaths; use rocm_engine_protocol::{ DetectRequest, DetectResponse, DevicePolicy, EngineCapabilities, EngineDeviceAvailability, EngineMethod, EngineRecipeHint, EngineRequestEnvelope, EngineResponseEnvelope, InstallRequest, @@ -99,15 +98,7 @@ pub(crate) struct ServeHttpRequest { pub port: u16, pub device_policy: DevicePolicy, pub gpu_indices: Vec, - #[allow(dead_code)] // ponytail: accepted for protocol-shape parity; AIM has no runtime concept - pub runtime_id: Option, - #[allow(dead_code)] - // ponytail: accepted for protocol-shape parity; AIM has no managed-env concept - pub env_id: Option, pub state_path: PathBuf, - #[allow(dead_code)] - // ponytail: no host-side log file in `docker run -d` mode; logs come from `docker logs` - pub log_path: Option, pub engine_recipe: Option, } @@ -176,17 +167,13 @@ pub fn run_cli() -> Result<()> { port, device_policy, gpu, - runtime_id, - env_id, + // AIM has no managed runtime or env; accepted for protocol-shape parity. + runtime_id: _, + env_id: _, state_path, engine_recipe_json, } => { let gpu_indices = crate::process::parse_gpu_indices_arg(gpu.as_deref())?; - let log_path = Some( - AppPaths::discover()? - .engine_logs_dir(ENGINE_NAME) - .join(format!("{service_id}.log")), - ); crate::process::serve_http(ServeHttpRequest { service_id, model_ref, @@ -194,10 +181,7 @@ pub fn run_cli() -> Result<()> { port, device_policy: crate::process::parse_device_policy_arg(&device_policy)?, gpu_indices, - runtime_id, - env_id, state_path, - log_path, engine_recipe: crate::process::parse_engine_recipe_json(engine_recipe_json)?, }) } @@ -224,10 +208,12 @@ pub fn builtin_serve_http( port: u16, device_policy: DevicePolicy, gpu_indices: Vec, - runtime_id: Option, - env_id: Option, + // AIM has no managed runtime, env or host-side log file (container logs come + // from `docker logs`); accepted for parity with the other built-in engines. + _runtime_id: Option, + _env_id: Option, state_path: PathBuf, - log_path: Option, + _log_path: Option, engine_recipe: Option, ) -> Result<()> { crate::process::serve_http(ServeHttpRequest { @@ -237,10 +223,7 @@ pub fn builtin_serve_http( port, device_policy, gpu_indices, - runtime_id, - env_id, state_path, - log_path, engine_recipe, }) } diff --git a/engines/aim/src/process.rs b/engines/aim/src/process.rs index 360ecf782..7a1ffa32d 100644 --- a/engines/aim/src/process.rs +++ b/engines/aim/src/process.rs @@ -3,19 +3,31 @@ // SPDX-License-Identifier: MIT use anyhow::{Context, Result, bail}; -use rocm_core::{AppPaths, require_nonempty}; +use rocm_core::require_nonempty; use rocm_engine_protocol::{ DevicePolicy, ENGINE_RECIPE_CONTRACT_VERSION, EngineRecipeHint, GpuSelection, LaunchRequest, LaunchResponse, ResolveModelRequest, ResolveModelResponse, StopRequest, StopResponse, }; use serde_json::json; +use std::collections::BTreeMap; +use std::path::Path; use std::time::{Duration, Instant}; use crate::ServeHttpRequest; +use crate::docker::ContainerState; +use crate::registry::AimImageMetadata; const STOP_GRACE: Duration = Duration::from_secs(10); const READY_POLL_INTERVAL: Duration = Duration::from_millis(500); const DEFAULT_READY_TIMEOUT: Duration = Duration::from_mins(5); +/// Consecutive failed `docker inspect` calls (daemon down or erroring, as +/// opposed to a container that is definitively gone) tolerated before the +/// supervisor gives up: 60 polls is about 30 seconds. +const MAX_CONSECUTIVE_INSPECT_ERRORS: u32 = 60; + +/// Env vars this engine sets itself; a caller-supplied `extra_env` may not +/// override them. +const RESERVED_ENV_KEYS: [&str; 2] = ["HIP_VISIBLE_DEVICES", "VLLM_API_KEY"]; /// Time to wait for the AIM container to become ready before reporting failure. /// @@ -35,13 +47,10 @@ fn resolve_aim_ready_timeout(override_value: Option) -> Duration { } /// Deterministic, Docker-name-sanitized container identity for a service. -/// This *is* the service's identity for the AIM engine — no PID/start-ticks -/// tracking, since `docker run -d` leaves no long-lived client process to -/// supervise. Unlike vLLM/Lemonade, this engine never signals a client -/// process at all: the container is stopped by name (`docker stop`/`kill`), -/// which sidesteps the SIGKILL-orphan hazard those engines have to guard -/// against (a killed `docker run` client would otherwise leave the container -/// running, with nothing left to track it by). +/// The container is stopped and removed by name (`docker kill`/`rm -f`), never +/// by signalling a client process, so it cannot be orphaned by a killed +/// `docker run` client. `serve_http` is the long-lived process rocmd +/// supervises: it blocks until the container exits and records how it ended. pub(crate) fn container_name(service_id: &str) -> String { let sanitized: String = service_id .chars() @@ -56,23 +65,47 @@ pub(crate) fn container_name(service_id: &str) -> String { format!("rocm-aim-{sanitized}") } +/// Image metadata for the pre-launch `ResolveModel` step: the local daemon's +/// copy when the image is already present, otherwise a best-effort Docker Hub +/// lookup. Both are hints; callers must tolerate failure. +fn lookup_image_metadata(image_ref: &str) -> Result { + let image_ref = crate::registry::normalize_image_ref(image_ref); + if let Some(local) = crate::registry::local_image_metadata(&image_ref) { + return Ok(local); + } + crate::registry::fetch_image_metadata(&image_ref) +} + pub(crate) fn resolve_model_response(request: ResolveModelRequest) -> Result { + resolve_model_response_with(request, lookup_image_metadata) +} + +fn resolve_model_response_with( + request: ResolveModelRequest, + lookup: impl FnOnce(&str) -> Result, +) -> Result { require_nonempty(&request.model_ref, "model_ref")?; let engine_recipe = accepted_engine_recipe(request.engine_recipe)?; - let metadata = crate::registry::fetch_image_metadata(&request.model_ref)?; let mut warnings = Vec::new(); + let metadata = lookup(&request.model_ref).unwrap_or_else(|error| { + warnings.push(format!( + "could not read AIM image metadata for `{}` ({error:#}); it will be read from the \ + image after the pull", + request.model_ref + )); + AimImageMetadata::default() + }); - let device_policy = match metadata.accelerator_type.as_deref() { - Some("gpu") => DevicePolicy::GpuRequired, - Some("cpu") => DevicePolicy::CpuOnly, - _ => { - warnings.push( - "AIM image did not report AIM_ACCELERATOR_TYPE; assuming GPU-preferred".to_owned(), - ); - DevicePolicy::GpuPreferred - } - }; + // The caller's `--device` wins; image metadata only fills in when none was given. + let device_policy = + request + .device_policy + .unwrap_or(match metadata.accelerator_type.as_deref() { + Some("gpu") => DevicePolicy::GpuRequired, + Some("cpu") => DevicePolicy::CpuOnly, + _ => DevicePolicy::GpuPreferred, + }); if device_policy != DevicePolicy::CpuOnly && let Some(host_gfx) = rocm_core::detect_host_gfx_target() @@ -89,15 +122,14 @@ pub(crate) fn resolve_model_response(request: ResolveModelRequest) -> Result Result { let gpu_indices = rocm_engine_protocol::launch_gpu_indices(request.gpu_selection.as_ref()); let device_policy = request.device_policy.unwrap_or(DevicePolicy::GpuPreferred); - let paths = AppPaths::discover()?; - let state_path = paths - .engine_state_dir(crate::ENGINE_NAME) - .join(format!("{}.json", request.service_id)); + let state_path = crate::state::state_path(&request.service_id); - run_container( + let served_model_id = run_container( &request.service_id, &request.model_ref, &request.host, @@ -165,8 +194,10 @@ pub(crate) fn launch_service(request: LaunchRequest) -> Result { )?; let endpoint_url = endpoint_url(&request.host, request.port); crate::state::write_running_state( + &state_path, &request.service_id, &request.model_ref, + served_model_id.as_deref(), &request.host, request.port, device_policy, @@ -176,8 +207,8 @@ pub(crate) fn launch_service(request: LaunchRequest) -> Result { let log_path = format!("docker://{}", container_name(&request.service_id)); Ok(LaunchResponse { service_id: request.service_id, - // ponytail: `docker run -d` backgrounds the container itself — there is - // no long-lived client process to report a real PID for. + // ponytail: `docker run -d` backgrounds the container itself, so there is + // no client process to report a real PID for. pid: 0, endpoint_url, log_path, @@ -185,55 +216,107 @@ pub(crate) fn launch_service(request: LaunchRequest) -> Result { }) } -/// Starts the named service's container if it isn't already running. -/// Container names are globally unique in Docker (unlike host PIDs), so -/// unconditionally calling `docker run --name X` twice for the same service -/// fails with a name conflict — `launch_service` and `serve_http` must agree -/// on "already running? skip `docker run`" instead of vLLM's host-process -/// model, where two independent spawns are safe because each gets its own PID. +/// Which GPUs, if any, the container may use. +#[derive(Debug, PartialEq, Eq)] +struct GpuPlan { + available: bool, + indices: Vec, +} + +/// Resolves the caller's GPU request against what the host can actually use, +/// the way `engines/vllm` does: an authoritative empty probe means no usable +/// GPU, an empty request picks the first usable device, and an unknown probe +/// (WSL and other unprobeable surfaces) stays permissive. Native Windows has no +/// GPU passthrough shape (`docker::gpu_device_args`), so it never has a GPU. +fn plan_gpus(requested: &[u32], usable: Option>, windows: bool) -> Result { + let none = GpuPlan { + available: false, + indices: Vec::new(), + }; + if windows { + return Ok(none); + } + let Some(usable) = usable else { + return Ok(GpuPlan { + available: true, + indices: requested.to_vec(), + }); + }; + if usable.is_empty() { + return Ok(none); + } + if requested.is_empty() { + return Ok(GpuPlan { + available: true, + indices: vec![usable[0]], + }); + } + if let Some(missing) = requested.iter().find(|index| !usable.contains(index)) { + bail!( + "requested GPU {missing} is not available on this host; usable GPU indices: {usable:?}" + ); + } + Ok(GpuPlan { + available: true, + indices: requested.to_vec(), + }) +} + +/// Starts the named service's container unless it is already running, and +/// returns the model name the image says it serves under, if known. +/// +/// Container names are unique in Docker, so `docker run --name X` against a +/// name still held by an exited container (a crash, a reboot, a failed +/// readiness check) would fail with "name already in use" forever; a stale +/// container is therefore removed first. fn run_container( service_id: &str, model_ref: &str, host: &str, port: u16, device_policy: &DevicePolicy, - gpu_indices: &[u32], - extra_env: Option<&std::collections::BTreeMap>, + requested_gpus: &[u32], + extra_env: Option<&BTreeMap>, ) -> Result> { let name = container_name(service_id); - // Normalize once and reuse everywhere below: metadata and the actual - // `docker pull`/`run` must resolve the exact same image, or the fetched - // accelerator/GPU-arch metadata silently describes a different image - // than the one that runs (ROCMAI-554 review finding). + // Normalize once so the pull, the metadata and the run all use the same image. let image_ref = crate::registry::normalize_image_ref(model_ref); - let metadata = crate::registry::fetch_image_metadata(&image_ref).unwrap_or_else(|error| { - eprintln!( - "Warning: could not fetch AIM image metadata for `{image_ref}`: {error:#}; \ - assuming GPU-preferred with no accelerator-arch check." - ); - crate::registry::AimImageMetadata::default() - }); - if crate::docker::inspect_running(&name)?.unwrap_or(false) { - return Ok(metadata.canonical_name); + + match crate::docker::container_state(&name)? { + ContainerState::Running => { + return Ok(crate::registry::local_image_metadata(&image_ref) + .and_then(|metadata| metadata.canonical_name)); + } + ContainerState::Exited(_) => crate::docker::remove(&name)?, + ContainerState::Missing => {} } if !crate::docker::image_present(&image_ref)? { crate::docker::pull(&image_ref)?; } - - // `crate::docker::gpu_device_args` has no native-Windows device shape, and - // an image reporting AIM_ACCELERATOR_TYPE=cpu has nothing to pass GPU args - // to, so neither counts as "a GPU is usable" regardless of what the caller - // selected. - let gpu_usable = !cfg!(windows) && !gpu_indices.is_empty(); + // Read from the pulled image itself, so this works for any registry and + // offline, unlike the Docker Hub lookup. + let metadata = crate::registry::local_image_metadata(&image_ref).unwrap_or_default(); + + let gpus = plan_gpus( + requested_gpus, + rocm_core::usable_amd_gpu_indices(), + cfg!(windows), + )?; let image_wants_cpu = metadata.accelerator_type.as_deref() == Some("cpu"); let use_gpu = resolve_use_gpu( service_id, &image_ref, device_policy, - gpu_usable, + gpus.available, image_wants_cpu, )?; + if !use_gpu && matches!(device_policy, DevicePolicy::GpuPreferred) { + eprintln!( + "Warning: AIM service `{service_id}` will run on CPU (device_policy=gpu_preferred \ + and no usable GPU for `{image_ref}`); pass `--device gpu` to require a GPU." + ); + } let render_gid = if use_gpu { crate::docker::render_group_gid() } else { @@ -241,24 +324,28 @@ fn run_container( }; let mut env = extra_env.cloned().unwrap_or_default(); - // AIM images bundle vLLM's OpenAI server, so it is `VLLM_API_KEY` — the - // same variable `engines/vllm` sets on its own child process — that this + if let Some(key) = RESERVED_ENV_KEYS.iter().find(|key| env.contains_key(**key)) { + bail!("engine env `{key}` is managed by rocm-cli and cannot be overridden"); + } + // AIM images bundle vLLM's OpenAI server, so it is `VLLM_API_KEY` that this // engine's actual server enforces. The container is the only process this - // engine ever runs, so without this the require-api-key setting a public - // `rocm serve --engine aim` was given is silently never enforced. + // engine runs, so without this the require-api-key setting a public + // `rocm serve --engine aim` was given is silently never enforced. The value + // reaches docker through a private env file, never argv. if let Some(api_key) = rocm_engine_protocol::resolve_endpoint_api_key() { env.insert("VLLM_API_KEY".to_owned(), api_key); } - crate::docker::run( - &name, - &image_ref, + crate::docker::run(&crate::docker::RunSpec { + name: &name, + image: &image_ref, host, port, - crate::DEFAULT_AIM_PORT, - if use_gpu { gpu_indices } else { &[] }, + container_port: crate::DEFAULT_AIM_PORT, + use_gpu, + gpu_indices: if use_gpu { &gpus.indices } else { &[] }, render_gid, - &env, - )?; + env: &env, + })?; Ok(metadata.canonical_name) } @@ -278,8 +365,8 @@ fn resolve_use_gpu( if !gpu_usable || image_wants_cpu { bail!( "AIM service `{service_id}` has device_policy=gpu_required but no GPU is \ - usable for `{model_ref}` (native Windows has no GPU passthrough, no GPU \ - indices were selected, or the image reports AIM_ACCELERATOR_TYPE=cpu)" + usable for `{model_ref}` (native Windows has no GPU passthrough, the host \ + reports no usable GPU, or the image reports AIM_ACCELERATOR_TYPE=cpu)" ); } true @@ -289,7 +376,8 @@ fn resolve_use_gpu( } pub(crate) fn serve_http(request: ServeHttpRequest) -> Result<()> { - let canonical_name = run_container( + let name = container_name(&request.service_id); + let label_name = run_container( &request.service_id, &request.model_ref, &request.host, @@ -300,76 +388,140 @@ pub(crate) fn serve_http(request: ServeHttpRequest) -> Result<()> { )?; crate::state::write_running_state( + &request.state_path, &request.service_id, &request.model_ref, + label_name.as_deref(), &request.host, request.port, request.device_policy, request.engine_recipe.as_ref(), )?; - let name = container_name(&request.service_id); - let timeout = aim_ready_timeout(); - // The container self-reports under the registry's HF-style canonical name - // (e.g. `meta-llama/Llama-3.2-1B-Instruct`), not the Docker image reference - // used for `docker pull`/`run` — matching `/v1/models` against the image ref - // would never find it. Fall back to the image ref only when the registry - // didn't advertise a canonical name. - let expected_model_id = canonical_name.as_deref().unwrap_or(&request.model_ref); + // The container serves under the image's HF-style canonical name, not the + // Docker image reference used for `docker pull`/`run`. When the image didn't + // label one, wait for the container to list whatever it serves instead. match wait_for_container_ready( &name, &request.host, request.port, - expected_model_id, - timeout, + label_name.as_deref(), + aim_ready_timeout(), ) { - Ok(()) => block_until_container_stops(&name, &request.state_path), + Ok(served_model_id) => { + crate::state::record_served_model_id(&request.state_path, &served_model_id)?; + let exit = wait_for_exit( + || crate::docker::container_state(&name), + || std::thread::sleep(READY_POLL_INTERVAL), + ); + finish_after_exit(&name, &request.state_path, exit) + } Err(error) => { - // A readiness timeout (unlike an early container exit) leaves the - // container running; stop it so a failed launch doesn't leak an - // orphaned container nothing will ever signal again (the service - // record is about to be marked `failed`, not `running`, so - // `apps/rocmd`'s own stop path will never find it). - let _ = crate::docker::stop(&name, true, STOP_GRACE); + // The logs are already in `error`; remove the container so a failed + // launch leaks neither a running container nothing will signal + // again nor the stopped one that would hold its name. + let _ = crate::docker::remove(&name); crate::state::write_terminal_state(&request.state_path, "failed")?; Err(error) } } } -/// Blocks for the container's remaining lifetime, mirroring `engines/vllm` -/// and `engines/lemonade`'s blocking `child.wait()`: `apps/rocmd`'s generic -/// supervisor (`supervise_service`) treats this process's exit as the -/// service's exit, so returning right after readiness (as this used to) -/// made a still-running, actively-serving container read `stopped` the -/// instant it became ready. A transient `docker inspect` failure is treated -/// as "still running" rather than ending the loop, matching this file's -/// other `unwrap_or`-style tolerance of CLI hiccups elsewhere. -fn block_until_container_stops(name: &str, state_path: &std::path::Path) -> Result<()> { - while crate::docker::inspect_running(name).unwrap_or(Some(true)) != Some(false) { - std::thread::sleep(READY_POLL_INTERVAL); +/// How a supervised container ended. +#[derive(Debug, PartialEq, Eq)] +enum ContainerExit { + Code(i32), + /// Removed out from under the supervisor (`docker rm -f`, `docker system prune`). + Gone, + /// The daemon stayed unreachable past the tolerance. + DaemonLost, +} + +/// Blocks until the container stops, mirroring `engines/vllm` and +/// `engines/lemonade`'s blocking `child.wait()`: `apps/rocmd`'s generic +/// supervisor treats this process's exit as the service's exit, so returning +/// early would make a live container read `stopped`. A definitively missing +/// container ends the wait; transient `docker inspect` errors are tolerated +/// only [`MAX_CONSECUTIVE_INSPECT_ERRORS`] times in a row. +fn wait_for_exit( + mut poll: impl FnMut() -> Result, + mut pause: impl FnMut(), +) -> ContainerExit { + let mut consecutive_errors = 0; + loop { + match poll() { + Ok(ContainerState::Running) => consecutive_errors = 0, + Ok(ContainerState::Exited(code)) => return ContainerExit::Code(code), + Ok(ContainerState::Missing) => return ContainerExit::Gone, + Err(_) => { + consecutive_errors += 1; + if consecutive_errors >= MAX_CONSECUTIVE_INSPECT_ERRORS { + return ContainerExit::DaemonLost; + } + } + } + pause(); + } +} + +/// Records how the container ended, as `engines/vllm` does for its child: a +/// deliberate stop or clean exit is `stopped`; anything else is `failed` and +/// surfaces as an error (non-zero exit) so rocmd's recovery sees a crash. +fn finish_after_exit(name: &str, state_path: &Path, exit: ContainerExit) -> Result<()> { + // `stop_service` records its intent before killing, so a kill (exit 137) + // it caused is not mistaken for a crash. + let deliberate = matches!( + crate::state::read_status(state_path).as_deref(), + Some("stopping" | "stopped") + ); + if deliberate || exit == ContainerExit::Code(0) { + let _ = crate::docker::remove(name); + return crate::state::write_terminal_state(state_path, "stopped"); + } + crate::state::write_terminal_state(state_path, "failed")?; + match exit { + ContainerExit::Code(code) => bail!("AIM container `{name}` exited with code {code}"), + ContainerExit::Gone => bail!( + "AIM container `{name}` was removed while it was serving (`docker rm` or a prune?)" + ), + ContainerExit::DaemonLost => { + bail!("lost contact with the Docker daemon while supervising AIM container `{name}`") + } } - crate::state::write_terminal_state(state_path, "stopped") } /// Mirrors `engines/vllm/src/process.rs`'s `wait_for_vllm_ready` architecture: /// early-exit check, monotonic timeout, listed-then-inference-verified gate. -/// `docker::inspect_running` replaces `child.try_wait()` since there is no -/// `std::process::Child` for a detached container. +/// Returns the model name the server serves under (`expected_model` when the +/// image labelled one, otherwise whatever the server lists). fn wait_for_container_ready( name: &str, host: &str, port: u16, - model_ref: &str, + expected_model: Option<&str>, timeout: Duration, -) -> Result<()> { +) -> Result { let start = Instant::now(); let endpoint = endpoint_url(host, port); + let mut consecutive_errors = 0; loop { - if crate::docker::inspect_running(name)?.is_some_and(|running| !running) { - let exit_code = crate::docker::inspect_exit_code(name)?.unwrap_or(-1); - let tail = crate::docker::logs(name, 80).unwrap_or_default().join("\n"); - bail!("AIM container `{name}` exited early (exit code {exit_code}):\n{tail}"); + match crate::docker::container_state(name) { + Ok(ContainerState::Running) => consecutive_errors = 0, + Ok(ContainerState::Exited(exit_code)) => { + let tail = crate::docker::logs(name, 80).unwrap_or_default().join("\n"); + bail!("AIM container `{name}` exited early (exit code {exit_code}):\n{tail}"); + } + Ok(ContainerState::Missing) => { + bail!("AIM container `{name}` no longer exists; it was removed during startup") + } + Err(error) => { + consecutive_errors += 1; + if consecutive_errors >= MAX_CONSECUTIVE_INSPECT_ERRORS { + return Err(error.context(format!( + "lost contact with the Docker daemon while starting AIM container `{name}`" + ))); + } + } } if start.elapsed() > timeout { let tail = crate::docker::logs(name, 80).unwrap_or_default().join("\n"); @@ -379,12 +531,16 @@ fn wait_for_container_ready( ); } - let listed = - crate::state::query_loaded_model_endpoint(&endpoint, Some(model_ref)).unwrap_or(false); - if listed - && crate::state::query_inference_probe_endpoint(&endpoint, model_ref).unwrap_or(false) + let served = match expected_model { + Some(model) => crate::state::query_loaded_model_endpoint(&endpoint, Some(model)) + .unwrap_or(false) + .then(|| model.to_owned()), + None => crate::state::discover_served_model(&endpoint), + }; + if let Some(model) = served + && crate::state::query_inference_probe_endpoint(&endpoint, &model).unwrap_or(false) { - return Ok(()); + return Ok(model); } std::thread::sleep(READY_POLL_INTERVAL); @@ -394,13 +550,24 @@ fn wait_for_container_ready( pub(crate) fn stop_service(request: StopRequest) -> Result { require_nonempty(&request.service_id, "service_id")?; let name = container_name(&request.service_id); - let stopped = crate::docker::stop(&name, request.force, STOP_GRACE)?; - if stopped { - let state_path = AppPaths::discover()? - .engine_state_dir(crate::ENGINE_NAME) - .join(format!("{}.json", request.service_id)); - crate::state::write_terminal_state(&state_path, "stopped")?; - } + let state_path = crate::state::state_path(&request.service_id); + let prior_status = crate::state::read_status(&state_path); + // Record the intent first: the container's supervisor (`serve_http`) sees the + // exit within a poll interval and must read a deliberate stop, not a crash. + crate::state::write_status(&state_path, "stopping")?; + let stopped = match crate::docker::stop(&name, request.force, STOP_GRACE) { + Ok(stopped) => stopped, + Err(error) => { + if let Some(status) = prior_status { + let _ = crate::state::write_status(&state_path, &status); + } + return Err(error); + } + }; + // Release the name and the writable layer (downloaded weights) so a later + // restart can `docker run` the same name again. + crate::docker::remove(&name)?; + crate::state::write_terminal_state(&state_path, "stopped")?; Ok(StopResponse { stopped, graceful: stopped && !request.force, @@ -435,6 +602,25 @@ pub(crate) fn parse_gpu_indices_arg(value: Option<&str>) -> Result> { #[cfg(test)] mod tests { use super::*; + use std::cell::Cell; + + fn resolve_request(device_policy: Option) -> ResolveModelRequest { + ResolveModelRequest { + model_ref: "org/aim-foo:1".to_owned(), + runtime_id: None, + device_policy, + recipe_override: None, + engine_recipe: None, + } + } + + fn metadata(accelerator_type: Option<&str>, canonical_name: Option<&str>) -> AimImageMetadata { + AimImageMetadata { + accelerator_type: accelerator_type.map(ToOwned::to_owned), + canonical_name: canonical_name.map(ToOwned::to_owned), + ..AimImageMetadata::default() + } + } #[test] fn aim_ready_timeout_uses_default_without_override() { @@ -516,145 +702,486 @@ mod tests { assert!(!resolve_use_gpu("svc", "model", &DevicePolicy::GpuPreferred, true, true).unwrap()); assert!(resolve_use_gpu("svc", "model", &DevicePolicy::GpuPreferred, true, false).unwrap()); } + + #[test] + fn plan_gpus_resolves_an_empty_request_to_the_first_usable_gpu() { + assert_eq!( + plan_gpus(&[], Some(vec![1, 2]), false).unwrap(), + GpuPlan { + available: true, + indices: vec![1] + } + ); + } + + #[test] + fn plan_gpus_keeps_a_valid_request_and_rejects_an_unusable_one() { + assert_eq!( + plan_gpus(&[2], Some(vec![1, 2]), false).unwrap().indices, + vec![2] + ); + assert!(plan_gpus(&[5], Some(vec![1, 2]), false).is_err()); + } + + #[test] + fn plan_gpus_reports_no_gpu_for_an_empty_probe_and_on_windows() { + assert!(!plan_gpus(&[], Some(vec![]), false).unwrap().available); + assert!(!plan_gpus(&[0], Some(vec![0]), true).unwrap().available); + } + + #[test] + fn plan_gpus_stays_permissive_when_the_probe_is_unavailable() { + assert_eq!( + plan_gpus(&[], None, false).unwrap(), + GpuPlan { + available: true, + indices: vec![] + } + ); + } + + #[test] + fn resolve_reports_the_served_name_and_keeps_the_image_ref_as_canonical() { + let response = resolve_model_response_with(resolve_request(None), |_| { + Ok(metadata(Some("gpu"), Some("Org/Served-Model"))) + }) + .unwrap(); + assert_eq!(response.canonical_model_id, "org/aim-foo:1"); + assert_eq!( + response.served_model_id.as_deref(), + Some("Org/Served-Model") + ); + assert_eq!(response.device_policy, DevicePolicy::GpuRequired); + } + + #[test] + fn resolve_honors_the_callers_device_policy_over_image_metadata() { + let response = + resolve_model_response_with(resolve_request(Some(DevicePolicy::CpuOnly)), |_| { + Ok(metadata(Some("gpu"), None)) + }) + .unwrap(); + assert_eq!(response.device_policy, DevicePolicy::CpuOnly); + } + + #[test] + fn resolve_tolerates_a_failed_metadata_lookup_with_a_warning() { + let response = + resolve_model_response_with(resolve_request(None), |_| anyhow::bail!("offline")) + .unwrap(); + assert_eq!(response.served_model_id, None); + assert_eq!(response.device_policy, DevicePolicy::GpuPreferred); + assert!( + response + .warnings + .iter() + .any(|warning| warning.contains("offline")), + "{:?}", + response.warnings + ); + } + + #[test] + fn wait_for_exit_returns_the_exit_code() { + let mut states = vec![ + Ok(ContainerState::Running), + Ok(ContainerState::Running), + Ok(ContainerState::Exited(137)), + ] + .into_iter(); + let exit = wait_for_exit(|| states.next().unwrap(), || {}); + assert_eq!(exit, ContainerExit::Code(137)); + } + + #[test] + fn wait_for_exit_ends_when_the_container_is_removed_instead_of_spinning() { + let mut states = vec![Ok(ContainerState::Running), Ok(ContainerState::Missing)].into_iter(); + assert_eq!( + wait_for_exit(|| states.next().unwrap(), || {}), + ContainerExit::Gone + ); + } + + #[test] + fn wait_for_exit_tolerates_a_few_transient_errors_but_not_a_dead_daemon() { + let calls = Cell::new(0); + let transient = wait_for_exit( + || { + calls.set(calls.get() + 1); + match calls.get() { + 1..=3 => anyhow::bail!("daemon hiccup"), + 4 => Ok(ContainerState::Running), + _ => Ok(ContainerState::Exited(0)), + } + }, + || {}, + ); + assert_eq!(transient, ContainerExit::Code(0)); + + let polls = Cell::new(0); + let dead = wait_for_exit( + || { + polls.set(polls.get() + 1); + anyhow::bail!("daemon down") + }, + || {}, + ); + assert_eq!(dead, ContainerExit::DaemonLost); + assert_eq!(polls.get(), MAX_CONSECUTIVE_INSPECT_ERRORS); + } + + #[test] + fn finish_after_exit_records_failed_for_a_crash_and_stopped_for_a_clean_exit() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("svc.json"); + crate::state::write_status(&path, "running").unwrap(); + + // `docker rm` is best-effort here and may not exist on the test host. + let error = finish_after_exit("aim-x", &path, ContainerExit::Code(137)) + .expect_err("a crash must surface as an error so recovery runs"); + assert!(error.to_string().contains("137"), "{error}"); + assert_eq!(crate::state::read_status(&path).as_deref(), Some("failed")); + + finish_after_exit("aim-x", &path, ContainerExit::Code(0)).unwrap(); + assert_eq!(crate::state::read_status(&path).as_deref(), Some("stopped")); + } + + #[test] + fn finish_after_exit_treats_a_deliberate_stop_as_stopped_even_after_a_kill() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("svc.json"); + crate::state::write_status(&path, "stopping").unwrap(); + finish_after_exit("aim-x", &path, ContainerExit::Code(137)).unwrap(); + assert_eq!(crate::state::read_status(&path).as_deref(), Some("stopped")); + + crate::state::write_status(&path, "stopping").unwrap(); + finish_after_exit("aim-x", &path, ContainerExit::Gone).unwrap(); + assert_eq!(crate::state::read_status(&path).as_deref(), Some("stopped")); + } + + #[test] + fn finish_after_exit_fails_when_the_container_vanished_unexpectedly() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("svc.json"); + crate::state::write_status(&path, "running").unwrap(); + assert!(finish_after_exit("aim-x", &path, ContainerExit::Gone).is_err()); + assert_eq!(crate::state::read_status(&path).as_deref(), Some("failed")); + assert!(finish_after_exit("aim-x", &path, ContainerExit::DaemonLost).is_err()); + } } /// `run_container`/`wait_for_container_ready`/`serve_http`/`stop_service` all -/// shell out to `docker`, so (same reasoning as `docker.rs`'s -/// `subprocess_tests`) each scenario below re-spawns this test binary with a -/// fake `docker` on `PATH`. `fetch_image_metadata` (called unconditionally by -/// `run_container`) hits the real Docker Hub registry for a deliberately -/// bogus image and is expected to 404; that failure is already handled by a -/// production fallback (`unwrap_or_else` to `AimImageMetadata::default()`), -/// so no mock/seam is needed, just the real ~1s network round trip. +/// shell out to `docker`, so each scenario re-spawns this test binary with the +/// shared fake `docker` on `PATH` (see `docker::subprocess_tests`), a data dir +/// under the scratch dir, and, where readiness is reached, a minimal in-process +/// OpenAI-compatible server on loopback. `docker image inspect` is answered by +/// the fake, so nothing here touches the network or the real data dir. #[cfg(all(test, unix))] mod subprocess_tests { use super::*; - use rocm_core::fake_executable::write_fake_executable; - use std::process::{Command, Stdio}; - - const SENTINEL: &str = "ROCM_AIM_PROCESS_SUBPROCESS_TEST"; - - const FAKE_DOCKER_SCRIPT: &str = r#"#!/bin/sh -case "$1" in - run) - [ -n "$FAKE_DOCKER_RUN_STDERR" ] && echo "$FAKE_DOCKER_RUN_STDERR" >&2 - exit "${FAKE_DOCKER_RUN_EXIT:-0}" - ;; - kill|stop) - if [ -n "$FAKE_DOCKER_STOP_NO_SUCH_CONTAINER" ]; then - echo "Error response from daemon: No such container: fake" >&2 - exit 1 - fi - exit "${FAKE_DOCKER_STOP_EXIT:-0}" - ;; - inspect) - case "$3" in - *Running*) echo "${FAKE_DOCKER_RUNNING:-false}" ;; - *ExitCode*) echo "${FAKE_DOCKER_EXIT_CODE:-0}" ;; - esac - exit 0 - ;; - image) - exit "${FAKE_DOCKER_IMAGE_PRESENT_EXIT:-1}" - ;; - pull) - exit "${FAKE_DOCKER_PULL_EXIT:-0}" - ;; - logs) - [ -n "$FAKE_DOCKER_LOGS" ] && echo "$FAKE_DOCKER_LOGS" - exit 0 - ;; -esac -exit 0 -"#; - - fn in_subprocess() -> bool { - std::env::var_os(SENTINEL).is_some() - } - - fn run_with_fake_docker(test_path: &str, extra_env: &[(&str, &str)]) { - let dir = std::env::temp_dir().join(format!( - "aim-process-test-{}-{}", - std::process::id(), - std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap() - .as_nanos() - )); - let fake_bin = dir.join("bin"); - std::fs::create_dir_all(&fake_bin).unwrap(); - write_fake_executable(&fake_bin.join("docker"), FAKE_DOCKER_SCRIPT).unwrap(); + use crate::docker::subprocess_tests::{in_subprocess, run_with_fake_docker}; + use std::io::{BufRead, BufReader, Read, Write}; + use std::net::TcpListener; + + const IMAGE: &str = "org/aim-fixture:1"; + const SERVED: &str = "Org/Served-Model"; + + /// Serves `GET /v1/models` (listing `model`) and `POST /v1/chat/completions` + /// (always 200) on an ephemeral loopback port; returns the port. + fn spawn_fake_openai(model: &'static str) -> u16 { + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let port = listener.local_addr().unwrap().port(); + std::thread::spawn(move || { + for stream in listener.incoming().flatten() { + std::thread::spawn(move || { + let mut reader = BufReader::new(stream.try_clone().unwrap()); + let mut request_line = String::new(); + reader.read_line(&mut request_line).ok(); + let mut content_length = 0usize; + loop { + let mut header = String::new(); + if reader.read_line(&mut header).unwrap_or(0) == 0 || header == "\r\n" { + break; + } + if let Some(value) = + header.to_ascii_lowercase().strip_prefix("content-length:") + { + content_length = value.trim().parse().unwrap_or(0); + } + } + let mut body = vec![0; content_length]; + reader.read_exact(&mut body).ok(); + let payload = if request_line.starts_with("GET /v1/models") { + format!(r#"{{"data":[{{"id":"{model}"}}]}}"#) + } else { + r#"{"choices":[]}"#.to_owned() + }; + let mut stream = stream; + let _ = write!( + stream, + "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{payload}", + payload.len() + ); + }); + } + }); + port + } - let mut entries = vec![fake_bin]; - entries.extend(std::env::split_paths( - &std::env::var_os("PATH").unwrap_or_default(), - )); + fn data_dir() -> std::path::PathBuf { + std::path::PathBuf::from(std::env::var_os("ROCM_CLI_DATA_DIR").expect("data dir")) + } + + fn serve_request( + service_id: &str, + port: u16, + state_path: std::path::PathBuf, + ) -> ServeHttpRequest { + ServeHttpRequest { + service_id: service_id.to_owned(), + model_ref: IMAGE.to_owned(), + host: "127.0.0.1".to_owned(), + port, + device_policy: DevicePolicy::CpuOnly, + gpu_indices: vec![], + state_path, + engine_recipe: None, + } + } + + fn status_of(path: &Path) -> Option { + crate::state::read_status(path) + } + + const LABELLED_IMAGE: (&str, &str) = ( + "FAKE_DOCKER_IMAGE_CONFIG", + r#"{"Env":["AIM_ACCELERATOR_TYPE=cpu"],"Labels":{"com.amd.aim.model.canonicalName":"Org/Served-Model"}}"#, + ); + + #[test] + fn serve_http_reaches_ready_then_marks_a_clean_container_exit_stopped() { + if in_subprocess() { + let port = spawn_fake_openai(SERVED); + let state_path = data_dir().join("svc.json"); + serve_http(serve_request("svc-clean", port, state_path.clone())).unwrap(); + assert_eq!(status_of(&state_path).as_deref(), Some("stopped")); + let state = crate::state::read_service_state(&state_path).unwrap(); + assert_eq!(state["served_model_id"], SERVED); + return; + } + let fixture = run_with_fake_docker( + "process::subprocess_tests::serve_http_reaches_ready_then_marks_a_clean_container_exit_stopped", + &[ + LABELLED_IMAGE, + ("FAKE_DOCKER_IMAGE_PRESENT_EXIT", "0"), + ( + "FAKE_DOCKER_INSPECT_SEQUENCE", + "missing,true 0,true 0,false 0", + ), + ], + ); + let argv = fixture.argv(); + assert!( + argv.iter() + .any(|line| line.starts_with("run -d --name rocm-aim-svc-clean")), + "{argv:?}" + ); + assert!( + argv.contains(&"rm -f rocm-aim-svc-clean".to_owned()), + "{argv:?}" + ); + } + + #[test] + fn serve_http_discovers_the_served_name_when_the_image_has_no_label() { + if in_subprocess() { + let port = spawn_fake_openai(SERVED); + let state_path = data_dir().join("svc.json"); + serve_http(serve_request("svc-discover", port, state_path.clone())).unwrap(); + let state = crate::state::read_service_state(&state_path).unwrap(); + assert_eq!(state["served_model_id"], SERVED); + return; + } + run_with_fake_docker( + "process::subprocess_tests::serve_http_discovers_the_served_name_when_the_image_has_no_label", + &[ + ("FAKE_DOCKER_IMAGE_CONFIG", r#"{"Env":[]}"#), + ("FAKE_DOCKER_IMAGE_PRESENT_EXIT", "0"), + ( + "FAKE_DOCKER_INSPECT_SEQUENCE", + "missing,true 0,true 0,false 0", + ), + ], + ); + } + + #[test] + fn serve_http_records_failed_and_errors_when_the_container_crashes_after_ready() { + if in_subprocess() { + let port = spawn_fake_openai(SERVED); + let state_path = data_dir().join("svc.json"); + let error = serve_http(serve_request("svc-crash", port, state_path.clone())) + .expect_err("a crash after readiness must be an error so rocmd recovers it"); + assert!(error.to_string().contains("137"), "{error}"); + assert_eq!(status_of(&state_path).as_deref(), Some("failed")); + return; + } + run_with_fake_docker( + "process::subprocess_tests::serve_http_records_failed_and_errors_when_the_container_crashes_after_ready", + &[ + LABELLED_IMAGE, + ("FAKE_DOCKER_IMAGE_PRESENT_EXIT", "0"), + ( + "FAKE_DOCKER_INSPECT_SEQUENCE", + "missing,true 0,true 0,false 137", + ), + ], + ); + } - let mut command = Command::new(std::env::current_exe().unwrap()); - command - .args(["--exact", test_path, "--nocapture"]) - .env(SENTINEL, "1") - .env("PATH", std::env::join_paths(entries).unwrap()) - .stdin(Stdio::null()); - for (key, value) in extra_env { - command.env(key, value); + #[test] + fn serve_http_records_failed_when_the_container_is_removed_while_serving() { + if in_subprocess() { + let port = spawn_fake_openai(SERVED); + let state_path = data_dir().join("svc.json"); + serve_http(serve_request("svc-gone", port, state_path.clone())) + .expect_err("a vanished container must end the supervisor, not spin"); + assert_eq!(status_of(&state_path).as_deref(), Some("failed")); + return; } - let status = command.status().unwrap(); - assert!(status.success(), "helper subprocess failed for {test_path}"); + run_with_fake_docker( + "process::subprocess_tests::serve_http_records_failed_when_the_container_is_removed_while_serving", + &[ + LABELLED_IMAGE, + ("FAKE_DOCKER_IMAGE_PRESENT_EXIT", "0"), + ( + "FAKE_DOCKER_INSPECT_SEQUENCE", + "missing,true 0,true 0,missing", + ), + ], + ); + } - std::fs::remove_dir_all(&dir).ok(); + #[test] + fn serve_http_removes_the_container_and_marks_failed_on_readiness_failure() { + if in_subprocess() { + let state_path = data_dir().join("svc.json"); + serve_http(serve_request("svc-not-ready", 18084, state_path.clone())) + .expect_err("readiness failure should surface as an error"); + assert_eq!(status_of(&state_path).as_deref(), Some("failed")); + return; + } + let fixture = run_with_fake_docker( + "process::subprocess_tests::serve_http_removes_the_container_and_marks_failed_on_readiness_failure", + &[ + LABELLED_IMAGE, + ("FAKE_DOCKER_IMAGE_PRESENT_EXIT", "0"), + ("FAKE_DOCKER_INSPECT_SEQUENCE", "missing,false 1"), + ("FAKE_DOCKER_LOGS", "widget crashed"), + ], + ); + assert!( + fixture + .argv() + .contains(&"rm -f rocm-aim-svc-not-ready".to_owned()), + "{:?}", + fixture.argv() + ); } - const BOGUS_IMAGE: &str = "rocm-cli-test/does-not-exist-on-docker-hub-aim-fixture"; + #[test] + fn run_container_removes_a_stale_exited_container_before_running_again() { + if in_subprocess() { + run_container( + "svc-stale", + IMAGE, + "127.0.0.1", + 18085, + &DevicePolicy::CpuOnly, + &[], + None, + ) + .unwrap(); + return; + } + let fixture = run_with_fake_docker( + "process::subprocess_tests::run_container_removes_a_stale_exited_container_before_running_again", + &[ + ("FAKE_DOCKER_IMAGE_PRESENT_EXIT", "0"), + ("FAKE_DOCKER_INSPECT_SEQUENCE", "false 137"), + ], + ); + let argv = fixture.argv(); + let rm = argv + .iter() + .position(|line| line == "rm -f rocm-aim-svc-stale"); + let run = argv + .iter() + .position(|line| line.starts_with("run -d --name rocm-aim-svc-stale")); + assert!( + matches!((rm, run), (Some(rm), Some(run)) if rm < run), + "the stale container must be removed before `docker run`: {argv:?}" + ); + } #[test] fn run_container_skips_docker_run_when_already_running() { if in_subprocess() { - let result = run_container( + run_container( "svc-already-running", - BOGUS_IMAGE, + IMAGE, "127.0.0.1", 18080, &DevicePolicy::GpuPreferred, &[], None, - ); - assert!(result.is_ok(), "{result:?}"); + ) + .unwrap(); return; } - run_with_fake_docker( + let fixture = run_with_fake_docker( "process::subprocess_tests::run_container_skips_docker_run_when_already_running", &[ - ("FAKE_DOCKER_RUNNING", "true"), + ("FAKE_DOCKER_INSPECT_SEQUENCE", "true 0"), ("FAKE_DOCKER_PULL_EXIT", "1"), ("FAKE_DOCKER_RUN_EXIT", "1"), ], ); + assert!( + !fixture.argv().iter().any(|line| line.starts_with("run ")), + "{:?}", + fixture.argv() + ); } #[test] fn run_container_pulls_when_image_absent_then_runs() { if in_subprocess() { - let result = run_container( + run_container( "svc-pull-then-run", - BOGUS_IMAGE, + IMAGE, "127.0.0.1", 18081, - &DevicePolicy::GpuPreferred, + &DevicePolicy::CpuOnly, &[], None, - ); - assert!(result.is_ok(), "{result:?}"); + ) + .unwrap(); return; } - run_with_fake_docker( + let fixture = run_with_fake_docker( "process::subprocess_tests::run_container_pulls_when_image_absent_then_runs", &[ + ("FAKE_DOCKER_INSPECT_SEQUENCE", "missing"), ("FAKE_DOCKER_IMAGE_PRESENT_EXIT", "1"), - ("FAKE_DOCKER_PULL_EXIT", "0"), - ("FAKE_DOCKER_RUN_EXIT", "0"), ], ); + assert!( + fixture + .argv() + .iter() + .any(|line| line == "pull org/aim-fixture:1") + ); } #[test] @@ -662,7 +1189,7 @@ exit 0 if in_subprocess() { let error = run_container( "svc-pull-fails", - BOGUS_IMAGE, + IMAGE, "127.0.0.1", 18082, &DevicePolicy::GpuPreferred, @@ -676,6 +1203,7 @@ exit 0 run_with_fake_docker( "process::subprocess_tests::run_container_bails_when_pull_fails", &[ + ("FAKE_DOCKER_INSPECT_SEQUENCE", "missing"), ("FAKE_DOCKER_IMAGE_PRESENT_EXIT", "1"), ("FAKE_DOCKER_PULL_EXIT", "1"), ], @@ -687,10 +1215,10 @@ exit 0 if in_subprocess() { let error = run_container( "svc-run-fails", - BOGUS_IMAGE, + IMAGE, "127.0.0.1", 18083, - &DevicePolicy::GpuPreferred, + &DevicePolicy::CpuOnly, &[], None, ) @@ -703,6 +1231,7 @@ exit 0 run_with_fake_docker( "process::subprocess_tests::run_container_bails_when_docker_run_fails", &[ + ("FAKE_DOCKER_INSPECT_SEQUENCE", "missing"), ("FAKE_DOCKER_IMAGE_PRESENT_EXIT", "0"), ("FAKE_DOCKER_RUN_EXIT", "1"), ("FAKE_DOCKER_RUN_STDERR", "boom"), @@ -710,6 +1239,65 @@ exit 0 ); } + #[test] + fn run_container_rejects_reserved_env_overrides_before_running_docker() { + if in_subprocess() { + for key in ["VLLM_API_KEY", "HIP_VISIBLE_DEVICES"] { + let env = BTreeMap::from([(key.to_owned(), "x".to_owned())]); + let error = run_container( + "svc-reserved", + IMAGE, + "127.0.0.1", + 18086, + &DevicePolicy::CpuOnly, + &[], + Some(&env), + ) + .expect_err("reserved key must be rejected"); + assert!(error.to_string().contains(key), "{error}"); + } + return; + } + let fixture = run_with_fake_docker( + "process::subprocess_tests::run_container_rejects_reserved_env_overrides_before_running_docker", + &[ + ("FAKE_DOCKER_INSPECT_SEQUENCE", "missing"), + ("FAKE_DOCKER_IMAGE_PRESENT_EXIT", "0"), + ], + ); + assert!(!fixture.argv().iter().any(|line| line.starts_with("run "))); + } + + #[test] + fn run_container_passes_the_catalog_env_through_the_env_file_not_argv() { + if in_subprocess() { + let env = BTreeMap::from([("AIM_PROFILE_ID".to_owned(), "vllm-profile".to_owned())]); + run_container( + "svc-env", + IMAGE, + "127.0.0.1", + 18087, + &DevicePolicy::CpuOnly, + &[], + Some(&env), + ) + .unwrap(); + return; + } + let fixture = run_with_fake_docker( + "process::subprocess_tests::run_container_passes_the_catalog_env_through_the_env_file_not_argv", + &[ + ("FAKE_DOCKER_INSPECT_SEQUENCE", "missing"), + ("FAKE_DOCKER_IMAGE_PRESENT_EXIT", "0"), + ], + ); + assert!(!fixture.argv().join("\n").contains("vllm-profile")); + assert_eq!( + fixture.env_copy().as_deref(), + Some("AIM_PROFILE_ID=vllm-profile\n") + ); + } + #[test] fn wait_for_container_ready_bails_when_container_exits_early() { if in_subprocess() { @@ -717,7 +1305,7 @@ exit 0 "fake-container", "127.0.0.1", 19999, - "model-x", + Some("model-x"), Duration::from_secs(5), ) .expect_err("early container exit should bail"); @@ -737,18 +1325,35 @@ exit 0 ); } + #[test] + fn wait_for_container_ready_bails_when_the_container_is_removed() { + if in_subprocess() { + let error = wait_for_container_ready( + "fake-container", + "127.0.0.1", + 19998, + Some("model-x"), + Duration::from_secs(5), + ) + .expect_err("a missing container should bail"); + assert!(error.to_string().contains("no longer exists"), "{error}"); + return; + } + run_with_fake_docker( + "process::subprocess_tests::wait_for_container_ready_bails_when_the_container_is_removed", + &[("FAKE_DOCKER_INSPECT_MISSING", "1")], + ); + } + #[test] fn wait_for_container_ready_bails_on_timeout() { if in_subprocess() { - // Unbound port: the readiness HTTP probe fails fast (connection - // refused), so this only exercises the timeout bookkeeping, not a - // real OpenAI-compatible server (left as a judgment call — see the - // PR summary). + // Unbound port: the readiness HTTP probe fails fast (connection refused). let error = wait_for_container_ready( "fake-container", "127.0.0.1", 59123, - "model-x", + Some("model-x"), Duration::from_millis(10), ) .expect_err("timeout should bail"); @@ -765,49 +1370,7 @@ exit 0 } #[test] - fn serve_http_stops_container_and_marks_state_failed_on_readiness_failure() { - if in_subprocess() { - let dir = std::env::temp_dir().join(format!( - "aim-serve-http-test-{}", - std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap() - .as_nanos() - )); - let state_path = dir.join("svc.json"); - - let request = ServeHttpRequest { - service_id: "svc-serve-http-fails".to_owned(), - model_ref: BOGUS_IMAGE.to_owned(), - host: "127.0.0.1".to_owned(), - port: 18084, - device_policy: DevicePolicy::GpuPreferred, - gpu_indices: vec![], - runtime_id: None, - env_id: None, - state_path: state_path.clone(), - log_path: None, - engine_recipe: None, - }; - serve_http(request).expect_err("readiness failure should surface as an error"); - - let state = crate::state::read_service_state(&state_path).unwrap(); - assert_eq!(state.get("status").and_then(|v| v.as_str()), Some("failed")); - - std::fs::remove_dir_all(&dir).ok(); - return; - } - run_with_fake_docker( - "process::subprocess_tests::serve_http_stops_container_and_marks_state_failed_on_readiness_failure", - &[ - ("FAKE_DOCKER_RUNNING", "false"), - ("FAKE_DOCKER_EXIT_CODE", "1"), - ], - ); - } - - #[test] - fn stop_service_force_stop_succeeds_and_writes_stopped_state() { + fn stop_service_records_intent_kills_removes_and_writes_stopped_state() { if in_subprocess() { let response = stop_service(StopRequest { service_id: "svc-stop-succeeds".to_owned(), @@ -816,32 +1379,25 @@ exit 0 .unwrap(); assert!(response.stopped); assert!(!response.graceful); - let state_path = crate::state::state_path("svc-stop-succeeds"); - let state = crate::state::read_service_state(&state_path).unwrap(); - assert_eq!( - state.get("status").and_then(|v| v.as_str()), - Some("stopped") - ); + assert_eq!(status_of(&state_path).as_deref(), Some("stopped")); return; } - let dir = std::env::temp_dir().join(format!( - "aim-stop-service-test-{}", - std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap() - .as_nanos() - )); - std::fs::create_dir_all(&dir).unwrap(); - run_with_fake_docker( - "process::subprocess_tests::stop_service_force_stop_succeeds_and_writes_stopped_state", - &[("ROCM_CLI_DATA_DIR", dir.to_str().unwrap())], + let fixture = run_with_fake_docker( + "process::subprocess_tests::stop_service_records_intent_kills_removes_and_writes_stopped_state", + &[], + ); + assert_eq!( + fixture.argv(), + vec![ + "kill rocm-aim-svc-stop-succeeds", + "rm -f rocm-aim-svc-stop-succeeds" + ] ); - std::fs::remove_dir_all(&dir).ok(); } #[test] - fn stop_service_missing_container_reports_not_stopped_without_writing_state() { + fn stop_service_of_a_missing_container_reports_not_stopped_but_still_cleans_up() { if in_subprocess() { let response = stop_service(StopRequest { service_id: "svc-stop-missing".to_owned(), @@ -850,26 +1406,37 @@ exit 0 .unwrap(); assert!(!response.stopped); assert!(!response.graceful); - - let state_path = crate::state::state_path("svc-stop-missing"); - assert!(!state_path.exists()); return; } - let dir = std::env::temp_dir().join(format!( - "aim-stop-service-missing-test-{}", - std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap() - .as_nanos() - )); - std::fs::create_dir_all(&dir).unwrap(); run_with_fake_docker( - "process::subprocess_tests::stop_service_missing_container_reports_not_stopped_without_writing_state", + "process::subprocess_tests::stop_service_of_a_missing_container_reports_not_stopped_but_still_cleans_up", &[ - ("ROCM_CLI_DATA_DIR", dir.to_str().unwrap()), ("FAKE_DOCKER_STOP_NO_SUCH_CONTAINER", "1"), + ("FAKE_DOCKER_RM_NO_SUCH_CONTAINER", "1"), + ], + ); + } + + #[test] + fn stop_service_failure_restores_the_prior_status_and_keeps_the_container() { + if in_subprocess() { + let state_path = crate::state::state_path("svc-stop-fails"); + crate::state::write_status(&state_path, "running").unwrap(); + stop_service(StopRequest { + service_id: "svc-stop-fails".to_owned(), + force: true, + }) + .expect_err("a failed kill must be reported, not swallowed"); + assert_eq!(status_of(&state_path).as_deref(), Some("running")); + return; + } + let fixture = run_with_fake_docker( + "process::subprocess_tests::stop_service_failure_restores_the_prior_status_and_keeps_the_container", + &[ + ("FAKE_DOCKER_STOP_EXIT", "1"), + ("FAKE_DOCKER_STOP_STDERR", "permission denied"), ], ); - std::fs::remove_dir_all(&dir).ok(); + assert!(!fixture.argv().iter().any(|line| line.starts_with("rm "))); } } diff --git a/engines/aim/src/registry.rs b/engines/aim/src/registry.rs index 47b74d9ce..29e0be5d0 100644 --- a/engines/aim/src/registry.rs +++ b/engines/aim/src/registry.rs @@ -2,9 +2,11 @@ // // SPDX-License-Identifier: MIT -//! Docker Hub v2 registry client for AIM image metadata (ROCMAI-553 finding): -//! the `amdenterpriseai` org's image config blob carries structured -//! `Env`/`Labels` fetchable via two plain HTTP calls, without pulling the image. +//! AIM image metadata (ROCMAI-553 finding): the image config carries structured +//! `Env`/`Labels`. Read it from the local daemon once the image is present (any +//! registry), or, as a best-effort pre-pull hint, from Docker Hub's v2 API via +//! two plain HTTP calls. The Docker Hub lookup is only ever a hint: callers must +//! tolerate its failure (offline, rate-limited, or a non-Docker-Hub registry). use anyhow::{Context, Result}; use serde_json::Value; @@ -17,7 +19,6 @@ pub(crate) const DEFAULT_ORG: &str = "amdenterpriseai"; #[derive(Debug, Default, Clone)] pub(crate) struct AimImageMetadata { - pub accelerator_family: Option, pub accelerator_type: Option, pub pytorch_rocm_arch: Vec, pub canonical_name: Option, @@ -36,21 +37,34 @@ pub(crate) fn normalize_image_ref(image_ref: &str) -> String { } } -/// Splits a normalized (see [`normalize_image_ref`]) `org/repo[:tag|@digest]` -/// ref into `(repo, reference)`. Treats `@sha256:...` as a digest suffix, and -/// only looks for a tag's `:` after the last `/`, so a `host:port/repo` or -/// `repo@sha256:...` ref isn't mis-split by a naive first-`:` search. +/// Splits a normalized (see [`normalize_image_ref`]) `org/repo[:tag][@digest]` +/// ref into `(repo, reference)`. A digest wins over a tag (registries resolve +/// by digest), and the tag is dropped from `repo` so `repo:tag@sha256:...` +/// doesn't leave `:tag` in the name used for the token scope and manifest URL. +/// A tag's `:` is only looked for after the last `/`, so a `host:port/repo` ref +/// isn't mis-split by a naive first-`:` search. fn split_image_ref(image_ref: &str) -> (String, String) { - if let Some((name, digest)) = image_ref.split_once('@') { - return (name.to_owned(), digest.to_owned()); - } - let tag_sep = image_ref + let (name, digest) = match image_ref.split_once('@') { + Some((name, digest)) => (name, Some(digest)), + None => (image_ref, None), + }; + let tag_sep = name .rfind(':') - .filter(|&i| i > image_ref.rfind('/').unwrap_or(0)); - match tag_sep { - Some(i) => (image_ref[..i].to_owned(), image_ref[i + 1..].to_owned()), - None => (image_ref.to_owned(), "latest".to_owned()), - } + .filter(|&i| i > name.rfind('/').unwrap_or(0)); + let (repo, tag) = match tag_sep { + Some(i) => (&name[..i], Some(&name[i + 1..])), + None => (name, None), + }; + let reference = digest.or(tag).unwrap_or("latest"); + (repo.to_owned(), reference.to_owned()) +} + +/// Whether the (normalized) ref lives on Docker Hub, the only registry +/// [`fetch_image_metadata`] can read. The first path component is a registry +/// host when it looks like one (contains `.` or `:`, or is `localhost`). +fn is_docker_hub_ref(normalized: &str) -> bool { + let first = normalized.split('/').next().unwrap_or_default(); + !(first.contains('.') || first.contains(':') || first == "localhost") } fn agent() -> ureq::Agent { @@ -77,8 +91,13 @@ fn anonymous_pull_token(repo: &str) -> Result { .with_context(|| format!("registry auth response for {repo} had no token")) } +/// Best-effort pre-pull metadata hint from Docker Hub. Errors for any other +/// registry; callers treat every error as "no hint". pub(crate) fn fetch_image_metadata(image_ref: &str) -> Result { let normalized = normalize_image_ref(image_ref); + if !is_docker_hub_ref(&normalized) { + anyhow::bail!("`{normalized}` is not a Docker Hub image; skipping the registry lookup"); + } let (repo, reference) = split_image_ref(&normalized); let token = anonymous_pull_token(&repo)?; @@ -112,21 +131,32 @@ pub(crate) fn fetch_image_metadata(image_ref: &str) -> Result .into_json() .with_context(|| format!("failed to parse config blob for {image_ref}"))?; - Ok(parse_config_blob(&config)) + Ok(parse_image_config( + config.get("config").unwrap_or(&Value::Null), + )) +} + +/// Metadata from the local daemon's copy of the image, or `None` if the image +/// isn't present locally. +pub(crate) fn local_image_metadata(image_ref: &str) -> Option { + let config = crate::docker::image_config(image_ref).ok().flatten()?; + Some(parse_image_config(&config)) } -fn parse_config_blob(config: &Value) -> AimImageMetadata { +/// Parses an image `Config` object (`{"Env": [...], "Labels": {...}}`), the +/// shape shared by Docker Hub's config blob (under `config`) and +/// `docker image inspect`'s `.Config`. +fn parse_image_config(config: &Value) -> AimImageMetadata { let mut metadata = AimImageMetadata::default(); let env_entries = config - .pointer("/config/Env") + .get("Env") .and_then(Value::as_array) .map(|entries| entries.iter().filter_map(Value::as_str).collect::>()) .unwrap_or_default(); for entry in env_entries { if let Some((key, value)) = entry.split_once('=') { match key { - "AIM_ACCELERATOR_FAMILY" => metadata.accelerator_family = Some(value.to_owned()), "AIM_ACCELERATOR_TYPE" => metadata.accelerator_type = Some(value.to_owned()), "PYTORCH_ROCM_ARCH" => { metadata.pytorch_rocm_arch = value.split(';').map(ToOwned::to_owned).collect(); @@ -136,7 +166,7 @@ fn parse_config_blob(config: &Value) -> AimImageMetadata { } } - if let Some(labels) = config.pointer("/config/Labels").and_then(Value::as_object) { + if let Some(labels) = config.get("Labels").and_then(Value::as_object) { metadata.canonical_name = labels .get("com.amd.aim.model.canonicalName") .and_then(Value::as_str) @@ -197,6 +227,37 @@ mod tests { ); } + #[test] + fn split_image_ref_strips_the_tag_when_a_digest_is_present() { + // `repo:tag@sha256:...` must not leave `:tag` in the repo name, or the + // token scope and manifest URL point at a repository that doesn't exist. + assert_eq!( + split_image_ref("amdenterpriseai/aim-foo:0.12.0-preview@sha256:abcd1234"), + ( + "amdenterpriseai/aim-foo".to_owned(), + "sha256:abcd1234".to_owned() + ) + ); + } + + #[test] + fn is_docker_hub_ref_rejects_other_registries() { + assert!(is_docker_hub_ref("amdenterpriseai/aim-foo:1")); + assert!(!is_docker_hub_ref("myregistry.example.com:5000/aim-foo:1")); + assert!(!is_docker_hub_ref("ghcr.io/org/aim-foo:1")); + assert!(!is_docker_hub_ref("localhost/aim-foo")); + } + + #[test] + fn fetch_image_metadata_skips_non_docker_hub_registries_without_network() { + let error = fetch_image_metadata("myregistry.example.com:5000/aim-foo:1") + .expect_err("a private registry is not looked up"); + assert!( + error.to_string().contains("not a Docker Hub image"), + "{error}" + ); + } + #[test] fn split_image_ref_handles_registry_host_with_port() { // The first `:` here belongs to the registry host's port, not a tag @@ -212,22 +273,18 @@ mod tests { } #[test] - fn parse_config_blob_extracts_known_fields() { + fn parse_image_config_extracts_known_fields() { let config = serde_json::json!({ - "config": { - "Env": [ - "PATH=/usr/bin", - "AIM_ACCELERATOR_FAMILY=radeon", - "AIM_ACCELERATOR_TYPE=gpu", - "PYTORCH_ROCM_ARCH=gfx1100;gfx1151", - ], - "Labels": { - "com.amd.aim.model.canonicalName": "Llama-3.2-1B-Instruct", - }, + "Env": [ + "PATH=/usr/bin", + "AIM_ACCELERATOR_TYPE=gpu", + "PYTORCH_ROCM_ARCH=gfx1100;gfx1151", + ], + "Labels": { + "com.amd.aim.model.canonicalName": "Llama-3.2-1B-Instruct", }, }); - let metadata = parse_config_blob(&config); - assert_eq!(metadata.accelerator_family.as_deref(), Some("radeon")); + let metadata = parse_image_config(&config); assert_eq!(metadata.accelerator_type.as_deref(), Some("gpu")); assert_eq!(metadata.pytorch_rocm_arch, vec!["gfx1100", "gfx1151"]); assert_eq!( diff --git a/engines/aim/src/state.rs b/engines/aim/src/state.rs index 31adc2453..84e329454 100644 --- a/engines/aim/src/state.rs +++ b/engines/aim/src/state.rs @@ -24,7 +24,11 @@ pub(crate) fn healthcheck_service(request: HealthcheckRequest) -> Result EngineRecipeHint { hint } +#[allow(clippy::too_many_arguments)] pub(crate) fn write_running_state( + state_path: &Path, service_id: &str, model_ref: &str, + served_model_id: Option<&str>, host: &str, port: u16, device_policy: DevicePolicy, @@ -114,13 +121,14 @@ pub(crate) fn write_running_state( // before writing it to disk while keeping the keys for visibility. let engine_recipe = engine_recipe.map(redact_engine_recipe_extra_env); write_state( - &state_path(service_id), + state_path, &json!({ "service_id": service_id, "engine": crate::ENGINE_NAME, "status": "running", "container_name": crate::process::container_name(service_id), "model_ref": model_ref, + "served_model_id": served_model_id, "host": host, "port": port, "endpoint_url": endpoint, @@ -131,6 +139,28 @@ pub(crate) fn write_running_state( ) } +/// Records the model name the running container actually serves under. +pub(crate) fn record_served_model_id(state_path: &Path, served_model_id: &str) -> Result<()> { + merge_state_field(state_path, "served_model_id", json!(served_model_id)) +} + +/// Sets a non-terminal status (`stopping`) without stamping `stopped_at`. +pub(crate) fn write_status(state_path: &Path, status: &str) -> Result<()> { + merge_state_field(state_path, "status", json!(status)) +} + +pub(crate) fn read_status(state_path: &Path) -> Option { + value_string(&read_service_state(state_path).ok()?, "status") +} + +fn merge_state_field(state_path: &Path, key: &str, value: Value) -> Result<()> { + let mut state = read_service_state(state_path).unwrap_or_else(|_| json!({})); + if let Some(object) = state.as_object_mut() { + object.insert(key.to_owned(), value); + } + write_state(state_path, &state) +} + pub(crate) fn write_terminal_state(state_path: &Path, status: &str) -> Result<()> { let mut state = read_service_state(state_path).unwrap_or_else(|_| json!({})); if let Some(object) = state.as_object_mut() { @@ -180,6 +210,29 @@ pub(crate) fn query_loaded_model_endpoint( ) } +/// The first model the endpoint lists, for images that don't label the name +/// they serve under. +pub(crate) fn discover_served_model(endpoint_url: &str) -> Option { + let endpoint_api_key = rocm_engine_protocol::resolve_endpoint_api_key(); + let body = rocm_core::http_get_text_with_auth( + endpoint_url, + "/v1/models", + endpoint_api_key.as_deref(), + Duration::from_millis(HEALTHCHECK_TIMEOUT_MS), + ) + .ok()?; + first_listed_model_id(&body) +} + +fn first_listed_model_id(body: &str) -> Option { + serde_json::from_str::(body.trim()) + .ok()? + .get("data")? + .as_array()? + .iter() + .find_map(|item| value_string(item, "id").filter(|id| !id.trim().is_empty())) +} + pub(crate) fn query_inference_probe_endpoint(endpoint_url: &str, model_ref: &str) -> Result { if model_ref.trim().is_empty() { return Ok(false); @@ -275,12 +328,8 @@ mod tests { #[test] fn write_running_state_never_persists_extra_env_secrets_to_disk() { - let dir = std::env::temp_dir().join(format!( - "rocm-aim-state-secret-test-{}-{}", - std::process::id(), - current_unix_millis() - )); - let service_id = format!("svc-secret-{}", current_unix_millis()); + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("svc-secret.json"); let hint = EngineRecipeHint { contract_version: "0.1.0".to_owned(), engine: "aim".to_owned(), @@ -290,13 +339,11 @@ mod tests { )]), ..EngineRecipeHint::default() }; - let path = state_path(&service_id); - // `write_running_state` always writes under `state_path(service_id)`, which - // lives outside `dir`; only `dir` is cleaned up below, matching the other - // tests in this module that leave the real state dir alone. write_running_state( - &service_id, + &path, + "svc-secret", "some/model", + Some("Org/Served"), "127.0.0.1", 8000, DevicePolicy::GpuPreferred, @@ -306,8 +353,32 @@ mod tests { let on_disk = fs::read_to_string(&path).unwrap(); assert!(!on_disk.contains("hf_super_secret")); assert!(on_disk.contains("HF_TOKEN")); - let _ = fs::remove_file(&path); - let _ = fs::remove_dir_all(&dir); + } + + #[test] + fn status_helpers_round_trip_and_keep_other_fields() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("svc.json"); + assert_eq!(read_status(&path), None); + write_status(&path, "stopping").unwrap(); + assert_eq!(read_status(&path).as_deref(), Some("stopping")); + record_served_model_id(&path, "Org/Served").unwrap(); + let state = read_service_state(&path).unwrap(); + assert_eq!(value_string(&state, "status").as_deref(), Some("stopping")); + assert_eq!( + value_string(&state, "served_model_id").as_deref(), + Some("Org/Served") + ); + } + + #[test] + fn first_listed_model_id_reads_the_first_nonempty_id() { + assert_eq!( + first_listed_model_id(r#"{"data":[{"id":""},{"id":"a/b"},{"id":"c"}]}"#).as_deref(), + Some("a/b") + ); + assert_eq!(first_listed_model_id(r#"{"data":[]}"#), None); + assert_eq!(first_listed_model_id("not json"), None); } #[test] diff --git a/engines/lemonade/src/lib.rs b/engines/lemonade/src/lib.rs index 7f8fac293..7d8a91e5c 100644 --- a/engines/lemonade/src/lib.rs +++ b/engines/lemonade/src/lib.rs @@ -412,6 +412,7 @@ fn resolve_model_response(request: ResolveModelRequest) -> Result Result Date: Tue, 27 Oct 2026 13:32:08 +0200 Subject: [PATCH 19/22] fix(aim): address round-2 review findings - docker: match "no such ..." case-insensitively so a missing container is recognised under Docker's lowercase `error: no such object` wording - providers: local chat selection also matches the service's served model id - services stop: share the container-engine stop-confirmation rule between rocm and rocmd (rocm_core::engine_stop_decides_confirmation) so a failed container stop is not recorded as stopped - aim: remove the container when `docker run` fails; GpuPreferred no longer falls back to CPU - protocol: omit an empty recipe `extra_env` when serialising - tests: VLLM_API_KEY forwarding, stopping-before-kill, run-failure cleanup, container gone after stop in serve-27; model-catalog-02 is Linux only - docs: README/manual-testing cover bare --engine-env KEY, --yes, and ROCM_CLI_AIM_READY_TIMEOUT_SECS Signed-off-by: Juho Vainio --- README.md | 12 ++- apps/rocm/src/main.rs | 43 ++++++++- apps/rocm/src/providers.rs | 31 +++++++ apps/rocmd/src/service.rs | 10 +-- crates/rocm-core/src/lib.rs | 2 +- crates/rocm-core/src/service_record.rs | 12 +++ crates/rocm-engine-protocol/src/lib.rs | 4 +- docs/manual-testing.md | 10 ++- engines/aim/src/docker.rs | 38 ++++++-- engines/aim/src/process.rs | 89 ++++++++++++++++--- .../features/model_catalog.feature | 4 +- tests/e2e-cucumber/tests/e2e/serving_steps.rs | 11 +++ 12 files changed, 226 insertions(+), 40 deletions(-) diff --git a/README.md b/README.md index 610c54525..53ece1406 100644 --- a/README.md +++ b/README.md @@ -771,9 +771,15 @@ rocm serve [--engine lemonade|vllm|aim] `--engine aim` runs the model in a Docker container instead of a host process; it requires a reachable Docker daemon and is Linux/WSL2 only. -`--engine-env` passes an extra `KEY=VALUE` environment variable through to -the engine's own process/container (repeatable); currently only `aim` -consumes these (e.g. `AIM_PROFILE_ID`). +`--engine-env` passes an extra environment variable through to the engine's +own process/container (repeatable); currently only `aim` consumes these (e.g. +`AIM_PROFILE_ID`). `KEY=VALUE` values are visible in the detached child's +argv and in `rocm services list --json`, so do not put secrets there; pass a +bare `KEY` (e.g. `--engine-env HF_TOKEN`) to forward the value from your +current environment without storing it. A service restarted by the daemon +after a crash inherits the daemon's environment, not yours, so a forwarded +secret must also be set there. `ROCM_CLI_AIM_READY_TIMEOUT_SECS` raises the +5-minute readiness wait for a cold-start model download. `--gpu-memory-utilization` (a fraction above 0.0 and up to 1.0 of each GPU's total VRAM) and `--tool-call-parser` (for example `hermes`, `llama3_json`, or diff --git a/apps/rocm/src/main.rs b/apps/rocm/src/main.rs index 36f8a850d..d95f404bb 100644 --- a/apps/rocm/src/main.rs +++ b/apps/rocm/src/main.rs @@ -7095,7 +7095,12 @@ fn stop_internal_managed_service_confirmed_with( // unverifiable PID), leave the status so the standard liveness refresh // reconciles it to "stopped" once the process actually dies — rather than // asserting a stop that did not happen. - let all_stopped = terminated_all && stopping.names_same_processes(&record); + // A container engine is stopped by its engine `Stop`, not by the PID of the + // supervisor that follows it, so a failed `Stop` leaves the stop unconfirmed. + let engine_stop_confirms = + !rocm_core::engine_stop_decides_confirmation(&stopping.engine) || engine_stop.is_ok(); + let all_stopped = + terminated_all && stopping.names_same_processes(&record) && engine_stop_confirms; if all_stopped { record.status = "stopped".to_owned(); record.stop_requested_unix_ms = None; @@ -15800,6 +15805,32 @@ pub(crate) mod tests { assert_eq!(key.as_deref(), Some("key-123"), "the key is kept"); } + #[test] + fn stop_of_a_container_engine_is_unconfirmed_when_the_engine_stop_fails() { + // For `aim` the recorded PID is only the supervisor; killing it says + // nothing about the container. With no aim engine to answer `Stop`, the + // PID verdict alone must not mark the service stopped or drop its key. + let (root, paths) = test_paths("stop-aim-engine-stop-fails"); + paths.ensure().unwrap(); + let service_id = "svc-stop-aim-engine-stop-fails"; + let mut record = stop_test_record(&paths, service_id); + record.engine = "aim".to_owned(); + record.write().unwrap(); + endpoint_keys::store_endpoint_api_key(&paths, service_id, "key-123").unwrap(); + + let report = + stop_internal_managed_service_with(&paths, service_id, |_| (Vec::new(), true)).unwrap(); + + let on_disk = raw_manifest(&paths, service_id); + let key = endpoint_keys::endpoint_api_key(&paths, service_id); + let _ = fs::remove_dir_all(root); + assert!(report["engine_stop"]["error"].is_string(), "{report}"); + // Unconfirmed: the marker is re-stamped and the key kept for the liveness + // refresh to settle, rather than cleared as for a confirmed stop. + assert!(on_disk.stop_requested_unix_ms.is_some()); + assert_eq!(key.as_deref(), Some("key-123"), "the key is kept"); + } + #[test] fn stop_final_write_keeps_what_others_wrote_while_it_ran() { // The stop's last write applies its own fields to the record as it is @@ -19357,9 +19388,6 @@ pub(crate) mod tests { ); } - /// When the platform gate allows an engine, there is nothing to report -- - /// asserted directly against the pure helper, not by relying on whichever - /// OS this test happens to run on. /// AIM's Docker GPU passthrough only exists on Linux/WSL2, so native /// Windows must be ruled out (the Docker-reachability gate alone would pass /// on Docker Desktop and launch on CPU). @@ -19373,6 +19401,9 @@ pub(crate) mod tests { ); } + /// When the platform gate allows an engine, there is nothing to report -- + /// asserted directly against the pure helper, not by relying on whichever + /// OS this test happens to run on. #[test] fn unsupported_here_for_is_none_when_allowed() { assert_eq!(unsupported_here_for("vllm", PlatformRuleOut::Allowed), None); @@ -19384,6 +19415,10 @@ pub(crate) mod tests { /// `aim` is ruled out for a different reason than `vllm` on Windows /// (Docker, not the OS), so the message has to say Docker, not WSL. + /// + /// Pins the wording only; that the gate actually fires on an unreachable + /// daemon is proven by + /// `engine_ruled_out_by_platform_gates_aim_on_docker_reachability`. #[test] fn unsupported_here_for_names_docker_for_aim() { let message = unsupported_here_for("aim", PlatformRuleOut::AimDockerUnavailable) diff --git a/apps/rocm/src/providers.rs b/apps/rocm/src/providers.rs index 91834b539..af5273d7a 100644 --- a/apps/rocm/src/providers.rs +++ b/apps/rocm/src/providers.rs @@ -476,6 +476,10 @@ fn select_local_chat_service( services.retain(|service| { service.canonical_model_id.eq_ignore_ascii_case(model) || service.model_ref.eq_ignore_ascii_case(model) + || service + .served_model_id + .as_deref() + .is_some_and(|served| served.eq_ignore_ascii_case(model)) }); return services .into_iter() @@ -2253,6 +2257,33 @@ mod tests { Ok(()) } + #[test] + fn local_provider_selects_a_service_by_its_served_model_id() -> Result<()> { + let (root, paths) = temp_app_paths("local-provider-served-id"); + paths.ensure()?; + let mut aim = ManagedServiceRecord::new( + &paths, + "svc-aim", + "aim", + "aim-qwen3.5-9b", + "docker.io/org/aim-image:1", + "127.0.0.1", + 8000, + "managed", + 123, + None, + None, + None, + ); + aim.status = "ready".to_owned(); + aim.served_model_id = Some("Org/Served".to_owned()); + aim.write()?; + let selected = select_local_chat_service(&paths, Some("org/served"))?; + fs::remove_dir_all(root).ok(); + assert_eq!(selected.service_id, "svc-aim"); + Ok(()) + } + #[test] fn local_provider_default_chat_requires_builtin_qwen_assistant() -> Result<()> { let (root, paths) = temp_app_paths("local-provider-default-qwen"); diff --git a/apps/rocmd/src/service.rs b/apps/rocmd/src/service.rs index 934a75771..79793d1f8 100644 --- a/apps/rocmd/src/service.rs +++ b/apps/rocmd/src/service.rs @@ -55,14 +55,6 @@ fn request_engine_stop(paths: &AppPaths, record: &ManagedServiceRecord) -> Resul ) } -/// Engines whose service is a container rather than a recorded process: the -/// engine `Stop` is the only thing that stops one, so its verdict, not the PID -/// signalling, decides whether the stop is confirmed. For a PID-supervised -/// engine the PIDs are authoritative and the engine `Stop` is reported only. -fn engine_stop_decides_confirmation(engine: &str) -> bool { - engine == "aim" -} - /// [`stop_managed_service_with`], with the engine `Stop` request supplied by the /// caller as well, so a test can see what was asked of the engine and make it /// fail. @@ -126,7 +118,7 @@ pub(crate) fn stop_managed_service_via( // stop was running — those are processes it never touched, and claiming the // service stopped, clearing their PIDs or dropping the key they serve with // would each be false. - let engine_stop_confirms = !engine_stop_decides_confirmation(&stopping.engine) + let engine_stop_confirms = !rocm_core::engine_stop_decides_confirmation(&stopping.engine) || engine_stop.as_ref().is_some_and(Result::is_ok); let all_stopped = stops.iter().all(|stop| stop.stopped) && stopping.names_same_processes(&record) diff --git a/crates/rocm-core/src/lib.rs b/crates/rocm-core/src/lib.rs index ec4163dc2..0f5df9c21 100644 --- a/crates/rocm-core/src/lib.rs +++ b/crates/rocm-core/src/lib.rs @@ -174,7 +174,7 @@ pub use runtime::{ runtime_python_env_bin_dir, runtime_python_executable_in_env, runtime_python_executable_name, runtime_rocm_library_filename, shell_command_for_host, user_runtime_dir, }; -pub use service_record::ManagedServiceRecord; +pub use service_record::{ManagedServiceRecord, engine_stop_decides_confirmation}; pub use uv::{ DEFAULT_UV_TIMEOUT_SECS, DependencyViolation, UV_CACHE_DIR_ENV, UV_CACHE_DIR_OVERRIDE_ENV, UvCacheSource, ViolationSubject, check_dependencies, ensure_uv_binary, split_local_version, diff --git a/crates/rocm-core/src/service_record.rs b/crates/rocm-core/src/service_record.rs index 65cf009f6..37bb07d94 100644 --- a/crates/rocm-core/src/service_record.rs +++ b/crates/rocm-core/src/service_record.rs @@ -112,6 +112,18 @@ pub struct ManagedServiceRecord { pub created_at_unix_ms: u128, } +/// Whether the engine `Stop` verdict decides that a stop is confirmed. +/// +/// True for engines whose service is a container rather than a recorded +/// process: the engine `Stop` is the only thing that stops one, so its verdict, +/// not the PID signalling, decides. For a PID-supervised +/// engine the PIDs are authoritative and the engine `Stop` is reported only. +/// Shared by `rocm` and `rocmd` so the two stop paths cannot disagree. +#[must_use] +pub fn engine_stop_decides_confirmation(engine: &str) -> bool { + engine == "aim" +} + impl ManagedServiceRecord { #[allow(clippy::too_many_arguments)] pub fn new( diff --git a/crates/rocm-engine-protocol/src/lib.rs b/crates/rocm-engine-protocol/src/lib.rs index 924a49be4..6ca773ccd 100644 --- a/crates/rocm-engine-protocol/src/lib.rs +++ b/crates/rocm-engine-protocol/src/lib.rs @@ -532,7 +532,9 @@ pub struct EngineRecipeHint { /// Engine-specific environment variables to pass through verbatim to the /// engine's own process/container (e.g. AIM's `AIM_PROFILE_ID` override). /// Generic by design — only `engines/aim` currently reads this. - #[serde(default)] + /// Omitted when empty so a recipe hint serialises as it did before the + /// field existed (duplicate-serve reuse compares the JSON as a string). + #[serde(default, skip_serializing_if = "BTreeMap::is_empty")] pub extra_env: BTreeMap, #[serde(default, skip_serializing_if = "Option::is_none")] pub preferred_endpoint: Option, diff --git a/docs/manual-testing.md b/docs/manual-testing.md index 1db4681e0..9e007b188 100644 --- a/docs/manual-testing.md +++ b/docs/manual-testing.md @@ -774,11 +774,17 @@ succeed). ```bash rocm serve aim-qwen3.5-9b --engine aim --managed rocm services list --all # note the service id, e.g. aim-- -rocm services stop +rocm services stop --yes docker ps -a --filter name=rocm-aim- -rocm services restart +rocm services restart --yes ``` +A cold start downloads the model weights inside the container, which can exceed +the engine's default 5-minute readiness wait. Raise it with +`ROCM_CLI_AIM_READY_TIMEOUT_SECS=` (a positive integer) in the +environment of `rocm serve`. Stopping or failing a start removes the container +and its downloaded weights, so a retry downloads them again. + `aim-qwen3.5-9b` is the catalog entry (`rocm model --verbose`). It supplies the pinned image and `AIM_PROFILE_ID`; to serve a raw image reference instead, pass the profile yourself: diff --git a/engines/aim/src/docker.rs b/engines/aim/src/docker.rs index 211f4ef53..12e01e427 100644 --- a/engines/aim/src/docker.rs +++ b/engines/aim/src/docker.rs @@ -210,8 +210,15 @@ pub(crate) fn stop(name: &str, force: bool, grace: Duration) -> Result { Ok(true) } +/// Docker words a missing object as `No such container` or, in newer +/// releases, `error: no such object`; match both regardless of case. +fn is_no_such(detail: &str) -> bool { + detail.to_ascii_lowercase().contains("no such") +} + fn is_nothing_to_stop(detail: &str) -> bool { - detail.contains("No such container") || detail.contains("is not running") + let detail = detail.to_ascii_lowercase(); + detail.contains("no such container") || detail.contains("is not running") } /// Force-removes the container (stopping it first if needed) so its name and @@ -224,7 +231,7 @@ pub(crate) fn remove(name: &str) -> Result<()> { .with_context(|| format!("failed to run `docker rm` for container `{name}`"))?; if !output.status.success() { let detail = rocm_core::command_failure_detail(&output); - if detail.contains("No such container") { + if is_no_such(&detail) { return Ok(()); } bail!("failed to remove container `{name}`: {detail}"); @@ -254,7 +261,7 @@ pub(crate) fn container_state(name: &str) -> Result { .with_context(|| format!("failed to run `docker inspect` for container `{name}`"))?; if !output.status.success() { let detail = rocm_core::command_failure_detail(&output); - if detail.contains("No such") { + if is_no_such(&detail) { return Ok(ContainerState::Missing); } bail!("`docker inspect {name}` failed: {detail}"); @@ -359,8 +366,8 @@ pub(crate) fn daemon_reachable() -> bool { #[cfg(test)] mod tests { use super::{ - ContainerState, RunSpec, gpu_device_args, is_nothing_to_stop, parse_container_state, - parse_getent_gid, run_args, split_env, + ContainerState, RunSpec, gpu_device_args, is_no_such, is_nothing_to_stop, + parse_container_state, parse_getent_gid, run_args, split_env, }; use std::collections::BTreeMap; use std::path::Path; @@ -496,6 +503,17 @@ mod tests { assert!(parse_container_state("false oops").is_err()); } + #[test] + fn is_no_such_matches_every_docker_wording() { + assert!(is_no_such("Error: No such object: x")); + assert!(is_no_such("error: no such object: x")); + assert!(is_no_such( + "Error response from daemon: No such container: x" + )); + assert!(!is_no_such("Cannot connect to the Docker daemon")); + assert!(is_nothing_to_stop("error: no such container: x")); + } + #[test] fn is_nothing_to_stop_covers_missing_and_already_exited_containers() { assert!(is_nothing_to_stop("Error: No such container: x")); @@ -540,6 +558,7 @@ case "$1" in exit "${FAKE_DOCKER_RUN_EXIT:-0}" ;; kill|stop) + find "$FAKE_DOCKER_DIR/data" -name '*.json' -exec cat {} + > "$FAKE_DOCKER_DIR/state_at_stop" 2>/dev/null if [ -n "$FAKE_DOCKER_STOP_NO_SUCH_CONTAINER" ]; then echo "Error response from daemon: No such container: fake" >&2 exit 1 @@ -565,14 +584,14 @@ case "$1" in entry=$(echo "$FAKE_DOCKER_INSPECT_SEQUENCE" | tr ',' '\n' | sed -n "${n}p") [ -z "$entry" ] && entry=$(echo "$FAKE_DOCKER_INSPECT_SEQUENCE" | tr ',' '\n' | tail -n 1) if [ "$entry" = "missing" ]; then - echo "Error: No such object: fake" >&2 + echo "error: no such object: fake" >&2 exit 1 fi echo "$entry" exit 0 fi if [ -n "$FAKE_DOCKER_INSPECT_MISSING" ]; then - echo "Error: No such object: fake" >&2 + echo "error: no such object: fake" >&2 exit 1 fi if [ -n "$FAKE_DOCKER_INSPECT_ERROR" ]; then @@ -622,6 +641,11 @@ exit 0 .collect() } + /// The engine state files as they were when `docker kill`/`stop` ran. + pub(crate) fn state_at_stop(&self) -> Option { + std::fs::read_to_string(self.dir.join("state_at_stop")).ok() + } + pub(crate) fn env_copy(&self) -> Option { std::fs::read_to_string(self.dir.join("env.copy")).ok() } diff --git a/engines/aim/src/process.rs b/engines/aim/src/process.rs index 7a1ffa32d..bfb6eef68 100644 --- a/engines/aim/src/process.rs +++ b/engines/aim/src/process.rs @@ -311,12 +311,6 @@ fn run_container( gpus.available, image_wants_cpu, )?; - if !use_gpu && matches!(device_policy, DevicePolicy::GpuPreferred) { - eprintln!( - "Warning: AIM service `{service_id}` will run on CPU (device_policy=gpu_preferred \ - and no usable GPU for `{image_ref}`); pass `--device gpu` to require a GPU." - ); - } let render_gid = if use_gpu { crate::docker::render_group_gid() } else { @@ -361,7 +355,10 @@ fn resolve_use_gpu( ) -> Result { Ok(match device_policy { DevicePolicy::CpuOnly => false, - DevicePolicy::GpuRequired => { + // `GpuPreferred` is normalised to `GpuRequired`, as the sibling engines + // do: no silent CPU fallback (AGENTS.md section 6). CPU is reachable + // only through an explicit `CpuOnly`. + DevicePolicy::GpuRequired | DevicePolicy::GpuPreferred => { if !gpu_usable || image_wants_cpu { bail!( "AIM service `{service_id}` has device_policy=gpu_required but no GPU is \ @@ -371,7 +368,6 @@ fn resolve_use_gpu( } true } - DevicePolicy::GpuPreferred => gpu_usable && !image_wants_cpu, }) } @@ -385,7 +381,13 @@ pub(crate) fn serve_http(request: ServeHttpRequest) -> Result<()> { &request.device_policy, &request.gpu_indices, request.engine_recipe.as_ref().map(|hint| &hint.extra_env), - )?; + ) + .inspect_err(|_| { + // `docker run -d` can create the container and then fail to start it + // (port already allocated); service ids are unique, so nothing would + // ever reclaim it. + let _ = crate::docker::remove(&name); + })?; crate::state::write_running_state( &request.state_path, @@ -695,11 +697,11 @@ mod tests { } #[test] - fn resolve_use_gpu_gpu_preferred_falls_back_to_cpu_without_failing() { + fn resolve_use_gpu_gpu_preferred_never_falls_back_to_cpu() { assert!( - !resolve_use_gpu("svc", "model", &DevicePolicy::GpuPreferred, false, false).unwrap() + resolve_use_gpu("svc", "model", &DevicePolicy::GpuPreferred, false, false).is_err() ); - assert!(!resolve_use_gpu("svc", "model", &DevicePolicy::GpuPreferred, true, true).unwrap()); + assert!(resolve_use_gpu("svc", "model", &DevicePolicy::GpuPreferred, true, true).is_err()); assert!(resolve_use_gpu("svc", "model", &DevicePolicy::GpuPreferred, true, false).unwrap()); } @@ -1298,6 +1300,65 @@ mod subprocess_tests { ); } + #[test] + fn run_container_forwards_the_endpoint_api_key_as_vllm_api_key() { + if in_subprocess() { + run_container( + "svc-key", + IMAGE, + "127.0.0.1", + 18088, + &DevicePolicy::CpuOnly, + &[], + None, + ) + .unwrap(); + return; + } + let key_file = std::env::temp_dir().join(format!("aim-key-{}", std::process::id())); + std::fs::write(&key_file, "s3cret\n").unwrap(); + let fixture = run_with_fake_docker( + "process::subprocess_tests::run_container_forwards_the_endpoint_api_key_as_vllm_api_key", + &[ + ("FAKE_DOCKER_INSPECT_SEQUENCE", "missing"), + ("FAKE_DOCKER_IMAGE_PRESENT_EXIT", "0"), + ( + rocm_engine_protocol::ENDPOINT_API_KEY_FILE_ENV, + key_file.to_str().unwrap(), + ), + ], + ); + std::fs::remove_file(&key_file).ok(); + assert!(!fixture.argv().join("\n").contains("s3cret")); + assert_eq!(fixture.env_copy().as_deref(), Some("VLLM_API_KEY=s3cret\n")); + } + + #[test] + fn serve_http_removes_the_container_when_docker_run_fails() { + if in_subprocess() { + let state_path = data_dir().join("svc.json"); + serve_http(serve_request("svc-run-leak", 18089, state_path)) + .expect_err("docker run failure should surface"); + return; + } + let fixture = run_with_fake_docker( + "process::subprocess_tests::serve_http_removes_the_container_when_docker_run_fails", + &[ + ("FAKE_DOCKER_INSPECT_SEQUENCE", "missing"), + ("FAKE_DOCKER_IMAGE_PRESENT_EXIT", "0"), + ("FAKE_DOCKER_RUN_EXIT", "1"), + ("FAKE_DOCKER_RUN_STDERR", "port is already allocated"), + ], + ); + assert!( + fixture + .argv() + .contains(&"rm -f rocm-aim-svc-run-leak".to_owned()), + "{:?}", + fixture.argv() + ); + } + #[test] fn wait_for_container_ready_bails_when_container_exits_early() { if in_subprocess() { @@ -1394,6 +1455,10 @@ mod subprocess_tests { "rm -f rocm-aim-svc-stop-succeeds" ] ); + // The deliberate-stop intent must be on disk before the kill, or the + // supervisor reads the container's exit 137 as a crash. + let at_stop = fixture.state_at_stop().unwrap_or_default(); + assert!(at_stop.contains("stopping"), "{at_stop}"); } #[test] diff --git a/tests/e2e-cucumber/features/model_catalog.feature b/tests/e2e-cucumber/features/model_catalog.feature index e738029d5..a7acdee31 100644 --- a/tests/e2e-cucumber/features/model_catalog.feature +++ b/tests/e2e-cucumber/features/model_catalog.feature @@ -18,7 +18,9 @@ Feature: Model catalog listing # unavailable branch is exercised deterministically on every host, including # ones where Docker really is reachable (an earlier variant compared against # the host's real Docker state, which asserted a different outcome per host). - @id:model-catalog-flags-aim-unavailable-with-no-docker-daemon + # Linux only: native Windows rules AIM out for the OS before it asks Docker, + # so it reports `unsupported_native_windows`, not the Docker reason asserted here. + @id:model-catalog-flags-aim-unavailable-with-no-docker-daemon @requires-os:linux Scenario: model-catalog-02 - AIM image is flagged unavailable with no Docker daemon Given no Docker daemon is reachable When the user asks for the model catalog in verbose detail diff --git a/tests/e2e-cucumber/tests/e2e/serving_steps.rs b/tests/e2e-cucumber/tests/e2e/serving_steps.rs index 8cb551853..e8dc2b028 100644 --- a/tests/e2e-cucumber/tests/e2e/serving_steps.rs +++ b/tests/e2e-cucumber/tests/e2e/serving_steps.rs @@ -688,6 +688,17 @@ async fn stop_aim_service(world: &mut E2eWorld) { status.contains(": stopped"), "expected the AIM service to stop cleanly, got: {status}" ); + // A stop that leaves the container running would let the next serve reuse + // it and still pass inference, so check the container is actually gone. + let remaining = std::process::Command::new("docker") + .args(["ps", "-q", "--filter", "name=rocm-aim-"]) + .output() + .expect("docker ps should run on an AIM host"); + let remaining = String::from_utf8_lossy(&remaining.stdout); + assert!( + remaining.trim().is_empty(), + "an AIM container is still running after the stop: {remaining}" + ); } /// Serves the same AIM canary again, proving the engine can be launched a From 375a5eb415aee28919afab23148d9bf2b05b8d05 Mon Sep 17 00:00:00 2001 From: Juho Vainio Date: Tue, 27 Oct 2026 13:43:57 +0200 Subject: [PATCH 20/22] test(providers): serve a models endpoint for the served-id selection test Signed-off-by: Juho Vainio --- apps/rocm/src/providers.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/apps/rocm/src/providers.rs b/apps/rocm/src/providers.rs index af5273d7a..5d9acce78 100644 --- a/apps/rocm/src/providers.rs +++ b/apps/rocm/src/providers.rs @@ -2261,6 +2261,7 @@ mod tests { fn local_provider_selects_a_service_by_its_served_model_id() -> Result<()> { let (root, paths) = temp_app_paths("local-provider-served-id"); paths.ensure()?; + let server = spawn_models_only_server("Org/Served")?; let mut aim = ManagedServiceRecord::new( &paths, "svc-aim", @@ -2268,7 +2269,7 @@ mod tests { "aim-qwen3.5-9b", "docker.io/org/aim-image:1", "127.0.0.1", - 8000, + server.port(), "managed", 123, None, From eef0064c11401919a1065c99f28e057b8d7009a4 Mon Sep 17 00:00:00 2001 From: Juho Vainio Date: Tue, 27 Oct 2026 17:28:25 +0200 Subject: [PATCH 21/22] fix(aim): address round-3 review findings Re-read the served model name during the readiness wait, publish localhost/IPv6 hosts in a form docker accepts, send the served name in local chat requests and status, make the container-stop test independent of a reachable docker daemon, and list AIM among the serving engines. Signed-off-by: Juho Vainio --- CONTRIBUTING.md | 2 +- apps/rocm/src/chat.rs | 2 +- apps/rocm/src/main.rs | 13 +++++++++++-- apps/rocm/src/providers.rs | 22 +++++++-------------- apps/rocm/src/service_launch.rs | 6 +++++- docs/llm-tool-use.md | 4 ++-- engines/aim/src/docker.rs | 34 ++++++++++++++++++++++++++++++++- 7 files changed, 60 insertions(+), 23 deletions(-) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 011af6ad1..3e6442a63 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -79,7 +79,7 @@ cargo install lychee@0.24.2 --locked # pinned to match the CI docs-links job ( | `crates/rocm-core` | Core library | | `crates/rocm-dash-*` | Dashboard terminal user interface (TUI) libraries | | `crates/rocm-engine-protocol` | Engine IPC protocol | -| `engines/` | Inference engine adapters (lemonade, vllm) | +| `engines/` | Inference engine adapters (lemonade, vllm, aim) | ### Module organization diff --git a/apps/rocm/src/chat.rs b/apps/rocm/src/chat.rs index a2a7d8ea3..206e2286e 100644 --- a/apps/rocm/src/chat.rs +++ b/apps/rocm/src/chat.rs @@ -1798,7 +1798,7 @@ fn find_ascii_case_insensitive(haystack: &str, needle: &str) -> Option { /// to, and a WSL user was told "on native Windows, vLLM is skipped" while vLLM /// was in fact their supported path. Anything host-dependent now comes from /// [`chat_host_facts::HostFacts`], which knows which machine it is describing. -const ROCM_CHAT_TOOL_SYSTEM_PROMPT: &str = "You are ROCm CLI's local assistant. Speak in simple English for non-technical users. Use the provided ROCm tools when you need to inspect this machine, preview setup, read service logs, check updates, inspect automations, install or start ROCm-managed apps, or request ROCm/TheRock, config, engine, app, and local model server changes. For simple greetings or thanks like hello, hi, hey, ok, or thank you, reply normally; do not inspect ROCm, do not call tools, and do not launch or propose a model server. Tool-use rules: inspect first with read-only tools; call rocm_command only with argv-style args and no shell text; use natural_language_plan for ROCm requests that do not fit another read-only tool; ask for a mutating tool call only after explaining why it is needed; summarize tool results after they are returned. Read-only tools may run immediately. Tools that install, launch, stop, delete, or change state require user approval; request rocm_command and explain why. For 'is X running?', 'what is running?', status, or port questions, inspect before answering and do not start, stop, install, or serve anything. For ComfyUI or port 8188 use [\"comfyui\",\"status\"] or port_status. For vLLM, Lemonade, qwen, or local model servers use [\"services\",\"list\",\"--all\"] for running state and [\"engines\",\"list\"] for installed/available engine state. Treat ready/running as running, starting/recovering as starting, failed/stopped as not running, and no matching record as unknown or not managed by ROCm CLI. Interpret Examine carefully: active_runtime_status=ready means ROCm CLI has an active managed TheRock/ROCm runtime; legacy_rocm_status=not_detected only means no global system ROCm install was found. If active_runtime_status=ready, tell the user ROCm/TheRock is installed and active for ROCm CLI. For 'is TheRock installed', 'is ROCm installed', or 'which GPU is on this machine', use examine or gpu_snapshot before answering. For 'how do I setup TheRock' or install/setup requests, guide the user to choose an install folder first; do not answer with only a status check. For 'which LLMs can this machine support', use rocm_command args [\"model\"] or natural_language_plan before answering. For TheRock installs, always let the user choose the install folder. If the user names a folder or prefix, preserve that exact folder with [\"--prefix\",\"PATH\"]; you may call path_exists first to check whether that user-provided folder or its parent exists. If the user asks you to install TheRock/ROCm but has not named a folder, ask for the folder or let the guided setup folder picker collect it; do not invent a hidden default folder and do not request an install command without --prefix. Use rocm_command args [\"install\",\"sdk\",\"--channel\",\"release\",\"--format\",\"wheel\",\"--prefix\",\"PATH\"] only when the user asks you to install it and a folder is known; for a requested build date add [\"--build-date\",\"YYYY-MM-DD\"] and for a requested exact version add [\"--version\",\"VERSION\"]. For config changes, inspect with [\"config\",\"show\"] first when useful, then request config subcommands such as [\"config\",\"set-default-engine\",\"lemonade\"], [\"config\",\"set-default-runtime\",\"RUNTIME_KEY\"], or [\"config\",\"set-telemetry\",\"local\"] only after explaining why. For ComfyUI, use rocm_command with args like [\"comfyui\",\"status\"], [\"comfyui\",\"logs\"], [\"comfyui\",\"install\"], [\"comfyui\",\"start\"], or [\"comfyui\",\"stop\"]. First-time setup is the same thing as bootstrap in ROCm CLI; it is a deterministic ROCm setup flow, not a separate model chat. The built-in local assistant is fixed to qwen, which maps to Qwen3-4B-Instruct-2507-GGUF served by Lemonade with gpu_required. vLLM and Lemonade are the general serving engines; inspect or manage them when the user asks about general model serving, but do not switch the built-in assistant away from Lemonade. Use qwen-smoke only for a quick server smoke test. For vLLM management, inspect engines first and use [\"engines\",\"install\",\"vllm\"] or [\"serve\",\"MODEL\",\"--engine\",\"vllm\",\"--device\",\"gpu_required\",\"--managed\"] only where the host supports it. Do not invent shell commands and do not request CPU fallback."; +const ROCM_CHAT_TOOL_SYSTEM_PROMPT: &str = "You are ROCm CLI's local assistant. Speak in simple English for non-technical users. Use the provided ROCm tools when you need to inspect this machine, preview setup, read service logs, check updates, inspect automations, install or start ROCm-managed apps, or request ROCm/TheRock, config, engine, app, and local model server changes. For simple greetings or thanks like hello, hi, hey, ok, or thank you, reply normally; do not inspect ROCm, do not call tools, and do not launch or propose a model server. Tool-use rules: inspect first with read-only tools; call rocm_command only with argv-style args and no shell text; use natural_language_plan for ROCm requests that do not fit another read-only tool; ask for a mutating tool call only after explaining why it is needed; summarize tool results after they are returned. Read-only tools may run immediately. Tools that install, launch, stop, delete, or change state require user approval; request rocm_command and explain why. For 'is X running?', 'what is running?', status, or port questions, inspect before answering and do not start, stop, install, or serve anything. For ComfyUI or port 8188 use [\"comfyui\",\"status\"] or port_status. For vLLM, Lemonade, AIM, qwen, or local model servers use [\"services\",\"list\",\"--all\"] for running state and [\"engines\",\"list\"] for installed/available engine state. Treat ready/running as running, starting/recovering as starting, failed/stopped as not running, and no matching record as unknown or not managed by ROCm CLI. Interpret Examine carefully: active_runtime_status=ready means ROCm CLI has an active managed TheRock/ROCm runtime; legacy_rocm_status=not_detected only means no global system ROCm install was found. If active_runtime_status=ready, tell the user ROCm/TheRock is installed and active for ROCm CLI. For 'is TheRock installed', 'is ROCm installed', or 'which GPU is on this machine', use examine or gpu_snapshot before answering. For 'how do I setup TheRock' or install/setup requests, guide the user to choose an install folder first; do not answer with only a status check. For 'which LLMs can this machine support', use rocm_command args [\"model\"] or natural_language_plan before answering. For TheRock installs, always let the user choose the install folder. If the user names a folder or prefix, preserve that exact folder with [\"--prefix\",\"PATH\"]; you may call path_exists first to check whether that user-provided folder or its parent exists. If the user asks you to install TheRock/ROCm but has not named a folder, ask for the folder or let the guided setup folder picker collect it; do not invent a hidden default folder and do not request an install command without --prefix. Use rocm_command args [\"install\",\"sdk\",\"--channel\",\"release\",\"--format\",\"wheel\",\"--prefix\",\"PATH\"] only when the user asks you to install it and a folder is known; for a requested build date add [\"--build-date\",\"YYYY-MM-DD\"] and for a requested exact version add [\"--version\",\"VERSION\"]. For config changes, inspect with [\"config\",\"show\"] first when useful, then request config subcommands such as [\"config\",\"set-default-engine\",\"lemonade\"], [\"config\",\"set-default-runtime\",\"RUNTIME_KEY\"], or [\"config\",\"set-telemetry\",\"local\"] only after explaining why. For ComfyUI, use rocm_command with args like [\"comfyui\",\"status\"], [\"comfyui\",\"logs\"], [\"comfyui\",\"install\"], [\"comfyui\",\"start\"], or [\"comfyui\",\"stop\"]. First-time setup is the same thing as bootstrap in ROCm CLI; it is a deterministic ROCm setup flow, not a separate model chat. The built-in local assistant is fixed to qwen, which maps to Qwen3-4B-Instruct-2507-GGUF served by Lemonade with gpu_required. vLLM, Lemonade, and AIM (Docker-based) are the general serving engines; inspect or manage them when the user asks about general model serving, but do not switch the built-in assistant away from Lemonade. Use qwen-smoke only for a quick server smoke test. For vLLM management, inspect engines first and use [\"engines\",\"install\",\"vllm\"] or [\"serve\",\"MODEL\",\"--engine\",\"vllm\",\"--device\",\"gpu_required\",\"--managed\"] only where the host supports it. Do not invent shell commands and do not request CPU fallback."; const ROCM_CHAT_TOOL_SKILL: &str = include_str!("../../../skills/rocm-cli-assistant/SKILL.md"); fn rocm_chat_tool_system_prompt() -> String { diff --git a/apps/rocm/src/main.rs b/apps/rocm/src/main.rs index 62953169f..a42bdd190 100644 --- a/apps/rocm/src/main.rs +++ b/apps/rocm/src/main.rs @@ -10291,7 +10291,7 @@ fn wait_for_service_http_ready( engine, host, port, - canonical_model_id, + &|| canonical_model_id.to_owned(), endpoint_api_key, timeout, &mut |_elapsed| {}, @@ -10317,7 +10317,7 @@ fn wait_for_service_http_ready_with_progress( engine: &str, host: &str, port: u16, - canonical_model_id: &str, + model_name: &dyn Fn() -> String, endpoint_api_key: Option<&str>, timeout: Duration, on_tick: &mut dyn FnMut(Duration), @@ -10326,6 +10326,10 @@ fn wait_for_service_http_ready_with_progress( let endpoint = format_http_base_url(host, port); let mut best = EndpointReadiness::Unreachable; while start.elapsed() < timeout { + // Re-evaluated each pass: a container engine reports the name its server + // answers to only after it starts, so a name fixed up front can be wrong. + let canonical_model_id = model_name(); + let canonical_model_id = canonical_model_id.as_str(); let listed = service_http_readiness_paths(engine).iter().any(|path| { http_get_local_service( host, @@ -15826,6 +15830,10 @@ pub(crate) mod tests { // PID verdict alone must not mark the service stopped or drop its key. let (root, paths) = test_paths("stop-aim-engine-stop-fails"); paths.ensure().unwrap(); + // The built-in aim engine answers `Stop` through `docker`; a dead daemon + // address makes that fail the same way on a host with or without docker. + let mut env = ScopedTestEnv::new(); + env.set("DOCKER_HOST", "unix:///nonexistent/rocm-test-docker.sock"); let service_id = "svc-stop-aim-engine-stop-fails"; let mut record = stop_test_record(&paths, service_id); record.engine = "aim".to_owned(); @@ -15834,6 +15842,7 @@ pub(crate) mod tests { let report = stop_internal_managed_service_with(&paths, service_id, |_| (Vec::new(), true)).unwrap(); + drop(env); let on_disk = raw_manifest(&paths, service_id); let key = endpoint_keys::endpoint_api_key(&paths, service_id); diff --git a/apps/rocm/src/providers.rs b/apps/rocm/src/providers.rs index 598e19fcc..3beb07c48 100644 --- a/apps/rocm/src/providers.rs +++ b/apps/rocm/src/providers.rs @@ -283,7 +283,7 @@ impl ProviderAdapter for LocalProvider<'_> { let builtin_ready = services.iter().any(local_service_is_builtin_assistant); let mut models = services .into_iter() - .map(|record| record.canonical_model_id) + .map(|record| record.served_model_name().to_owned()) .collect::>(); models.sort(); models.dedup(); @@ -303,13 +303,7 @@ impl ProviderAdapter for LocalProvider<'_> { fn chat(&self, request: &ChatRequest) -> Result { let service = select_local_chat_service(self.paths, request.model.as_deref())?; - let body = openai_local_chat_request_body( - request - .model - .as_deref() - .unwrap_or(&service.canonical_model_id), - request, - ); + let body = openai_local_chat_request_body(service.served_model_name(), request); let endpoint_api_key = crate::endpoint_keys::endpoint_api_key(self.paths, &service.service_id); let response = post_json_to_local_endpoint( @@ -335,13 +329,7 @@ impl ProviderAdapter for LocalProvider<'_> { ) -> Result { let service = select_local_chat_service(self.paths, request.model.as_deref())?; let model = service.canonical_model_id.clone(); - let body = openai_local_stream_chat_request_body( - request - .model - .as_deref() - .unwrap_or(&service.canonical_model_id), - request, - ); + let body = openai_local_stream_chat_request_body(service.served_model_name(), request); let endpoint_api_key = crate::endpoint_keys::endpoint_api_key(self.paths, &service.service_id); stream_json_from_local_endpoint( @@ -2506,8 +2494,12 @@ mod tests { aim.served_model_id = Some("Org/Served".to_owned()); aim.write()?; let selected = select_local_chat_service(&paths, Some("org/served"))?; + let status = provider_status(&paths, "local")?; fs::remove_dir_all(root).ok(); assert_eq!(selected.service_id, "svc-aim"); + // Chat sends `served_model_name()`, so `status` must list that name too. + assert_eq!(status.models, vec!["Org/Served".to_owned()]); + assert_eq!(selected.served_model_name(), "Org/Served"); Ok(()) } diff --git a/apps/rocm/src/service_launch.rs b/apps/rocm/src/service_launch.rs index 400e2168d..c54d55b6e 100644 --- a/apps/rocm/src/service_launch.rs +++ b/apps/rocm/src/service_launch.rs @@ -587,7 +587,11 @@ pub(crate) fn start_managed_service( engine, host, port, - resolve.served_model_name(), + &|| { + let mut latest = record.clone(); + let _ = latest.refresh_from_engine_state(); + latest.served_model_name().to_owned() + }, endpoint_api_key, Duration::from_secs(45), on_wait_tick, diff --git a/docs/llm-tool-use.md b/docs/llm-tool-use.md index e040c3fb8..2e931760f 100644 --- a/docs/llm-tool-use.md +++ b/docs/llm-tool-use.md @@ -23,8 +23,8 @@ rocm-cli local assistants use structured tools, not shell commands. - CPU fallback is not a supported path. GPU-required ROCm commands must fail loudly when the ROCm GPU path is not ready. - The built-in local assistant is fixed to `qwen` - (`Qwen3-4B-Instruct-2507-GGUF`) served by Lemonade. vLLM and Lemonade are the - general serving engines; the assistant may inspect or manage them for model + (`Qwen3-4B-Instruct-2507-GGUF`) served by Lemonade. vLLM, Lemonade and AIM (a + Docker-based engine) are the general serving engines; the assistant may inspect or manage them for model serving, but it should not switch its own built-in chat engine away from Lemonade. - The assistant prompt is grounded in the host. rocm-cli detects the operating diff --git a/engines/aim/src/docker.rs b/engines/aim/src/docker.rs index 12e01e427..14a1c264b 100644 --- a/engines/aim/src/docker.rs +++ b/engines/aim/src/docker.rs @@ -54,6 +54,17 @@ fn split_env(env: &BTreeMap) -> (Vec<(&str, &str)>, Vec<&str>) { (file_entries, forwarded) } +/// `docker run -p` takes an IP, not a hostname, and wants IPv6 in brackets. +fn publish_host(host: &str) -> String { + if host.eq_ignore_ascii_case("localhost") { + "127.0.0.1".to_owned() + } else if host.contains(':') && !host.starts_with('[') { + format!("[{host}]") + } else { + host.to_owned() + } +} + fn run_args( spec: &RunSpec<'_>, wsl: bool, @@ -66,7 +77,12 @@ fn run_args( "--name".to_owned(), spec.name.to_owned(), "-p".to_owned(), - format!("{}:{}:{}", spec.host, spec.port, spec.container_port), + format!( + "{}:{}:{}", + publish_host(spec.host), + spec.port, + spec.container_port + ), ]; if spec.use_gpu { if let Some(csv) = rocm_engine_protocol::gpu_indices_to_csv(spec.gpu_indices) { @@ -386,6 +402,22 @@ mod tests { } } + #[test] + fn run_args_publishes_localhost_and_ipv6_hosts_in_a_form_docker_accepts() { + let env = BTreeMap::new(); + for (host, published) in [ + ("localhost", "127.0.0.1:8080:8000"), + ("::1", "[::1]:8080:8000"), + ("0.0.0.0", "0.0.0.0:8080:8000"), + ] { + let mut s = spec(&[], &env); + s.host = host; + let args = run_args(&s, false, None, &[]); + assert_eq!(args[4], "-p"); + assert_eq!(args[5], published); + } + } + #[test] fn gpu_device_args_native_linux_uses_kfd_dri_and_render_gid() { assert_eq!( From 22cb68abe188eb88599ea4c7ce5e1db9edbaab22 Mon Sep 17 00:00:00 2001 From: Juho Vainio Date: Wed, 28 Oct 2026 15:43:21 +0200 Subject: [PATCH 22/22] fix(aim): try all resolved addresses on probe connect; report served name in chat stream summary Signed-off-by: Juho Vainio --- apps/rocm/src/providers.rs | 2 +- crates/rocm-core/src/net.rs | 36 ++++++++++++++++++++++++++++++------ 2 files changed, 31 insertions(+), 7 deletions(-) diff --git a/apps/rocm/src/providers.rs b/apps/rocm/src/providers.rs index 12b611de9..ce5d3a04a 100644 --- a/apps/rocm/src/providers.rs +++ b/apps/rocm/src/providers.rs @@ -328,7 +328,7 @@ impl ProviderAdapter for LocalProvider<'_> { on_event: &mut dyn FnMut(ProviderStreamEvent) -> Result<()>, ) -> Result { let service = select_local_chat_service(self.paths, request.model.as_deref())?; - let model = service.canonical_model_id.clone(); + let model = service.served_model_name().to_owned(); let body = openai_local_stream_chat_request_body(service.served_model_name(), request); let endpoint_api_key = crate::endpoint_keys::endpoint_api_key(self.paths, &service.service_id); diff --git a/crates/rocm-core/src/net.rs b/crates/rocm-core/src/net.rs index dbcb31378..b60a2a7d7 100644 --- a/crates/rocm-core/src/net.rs +++ b/crates/rocm-core/src/net.rs @@ -962,16 +962,32 @@ fn model_ref_tokens(value: &str) -> Vec { } pub fn connect_tcp_stream(host: &str, port: u16, timeout: Duration) -> Result { - let addr = (host, port) + let addrs: Vec<_> = (host, port) .to_socket_addrs() .with_context(|| format!("failed to resolve {host}:{port}"))? - .next() - .with_context(|| format!("no socket addresses resolved for {host}:{port}"))?; + .collect(); + if addrs.is_empty() { + bail!("no socket addresses resolved for {host}:{port}"); + } // Bound the connect as well as the reads: a probe against an engine that is // pinned solid must fail within the caller's timeout, not sit in the OS - // default SYN retry window. - let stream = TcpStream::connect_timeout(&addr, timeout) - .with_context(|| format!("failed to connect to {host}:{port}"))?; + // default SYN retry window. Try every resolved address so `localhost` + // resolving to `::1` first still reaches an IPv4-only listener. + let mut last_err = None; + let mut connected = None; + for addr in &addrs { + match TcpStream::connect_timeout(addr, timeout) { + Ok(s) => { + connected = Some(s); + break; + } + Err(e) => last_err = Some(e), + } + } + let Some(stream) = connected else { + let e = last_err.expect("addrs non-empty"); + return Err(anyhow::Error::new(e).context(format!("failed to connect to {host}:{port}"))); + }; stream.set_read_timeout(Some(timeout)).ok(); stream.set_write_timeout(Some(timeout)).ok(); Ok(stream) @@ -1071,6 +1087,14 @@ fn http_response_is_complete(response: &[u8]) -> bool { mod tests { use super::*; + #[test] + fn connect_tcp_stream_reaches_ipv4_listener_via_localhost() { + // `localhost` may resolve to ::1 first; every address must be tried. + let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + let port = listener.local_addr().unwrap().port(); + connect_tcp_stream("localhost", port, Duration::from_secs(2)).unwrap(); + } + use crate::test_support::{probe_test_record, workspace_test_artifact_dir}; use std::io::{Read, Write};