diff --git a/AGENTS.md b/AGENTS.md index 8825a0519..f1e7053e1 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. The CLI (`apps/rocm`) and the dash-TUI (`crates/rocm-dash-tui`) each have their own set; they are not interchangeable (see the naming-collision warning below). 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/Cargo.lock b/Cargo.lock index f4457305d..bcbd0022c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3672,6 +3672,7 @@ dependencies = [ "rocm-dash-daemon", "rocm-dash-tui", "rocm-deps", + "rocm-engine-aim", "rocm-engine-lemonade", "rocm-engine-protocol", "rocm-engine-vllm", @@ -3808,6 +3809,20 @@ dependencies = [ "toml", ] +[[package]] +name = "rocm-engine-aim" +version = "0.1.0" +dependencies = [ + "anyhow", + "clap", + "rocm-core", + "rocm-engine-protocol", + "serde", + "serde_json", + "tempfile", + "ureq", +] + [[package]] name = "rocm-engine-lemonade" version = "0.1.0" diff --git a/Cargo.toml b/Cargo.toml index dfa77fd49..65f187950 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/README.md b/README.md index 94592919f..bfdba61ea 100644 --- a/README.md +++ b/README.md @@ -748,7 +748,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`. `--yes` approves required system-package installs, such as OpenMPI for vLLM, without asking. @@ -796,7 +796,7 @@ GPU memory, and preferred engines. 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] @@ -806,7 +806,20 @@ rocm serve [--engine lemonade|vllm] [--allow-public-bind] [--require-api-key] [--api-key KEY] [--temperature TEMP] [--top-p PROB] [--max-tokens N] [--gpu-memory-utilization FRACTION] [--tool-call-parser NAME] -``` + [--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 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/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/chat.rs b/apps/rocm/src/chat.rs index 28cfca266..ed6a054a9 100644 --- a/apps/rocm/src/chat.rs +++ b/apps/rocm/src/chat.rs @@ -25,7 +25,8 @@ //! above. use crate::freeform_planner::{ - format_structured_tool_call, plan_freeform_request_with_context, serve_args_request_cpu_device, + format_structured_tool_call, infer_engine_from_request, plan_freeform_request_with_context, + serve_args_request_cpu_device, }; use crate::managed_service_status::{ friendly_engine_label, load_managed_services, managed_service_is_live, @@ -326,7 +327,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") ) } @@ -437,7 +438,7 @@ 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) @@ -1417,17 +1418,16 @@ fn fallback_config_tool_call_for_prompt(normalized: &str) -> Option 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 { @@ -5148,7 +5148,7 @@ endpoint: 127.0.0.1:8001 // Stated only where it is true: present on Windows, absent elsewhere. let grounded = rocm_chat_tool_system_prompt_for_host(None); assert_eq!( - grounded.contains("vLLM is skipped on native Windows"), + grounded.contains("vLLM and AIM are skipped on native Windows"), cfg!(windows), "the vLLM caveat must track the host:\n{grounded}" ); @@ -5524,6 +5524,32 @@ model recipes assert_eq!(approval.command_title, "Config"); } + #[test] + fn default_engine_fallback_matches_aim_only_as_a_word() { + let proposed = |prompt: &str| { + fallback_config_tool_call_for_prompt(prompt).map(|call| { + call.arguments["args"][2] + .as_str() + .expect("engine arg") + .to_owned() + }) + }; + assert_eq!( + proposed("change the default engine to aim").as_deref(), + Some("aim") + ); + assert_eq!( + proposed("change the default engine to vllm, aiming for low latency").as_deref(), + Some("vllm") + ); + assert_eq!( + proposed( + "change the default engine so I can maintain my claim, aiming for low latency" + ), + None + ); + } + #[test] fn setup_status_is_read_only() { for args in [ diff --git a/apps/rocm/src/chat_host_facts.rs b/apps/rocm/src/chat_host_facts.rs index 3e2f4f31c..84dc740a7 100644 --- a/apps/rocm/src/chat_host_facts.rs +++ b/apps/rocm/src/chat_host_facts.rs @@ -117,10 +117,11 @@ impl HostFacts { format!("- ROCm/TheRock family for this GPU: {f}\n") }); let engines = if self.vllm_available() { - "- Serving engines that run here: Lemonade and vLLM.\n" + "- Serving engines that run here: Lemonade, vLLM and AIM (AIM runs models from Docker \ + containers and needs a running Docker daemon).\n" } else { - "- Serving engines that run here: Lemonade. vLLM is skipped on native Windows; \ - tell the user to use WSL2 or Linux for that ROCm GPU engine.\n" + "- Serving engines that run here: Lemonade. vLLM and AIM are skipped on native Windows; \ + tell the user to use WSL2 or Linux for those ROCm GPU engines.\n" }; format!( "Facts about the machine you are running on, detected by ROCm CLI. Trust these \ @@ -171,8 +172,8 @@ mod tests { "the detected adapter and target must both appear:\n{block}" ); assert!( - block.contains("Lemonade and vLLM"), - "vLLM is available under WSL:\n{block}" + block.contains("Lemonade, vLLM and AIM"), + "vLLM and AIM are available under WSL:\n{block}" ); assert!( !block.contains("ROCm/TheRock family"), @@ -194,7 +195,7 @@ mod tests { let block = facts.render(); assert!(block.contains("- Operating system: Windows"), "{block}"); assert!( - block.contains("vLLM is skipped on native Windows"), + block.contains("vLLM and AIM are skipped on native Windows"), "the caveat must survive, now stated only where it is true:\n{block}" ); assert!( diff --git a/apps/rocm/src/engine_ipc.rs b/apps/rocm/src/engine_ipc.rs index b26dbd873..05409c16f 100644 --- a/apps/rocm/src/engine_ipc.rs +++ b/apps/rocm/src/engine_ipc.rs @@ -324,6 +324,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, } } @@ -342,7 +343,7 @@ pub(crate) fn run_builtin_engine_stdio(engine: &str) -> Result<()> { } pub(crate) fn builtin_engine_available(engine: &str) -> bool { - matches!(engine, "lemonade" | "vllm") + matches!(engine, "lemonade" | "vllm" | "aim") } #[allow(clippy::too_many_arguments)] @@ -388,6 +389,19 @@ pub(crate) 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/engines_cmd.rs b/apps/rocm/src/engines_cmd.rs index b17443892..9a2532937 100644 --- a/apps/rocm/src/engines_cmd.rs +++ b/apps/rocm/src/engines_cmd.rs @@ -179,7 +179,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( @@ -202,15 +202,24 @@ pub(crate) fn env_root_for_engine_install( runtime_id: &str, ) -> 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 [ @@ -335,7 +344,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/freeform_planner.rs b/apps/rocm/src/freeform_planner.rs index add556fa0..ce2d18c62 100644 --- a/apps/rocm/src/freeform_planner.rs +++ b/apps/rocm/src/freeform_planner.rs @@ -1204,10 +1204,17 @@ fn contains_planner_word(text: &str, expected: &str) -> bool { .any(|word| word == expected) } -fn infer_engine_from_request(lower: &str) -> Option<&'static str> { - ["lemonade", "vllm"] +pub(crate) 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> { @@ -1279,6 +1286,26 @@ fn quote_tool_arg(value: &str) -> String { #[cfg(test)] mod tests { + use super::infer_engine_from_request; + + /// "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") + ); + } + use super::*; use crate::tests::test_paths; diff --git a/apps/rocm/src/main.rs b/apps/rocm/src/main.rs index ed0ad9df7..b5ae462ca 100644 --- a/apps/rocm/src/main.rs +++ b/apps/rocm/src/main.rs @@ -148,7 +148,7 @@ use std::time::{Duration, SystemTime}; /// 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( @@ -587,6 +587,14 @@ rocm serve qwen --verbose --device gpu_required")] /// authenticated too. #[arg(long)] api_key: Option, + /// 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. #[command(alias = "comfy")] @@ -2395,6 +2403,7 @@ fn dispatch(cli: Cli) -> Result<()> { top_p, max_tokens, api_key, + engine_env, }) => serve(ServeArgs { model, engine, @@ -2416,6 +2425,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), @@ -2649,6 +2659,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", + ), ] } @@ -3052,26 +3066,55 @@ 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. -const fn engine_ruled_out_by_platform(engine: &str) -> bool { - rocm_core::runtime_is_windows() && engine.eq_ignore_ascii_case("vllm") +/// 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, + /// AIM's Docker GPU passthrough exists only for Linux and WSL2. + NativeWindowsAim, + 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 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() + { + 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(|| 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 | PlatformRuleOut::NativeWindowsAim => 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. @@ -5943,13 +5986,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) { - Some( +fn model_registry_adapter_availability_note(engine: &str) -> Option<&'static str> { + 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", + ), + 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", - ) - } else { - None + ), + PlatformRuleOut::NativeWindowsAim => Some( + "runtime_status=unsupported_native_windows reason=native Windows skipped; AIM GPU serving is Linux/WSL2 only; run /engine for adapter details", + ), } } @@ -6925,7 +6973,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; @@ -7199,7 +7252,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), ); @@ -8032,6 +8085,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", + ), ] } @@ -9036,7 +9093,7 @@ fn wait_for_service_http_ready( engine, host, port, - canonical_model_id, + &|| canonical_model_id.to_owned(), endpoint_api_key, timeout, &mut |_elapsed| {}, @@ -9062,7 +9119,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), @@ -9071,6 +9128,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, @@ -13815,6 +13876,37 @@ 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(); + // 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(); + 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(); + drop(env); + + 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 @@ -16018,7 +16110,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")); @@ -16158,6 +16250,8 @@ 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(); @@ -17322,8 +17416,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:?}" @@ -17334,13 +17428,118 @@ pub(crate) mod tests { ); } + /// 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:?}" + ); + } + /// 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", 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. + /// + /// 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) + .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` diff --git a/apps/rocm/src/managed_service_status.rs b/apps/rocm/src/managed_service_status.rs index c1083782c..4503d8096 100644 --- a/apps/rocm/src/managed_service_status.rs +++ b/apps/rocm/src/managed_service_status.rs @@ -44,6 +44,7 @@ pub(crate) fn friendly_engine_label(engine: &str) -> &str { match engine { "lemonade" => "Lemonade", "vllm" => "vLLM", + "aim" => "AIM", other => other, } } diff --git a/apps/rocm/src/providers.rs b/apps/rocm/src/providers.rs index 328ad02ce..ce5d3a04a 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( @@ -334,14 +328,8 @@ 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 body = openai_local_stream_chat_request_body( - request - .model - .as_deref() - .unwrap_or(&service.canonical_model_id), - request, - ); + 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); stream_json_from_local_endpoint( @@ -476,6 +464,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() @@ -2479,6 +2471,38 @@ 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 server = spawn_models_only_server("Org/Served")?; + let mut aim = ManagedServiceRecord::new( + &paths, + "svc-aim", + "aim", + "aim-qwen3.5-9b", + "docker.io/org/aim-image:1", + "127.0.0.1", + server.port(), + "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"))?; + 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(()) + } + #[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/rocm/src/serve_cmd.rs b/apps/rocm/src/serve_cmd.rs index 5c2fbf49c..404067f91 100644 --- a/apps/rocm/src/serve_cmd.rs +++ b/apps/rocm/src/serve_cmd.rs @@ -130,16 +130,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,6 +186,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: engine_recipe.extra_env.clone(), preferred_endpoint: engine_recipe.preferred_endpoint.as_ref().map(|endpoint| { EngineRecipeEndpointHint { endpoint_mode: endpoint.endpoint_mode.clone(), @@ -483,6 +489,58 @@ fn parse_gpu_memory_utilization(value: Option<&str>) -> Result> { Ok(Some(parsed)) } +/// 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() || !engine.eq_ignore_ascii_case("aim") { + 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('=').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()); + } + 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| { @@ -515,6 +573,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 +598,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 +687,8 @@ 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 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}'" @@ -842,6 +904,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(); @@ -905,13 +970,13 @@ 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() }; spinner.clear(); - let notes = collect_serve_notes( + let mut notes = collect_serve_notes( cpu_only, &gpu_selection, rocr_visible_devices_set, @@ -921,10 +986,13 @@ 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, - 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(), @@ -1073,6 +1141,59 @@ mod tests { ); } + #[test] + 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!( + 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] 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). @@ -1090,6 +1211,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"); @@ -1136,6 +1273,8 @@ 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(), @@ -1145,6 +1284,8 @@ mod tests { unsupported_combinations: Vec::new(), notes: Vec::new(), model_id_override: None, + image_digest: None, + extra_env: BTreeMap::new(), }, ]; @@ -1509,6 +1650,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( @@ -1563,6 +1724,107 @@ 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_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_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] + 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/rocm/src/service_launch.rs b/apps/rocm/src/service_launch.rs index 6586b9a73..46e6431c6 100644 --- a/apps/rocm/src/service_launch.rs +++ b/apps/rocm/src/service_launch.rs @@ -415,6 +415,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. // @@ -587,7 +588,11 @@ pub(crate) fn start_managed_service( engine, host, port, - &resolve.canonical_model_id, + &|| { + 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, @@ -1218,6 +1223,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(), @@ -1307,6 +1313,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(), @@ -1391,6 +1398,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 1a8c3335e..79793d1f8 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; @@ -37,6 +38,31 @@ 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, + }, + ) +} + +/// [`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 @@ -48,6 +74,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. 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 @@ -86,8 +118,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 = !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) + && 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 @@ -150,10 +185,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. @@ -1845,6 +1898,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/coverage-floors.toml b/coverage-floors.toml index b1b9d7a95..aa2ece4ed 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" = 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 diff --git a/crates/rocm-core/src/lib.rs b/crates/rocm-core/src/lib.rs index 7a5f5fd9a..b3ca27796 100644 --- a/crates/rocm-core/src/lib.rs +++ b/crates/rocm-core/src/lib.rs @@ -185,7 +185,7 @@ pub use runtime::{ runtime_rocm_library_filename, shell_command_for_host, user_runtime_dir, }; pub use service_identity::{ServiceId, generate_service_id}; -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/model_catalog.json b/crates/rocm-core/src/model_catalog.json index f970bbbbf..c107637f3 100644 --- a/crates/rocm-core/src/model_catalog.json +++ b/crates/rocm-core/src/model_catalog.json @@ -4,13 +4,13 @@ "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)", - "engines": ["lemonade", "llama.cpp", "llamacpp"], + "label": "AMD Radeon - RX 7000/9000, Radeon PRO (Lemonade / llama.cpp / AIM)", + "engines": ["lemonade", "llama.cpp", "llamacpp", "aim"], "gfx_families": ["gfx110X-all", "gfx120X-all"] }, { @@ -508,6 +508,50 @@ "warnings": [ "multimodal model; text chat is the default serve path" ] + }, + { + "canonical_model_id": "aim-qwen3.5-9b", + "aliases": [ + "amdenterpriseai/aim-radeon-qwen-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": [ + "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", + "image_digest": "sha256:ee1201261f358be881318de5930ce9505d17309486f747a83147d7d042fc95ce", + "extra_env": { + "AIM_PROFILE_ID": "vllm-w7900-bf16-tp1-latency" + } + } + ], + "manual_alternatives": [], + "chat_template_mode": "auto", + "preferred_engines": [ + "aim" + ], + "warnings": [ + "Radeon-targeted AIM image; requires Docker and an explicit AIM_PROFILE_ID" + ] } ] } diff --git a/crates/rocm-core/src/model_catalog.rs b/crates/rocm-core/src/model_catalog.rs index 5bd43f3eb..18c412604 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)] @@ -80,6 +85,13 @@ 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:` 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, } #[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq)] @@ -641,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, @@ -653,6 +671,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 { @@ -670,12 +692,31 @@ 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() && engine_recipe.unsupported_combinations.is_empty() && engine_recipe.notes.is_empty() && engine_recipe.model_id_override.is_none() + && engine_recipe.image_digest.is_none() + && engine_recipe.extra_env.is_empty() { bail!( "engine recipe for `{}` on `{canonical_model_id}` must not be empty", @@ -837,7 +878,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!( @@ -921,6 +962,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"]); @@ -935,6 +1017,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())]), @@ -945,6 +1028,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, diff --git a/crates/rocm-core/src/net.rs b/crates/rocm-core/src/net.rs index 03b6616cb..b60a2a7d7 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, ) @@ -972,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) @@ -1081,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}; diff --git a/crates/rocm-core/src/service_record.rs b/crates/rocm-core/src/service_record.rs index ce0459001..37bb07d94 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, @@ -106,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( @@ -134,6 +152,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 +182,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 +311,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 +441,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-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/crates/rocm-dash-tui/src/ui/serve_wizard.rs b/crates/rocm-dash-tui/src/ui/serve_wizard.rs index f4af2b5d3..ed3a5da46 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/crates/rocm-engine-protocol/src/lib.rs b/crates/rocm-engine-protocol/src/lib.rs index fbca3104f..6ca773ccd 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, @@ -511,6 +529,13 @@ 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. + /// 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, #[serde(default)] @@ -621,6 +646,32 @@ pub struct LogsResponse { #[cfg(test)] mod tests { + fn resolve_response(served: Option<&str>) -> ResolveModelResponse { + ResolveModelResponse { + canonical_model_id: "image:tag".to_owned(), + served_model_id: served.map(str::to_owned), + task: "chat".to_owned(), + source: "s".to_owned(), + revision: "r".to_owned(), + loader: "l".to_owned(), + trust_remote_code: false, + chat_template_mode: "auto".to_owned(), + dtype: "auto".to_owned(), + device_policy: DevicePolicy::GpuRequired, + estimated_memory: "1 GiB".to_owned(), + launch_defaults: Value::Null, + engine_recipe: None, + warnings: Vec::new(), + } + } + + #[test] + fn served_model_name_prefers_a_non_blank_served_id() { + assert_eq!(resolve_response(Some("qwen")).served_model_name(), "qwen"); + assert_eq!(resolve_response(Some(" ")).served_model_name(), "image:tag"); + assert_eq!(resolve_response(None).served_model_name(), "image:tag"); + } + use super::*; #[test] @@ -888,6 +939,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 +973,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 +995,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/docs/architecture.md b/docs/architecture.md index 4f2244be8..1f6b98eb3 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -160,7 +160,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 expectation `Grid` and scenario-reference map built from the same `inputs` means a caller producing both outputs can no longer pass them out of sync with each other. Building them there also means the grid is built once per run rather than once per report section that renders it (ROCMAI-660). `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: @@ -181,3 +181,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) — 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/engine-plugins.md b/docs/engine-plugins.md index b7d7e2b40..682594302 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..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 @@ -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 f2bbf984f..b9c066d36 100644 --- a/docs/manual-testing.md +++ b/docs/manual-testing.md @@ -824,3 +824,54 @@ 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 aim-qwen3.5-9b --engine aim --managed +rocm services list --all # note the service id, e.g. aim-- +rocm services stop --yes +docker ps -a --filter name=rocm-aim- +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: + +```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`. 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, 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 51c6f8c04..176cab068 100644 --- a/docs/testing.md +++ b/docs/testing.md @@ -830,6 +830,8 @@ 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 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 new file mode 100644 index 000000000..97b7d0265 --- /dev/null +++ b/engines/aim/Cargo.toml @@ -0,0 +1,28 @@ +[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 +tempfile = "3" +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..14a1c264b --- /dev/null +++ b/engines/aim/src/docker.rs @@ -0,0 +1,986 @@ +// 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::collections::BTreeMap; +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"; + +/// 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) +} + +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())); + } + } + (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, + 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!( + "{}:{}:{}", + 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) { + 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 +} + +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) +} + +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 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 {}` failed: {}", + spec.name, + rocm_core::command_failure_detail(&output) + ); + } + 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 { + let output = Command::new("getent") + .arg("group") + .arg("render") + .output() + .ok()?; + if !output.status.success() { + return None; + } + 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() +} + +/// 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 { + command.arg("kill").arg(name); + } else { + command + .arg("stop") + .arg("-t") + .arg(grace.as_secs().to_string()) + .arg(name); + } + 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() { + let detail = rocm_core::command_failure_detail(&output); + if is_nothing_to_stop(&detail) { + return Ok(false); + } + bail!("failed to stop container `{name}`: {detail}"); + } + 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 { + 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 +/// 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 is_no_such(&detail) { + return Ok(()); + } + bail!("failed to remove container `{name}`: {detail}"); + } + 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, +} + +/// 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("{{.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() { + let detail = rocm_core::command_failure_detail(&output); + if is_no_such(&detail) { + 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()), + } +} + +pub(crate) fn logs(name: &str, tail: usize) -> Result> { + let mut command = Command::new(DOCKER_BIN); + command + .arg("logs") + .arg("--tail") + .arg(tail.to_string()) + .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!( + "`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) +} + +/// 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") + .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(()) +} + +/// 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 { + 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::{ + 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; + + 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 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!( + 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() { + 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] + 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")); + } + + #[test] + fn parse_container_state_reads_running_and_exit_code() { + assert_eq!( + parse_container_state("true 0\n").unwrap(), + ContainerState::Running + ); + assert_eq!( + 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_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")); + 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 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))] +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; + + 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. + 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}" + ;; + 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 + 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) + 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) + 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 +"#; + + 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() + } + + /// 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() + } + } + + 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. + 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(), + 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 fixture = Fixture { dir: dir.clone() }; + + 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()) + .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 + } + + 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(&plain_spec(&BTreeMap::new())).expect("fake docker run should succeed"); + return; + } + 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(&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}"); + 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; + } + 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_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_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] + 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 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() { + 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::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] + 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!(super::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")], + ); + } + + #[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/install.rs b/engines/aim/src/install.rs new file mode 100644 index 000000000..d15f1a7cb --- /dev/null +++ b/engines/aim/src/install.rs @@ -0,0 +1,109 @@ +// 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(), + ], + }) +} + +#[cfg(all(test, unix))] +mod tests { + use super::install_response; + use rocm_core::fake_executable::write_fake_executable; + use rocm_engine_protocol::InstallRequest; + use std::path::PathBuf; + use std::process::{Command, Stdio}; + + const DOCKER_UNREACHABLE_TEST_DIR: &str = "ROCM_AIM_DOCKER_UNREACHABLE_TEST_DIR"; + + /// Child half of [`install_response_bails_clearly_when_docker_is_unreachable`]. + /// + /// 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(DOCKER_UNREACHABLE_TEST_DIR).is_none() { + return; + } + let request = InstallRequest { + runtime_id: "aim".to_owned(), + python_version: None, + reinstall: false, + env_root: None, + }; + let error = install_response(request).expect_err("docker is unreachable"); + assert_eq!( + error.to_string(), + "docker daemon is not reachable; install Docker and ensure it is running before using the AIM engine" + ); + } + + #[test] + fn install_response_bails_clearly_when_docker_is_unreachable() { + let dir = std::env::temp_dir().join(format!( + "aim-install-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(); + 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 = Command::new(std::env::current_exe().unwrap()) + .args([ + "--exact", + "install::tests::docker_unreachable_subprocess_helper", + "--nocapture", + ]) + .env(DOCKER_UNREACHABLE_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(); + } +} diff --git a/engines/aim/src/lib.rs b/engines/aim/src/lib.rs new file mode 100644 index 000000000..3bab19210 --- /dev/null +++ b/engines/aim/src/lib.rs @@ -0,0 +1,430 @@ +// Copyright © Advanced Micro Devices, Inc., or its affiliates. +// +// SPDX-License-Identifier: MIT + +use anyhow::{Context, Result}; +use clap::{Parser, Subcommand}; +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, + pub state_path: PathBuf, + pub engine_recipe: Option, +} + +pub fn run_cli() -> Result<()> { + let cli = Cli::parse(); + match cli.command { + CommandKind::Detect => print_json(&detect_response(None)), + 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, + // 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())?; + crate::process::serve_http(ServeHttpRequest { + service_id, + model_ref, + host, + port, + device_policy: crate::process::parse_device_policy_arg(&device_policy)?, + gpu_indices, + state_path, + engine_recipe: crate::process::parse_engine_recipe_json(engine_recipe_json)?, + }) + } + } +} + +pub fn builtin_handle_envelope(envelope: EngineRequestEnvelope) -> EngineResponseEnvelope { + 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, + model_ref: String, + host: String, + port: u16, + device_policy: DevicePolicy, + gpu_indices: Vec, + // 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, + engine_recipe: Option, +) -> Result<()> { + crate::process::serve_http(ServeHttpRequest { + service_id, + model_ref, + host, + port, + device_policy, + gpu_indices, + state_path, + engine_recipe, + }) +} + +fn handle_envelope(envelope: EngineRequestEnvelope) -> EngineResponseEnvelope { + match envelope.method { + EngineMethod::Detect => { + deserialize_and_respond::(envelope.payload, |request| { + Ok(detect_response(request.runtime_id)) + }) + } + 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()), + } +} + +/// `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 { + 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: reachable.then_some(runtime_id).flatten(), + 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 && !cfg!(windows), + reason: if cfg!(windows) { + Some(windows_gpu_unsupported_message().to_owned()) + } else if reachable { + None + } else { + Some("docker daemon unreachable".to_owned()) + }, + }], + capabilities: capabilities(), + notes, + } +} + +#[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, + // `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(), + reasoning_parser: false, + } +} + +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() + .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..bfb6eef68 --- /dev/null +++ b/engines/aim/src/process.rs @@ -0,0 +1,1507 @@ +// Copyright © Advanced Micro Devices, Inc., or its affiliates. +// +// SPDX-License-Identifier: MIT + +use anyhow::{Context, Result, bail}; +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. +/// +/// 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. +/// 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() + .map(|c| { + if c.is_ascii_alphanumeric() || c == '_' || c == '.' || c == '-' { + c + } else { + '-' + } + }) + .collect(); + 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 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() + }); + + // 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() + && !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 { + // Must stay the Docker image reference: `apps/rocm` re-threads + // `canonical_model_id` back into Launch/ServeHttp as `model_ref` + // (service id generation, `__engine-serve-http` argv, `docker pull`/`run`) + // for every built-in engine uniformly. The name the container serves + // under (the image's HF-style canonical name) is reported separately in + // `served_model_id`, which readiness, the smoke test and the summary use. + canonical_model_id: request.model_ref, + served_model_id: metadata.canonical_name, + 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 state_path = crate::state::state_path(&request.service_id); + + let served_model_id = run_container( + &request.service_id, + &request.model_ref, + &request.host, + request.port, + &device_policy, + &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( + &state_path, + &request.service_id, + &request.model_ref, + served_model_id.as_deref(), + &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, so there is + // no client process to report a real PID for. + pid: 0, + endpoint_url, + log_path, + state_path: state_path.display().to_string(), + }) +} + +/// 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, + requested_gpus: &[u32], + extra_env: Option<&BTreeMap>, +) -> Result> { + let name = container_name(service_id); + // 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); + + 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)?; + } + // 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, + gpus.available, + image_wants_cpu, + )?; + let render_gid = if use_gpu { + crate::docker::render_group_gid() + } else { + None + }; + + let mut env = extra_env.cloned().unwrap_or_default(); + 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 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(&crate::docker::RunSpec { + name: &name, + image: &image_ref, + host, + port, + container_port: crate::DEFAULT_AIM_PORT, + use_gpu, + gpu_indices: if use_gpu { &gpus.indices } else { &[] }, + render_gid, + 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, + // `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 \ + 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 + } + }) +} + +pub(crate) fn serve_http(request: ServeHttpRequest) -> Result<()> { + let name = container_name(&request.service_id); + let label_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), + ) + .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, + &request.service_id, + &request.model_ref, + label_name.as_deref(), + &request.host, + request.port, + request.device_policy, + request.engine_recipe.as_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, + label_name.as_deref(), + aim_ready_timeout(), + ) { + 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) => { + // 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) + } + } +} + +/// 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}`") + } + } +} + +/// Mirrors `engines/vllm/src/process.rs`'s `wait_for_vllm_ready` architecture: +/// early-exit check, monotonic timeout, listed-then-inference-verified gate. +/// 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, + expected_model: Option<&str>, + timeout: Duration, +) -> Result { + let start = Instant::now(); + let endpoint = endpoint_url(host, port); + let mut consecutive_errors = 0; + loop { + 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"); + bail!( + "AIM container `{name}` did not become ready within {}s:\n{tail}", + timeout.as_secs() + ); + } + + 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(model); + } + + 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 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, + }) +} + +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::*; + 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() { + 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"); + } + + #[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()); + } + + #[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_never_falls_back_to_cpu() { + assert!( + resolve_use_gpu("svc", "model", &DevicePolicy::GpuPreferred, false, false).is_err() + ); + assert!(resolve_use_gpu("svc", "model", &DevicePolicy::GpuPreferred, true, true).is_err()); + 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 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 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 + } + + 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", + ), + ], + ); + } + + #[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; + } + 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", + ), + ], + ); + } + + #[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() + ); + } + + #[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() { + run_container( + "svc-already-running", + IMAGE, + "127.0.0.1", + 18080, + &DevicePolicy::GpuPreferred, + &[], + None, + ) + .unwrap(); + return; + } + let fixture = run_with_fake_docker( + "process::subprocess_tests::run_container_skips_docker_run_when_already_running", + &[ + ("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() { + run_container( + "svc-pull-then-run", + IMAGE, + "127.0.0.1", + 18081, + &DevicePolicy::CpuOnly, + &[], + None, + ) + .unwrap(); + return; + } + 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"), + ], + ); + assert!( + fixture + .argv() + .iter() + .any(|line| line == "pull org/aim-fixture:1") + ); + } + + #[test] + fn run_container_bails_when_pull_fails() { + if in_subprocess() { + let error = run_container( + "svc-pull-fails", + 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_INSPECT_SEQUENCE", "missing"), + ("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", + IMAGE, + "127.0.0.1", + 18083, + &DevicePolicy::CpuOnly, + &[], + 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_INSPECT_SEQUENCE", "missing"), + ("FAKE_DOCKER_IMAGE_PRESENT_EXIT", "0"), + ("FAKE_DOCKER_RUN_EXIT", "1"), + ("FAKE_DOCKER_RUN_STDERR", "boom"), + ], + ); + } + + #[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 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() { + let error = wait_for_container_ready( + "fake-container", + "127.0.0.1", + 19999, + Some("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_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). + let error = wait_for_container_ready( + "fake-container", + "127.0.0.1", + 59123, + Some("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 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(), + force: true, + }) + .unwrap(); + assert!(response.stopped); + assert!(!response.graceful); + let state_path = crate::state::state_path("svc-stop-succeeds"); + assert_eq!(status_of(&state_path).as_deref(), Some("stopped")); + return; + } + 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" + ] + ); + // 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] + 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(), + force: true, + }) + .unwrap(); + assert!(!response.stopped); + assert!(!response.graceful); + return; + } + run_with_fake_docker( + "process::subprocess_tests::stop_service_of_a_missing_container_reports_not_stopped_but_still_cleans_up", + &[ + ("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"), + ], + ); + assert!(!fixture.argv().iter().any(|line| line.starts_with("rm "))); + } +} diff --git a/engines/aim/src/registry.rs b/engines/aim/src/registry.rs new file mode 100644 index 000000000..29e0be5d0 --- /dev/null +++ b/engines/aim/src/registry.rs @@ -0,0 +1,295 @@ +// Copyright © Advanced Micro Devices, Inc., or its affiliates. +// +// SPDX-License-Identifier: MIT + +//! 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; +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_type: Option, + pub pytorch_rocm_arch: Vec, + pub canonical_name: Option, +} + +/// 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}/{image_ref}") + } +} + +/// 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) { + 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 > 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 { + 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")) +} + +/// 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)?; + + let manifest_url = format!("https://{REGISTRY_HOST}/v2/{repo}/manifests/{reference}"); + 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_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)) +} + +/// 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 + .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_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.get("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 normalize_image_ref_defaults_org_for_bare_name() { + assert_eq!( + 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" + ); + } + + #[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 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_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 + // 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_image_config_extracts_known_fields() { + let config = serde_json::json!({ + "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_image_config(&config); + 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..84e329454 --- /dev/null +++ b/engines/aim/src/state.rs @@ -0,0 +1,399 @@ +// 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(); + // The name the container lists in `/v1/models`, which is not the image ref + // `model_ref` holds (see `process::resolve_model_response`). + let model_ref = state + .as_ref() + .and_then(|s| value_string(s, "served_model_id").or_else(|| 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")) + }, + ) +} + +/// 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 +} + +#[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, + 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, + &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, + "device_policy": device_policy, + "engine_recipe": engine_recipe, + "started_at_unix_ms": current_unix_millis(), + }), + ) +} + +/// 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() { + 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), + ) +} + +/// 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); + } + 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 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 = 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(), + extra_env: std::collections::BTreeMap::from([( + "HF_TOKEN".to_owned(), + "hf_super_secret".to_owned(), + )]), + ..EngineRecipeHint::default() + }; + write_running_state( + &path, + "svc-secret", + "some/model", + Some("Org/Served"), + "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")); + } + + #[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] + 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/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` 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 b01dfa30f..c0385d8e0 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 new file mode 100644 index 000000000..a7acdee31 --- /dev/null +++ b/tests/e2e-cucumber/features/model_catalog.feature @@ -0,0 +1,27 @@ +Feature: Model catalog listing + + # 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 + # models" wording either way. + @id:model-catalog-lists-aim-recipe + 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 + + # 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 plants a `docker` that always fails so the + # 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). + # 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 + 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 90caa1aee..25fed56b7 100644 --- a/tests/e2e-cucumber/features/model_serving.feature +++ b/tests/e2e-cucumber/features/model_serving.feature @@ -341,3 +341,41 @@ 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). 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` + # 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 + # 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 + 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/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 548b4700a..3f4a03727 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,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'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), } } @@ -647,7 +681,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/src/expectation.rs b/tests/e2e-cucumber/src/expectation.rs index 71c1c93bc..380d8210c 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.rs b/tests/e2e-cucumber/tests/e2e.rs index 18c82bbed..d0d368577 100644 --- a/tests/e2e-cucumber/tests/e2e.rs +++ b/tests/e2e-cucumber/tests/e2e.rs @@ -34,6 +34,7 @@ mod e2e { pub mod host_guard; 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; @@ -403,14 +404,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/examine_steps.rs b/tests/e2e-cucumber/tests/e2e/examine_steps.rs index 42613a6ce..67898e407 100644 --- a/tests/e2e-cucumber/tests/e2e/examine_steps.rs +++ b/tests/e2e-cucumber/tests/e2e/examine_steps.rs @@ -286,7 +286,7 @@ async fn assert_host_default_engine_reported(world: &mut E2eWorld) { #[then("all supported engines are listed")] async fn assert_all_engines_listed(world: &mut E2eWorld) { let output = world.cli_output.as_ref().expect("no command was run"); - for engine in ["lemonade", "vllm"] { + for engine in ["lemonade", "vllm", "aim"] { assert!( output.contains(engine), "engine '{engine}' not found in:\n{output}" 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..3e0e077d5 --- /dev/null +++ b/tests/e2e-cucumber/tests/e2e/model_catalog_steps.rs @@ -0,0 +1,62 @@ +// Copyright © Advanced Micro Devices, Inc., or its affiliates. +// +// SPDX-License-Identifier: MIT + +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) { + // `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); +} + +#[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("aim-qwen3.5-9b") && output.contains("engines=[aim]"), + "expected the AIM recipe listed by canonical id and engine:\n{output}" + ); +} + +/// 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}" + ); +} diff --git a/tests/e2e-cucumber/tests/e2e/serving_steps.rs b/tests/e2e-cucumber/tests/e2e/serving_steps.rs index 7ce18b326..e8dc2b028 100644 --- a/tests/e2e-cucumber/tests/e2e/serving_steps.rs +++ b/tests/e2e-cucumber/tests/e2e/serving_steps.rs @@ -633,6 +633,81 @@ 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) { + 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. + // + // 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", + "--engine-env", + "AIM_PROFILE_ID=vllm-w7900-bf16-tp1-latency", + "--managed", + ], + model, + "Qwen3.5-9B", + ) + .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}" + ); + // 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 +/// 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 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"), 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"),