From f74180e484e03bd336b01749bca3c335c91c46e7 Mon Sep 17 00:00:00 2001 From: Roman Inflianskas Date: Fri, 3 Jul 2026 18:34:23 +0000 Subject: [PATCH 1/2] Limit serving engines to Lemonade and vLLM Narrow the supported model backends to Lemonade (with its bundled engines) and vLLM, and remove the standalone atom, llama-cpp, pytorch, and sglang engine integrations. This shrinks the maintenance, testing, and packaging surface. vLLM is kept as a separate backend since Lemonade's bundled vLLM is not yet tuned for Instinct. - Delete the four engine crates and their workspace members, path deps, and serve/registry/inventory/dispatch wiring. - Constrain the user-facing `--engine` / `engines install|shell` / `config set[-default]-engine` arguments to `lemonade`/`vllm` via a clap value_parser, so `--help` and shell completion advertise exactly those two. - Retarget built-in model recipes off the removed engines: GGUF recipes serve via Lemonade; HF safetensors recipes serve via vLLM; existing vLLM preferences are unchanged. - Prune the EngineKind registry, TUI engine lists, the local-assistant system prompt, docs, CI, and acceptance/smoke scripts (Linux and Windows). Framework diagnostics for PyTorch and llama.cpp (as user frameworks that `rocm diagnose`/`examine` help debug) and Lemonade's own llama.cpp backend are intentionally retained. Signed-off-by: Roman Inflianskas --- .github/workflows/ci.yml | 10 +- AGENTS.md | 4 +- Cargo.lock | 55 - Cargo.toml | 4 - MANIFEST.md | 15 +- README.md | 10 +- apps/rocm/Cargo.toml | 4 - apps/rocm/src/main.rs | 599 +-- apps/rocm/src/providers.rs | 20 +- apps/rocmd/src/lib.rs | 69 +- crates/rocm-core/src/lib.rs | 42 +- .../src/engine_registry.rs | 19 +- crates/rocm-dash-collectors/src/lib.rs | 1 - .../rocm-dash-collectors/src/llama_slots.rs | 30 - crates/rocm-dash-collectors/src/parallel.rs | 6 +- crates/rocm-dash-tui/src/agent.rs | 4 +- crates/rocm-dash-tui/src/app/chat.rs | 2 +- crates/rocm-dash-tui/src/ui/engine_manager.rs | 26 +- crates/rocm-dash-tui/src/ui/serve_wizard.rs | 2 +- crates/rocm-dash-tui/src/ui/tabs/serving.rs | 4 +- crates/rocm-engine-protocol/src/lib.rs | 49 +- docs/atom.md | 82 - docs/engine-plugins.md | 16 +- docs/llm-tool-use.md | 28 +- docs/manual-testing.md | 70 +- docs/sglang.md | 115 - docs/testing.md | 213 +- docs/wsl.md | 12 +- engines/atom/Cargo.toml | 23 - engines/atom/src/lib.rs | 1685 -------- engines/atom/src/main.rs | 7 - engines/llama-cpp/Cargo.toml | 23 - engines/llama-cpp/src/lib.rs | 2818 ------------ engines/llama-cpp/src/main.rs | 7 - engines/pytorch/Cargo.toml | 22 - engines/pytorch/src/lib.rs | 3809 ----------------- engines/pytorch/src/main.rs | 7 - engines/pytorch/src/python_worker.py | 636 --- engines/sglang/Cargo.toml | 23 - engines/sglang/src/lib.rs | 1856 -------- engines/sglang/src/main.rs | 7 - engines/vllm/src/lib.rs | 2 +- install.ps1 | 2 +- install.sh | 2 +- ...ceptance-install-upgrade-tui-uninstall.ps1 | 6 +- ...cceptance-install-upgrade-tui-uninstall.sh | 10 +- scripts/atom_therock_gpu_test.py | 481 --- scripts/build_single_exe_release.py | 4 - scripts/llama_cpp_therock_gpu_test.py | 689 --- scripts/package-linux-release.sh | 6 +- scripts/pytorch_therock_gpu_test.py | 880 ---- scripts/sglang_therock_gpu_test.py | 560 --- scripts/smoke_local.py | 266 +- scripts/therock_sdk_install_test.py | 22 +- 54 files changed, 350 insertions(+), 15014 deletions(-) delete mode 100644 crates/rocm-dash-collectors/src/llama_slots.rs delete mode 100644 docs/atom.md delete mode 100644 docs/sglang.md delete mode 100644 engines/atom/Cargo.toml delete mode 100644 engines/atom/src/lib.rs delete mode 100644 engines/atom/src/main.rs delete mode 100644 engines/llama-cpp/Cargo.toml delete mode 100644 engines/llama-cpp/src/lib.rs delete mode 100644 engines/llama-cpp/src/main.rs delete mode 100644 engines/pytorch/Cargo.toml delete mode 100644 engines/pytorch/src/lib.rs delete mode 100644 engines/pytorch/src/main.rs delete mode 100644 engines/pytorch/src/python_worker.py delete mode 100644 engines/sglang/Cargo.toml delete mode 100644 engines/sglang/src/lib.rs delete mode 100644 engines/sglang/src/main.rs delete mode 100644 scripts/atom_therock_gpu_test.py delete mode 100644 scripts/llama_cpp_therock_gpu_test.py delete mode 100644 scripts/pytorch_therock_gpu_test.py delete mode 100644 scripts/sglang_therock_gpu_test.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index a86f56ea..0c8275f7 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -133,13 +133,9 @@ jobs: - name: Acceptance harness self-tests if: needs.changes.outputs.heavy == 'true' run: | - python scripts/pytorch_therock_gpu_test.py --self-test - python scripts/llama_cpp_therock_gpu_test.py --self-test python scripts/comfyui_therock_gpu_test.py --self-test python scripts/local_assistant_therock_gpu_test.py --self-test python scripts/vllm_therock_gpu_test.py --self-test - python scripts/sglang_therock_gpu_test.py --self-test - python scripts/atom_therock_gpu_test.py --self-test python scripts/wsl_preflight.py --self-test - name: Portable WSL build deps self-test @@ -173,7 +169,7 @@ jobs: shell: pwsh run: | cargo build --workspace --all-targets - cargo build --release -p rocm -p rocmd -p rocm-engine-pytorch -p rocm-engine-llama-cpp -p rocm-engine-lemonade -p rocm-engine-atom -p rocm-engine-vllm -p rocm-engine-sglang -p xtask + cargo build --release -p rocm -p rocmd -p rocm-engine-lemonade -p rocm-engine-vllm -p xtask - name: Test if: needs.changes.outputs.heavy == 'true' @@ -189,13 +185,9 @@ jobs: if: needs.changes.outputs.heavy == 'true' shell: pwsh run: | - python .\scripts\pytorch_therock_gpu_test.py --self-test - python .\scripts\llama_cpp_therock_gpu_test.py --self-test python .\scripts\comfyui_therock_gpu_test.py --self-test python .\scripts\local_assistant_therock_gpu_test.py --self-test python .\scripts\vllm_therock_gpu_test.py --self-test - python .\scripts\sglang_therock_gpu_test.py --self-test - python .\scripts\atom_therock_gpu_test.py --self-test python .\scripts\wsl_preflight.py --self-test - name: Release readiness self-test diff --git a/AGENTS.md b/AGENTS.md index e751812d..9ce34d58 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -128,13 +128,13 @@ Current workspace members: - apps: `apps/rocm`, `apps/rocmd` - shared crates: `crates/rocm-core`, `crates/rocm-engine-protocol` -- engine crates: `engines/atom`, `engines/lemonade`, `engines/llama-cpp`, `engines/pytorch`, `engines/sglang`, `engines/vllm` +- engine crates: `engines/lemonade`, `engines/vllm` Guardrails: - `crates/rocm-engine-protocol` is a contract surface; verify all impacted engines after protocol changes - preserve strict GPU-required behavior; do not introduce silent CPU fallback -- respect platform gates (for example, native Windows handling for vLLM/SGLang) +- respect platform gates (for example, native Windows handling for vLLM) - pin third-party GitHub Actions to a full commit SHA with a trailing `# vX.Y.Z` comment, never a moving tag (`@v2`, `@main`); a retagged or compromised action otherwise enters CI silently. Bump the SHA and comment together when upgrading - supported host platforms are Windows and Linux only (including WSL where documented) - platforms outside Windows/Linux are unsupported; do not implement, debug, or "fix" unsupported-platform behavior diff --git a/Cargo.lock b/Cargo.lock index dcb33fa3..59dd53de 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3211,12 +3211,8 @@ dependencies = [ "rocm-core", "rocm-dash-daemon", "rocm-dash-tui", - "rocm-engine-atom", "rocm-engine-lemonade", - "rocm-engine-llama-cpp", "rocm-engine-protocol", - "rocm-engine-pytorch", - "rocm-engine-sglang", "rocm-engine-vllm", "rocmd", "rpassword", @@ -3325,18 +3321,6 @@ dependencies = [ "tracing-subscriber", ] -[[package]] -name = "rocm-engine-atom" -version = "0.3.0" -dependencies = [ - "anyhow", - "clap", - "rocm-core", - "rocm-engine-protocol", - "serde", - "serde_json", -] - [[package]] name = "rocm-engine-lemonade" version = "0.3.0" @@ -3350,18 +3334,6 @@ dependencies = [ "ureq", ] -[[package]] -name = "rocm-engine-llama-cpp" -version = "0.3.0" -dependencies = [ - "anyhow", - "clap", - "rocm-core", - "rocm-engine-protocol", - "serde", - "serde_json", -] - [[package]] name = "rocm-engine-protocol" version = "0.3.0" @@ -3371,33 +3343,6 @@ dependencies = [ "serde_json", ] -[[package]] -name = "rocm-engine-pytorch" -version = "0.3.0" -dependencies = [ - "anyhow", - "async-stream", - "axum", - "clap", - "rocm-core", - "rocm-engine-protocol", - "serde", - "serde_json", - "tokio", -] - -[[package]] -name = "rocm-engine-sglang" -version = "0.3.0" -dependencies = [ - "anyhow", - "clap", - "rocm-core", - "rocm-engine-protocol", - "serde", - "serde_json", -] - [[package]] name = "rocm-engine-vllm" version = "0.3.0" diff --git a/Cargo.toml b/Cargo.toml index 2a1f56e3..3ff894a6 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -9,11 +9,7 @@ members = [ "crates/rocm-dash-collectors", "crates/rocm-dash-daemon", "crates/rocm-dash-tui", - "engines/llama-cpp", - "engines/atom", - "engines/pytorch", "engines/lemonade", - "engines/sglang", "engines/vllm", "xtask", ] diff --git a/MANIFEST.md b/MANIFEST.md index 433c881b..dd85043a 100644 --- a/MANIFEST.md +++ b/MANIFEST.md @@ -577,7 +577,7 @@ from the workspace source. The user-facing binaries are: processes and expose a local HTTP API In addition, each engine crate under `engines/` builds its own -`rocm-engine-` host binary (for example `rocm-engine-pytorch`, +`rocm-engine-` host binary (for example `rocm-engine-lemonade`, `rocm-engine-vllm`), spawned by `rocmd` to run a specific inference engine. All binaries are compiled from the workspace source using the standard Cargo @@ -629,13 +629,12 @@ distribution; no Python packages are installed for Lemonade by rocm-cli. ### Engine-Specific Python Dependencies -The PyTorch engine manages its own Python virtual environment using the `uv` -binary described above. Python packages are installed from the TheRock PyPI -index and, where applicable, from public PyPI (`https://pypi.org`). No Python -packages are bundled in the repository. - -The vLLM, SGLang, ATOM, and llama.cpp engines do not install Python packages -automatically; they record externally-provided runtimes supplied by the user. +The vLLM engine manages Python packages through the `uv` binary described +above, installed from the TheRock PyPI index and, where applicable, from public +PyPI (`https://pypi.org`). No Python packages are bundled in the repository. The +vLLM engine can also record an externally-provided runtime supplied by the user +instead of installing one. Lemonade ships as a self-contained embeddable +distribution and installs no Python packages through rocm-cli. ## Repository Structure Notes diff --git a/README.md b/README.md index 2ad9509a..9b79f016 100644 --- a/README.md +++ b/README.md @@ -24,8 +24,8 @@ ROCm CLI is a command-line tool for setting up and running local AI on AMD GPUs, full-screen TUI dashboard for GPU telemetry, model serving, and chat. A single prebuilt binary for Linux and Windows. No Python, Rust, or existing -ROCm install required. Ships with inference engine adapters for PyTorch, -llama.cpp, Lemonade, ATOM, vLLM, and SGLang. +ROCm install required. Ships with inference engine adapters for Lemonade and +vLLM. > [!IMPORTANT] > **Tech Preview** -- This software is provided as-is, without warranty or @@ -166,14 +166,14 @@ 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`, `pytorch`, `llama.cpp`, `atom`, `vllm`, `sglang`. +Supported engines: `lemonade`, `vllm`. ### Model serving Start a local OpenAI-compatible model server: ``` -rocm serve [--engine lemonade|pytorch|llama.cpp|atom|vllm|sglang] +rocm serve [--engine lemonade|vllm] [--device gpu_required|gpu_preferred|cpu_only] [--gpu auto|] [--runtime-id KEY | --env-id ID] @@ -369,6 +369,4 @@ See `docs/commit-signatures.md` for details (GPG signing, GitHub "Verified", and - Testing and verification: `docs/testing.md` - Developer manual QA: `docs/manual-testing.md` - Engine plugin policy: `docs/engine-plugins.md` -- ATOM adapter: `docs/atom.md` - vLLM adapter: `docs/vllm.md` -- SGLang adapter: `docs/sglang.md` diff --git a/apps/rocm/Cargo.toml b/apps/rocm/Cargo.toml index 04f34235..e35336d7 100644 --- a/apps/rocm/Cargo.toml +++ b/apps/rocm/Cargo.toml @@ -26,12 +26,8 @@ rocm-core = { path = "../../crates/rocm-core" } rocm-dash-daemon = { path = "../../crates/rocm-dash-daemon" } rocm-dash-tui = { path = "../../crates/rocm-dash-tui" } tokio = { workspace = true } -rocm-engine-atom = { path = "../../engines/atom" } rocm-engine-lemonade = { path = "../../engines/lemonade" } -rocm-engine-llama-cpp = { path = "../../engines/llama-cpp" } -rocm-engine-pytorch = { path = "../../engines/pytorch" } rocm-engine-protocol = { path = "../../crates/rocm-engine-protocol" } -rocm-engine-sglang = { path = "../../engines/sglang" } rocmd = { path = "../rocmd" } rocm-engine-vllm = { path = "../../engines/vllm" } rpassword.workspace = true diff --git a/apps/rocm/src/main.rs b/apps/rocm/src/main.rs index d427c0a4..fa663c02 100644 --- a/apps/rocm/src/main.rs +++ b/apps/rocm/src/main.rs @@ -65,6 +65,11 @@ use std::time::Duration; static BUILTIN_ENGINE_ENV_LOCK: OnceLock> = OnceLock::new(); +/// User-selectable serving engines, in the order shown in `--help` and shell +/// completions. Used as the clap `value_parser` for every user-facing `--engine` +/// argument so the possible values stay in sync across the CLI surface. +const SUPPORTED_ENGINES: [&str; 2] = ["lemonade", "vllm"]; + #[derive(Parser, Debug)] #[command( name = "rocm", @@ -277,13 +282,13 @@ rocm model --verbose" /// logs in this terminal instead. Inspect or stop servers later with `rocm services`. #[command(after_help = "EXAMPLES:\n \ rocm serve qwen2.5-7b-instruct\n \ -rocm serve ./models/model.gguf --engine llama.cpp --port 8080\n \ +rocm serve qwen2.5-7b-instruct --engine vllm --port 8000\n \ rocm serve qwen2.5-7b-instruct --verbose --device gpu_required")] Serve { /// Model name, alias, or local model file path. model: String, - /// Engine to use [possible values: lemonade, pytorch, llama.cpp, vllm, sglang, atom]. - #[arg(long)] + /// Engine to use. + #[arg(long, value_parser = SUPPORTED_ENGINES)] engine: Option, /// Device policy [possible values: gpu_required, gpu_preferred, cpu_only]. #[arg(long)] @@ -363,7 +368,7 @@ rocm logs --search error timeout")] /// - executes enabled automation watchers (update checks, driver-plan /// checks, artifact prefetch) on a 5s tick in a sandboxed subprocess /// - health-checks and auto-recovers managed local model servers - /// (vLLM, SGLang, Lemonade, llama.cpp) + /// (Lemonade, vLLM) /// - collects GPU thermal/VRAM metrics every 60s for the TUI dashboard /// - listens on a local webhook port for automation events from other /// `rocm` commands @@ -473,9 +478,10 @@ enum EnginesCommand { /// Install the selected engine into ROCm CLI's managed engine folder. #[command(after_help = "EXAMPLES:\n \ rocm engines install lemonade\n \ -rocm engines install pytorch --reinstall")] +rocm engines install vllm --reinstall")] Install { - /// Engine name, such as lemonade, pytorch, or llama.cpp. + /// Engine name. + #[arg(value_parser = SUPPORTED_ENGINES)] engine: String, /// ROCm runtime key to install against. #[arg(long)] @@ -493,6 +499,7 @@ rocm engines install pytorch --reinstall")] /// Open a shell with the selected engine environment activated. Shell { /// Engine name. + #[arg(value_parser = SUPPORTED_ENGINES)] engine: String, /// ROCm runtime key to use. #[arg(long, conflicts_with = "env_id")] @@ -652,6 +659,7 @@ enum ConfigCommand { /// Set the preferred ROCm install for one engine. SetEngine { /// Engine name. + #[arg(value_parser = SUPPORTED_ENGINES)] engine: String, /// ROCm runtime key to use. #[arg(long, conflicts_with = "env_id")] @@ -666,6 +674,7 @@ enum ConfigCommand { /// Choose the default local model engine. SetDefaultEngine { /// Engine name. + #[arg(value_parser = SUPPORTED_ENGINES)] engine: String, }, /// Clear the saved default engine. @@ -1702,11 +1711,6 @@ fn builtin_codex_bridge_engine_inventory() -> Vec { const fn builtin_engine_inventory() -> &'static [(&'static str, &'static str)] { &[ - ("pytorch", "TheRock PyTorch local serving engine"), - ( - "llama.cpp", - "GGUF serving with ROCm GPU required by rocm-cli", - ), ( "lemonade", "default embedded Lemonade server with ROCm llama.cpp backend", @@ -1715,14 +1719,6 @@ const fn builtin_engine_inventory() -> &'static [(&'static str, &'static str)] { "vllm", "Linux/WSL ROCm GPU serving engine through external vLLM", ), - ( - "sglang", - "Linux/WSL ROCm GPU serving engine through external SGLang", - ), - ( - "atom", - "Linux/WSL ROCm GPU serving engine through external ATOM Python", - ), ] } @@ -3116,33 +3112,6 @@ fn engines(command: EnginesCommand) -> Result<()> { let mut config = RocmCliConfig::load(&paths)?; let runtime_id = resolve_engine_install_runtime_id(&paths, &config, &engine, runtime_id)?; - if engine == "pytorch" - && !reinstall - && python_version.is_none() - && let Some(manifest) = active_pytorch_runtime(&paths, &runtime_id)? - { - let engine_config = config.engine_config_mut(&engine); - engine_config.last_installed_runtime_id = Some(runtime_id.clone()); - if engine_config.preferred_runtime_id.is_none() - && engine_config.preferred_env_id.is_none() - { - engine_config.preferred_runtime_id = Some(runtime_id.clone()); - } - config.save(&paths)?; - println!("engine ready"); - println!(" engine: {engine}"); - println!(" runtime_id: {runtime_id}"); - println!(" env_path: {}", manifest.install_root.display()); - record_cli_audit_event( - &paths, - "engine", - "engine_install", - "info", - format!("ready engine={engine} runtime_id={runtime_id}"), - None, - ); - return Ok(()); - } let env_root = env_root_for_engine_install(&paths, &config, &engine, &runtime_id)?; if engine == "vllm" { ensure_openmpi_for_vllm(yes)?; @@ -3324,27 +3293,6 @@ fn env_root_for_service( } } -fn active_pytorch_runtime( - paths: &AppPaths, - runtime_id: &str, -) -> Result> { - let manifests = therock::load_runtime_manifests(paths)?; - let manifest = select_runtime_manifest(&manifests, runtime_id)?; - if !pytorch_runtime_ready(manifest) { - return Ok(None); - } - Ok(Some(manifest.clone())) -} - -fn pytorch_runtime_ready(manifest: &therock::InstalledRuntimeManifest) -> bool { - manifest.format == "wheel" - && validate_runtime_manifest_for_activation(manifest).is_ok() - && manifest - .rocm_sdk - .as_ref() - .is_some_and(|sdk| sdk.import_ok && sdk.root_path.is_some()) -} - fn managed_engine_runtime_id(engine: &str) -> &'static str { match engine { "lemonade" => "lemonade-embeddable-10.6.0", @@ -3428,7 +3376,6 @@ struct ManagedEngineEnvManifest { #[derive(Debug, Clone)] struct ResolvedEngineEnv { env_id: String, - managed_env_id: Option, runtime_id: String, python_executable: String, env_path: PathBuf, @@ -3447,11 +3394,6 @@ fn engine_shell( let paths = AppPaths::discover()?; let config = RocmCliConfig::load(&paths)?; - if engine == "llama.cpp" { - bail!( - "`rocm engines shell llama.cpp` is not available because llama.cpp uses an external llama-server binary, not a managed Python environment" - ); - } let resolved = resolve_engine_env(&paths, &config, engine, runtime_id, env_id)?; let shell_program = shell_override .map(str::to_owned) @@ -3516,7 +3458,6 @@ fn resolve_engine_env( if let Some(env_id) = selection.env_id.as_deref() { let manifest = load_engine_env_manifest(paths, engine, env_id)?; return Ok(ResolvedEngineEnv { - managed_env_id: Some(manifest.env_id.clone()), env_id: manifest.env_id, runtime_id: manifest.runtime_id, python_executable: manifest.python_executable, @@ -3530,25 +3471,6 @@ fn resolve_engine_env( let runtime_id = selection.runtime_id.with_context(|| { "no active ROCm runtime is configured; run `rocm runtimes list` and `rocm runtimes activate `, or pass --runtime-id" })?; - if engine == "pytorch" - && let Some(manifest) = active_pytorch_runtime(paths, &runtime_id)? - { - let python_executable = manifest.python_executable.clone().unwrap_or_else(|| { - runtime_python_executable_in_env(&manifest.install_root) - .display() - .to_string() - }); - return Ok(ResolvedEngineEnv { - env_id: manifest.runtime_key.clone(), - managed_env_id: None, - runtime_id, - python_executable, - env_path: manifest.install_root, - source: selection - .source - .unwrap_or_else(|| "active_therock_runtime".to_owned()), - }); - } let env_root = env_root_for_engine_install(paths, config, engine, &runtime_id)?; let response = engine_request_with_env_root::<_, InstallResponse>( Some(paths), @@ -3563,7 +3485,6 @@ fn resolve_engine_env( env_root.as_deref(), )?; Ok(ResolvedEngineEnv { - managed_env_id: Some(response.env_id.clone()), env_id: response.env_id, runtime_id, python_executable: response.python_executable, @@ -3917,22 +3838,8 @@ fn serve(args: ServeArgs) -> Result<()> { } } - let mut managed_runtime_id = resolved_selection.runtime_id.clone(); - let mut managed_env_id = resolved_selection.env_id.clone(); - if background && selected_engine == "pytorch" { - let engine_env = resolve_engine_env( - &paths, - &config, - &selected_engine, - resolved_selection.runtime_id.as_deref(), - resolved_selection.env_id.as_deref(), - )?; - if !summary_mode { - println!(" engine_env_id: {}", engine_env.env_id); - } - managed_runtime_id = Some(engine_env.runtime_id); - managed_env_id = engine_env.managed_env_id; - } + let managed_runtime_id = resolved_selection.runtime_id.clone(); + let managed_env_id = resolved_selection.env_id.clone(); if background { let mut spinner = @@ -7000,11 +6907,9 @@ fn prompt_can_use_read_only_without_local_assistant(prompt: &str) -> bool { "comfy ui", "comfy", "vllm", - "sglang", "lemonade", "llama.cpp", "llama cpp", - "pytorch", "qwen", "model server", "local server", @@ -7029,11 +6934,9 @@ fn prompt_mentions_serving_engine_or_service(normalized_prompt: &str) -> bool { normalized_prompt, &[ "vllm", - "sglang", "lemonade", "llama.cpp", "llama cpp", - "pytorch", "qwen", "model server", "local server", @@ -7335,7 +7238,7 @@ fn fallback_config_tool_call_for_prompt(normalized: &str) -> Option Option { .position(|window| window.eq_ignore_ascii_case(needle)) } -const ROCM_CHAT_TOOL_SYSTEM_PROMPT: &str = "You are ROCm CLI's local assistant. Speak in simple English for non-technical Windows 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, SGLang, Lemonade, PyTorch, llama.cpp, 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, SGLang, PyTorch, and Lemonade are 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 llama.cpp, use the llama.cpp engine backed by upstream llama-server: request rocm_command args like [\"engines\",\"install\",\"llama.cpp\"] or [\"serve\",\"MODEL.gguf\",\"--engine\",\"llama.cpp\",\"--device\",\"gpu_required\",\"--managed\"]. On native Windows, vLLM and SGLang are skipped; use WSL/Linux for those ROCm GPU engines. 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 Windows 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. On native Windows, vLLM is skipped; use WSL/Linux for that ROCm GPU engine. 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 { @@ -8932,7 +8835,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("pytorch" | "llama.cpp" | "lemonade" | "vllm" | "sglang" | "atom") + Some("lemonade" | "vllm") ) } @@ -9209,14 +9112,12 @@ fn parse_deterministic_engine_rows(tool_text: &str) -> Vec Option<&'static str> { - ["pytorch", "llama.cpp", "lemonade", "vllm", "sglang", "atom"] - .into_iter() - .find(|engine| { - line == *engine - || line - .strip_prefix(*engine) - .is_some_and(|rest| rest.starts_with(char::is_whitespace)) - }) + ["lemonade", "vllm"].into_iter().find(|engine| { + line == *engine + || line + .strip_prefix(*engine) + .is_some_and(|rest| rest.starts_with(char::is_whitespace)) + }) } fn should_request_local_tool_follow_up( @@ -11093,14 +10994,6 @@ const fn model_registry_adapter_availability_note(engine: &str) -> Option<&'stat 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 if rocm_core::runtime_is_windows() && engine.eq_ignore_ascii_case("sglang") { - Some( - "runtime_status=unsupported_native_windows reason=native Windows skipped; use WSL/Linux SGLang ROCm; gpu_execution_required=true; run /engine for adapter details", - ) - } else if rocm_core::runtime_is_windows() && engine.eq_ignore_ascii_case("atom") { - Some( - "runtime_status=unsupported_native_windows reason=use WSL/Linux ATOM ROCm; gpu_execution_required=true; run /engine for adapter details", - ) } else { None } @@ -11212,21 +11105,6 @@ fn friendly_engine_detect_notes(engine: &str, notes: &[String]) -> Option Option Option String { let lower = note.to_ascii_lowercase(); - if lower.contains("torch probe failed") || lower.contains("torch probe import failed") { - return "PyTorch engine is not ready yet; reinstall the PyTorch engine from the TUI or run `rocm engines install pytorch`.".to_owned(); - } - if lower.contains("no managed pytorch envs found") { - return "PyTorch engine is not installed yet.".to_owned(); - } - if lower.contains("torch probe:") - || lower.contains("rocm_sdk:") - || lower.contains("torch._rocm_init") - { - return "PyTorch is ready on your AMD GPU.".to_owned(); - } if rocm_core::runtime_is_windows() && (lower.contains("unsupported_native_windows") || lower.contains("native windows") @@ -11303,11 +11156,7 @@ fn engine_runtime_status_label(engine: &str, detect: &DetectResponse) -> &'stati } fn engine_runtime_is_native_windows_unsupported(engine: &str, detect: &DetectResponse) -> bool { - if !rocm_core::runtime_is_windows() - || !(engine.eq_ignore_ascii_case("vllm") - || engine.eq_ignore_ascii_case("sglang") - || engine.eq_ignore_ascii_case("atom")) - { + if !rocm_core::runtime_is_windows() || !engine.eq_ignore_ascii_case("vllm") { return false; } @@ -12784,11 +12633,7 @@ pub(crate) fn render_daemon_text(paths: &AppPaths, config: &RocmCliConfig) -> St fn friendly_engine_label(engine: &str) -> &str { match engine { "lemonade" => "Lemonade", - "pytorch" => "PyTorch", - "llama.cpp" => "llama.cpp", "vllm" => "vLLM", - "sglang" => "SGLang", - "atom" => "ATOM", other => other, } } @@ -14219,8 +14064,6 @@ fn remove_path(path: &Path) -> Result<()> { pub(crate) const fn engine_inventory() -> &'static [(&'static str, &'static str)] { &[ - ("pytorch", "TheRock PyTorch local serving engine"), - ("llama.cpp", "external GGUF serving engine for llama-server"), ( "lemonade", "default embedded Lemonade server with ROCm llama.cpp backend", @@ -14229,27 +14072,13 @@ pub(crate) const fn engine_inventory() -> &'static [(&'static str, &'static str) "vllm", "Linux/WSL ROCm GPU serving engine through external vLLM", ), - ( - "sglang", - "Linux/WSL ROCm GPU serving engine through external SGLang", - ), - ( - "atom", - "Linux/WSL ROCm GPU serving engine through external ATOM Python", - ), ] } -fn infer_engine_from_request(lower: &str) -> Option<&str> { - for engine in ["pytorch", "llama.cpp", "lemonade", "vllm", "sglang", "atom"] { - if lower.contains(engine) { - return Some(engine); - } - } - if lower.contains("llama cpp") { - return Some("llama.cpp"); - } - None +fn infer_engine_from_request(lower: &str) -> Option<&'static str> { + ["lemonade", "vllm"] + .into_iter() + .find(|engine| lower.contains(*engine)) } fn infer_model_from_request(request: &str) -> Option<&str> { @@ -14522,19 +14351,9 @@ fn builtin_engine_request( envelope: &EngineRequestEnvelope, ) -> Option { match engine { - "atom" => Some(rocm_engine_atom::builtin_handle_envelope(envelope.clone())), "lemonade" => Some(rocm_engine_lemonade::builtin_handle_envelope( envelope.clone(), )), - "llama.cpp" => Some(rocm_engine_llama_cpp::builtin_handle_envelope( - envelope.clone(), - )), - "pytorch" => Some(rocm_engine_pytorch::builtin_handle_envelope( - envelope.clone(), - )), - "sglang" => Some(rocm_engine_sglang::builtin_handle_envelope( - envelope.clone(), - )), "vllm" => Some(rocm_engine_vllm::builtin_handle_envelope(envelope.clone())), _ => None, } @@ -14554,10 +14373,7 @@ fn run_builtin_engine_stdio(engine: &str) -> Result<()> { } fn builtin_engine_available(engine: &str) -> bool { - matches!( - engine, - "atom" | "lemonade" | "llama.cpp" | "pytorch" | "sglang" | "vllm" - ) + matches!(engine, "lemonade" | "vllm") } #[allow(clippy::too_many_arguments)] @@ -14577,18 +14393,6 @@ fn run_builtin_engine_serve_http( ) -> Result<()> { let parsed_policy = parse_device_policy(Some(device_policy))?; match engine { - "atom" => rocm_engine_atom::builtin_serve_http( - service_id, - model_ref, - host, - port, - parsed_policy, - gpu_indices, - runtime_id, - env_id, - state_path, - engine_recipe, - ), "lemonade" => rocm_engine_lemonade::builtin_serve_http( service_id, model_ref, @@ -14602,43 +14406,6 @@ fn run_builtin_engine_serve_http( log_path, engine_recipe, ), - "llama.cpp" => rocm_engine_llama_cpp::builtin_serve_http( - service_id, - model_ref, - host, - port, - Some(device_policy_name(&parsed_policy).to_owned()), - gpu_indices, - runtime_id, - env_id, - state_path, - log_path, - engine_recipe, - ), - "pytorch" => rocm_engine_pytorch::builtin_serve_http( - service_id, - model_ref, - host, - port, - parsed_policy, - gpu_indices, - env_id, - runtime_id, - state_path, - engine_recipe, - ), - "sglang" => rocm_engine_sglang::builtin_serve_http( - service_id, - model_ref, - host, - port, - parsed_policy, - gpu_indices, - runtime_id, - env_id, - state_path, - engine_recipe, - ), "vllm" => rocm_engine_vllm::builtin_serve_http( service_id, model_ref, @@ -15271,8 +15038,6 @@ fn wait_for_service_http_ready_with_progress( fn service_http_readiness_paths(engine: &str) -> &'static [&'static str] { match engine { "lemonade" => &["/v1/health", "/v1/models"], - "llama.cpp" => &["/v1/models", "/health"], - "pytorch" => &["/v1/models", "/healthz"], _ => &["/v1/models", "/v1/health", "/health", "/healthz"], } } @@ -15514,15 +15279,27 @@ mod tests { .unwrap_or_default() } - // These guard against the hand-written `[possible values: ...]` help on the - // free-form `--engine`/`--device` flags drifting from their real source of - // truth. The flags stay `Option` on purpose (`--engine` is - // free-form and `--device` accepts aliases such as `auto`/`gpu` plus the - // intentional `cpu_only` rejection), so a ValueEnum would change accepted - // input; a sync test keeps the advertised list honest without that. + // `--engine` restricts its input to `SUPPORTED_ENGINES` via a clap + // `value_parser`, so the possible values are advertised in `--help` and shell + // completion structurally (not a hand-written doc string). `--device` stays + // free-form (it accepts aliases such as `auto`/`gpu` plus the intentional + // `cpu_only` rejection), so its advertised list is hand-written and guarded + // by a sync test below. Both guards keep the advertised lists honest. #[test] fn serve_engine_help_lists_match_engine_inventory() { - let mut listed = possible_values_listed_in_help(&serve_arg_help("engine")); + let cli = Cli::command(); + let serve = cli + .find_subcommand("serve") + .expect("serve subcommand exists"); + let engine_arg = serve + .get_arguments() + .find(|arg| arg.get_id() == "engine") + .expect("serve has an `engine` argument"); + let mut listed: Vec = engine_arg + .get_possible_values() + .iter() + .map(|value| value.get_name().to_owned()) + .collect(); let mut expected: Vec = builtin_engine_inventory() .iter() .map(|(name, _)| (*name).to_owned()) @@ -15531,7 +15308,7 @@ mod tests { expected.sort(); assert_eq!( listed, expected, - "serve --engine help possible-values must stay in sync with builtin_engine_inventory()" + "serve --engine possible-values must stay in sync with builtin_engine_inventory()" ); } @@ -15766,7 +15543,7 @@ mod tests { fn service_http_readiness_requires_model_list_entry() { let empty = json!({ "data": [] }).to_string(); assert!(!service_http_readiness_response_ready( - "llama.cpp", + "vllm", "/v1/models", 200, &empty, @@ -15775,7 +15552,7 @@ mod tests { let models = json!({ "data": [{ "id": "tiny.gguf" }] }).to_string(); assert!(service_http_readiness_response_ready( - "llama.cpp", + "vllm", "/v1/models", 200, &models, @@ -15813,14 +15590,14 @@ mod tests { )); assert!(!service_http_readiness_response_ready( - "llama.cpp", + "vllm", "/health", 200, "OK", "tiny.gguf" )); assert!(!service_http_readiness_response_ready( - "pytorch", + "vllm", "/healthz", 200, "OK", @@ -15900,7 +15677,7 @@ mod tests { cpu: Some("AMD Ryzen".to_owned()), system_ram_gib: Some(64.0), interactive_terminal: false, - default_engine: "pytorch".to_owned(), + default_engine: "vllm".to_owned(), detected_gfx_target: Some("gfx1201".to_owned()), compatible_therock_family: Some("gfx120X-all".to_owned()), detected_therock_family: None, @@ -16013,7 +15790,7 @@ mod tests { recipe.canonical_model_id = "Acme/SignedTiny".to_owned(); recipe.aliases = vec!["signedtiny".to_owned()]; recipe.source = "signed_recipe_index".to_owned(); - recipe.preferred_engines = vec!["llama.cpp".to_owned()]; + recipe.preferred_engines = vec!["vllm".to_owned()]; recipe.device_policy = "cpu_only".to_owned(); recipe.dtype = "float16".to_owned(); @@ -16191,7 +15968,7 @@ mod tests { #[test] fn freeform_plan_next_action_surfaces_approval_action() { let action = - freeform_plan_next_action("serve qwen3.5 with llama.cpp", &RocmCliConfig::default()) + freeform_plan_next_action("serve qwen3.5 with vllm", &RocmCliConfig::default()) .expect("serve request should have next action"); assert_eq!(action.title, "Launch local endpoint"); @@ -16203,7 +15980,7 @@ mod tests { "serve".to_owned(), "Qwen/Qwen3.5-4B".to_owned(), "--engine".to_owned(), - "llama.cpp".to_owned(), + "vllm".to_owned(), "--device".to_owned(), "gpu_required".to_owned(), "--managed".to_owned(), @@ -16219,7 +15996,7 @@ mod tests { "serve".to_owned(), "qwen3.5".to_owned(), "with".to_owned(), - "llama.cpp".to_owned(), + "vllm".to_owned(), ]); assert!(invocation.approve); @@ -16231,7 +16008,7 @@ mod tests { "serve".to_owned(), "qwen3.5".to_owned(), "with".to_owned(), - "llama.cpp".to_owned(), + "vllm".to_owned(), ] ); @@ -16259,7 +16036,7 @@ mod tests { "serve".to_owned(), "qwen3.5".to_owned(), "with".to_owned(), - "llama.cpp".to_owned(), + "vllm".to_owned(), ]); assert!(!should_treat_as_freeform(&invalid_install)); @@ -16292,13 +16069,13 @@ mod tests { #[test] fn freeform_execution_validation_accepts_fully_structured_tool_call() -> Result<()> { let action = - freeform_plan_next_action("serve qwen3.5 with llama.cpp", &RocmCliConfig::default()) + freeform_plan_next_action("serve qwen3.5 with vllm", &RocmCliConfig::default()) .expect("serve request should have next action"); validate_freeform_execution_action(&action)?; assert_eq!( format_structured_tool_call("rocm", &action.args), - "rocm serve Qwen/Qwen3.5-4B --engine llama.cpp --device gpu_required --managed" + "rocm serve Qwen/Qwen3.5-4B --engine vllm --device gpu_required --managed" ); Ok(()) } @@ -16306,14 +16083,14 @@ mod tests { #[test] fn freeform_execution_header_surfaces_explicit_approval_and_tool_call() { let action = - freeform_plan_next_action("serve qwen3.5 with llama.cpp", &RocmCliConfig::default()) + freeform_plan_next_action("serve qwen3.5 with vllm", &RocmCliConfig::default()) .expect("serve request should have next action"); let rendered = render_freeform_execution_header(&action); assert!(rendered.contains("execution")); assert!(rendered.contains("approval: granted by --yes")); assert!(rendered.contains( - "tool_call: rocm serve Qwen/Qwen3.5-4B --engine llama.cpp --device gpu_required --managed" + "tool_call: rocm serve Qwen/Qwen3.5-4B --engine vllm --device gpu_required --managed" )); } @@ -16486,7 +16263,7 @@ mod tests { "confidence": "high", "tool_call": { "tool": "rocm", - "args": ["serve", "sshleifer/tiny-gpt2", "--engine", "pytorch", "--device", "gpu_required", "--managed"] + "args": ["serve", "sshleifer/tiny-gpt2", "--engine", "vllm", "--device", "gpu_required", "--managed"] }, "notes": ["resolved the missing model to a tiny test model"] }"#; @@ -16504,7 +16281,7 @@ mod tests { "serve".to_owned(), "sshleifer/tiny-gpt2".to_owned(), "--engine".to_owned(), - "pytorch".to_owned(), + "vllm".to_owned(), "--device".to_owned(), "gpu_required".to_owned(), "--managed".to_owned(), @@ -16525,14 +16302,14 @@ mod tests { "intent": "serve", "tool_call": { "tool": "rocm", - "args": ["serve", "tiny.gguf", "--engine", "llama.cpp", "--allow-public-bind", "--managed"] + "args": ["serve", "tiny.gguf", "--engine", "vllm", "--allow-public-bind", "--managed"] } }"#, r#"{ "intent": "serve", "tool_call": { "tool": "rocm", - "args": ["serve", "tiny.gguf", "--engine", "llama.cpp", "--host", "0.0.0.0", "--managed"] + "args": ["serve", "tiny.gguf", "--engine", "vllm", "--host", "0.0.0.0", "--managed"] } }"#, ] { @@ -16550,8 +16327,8 @@ mod tests { #[test] fn provider_planner_requires_managed_serve_requests() { for args in [ - vec!["serve", "qwen", "--engine", "pytorch"], - vec!["serve", "qwen", "--engine", "pytorch", "--foreground"], + vec!["serve", "qwen", "--engine", "vllm"], + vec!["serve", "qwen", "--engine", "vllm", "--foreground"], ] { let call = ProviderPlannerToolCall { tool: "rocm".to_owned(), @@ -16592,7 +16369,7 @@ mod tests { "serve", "sshleifer/tiny-gpt2", "--engine", - "pytorch", + "vllm", "--device", "cpu", "--managed", @@ -16601,7 +16378,7 @@ mod tests { "serve", "sshleifer/tiny-gpt2", "--engine", - "pytorch", + "vllm", "--device=cpu", "--managed", ], @@ -16609,7 +16386,7 @@ mod tests { "serve", "sshleifer/tiny-gpt2", "--engine", - "pytorch", + "vllm", "--device", "cpu_only", "--managed", @@ -16897,7 +16674,6 @@ mod tests { "port_status", "[\"services\",\"list\",\"--all\"]", "qwen-smoke", - "llama-server", "Do not invent shell commands", "ROCm CLI Assistant Skill", "Treat `localhost` and `127.0.0.1` as the same loopback endpoint", @@ -16995,16 +16771,16 @@ model recipes engine_support: lemonade: available path=D:\\rocm\\rocm-engine-lemonade.exe warning: tiny Lemonade GGUF smoke-test model; not the default assistant - Qwen/Qwen2.5-0.5B-Instruct aliases=[qwen-tiny] task=chat dtype=float16 device=gpu_required min_gpu_mem=4 GiB engines=[pytorch] + Qwen/Qwen2.5-0.5B-Instruct aliases=[qwen-tiny] task=chat dtype=float16 device=gpu_required min_gpu_mem=4 GiB engines=[lemonade] engine_support: - pytorch: available path=D:\\rocm\\rocm-engine-pytorch.exe + lemonade: available path=D:\\rocm\\rocm-engine-lemonade.exe Qwen/Qwen3.5-4B aliases=[qwen3.5] task=chat dtype=bfloat16 device=gpu_preferred min_gpu_mem=12 GiB engines=[vllm] engine_support: vllm: adapter_available path=D:\\rocm\\rocm-engine-vllm.exe runtime_status=unsupported_native_windows reason=native Windows skipped; use WSL/Linux vLLM ROCm - meta-llama/Llama-3.2-3B-Instruct aliases=[llama] task=chat dtype=bfloat16 device=gpu_preferred min_gpu_mem=8 GiB engines=[pytorch, llama.cpp] + meta-llama/Llama-3.2-3B-Instruct aliases=[llama] task=chat dtype=bfloat16 device=gpu_preferred min_gpu_mem=8 GiB engines=[lemonade, vllm] engine_support: - pytorch: available path=D:\\rocm\\rocm-engine-pytorch.exe - llama.cpp: available path=D:\\rocm\\rocm-engine-llama-cpp.exe + lemonade: available path=D:\\rocm\\rocm-engine-lemonade.exe + vllm: available path=D:\\rocm\\rocm-engine-vllm.exe ", ) .expect("model output should summarize"); @@ -17015,7 +16791,7 @@ model recipes assert!(summary.contains("Tiny smoke test: qwen-smoke")); assert!(summary.contains("Qwen3-0.6B-GGUF")); assert!(summary.contains("8 GiB-class option: llama")); - assert!(summary.contains("pytorch, llama.cpp")); + assert!(summary.contains("lemonade, vllm")); assert!(summary.contains("Qwen/Qwen3.5-4B asks for 12 GiB")); assert!(summary.contains("Native Windows note")); assert!(summary.contains("Run `rocm examine`")); @@ -17107,11 +16883,11 @@ model recipes name: "launch_server".to_owned(), arguments: serde_json::json!({ "model": "qwen", - "engine": "pytorch", + "engine": "vllm", "device": "gpu_required" }), }, - Some("rocm serve qwen --managed --engine pytorch --device gpu_required"), + Some("rocm serve qwen --managed --engine vllm --device gpu_required"), false, ), ] { @@ -17132,7 +16908,7 @@ model recipes } #[test] - fn chat_rocm_command_routes_comfyui_and_llama_cpp_actions() { + fn chat_rocm_command_routes_comfyui_and_engine_actions() { let comfy_install = providers::ChatToolCall { id: Some("call-comfy".to_owned()), name: "rocm_command".to_owned(), @@ -17156,21 +16932,21 @@ model recipes vec!["comfyui".to_owned(), "install".to_owned()] ); - let llama = providers::ChatToolCall { - id: Some("call-llama".to_owned()), + let lemonade = providers::ChatToolCall { + id: Some("call-lemonade".to_owned()), name: "rocm_command".to_owned(), arguments: serde_json::json!({ - "args": ["engines", "install", "llama.cpp"] + "args": ["engines", "install", "lemonade"] }), }; - validate_chat_tool_call(&llama).expect("llama.cpp engine install should validate"); - assert!(!chat_tool_call_is_read_only(&llama)); + validate_chat_tool_call(&lemonade).expect("lemonade engine install should validate"); + assert!(!chat_tool_call_is_read_only(&lemonade)); assert_eq!( - rocm_chat_tool_requested_command(&llama).as_deref(), - Some("rocm engines install llama.cpp") + rocm_chat_tool_requested_command(&lemonade).as_deref(), + Some("rocm engines install lemonade") ); let approval = - chat_tool_approval_request(&llama, Some("Install llama.cpp for GGUF serving.")) + chat_tool_approval_request(&lemonade, Some("Install Lemonade for local serving.")) .expect("approval should be built"); assert_eq!(approval.pending_title, "Install engine"); assert_eq!(approval.command_title, "Engine"); @@ -17211,14 +16987,14 @@ model recipes id: Some("call-serve".to_owned()), name: "rocm_command".to_owned(), arguments: serde_json::json!({ - "args": ["serve", "qwen", "--engine", "pytorch", "--device", "gpu_required", "--managed"] + "args": ["serve", "qwen", "--engine", "vllm", "--device", "gpu_required", "--managed"] }), }; validate_chat_tool_call(&serve).expect("managed serve should validate"); assert!(!chat_tool_call_is_read_only(&serve)); assert_eq!( rocm_chat_tool_requested_command(&serve).as_deref(), - Some("rocm serve qwen --engine pytorch --device gpu_required --managed") + Some("rocm serve qwen --engine vllm --device gpu_required --managed") ); let approval = chat_tool_approval_request(&serve, Some("Start the recommended assistant.")) .expect("approval should be built"); @@ -17243,14 +17019,13 @@ model recipes id: Some("call-config".to_owned()), name: "rocm_command".to_owned(), arguments: serde_json::json!({ - "args": ["config", "set-default-engine", "pytorch"] + "args": ["config", "set-default-engine", "vllm"] }), }; validate_chat_tool_call(&config).expect("config change should validate"); assert!(!chat_tool_call_is_read_only(&config)); - let approval = - chat_tool_approval_request(&config, Some("Use PyTorch as the default engine.")) - .expect("approval should be built"); + let approval = chat_tool_approval_request(&config, Some("Use vLLM as the default engine.")) + .expect("approval should be built"); assert_eq!(approval.pending_title, "Change settings"); assert_eq!(approval.command_title, "Config"); } @@ -17610,7 +17385,7 @@ model recipes id: None, name: "rocm_command".to_owned(), arguments: serde_json::json!({ - "args": ["serve", "tiny.gguf", "--engine", "llama.cpp", "--device", "cpu"] + "args": ["serve", "tiny.gguf", "--engine", "vllm", "--device", "cpu"] }), }; let error = validate_chat_tool_call(&cpu).unwrap_err().to_string(); @@ -17620,7 +17395,7 @@ model recipes id: None, name: "rocm_command".to_owned(), arguments: serde_json::json!({ - "args": ["serve", "tiny.gguf", "--engine", "llama.cpp", "--allow-public-bind", "--managed"] + "args": ["serve", "tiny.gguf", "--engine", "vllm", "--allow-public-bind", "--managed"] }), }; let error = validate_chat_tool_call(&public_flag) @@ -17632,7 +17407,7 @@ model recipes id: None, name: "rocm_command".to_owned(), arguments: serde_json::json!({ - "args": ["serve", "qwen", "--engine", "pytorch", "--foreground"] + "args": ["serve", "qwen", "--engine", "vllm", "--foreground"] }), }; let error = validate_chat_tool_call(&foreground) @@ -17833,7 +17608,7 @@ model recipes name: "launch_server".to_owned(), arguments: serde_json::json!({ "model": "tiny.gguf", - "engine": "llama.cpp", + "engine": "vllm", "device": "cpu" }), }; @@ -17845,7 +17620,7 @@ model recipes name: "launch_server".to_owned(), arguments: serde_json::json!({ "model": "tiny.gguf", - "engine": "llama.cpp", + "engine": "vllm", "host": "0.0.0.0", "allow_public_bind": true }), @@ -17858,7 +17633,7 @@ model recipes name: "launch_server".to_owned(), arguments: serde_json::json!({ "model": "tiny.gguf", - "engine": "llama.cpp", + "engine": "vllm", "host": "0.0.0.0" }), }; @@ -18023,7 +17798,7 @@ model recipes for prompt in [ "Is vLLM running?", - "is sglang running?", + "is lemonade running?", "is the model server running?", "is qwen running?", ] { @@ -18048,7 +17823,7 @@ model recipes #[test] fn fallback_tool_call_routes_engine_install_state_to_engines_list() { - for prompt in ["is vLLM installed?", "is SGLang available?"] { + for prompt in ["is vLLM installed?", "is Lemonade available?"] { let call = fallback_rocm_tool_call_for_prompt(prompt).unwrap(); assert_eq!(call.name, "rocm_command", "{prompt}"); assert_eq!( @@ -18472,14 +18247,13 @@ install therock"; ); assert!(chat_tool_call_is_read_only(&show)); - let engine = - fallback_rocm_tool_call_for_prompt("Set the default engine to pytorch").unwrap(); + let engine = fallback_rocm_tool_call_for_prompt("Set the default engine to vllm").unwrap(); assert_eq!( normalized_chat_rocm_command_args(&engine).unwrap(), vec![ "config".to_owned(), "set-default-engine".to_owned(), - "pytorch".to_owned(), + "vllm".to_owned(), ] ); assert!(!chat_tool_call_is_read_only(&engine)); @@ -18754,7 +18528,7 @@ install therock"; "intent": "serve", "tool_call": { "tool": "rocm", - "args": ["serve", "sshleifer/tiny-gpt2", "--engine", "pytorch", "--managed"] + "args": ["serve", "sshleifer/tiny-gpt2", "--engine", "vllm", "--managed"] } }"#; let plan = provider_planner_response_to_plan("start a local model", "local", content)?; @@ -18957,7 +18731,7 @@ install therock"; let mut record = ManagedServiceRecord::new( &paths, "svc_qwen35_primary", - "pytorch", + "vllm", "qwen3.5", "Qwen/Qwen3.5", "127.0.0.1", @@ -18980,7 +18754,7 @@ install therock"; let rendered = render_service_logs_text(&paths, "svc_qwen35_primary")?; assert!(rendered.contains("Service Log")); assert!(rendered.contains("Service: svc_qwen35_primary")); - assert!(rendered.contains("Engine: pytorch")); + assert!(rendered.contains("Engine: vllm")); assert!(rendered.contains("Status: starting")); assert!(rendered.contains("File locations: shown")); assert!(rendered.contains(&format!( @@ -19029,7 +18803,7 @@ install therock"; let mut record = ManagedServiceRecord::new( &paths, service_id, - "pytorch", + "vllm", "qwen", "Qwen/Qwen3.5", "127.0.0.1", @@ -19596,15 +19370,18 @@ install therock"; #[test] fn explicit_engine_override_keeps_alias_when_shared_recipe_is_for_another_engine() { - let recipe = resolve_builtin_model_recipe("qwen").expect("qwen recipe"); + // `qwen-smoke` is a Lemonade-only GGUF recipe (no vLLM engine recipe). + let recipe = resolve_builtin_model_recipe("qwen-smoke").expect("qwen-smoke recipe"); + // Served under the engine it targets, the alias resolves to the canonical id. assert_eq!( - serve_model_ref_for_engine("qwen", Some(&recipe), "lemonade"), - "Qwen3-4B-Instruct-2507-GGUF" + serve_model_ref_for_engine("qwen-smoke", Some(&recipe), "lemonade"), + "Qwen3-0.6B-GGUF" ); + // Under an engine the recipe does not support, the raw alias flows through unchanged. assert_eq!( - serve_model_ref_for_engine("qwen", Some(&recipe), "pytorch"), - "qwen" + serve_model_ref_for_engine("qwen-smoke", Some(&recipe), "vllm"), + "qwen-smoke" ); } @@ -19612,20 +19389,20 @@ install therock"; fn serve_engine_selection_respects_explicit_and_configured_engines() { let recipe = resolve_builtin_model_recipe("qwen32b").expect("qwen32b recipe"); - let explicit = select_serve_engine(Some("llama.cpp"), Some("pytorch"), Some(&recipe), None); - let configured = select_serve_engine(None, Some("pytorch"), Some(&recipe), None); + let explicit = select_serve_engine(Some("vllm"), Some("lemonade"), Some(&recipe), None); + let configured = select_serve_engine(None, Some("lemonade"), Some(&recipe), None); assert_eq!( explicit, ServeEngineSelection { - engine: "llama.cpp".to_owned(), + engine: "vllm".to_owned(), source: "explicit --engine", } ); assert_eq!( configured, ServeEngineSelection { - engine: "pytorch".to_owned(), + engine: "lemonade".to_owned(), source: "configured default_engine", } ); @@ -19656,7 +19433,7 @@ install therock"; model_id_override: None, }, rocm_core::ModelRecipeEngineRecord { - engine: "sglang".to_owned(), + engine: "lemonade".to_owned(), required_flags: vec!["--reasoning-parser".to_owned(), "qwen3".to_owned()], parser_settings: BTreeMap::new(), preferred_endpoint: None, @@ -19700,7 +19477,7 @@ install therock"; "engine_recipe_policy: selected-engine required_flags are applied at launch" )); assert!(serve_lines.contains("engine_recipe_required_flags: --enable-auto-tool-choice")); - assert!(protocol_engine_recipe_hint(&recipe, "pytorch").is_none()); + assert!(protocol_engine_recipe_hint(&recipe, "unknown-engine").is_none()); } #[test] @@ -20163,7 +19940,7 @@ VERSION_ID="41" ..RocmCliConfig::default() }; - let selection = resolve_engine_selection(&config, "pytorch", None, None); + let selection = resolve_engine_selection(&config, "vllm", None, None); assert_eq!( selection.runtime_id.as_deref(), Some("therock-release:gfx120X-all") @@ -20174,7 +19951,7 @@ VERSION_ID="41" ); config.active_runtime_key = Some("release-pip-gfx120x-all-7-13-0".to_owned()); - let selection = resolve_engine_selection(&config, "pytorch", None, None); + let selection = resolve_engine_selection(&config, "vllm", None, None); assert_eq!( selection.runtime_id.as_deref(), Some("release-pip-gfx120x-all-7-13-0") @@ -20184,9 +19961,9 @@ VERSION_ID="41" Some("config_active_runtime_key") ); - config.engine_config_mut("pytorch").preferred_runtime_id = + config.engine_config_mut("vllm").preferred_runtime_id = Some("therock-nightly:gfx120X-all".to_owned()); - let selection = resolve_engine_selection(&config, "pytorch", None, None); + let selection = resolve_engine_selection(&config, "vllm", None, None); assert_eq!( selection.runtime_id.as_deref(), Some("release-pip-gfx120x-all-7-13-0") @@ -20197,7 +19974,7 @@ VERSION_ID="41" ); config.active_runtime_key = None; - let selection = resolve_engine_selection(&config, "pytorch", None, None); + let selection = resolve_engine_selection(&config, "vllm", None, None); assert_eq!( selection.runtime_id.as_deref(), Some("therock-nightly:gfx120X-all") @@ -20220,7 +19997,7 @@ VERSION_ID="41" )?; let selection = validate_engine_selection_runtime( &paths, - resolve_engine_selection(&RocmCliConfig::default(), "pytorch", None, None), + resolve_engine_selection(&RocmCliConfig::default(), "vllm", None, None), )?; assert_eq!( @@ -20251,7 +20028,7 @@ VERSION_ID="41" )?; let selection = validate_engine_selection_runtime( &paths, - resolve_engine_selection(&RocmCliConfig::default(), "pytorch", None, None), + resolve_engine_selection(&RocmCliConfig::default(), "vllm", None, None), )?; assert!(selection.runtime_id.is_none()); @@ -20302,7 +20079,7 @@ VERSION_ID="41" fn engine_install_runtime_selection_requires_configured_runtime() -> Result<()> { let (root, paths) = test_paths("engine-install-runtime-selection"); let error = - resolve_engine_install_runtime_id(&paths, &RocmCliConfig::default(), "pytorch", None) + resolve_engine_install_runtime_id(&paths, &RocmCliConfig::default(), "vllm", None) .unwrap_err() .to_string(); assert!(error.contains("no active ROCm runtime is configured")); @@ -20323,14 +20100,14 @@ VERSION_ID="41" ..RocmCliConfig::default() }; assert_eq!( - resolve_engine_install_runtime_id(&paths, &config, "pytorch", None)?, + resolve_engine_install_runtime_id(&paths, &config, "vllm", None)?, "release-pip-gfx120x-all" ); assert_eq!( resolve_engine_install_runtime_id( &paths, &config, - "pytorch", + "vllm", Some("therock-release:gfx120X-all".to_owned()) )?, "release-pip-gfx120x-all" @@ -20446,7 +20223,7 @@ VERSION_ID="41" 1, )?; - let engine_root = env_root_for_runtime(&paths, "pytorch", &manifest.runtime_key)?; + let engine_root = env_root_for_runtime(&paths, "vllm", &manifest.runtime_key)?; assert_eq!(engine_root, Some(manifest.install_root.join("engines"))); assert_eq!( @@ -20480,31 +20257,6 @@ VERSION_ID="41" Ok(()) } - #[test] - fn resolve_pytorch_env_uses_active_therock_runtime() -> Result<()> { - let (root, paths) = test_paths("pytorch-active-runtime-env"); - let manifest = write_test_pip_runtime( - &paths, - "release-pip-gfx120x-all", - "therock-release:gfx120X-all", - "7.13.0", - 1, - )?; - let config = RocmCliConfig { - active_runtime_key: Some(manifest.runtime_key.clone()), - ..RocmCliConfig::default() - }; - - let resolved = resolve_engine_env(&paths, &config, "pytorch", None, None)?; - - assert_eq!(resolved.runtime_id, manifest.runtime_key); - assert_eq!(resolved.env_path, manifest.install_root); - assert_eq!(resolved.managed_env_id, None); - assert!(!resolved.env_path.join("engines").exists()); - let _ = fs::remove_dir_all(root); - Ok(()) - } - #[test] fn engine_runtime_selection_rejects_ambiguous_default_runtime_id() -> Result<()> { let (root, paths) = test_paths("engine-runtime-ambiguous-default"); @@ -20527,20 +20279,20 @@ VERSION_ID="41" ..RocmCliConfig::default() }; - let error = resolve_engine_install_runtime_id(&paths, &config, "pytorch", None) + let error = resolve_engine_install_runtime_id(&paths, &config, "vllm", None) .unwrap_err() .to_string(); assert!(error.contains("matches multiple installed runtimes")); assert!(error.contains("rocm runtimes activate ")); - let selection = resolve_engine_selection(&config, "pytorch", None, None); + let selection = resolve_engine_selection(&config, "vllm", None, None); let error = validate_engine_selection_runtime(&paths, selection) .unwrap_err() .to_string(); assert!(error.contains("matches multiple installed runtimes")); let selection = - resolve_engine_selection(&config, "pytorch", Some("release-pip-gfx120x-all"), None); + resolve_engine_selection(&config, "vllm", Some("release-pip-gfx120x-all"), None); let selection = validate_engine_selection_runtime(&paths, selection)?; assert_eq!( selection.runtime_id.as_deref(), @@ -21260,11 +21012,11 @@ VERSION_ID="41" let plugin_dir = paths.primary_engine_plugin_dir(); fs::create_dir_all(&plugin_dir)?; let plugin_path = plugin_dir.join( - rocm_engine_protocol::platform_engine_plugin_binary_name("pytorch"), + rocm_engine_protocol::platform_engine_plugin_binary_name("vllm"), ); fs::write(&plugin_path, "plugin")?; - let discovered = find_engine_plugin_binary("pytorch", engine_plugin_dirs(&paths))?; + let discovered = find_engine_plugin_binary("vllm", engine_plugin_dirs(&paths))?; let _ = fs::remove_dir_all(root); assert_eq!(discovered, Some(plugin_path)); @@ -21278,12 +21030,12 @@ VERSION_ID="41" let compatibility_dir = paths.data_dir.join("engines"); fs::create_dir_all(&primary_dir)?; fs::create_dir_all(&compatibility_dir)?; - let name = rocm_engine_protocol::platform_engine_plugin_binary_name("pytorch"); + let name = rocm_engine_protocol::platform_engine_plugin_binary_name("vllm"); let primary_path = primary_dir.join(&name); fs::write(&primary_path, "primary")?; fs::write(compatibility_dir.join(&name), "compatibility")?; - let discovered = find_engine_plugin_binary("pytorch", engine_plugin_dirs(&paths))?; + let discovered = find_engine_plugin_binary("vllm", engine_plugin_dirs(&paths))?; let _ = fs::remove_dir_all(root); assert_eq!(discovered, Some(primary_path)); @@ -21306,13 +21058,6 @@ VERSION_ID="41" #[test] fn friendly_engine_detect_notes_hide_probe_and_path_noise() { - let pytorch = friendly_engine_detect_notes( - "pytorch", - &[r"torch probe: cuda_available=true device_count=1; managed env detected at C:\Users\user\.rocm\engines\pytorch\envs\release; rocm_sdk: version=7.13".to_owned()], - ) - .expect("pytorch note"); - assert_eq!(pytorch, "PyTorch is ready on your AMD GPU."); - let lemonade = friendly_engine_detect_notes( "lemonade", &["Lemonade embeddable 10.6.0 is installed at D:/ROCm/temp/runtime; Lemonade is configured for llamacpp:rocm; no CPU fallback is used".to_owned()], @@ -21320,12 +21065,6 @@ VERSION_ID="41" .expect("lemonade note"); assert_eq!(lemonade, "Lemonade is ready on your AMD GPU."); - let llama = friendly_engine_detect_notes( - "llama.cpp", - &["llama-server not found; TheRock HIP runtime env available: root=D:\\ROCm\\temp\\therock".to_owned()], - ) - .expect("llama.cpp note"); - assert_eq!(llama, "llama.cpp server is not installed yet."); let vllm = friendly_engine_detect_notes( "vllm", &["vLLM is not installed in a Linux/WSL ROCm Python environment. Native Windows is skipped; no CPU fallback is used.".to_owned()], @@ -21335,22 +21074,14 @@ VERSION_ID="41" vllm, "vLLM is not installed in a Linux/WSL ROCm Python environment." ); - let atom = friendly_engine_detect_note_fallback( - "ATOM Python environment was not found; set ROCM_CLI_ATOM_PYTHON.", - ); - assert!(atom.contains("ROCM_CLI_ATOM_PYTHON")); - assert!(!pytorch.contains("torch probe")); + // Raw install paths from the probe body must not leak into the friendly note. assert!(!lemonade.contains("D:/")); - assert!(!llama.contains("D:\\")); } #[test] fn missing_packaged_engine_reason_has_no_deferred_first_party_engines() { - assert!(missing_packaged_engine_reason("atom").is_none()); + assert!(missing_packaged_engine_reason("lemonade").is_none()); assert!(missing_packaged_engine_reason("vllm").is_none()); - assert!(missing_packaged_engine_reason("sglang").is_none()); - assert!(missing_packaged_engine_reason("pytorch").is_none()); - assert!(missing_packaged_engine_reason("llama.cpp").is_none()); } #[test] @@ -21543,7 +21274,7 @@ VERSION_ID="41" let detect = DetectResponse { installed: false, env_id: None, - runtime_kind: Some("external_sglang".to_owned()), + runtime_kind: Some("external_vllm".to_owned()), runtime_executable: None, managed_env: Some(false), python_version: None, @@ -21553,7 +21284,7 @@ VERSION_ID="41" kind: "rocm_gpu".to_owned(), available: false, reason: Some( - "SGLang ROCm serving is supported by rocm-cli only on Linux/WSL; native Windows SGLang is skipped. No CPU fallback is used." + "vLLM ROCm serving is supported by rocm-cli only on Linux/WSL; native Windows vLLM is skipped. No CPU fallback is used." .to_owned(), ), }], @@ -21562,7 +21293,7 @@ VERSION_ID="41" rocm_gpu: false, openai_compatible: true, tool_calling: false, - quantized_models: "sglang-supported".to_owned(), + quantized_models: "vllm-supported".to_owned(), reasoning_parser: false, }, notes: Vec::new(), @@ -21570,26 +21301,16 @@ VERSION_ID="41" if cfg!(windows) { assert_eq!( - engine_runtime_status_label("sglang", &detect), - "unsupported_native_windows" - ); - let mut atom_detect = detect; - atom_detect.runtime_kind = Some("external_atom".to_owned()); - atom_detect.available_devices[0].reason = Some( - "ATOM ROCm serving is supported by rocm-cli only on Linux/WSL; native Windows ATOM is not enabled. No CPU fallback is used." - .to_owned(), - ); - assert_eq!( - engine_runtime_status_label("atom", &atom_detect), + engine_runtime_status_label("vllm", &detect), "unsupported_native_windows" ); assert!( - model_registry_adapter_availability_note("atom") + model_registry_adapter_availability_note("vllm") .is_some_and(|note| note.contains("unsupported_native_windows")) ); } else { - assert_eq!(engine_runtime_status_label("sglang", &detect), "not found"); - assert!(model_registry_adapter_availability_note("atom").is_none()); + assert_eq!(engine_runtime_status_label("vllm", &detect), "not found"); + assert!(model_registry_adapter_availability_note("vllm").is_none()); } } @@ -21695,7 +21416,7 @@ VERSION_ID="41" license: Some("apache-2.0".to_owned()), gated: Some(false), quantization: Some("bfloat16".to_owned()), - engines: vec!["pytorch".to_owned()], + engines: vec!["vllm".to_owned()], source_policy: Some(rocm_core::ModelRecipeArtifactSourcePolicyRecord { policy: "huggingface_public".to_owned(), required_hosts: vec!["huggingface.co".to_owned()], @@ -21730,7 +21451,7 @@ VERSION_ID="41" assert!(output.contains("note: test metadata only")); assert!(!output.contains("source_policy=huggingface_public")); assert!(output.contains("size=2.0 GiB")); - assert!(output.contains("engines=[pytorch]")); + assert!(output.contains("engines=[vllm]")); assert!(output.contains("artifact_cache hf-main status=missing")); assert!(output.contains("prefetch requires an approved source policy")); assert!(output.contains( @@ -21906,19 +21627,19 @@ VERSION_ID="41" fn examine_engine_inventory_reports_config_without_engine_detect() { let (root, paths) = test_paths("examine-engine-inventory"); let mut config = RocmCliConfig { - default_engine: Some("llama.cpp".to_owned()), + default_engine: Some("vllm".to_owned()), ..RocmCliConfig::default() }; - config.engine_config_mut("llama.cpp").preferred_runtime_id = + config.engine_config_mut("vllm").preferred_runtime_id = Some("therock-release:gfx120X-all".to_owned()); let mut output = String::new(); append_examine_engine_inventory(&mut output, &paths, &config); assert!(output.contains("engine_inventory:")); - assert!(output.contains("configured_default_engine: llama.cpp")); - assert!(output.contains("effective_default_engine: llama.cpp")); - assert!(output.contains("* llama.cpp")); + assert!(output.contains("configured_default_engine: vllm")); + assert!(output.contains("effective_default_engine: vllm")); + assert!(output.contains("* vllm")); assert!(output.contains("runtime_pref=therock-release:gfx120X-all")); assert!(output.contains("plugin_dirs:")); let _ = fs::remove_dir_all(root); diff --git a/apps/rocm/src/providers.rs b/apps/rocm/src/providers.rs index b2bae5f8..f0be13c1 100644 --- a/apps/rocm/src/providers.rs +++ b/apps/rocm/src/providers.rs @@ -1526,7 +1526,7 @@ mod tests { let mut ready = ManagedServiceRecord::new( &paths, "svc-ready", - "pytorch", + "vllm", "qwen", BUILTIN_ASSISTANT_MODEL_ID, "127.0.0.1", @@ -1543,7 +1543,7 @@ mod tests { let mut stopped = ManagedServiceRecord::new( &paths, "svc-stopped", - "pytorch", + "vllm", "llama", "meta-llama/Llama", "127.0.0.1", @@ -2025,7 +2025,7 @@ mod tests { let mut ready = ManagedServiceRecord::new( &paths, "svc-ready", - "pytorch", + "vllm", "qwen", "Qwen/Qwen3.5", "127.0.0.1", @@ -2088,7 +2088,7 @@ mod tests { let mut custom = ManagedServiceRecord::new( &paths, "svc-custom", - "pytorch", + "vllm", "tiny-gpt2", "sshleifer/tiny-gpt2", "127.0.0.1", @@ -2118,7 +2118,7 @@ mod tests { let mut qwen = ManagedServiceRecord::new( &paths, "svc-qwen", - "pytorch", + "vllm", BUILTIN_ASSISTANT_MODEL_ALIAS, BUILTIN_ASSISTANT_MODEL_ID, "127.0.0.1", @@ -2191,7 +2191,7 @@ mod tests { let mut custom = ManagedServiceRecord::new( &paths, "svc-custom-qwen", - "pytorch", + "vllm", "Qwen/Qwen3.5-4B", "Qwen/Qwen3.5-4B", "127.0.0.1", @@ -2259,7 +2259,7 @@ mod tests { let mut ready = ManagedServiceRecord::new( &paths, "svc-ready", - "llama.cpp", + "vllm", "tiny.gguf", "tiny.gguf", "127.0.0.1", @@ -2322,7 +2322,7 @@ mod tests { let mut ready = ManagedServiceRecord::new( &paths, "svc-ready", - "pytorch", + "vllm", "qwen", "Qwen/Qwen3.5", "127.0.0.1", @@ -2425,7 +2425,7 @@ mod tests { let mut ready = ManagedServiceRecord::new( &paths, "svc-ready", - "pytorch", + "vllm", "qwen", "Qwen/Qwen3.5", "127.0.0.1", @@ -2510,7 +2510,7 @@ mod tests { let mut ready = ManagedServiceRecord::new( &paths, "svc-ready", - "pytorch", + "vllm", "qwen", "Qwen/Qwen3.5", "127.0.0.1", diff --git a/apps/rocmd/src/lib.rs b/apps/rocmd/src/lib.rs index 8f205c32..62e38edf 100644 --- a/apps/rocmd/src/lib.rs +++ b/apps/rocmd/src/lib.rs @@ -1383,11 +1383,6 @@ fn bridge_engine_inventory() -> Vec { const fn rocmd_engine_inventory() -> &'static [(&'static str, &'static str)] { &[ - ("pytorch", "default local serving engine"), - ( - "llama.cpp", - "GGUF serving with ROCm GPU required by rocm-cli", - ), ( "lemonade", "embedded Lemonade server with ROCm llama.cpp backend", @@ -1396,14 +1391,6 @@ const fn rocmd_engine_inventory() -> &'static [(&'static str, &'static str)] { "vllm", "Linux/WSL ROCm GPU serving engine through external vLLM", ), - ( - "sglang", - "Linux/WSL ROCm GPU serving engine through external SGLang", - ), - ( - "atom", - "Linux/WSL ROCm GPU serving engine through external ATOM Python", - ), ] } @@ -5288,11 +5275,11 @@ mod tests { let plugin_dir = paths.data_dir.join("engines").join("plugins"); fs::create_dir_all(&plugin_dir)?; let plugin_path = plugin_dir.join( - rocm_engine_protocol::platform_engine_plugin_binary_name("pytorch"), + rocm_engine_protocol::platform_engine_plugin_binary_name("vllm"), ); fs::write(&plugin_path, "plugin")?; - let discovered = find_engine_plugin_binary("pytorch", engine_plugin_dirs(&paths))?; + let discovered = find_engine_plugin_binary("vllm", engine_plugin_dirs(&paths))?; fs::remove_dir_all(root).ok(); assert_eq!(discovered, Some(plugin_path)); @@ -5331,7 +5318,7 @@ mod tests { "supervise", "svc", "--engine", - "pytorch", + "vllm", "--model-ref", "qwen", "--canonical-model-id", @@ -5848,7 +5835,7 @@ mod tests { let mut failed = ManagedServiceRecord::new( &paths, "svc-failed", - "pytorch", + "vllm", "qwen", "Qwen/Qwen3.5", "127.0.0.1", @@ -6095,7 +6082,7 @@ mod tests { let mut record = ManagedServiceRecord::new( &paths, "svc-hot", - "llama.cpp", + "vllm", "tiny", "Tiny/Test", "127.0.0.1", @@ -6163,7 +6150,7 @@ mod tests { let mut record = ManagedServiceRecord::new( &paths, "svc-hot", - "pytorch", + "vllm", "qwen", "Qwen/Test", "127.0.0.1", @@ -6250,7 +6237,7 @@ mod tests { let mut record = ManagedServiceRecord::new( &paths, service_id, - "pytorch", + "vllm", "qwen", "Qwen/Test", "127.0.0.1", @@ -6300,7 +6287,7 @@ mod tests { let mut record = ManagedServiceRecord::new( &paths, "svc-hot", - "llama.cpp", + "vllm", "tiny", "Tiny/Test", "127.0.0.1", @@ -6743,7 +6730,7 @@ mod tests { let mut failed = ManagedServiceRecord::new( &paths, "svc-failed", - "pytorch", + "vllm", "qwen", "Qwen/Qwen3.5", "127.0.0.1", @@ -6795,7 +6782,7 @@ mod tests { let mut healthy = ManagedServiceRecord::new( &paths, "svc-healthy", - "pytorch", + "vllm", "qwen", "Qwen/Qwen3.5", "127.0.0.1", @@ -6906,7 +6893,7 @@ mod tests { let mut record = ManagedServiceRecord::new( &paths, "svc-stale", - "pytorch", + "vllm", "qwen", "Qwen/Qwen3.5", "127.0.0.1", @@ -6955,7 +6942,7 @@ mod tests { let mut failed = ManagedServiceRecord::new( &paths, "svc-failed", - "pytorch", + "vllm", "qwen", "Qwen/Qwen3.5", "127.0.0.1", @@ -6983,7 +6970,7 @@ mod tests { let mut stale = ManagedServiceRecord::new( &paths, "svc-starting", - "pytorch", + "vllm", "qwen", "Qwen/Qwen3.5", "127.0.0.1", @@ -7059,7 +7046,7 @@ mod tests { let mut record = ManagedServiceRecord::new( &paths, "svc-1", - "pytorch", + "vllm", "qwen", "Qwen/Qwen3.5", "127.0.0.1", @@ -7521,7 +7508,7 @@ mod tests { let record = ManagedServiceRecord::new( &paths, "svc-1", - "pytorch", + "vllm", "qwen", "Qwen/Qwen3.5", "127.0.0.1", @@ -7674,7 +7661,7 @@ mod tests { let mut record = ManagedServiceRecord::new( &paths, "svc-current", - "pytorch", + "vllm", "qwen", "Qwen/Qwen3.5", "127.0.0.1", @@ -7863,7 +7850,7 @@ mod tests { license: Some("apache-2.0".to_owned()), gated: Some(false), quantization: Some("bf16".to_owned()), - engines: vec!["pytorch".to_owned()], + engines: vec!["vllm".to_owned()], source_policy: None, }, SandboxToolPolicy::default(), @@ -7919,7 +7906,7 @@ mod tests { license: Some("test-only".to_owned()), gated: Some(false), quantization: None, - engines: vec!["pytorch".to_owned()], + engines: vec!["vllm".to_owned()], source_policy: None, }; let cache = model_artifact_cache_status(&paths, "Qwen/Test-1B", &artifact); @@ -7960,7 +7947,7 @@ mod tests { license: Some("test-only".to_owned()), gated: Some(false), quantization: None, - engines: vec!["pytorch".to_owned()], + engines: vec!["vllm".to_owned()], source_policy: None, }; @@ -8004,7 +7991,7 @@ mod tests { license: Some("test-only".to_owned()), gated: Some(false), quantization: None, - engines: vec!["pytorch".to_owned()], + engines: vec!["vllm".to_owned()], source_policy: None, }; @@ -8048,7 +8035,7 @@ mod tests { license: Some("test-only".to_owned()), gated: Some(true), quantization: None, - engines: vec!["pytorch".to_owned()], + engines: vec!["vllm".to_owned()], source_policy: None, }; @@ -8092,7 +8079,7 @@ mod tests { license: Some("test-only".to_owned()), gated: Some(true), quantization: None, - engines: vec!["pytorch".to_owned()], + engines: vec!["vllm".to_owned()], source_policy: None, }; @@ -8137,7 +8124,7 @@ mod tests { license: Some("test-only".to_owned()), gated: Some(false), quantization: None, - engines: vec!["pytorch".to_owned()], + engines: vec!["vllm".to_owned()], source_policy: Some(ModelRecipeArtifactSourcePolicyRecord { policy: "manual_only".to_owned(), required_hosts: Vec::new(), @@ -8185,7 +8172,7 @@ mod tests { license: Some("test-only".to_owned()), gated: Some(false), quantization: None, - engines: vec!["pytorch".to_owned()], + engines: vec!["vllm".to_owned()], source_policy: Some(ModelRecipeArtifactSourcePolicyRecord { policy: "huggingface_authenticated".to_owned(), required_hosts: vec!["huggingface.co".to_owned()], @@ -8233,7 +8220,7 @@ mod tests { license: Some("test-only".to_owned()), gated: Some(true), quantization: None, - engines: vec!["pytorch".to_owned()], + engines: vec!["vllm".to_owned()], source_policy: None, }; @@ -8282,7 +8269,7 @@ mod tests { license: Some("test-only".to_owned()), gated: Some(true), quantization: None, - engines: vec!["pytorch".to_owned()], + engines: vec!["vllm".to_owned()], source_policy: None, }; @@ -8334,7 +8321,7 @@ mod tests { license: Some("test-only".to_owned()), gated: Some(false), quantization: None, - engines: vec!["pytorch".to_owned()], + engines: vec!["vllm".to_owned()], source_policy: None, }; @@ -8397,7 +8384,7 @@ mod tests { license: Some("test-only".to_owned()), gated: Some(false), quantization: None, - engines: vec!["pytorch".to_owned()], + engines: vec!["vllm".to_owned()], source_policy: None, }; diff --git a/crates/rocm-core/src/lib.rs b/crates/rocm-core/src/lib.rs index 1bf707f0..3c931359 100644 --- a/crates/rocm-core/src/lib.rs +++ b/crates/rocm-core/src/lib.rs @@ -5342,7 +5342,7 @@ pub fn builtin_model_recipes() -> Vec { engine_recipes: Vec::new(), manual_alternatives: vec!["qwen-tiny".to_owned(), "tiny-gpt2".to_owned()], chat_template_mode: "auto".to_owned(), - preferred_engines: vec!["pytorch".to_owned()], + preferred_engines: vec!["lemonade".to_owned()], warnings: vec![ "recommended local assistant path for low-VRAM ROCm machines".to_owned(), ], @@ -5366,14 +5366,14 @@ pub fn builtin_model_recipes() -> Vec { recommended_system_ram_gb: Some(8), quantization: Some("none; tiny instruct smoke recipe".to_owned()), artifact_hint: Some( - "Hugging Face model id; verified with the managed PyTorch Transformers line" + "Hugging Face safetensors model id; served on GPU through vLLM (Linux/WSL)" .to_owned(), ), artifacts: Vec::new(), engine_recipes: Vec::new(), manual_alternatives: vec!["qwen".to_owned(), "tiny-gpt2".to_owned()], chat_template_mode: "auto".to_owned(), - preferred_engines: vec!["pytorch".to_owned()], + preferred_engines: vec!["vllm".to_owned()], warnings: vec![ "tiny smoke path; use qwen for the smarter low-VRAM assistant".to_owned(), ], @@ -5504,7 +5504,7 @@ pub fn builtin_model_recipes() -> Vec { "tiny-gpt2".to_owned(), ], chat_template_mode: "auto".to_owned(), - preferred_engines: vec!["vllm".to_owned(), "pytorch".to_owned()], + preferred_engines: vec!["vllm".to_owned()], warnings: vec![ "this recipe prefers ROCm GPU execution and may span multiple accelerators" .to_owned(), @@ -5541,7 +5541,7 @@ pub fn builtin_model_recipes() -> Vec { "llama-3.2-3b-instruct".to_owned(), ], chat_template_mode: "auto".to_owned(), - preferred_engines: vec!["vllm".to_owned(), "pytorch".to_owned()], + preferred_engines: vec!["vllm".to_owned()], warnings: vec![ "this model family is configured with trust_remote_code enabled by recipe" .to_owned(), @@ -5566,14 +5566,14 @@ pub fn builtin_model_recipes() -> Vec { recommended_system_ram_gb: Some(16), quantization: Some("none; bfloat16 weights".to_owned()), artifact_hint: Some( - "Hugging Face model id for PyTorch; llama.cpp serving requires an explicit GGUF path" + "Hugging Face safetensors model id; served on GPU through vLLM (Linux/WSL)" .to_owned(), ), artifacts: Vec::new(), engine_recipes: Vec::new(), manual_alternatives: vec!["qwen".to_owned(), "tiny-gpt2".to_owned()], chat_template_mode: "auto".to_owned(), - preferred_engines: vec!["pytorch".to_owned(), "llama.cpp".to_owned()], + preferred_engines: vec!["vllm".to_owned()], warnings: Vec::new(), }, ModelRecipeRecord { @@ -5594,7 +5594,7 @@ pub fn builtin_model_recipes() -> Vec { engine_recipes: Vec::new(), manual_alternatives: Vec::new(), chat_template_mode: "auto".to_owned(), - preferred_engines: vec!["pytorch".to_owned()], + preferred_engines: vec!["vllm".to_owned()], warnings: Vec::new(), }, ] @@ -5878,11 +5878,7 @@ fn sibling_binary_candidates(current_exe: &Path, binary_name: &str) -> Result Result { - let binary_engine = match engine { - "llama.cpp" => "llama-cpp", - other => other, - }; - sibling_binary_path(&format!("rocm-engine-{binary_engine}")) + sibling_binary_path(&format!("rocm-engine-{engine}")) } pub fn daemon_binary_path() -> Result { @@ -6516,7 +6512,7 @@ Class Name: Display #[test] fn managed_therock_family_falls_back_to_engine_env_manifest() -> Result<()> { let (root, paths) = temp_app_paths("engine-therock-family"); - let manifests = paths.engine_manifests_dir("pytorch"); + let manifests = paths.engine_manifests_dir("vllm"); fs::create_dir_all(&manifests)?; fs::write( manifests.join("env.json"), @@ -6719,7 +6715,7 @@ Class Name: Display cpu: Some("AMD Ryzen".to_owned()), system_ram_gib: Some(64.0), interactive_terminal: false, - default_engine: "pytorch".to_owned(), + default_engine: "vllm".to_owned(), detected_gfx_target: None, compatible_therock_family: Some("gfx120X-all".to_owned()), detected_therock_family: None, @@ -6771,7 +6767,7 @@ Class Name: Display cpu: None, system_ram_gib: None, interactive_terminal: false, - default_engine: "pytorch".to_owned(), + default_engine: "vllm".to_owned(), detected_gfx_target: None, compatible_therock_family: None, detected_therock_family: None, @@ -7078,7 +7074,7 @@ Class Name: Display assert_eq!(artifact.kind, "huggingface"); let expected_sha = "a".repeat(64); assert_eq!(artifact.sha256.as_deref(), Some(expected_sha.as_str())); - assert_eq!(artifact.engines, vec!["pytorch"]); + assert_eq!(artifact.engines, vec!["vllm"]); assert_eq!( artifact .source_policy @@ -7569,13 +7565,13 @@ Class Name: Display license: Some("apache-2.0".to_owned()), gated: Some(false), quantization: Some("none".to_owned()), - engines: vec!["pytorch".to_owned()], + engines: vec!["vllm".to_owned()], source_policy: None, }], engine_recipes: Vec::new(), manual_alternatives: Vec::new(), chat_template_mode: "auto".to_owned(), - preferred_engines: vec!["pytorch".to_owned()], + preferred_engines: vec!["vllm".to_owned()], warnings: Vec::new(), } } @@ -7681,9 +7677,9 @@ Class Name: Display } assert_eq!( - paths.engine_envs_dir("pytorch"), + paths.engine_envs_dir("vllm"), normalize_runtime_path_for_host(&override_root) - .join("pytorch") + .join("vllm") .join("envs") ); @@ -7699,10 +7695,10 @@ Class Name: Display #[test] fn legacy_config_without_telemetry_uses_default_policy() -> Result<()> { let config = serde_json::from_value::(serde_json::json!({ - "default_engine": "pytorch" + "default_engine": "vllm" }))?; - assert_eq!(config.default_engine.as_deref(), Some("pytorch")); + assert_eq!(config.default_engine.as_deref(), Some("vllm")); assert_eq!(config.telemetry.mode_label(), TELEMETRY_MODE_LOCAL); Ok(()) } diff --git a/crates/rocm-dash-collectors/src/engine_registry.rs b/crates/rocm-dash-collectors/src/engine_registry.rs index 41473ff1..208e75af 100644 --- a/crates/rocm-dash-collectors/src/engine_registry.rs +++ b/crates/rocm-dash-collectors/src/engine_registry.rs @@ -19,8 +19,6 @@ pub enum EngineKind { Vllm, /// Lemonade Server — JSON `/api/v1/stats`. Lemonade, - /// llama.cpp `llama-server` — `/slots` (parser not yet wired). - LlamaCpp, } impl EngineKind { @@ -29,7 +27,6 @@ impl EngineKind { match self { Self::Vllm => "vllm", Self::Lemonade => "lemonade", - Self::LlamaCpp => "llama.cpp", } } @@ -41,30 +38,25 @@ impl EngineKind { match self { Self::Vllm => 8000, Self::Lemonade => crate::lemonade::LEMONADE_PORT, // 13305 - Self::LlamaCpp => 8080, } } - /// Map an engine label (config / discovery) to a kind. Case-insensitive; - /// accepts the `llama.cpp` / `llamacpp` / `llama_cpp` spellings. + /// Map an engine label (config / discovery) to a kind. Case-insensitive. pub fn from_label(label: &str) -> Option { match label.trim().to_ascii_lowercase().as_str() { "vllm" => Some(Self::Vllm), "lemonade" => Some(Self::Lemonade), - "llama.cpp" | "llamacpp" | "llama_cpp" => Some(Self::LlamaCpp), _ => None, } } /// Parse this engine's raw scrape body into an [`InstanceSample`] using the /// engine-appropriate parser. vLLM dispatches to the **unchanged** - /// `vllm_prom::parse`; Lemonade to `lemonade::parse_stats`. Unwired engines - /// return an empty sample (never panic). + /// `vllm_prom::parse`; Lemonade to `lemonade::parse_stats`. pub fn parse_sample(self, body: &str) -> InstanceSample { match self { Self::Vllm => crate::vllm_prom::parse(body), Self::Lemonade => crate::lemonade::parse_stats(body), - Self::LlamaCpp => InstanceSample::default(), } } } @@ -75,15 +67,12 @@ mod tests { #[test] fn engine_kind_labels_and_ports_and_from_label_roundtrip() { - for kind in [EngineKind::Vllm, EngineKind::Lemonade, EngineKind::LlamaCpp] { + for kind in [EngineKind::Vllm, EngineKind::Lemonade] { assert_eq!(EngineKind::from_label(kind.label()), Some(kind)); } assert_eq!(EngineKind::Lemonade.default_port(), 13305); assert_eq!(EngineKind::Vllm.default_port(), 8000); - assert_eq!( - EngineKind::from_label("LLAMACPP"), - Some(EngineKind::LlamaCpp) - ); + assert_eq!(EngineKind::from_label("VLLM"), Some(EngineKind::Vllm)); assert_eq!(EngineKind::from_label("nope"), None); } diff --git a/crates/rocm-dash-collectors/src/lib.rs b/crates/rocm-dash-collectors/src/lib.rs index a4ac4192..f967a6fa 100644 --- a/crates/rocm-dash-collectors/src/lib.rs +++ b/crates/rocm-dash-collectors/src/lib.rs @@ -14,7 +14,6 @@ pub mod docker; pub mod engine_registry; pub mod host; pub mod lemonade; -pub mod llama_slots; pub mod parallel; pub mod proc_scan; pub mod sysfs; diff --git a/crates/rocm-dash-collectors/src/llama_slots.rs b/crates/rocm-dash-collectors/src/llama_slots.rs deleted file mode 100644 index dde21ea2..00000000 --- a/crates/rocm-dash-collectors/src/llama_slots.rs +++ /dev/null @@ -1,30 +0,0 @@ -// Copyright © Advanced Micro Devices, Inc., or its affiliates. -// -// SPDX-License-Identifier: MIT - -//! llama-server /slots collector. Stub. - -use rocm_dash_core::traits::{ - CollectorError, DiscoveredService, InstanceMetrics, InstanceSample, Result, -}; - -#[derive(Debug, Default)] -pub struct LlamaSlotsCollector; - -impl LlamaSlotsCollector { - pub const fn new() -> Self { - Self - } -} - -impl InstanceMetrics for LlamaSlotsCollector { - fn name(&self) -> &'static str { - "llama-slots" - } - - fn fetch(&self, _svc: &DiscoveredService) -> Result { - // TODO: GET http://localhost:{port}/slots; count non-idle entries. - // llama.cpp has no waiting queue → set waiting_reqs = Some(0). - Err(CollectorError::Unsupported("llama-slots stub".into())) - } -} diff --git a/crates/rocm-dash-collectors/src/parallel.rs b/crates/rocm-dash-collectors/src/parallel.rs index dc0e447f..0ad1d68f 100644 --- a/crates/rocm-dash-collectors/src/parallel.rs +++ b/crates/rocm-dash-collectors/src/parallel.rs @@ -5,9 +5,9 @@ //! Shared primitives for async collectors that scrape N targets concurrently. //! //! Today: [[VllmPrometheusCollector]] uses this pattern hand-rolled in the -//! daemon's runner. As Strix-Halo's `LlamaSlotsCollector` and any future -//! per-instance scrapers come online, they'll reuse `parallel_scrape` and -//! the `WarningBus` instead of re-implementing the JoinSet glue. +//! daemon's runner. As future per-instance scrapers come online, they'll reuse +//! `parallel_scrape` and the `WarningBus` instead of re-implementing the +//! JoinSet glue. use std::future::Future; use std::sync::Arc; diff --git a/crates/rocm-dash-tui/src/agent.rs b/crates/rocm-dash-tui/src/agent.rs index 0ffc759b..9bae1c69 100644 --- a/crates/rocm-dash-tui/src/agent.rs +++ b/crates/rocm-dash-tui/src/agent.rs @@ -530,7 +530,7 @@ rocm_read_tool!( rocm_read_tool!( EnginesRocmTool, "engines", - "List the available inference engines (e.g. vLLM, llama.cpp, ComfyUI) and \ + "List the available inference engines (e.g. Lemonade, vLLM, ComfyUI) and \ their install/availability status. Read-only.", { "type": "object", "properties": {} } ); @@ -753,7 +753,7 @@ rocm_mutating_tool!( rocm_mutating_tool!( InstallEngineRocmTool, "install_engine", - "Install an inference engine (e.g. vllm, llama-cpp, comfyui). MUTATING — \ + "Install an inference engine (e.g. lemonade, vllm, comfyui). MUTATING — \ surfaced for operator approval before anything runs.", { "type": "object", diff --git a/crates/rocm-dash-tui/src/app/chat.rs b/crates/rocm-dash-tui/src/app/chat.rs index 2ddb1e5c..63b058d1 100644 --- a/crates/rocm-dash-tui/src/app/chat.rs +++ b/crates/rocm-dash-tui/src/app/chat.rs @@ -90,7 +90,7 @@ pub(super) fn build_chat_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", "llama.cpp", "sglang", "pytorch"]; +const OPENAI_COMPATIBLE_ENGINES: &[&str] = &["vllm", "lemonade"]; /// 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/engine_manager.rs b/crates/rocm-dash-tui/src/ui/engine_manager.rs index 112495ab..efcf26d7 100644 --- a/crates/rocm-dash-tui/src/ui/engine_manager.rs +++ b/crates/rocm-dash-tui/src/ui/engine_manager.rs @@ -38,20 +38,10 @@ pub const ENGINE_CATALOG: &[(&str, &str)] = &[ "lemonade", "default embedded Lemonade server with ROCm llama.cpp backend", ), - ("pytorch", "TheRock PyTorch local serving engine"), - ("llama.cpp", "external GGUF serving engine for llama-server"), ( "vllm", "Linux/WSL ROCm GPU serving engine through external vLLM", ), - ( - "sglang", - "Linux/WSL ROCm GPU serving engine through external SGLang", - ), - ( - "atom", - "Linux/WSL ROCm GPU serving engine through external ATOM Python", - ), ]; /// A lifecycle operation on an engine. @@ -205,7 +195,7 @@ fn spawn_engine_op( jobs: &mut State, pending: PendingEngineOp, ) -> Vec { - // Sanitize the engine name for the job id (e.g. `llama.cpp` → `llama-cpp`). + // Sanitize the engine name for the job id (any non-alphanumeric char → `-`). let key: String = pending .engine .chars() @@ -361,20 +351,6 @@ mod tests { assert_eq!(jobs.jobs.len(), 1); } - #[test] - fn job_id_sanitizes_dotted_engine_name() { - let mut em = Some(EngineManagerState::default()); - let mut jobs = State::default(); - // Select llama.cpp (index 2). - em.as_mut().unwrap().selected = 2; - on_key(&mut em, &mut jobs, key(KeyCode::Char('r'))); - on_key(&mut em, &mut jobs, key(KeyCode::Char('y'))); - assert_eq!( - em.as_ref().unwrap().active_job.as_deref(), - Some("engine-reinstall-llama-cpp") - ); - } - #[test] fn relaunch_while_prior_job_running_surfaces_message() { // The reducer no-ops a StartJob for a still-running id. spawn_engine_op diff --git a/crates/rocm-dash-tui/src/ui/serve_wizard.rs b/crates/rocm-dash-tui/src/ui/serve_wizard.rs index 236017d1..4e474804 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", "pytorch", "llama.cpp", "vllm", "sglang", "atom"]; +pub const ENGINES: &[&str] = &["lemonade", "vllm"]; /// 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-dash-tui/src/ui/tabs/serving.rs b/crates/rocm-dash-tui/src/ui/tabs/serving.rs index c8dc31c3..6d66a249 100644 --- a/crates/rocm-dash-tui/src/ui/tabs/serving.rs +++ b/crates/rocm-dash-tui/src/ui/tabs/serving.rs @@ -25,7 +25,7 @@ pub const VERBS: &[Verb] = &[ summary: "Launch a model on a serving engine and expose an OpenAI-style endpoint.", steps: &[ "Pick a model", - "Choose an engine — vLLM · SGLang · llama.cpp · PyTorch", + "Choose an engine — Lemonade · vLLM", "Set GPU placement (required / preferred / CPU-only)", "Launch on 127.0.0.1:11435 and watch it come up", ], @@ -39,7 +39,7 @@ pub const VERBS: &[Verb] = &[ action: KeyAction::OpenEngineManager, summary: "Install, reinstall, or configure serving engines.", steps: &[ - "Browse engines — vLLM · SGLang · llama.cpp · PyTorch · lemonade", + "Browse engines — Lemonade · vLLM", "See install status per engine", "Install / reinstall an engine", "Adjust engine config", diff --git a/crates/rocm-engine-protocol/src/lib.rs b/crates/rocm-engine-protocol/src/lib.rs index afdfea0d..72c51047 100644 --- a/crates/rocm-engine-protocol/src/lib.rs +++ b/crates/rocm-engine-protocol/src/lib.rs @@ -21,30 +21,16 @@ pub struct EnginePluginDescriptor { } pub fn platform_engine_plugin_binary_name(engine_id: &str) -> String { - let binary_name = format!( - "{ENGINE_PLUGIN_BINARY_PREFIX}{}", - engine_id_to_plugin_binary_component(engine_id) - ); + let binary_name = format!("{ENGINE_PLUGIN_BINARY_PREFIX}{engine_id}"); rocm_core::platform_binary_name(&binary_name) } -fn engine_id_to_plugin_binary_component(engine_id: &str) -> &str { - match engine_id { - "llama.cpp" => "llama-cpp", - other => other, - } -} - pub fn engine_id_from_plugin_binary_name(name: &str) -> Option { if name.contains('/') || name.contains('\\') { return None; } let name = strip_optional_exe_suffix(name); let engine_id = name.strip_prefix(ENGINE_PLUGIN_BINARY_PREFIX)?; - let engine_id = match engine_id { - "llama-cpp" => "llama.cpp", - other => other, - }; is_valid_engine_id(engine_id).then(|| engine_id.to_owned()) } @@ -629,12 +615,12 @@ mod tests { #[test] fn plugin_binary_names_normalize_to_engine_ids() { assert_eq!( - engine_id_from_plugin_binary_name("rocm-engine-pytorch"), - Some("pytorch".to_owned()) + engine_id_from_plugin_binary_name("rocm-engine-lemonade"), + Some("lemonade".to_owned()) ); assert_eq!( - engine_id_from_plugin_binary_name("rocm-engine-llama-cpp.EXE"), - Some("llama.cpp".to_owned()) + engine_id_from_plugin_binary_name("rocm-engine-vllm.EXE"), + Some("vllm".to_owned()) ); assert_eq!(engine_id_from_plugin_binary_name("rocm-engine-"), None); assert_eq!(engine_id_from_plugin_binary_name("rocm"), None); @@ -643,7 +629,7 @@ mod tests { None ); assert_eq!( - engine_id_from_plugin_binary_name("nested/rocm-engine-pytorch"), + engine_id_from_plugin_binary_name("nested/rocm-engine-lemonade"), None ); } @@ -656,16 +642,17 @@ mod tests { fs::create_dir_all(&first_dir).unwrap(); fs::create_dir_all(&second_dir).unwrap(); - let first_pytorch = first_dir.join(platform_engine_plugin_binary_name("pytorch")); - let first_vllm = first_dir.join(platform_engine_plugin_binary_name("vllm")); - let second_atom = second_dir.join(platform_engine_plugin_binary_name("atom")); - let second_pytorch = second_dir.join(platform_engine_plugin_binary_name("pytorch")); - fs::write(&first_pytorch, b"pytorch").unwrap(); + // The `.exe` suffix is optional in `engine_id_from_plugin_binary_name`, so + // the bare `rocm-engine-` names are discovered on every platform. + let first_lemonade = first_dir.join("rocm-engine-lemonade"); + let first_vllm = first_dir.join("rocm-engine-vllm"); + let second_lemonade = second_dir.join("rocm-engine-lemonade"); + fs::write(&first_lemonade, b"lemonade").unwrap(); fs::write(&first_vllm, b"vllm").unwrap(); - fs::write(&second_atom, b"atom").unwrap(); - fs::write(&second_pytorch, b"duplicate pytorch").unwrap(); + fs::write(&second_lemonade, b"duplicate lemonade").unwrap(); fs::write(first_dir.join("not-an-engine"), b"ignore").unwrap(); - fs::create_dir_all(first_dir.join(platform_engine_plugin_binary_name("sglang"))).unwrap(); + // A directory matching the plugin naming must be ignored (only files count). + fs::create_dir_all(second_dir.join("rocm-engine-ghost")).unwrap(); let plugins = discover_engine_plugins([root.join("missing"), first_dir, second_dir]).unwrap(); @@ -673,14 +660,14 @@ mod tests { .iter() .map(|plugin| plugin.id.as_str()) .collect::>(); - assert_eq!(ids, vec!["atom", "pytorch", "vllm"]); + assert_eq!(ids, vec!["lemonade", "vllm"]); assert_eq!( plugins .iter() - .find(|plugin| plugin.id == "pytorch") + .find(|plugin| plugin.id == "lemonade") .unwrap() .executable_path, - first_pytorch + first_lemonade ); fs::remove_dir_all(root).unwrap(); diff --git a/docs/atom.md b/docs/atom.md deleted file mode 100644 index db121c12..00000000 --- a/docs/atom.md +++ /dev/null @@ -1,82 +0,0 @@ - - -# ATOM Adapter - -`rocm-engine-atom` is the first-party rocm-cli adapter for AMD's ATOM serving -runtime. - -Current behavior: - -- Linux/WSL only for ROCm GPU serving. -- Native Windows is explicitly gated and does not fall back to CPU. -- CPU policy is rejected. `gpu_preferred` is treated as `gpu_required`. -- The adapter can use: - - `ROCM_CLI_ATOM_COMMAND` or `ATOM_COMMAND` - - `ROCM_CLI_ATOM_PYTHON` or `ATOM_PYTHON` - - an active rocm-cli managed TheRock runtime if its Python environment has - the `atom` package installed -- TheRock runtime environment variables and library paths are applied before - launching ATOM, so non-PyTorch HIP processes can load libraries from the - managed SDK root. - -Live acceptance status: - -- Upstream ATOM documentation currently lists AMD Instinct MI355X (`gfx950`), - MI300X (`gfx942`), and MI250X (`gfx90a`) as supported GPUs. -- This development machine is an RDNA4 Radeon GPU (`gfx1201`), so rocm-cli - does not claim live ATOM GPU acceptance on this host. -- The adapter remains packaged and test-covered for protocol behavior, - Windows no-fallback gating, CPU-policy rejection, and managed TheRock - environment propagation. Live acceptance should be run only on a supported - ATOM GPU/runtime/model combination. - -Acceptance harness: - -```bash -python -m py_compile scripts/atom_therock_gpu_test.py -python scripts/atom_therock_gpu_test.py --self-test -``` - -On supported Linux/WSL ATOM hardware with ATOM installed in the active -rocm-cli managed TheRock runtime: - -```bash -python3 scripts/atom_therock_gpu_test.py \ - --engine /home/user/.cache/rocm-cli-target/debug/rocm-engine-atom \ - --model Qwen/Qwen3-0.6B -``` - -The harness defaults to the active exact runtime key. An explicit -`--runtime-id` may be an exact runtime key or an unambiguous runtime id, but it -never picks the newest manifest on ambiguity. It rejects -`ROCM_CLI_ATOM_COMMAND`, `ATOM_COMMAND`, `ROCM_CLI_ATOM_PYTHON`, and -`ATOM_PYTHON` so live acceptance proves the managed TheRock runtime path. It -also checks `cpu_only` rejection, `gpu_required` launch state, OpenAI-compatible -serving, managed TheRock environment variables, and loaded ROCm HIP/math -libraries from the managed SDK wheel roots. - -The upstream ATOM serving command documented by ROCm is: - -```bash -python -m atom.entrypoints.openai_server --model Qwen/Qwen3-0.6B --kv_cache_dtype fp8 -``` - -rocm-cli currently launches the same Python module form and passes `--model`, -`--host`, and `--port`. - -Useful checks: - -```bash -rocm-engine-atom detect -rocm-engine-atom capabilities -rocm-engine-atom resolve-model Qwen/Qwen3-0.6B -``` - -Sources: - -- https://github.com/ROCm/ATOM -- https://rocm.github.io/ATOM/docs/ diff --git a/docs/engine-plugins.md b/docs/engine-plugins.md index 8bb38a2e..5b0d2ac9 100644 --- a/docs/engine-plugins.md +++ b/docs/engine-plugins.md @@ -16,13 +16,17 @@ Search order: 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. The `llama.cpp` engine uses -`rocm-engine-llama-cpp` on Linux/WSL and `rocm-engine-llama-cpp.exe` on -Windows. +`rocm-engine-.exe` on Windows. -Packaged first-party adapters are `pytorch`, `llama.cpp`, `lemonade`, `atom`, -`vllm`, and `sglang`. Linux/WSL-only ROCm GPU adapters fail explicitly on -native Windows instead of selecting a CPU fallback. +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. + +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. The `lemonade` adapter uses Lemonade embeddable and requires Lemonade's `llamacpp:rocm` backend. Windows ROCm serving is validated. WSL is currently diff --git a/docs/llm-tool-use.md b/docs/llm-tool-use.md index 1d6109e7..d32905f8 100644 --- a/docs/llm-tool-use.md +++ b/docs/llm-tool-use.md @@ -23,13 +23,12 @@ 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, SGLang, - PyTorch, llama.cpp, and Lemonade are 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. -- On native Windows, vLLM and SGLang live serving/install checks are skipped. - The assistant should direct those requests to WSL/Linux and should not suggest - CPU fallback. + (`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 + serving, but it should not switch its own built-in chat engine away from + Lemonade. +- On native Windows, vLLM live serving/install checks are skipped. The assistant + should direct those requests to WSL/Linux and should not suggest CPU fallback. This follows the same shape described by current tool-use docs: the application defines tool schemas, the model requests a tool, the application executes the @@ -76,19 +75,18 @@ before running it: {"name":"rocm_command","arguments":{"args":["comfyui","install"],"reason":"Install ComfyUI into ROCm CLI's app folder."}} ``` -The assistant can request installing llama.cpp through the existing engine -surface. The `llama.cpp` engine is backed by upstream `llama-server`; do not -replace it with one-off llama.cpp command runners: +The assistant can request installing a serving engine through the existing +engine surface. The supported engines are `lemonade` and `vllm`: ```json -{"name":"rocm_command","arguments":{"args":["engines","install","llama.cpp"],"reason":"Install the GGUF serving engine."}} +{"name":"rocm_command","arguments":{"args":["engines","install","vllm"],"reason":"Install the vLLM serving engine."}} ``` -The assistant can request serving a GGUF model through `llama-server` by asking -rocm-cli to start the managed `llama.cpp` engine. GPU execution is required: +The assistant can request serving a model through vLLM by asking rocm-cli to +start the managed `vllm` engine. GPU execution is required: ```json -{"name":"rocm_command","arguments":{"args":["serve","D:\\models\\tiny.gguf","--engine","llama.cpp","--device","gpu_required","--managed"],"reason":"Start a local GPU llama-server for this GGUF model."}} +{"name":"rocm_command","arguments":{"args":["serve","Qwen/Qwen3.5-4B","--engine","vllm","--device","gpu_required","--managed"],"reason":"Start a local GPU vLLM server for this model."}} ``` To target a specific GPU, add `--gpu` with `auto` (default; first free GPU) or a @@ -96,7 +94,7 @@ single index. Serving one model across multiple GPUs is not supported. CPU fallback is never used when a GPU is busy or out of range: ```json -{"name":"rocm_command","arguments":{"args":["serve","D:\\models\\tiny.gguf","--engine","llama.cpp","--device","gpu_required","--gpu","1","--managed"],"reason":"Serve this GGUF model on GPU 1."}} +{"name":"rocm_command","arguments":{"args":["serve","Qwen/Qwen3.5-4B","--engine","vllm","--device","gpu_required","--gpu","1","--managed"],"reason":"Serve this model on GPU 1."}} ``` The assistant can request starting ComfyUI. rocm-cli shows the local URL and diff --git a/docs/manual-testing.md b/docs/manual-testing.md index d7a91c9f..ca37de60 100644 --- a/docs/manual-testing.md +++ b/docs/manual-testing.md @@ -153,35 +153,7 @@ Expected result: Stop the foreground server with `Ctrl+C`. -## 4. PyTorch GPU Verification - -Install or refresh the PyTorch engine: - -```powershell -rocm engines install pytorch -``` - -Serve a small model with the managed runtime: - -```powershell -rocm serve qwen --engine pytorch --device gpu_required --managed --foreground --port 11435 -``` - -Expected result: - -- The engine uses the active rocm-cli managed TheRock runtime. -- PyTorch detects the AMD GPU. -- The run does not fall back to CPU. -- If the GPU cannot be used, the command fails with a clear error. - -Current known caveat: explicit `Qwen/Qwen3.5-4B` PyTorch requests are gated -before launch because the installed Transformers package does not recognize -the checkpoint's `qwen3_5` architecture. The short alias `qwen` resolves to -the recommended low-VRAM `Qwen/Qwen2.5-1.5B-Instruct` assistant recipe. - -Stop the foreground server with `Ctrl+C`. - -## 5. Local Server Records +## 4. Local Server Records After a managed or foreground serve attempt, inspect local server records: @@ -204,41 +176,7 @@ rocm services stop --yes rocm services restart --yes ``` -## 6. llama.cpp GPU Verification - -Install or refresh the llama.cpp engine: - -```powershell -rocm engines install llama.cpp -``` - -Run a GGUF model through llama.cpp with the managed runtime: - -```powershell -rocm serve target\models\stories260K.gguf --engine llama.cpp --managed --foreground --port 11450 -``` - -Expected result: - -- llama.cpp starts with HIP enabled. -- ROCm libraries load from the active rocm-cli managed TheRock runtime. -- The run does not fall back to CPU. -- If `llama-server` or the model file is missing, the command says what is - missing. - -Stop the foreground server with `Ctrl+C`. - -For the stricter developer GPU test: - -```powershell -python scripts\llama_cpp_therock_gpu_test.py --timeout 120 -``` - -This test downloads or reuses a tiny GGUF model, launches llama.cpp with GPU -required, checks the HTTP endpoint, and verifies that the loaded ROCm libraries -come from the managed TheRock runtime. - -## 7. ComfyUI Verification +## 5. ComfyUI Verification ComfyUI is managed as an app surface. It should start a local web server and show the URL to open: @@ -266,7 +204,7 @@ python scripts\comfyui_therock_gpu_test.py This test may download a small checkpoint and submit a cat image workflow through the ComfyUI HTTP API. -## 8. Optional Cloud Provider Key +## 6. Optional Cloud Provider Key Local ROCm use does not need a cloud provider key. If you want to test OpenAI or Anthropic provider setup, save the key through stdin so it does not land in @@ -291,7 +229,7 @@ To remove the saved key: rocm config clear-provider-key openai ``` -## 9. Optional Provider-Assisted Planning +## 7. Optional Provider-Assisted Planning Most users should leave this off. To test ambiguity resolution with an already running local provider service: diff --git a/docs/sglang.md b/docs/sglang.md deleted file mode 100644 index b43a7881..00000000 --- a/docs/sglang.md +++ /dev/null @@ -1,115 +0,0 @@ - - -# SGLang Adapter - -`rocm-engine-sglang` is a first-party adapter around an existing SGLang -installation. It does not install SGLang automatically and does not run CPU -mode. - -Use Linux or WSL with a ROCm-capable SGLang Python environment, then expose one -of these launchers to rocm-cli: - -- `ROCM_CLI_SGLANG_COMMAND=/path/to/sglang` -- `ROCM_CLI_SGLANG_PYTHON=/path/to/python` -- an active rocm-cli managed TheRock runtime that contains SGLang -- `sglang` on `PATH` - -The adapter launches either: - -```bash -sglang serve --model-path --host --port --attention-backend triton -``` - -or, when only a Python environment is configured: - -```bash -python -m sglang.launch_server --model-path --host --port --attention-backend triton -``` - -The Triton attention backend is the rocm-cli default so basic SGLang serving -does not require a separately built AITER package. AITER remains an upstream -SGLang ROCm dependency for AITER-specific attention, MoE, and quantized paths. - -Smoke commands: - -```bash -rocm-engine-sglang detect -rocm-engine-sglang capabilities -rocm-engine-sglang resolve-model Qwen/Qwen3.5-4B --device-policy gpu_required -python scripts/sglang_therock_gpu_test.py --self-test -``` - -Managed serving: - -```bash -rocm serve Qwen/Qwen3.5-4B --engine sglang --device gpu_required --managed -``` - -### GPU selection - -Use `--gpu` to choose the AMD GPU SGLang runs on: - -```bash -# Default: first free GPU (auto) -rocm serve Qwen/Qwen3.5-4B --engine sglang --managed - -# Pin a specific GPU -rocm serve Qwen/Qwen3.5-4B --engine sglang --gpu 1 --managed -``` - -rocm-cli pins the device via `HIP_VISIBLE_DEVICES`. Serving one model across -multiple GPUs is not supported. - -Native Windows SGLang serving is skipped in this adapter. Use WSL/Linux for -SGLang ROCm serving, or choose a different engine explicitly. No CPU fallback is -used. - -## TheRock And RDNA4 Status - -The adapter can resolve SGLang from a rocm-cli managed TheRock runtime and will -report that as `managed_env: true`. When a managed runtime is used, service -state records the TheRock SDK root/bin paths so acceptance tests can verify -that HIP libraries were loaded from the managed SDK wheel directories. - -Live SGLang ROCm acceptance is currently gated on upstream SGLang kernel -support for the host GPU. On this WSL test host (`gfx1201`, Radeon RX 9070 XT), -both SGLang `v0.5.12` and current `origin/main` reject the ROCm kernel build in -`sgl-kernel/setup_rocm.py` with a supported-architecture check for `gfx942` and -`gfx950` only. rocm-cli must not force a different target or fall back to CPU -mode. Re-run live SGLang GPU acceptance only on a supported SGLang ROCm target, -or after upstream adds support for the host gfx target. - -On the MI300X/gfx942 TheRock 7.13 runtime, SGLang `v0.5.12.post1` was installed -from source with `python/pyproject_other.toml`, `sgl-kernel/setup_rocm.py`, and -the active TheRock SDK compiler/library paths. The generic PyPI package is not a -safe automatic install path because it can resolve CUDA/NVIDIA packages. - -The AITER-free MI300X smoke path also needed source guards so optional Quark -quantization imports do not require AITER for unquantized models, and needed -SGLang's HIP layernorm path to use its native fallback instead of the vLLM -RMSNorm op when AITER is absent. With those source adjustments plus the adapter -Triton attention default, the live harness passed on -`Qwen/Qwen2.5-1.5B-Instruct` and verified HIP/BLAS libraries loaded from the -managed TheRock SDK wheel directories. - -On a supported ROCm target, run the live acceptance harness from the repo root: - -```bash -python3 scripts/sglang_therock_gpu_test.py \ - --engine /home/user/.cache/rocm-cli-target/debug/rocm-engine-sglang \ - --model Qwen/Qwen2.5-1.5B-Instruct -``` - -The harness uses the active rocm-cli managed TheRock runtime by default, -rejects external SGLang command/Python overrides, launches with -`gpu_required`, and verifies loaded HIP libraries came from the managed SDK -wheel directories. - -References: - -- SGLang launch server: https://sgl-project.github.io/basic_usage/send_request.html -- SGLang serve command: https://sgl-project-sglang-93.mintlify.app/backend/launch-server diff --git a/docs/testing.md b/docs/testing.md index 8fc28ecc..464b20ed 100644 --- a/docs/testing.md +++ b/docs/testing.md @@ -42,9 +42,7 @@ not require a managed runtime, and verifies: - telemetry-off config behavior - GPU-required recipe planning for `tiny-gpt2` - exact-runtime rejection for `rocm engines install` before any runtime exists -- direct llama.cpp external-runtime probing with a fake local `llama-server` -- explicit PyTorch `tiny-gpt2` GPU-required recipe resolution -- direct adapter and `rocm serve` CPU policy rejection +- direct vLLM adapter and `rocm serve` CPU policy rejection - GPU-required paths fail loudly instead of silently falling back - no first-run pip cache or runtime registry is created @@ -89,7 +87,7 @@ cargo test -p rocm --bin rocm permissions_ Self-hosted GPU CI smoke is intentionally non-mutating. The MI300X job builds the workspace, runs `rocm examine`, then runs `detect` and `capabilities` for -all first-party engine adapters: PyTorch, llama.cpp, ATOM, vLLM, and SGLang. +all first-party engine adapters: Lemonade and vLLM. Live serving acceptance remains separate because it needs engine-specific runtime installs, model artifacts, and supported upstream GPU targets. @@ -116,8 +114,8 @@ Ambiguous runtime selectors should also fail before engine install or serve launches: ```bash -rocm engines install pytorch -rocm serve target/models/stories260K.gguf --engine llama.cpp --managed +rocm engines install vllm +rocm serve qwen --engine vllm --managed ``` On a host where `default_runtime_id` matches multiple registered runtime keys, @@ -157,7 +155,7 @@ Then it verifies: - `python -m rocm_sdk targets` - runtime-only TheRock wheel discovery through `rocm_sdk.find_libraries("amdhip64", "hipblas")` -- llama.cpp adapter discovery of the managed TheRock HIP runtime environment +- vLLM adapter discovery of the managed TheRock runtime environment Run the full live test with auto-detected GPU family: @@ -185,72 +183,6 @@ selection: python scripts/therock_sdk_install_test.py --root .rocm-work/tests/therock-sdk-install --fresh --family gfx120X-all ``` -## PyTorch TheRock GPU Acceptance - -This opt-in test verifies that PyTorch can run a tiny model on your AMD GPU -using the TheRock ROCm wheels managed by rocm-cli. It does not use an external -Python environment and it does not fall back to CPU. - -Simple path: - -```bash -rocm -``` - -In the first-time setup screen, choose the install folder and choose -`Install ROCm`. After setup finishes, install the managed PyTorch engine: - -```bash -rocm engines install pytorch -``` - -The engine installer should use the Python version from the selected managed -TheRock runtime, pin `torch`, `torchvision`, and `torchaudio` to the exact -versions already installed in that runtime, and keep its pip cache under -rocm-cli's cache directory. It must not ask pip to solve an unbounded TheRock -torch stack. - -Then run the GPU acceptance script: - -```bash -python scripts/pytorch_therock_gpu_test.py -``` - -The script uses the active rocm-cli runtime key. If no runtime is active, run -`rocm runtimes activate ` first or pass `--runtime-id -`. A broad runtime id is only accepted when it matches one saved -runtime unambiguously. - -Developer direct path: - -```bash -rocm install sdk --channel release --format wheel -rocm runtimes activate -rocm engines install pytorch -python scripts/pytorch_therock_gpu_test.py -``` - -The script starts a local PyTorch test server in AMD GPU mode, loads -`hf-internal-testing/tiny-random-gpt2`, sends a tiny prompt, verifies the -service reports AMD GPU execution, and checks the worker process for loaded HIP -modules from the managed PyTorch folder. PyTorch reports ROCm GPUs through its -`cuda` device API; seeing `device: cuda` in the test output is expected on AMD -ROCm. Model downloads are localized under `target/test-cache/huggingface` -unless Hugging Face cache variables are already set. When -`ROCM_CLI_CACHE_DIR` or `ROCM_CLI_DATA_DIR` are set for isolated acceptance -tests, the script keeps its Hugging Face cache, state file, and logs under -those rocm-cli directories. If `CARGO_TARGET_DIR` is set, the default engine -binary is resolved from that target directory. If the managed PyTorch folder is -missing, the script stops with the exact -`rocm engines install pytorch` prerequisite instead of creating a surprise -install or using CPU. - -Offline selector sanity check: - -```bash -python scripts/pytorch_therock_gpu_test.py --self-test -``` - ## ComfyUI TheRock GPU Acceptance This opt-in test verifies that rocm-cli can install or reuse its managed @@ -391,11 +323,9 @@ cargo test -p rocm --bin rocm logs cargo test -p rocm-core model_recipe cargo test -p rocm-core load_model_recipe_index cargo test -p rocm --bin rocm update_report_policy -cargo test -p rocm-engine-pytorch stdio_protocol_routes_all_methods_without_side_effects -cargo test -p rocm-engine-llama-cpp stdio_protocol_routes_all_methods_without_side_effects -cargo test -p rocm-engine-atom +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-sglang cargo test -p rocmd event_collector cargo test -p rocmd event_dispatcher ``` @@ -458,10 +388,10 @@ Engine inventory smoke on native Windows: rocm engines list ``` -The packaged Linux/WSL-only ATOM, vLLM, and SGLang adapters should render +The packaged Linux/WSL-only vLLM adapter should render `runtime: unsupported_native_windows`, not `runtime: not found`. -The vLLM and SGLang live GPU acceptance scripts should return a clean skip on -native Windows. They remain strict GPU-required tests on Linux/WSL. +The vLLM live GPU acceptance script should return a clean skip on native +Windows. It remains a strict GPU-required test on Linux/WSL. Serve resolver focused tests: @@ -477,11 +407,8 @@ Engine recipe adapter contract focused tests: ```bash cargo test -p rocm-engine-protocol engine_recipe_hint_roundtrips_through_resolve_request cargo test -p rocm --bin rocm engine_recipe -cargo test -p rocm-engine-pytorch engine_recipe -cargo test -p rocm-engine-llama-cpp engine_recipe -cargo test -p rocm-engine-atom engine_recipe +cargo test -p rocm-engine-lemonade engine_recipe cargo test -p rocm-engine-vllm engine_recipe -cargo test -p rocm-engine-sglang engine_recipe ``` These tests verify that signed-index engine-specific recipe metadata is mapped @@ -866,34 +793,6 @@ same watcher policy paths as scheduler, managed-service recovery, GPU telemetry, and cache-warm proposal events. See `docs/automations.md` for the local curl smoke and accepted fields. -ATOM TheRock GPU acceptance: - -```bash -cargo test -p rocm-engine-atom managed_env_reflects_managed_runtime_manifest_source -cargo test -p rocm-engine-atom running_state_records_managed_therock_env_for_gpu_verification -python -m py_compile scripts/atom_therock_gpu_test.py -python scripts/atom_therock_gpu_test.py --self-test -``` - -Live ATOM ROCm serving is only valid on a GPU/runtime/model combination -supported by upstream ATOM. The current WSL validation host is `gfx1201`, while -upstream ATOM documentation currently lists Instinct `gfx950`, `gfx942`, and -`gfx90a` targets. Do not force a different target and do not use CPU fallback. -On a supported ATOM ROCm target, install/build ATOM inside a rocm-cli managed -TheRock runtime, then verify `rocm-engine-atom detect` reports -`managed_env: true`, launch with `gpu_required`, and check loaded HIP libraries -from the managed TheRock SDK wheel directories: - -```bash -python3 scripts/atom_therock_gpu_test.py \ - --engine /home/user/.cache/rocm-cli-target/debug/rocm-engine-atom \ - --model Qwen/Qwen3-0.6B -``` - -The script defaults to the active exact runtime key. If `--runtime-id` is used, -it must be an exact runtime key or an unambiguous runtime id. It rejects -external ATOM command/Python overrides and does not allow CPU fallback. - vLLM TheRock GPU acceptance: ```bash @@ -920,41 +819,6 @@ For TheRock 7.13, patch vLLM's GPTQ ROCm compatibility guard to include HIP On native Windows this script prints a JSON skip result; run it from WSL/Linux for live ROCm GPU acceptance. -SGLang TheRock GPU acceptance: - -```bash -cargo test -p rocm-engine-sglang managed_env_reflects_managed_runtime_manifest_source -cargo test -p rocm-engine-sglang running_state_records_managed_therock_env_for_gpu_verification -python -m py_compile scripts/sglang_therock_gpu_test.py -python scripts/sglang_therock_gpu_test.py --self-test -``` - -Live SGLang ROCm serving is only valid on a GPU target supported by upstream -SGLang ROCm kernels. The current WSL validation host is `gfx1201`; SGLang -`v0.5.12` and current `origin/main` reject `sgl-kernel/setup_rocm.py` for that -target and accept only `gfx942`/`gfx950`. Do not force a different target and -do not use CPU fallback. On a supported SGLang ROCm target, install/build -SGLang inside a rocm-cli managed TheRock runtime, then verify -`rocm-engine-sglang detect` reports `managed_env: true`, launch with -`gpu_required`, and check loaded HIP libraries from the managed TheRock SDK -wheel directories: - -```bash -python3 scripts/sglang_therock_gpu_test.py \ - --engine /home/user/.cache/rocm-cli-target/debug/rocm-engine-sglang \ - --model Qwen/Qwen2.5-1.5B-Instruct -``` - -The script defaults to the active exact runtime key. If `--runtime-id` is used, -it must be an exact runtime key or an unambiguous runtime id. It rejects -external SGLang command/Python overrides and does not allow CPU fallback. -For MI300X/gfx942 TheRock 7.13, install SGLang from source with -`python/pyproject_other.toml` and the ROCm `sgl-kernel` wheel. Use the -rocm-cli adapter's Triton attention default unless AITER has been built and -verified for the runtime. -On native Windows this script prints a JSON skip result; run it from WSL/Linux -for live ROCm GPU acceptance. - ## Windows Tool Notes The TheRock SDK wheel install path should not require users to install global @@ -968,61 +832,6 @@ SDK wheel setup should avoid global source-build tools. Reference: [TheRock Windows install tools](https://github.com/ROCm/TheRock/blob/main/docs/development/windows_support.md#install-tools) -## llama.cpp TheRock GPU Test - -This opt-in test requires a HIP-enabled `llama-server`, a rocm-cli managed -TheRock runtime, and an AMD GPU. It does not allow CPU fallback. - -```bash -python scripts/llama_cpp_therock_gpu_test.py --llama-server target/llama.cpp-build-hip/bin/llama-server.exe -``` - -To verify the background `launch` command also returns promptly when its JSON -output is captured by a shell or test harness: - -```bash -python scripts/llama_cpp_therock_gpu_test.py --launch-mode launch --llama-server target/llama.cpp-build-hip/bin/llama-server.exe -``` - -WSL command shape, using the WSL-built adapter and server: - -```bash -python3 scripts/llama_cpp_therock_gpu_test.py \ - --engine /home/user/.cache/rocm-cli-target/debug/rocm-engine-llama-cpp \ - --llama-server /home/user/.cache/rocm-cli-llama.cpp-build-hip/bin/llama-server \ - --model-path /mnt/d/path/to/rocm-cli/target/models/stories260K.gguf \ - --timeout 120 -``` - -The script uses the active rocm-cli runtime key. If no runtime is active, run -`rocm runtimes activate ` first or pass `--runtime-id -`. A broad runtime id is only accepted when it matches one saved -runtime unambiguously; the test no longer guesses by picking the newest -manifest. When testing isolated state, set `ROCM_CLI_CONFIG_DIR` and -`ROCM_CLI_DATA_DIR`; the script should honor those directories instead of -looking in the default `~/.rocm` state. If `CARGO_TARGET_DIR` is set, the -default adapter binary is resolved from that target directory. It downloads or -reuses the tiny `stories260K.gguf` model, launches the llama.cpp adapter with -`gpu_required`, checks `/health` and `/v1/completions`, and on Windows verifies -that HIP DLLs load from the staged rocm-cli TheRock runtime rather than -`System32`. On Linux/WSL it verifies HIP and BLAS shared objects are loaded -from the managed TheRock SDK wheel roots via `/proc//maps`, including -split `_rocm_sdk_core` and `_rocm_sdk_libraries_*` wheel directories. - -Offline selector sanity check: - -```bash -python scripts/llama_cpp_therock_gpu_test.py --self-test -``` - -Top-level managed serving smoke: - -```bash -# Set ROCM_CLI_LLAMA_CPP_SERVER if llama-server is not on PATH. -rocm serve target/models/stories260K.gguf --engine llama.cpp --managed --port 11450 -rocmd sandbox-run stop_server --service-id --allow-native-fallback -``` - ## WSL Preflight Read-only WSL/ROCDXG preflight: diff --git a/docs/wsl.md b/docs/wsl.md index 692694c6..1b9a90a1 100644 --- a/docs/wsl.md +++ b/docs/wsl.md @@ -121,8 +121,8 @@ export LD_LIBRARY_PATH=":/usr/lib/wsl/lib${LD_LIBR ``` For `rocm-cli`, the command itself should resolve the managed runtime manifest -and apply that environment before launching non-PyTorch HIP apps such as -`llama.cpp`. Users should not have to hand-export these values. +and apply that environment before launching HIP apps such as Lemonade's bundled +`llama.cpp` backend. Users should not have to hand-export these values. ## Examine And Install UX Recommendations @@ -154,8 +154,8 @@ and apply that environment before launching non-PyTorch HIP apps such as 2. Explain missing ROCDXG if absent. 3. Ask before any `sudo apt install` or `sudo make install`. 4. Install TheRock into a managed venv. -5. Install/build HIP `llama.cpp`. -6. Run a tiny GGUF GPU smoke test with CPU fallback disabled. +5. Install a serving engine (Lemonade or vLLM). +6. Run a tiny GPU smoke test with CPU fallback disabled. ## Non-Destructive Tests @@ -173,5 +173,5 @@ need `sudo`: - install ROCDXG `.deb` - build ROCDXG from source - install TheRock wheels into a fresh managed WSL venv -- build HIP `llama.cpp` -- run tiny GGUF inference on GPU +- install a serving engine (Lemonade or vLLM) +- run tiny inference on GPU diff --git a/engines/atom/Cargo.toml b/engines/atom/Cargo.toml deleted file mode 100644 index 131994e9..00000000 --- a/engines/atom/Cargo.toml +++ /dev/null @@ -1,23 +0,0 @@ -[package] -name = "rocm-engine-atom" -version.workspace = true -edition.workspace = true -license.workspace = true -repository.workspace = true -rust-version.workspace = true -publish.workspace = true - -[lints] -workspace = true - -[[bin]] -name = "rocm-engine-atom" -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 diff --git a/engines/atom/src/lib.rs b/engines/atom/src/lib.rs deleted file mode 100644 index 4e63d626..00000000 --- a/engines/atom/src/lib.rs +++ /dev/null @@ -1,1685 +0,0 @@ -// Copyright © Advanced Micro Devices, Inc., or its affiliates. -// -// SPDX-License-Identifier: MIT - -use anyhow::{Context, Result, bail}; -use clap::{Parser, Subcommand}; -use rocm_core::{ - AppPaths, DEFAULT_LOCAL_PORT, format_http_base_url, openai_models_endpoint_has_model, - require_nonempty, -}; -use rocm_engine_protocol::{ - DEFAULT_LOG_TAIL_LINES, DetectRequest, DetectResponse, DevicePolicy, - ENGINE_RECIPE_CONTRACT_VERSION, EndpointRequest, EndpointResponse, EngineCapabilities, - EngineDeviceAvailability, EngineMethod, EngineRecipeHint, EngineRequestEnvelope, - EngineResponseEnvelope, GpuSelection, HealthcheckRequest, HealthcheckResponse, InstallRequest, - InstallResponse, LaunchRequest, LaunchResponse, LogsRequest, LogsResponse, ResolveModelRequest, - ResolveModelResponse, StopRequest, StopResponse, -}; -use serde::de::DeserializeOwned; -use serde::{Deserialize, Serialize}; -use serde_json::{Value, json}; -use std::collections::hash_map::DefaultHasher; -use std::ffi::OsString; -use std::fs; -use std::hash::{Hash, Hasher}; -use std::io::{Read, Write}; -use std::path::{Path, PathBuf}; -use std::process::{Command as ProcessCommand, Stdio}; -use std::time::{Duration, SystemTime, UNIX_EPOCH}; - -const ENGINE_NAME: &str = "atom"; -const DEFAULT_HOST: &str = "127.0.0.1"; -const HEALTHCHECK_TIMEOUT_MS: u64 = 700; - -#[derive(Parser, Debug)] -#[command(name = "rocm-engine-atom", about = "rocm-cli ATOM engine adapter")] -struct Cli { - #[command(subcommand)] - command: CommandKind, -} - -#[derive(Subcommand, Debug)] -enum CommandKind { - Detect, - Capabilities, - Install { - #[arg(long, default_value = "external-atom")] - runtime_id: String, - #[arg(long)] - reinstall: bool, - }, - ResolveModel { - model_ref: String, - #[arg(long)] - device_policy: Option, - }, - Launch { - service_id: String, - model_ref: String, - #[arg(long, default_value = DEFAULT_HOST)] - host: String, - #[arg(long, default_value_t = DEFAULT_LOCAL_PORT)] - port: u16, - #[arg(long)] - device_policy: Option, - #[arg(long)] - runtime_id: Option, - #[arg(long)] - env_id: Option, - #[arg(long)] - gpu: Option, - }, - Stdio, - ServeHttp { - service_id: String, - model_ref: String, - #[arg(long, default_value = DEFAULT_HOST)] - host: String, - #[arg(long, default_value_t = DEFAULT_LOCAL_PORT)] - port: u16, - #[arg(long, default_value = "gpu_required")] - device_policy: String, - #[arg(long)] - runtime_id: Option, - #[arg(long)] - env_id: Option, - #[arg(long)] - state_path: PathBuf, - #[arg(long)] - engine_recipe_json: Option, - #[arg(long)] - gpu: Option, - }, -} - -#[derive(Debug, Clone)] -struct AtomRuntime { - runtime_id: String, - env_id: String, - command: PathBuf, - launcher: AtomLauncher, - python_executable: Option, - version: Option, - source: String, - sdk_root: Option, - sdk_bin: Option, - sdk_bin_paths: Vec, - sdk_library_paths: Vec, -} - -#[derive(Debug, Clone, Copy, Eq, PartialEq)] -enum AtomLauncher { - Command, - PythonModule, -} - -#[derive(Debug, Clone, Deserialize)] -struct TheRockRuntimeManifest { - #[serde(default)] - runtime_key: Option, - #[serde(default)] - runtime_id: Option, - #[serde(default)] - python_executable: Option, - #[serde(default)] - rocm_sdk: Option, - #[serde(default)] - installed_at_unix_ms: Option, -} - -#[derive(Debug, Clone, Deserialize)] -struct RocmSdkRuntimeProbe { - #[serde(default)] - import_ok: bool, - #[serde(default)] - root_path: Option, - #[serde(default)] - bin_path: Option, - #[serde(default)] - bin_paths: Vec, - #[serde(default)] - library_paths: Vec, -} - -#[derive(Debug, Clone)] -struct ServiceFiles { - state_path: PathBuf, - log_path: PathBuf, -} - -pub fn run_cli() -> Result<()> { - let cli = Cli::parse(); - match cli.command { - CommandKind::Detect => print_json(&detect_response())?, - CommandKind::Capabilities => print_json(&capabilities())?, - CommandKind::Install { - runtime_id, - reinstall, - } => print_json(&install_response(InstallRequest { - runtime_id, - python_version: None, - env_root: None, - reinstall, - })?)?, - CommandKind::ResolveModel { - model_ref, - device_policy, - } => print_json(&resolve_model_response(ResolveModelRequest { - model_ref, - runtime_id: None, - device_policy: device_policy - .as_deref() - .map(|value| parse_device_policy_arg(Some(value))) - .transpose()?, - recipe_override: None, - engine_recipe: None, - })?)?, - CommandKind::Launch { - service_id, - model_ref, - host, - port, - device_policy, - runtime_id, - env_id, - gpu, - } => print_json(&launch_service(LaunchRequest { - service_id, - env_id, - runtime_id, - model_ref, - host, - port, - device_policy: Some(parse_device_policy_arg(device_policy.as_deref())?), - endpoint_mode: Some("openai".to_owned()), - engine_recipe: None, - gpu_selection: parse_gpu_selection_arg(gpu.as_deref())?, - })?)?, - CommandKind::Stdio => { - let envelope = read_request()?; - print_json(&handle_envelope(envelope))?; - } - CommandKind::ServeHttp { - service_id, - model_ref, - host, - port, - device_policy, - runtime_id, - env_id, - state_path, - engine_recipe_json, - gpu, - } => serve_http(ServeHttpRequest { - service_id, - model_ref, - host, - port, - device_policy: parse_device_policy_arg(Some(&device_policy))?, - gpu_indices: parse_gpu_indices_arg(gpu.as_deref())?, - runtime_id, - env_id, - state_path, - engine_recipe: parse_engine_recipe_json(engine_recipe_json)?, - })?, - } - Ok(()) -} - -pub fn builtin_handle_envelope(envelope: EngineRequestEnvelope) -> EngineResponseEnvelope { - handle_envelope(envelope) -} - -#[allow(clippy::too_many_arguments)] -pub fn builtin_serve_http( - service_id: String, - model_ref: String, - host: String, - port: u16, - device_policy: DevicePolicy, - gpu_indices: Vec, - runtime_id: Option, - env_id: Option, - state_path: PathBuf, - engine_recipe: Option, -) -> Result<()> { - serve_http(ServeHttpRequest { - service_id, - model_ref, - host, - port, - device_policy, - gpu_indices, - runtime_id, - env_id, - state_path, - engine_recipe, - }) -} - -fn handle_envelope(envelope: EngineRequestEnvelope) -> EngineResponseEnvelope { - match envelope.method { - EngineMethod::Detect => { - deserialize_and_respond::(envelope.payload, |_| { - Ok(detect_response()) - }) - } - EngineMethod::Capabilities => EngineResponseEnvelope::success(capabilities()), - EngineMethod::Install => { - deserialize_and_respond::(envelope.payload, install_response) - } - EngineMethod::ResolveModel => deserialize_and_respond::( - envelope.payload, - resolve_model_response, - ), - EngineMethod::Launch => { - deserialize_and_respond::(envelope.payload, launch_service) - } - EngineMethod::Healthcheck => deserialize_and_respond::( - envelope.payload, - healthcheck_service, - ), - EngineMethod::Endpoint => { - deserialize_and_respond::(envelope.payload, endpoint_response) - } - EngineMethod::Stop => { - deserialize_and_respond::(envelope.payload, stop_service) - } - EngineMethod::Logs => { - deserialize_and_respond::(envelope.payload, logs_response) - } - } -} - -fn deserialize_and_respond(payload: Value, handler: F) -> EngineResponseEnvelope -where - T: DeserializeOwned, - F: FnOnce(T) -> Result, - U: Serialize, -{ - match serde_json::from_value::(payload) { - Ok(request) => match handler(request) { - Ok(response) => EngineResponseEnvelope::success(response), - Err(error) => EngineResponseEnvelope::failure("request_failed", error.to_string()), - }, - Err(error) => EngineResponseEnvelope::failure("invalid_payload", error.to_string()), - } -} - -fn detect_response() -> DetectResponse { - let runtime = resolve_atom_runtime(None); - let installed = runtime.is_ok(); - let mut notes = Vec::new(); - if cfg!(windows) { - notes.push(windows_unsupported_message().to_owned()); - } else if let Err(error) = runtime.as_ref() { - notes.push(error.to_string()); - } - let runtime = runtime.ok(); - if let Some(runtime) = runtime.as_ref() { - notes.push(format!( - "ATOM Python launcher resolved from {}; no CPU fallback is used", - runtime.source - )); - } - - DetectResponse { - installed, - env_id: runtime.as_ref().map(|runtime| runtime.env_id.clone()), - runtime_kind: Some("external_atom".to_owned()), - runtime_executable: runtime - .as_ref() - .map(|runtime| runtime.command.display().to_string()), - managed_env: Some(runtime.as_ref().is_some_and(runtime_is_managed)), - python_version: runtime - .as_ref() - .and_then(|runtime| runtime.version.as_ref()) - .map(|version| format!("ATOM {version}")), - torch_version: None, - transformers_version: None, - available_devices: vec![EngineDeviceAvailability { - kind: "rocm_gpu".to_owned(), - available: installed && !cfg!(windows), - reason: if cfg!(windows) { - Some(windows_unsupported_message().to_owned()) - } else if installed { - None - } else { - Some("ATOM Python environment was not found on Linux/WSL".to_owned()) - }, - }], - capabilities: capabilities(), - notes, - } -} - -fn capabilities() -> EngineCapabilities { - EngineCapabilities { - cpu: false, - rocm_gpu: !cfg!(windows), - openai_compatible: true, - tool_calling: false, - quantized_models: "ATOM-supported".to_owned(), - reasoning_parser: false, - } -} - -fn install_response(request: InstallRequest) -> Result { - let runtime = resolve_atom_runtime(Some(&request.runtime_id))?; - let env_path = runtime - .command - .parent() - .map_or_else(|| PathBuf::from("."), Path::to_path_buf); - Ok(InstallResponse { - env_id: runtime.env_id.clone(), - env_path: env_path.display().to_string(), - python_executable: runtime - .python_executable - .as_ref() - .unwrap_or(&runtime.command) - .display() - .to_string(), - runtime_kind: Some("external_atom".to_owned()), - runtime_executable: Some(runtime.command.display().to_string()), - managed_env: Some(runtime_is_managed(&runtime)), - installed_packages: vec![format!( - "ATOM{}", - runtime - .version - .as_deref() - .map(|version| format!("=={version}")) - .unwrap_or_default() - )], - capabilities: capabilities(), - lock_hash: runtime_lock_hash(&runtime), - warnings: atom_runtime_warnings(&runtime), - }) -} - -fn runtime_is_managed(runtime: &AtomRuntime) -> bool { - runtime.source.starts_with("managed_runtime_manifest") -} - -fn atom_runtime_warnings(runtime: &AtomRuntime) -> Vec { - let runtime_scope = if runtime_is_managed(runtime) { - "rocm-cli records this ATOM Python environment from a managed TheRock runtime; it does not pip install ATOM automatically" - } else { - "rocm-cli records this as an external ATOM runtime; it does not pip install ATOM automatically" - }; - vec![ - runtime_scope.to_owned(), - "ATOM serving remains ROCm GPU required; no CPU fallback is used".to_owned(), - ] -} - -fn resolve_model_response(request: ResolveModelRequest) -> Result { - let device_policy = normalize_atom_device_policy(request.device_policy)?; - let engine_recipe = accepted_engine_recipe(request.engine_recipe)?; - Ok(ResolveModelResponse { - canonical_model_id: request.model_ref, - task: "text-generation".to_owned(), - source: "huggingface_or_local".to_owned(), - revision: "main".to_owned(), - loader: "atom".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": DEFAULT_HOST, - "port": DEFAULT_LOCAL_PORT - }), - engine_recipe, - warnings: vec![ - "ATOM is treated as a ROCm GPU engine in rocm-cli; select another engine explicitly for CPU serving".to_owned(), - ], - }) -} - -fn accepted_engine_recipe( - engine_recipe: Option, -) -> Result> { - if let Some(hint) = &engine_recipe { - if hint.engine != ENGINE_NAME { - bail!( - "engine_recipe target `{}` does not match adapter `{}`", - hint.engine, - ENGINE_NAME - ); - } - if hint.contract_version != ENGINE_RECIPE_CONTRACT_VERSION { - bail!( - "engine_recipe contract `{}` is unsupported; expected `{}`", - hint.contract_version, - ENGINE_RECIPE_CONTRACT_VERSION - ); - } - } - Ok(engine_recipe) -} - -fn parse_engine_recipe_json(value: Option) -> Result> { - value - .map(|text| { - serde_json::from_str::(&text) - .context("failed to parse engine recipe JSON") - }) - .transpose() - .and_then(accepted_engine_recipe) -} - -fn launch_service(request: LaunchRequest) -> Result { - let device_policy = normalize_atom_device_policy(request.device_policy)?; - let engine_recipe = accepted_engine_recipe(request.engine_recipe)?; - let runtime = resolve_atom_runtime(request.runtime_id.as_deref())?; - let state_path = AppPaths::discover()? - .engine_state_dir(ENGINE_NAME) - .join(format!("{}.json", request.service_id)); - let serve_request = ServeHttpRequest { - service_id: request.service_id.clone(), - model_ref: request.model_ref.clone(), - host: request.host.clone(), - port: request.port, - device_policy, - gpu_indices: rocm_engine_protocol::launch_gpu_indices(request.gpu_selection.as_ref()), - runtime_id: request.runtime_id.clone(), - env_id: request.env_id.clone(), - state_path: state_path.clone(), - engine_recipe, - }; - let log_path = AppPaths::discover()? - .engine_logs_dir(ENGINE_NAME) - .join(format!("{}.log", request.service_id)); - let child = spawn_atom_server(&serve_request, &runtime, Some(&log_path))?; - let pid = child.id(); - write_running_state(&serve_request, &runtime, pid)?; - Ok(LaunchResponse { - service_id: request.service_id, - pid, - endpoint_url: endpoint_url(&request.host, request.port), - log_path: log_path.display().to_string(), - state_path: state_path.display().to_string(), - }) -} - -#[derive(Debug, Clone)] -struct ServeHttpRequest { - service_id: String, - model_ref: String, - host: String, - port: u16, - device_policy: DevicePolicy, - gpu_indices: Vec, - runtime_id: Option, - env_id: Option, - state_path: PathBuf, - engine_recipe: Option, -} - -fn serve_http(request: ServeHttpRequest) -> Result<()> { - let runtime = resolve_atom_runtime(request.runtime_id.as_deref())?; - let mut child = spawn_atom_server(&request, &runtime, None)?; - write_running_state(&request, &runtime, child.id())?; - let status = child.wait().context("failed waiting for ATOM server")?; - write_terminal_state( - &request.state_path, - if status.success() { - "stopped" - } else { - "failed" - }, - )?; - if status.success() { - Ok(()) - } else { - std::process::exit(status.code().unwrap_or(1)); - } -} - -fn spawn_atom_server( - request: &ServeHttpRequest, - runtime: &AtomRuntime, - log_path: Option<&Path>, -) -> Result { - require_nonempty(&request.service_id, "service_id")?; - require_nonempty(&request.model_ref, "model_ref")?; - if !matches!(request.device_policy, DevicePolicy::GpuRequired) { - bail!("ATOM launch requires ROCm GPU execution; no CPU fallback is used"); - } - - if let Some(parent) = request.state_path.parent() { - fs::create_dir_all(parent) - .with_context(|| format!("failed to create {}", parent.display()))?; - } - if let Some(log_path) = log_path - && let Some(parent) = log_path.parent() - { - fs::create_dir_all(parent) - .with_context(|| format!("failed to create {}", parent.display()))?; - } - - let mut command = ProcessCommand::new(&runtime.command); - command - .args(atom_server_args( - runtime.launcher, - &request.model_ref, - &request.host, - request.port, - )) - .args(engine_recipe_launch_args(request.engine_recipe.as_ref())) - .stdin(Stdio::null()); - apply_therock_env(&mut command, runtime)?; - rocm_engine_protocol::apply_gpu_visibility(&mut command, &request.gpu_indices); - if let Some(log_path) = log_path { - let log = fs::File::create(log_path) - .with_context(|| format!("failed to create {}", log_path.display()))?; - command.stdout(Stdio::from( - log.try_clone().context("failed to clone log handle")?, - )); - command.stderr(Stdio::from(log)); - } - - command - .spawn() - .with_context(|| format!("failed to spawn ATOM command {}", runtime.command.display())) -} - -fn atom_server_args(launcher: AtomLauncher, model_ref: &str, host: &str, port: u16) -> Vec { - let mut args = Vec::new(); - match launcher { - AtomLauncher::Command => {} - AtomLauncher::PythonModule => { - args.push("-m".to_owned()); - args.push("atom.entrypoints.openai_server".to_owned()); - } - } - args.extend([ - "--model".to_owned(), - model_ref.to_owned(), - "--host".to_owned(), - host.to_owned(), - "--port".to_owned(), - port.to_string(), - ]); - args -} - -fn engine_recipe_launch_args(engine_recipe: Option<&EngineRecipeHint>) -> Vec { - engine_recipe - .map(|hint| hint.required_flags.clone()) - .unwrap_or_default() -} - -fn healthcheck_service(request: HealthcheckRequest) -> Result { - require_nonempty(&request.service_id, "service_id")?; - let files = service_files(&request.service_id)?; - let state = read_service_state(&files.state_path).ok(); - let endpoint_url = state.as_ref().and_then(endpoint_url_from_state); - let model_ref = state - .as_ref() - .and_then(|value| value_string(value, "model_ref")); - let ready = endpoint_url - .as_deref() - .map(|endpoint| query_loaded_model_endpoint(endpoint, model_ref.as_deref())) - .transpose() - .unwrap_or(None) - .unwrap_or(false); - let status = if ready { - "ready".to_owned() - } else { - state - .as_ref() - .and_then(|value| value_string(value, "status")) - .unwrap_or_else(|| "unknown".to_owned()) - }; - Ok(HealthcheckResponse { - status, - model_loaded: ready, - device: if state.is_some() { - "rocm_gpu".to_owned() - } else { - "unknown".to_owned() - }, - uptime_sec: 0, - queue_depth: 0, - last_error: None, - tokens_per_sec: None, - }) -} - -fn endpoint_response(request: EndpointRequest) -> Result { - require_nonempty(&request.service_id, "service_id")?; - let files = service_files(&request.service_id)?; - let state = read_service_state(&files.state_path) - .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(), - ], - }) -} - -fn logs_response(request: LogsRequest) -> Result { - require_nonempty(&request.service_id, "service_id")?; - let files = service_files(&request.service_id)?; - let limit = request.tail_lines.unwrap_or(DEFAULT_LOG_TAIL_LINES); - Ok(LogsResponse { - log_path: files.log_path.display().to_string(), - recent_lines: if files.log_path.is_file() { - tail_lines(&files.log_path, limit)? - } else { - Vec::new() - }, - }) -} - -fn stop_service(request: StopRequest) -> Result { - require_nonempty(&request.service_id, "service_id")?; - let files = service_files(&request.service_id)?; - let state = read_service_state(&files.state_path).ok(); - let stopped = match state.as_ref().and_then(pid_from_state) { - Some(pid) => terminate_pid(pid, request.force), - None => false, - }; - if stopped { - write_terminal_state(&files.state_path, "stopped")?; - } - Ok(StopResponse { - stopped, - graceful: stopped && !request.force, - }) -} - -fn resolve_atom_runtime(runtime_id: Option<&str>) -> Result { - if cfg!(windows) { - bail!("{}", windows_unsupported_message()); - } - - if let Some(command) = std::env::var_os("ROCM_CLI_ATOM_COMMAND") - .or_else(|| std::env::var_os("ATOM_COMMAND")) - .map(PathBuf::from) - { - let command = resolve_command_path(&command)?; - return Ok(AtomRuntime { - runtime_id: runtime_id.unwrap_or("external-atom").to_owned(), - env_id: "external-atom-command".to_owned(), - command, - launcher: AtomLauncher::Command, - python_executable: None, - version: None, - source: "environment command".to_owned(), - sdk_root: None, - sdk_bin: None, - sdk_bin_paths: Vec::new(), - sdk_library_paths: Vec::new(), - }); - } - - if let Some(python) = std::env::var_os("ROCM_CLI_ATOM_PYTHON") - .or_else(|| std::env::var_os("ATOM_PYTHON")) - .map(PathBuf::from) - .filter(|path| path.is_file()) - { - return runtime_from_python( - python, - runtime_id.unwrap_or("external-atom-python"), - "environment python", - None, - None, - Vec::new(), - Vec::new(), - ); - } - - if let Some(runtime) = resolve_managed_runtime(runtime_id)? { - return Ok(runtime); - } - - if let Some(command) = find_command_on_path("atom") { - return Ok(AtomRuntime { - runtime_id: runtime_id.unwrap_or("external-atom-path").to_owned(), - env_id: "external-atom-path".to_owned(), - command, - launcher: AtomLauncher::Command, - python_executable: None, - version: None, - source: "PATH".to_owned(), - sdk_root: None, - sdk_bin: None, - sdk_bin_paths: Vec::new(), - sdk_library_paths: Vec::new(), - }); - } - - bail!( - "ATOM Python environment was not found. Install/build ATOM in a Linux or WSL ROCm Python environment, then set ROCM_CLI_ATOM_PYTHON or install it into the active rocm-cli TheRock runtime. No CPU fallback is used." - ) -} - -fn runtime_from_python( - python: PathBuf, - runtime_id: &str, - source: &str, - sdk_root: Option, - sdk_bin: Option, - sdk_bin_paths: Vec, - sdk_library_paths: Vec, -) -> Result { - let version = probe_atom_version(&python) - .with_context(|| format!("ATOM package not found in {}", python.display()))? - .unwrap_or_else(|| "unknown".to_owned()); - let (command, launcher) = atom_command_from_python(&python).map_or_else( - || (python.clone(), AtomLauncher::PythonModule), - |command| (command, AtomLauncher::Command), - ); - Ok(AtomRuntime { - runtime_id: runtime_id.to_owned(), - env_id: format!("external-atom-{}", stable_id_component(runtime_id)), - command, - launcher, - python_executable: Some(python), - version: Some(version), - source: source.to_owned(), - sdk_root, - sdk_bin, - sdk_bin_paths, - sdk_library_paths, - }) -} - -fn resolve_managed_runtime(runtime_id: Option<&str>) -> Result> { - let paths = AppPaths::discover()?; - let registry = paths.data_dir.join("runtimes").join("registry"); - if !registry.is_dir() { - return Ok(None); - } - let mut manifests = Vec::new(); - for entry in - fs::read_dir(®istry).with_context(|| format!("failed to read {}", registry.display()))? - { - let path = entry?.path(); - if path.extension().and_then(|value| value.to_str()) != Some("json") { - continue; - } - let bytes = - fs::read(&path).with_context(|| format!("failed to read {}", path.display()))?; - let Ok(manifest) = serde_json::from_slice::(&bytes) else { - continue; - }; - if !runtime_matches(&manifest, runtime_id) { - continue; - } - manifests.push((manifest.installed_at_unix_ms.unwrap_or(0), manifest)); - } - manifests.sort_by_key(|(installed_at, _)| std::cmp::Reverse(*installed_at)); - - for (_, manifest) in manifests { - let Some(python) = manifest - .python_executable - .clone() - .filter(|path| path.is_file()) - else { - continue; - }; - let runtime_id = manifest - .runtime_id - .as_deref() - .unwrap_or("therock-atom-runtime") - .to_owned(); - let source = manifest.runtime_key.as_deref().map_or_else( - || "managed_runtime_manifest".to_owned(), - |key| format!("managed_runtime_manifest:{key}"), - ); - let (sdk_root, sdk_bin, sdk_bin_paths, sdk_library_paths) = manifest - .rocm_sdk - .as_ref() - .filter(|probe| probe.import_ok) - .map_or((None, None, Vec::new(), Vec::new()), |probe| { - ( - probe.root_path.clone(), - probe.bin_path.clone(), - probe.bin_paths.clone(), - probe.library_paths.clone(), - ) - }); - if let Ok(runtime) = runtime_from_python( - python, - &runtime_id, - &source, - sdk_root, - sdk_bin, - sdk_bin_paths, - sdk_library_paths, - ) { - return Ok(Some(runtime)); - } - } - Ok(None) -} - -fn runtime_matches(manifest: &TheRockRuntimeManifest, requested: Option<&str>) -> bool { - let Some(requested) = requested.map(str::trim).filter(|value| !value.is_empty()) else { - return true; - }; - let requested = requested.to_ascii_lowercase(); - if requested == "external" || requested == "external-atom" { - return false; - } - for candidate in [ - manifest.runtime_id.as_deref(), - manifest.runtime_key.as_deref(), - ] - .into_iter() - .flatten() - { - let candidate = candidate.to_ascii_lowercase(); - if candidate == requested || candidate.starts_with(&requested) { - return true; - } - } - false -} - -fn normalize_atom_device_policy(policy: Option) -> Result { - match policy.unwrap_or(DevicePolicy::GpuRequired) { - DevicePolicy::GpuRequired => Ok(DevicePolicy::GpuRequired), - DevicePolicy::GpuPreferred => Ok(DevicePolicy::GpuRequired), - DevicePolicy::CpuOnly => { - bail!("ATOM adapter is ROCm GPU-only in rocm-cli; no CPU fallback is used") - } - } -} - -fn parse_device_policy_arg(policy: Option<&str>) -> Result { - match policy.unwrap_or("gpu_required") { - "gpu" | "gpu_required" => Ok(DevicePolicy::GpuRequired), - "gpu_preferred" => Ok(DevicePolicy::GpuPreferred), - "cpu" | "cpu_only" => Ok(DevicePolicy::CpuOnly), - other => bail!("unsupported device policy: {other}"), - } -} - -/// Parse a `--gpu` CLI value into an optional `GpuSelection` for `LaunchRequest`. -fn parse_gpu_selection_arg(value: Option<&str>) -> Result> { - value - .map(|raw| GpuSelection::parse_cli_value(raw).map_err(anyhow::Error::msg)) - .transpose() -} - -/// Parse a `--gpu` CLI value into explicit device ordinals (empty for `auto`). -fn parse_gpu_indices_arg(value: Option<&str>) -> Result> { - Ok(rocm_engine_protocol::launch_gpu_indices( - parse_gpu_selection_arg(value)?.as_ref(), - )) -} - -fn atom_command_from_python(python: &Path) -> Option { - let dir = python.parent()?; - candidate_command_names("atom") - .into_iter() - .map(|name| dir.join(name)) - .find(|path| path.is_file()) -} - -fn find_command_on_path(name: &str) -> Option { - let path = std::env::var_os("PATH")?; - for dir in std::env::split_paths(&path) { - for candidate in candidate_command_names(name) { - let path = dir.join(candidate); - if path.is_file() { - return Some(path); - } - } - } - None -} - -fn resolve_command_path(command: &Path) -> Result { - if command.components().count() > 1 || command.is_absolute() { - if command.is_file() { - return Ok(command.to_path_buf()); - } - bail!( - "configured ATOM command is not a file: {}", - command.display() - ); - } - find_command_on_path(&command.display().to_string()).with_context(|| { - format!( - "configured ATOM command `{}` was not found on PATH", - command.display() - ) - }) -} - -fn candidate_command_names(name: &str) -> Vec { - if cfg!(windows) { - vec![ - format!("{name}.exe"), - format!("{name}.cmd"), - name.to_owned(), - ] - } else { - vec![name.to_owned()] - } -} - -fn probe_atom_version(python: &Path) -> Result> { - let script = r#"import importlib.metadata, importlib.util, json -spec = importlib.util.find_spec("atom") -version = None -if spec is not None: - try: - version = importlib.metadata.version("atom") - except importlib.metadata.PackageNotFoundError: - version = "unknown" -print(json.dumps({"present": spec is not None, "version": version})) -"#; - let output = ProcessCommand::new(python) - .arg("-c") - .arg(script) - .output() - .with_context(|| format!("failed to probe ATOM with {}", python.display()))?; - if !output.status.success() { - bail!( - "ATOM probe failed: {}", - String::from_utf8_lossy(&output.stderr).trim() - ); - } - let value: Value = serde_json::from_slice(&output.stdout).context("invalid ATOM probe JSON")?; - if value - .get("present") - .and_then(Value::as_bool) - .unwrap_or(false) - { - Ok(value - .get("version") - .and_then(Value::as_str) - .map(str::to_owned)) - } else { - bail!("Python environment does not contain the ATOM package") - } -} - -fn apply_therock_env(command: &mut ProcessCommand, runtime: &AtomRuntime) -> Result<()> { - let Some(root) = runtime.sdk_root.as_ref() else { - return Ok(()); - }; - let bin = runtime.sdk_bin.as_ref(); - command - .env("ROCM_SDK_ROOT", root) - .env("ROCM_PATH", root) - .env("ROCM_HOME", root) - .env("HIP_PATH", root) - .env("ROCM_CLI_THEROCK_RUNTIME_ID", &runtime.runtime_id); - if let Some(bin) = bin { - command.env("ROCM_CLI_THEROCK_SDK_BIN", bin).env( - "PATH", - prepend_path_entries(&runtime_bin_paths(runtime), std::env::var_os("PATH"))?, - ); - } else if !runtime.sdk_bin_paths.is_empty() { - command.env( - "PATH", - prepend_path_entries(&runtime_bin_paths(runtime), std::env::var_os("PATH"))?, - ); - } - if !cfg!(windows) { - command.env( - "LD_LIBRARY_PATH", - prepend_path_entries( - &therock_library_path_entries(runtime), - std::env::var_os("LD_LIBRARY_PATH"), - )?, - ); - } - Ok(()) -} - -fn runtime_bin_paths(runtime: &AtomRuntime) -> Vec { - let mut entries = Vec::new(); - if let Some(bin) = runtime.sdk_bin.as_ref() { - entries.push(bin.clone()); - } - entries.extend(runtime.sdk_bin_paths.iter().cloned()); - dedupe_paths(entries) -} - -fn therock_library_path_entries(runtime: &AtomRuntime) -> Vec { - let Some(root) = runtime.sdk_root.as_ref() else { - return dedupe_paths(runtime.sdk_library_paths.clone()); - }; - let mut entries = runtime.sdk_library_paths.clone(); - entries.extend([ - root.join("lib"), - root.join("lib64"), - root.join("lib").join("rocm_sysdeps").join("lib"), - ]); - if cfg!(target_os = "linux") { - let wsl_dxcore_lib = PathBuf::from("/usr/lib/wsl/lib"); - if wsl_dxcore_lib.is_dir() { - entries.push(wsl_dxcore_lib); - } - } - dedupe_paths(entries) -} - -fn dedupe_paths(entries: Vec) -> Vec { - let mut deduped = Vec::new(); - for entry in entries { - if !entry.as_os_str().is_empty() && !deduped.iter().any(|seen| seen == &entry) { - deduped.push(entry); - } - } - deduped -} - -fn prepend_path_entries(entries: &[PathBuf], current: Option) -> Result { - let mut parts = Vec::new(); - for entry in entries { - if !entry.as_os_str().is_empty() && !parts.iter().any(|part: &PathBuf| part == entry) { - parts.push(entry.clone()); - } - } - if let Some(current) = current { - for entry in std::env::split_paths(¤t) { - if !entry.as_os_str().is_empty() && !parts.iter().any(|part| part == &entry) { - parts.push(entry); - } - } - } - std::env::join_paths(parts).context("failed to compose runtime path") -} - -fn service_files(service_id: &str) -> Result { - let paths = AppPaths::discover()?; - Ok(ServiceFiles { - state_path: paths - .engine_state_dir(ENGINE_NAME) - .join(format!("{service_id}.json")), - log_path: paths - .engine_logs_dir(ENGINE_NAME) - .join(format!("{service_id}.log")), - }) -} - -fn write_running_state(request: &ServeHttpRequest, runtime: &AtomRuntime, pid: u32) -> Result<()> { - write_state( - &request.state_path, - &json!({ - "service_id": request.service_id, - "engine": ENGINE_NAME, - "status": "running", - "pid": pid, - "model_ref": request.model_ref, - "host": request.host, - "port": request.port, - "endpoint_url": endpoint_url(&request.host, request.port), - "device_policy": "gpu_required", - "runtime_id": request.runtime_id.as_deref().unwrap_or(runtime.runtime_id.as_str()), - "env_id": request.env_id.as_deref().unwrap_or(runtime.env_id.as_str()), - "runtime_executable": runtime.command, - "server_pid": pid, - "engine_recipe": request.engine_recipe, - "engine_recipe_required_flags": engine_recipe_launch_args(request.engine_recipe.as_ref()), - "therock_runtime_env": therock_runtime_env_state(runtime), - "started_at_unix_ms": current_unix_millis() - }), - ) -} - -fn therock_runtime_env_state(runtime: &AtomRuntime) -> Option { - let root = runtime.sdk_root.as_ref()?; - Some(json!({ - "runtime_id": runtime.runtime_id, - "env_id": runtime.env_id, - "root": root.display().to_string(), - "bin": runtime.sdk_bin.as_ref().map(|path| path.display().to_string()), - "bin_paths": runtime_bin_paths(runtime) - .into_iter() - .map(|path| path.display().to_string()) - .collect::>(), - "library_paths": therock_library_path_entries(runtime) - .into_iter() - .map(|path| path.display().to_string()) - .collect::>(), - "source": runtime.source, - })) -} - -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(), Value::String(status.to_owned())); - object.insert( - "stopped_at_unix_ms".to_owned(), - Value::from(current_unix_millis() as u64), - ); - } - write_state(state_path, &state) -} - -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).context("failed to serialize ATOM state")?, - ) - .with_context(|| format!("failed to write {}", path.display())) -} - -fn endpoint_url(host: &str, port: u16) -> String { - format!("{}/v1", format_http_base_url(host, port)) -} - -fn endpoint_url_from_state(state: &Value) -> Option { - value_string(state, "endpoint_url").or_else(|| { - let host = value_string(state, "host")?; - let port = state.get("port")?.as_u64()?; - let port = u16::try_from(port).ok()?; - Some(endpoint_url(&host, port)) - }) -} - -fn query_loaded_model_endpoint(endpoint_url: &str, model_ref: Option<&str>) -> Result { - openai_models_endpoint_has_model( - endpoint_url, - model_ref, - Duration::from_millis(HEALTHCHECK_TIMEOUT_MS), - ) -} - -fn pid_from_state(state: &Value) -> Option { - state - .get("pid")? - .as_u64() - .and_then(|pid| pid.try_into().ok()) -} - -fn terminate_pid(pid: u32, _force: bool) -> bool { - rocm_core::terminate_process(pid).is_ok() -} - -fn tail_lines(path: &Path, limit: usize) -> Result> { - let text = - fs::read_to_string(path).with_context(|| format!("failed to read {}", path.display()))?; - Ok(tail_lines_from_text(&text, limit)) -} - -fn tail_lines_from_text(text: &str, limit: usize) -> Vec { - if limit == 0 { - return Vec::new(); - } - let mut lines = text - .lines() - .rev() - .take(limit) - .map(str::to_owned) - .collect::>(); - lines.reverse(); - lines -} - -fn value_string(value: &Value, key: &str) -> Option { - value.get(key)?.as_str().map(str::to_owned) -} - -fn stable_id_component(value: &str) -> String { - value - .chars() - .map(|ch| { - if ch.is_ascii_alphanumeric() || ch == '-' || ch == '_' { - ch.to_ascii_lowercase() - } else { - '-' - } - }) - .collect() -} - -fn runtime_lock_hash(runtime: &AtomRuntime) -> String { - let mut hasher = DefaultHasher::new(); - runtime.runtime_id.hash(&mut hasher); - runtime.command.hash(&mut hasher); - runtime.version.hash(&mut hasher); - format!("{:016x}", hasher.finish()) -} - -fn current_unix_millis() -> u128 { - SystemTime::now() - .duration_since(UNIX_EPOCH) - .unwrap_or_default() - .as_millis() -} - -const fn windows_unsupported_message() -> &'static str { - "ATOM ROCm serving is supported by rocm-cli only on Linux/WSL; native Windows ATOM is not enabled. No CPU fallback is used." -} - -fn read_request() -> Result { - let mut buffer = String::new(); - std::io::stdin() - .read_to_string(&mut buffer) - .context("failed to read stdin for engine request")?; - serde_json::from_str(&buffer).context("failed to parse engine request envelope") -} - -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!(&mut handle)?; - Ok(()) -} - -#[cfg(test)] -mod tests { - use super::*; - - fn test_engine_recipe(engine: &str, contract_version: &str) -> EngineRecipeHint { - EngineRecipeHint { - contract_version: contract_version.to_owned(), - engine: engine.to_owned(), - required_flags: vec!["--reasoning-parser".to_owned(), "qwen3".to_owned()], - parser_settings: std::collections::BTreeMap::default(), - preferred_endpoint: None, - unsupported_combinations: Vec::new(), - notes: vec!["test recipe".to_owned()], - } - } - - #[test] - fn cpu_policy_is_rejected_without_fallback() { - let error = normalize_atom_device_policy(Some(DevicePolicy::CpuOnly)) - .expect_err("ATOM CPU policy must fail"); - assert!(error.to_string().contains("no CPU fallback is used")); - } - - #[test] - fn gpu_preferred_resolves_to_gpu_required() -> Result<()> { - assert_eq!( - normalize_atom_device_policy(Some(DevicePolicy::GpuPreferred))?, - DevicePolicy::GpuRequired - ); - Ok(()) - } - - #[test] - fn engine_recipe_launch_args_forward_required_flags() { - let hint = test_engine_recipe(ENGINE_NAME, ENGINE_RECIPE_CONTRACT_VERSION); - - assert_eq!( - engine_recipe_launch_args(Some(&hint)), - vec!["--reasoning-parser".to_owned(), "qwen3".to_owned()] - ); - } - - #[test] - fn server_args_support_command_and_python_module_launchers() { - assert_eq!( - atom_server_args(AtomLauncher::Command, "qwen", "127.0.0.1", 30000), - vec!["--model", "qwen", "--host", "127.0.0.1", "--port", "30000"] - ); - assert_eq!( - atom_server_args(AtomLauncher::PythonModule, "qwen", "0.0.0.0", 30001), - vec![ - "-m", - "atom.entrypoints.openai_server", - "--model", - "qwen", - "--host", - "0.0.0.0", - "--port", - "30001" - ] - ); - } - - #[test] - fn endpoint_response_errors_without_service_state() { - let error = endpoint_response(EndpointRequest { - service_id: format!("missing-{}", current_unix_millis()), - }) - .expect_err("missing service state should not produce a default endpoint"); - - assert!(error.to_string().contains("service state not found")); - } - - #[test] - fn endpoint_url_falls_back_to_host_and_port() { - let state = json!({ - "host": "127.0.0.1", - "port": 12345 - }); - assert_eq!( - endpoint_url_from_state(&state), - Some("http://127.0.0.1:12345/v1".to_owned()) - ); - let ipv6_state = json!({ - "host": "::1", - "port": 12345 - }); - assert_eq!( - endpoint_url_from_state(&ipv6_state), - Some("http://[::1]:12345/v1".to_owned()) - ); - } - - #[test] - fn resolve_model_echoes_matching_engine_recipe() -> Result<()> { - let hint = test_engine_recipe(ENGINE_NAME, ENGINE_RECIPE_CONTRACT_VERSION); - let response = resolve_model_response(ResolveModelRequest { - model_ref: "Qwen/Qwen3.5-4B".to_owned(), - runtime_id: None, - device_policy: Some(DevicePolicy::GpuRequired), - recipe_override: None, - engine_recipe: Some(hint.clone()), - })?; - - assert_eq!(response.engine_recipe, Some(hint)); - Ok(()) - } - - #[test] - fn resolve_model_rejects_mismatched_engine_recipe() { - let error = resolve_model_response(ResolveModelRequest { - model_ref: "Qwen/Qwen3.5-4B".to_owned(), - runtime_id: None, - device_policy: Some(DevicePolicy::GpuRequired), - recipe_override: None, - engine_recipe: Some(test_engine_recipe( - "pytorch", - ENGINE_RECIPE_CONTRACT_VERSION, - )), - }) - .expect_err("mismatched engine recipe should fail"); - - assert!(error.to_string().contains("does not match adapter")); - } - - #[test] - fn resolve_model_rejects_unsupported_engine_recipe_contract() { - let error = resolve_model_response(ResolveModelRequest { - model_ref: "Qwen/Qwen3.5-4B".to_owned(), - runtime_id: None, - device_policy: Some(DevicePolicy::GpuRequired), - recipe_override: None, - engine_recipe: Some(test_engine_recipe(ENGINE_NAME, "999.0.0")), - }) - .expect_err("unsupported recipe contract should fail"); - - assert!(error.to_string().contains("unsupported")); - } - - #[test] - fn tail_lines_returns_suffix() -> Result<()> { - let path = std::env::temp_dir().join(format!( - "rocm-atom-tail-{}-{}.log", - std::process::id(), - current_unix_millis() - )); - fs::write(&path, "a\nb\nc\n")?; - let lines = tail_lines(&path, 2)?; - fs::remove_file(path).ok(); - assert_eq!(lines, vec!["b".to_owned(), "c".to_owned()]); - Ok(()) - } - - #[test] - fn stdio_protocol_routes_all_methods_without_side_effects() { - let service_id = format!( - "missing-protocol-{}-{}", - std::process::id(), - current_unix_millis() - ); - let success_cases = [ - (EngineMethod::Detect, json!({})), - (EngineMethod::Capabilities, json!({})), - ( - EngineMethod::ResolveModel, - json!({ - "model_ref": "Qwen/Qwen3.5-4B", - "device_policy": "gpu_required" - }), - ), - ( - EngineMethod::Healthcheck, - json!({ - "service_id": service_id.as_str() - }), - ), - ( - EngineMethod::Logs, - json!({ - "service_id": service_id.as_str(), - "tail_lines": 4 - }), - ), - ( - EngineMethod::Stop, - json!({ - "service_id": service_id.as_str(), - "force": false - }), - ), - ]; - - for (method, payload) in success_cases { - let response = handle_envelope(EngineRequestEnvelope { method, payload }); - assert!( - response.ok, - "expected protocol method to return a typed success envelope: {:?}", - response.error - ); - } - - let endpoint = handle_envelope(EngineRequestEnvelope { - method: EngineMethod::Endpoint, - payload: json!({ - "service_id": service_id.as_str() - }), - }); - assert!(!endpoint.ok); - assert_eq!( - endpoint.error.as_ref().map(|error| error.code.as_str()), - Some("request_failed") - ); - - for method in [EngineMethod::Install, EngineMethod::Launch] { - let response = handle_envelope(EngineRequestEnvelope { - method, - payload: json!({}), - }); - assert!(!response.ok); - assert_eq!( - response.error.as_ref().map(|error| error.code.as_str()), - Some("invalid_payload") - ); - } - } - - #[test] - fn therock_library_path_entries_include_sysdeps_for_hip_apps() { - let root = PathBuf::from(if cfg!(windows) { - r"C:\rocm-sdk" - } else { - "/tmp/rocm-sdk" - }); - let runtime = AtomRuntime { - runtime_id: "therock-release:gfx120X-all".to_owned(), - env_id: "external-atom-therock".to_owned(), - command: PathBuf::from("atom"), - launcher: AtomLauncher::Command, - python_executable: None, - version: None, - source: "managed_runtime_manifest:test".to_owned(), - sdk_root: Some(root.clone()), - sdk_bin: Some(root.join("bin")), - sdk_bin_paths: vec![root.join("runtime").join("bin")], - sdk_library_paths: vec![root.join("runtime").join("lib")], - }; - let entries = therock_library_path_entries(&runtime); - assert!(entries.contains(&root.join("runtime").join("lib"))); - assert!(entries.contains(&root.join("lib"))); - assert!( - entries - .iter() - .any(|entry| entry.ends_with(Path::new("lib").join("rocm_sysdeps").join("lib"))) - ); - } - - #[test] - fn managed_env_reflects_managed_runtime_manifest_source() { - let runtime = AtomRuntime { - runtime_id: "therock-release:gfx120X-all".to_owned(), - env_id: "external-atom-therock".to_owned(), - command: PathBuf::from(if cfg!(windows) { - r"C:\venv\Scripts\python.exe" - } else { - "/home/user/.venv/bin/python" - }), - launcher: AtomLauncher::Command, - python_executable: None, - version: Some("0.5.12".to_owned()), - source: "managed_runtime_manifest:atom-source-pip-gfx120x-all".to_owned(), - sdk_root: Some(PathBuf::from(if cfg!(windows) { - r"C:\rocm-sdk" - } else { - "/home/user/.venv/lib/python/site-packages/_rocm_sdk_devel" - })), - sdk_bin: Some(PathBuf::from(if cfg!(windows) { - r"C:\rocm-sdk\bin" - } else { - "/home/user/.venv/lib/python/site-packages/_rocm_sdk_devel/bin" - })), - sdk_bin_paths: Vec::new(), - sdk_library_paths: Vec::new(), - }; - - assert!(runtime_is_managed(&runtime)); - assert!( - atom_runtime_warnings(&runtime) - .join("\n") - .contains("managed TheRock runtime") - ); - - let external = AtomRuntime { - source: "environment command".to_owned(), - sdk_root: None, - sdk_bin: None, - sdk_bin_paths: Vec::new(), - sdk_library_paths: Vec::new(), - ..runtime - }; - assert!(!runtime_is_managed(&external)); - assert!( - atom_runtime_warnings(&external) - .join("\n") - .contains("external ATOM runtime") - ); - } - - #[test] - fn running_state_records_managed_therock_env_for_gpu_verification() -> Result<()> { - let state_path = std::env::temp_dir().join(format!( - "rocm-atom-state-{}-{}.json", - std::process::id(), - current_unix_millis() - )); - let request = ServeHttpRequest { - service_id: "atom-test".to_owned(), - model_ref: "Qwen/Qwen2.5-0.5B-Instruct".to_owned(), - host: "127.0.0.1".to_owned(), - port: 11435, - device_policy: DevicePolicy::GpuRequired, - gpu_indices: Vec::new(), - runtime_id: Some("therock-release:gfx120X-all".to_owned()), - env_id: None, - state_path: state_path.clone(), - engine_recipe: None, - }; - let runtime = AtomRuntime { - runtime_id: "therock-release:gfx120X-all".to_owned(), - env_id: "external-atom-therock".to_owned(), - command: PathBuf::from(if cfg!(windows) { - r"C:\venv\Scripts\python.exe" - } else { - "/home/user/.venv/bin/python" - }), - launcher: AtomLauncher::Command, - python_executable: None, - version: Some("0.5.12".to_owned()), - source: "managed_runtime_manifest:test".to_owned(), - sdk_root: Some(PathBuf::from(if cfg!(windows) { - r"C:\rocm-sdk" - } else { - "/home/user/.venv/lib/python/site-packages/_rocm_sdk_devel" - })), - sdk_bin: Some(PathBuf::from(if cfg!(windows) { - r"C:\rocm-sdk\bin" - } else { - "/home/user/.venv/lib/python/site-packages/_rocm_sdk_devel/bin" - })), - sdk_bin_paths: vec![PathBuf::from(if cfg!(windows) { - r"C:\rocm-sdk\extra-bin" - } else { - "/home/user/.venv/lib/python/site-packages/_rocm_sdk_libraries/bin" - })], - sdk_library_paths: vec![PathBuf::from(if cfg!(windows) { - r"C:\rocm-sdk\extra-lib" - } else { - "/home/user/.venv/lib/python/site-packages/_rocm_sdk_libraries/lib" - })], - }; - - write_running_state(&request, &runtime, 12345)?; - let state = read_service_state(&state_path)?; - fs::remove_file(&state_path).ok(); - - assert_eq!(state.get("server_pid").and_then(Value::as_u64), Some(12345)); - let runtime_env = state - .get("therock_runtime_env") - .expect("runtime env should be recorded"); - assert_eq!( - runtime_env.get("runtime_id").and_then(Value::as_str), - Some("therock-release:gfx120X-all") - ); - assert!( - runtime_env - .get("root") - .and_then(Value::as_str) - .is_some_and(|root| root.contains("rocm")) - ); - assert!( - runtime_env - .get("bin_paths") - .and_then(Value::as_array) - .is_some_and(|paths| paths.len() >= 2) - ); - assert!( - runtime_env - .get("library_paths") - .and_then(Value::as_array) - .is_some_and(|paths| !paths.is_empty()) - ); - Ok(()) - } -} diff --git a/engines/atom/src/main.rs b/engines/atom/src/main.rs deleted file mode 100644 index e9e36344..00000000 --- a/engines/atom/src/main.rs +++ /dev/null @@ -1,7 +0,0 @@ -// Copyright © Advanced Micro Devices, Inc., or its affiliates. -// -// SPDX-License-Identifier: MIT - -fn main() -> anyhow::Result<()> { - rocm_engine_atom::run_cli() -} diff --git a/engines/llama-cpp/Cargo.toml b/engines/llama-cpp/Cargo.toml deleted file mode 100644 index 7048db0a..00000000 --- a/engines/llama-cpp/Cargo.toml +++ /dev/null @@ -1,23 +0,0 @@ -[package] -name = "rocm-engine-llama-cpp" -version.workspace = true -edition.workspace = true -license.workspace = true -repository.workspace = true -rust-version.workspace = true -publish.workspace = true - -[lints] -workspace = true - -[[bin]] -name = "rocm-engine-llama-cpp" -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 diff --git a/engines/llama-cpp/src/lib.rs b/engines/llama-cpp/src/lib.rs deleted file mode 100644 index 4025f558..00000000 --- a/engines/llama-cpp/src/lib.rs +++ /dev/null @@ -1,2818 +0,0 @@ -// Copyright © Advanced Micro Devices, Inc., or its affiliates. -// -// SPDX-License-Identifier: MIT - -use anyhow::{Context, Result, bail}; -use clap::{Parser, Subcommand}; -use rocm_core::{ - AppPaths, DEFAULT_LOCAL_PORT, format_http_base_url, openai_models_endpoint_has_model, - require_nonempty, -}; -use rocm_engine_protocol::{ - DetectRequest, DetectResponse, DevicePolicy, ENGINE_RECIPE_CONTRACT_VERSION, EndpointRequest, - EndpointResponse, EngineCapabilities, EngineDeviceAvailability, EngineMethod, EngineRecipeHint, - EngineRequestEnvelope, EngineResponseEnvelope, GpuSelection, HealthcheckRequest, - HealthcheckResponse, InstallRequest, InstallResponse, LaunchRequest, LaunchResponse, - LogsRequest, LogsResponse, ResolveModelRequest, ResolveModelResponse, StopRequest, - StopResponse, -}; -use serde::de::DeserializeOwned; -use serde::{Deserialize, Serialize}; -use serde_json::{Value, json}; -use std::collections::{VecDeque, hash_map::DefaultHasher}; -use std::ffi::OsString; -use std::fs; -use std::hash::{Hash, Hasher}; -use std::io::{BufRead, Read, Write}; -use std::path::{Path, PathBuf}; -use std::process::{Command as ProcessCommand, Stdio}; -use std::time::{Duration, SystemTime, UNIX_EPOCH}; - -const ENGINE_NAME: &str = "llama.cpp"; -const DEFAULT_HOST: &str = "127.0.0.1"; -const HEALTHCHECK_TIMEOUT_MS: u64 = 700; -const DEFAULT_LOG_TAIL_LINES: usize = 200; -const LLAMA_GPU_LAYERS_VALUE: &str = "-1"; -const REQUIRED_WINDOWS_THEROCK_EXACT_DLLS: &[&str] = &[ - "amdhip64_7.dll", - "hipblas.dll", - "libhipblaslt.dll", - "rocblas.dll", - "rocm_kpack.dll", - "rocsolver.dll", -]; -const REQUIRED_WINDOWS_THEROCK_VERSIONED_DLLS: &[(&str, &str)] = &[ - ("amd_comgr", ".dll"), - ("hiprtc-builtins", ".dll"), - ("hiprtc", ".dll"), -]; -const REQUIRED_WINDOWS_THEROCK_DATA_DIRS: &[&[&str]] = - &[&["rocblas", "library"], &["hipblaslt", "library"]]; - -#[derive(Parser)] -#[command(name = "rocm-engine-llama.cpp")] -struct Cli { - #[command(subcommand)] - command: CommandKind, -} - -#[derive(Subcommand)] -enum CommandKind { - Detect, - Capabilities, - Install { - #[arg(long)] - runtime_id: String, - #[arg(long)] - reinstall: bool, - }, - ResolveModel { - model_ref: String, - }, - Launch { - service_id: String, - model_ref: String, - #[arg(long, default_value = DEFAULT_HOST)] - host: String, - #[arg(long, default_value_t = DEFAULT_LOCAL_PORT)] - port: u16, - #[arg(long)] - device_policy: Option, - #[arg(long)] - runtime_id: Option, - #[arg(long)] - env_id: Option, - #[arg(long)] - gpu: Option, - }, - Stdio, - ServeHttp { - service_id: String, - model_ref: String, - #[arg(long, default_value = DEFAULT_HOST)] - host: String, - #[arg(long, default_value_t = DEFAULT_LOCAL_PORT)] - port: u16, - #[arg(long)] - device_policy: Option, - #[arg(long)] - runtime_id: Option, - #[arg(long)] - env_id: Option, - #[arg(long)] - state_path: PathBuf, - #[arg(long)] - log_path: Option, - #[arg(long)] - engine_recipe_json: Option, - #[arg(long)] - gpu: Option, - }, -} - -#[derive(Debug, Clone)] -struct LlamaServer { - program: String, - display: String, -} - -#[derive(Debug, Clone)] -struct PreparedLlamaServer { - program: PathBuf, - display: String, - staged_dir: Option, -} - -#[derive(Debug, Clone)] -struct ServiceFiles { - state_path: PathBuf, - log_path: PathBuf, -} - -#[derive(Debug, Clone)] -struct ServeHttpRequest { - service_id: String, - model_ref: String, - host: String, - port: u16, - device_policy: Option, - gpu_indices: Vec, - runtime_id: Option, - env_id: Option, - state_path: PathBuf, - log_path: Option, - engine_recipe: Option, -} - -#[derive(Debug, Clone)] -struct TheRockHipRuntimeEnv { - runtime_id: String, - runtime_key: Option, - root_path: PathBuf, - bin_path: PathBuf, - bin_paths: Vec, - library_paths: Vec, - source: String, -} - -#[derive(Debug, Clone, Deserialize)] -struct TheRockRuntimeManifest { - #[serde(default)] - runtime_key: Option, - #[serde(default)] - runtime_id: Option, - #[serde(default)] - format: Option, - #[serde(default)] - rocm_sdk: Option, - #[serde(default)] - installed_at_unix_ms: Option, -} - -#[derive(Debug, Clone, Deserialize)] -struct RocmSdkRuntimeProbe { - #[serde(default)] - import_ok: bool, - #[serde(default)] - root_path: Option, - #[serde(default)] - bin_path: Option, - #[serde(default)] - bin_paths: Vec, - #[serde(default)] - library_paths: Vec, -} - -pub fn run_cli() -> Result<()> { - let cli = Cli::parse(); - match cli.command { - CommandKind::Detect => print_json(&detect_response())?, - CommandKind::Capabilities => print_json(&capabilities())?, - CommandKind::Install { - runtime_id, - reinstall, - } => print_json(&install_response(InstallRequest { - runtime_id, - python_version: None, - env_root: None, - reinstall, - })?)?, - CommandKind::ResolveModel { model_ref } => { - print_json(&resolve_model_response(ResolveModelRequest { - model_ref, - runtime_id: None, - device_policy: None, - recipe_override: None, - engine_recipe: None, - })?)?; - } - CommandKind::Launch { - service_id, - model_ref, - host, - port, - device_policy, - runtime_id, - env_id, - gpu, - } => print_json(&launch_service(LaunchRequest { - service_id, - env_id, - runtime_id, - model_ref, - host, - port, - device_policy: Some(parse_device_policy_arg(device_policy.as_deref())?), - endpoint_mode: Some("openai".to_owned()), - engine_recipe: None, - gpu_selection: parse_gpu_selection_arg(gpu.as_deref())?, - })?)?, - CommandKind::Stdio => { - let envelope = read_request()?; - print_json(&handle_envelope(envelope))?; - } - CommandKind::ServeHttp { - service_id, - model_ref, - host, - port, - device_policy, - runtime_id, - env_id, - state_path, - log_path, - engine_recipe_json, - gpu, - } => serve_http(ServeHttpRequest { - service_id, - model_ref, - host, - port, - device_policy, - gpu_indices: parse_gpu_indices_arg(gpu.as_deref())?, - runtime_id, - env_id, - state_path, - log_path, - engine_recipe: parse_engine_recipe_json(engine_recipe_json)?, - })?, - } - Ok(()) -} - -pub fn builtin_handle_envelope(envelope: EngineRequestEnvelope) -> EngineResponseEnvelope { - handle_envelope(envelope) -} - -#[allow(clippy::too_many_arguments)] -pub fn builtin_serve_http( - service_id: String, - model_ref: String, - host: String, - port: u16, - device_policy: Option, - gpu_indices: Vec, - runtime_id: Option, - env_id: Option, - state_path: PathBuf, - log_path: Option, - engine_recipe: Option, -) -> Result<()> { - serve_http(ServeHttpRequest { - service_id, - model_ref, - host, - port, - device_policy, - gpu_indices, - runtime_id, - env_id, - state_path, - log_path, - engine_recipe, - }) -} - -fn handle_envelope(envelope: EngineRequestEnvelope) -> EngineResponseEnvelope { - match envelope.method { - EngineMethod::Detect => { - deserialize_and_respond::(envelope.payload, |_| { - Ok(detect_response()) - }) - } - EngineMethod::Capabilities => EngineResponseEnvelope::success(capabilities()), - EngineMethod::Install => { - deserialize_and_respond::(envelope.payload, install_response) - } - EngineMethod::ResolveModel => deserialize_and_respond::( - envelope.payload, - resolve_model_response, - ), - EngineMethod::Launch => { - deserialize_and_respond::(envelope.payload, launch_service) - } - EngineMethod::Healthcheck => deserialize_and_respond::( - envelope.payload, - healthcheck_service, - ), - EngineMethod::Endpoint => { - deserialize_and_respond::(envelope.payload, endpoint_response) - } - EngineMethod::Stop => { - deserialize_and_respond::(envelope.payload, stop_service) - } - EngineMethod::Logs => { - deserialize_and_respond::(envelope.payload, logs_response) - } - } -} - -fn deserialize_and_respond(payload: Value, handler: F) -> EngineResponseEnvelope -where - T: DeserializeOwned, - F: FnOnce(T) -> Result, - U: Serialize, -{ - match serde_json::from_value::(payload) { - Ok(request) => match handler(request) { - Ok(response) => EngineResponseEnvelope::success(response), - Err(error) => EngineResponseEnvelope::failure("request_failed", error.to_string()), - }, - Err(error) => EngineResponseEnvelope::failure("invalid_payload", error.to_string()), - } -} - -fn detect_response() -> DetectResponse { - let server = resolve_llama_server().ok(); - let runtime_executable = server.as_ref().map(|server| server.program.clone()); - let runtime_env = resolve_therock_hip_runtime_env(None); - let rocm_gpu_available = server.as_ref().is_some_and(llama_server_has_hip_backend) - && matches!(runtime_env.as_ref(), Ok(Some(_))); - let mut notes = Vec::new(); - match server.as_ref() { - Some(server) => notes.push(format!("llama-server detected: {}", server.display)), - None => notes.push( - "llama-server not found; set ROCM_CLI_LLAMA_CPP_SERVER or install llama.cpp".to_owned(), - ), - } - match runtime_env.as_ref() { - Ok(Some(runtime_env)) => notes.push(format!( - "TheRock HIP runtime env available: root={} bin={} source={}", - runtime_env.root_path.display(), - runtime_env.bin_path.display(), - runtime_env.source - )), - Ok(None) => notes.push( - "no managed TheRock HIP runtime manifest with rocm_sdk root/bin was found; install a rocm-cli managed SDK before GPU llama.cpp serving".to_owned(), - ), - Err(error) => notes.push(format!("TheRock HIP runtime env probe failed: {error}")), - } - DetectResponse { - installed: server.is_some(), - env_id: server.as_ref().map(|_| "external-llama.cpp".to_owned()), - runtime_kind: Some("external_llama_server".to_owned()), - runtime_executable, - managed_env: Some(false), - python_version: None, - torch_version: None, - transformers_version: None, - available_devices: vec![ - EngineDeviceAvailability { - kind: "cpu".to_owned(), - available: false, - reason: Some( - "rocm-cli does not offer llama.cpp CPU serving; a managed TheRock HIP runtime is required" - .to_owned(), - ), - }, - EngineDeviceAvailability { - kind: "rocm_gpu".to_owned(), - available: rocm_gpu_available, - reason: Some(rocm_gpu_reason(server.as_ref(), runtime_env.as_ref().ok())), - }, - ], - capabilities: engine_capabilities(rocm_gpu_available), - notes, - } -} - -fn capabilities() -> EngineCapabilities { - engine_capabilities(false) -} - -fn engine_capabilities(rocm_gpu: bool) -> EngineCapabilities { - EngineCapabilities { - cpu: false, - rocm_gpu, - openai_compatible: true, - tool_calling: false, - quantized_models: "gguf".to_owned(), - reasoning_parser: false, - } -} - -fn rocm_gpu_reason( - server: Option<&LlamaServer>, - runtime_env: Option<&Option>, -) -> String { - let Some(server) = server else { - return "llama-server was not found".to_owned(); - }; - if !llama_server_has_hip_backend(server) { - return "llama-server was found, but no sibling ggml-hip backend library was detected" - .to_owned(); - } - match runtime_env { - Some(Some(runtime_env)) => format!( - "llama-server has a HIP backend and managed TheRock SDK paths are available from {}", - runtime_env.source - ), - _ => "llama-server has a HIP backend, but no managed TheRock runtime manifest with rocm_sdk root/bin was found".to_owned(), - } -} - -fn llama_server_has_hip_backend(server: &LlamaServer) -> bool { - let Ok(program) = resolve_llama_server_program_path(server) else { - return false; - }; - let Some(dir) = program.parent() else { - return false; - }; - if cfg!(windows) { - dir.join("ggml-hip.dll").is_file() - } else if cfg!(target_os = "macos") { - dir.join("libggml-hip.dylib").is_file() - } else { - dir.join("libggml-hip.so").is_file() || dir.join("ggml-hip.so").is_file() - } -} - -fn install_response(request: InstallRequest) -> Result { - let server = resolve_llama_server().ok(); - let runtime_env = resolve_therock_hip_runtime_env(Some(&request.runtime_id))?; - let rocm_gpu_available = - server.as_ref().is_some_and(llama_server_has_hip_backend) && runtime_env.is_some(); - let paths = AppPaths::discover()?; - let runtime_executable = server.as_ref().map(|server| server.program.clone()); - let mut installed_packages = server - .as_ref() - .map(|server| vec![format!("llama-server={}", server.display)]) - .unwrap_or_default(); - if let Some(runtime_env) = runtime_env.as_ref() { - installed_packages.push(format!( - "therock-hip-runtime-root={}", - runtime_env.root_path.display() - )); - installed_packages.push(format!( - "therock-hip-runtime-bin={}", - runtime_env.bin_path.display() - )); - } - let mut warnings = Vec::new(); - if server.is_some() { - warnings.push( - "llama.cpp adapter uses an external llama-server binary; rocm-cli does not build llama.cpp yet".to_owned(), - ); - } else { - warnings.push( - "llama-server was not found; set ROCM_CLI_LLAMA_CPP_SERVER to a llama-server executable".to_owned(), - ); - } - match runtime_env.as_ref() { - Some(runtime_env) => warnings.push(format!( - "external HIP binaries will inherit TheRock SDK paths from managed runtime manifest {}", - runtime_env.root_path.display() - )), - None => warnings.push( - "no managed TheRock SDK root/bin was resolved; GPU llama.cpp serving will fail without a rocm-cli managed runtime manifest".to_owned(), - ), - } - Ok(InstallResponse { - env_id: "external-llama.cpp".to_owned(), - env_path: paths.engine_dir(ENGINE_NAME).display().to_string(), - python_executable: server.as_ref().map_or_else( - || "".to_owned(), - |server| server.display.clone(), - ), - runtime_kind: Some("external_llama_server".to_owned()), - runtime_executable, - managed_env: Some(false), - installed_packages, - capabilities: engine_capabilities(rocm_gpu_available), - lock_hash: "external".to_owned(), - warnings, - }) -} - -fn resolve_model_response(request: ResolveModelRequest) -> Result { - require_nonempty(&request.model_ref, "model_ref")?; - let engine_recipe = accepted_engine_recipe(request.engine_recipe)?; - let mut warnings = Vec::new(); - let device_policy = normalize_llama_device_policy(request.device_policy)?; - let canonical_model_id = resolve_llama_model_ref(&request.model_ref, &mut warnings); - Ok(ResolveModelResponse { - canonical_model_id, - task: "chat-completions".to_owned(), - source: "llama.cpp".to_owned(), - revision: "local".to_owned(), - loader: "llama.cpp".to_owned(), - trust_remote_code: false, - chat_template_mode: "llama.cpp".to_owned(), - dtype: "gguf".to_owned(), - device_policy, - estimated_memory: "depends on GGUF quantization and context size".to_owned(), - launch_defaults: json!({ - "host": DEFAULT_HOST, - "port": DEFAULT_LOCAL_PORT, - "endpoint_mode": "openai" - }), - engine_recipe, - warnings, - }) -} - -fn accepted_engine_recipe( - engine_recipe: Option, -) -> Result> { - if let Some(hint) = &engine_recipe { - if hint.engine != ENGINE_NAME { - bail!( - "engine_recipe target `{}` does not match adapter `{}`", - hint.engine, - ENGINE_NAME - ); - } - if hint.contract_version != ENGINE_RECIPE_CONTRACT_VERSION { - bail!( - "engine_recipe contract `{}` is unsupported; expected `{}`", - hint.contract_version, - ENGINE_RECIPE_CONTRACT_VERSION - ); - } - } - Ok(engine_recipe) -} - -fn parse_engine_recipe_json(value: Option) -> Result> { - value - .map(|text| { - serde_json::from_str::(&text) - .context("failed to parse engine recipe JSON") - }) - .transpose() - .and_then(accepted_engine_recipe) -} - -fn normalize_llama_device_policy(policy: Option) -> Result { - match policy.unwrap_or(DevicePolicy::GpuRequired) { - DevicePolicy::GpuRequired => Ok(DevicePolicy::GpuRequired), - DevicePolicy::GpuPreferred => Ok(DevicePolicy::GpuRequired), - DevicePolicy::CpuOnly => { - bail!("llama.cpp adapter requires ROCm GPU execution; no CPU fallback is used") - } - } -} - -fn resolve_llama_model_ref(model_ref: &str, warnings: &mut Vec) -> String { - let trimmed = model_ref.trim(); - if !trimmed.to_ascii_lowercase().ends_with(".gguf") { - warnings.push( - "llama.cpp adapter expects a local GGUF model path or a llama-server-compatible model reference".to_owned(), - ); - return trimmed.to_owned(); - } - - let expanded = expand_home_path(trimmed); - let mut candidates = Vec::new(); - if expanded.is_absolute() { - candidates.push(expanded); - } else { - if let Ok(current_dir) = std::env::current_dir() { - candidates.push(current_dir.join(&expanded)); - } - if let Ok(paths) = AppPaths::discover() { - candidates.push(paths.data_dir.join("models").join(&expanded)); - } - } - - for candidate in candidates { - if candidate.is_file() { - return candidate - .canonicalize() - .unwrap_or(candidate) - .display() - .to_string(); - } - } - - warnings.push(format!( - "GGUF model file `{trimmed}` was not found locally; llama-server may still resolve it if it supports this reference" - )); - trimmed.to_owned() -} - -fn expand_home_path(value: &str) -> PathBuf { - let Some(rest) = value - .strip_prefix("~/") - .or_else(|| value.strip_prefix("~\\")) - else { - return PathBuf::from(value); - }; - std::env::var_os("HOME") - .or_else(|| std::env::var_os("USERPROFILE")) - .map_or_else( - || PathBuf::from(value), - |home| PathBuf::from(home).join(rest), - ) -} - -fn launch_service(mut request: LaunchRequest) -> Result { - require_nonempty(&request.service_id, "service_id")?; - require_nonempty(&request.model_ref, "model_ref")?; - let warnings = Vec::::new(); - let device_policy = normalize_llama_device_policy(request.device_policy.clone())?; - require_managed_therock_hip_runtime_env( - resolve_therock_hip_runtime_env(request.runtime_id.as_deref())?, - request.runtime_id.as_deref(), - device_policy_name(&device_policy), - )?; - request.device_policy = Some(device_policy); - request.engine_recipe = accepted_engine_recipe(request.engine_recipe)?; - let paths = AppPaths::discover()?; - paths.ensure()?; - fs::create_dir_all(paths.engine_logs_dir(ENGINE_NAME))?; - fs::create_dir_all(paths.engine_state_dir(ENGINE_NAME))?; - - let log_path = paths - .engine_logs_dir(ENGINE_NAME) - .join(format!("{}.log", request.service_id)); - let state_path = paths - .engine_state_dir(ENGINE_NAME) - .join(format!("{}.json", request.service_id)); - let current_exe = - std::env::current_exe().context("failed to discover current engine binary")?; - let serve_args = serve_http_command_args(&request, &state_path, Some(&log_path)); - let endpoint_url = format!("{}/v1", format_http_base_url(&request.host, request.port)); - write_state( - &state_path, - &json!({ - "engine": ENGINE_NAME, - "service_id": request.service_id, - "model_ref": request.model_ref, - "host": request.host, - "port": request.port, - "status": "starting", - "endpoint_url": endpoint_url, - "log_path": log_path.display().to_string(), - "device_policy": request.device_policy.as_ref().map(device_policy_name), - "runtime_id": request.runtime_id, - "env_id": request.env_id, - "warnings": warnings, - }), - )?; - let wrapper_pid = spawn_serve_http_background(¤t_exe, &serve_args, &state_path) - .context("failed to spawn llama.cpp serve-http process")?; - - merge_json_state( - &state_path, - &json!({ - "pid": wrapper_pid, - }), - )?; - - Ok(LaunchResponse { - service_id: request.service_id, - pid: wrapper_pid, - endpoint_url, - log_path: log_path.display().to_string(), - state_path: state_path.display().to_string(), - }) -} - -fn serve_http_command_args( - request: &LaunchRequest, - state_path: &Path, - log_path: Option<&Path>, -) -> Vec { - let mut args = vec![ - "serve-http".to_owned(), - request.service_id.clone(), - request.model_ref.clone(), - "--host".to_owned(), - request.host.clone(), - "--port".to_owned(), - request.port.to_string(), - ]; - if let Some(device_policy) = request.device_policy.as_ref() { - args.push("--device-policy".to_owned()); - args.push(device_policy_name(device_policy).to_owned()); - } - if let Some(runtime_id) = request.runtime_id.as_deref() { - args.push("--runtime-id".to_owned()); - args.push(runtime_id.to_owned()); - } - if let Some(env_id) = request.env_id.as_deref() { - args.push("--env-id".to_owned()); - args.push(env_id.to_owned()); - } - args.push("--state-path".to_owned()); - args.push(state_path.display().to_string()); - if let Some(log_path) = log_path { - args.push("--log-path".to_owned()); - args.push(log_path.display().to_string()); - } - if let Some(csv) = rocm_engine_protocol::gpu_indices_to_csv( - &rocm_engine_protocol::launch_gpu_indices(request.gpu_selection.as_ref()), - ) { - args.push("--gpu".to_owned()); - args.push(csv); - } - if let Some(engine_recipe) = request.engine_recipe.as_ref() { - args.push("--engine-recipe-json".to_owned()); - args.push(serde_json::to_string(engine_recipe).expect("engine recipe serializes")); - } - args -} - -fn normalize_serve_http_request(request: ServeHttpRequest) -> Result { - match request.device_policy.as_deref() { - Some("gpu_required" | "gpu_preferred") | None => {} - Some("cpu" | "cpu_only") => { - bail!("llama.cpp adapter requires ROCm GPU execution; no CPU fallback is used") - } - Some(other) => bail!("unknown llama.cpp device policy `{other}`"), - } - Ok(request) -} - -fn serve_http(request: ServeHttpRequest) -> Result<()> { - let request = normalize_serve_http_request(request)?; - let server = resolve_llama_server()?; - let paths = AppPaths::discover()?; - let runtime_env = resolve_therock_hip_runtime_env(request.runtime_id.as_deref())?; - let gpu_requested = serve_http_gpu_requested(request.device_policy.as_deref()); - let runtime_env = if gpu_requested { - Some(require_managed_therock_hip_runtime_env( - runtime_env, - request.runtime_id.as_deref(), - request.device_policy.as_deref().unwrap_or("gpu_required"), - )?) - } else { - runtime_env - }; - let prepared_server = prepare_llama_server_for_launch( - &paths, - &server, - runtime_env.as_ref(), - &request.service_id, - )?; - let endpoint_url = format!("{}/v1", format_http_base_url(&request.host, request.port)); - write_state( - &request.state_path, - &json!({ - "engine": ENGINE_NAME, - "service_id": request.service_id, - "model_ref": request.model_ref, - "host": request.host, - "port": request.port, - "pid": std::process::id(), - "wrapper_pid": std::process::id(), - "status": "starting", - "endpoint_url": endpoint_url, - "log_path": request - .log_path - .as_ref() - .map(|path| path.display().to_string()), - "server": prepared_server.display, - "server_source": server.display, - "staged_runtime_dir": prepared_server - .staged_dir - .as_ref() - .map(|path| path.display().to_string()), - "device_policy": request.device_policy, - "runtime_id": request.runtime_id, - "env_id": request.env_id, - "engine_recipe": request.engine_recipe, - "engine_recipe_required_flags": engine_recipe_launch_args(request.engine_recipe.as_ref()), - "therock_runtime_env": runtime_env.as_ref().map(|runtime_env| json!({ - "runtime_id": runtime_env.runtime_id.clone(), - "runtime_key": runtime_env.runtime_key.clone(), - "root": runtime_env.root_path.display().to_string(), - "bin": runtime_env.bin_path.display().to_string(), - "bin_paths": runtime_bin_paths(runtime_env) - .into_iter() - .map(|path| path.display().to_string()) - .collect::>(), - "library_paths": runtime_env.library_paths - .iter() - .map(|path| path.display().to_string()) - .collect::>(), - "source": runtime_env.source.clone(), - })), - }), - )?; - let mut command = ProcessCommand::new(&prepared_server.program); - let server_args = llama_server_args(&request, gpu_requested); - command.args(&server_args).stdin(Stdio::null()); - rocm_engine_protocol::apply_gpu_visibility(&mut command, &request.gpu_indices); - if let Some(log_path) = request.log_path.as_ref() { - if let Some(parent) = log_path.parent() { - fs::create_dir_all(parent) - .with_context(|| format!("failed to create {}", parent.display()))?; - } - let log_file = fs::File::create(log_path) - .with_context(|| format!("failed to create {}", log_path.display()))?; - let log_file_err = log_file - .try_clone() - .context("failed to clone log file handle")?; - command - .stdout(Stdio::from(log_file)) - .stderr(Stdio::from(log_file_err)); - } else { - command.stdout(Stdio::inherit()).stderr(Stdio::inherit()); - } - if let Some(runtime_env) = runtime_env.as_ref() { - apply_therock_hip_runtime_env(&mut command, runtime_env)?; - } - let mut child = command.spawn().context("failed to start llama-server")?; - merge_json_state( - &request.state_path, - &json!({ - "status": "running", - "server_pid": child.id(), - }), - )?; - let status = child.wait().context("failed waiting for llama-server")?; - let final_status = if status.success() { - "stopped" - } else { - "failed" - }; - mark_json_status(&request.state_path, final_status)?; - if status.success() { - Ok(()) - } else { - bail!("llama-server exited with status {status}") - } -} - -fn serve_http_gpu_requested(policy: Option<&str>) -> bool { - !matches!(policy, Some("cpu" | "cpu_only")) -} - -#[cfg(windows)] -fn spawn_serve_http_background( - current_exe: &Path, - serve_args: &[String], - state_path: &Path, -) -> Result { - let launcher_path = background_launcher_script_path(state_path); - fs::write(&launcher_path, windows_background_launcher_script()) - .with_context(|| format!("failed to write {}", launcher_path.display()))?; - let output = ProcessCommand::new("powershell.exe") - .arg("-NoProfile") - .arg("-NonInteractive") - .arg("-ExecutionPolicy") - .arg("Bypass") - .arg("-File") - .arg(&launcher_path) - .arg(current_exe) - .args(serve_args) - .stdin(Stdio::null()) - .output(); - let _ = fs::remove_file(&launcher_path); - let output = output.context("failed to invoke PowerShell background launcher")?; - if !output.status.success() { - bail!( - "PowerShell background launcher failed with status {}\nstdout:\n{}\nstderr:\n{}", - output.status, - String::from_utf8_lossy(&output.stdout), - String::from_utf8_lossy(&output.stderr) - ); - } - let pid_text = String::from_utf8_lossy(&output.stdout); - let pid_text = pid_text.trim(); - pid_text.parse::().with_context(|| { - format!("PowerShell background launcher returned invalid pid `{pid_text}`") - }) -} - -#[cfg(windows)] -fn background_launcher_script_path(state_path: &Path) -> PathBuf { - state_path.with_extension("launch.ps1") -} - -#[cfg(windows)] -fn windows_background_launcher_script() -> &'static str { - r#"$ErrorActionPreference = 'Stop' -if ($args.Count -lt 1) { - throw 'missing executable path' -} -$exe = $args[0] -$childArgs = @() -if ($args.Count -gt 1) { - $childArgs = $args[1..($args.Count - 1)] -} -$p = Start-Process -FilePath $exe -ArgumentList $childArgs -WindowStyle Hidden -PassThru -[Console]::Out.Write($p.Id) -"# -} - -#[cfg(not(windows))] -fn spawn_serve_http_background( - current_exe: &Path, - serve_args: &[String], - _state_path: &Path, -) -> Result { - let child = ProcessCommand::new(current_exe) - .args(serve_args) - .stdin(Stdio::null()) - .stdout(Stdio::null()) - .stderr(Stdio::null()) - .spawn()?; - Ok(child.id()) -} - -fn require_managed_therock_hip_runtime_env( - runtime_env: Option, - runtime_id: Option<&str>, - policy: &str, -) -> Result { - let Some(runtime_env) = runtime_env else { - let runtime_hint = runtime_id - .map(|runtime_id| format!(" for runtime_id `{runtime_id}`")) - .unwrap_or_default(); - bail!( - "llama.cpp {policy} requires a rocm-cli managed TheRock runtime manifest with rocm_sdk.root_path and rocm_sdk.bin_path{runtime_hint}; no CPU fallback is applied" - ); - }; - if !runtime_env.source.starts_with("managed_runtime_manifest") { - bail!( - "llama.cpp {policy} requires TheRock SDK paths from a rocm-cli managed runtime manifest; source `{}` is not allowed and no CPU fallback is applied", - runtime_env.source - ); - } - Ok(runtime_env) -} - -fn llama_server_args(request: &ServeHttpRequest, gpu_requested: bool) -> Vec { - let mut args = vec![ - "-m".to_owned(), - request.model_ref.clone(), - "--host".to_owned(), - request.host.clone(), - "--port".to_owned(), - request.port.to_string(), - ]; - if gpu_requested { - args.push("--n-gpu-layers".to_owned()); - args.push(LLAMA_GPU_LAYERS_VALUE.to_owned()); - } - args.extend(engine_recipe_launch_args(request.engine_recipe.as_ref())); - args -} - -fn engine_recipe_launch_args(engine_recipe: Option<&EngineRecipeHint>) -> Vec { - engine_recipe - .map(|hint| hint.required_flags.clone()) - .unwrap_or_default() -} - -fn healthcheck_service(request: HealthcheckRequest) -> Result { - require_nonempty(&request.service_id, "service_id")?; - let files = service_files(&request.service_id)?; - let state = read_service_state(&files.state_path).ok(); - let endpoint_url = state.as_ref().and_then(endpoint_url_from_state); - let model_ref = state - .as_ref() - .and_then(|value| value_string(value, "model_ref")); - let ready = endpoint_url - .as_deref() - .map(|endpoint| query_loaded_model_endpoint(endpoint, model_ref.as_deref())) - .transpose() - .unwrap_or(None) - .unwrap_or(false); - let status = if ready { - "ready".to_owned() - } else { - state - .as_ref() - .and_then(|value| value_string(value, "status")) - .unwrap_or_else(|| "unknown".to_owned()) - }; - Ok(HealthcheckResponse { - status, - model_loaded: ready, - device: healthcheck_device_from_state(state.as_ref()), - uptime_sec: 0, - queue_depth: 0, - last_error: None, - tokens_per_sec: None, - }) -} - -fn healthcheck_device_from_state(state: Option<&Value>) -> String { - let Some(state) = state else { - return "unknown".to_owned(); - }; - match value_string(state, "device_policy").as_deref() { - Some("gpu_required" | "gpu_preferred") => "rocm_gpu".to_owned(), - _ if state.get("therock_runtime_env").is_some() => "rocm_gpu".to_owned(), - _ => "unknown".to_owned(), - } -} - -fn endpoint_response(request: EndpointRequest) -> Result { - require_nonempty(&request.service_id, "service_id")?; - let files = service_files(&request.service_id)?; - let state = read_service_state(&files.state_path) - .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(), - ], - }) -} - -fn logs_response(request: LogsRequest) -> Result { - require_nonempty(&request.service_id, "service_id")?; - let files = service_files(&request.service_id)?; - let limit = request.tail_lines.unwrap_or(DEFAULT_LOG_TAIL_LINES); - Ok(LogsResponse { - log_path: files.log_path.display().to_string(), - recent_lines: if files.log_path.is_file() { - tail_lines(&files.log_path, limit)? - } else { - Vec::new() - }, - }) -} - -fn stop_service(request: StopRequest) -> Result { - require_nonempty(&request.service_id, "service_id")?; - let files = service_files(&request.service_id)?; - let state = read_service_state(&files.state_path).ok(); - let stopped = match state.as_ref().and_then(pid_to_terminate_from_state) { - Some(pid) => terminate_pid(pid, request.force), - None => false, - }; - if stopped { - mark_json_status(&files.state_path, "stopped")?; - } - Ok(StopResponse { - stopped, - graceful: stopped && !request.force, - }) -} - -fn pid_to_terminate_from_state(state: &Value) -> Option { - value_u32(state, "server_pid").or_else(|| value_u32(state, "pid")) -} - -fn service_files(service_id: &str) -> Result { - let paths = AppPaths::discover()?; - Ok(ServiceFiles { - state_path: paths - .engine_state_dir(ENGINE_NAME) - .join(format!("{service_id}.json")), - log_path: paths - .engine_logs_dir(ENGINE_NAME) - .join(format!("{service_id}.log")), - }) -} - -fn resolve_llama_server() -> Result { - for env in ["ROCM_CLI_LLAMA_CPP_SERVER", "LLAMA_CPP_SERVER"] { - if let Some(value) = std::env::var_os(env) { - let path = PathBuf::from(value); - if path.is_file() { - let path = path.canonicalize().unwrap_or(path); - return Ok(LlamaServer { - program: path.display().to_string(), - display: format!("{env}={}", path.display()), - }); - } - } - } - if let Some(path) = find_program_on_path("llama-server") { - return Ok(LlamaServer { - program: path.display().to_string(), - display: format!("llama-server on PATH ({})", path.display()), - }); - } - bail!("unable to locate llama-server; set ROCM_CLI_LLAMA_CPP_SERVER") -} - -fn prepare_llama_server_for_launch( - paths: &AppPaths, - server: &LlamaServer, - runtime_env: Option<&TheRockHipRuntimeEnv>, - service_id: &str, -) -> Result { - if cfg!(windows) - && let Some(runtime_env) = runtime_env - { - return stage_windows_therock_llama_server(paths, server, runtime_env, service_id); - } - - Ok(PreparedLlamaServer { - program: PathBuf::from(&server.program), - display: server.display.clone(), - staged_dir: None, - }) -} - -fn stage_windows_therock_llama_server( - paths: &AppPaths, - server: &LlamaServer, - runtime_env: &TheRockHipRuntimeEnv, - service_id: &str, -) -> Result { - let server_program = resolve_llama_server_program_path(server)?; - let server_file_name = server_program - .file_name() - .context("llama-server path has no file name")?; - let stage_dir = paths - .engine_dir(ENGINE_NAME) - .join("runtime-bin") - .join(stage_dir_name(service_id, server, runtime_env)); - if stage_dir.exists() { - fs::remove_dir_all(&stage_dir).with_context(|| { - format!( - "failed to clear stale staged runtime {}", - stage_dir.display() - ) - })?; - } - fs::create_dir_all(&stage_dir) - .with_context(|| format!("failed to create {}", stage_dir.display()))?; - - stage_llama_server_files(&server_program, &stage_dir)?; - stage_therock_runtime_files(runtime_env, &stage_dir)?; - - let staged_program = stage_dir.join(server_file_name); - Ok(PreparedLlamaServer { - program: staged_program.clone(), - display: format!( - "staged llama-server with managed TheRock runtime: {}", - staged_program.display() - ), - staged_dir: Some(stage_dir), - }) -} - -fn stage_dir_name( - service_id: &str, - server: &LlamaServer, - runtime_env: &TheRockHipRuntimeEnv, -) -> String { - let mut hasher = DefaultHasher::new(); - server.program.hash(&mut hasher); - runtime_env.runtime_id.hash(&mut hasher); - runtime_env.bin_path.hash(&mut hasher); - for path in &runtime_env.bin_paths { - path.hash(&mut hasher); - } - format!( - "{}-{:016x}", - safe_path_component(service_id), - hasher.finish() - ) -} - -fn stage_llama_server_files(server_program: &Path, stage_dir: &Path) -> Result<()> { - let server_dir = server_program - .parent() - .context("llama-server path has no parent directory")?; - let server_file_name = server_program - .file_name() - .context("llama-server path has no file name")?; - stage_file(server_program, &stage_dir.join(server_file_name), false)?; - - for entry in fs::read_dir(server_dir) - .with_context(|| format!("failed to read {}", server_dir.display()))? - { - let entry = entry?; - let source = entry.path(); - if source == server_program || !is_windows_dll(&source) { - continue; - } - let destination = stage_dir.join(entry.file_name()); - stage_file(&source, &destination, false)?; - } - Ok(()) -} - -fn stage_therock_runtime_files(runtime_env: &TheRockHipRuntimeEnv, stage_dir: &Path) -> Result<()> { - validate_therock_windows_runtime_files(runtime_env)?; - for bin_path in runtime_bin_paths(runtime_env) { - for entry in fs::read_dir(&bin_path) - .with_context(|| format!("failed to read {}", bin_path.display()))? - { - let entry = entry?; - let source = entry.path(); - if !source.is_file() || !is_windows_dll(&source) { - continue; - } - let destination = stage_dir.join(entry.file_name()); - stage_file(&source, &destination, true)?; - } - - for dirname in ["rocblas", "hipblaslt"] { - let source = bin_path.join(dirname); - if source.is_dir() { - stage_runtime_tree(&source, &stage_dir.join(dirname))?; - } - } - } - Ok(()) -} - -fn validate_therock_windows_runtime_files(runtime_env: &TheRockHipRuntimeEnv) -> Result<()> { - for filename in REQUIRED_WINDOWS_THEROCK_EXACT_DLLS { - if find_runtime_bin_file(runtime_env, filename).is_none() { - bail!( - "managed TheRock runtime is missing required llama.cpp HIP DLL {filename}; no CPU fallback is applied" - ); - } - } - for (prefix, suffix) in REQUIRED_WINDOWS_THEROCK_VERSIONED_DLLS { - if find_runtime_bin_file_by_prefix_suffix(runtime_env, prefix, suffix).is_none() { - bail!( - "managed TheRock runtime is missing required llama.cpp HIP DLL matching {prefix}*{suffix}; no CPU fallback is applied" - ); - } - } - for components in REQUIRED_WINDOWS_THEROCK_DATA_DIRS { - if find_runtime_bin_dir(runtime_env, components).is_none() { - bail!( - "managed TheRock runtime is missing required llama.cpp HIP data directory {}; no CPU fallback is applied", - components.join("/") - ); - } - } - Ok(()) -} - -fn runtime_bin_paths(runtime_env: &TheRockHipRuntimeEnv) -> Vec { - let mut paths = Vec::new(); - for path in std::iter::once(&runtime_env.bin_path).chain(runtime_env.bin_paths.iter()) { - if path.is_dir() && !paths.iter().any(|existing| existing == path) { - paths.push(path.clone()); - } - } - paths -} - -fn find_runtime_bin_file(runtime_env: &TheRockHipRuntimeEnv, filename: &str) -> Option { - runtime_bin_paths(runtime_env) - .into_iter() - .map(|path| path.join(filename)) - .find(|path| path.is_file()) -} - -fn find_runtime_bin_file_by_prefix_suffix( - runtime_env: &TheRockHipRuntimeEnv, - prefix: &str, - suffix: &str, -) -> Option { - for bin_path in runtime_bin_paths(runtime_env) { - let Ok(entries) = fs::read_dir(&bin_path) else { - continue; - }; - for entry in entries.flatten() { - let path = entry.path(); - if !path.is_file() { - continue; - } - let Some(file_name) = path.file_name().and_then(|value| value.to_str()) else { - continue; - }; - if starts_with_ascii_case_insensitive(file_name, prefix) - && ends_with_ascii_case_insensitive(file_name, suffix) - { - return Some(path); - } - } - } - None -} - -fn starts_with_ascii_case_insensitive(value: &str, prefix: &str) -> bool { - value - .get(..prefix.len()) - .is_some_and(|head| head.eq_ignore_ascii_case(prefix)) -} - -fn ends_with_ascii_case_insensitive(value: &str, suffix: &str) -> bool { - value - .get(value.len().saturating_sub(suffix.len())..) - .is_some_and(|tail| tail.eq_ignore_ascii_case(suffix)) -} - -fn find_runtime_bin_dir( - runtime_env: &TheRockHipRuntimeEnv, - components: &[&str], -) -> Option { - runtime_bin_paths(runtime_env) - .into_iter() - .map(|path| { - components - .iter() - .fold(path, |path, component| path.join(component)) - }) - .find(|path| path.is_dir()) -} - -fn stage_runtime_tree(source_dir: &Path, destination_dir: &Path) -> Result<()> { - fs::create_dir_all(destination_dir) - .with_context(|| format!("failed to create {}", destination_dir.display()))?; - for entry in fs::read_dir(source_dir) - .with_context(|| format!("failed to read {}", source_dir.display()))? - { - let entry = entry?; - let source = entry.path(); - let destination = destination_dir.join(entry.file_name()); - if source.is_dir() { - stage_runtime_tree(&source, &destination)?; - } else if source.is_file() { - stage_file(&source, &destination, true)?; - } - } - Ok(()) -} - -fn stage_file(source: &Path, destination: &Path, prefer_hardlink: bool) -> Result<()> { - if staged_file_is_current(source, destination) { - return Ok(()); - } - if let Some(parent) = destination.parent() { - fs::create_dir_all(parent) - .with_context(|| format!("failed to create {}", parent.display()))?; - } - if destination.exists() { - fs::remove_file(destination) - .with_context(|| format!("failed to replace {}", destination.display()))?; - } - if prefer_hardlink && fs::hard_link(source, destination).is_ok() { - return Ok(()); - } - fs::copy(source, destination).with_context(|| { - format!( - "failed to stage {} as {}", - source.display(), - destination.display() - ) - })?; - Ok(()) -} - -fn staged_file_is_current(source: &Path, destination: &Path) -> bool { - let (Ok(source_meta), Ok(destination_meta)) = (source.metadata(), destination.metadata()) - else { - return false; - }; - if source_meta.len() != destination_meta.len() { - return false; - } - let (Ok(source_modified), Ok(destination_modified)) = - (source_meta.modified(), destination_meta.modified()) - else { - return true; - }; - destination_modified >= source_modified -} - -fn is_windows_dll(path: &Path) -> bool { - path.extension() - .and_then(|extension| extension.to_str()) - .is_some_and(|extension| extension.eq_ignore_ascii_case("dll")) -} - -fn safe_path_component(value: &str) -> String { - let mut safe = String::new(); - for ch in value.chars() { - if ch.is_ascii_alphanumeric() || matches!(ch, '.' | '-' | '_') { - safe.push(ch); - } else { - safe.push('_'); - } - } - let safe = safe.trim_matches('.'); - if safe.is_empty() { - "service".to_owned() - } else { - safe.to_owned() - } -} - -fn resolve_llama_server_program_path(server: &LlamaServer) -> Result { - let program = PathBuf::from(&server.program); - if program.is_file() { - return Ok(program.canonicalize().unwrap_or(program)); - } - find_program_on_path(&server.program).with_context(|| { - format!( - "unable to resolve llama-server executable `{}`", - server.program - ) - }) -} - -fn find_program_on_path(program: &str) -> Option { - let program_path = Path::new(program); - if program_path.components().count() > 1 && program_path.is_file() { - return Some(program_path.to_path_buf()); - } - - let path = std::env::var_os("PATH")?; - let extensions = executable_extensions(program); - for dir in std::env::split_paths(&path) { - for extension in &extensions { - let candidate = dir.join(format!("{program}{extension}")); - if candidate.is_file() { - return Some(candidate.canonicalize().unwrap_or(candidate)); - } - } - } - None -} - -fn executable_extensions(program: &str) -> Vec { - if Path::new(program).extension().is_some() { - return vec![String::new()]; - } - if cfg!(windows) { - let pathext = std::env::var_os("PATHEXT") - .and_then(|value| value.into_string().ok()) - .unwrap_or_else(|| ".COM;.EXE;.BAT;.CMD".to_owned()); - let mut extensions = vec![String::new()]; - for extension in pathext.split(';') { - if extension.is_empty() { - continue; - } - let extension = if extension.starts_with('.') { - extension.to_owned() - } else { - format!(".{extension}") - }; - if !extensions - .iter() - .any(|existing| existing.eq_ignore_ascii_case(&extension)) - { - extensions.push(extension); - } - } - extensions - } else { - vec![String::new()] - } -} - -fn resolve_therock_hip_runtime_env( - runtime_id: Option<&str>, -) -> Result> { - let paths = AppPaths::discover()?; - resolve_managed_therock_hip_runtime_env(&paths, runtime_id) -} - -fn resolve_managed_therock_hip_runtime_env( - paths: &AppPaths, - runtime_id: Option<&str>, -) -> Result> { - let registry_dir = paths.data_dir.join("runtimes").join("registry"); - if !registry_dir.is_dir() { - return Ok(None); - } - - let mut manifests = Vec::new(); - for entry in fs::read_dir(®istry_dir) - .with_context(|| format!("failed to read {}", registry_dir.display()))? - { - let entry = entry?; - let path = entry.path(); - if path.extension().and_then(|value| value.to_str()) != Some("json") { - continue; - } - let bytes = - fs::read(&path).with_context(|| format!("failed to read {}", path.display()))?; - let Ok(manifest) = serde_json::from_slice::(&bytes) else { - continue; - }; - if !therock_runtime_matches(&manifest, runtime_id) { - continue; - } - manifests.push((manifest.installed_at_unix_ms.unwrap_or(0), manifest)); - } - manifests.sort_by_key(|(installed_at, _)| std::cmp::Reverse(*installed_at)); - - for (_, manifest) in manifests { - if let Some(env) = therock_env_from_manifest(&manifest)? { - return Ok(Some(env)); - } - } - Ok(None) -} - -fn therock_runtime_matches(manifest: &TheRockRuntimeManifest, requested: Option<&str>) -> bool { - let Some(requested) = requested.map(str::trim).filter(|value| !value.is_empty()) else { - return true; - }; - let requested = requested.to_ascii_lowercase(); - if requested == "external" || requested == "external-llama.cpp" { - return false; - } - - for candidate in [ - manifest.runtime_id.as_deref(), - manifest.runtime_key.as_deref(), - manifest.format.as_deref(), - ] - .into_iter() - .flatten() - { - let candidate = candidate.to_ascii_lowercase(); - if candidate == requested || candidate.starts_with(&requested) { - return true; - } - } - false -} - -fn therock_env_from_manifest( - manifest: &TheRockRuntimeManifest, -) -> Result> { - let runtime_id = manifest - .runtime_id - .clone() - .unwrap_or_else(|| "therock-runtime".to_owned()); - let runtime_key = manifest.runtime_key.clone(); - let source = runtime_key.as_deref().map_or_else( - || "managed_runtime_manifest".to_owned(), - |key| format!("managed_runtime_manifest:{key}"), - ); - - if let Some(probe) = manifest.rocm_sdk.as_ref() - && probe.import_ok - && let (Some(root_path), Some(bin_path)) = - (probe.root_path.as_ref(), probe.bin_path.as_ref()) - && root_path.is_dir() - && bin_path.is_dir() - { - let bin_paths = probe - .bin_paths - .iter() - .filter(|path| path.is_dir()) - .cloned() - .collect::>(); - let library_paths = probe - .library_paths - .iter() - .filter(|path| path.is_dir()) - .cloned() - .collect::>(); - return Ok(Some(TheRockHipRuntimeEnv { - runtime_id, - runtime_key, - root_path: root_path.clone(), - bin_path: bin_path.clone(), - bin_paths, - library_paths, - source, - })); - } - - Ok(None) -} - -fn apply_therock_hip_runtime_env( - command: &mut ProcessCommand, - runtime_env: &TheRockHipRuntimeEnv, -) -> Result<()> { - command - .env("ROCM_SDK_ROOT", &runtime_env.root_path) - .env("ROCM_PATH", &runtime_env.root_path) - .env("ROCM_HOME", &runtime_env.root_path) - .env("HIP_PATH", &runtime_env.root_path) - .env("ROCM_CLI_THEROCK_RUNTIME_ID", &runtime_env.runtime_id) - .env("ROCM_CLI_THEROCK_SDK_BIN", &runtime_env.bin_path) - .env( - "PATH", - prepend_path_entries(&runtime_bin_paths(runtime_env), std::env::var_os("PATH"))?, - ); - - let cmake_path = runtime_env.root_path.join("lib").join("cmake"); - command.env( - "CMAKE_PREFIX_PATH", - prepend_path_entries( - &[runtime_env.root_path.clone(), cmake_path], - std::env::var_os("CMAKE_PREFIX_PATH"), - )?, - ); - - if !cfg!(windows) { - command.env( - "LD_LIBRARY_PATH", - prepend_path_entries( - &therock_library_path_entries(runtime_env), - std::env::var_os("LD_LIBRARY_PATH"), - )?, - ); - } - Ok(()) -} - -fn therock_library_path_entries(runtime_env: &TheRockHipRuntimeEnv) -> Vec { - let mut entries = runtime_env.library_paths.clone(); - entries.extend([ - runtime_env.root_path.join("lib"), - runtime_env.root_path.join("lib64"), - runtime_env - .root_path - .join("lib") - .join("rocm_sysdeps") - .join("lib"), - ]); - if cfg!(target_os = "linux") { - let wsl_dxcore_lib = PathBuf::from("/usr/lib/wsl/lib"); - if wsl_dxcore_lib.is_dir() { - entries.push(wsl_dxcore_lib); - } - } - entries -} - -fn prepend_path_entries(entries: &[PathBuf], current: Option) -> Result { - let mut parts = Vec::new(); - for entry in entries { - if !entry.as_os_str().is_empty() && !parts.iter().any(|part: &PathBuf| part == entry) { - parts.push(entry.clone()); - } - } - if let Some(current) = current { - for entry in std::env::split_paths(¤t) { - if !entry.as_os_str().is_empty() && !parts.iter().any(|part| part == &entry) { - parts.push(entry); - } - } - } - std::env::join_paths(parts).context("failed to compose TheRock runtime path") -} - -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).context("failed to serialize llama.cpp state")?, - ) - .with_context(|| format!("failed to write {}", path.display())) -} - -fn mark_json_status(path: &Path, status: &str) -> Result<()> { - merge_json_state( - path, - &json!({ - "engine": ENGINE_NAME, - "status": status, - "stopped_at_unix_ms": current_unix_millis(), - }), - ) -} - -fn merge_json_state(path: &Path, patch: &Value) -> Result<()> { - let mut value = read_service_state(path).unwrap_or_else(|_| json!({})); - if !value.is_object() { - value = json!({}); - } - let object = value.as_object_mut().expect("object checked above"); - if let Some(patch) = patch.as_object() { - for (key, value) in patch { - object.insert(key.clone(), value.clone()); - } - } - write_state(path, &value) -} - -fn endpoint_url_from_state(state: &Value) -> Option { - value_string(state, "endpoint_url").or_else(|| { - let host = value_string(state, "host")?; - let port = value_u32(state, "port")?; - let port = u16::try_from(port).ok()?; - Some(format!("{}/v1", format_http_base_url(&host, port))) - }) -} - -fn value_string(value: &Value, key: &str) -> Option { - value - .get(key) - .and_then(Value::as_str) - .filter(|value| !value.trim().is_empty()) - .map(ToOwned::to_owned) -} - -fn value_u32(value: &Value, key: &str) -> Option { - value - .get(key) - .and_then(Value::as_u64) - .and_then(|value| u32::try_from(value).ok()) -} - -fn query_loaded_model_endpoint(endpoint_url: &str, model_ref: Option<&str>) -> Result { - openai_models_endpoint_has_model( - endpoint_url, - model_ref, - Duration::from_millis(HEALTHCHECK_TIMEOUT_MS), - ) -} - -#[cfg(test)] -fn parse_http_endpoint(endpoint_url: &str) -> Option<(String, u16)> { - let without_scheme = endpoint_url.trim().strip_prefix("http://")?; - let authority = without_scheme.split('/').next()?.trim(); - if let Some(rest) = authority.strip_prefix('[') { - let end = rest.find(']')?; - let host = rest[..end].to_owned(); - let port = rest[end + 1..].strip_prefix(':')?.parse().ok()?; - return Some((host, port)); - } - let (host, port) = authority.rsplit_once(':')?; - Some((host.to_owned(), port.parse().ok()?)) -} - -fn tail_lines(path: &Path, limit: usize) -> Result> { - if limit == 0 { - return Ok(Vec::new()); - } - let file = - fs::File::open(path).with_context(|| format!("failed to open {}", path.display()))?; - let reader = std::io::BufReader::new(file); - let mut lines = VecDeque::with_capacity(limit); - for line in reader.lines() { - let line = line.with_context(|| format!("failed to read {}", path.display()))?; - if lines.len() == limit { - lines.pop_front(); - } - lines.push_back(line); - } - Ok(lines.into_iter().collect()) -} - -fn terminate_pid(pid: u32, _force: bool) -> bool { - rocm_core::terminate_process(pid).is_ok() -} - -const fn device_policy_name(policy: &DevicePolicy) -> &'static str { - match policy { - DevicePolicy::GpuRequired => "gpu_required", - DevicePolicy::GpuPreferred => "gpu_preferred", - DevicePolicy::CpuOnly => "cpu_only", - } -} - -fn parse_device_policy_arg(value: Option<&str>) -> Result { - match value.unwrap_or("gpu_required") { - "cpu" | "cpu_only" => Ok(DevicePolicy::CpuOnly), - "gpu" | "gpu_preferred" => Ok(DevicePolicy::GpuPreferred), - "gpu_required" => Ok(DevicePolicy::GpuRequired), - other => bail!("unknown device policy `{other}`"), - } -} - -/// Parse a `--gpu` CLI value into an optional `GpuSelection` for `LaunchRequest`. -fn parse_gpu_selection_arg(value: Option<&str>) -> Result> { - value - .map(|raw| GpuSelection::parse_cli_value(raw).map_err(anyhow::Error::msg)) - .transpose() -} - -/// Parse a `--gpu` CLI value into explicit device ordinals (empty for `auto`). -fn parse_gpu_indices_arg(value: Option<&str>) -> Result> { - Ok(rocm_engine_protocol::launch_gpu_indices( - parse_gpu_selection_arg(value)?.as_ref(), - )) -} - -fn current_unix_millis() -> u128 { - SystemTime::now() - .duration_since(UNIX_EPOCH) - .unwrap_or_default() - .as_millis() -} - -fn read_request() -> Result { - let mut buffer = String::new(); - std::io::stdin() - .read_to_string(&mut buffer) - .context("failed to read stdin for engine request")?; - serde_json::from_str(&buffer).context("failed to parse engine request envelope") -} - -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!(&mut handle)?; - Ok(()) -} - -#[cfg(test)] -mod tests { - use super::*; - - fn test_engine_recipe(engine: &str, contract_version: &str) -> EngineRecipeHint { - EngineRecipeHint { - contract_version: contract_version.to_owned(), - engine: engine.to_owned(), - required_flags: vec!["--jinja".to_owned()], - parser_settings: std::collections::BTreeMap::default(), - preferred_endpoint: None, - unsupported_combinations: Vec::new(), - notes: vec!["test recipe".to_owned()], - } - } - - #[test] - fn resolve_model_rejects_cpu_policy_without_fallback() { - let error = resolve_model_response(ResolveModelRequest { - model_ref: "Qwen/Qwen3.5".to_owned(), - runtime_id: None, - device_policy: Some(DevicePolicy::CpuOnly), - recipe_override: None, - engine_recipe: None, - }) - .expect_err("cpu policy should not resolve"); - assert!(error.to_string().contains("no CPU fallback is used")); - } - - #[test] - fn resolve_model_accepts_gpu_preferred_without_cpu_fallback() -> Result<()> { - let response = resolve_model_response(ResolveModelRequest { - model_ref: "tiny.gguf".to_owned(), - runtime_id: None, - device_policy: Some(DevicePolicy::GpuPreferred), - recipe_override: None, - engine_recipe: None, - })?; - - assert_eq!(response.device_policy, DevicePolicy::GpuRequired); - Ok(()) - } - - #[test] - fn resolve_model_accepts_gpu_required() -> Result<()> { - let response = resolve_model_response(ResolveModelRequest { - model_ref: "tiny.gguf".to_owned(), - runtime_id: None, - device_policy: Some(DevicePolicy::GpuRequired), - recipe_override: None, - engine_recipe: None, - })?; - - assert_eq!(response.device_policy, DevicePolicy::GpuRequired); - Ok(()) - } - - #[test] - fn resolve_model_canonicalizes_existing_gguf_path() -> Result<()> { - let path = std::env::temp_dir().join(format!( - "rocm-llama-model-{}-{}.gguf", - std::process::id(), - current_unix_millis() - )); - fs::write(&path, "fake")?; - let canonical = path.canonicalize()?.display().to_string(); - let response = resolve_model_response(ResolveModelRequest { - model_ref: path.display().to_string(), - runtime_id: None, - device_policy: Some(DevicePolicy::GpuRequired), - recipe_override: None, - engine_recipe: None, - })?; - fs::remove_file(&path).ok(); - - assert_eq!(response.canonical_model_id, canonical); - assert!(response.warnings.is_empty()); - Ok(()) - } - - #[test] - fn resolve_model_warns_for_missing_gguf_path() -> Result<()> { - let response = resolve_model_response(ResolveModelRequest { - model_ref: "missing-model.gguf".to_owned(), - runtime_id: None, - device_policy: Some(DevicePolicy::GpuRequired), - recipe_override: None, - engine_recipe: None, - })?; - - assert_eq!(response.canonical_model_id, "missing-model.gguf"); - assert!( - response - .warnings - .iter() - .any(|warning| warning.contains("not found locally")) - ); - Ok(()) - } - - #[test] - fn resolve_model_echoes_matching_engine_recipe() -> Result<()> { - let hint = test_engine_recipe(ENGINE_NAME, ENGINE_RECIPE_CONTRACT_VERSION); - let response = resolve_model_response(ResolveModelRequest { - model_ref: "tiny.gguf".to_owned(), - runtime_id: None, - device_policy: Some(DevicePolicy::GpuRequired), - recipe_override: None, - engine_recipe: Some(hint.clone()), - })?; - - assert_eq!(response.engine_recipe, Some(hint)); - Ok(()) - } - - #[test] - fn resolve_model_rejects_mismatched_engine_recipe() { - let error = resolve_model_response(ResolveModelRequest { - model_ref: "tiny.gguf".to_owned(), - runtime_id: None, - device_policy: Some(DevicePolicy::GpuRequired), - recipe_override: None, - engine_recipe: Some(test_engine_recipe( - "pytorch", - ENGINE_RECIPE_CONTRACT_VERSION, - )), - }) - .expect_err("mismatched engine recipe should fail"); - - assert!(error.to_string().contains("does not match adapter")); - } - - #[test] - fn resolve_model_rejects_unsupported_engine_recipe_contract() { - let error = resolve_model_response(ResolveModelRequest { - model_ref: "tiny.gguf".to_owned(), - runtime_id: None, - device_policy: Some(DevicePolicy::GpuRequired), - recipe_override: None, - engine_recipe: Some(test_engine_recipe(ENGINE_NAME, "999.0.0")), - }) - .expect_err("unsupported recipe contract should fail"); - - assert!(error.to_string().contains("unsupported")); - } - - #[test] - fn install_response_marks_llama_server_as_external_runtime() -> Result<()> { - let response = install_response(InstallRequest { - runtime_id: "external".to_owned(), - python_version: None, - env_root: None, - reinstall: false, - })?; - - assert_eq!( - response.runtime_kind.as_deref(), - Some("external_llama_server") - ); - assert_eq!(response.managed_env, Some(false)); - Ok(()) - } - - #[test] - fn managed_therock_env_uses_rocm_sdk_root_for_hip_apps() -> Result<()> { - let root = - std::env::temp_dir().join(format!("rocm-llama-therock-env-{}", current_unix_millis())); - let paths = AppPaths { - config_dir: root.join("config"), - data_dir: root.join("data"), - cache_dir: root.join("cache"), - }; - let sdk_root = root.join("_rocm_sdk_devel"); - let sdk_bin = sdk_root.join("bin"); - fs::create_dir_all(&sdk_bin)?; - let registry = paths.data_dir.join("runtimes").join("registry"); - fs::create_dir_all(®istry)?; - fs::write( - registry.join("release-pip-gfx120x-all.json"), - serde_json::to_vec_pretty(&json!({ - "runtime_key": "release-pip-gfx120x-all", - "runtime_id": "therock-release:gfx120X-all", - "format": "pip", - "installed_at_unix_ms": 2, - "rocm_sdk": { - "import_ok": true, - "root_path": sdk_root, - "bin_path": sdk_bin - } - }))?, - )?; - - let env = resolve_managed_therock_hip_runtime_env(&paths, Some("therock-release"))? - .expect("managed TheRock env"); - fs::remove_dir_all(root).ok(); - - assert_eq!(env.runtime_id, "therock-release:gfx120X-all"); - assert_eq!(env.runtime_key.as_deref(), Some("release-pip-gfx120x-all")); - assert!(env.root_path.ends_with("_rocm_sdk_devel")); - assert!(env.bin_path.ends_with("bin")); - Ok(()) - } - - #[test] - fn external_runtime_id_does_not_pick_managed_therock_env() -> Result<()> { - let root = - std::env::temp_dir().join(format!("rocm-llama-external-env-{}", current_unix_millis())); - let paths = AppPaths { - config_dir: root.join("config"), - data_dir: root.join("data"), - cache_dir: root.join("cache"), - }; - let registry = paths.data_dir.join("runtimes").join("registry"); - fs::create_dir_all(®istry)?; - fs::write( - registry.join("release-pip-gfx120x-all.json"), - serde_json::to_vec_pretty(&json!({ - "runtime_key": "release-pip-gfx120x-all", - "runtime_id": "therock-release:gfx120X-all", - "format": "pip", - "installed_at_unix_ms": 2, - "rocm_sdk": { - "import_ok": true, - "root_path": root.join("_rocm_sdk_devel"), - "bin_path": root.join("_rocm_sdk_devel").join("bin") - } - }))?, - )?; - - let env = resolve_managed_therock_hip_runtime_env(&paths, Some("external"))?; - fs::remove_dir_all(root).ok(); - - assert!(env.is_none()); - Ok(()) - } - - #[test] - fn prepend_path_entries_puts_therock_bin_first() -> Result<()> { - let root = if cfg!(windows) { - "C:\\rocm" - } else { - "/tmp/rocm" - }; - let current = std::env::join_paths([PathBuf::from("existing")])?; - let composed = prepend_path_entries(&[PathBuf::from(root).join("bin")], Some(current))?; - let parts = std::env::split_paths(&composed).collect::>(); - - assert_eq!(parts.first(), Some(&PathBuf::from(root).join("bin"))); - assert!(parts.iter().any(|part| part == &PathBuf::from("existing"))); - Ok(()) - } - - #[test] - fn therock_library_path_entries_include_sysdeps_for_hip_apps() { - let env = TheRockHipRuntimeEnv { - runtime_id: "therock-release:gfx120X-all".to_owned(), - runtime_key: Some("release-pip-gfx120x-all".to_owned()), - root_path: PathBuf::from(if cfg!(windows) { - "C:\\rocm" - } else { - "/tmp/rocm" - }), - bin_path: PathBuf::from(if cfg!(windows) { - "C:\\rocm\\bin" - } else { - "/tmp/rocm/bin" - }), - bin_paths: Vec::new(), - library_paths: Vec::new(), - source: "managed_runtime_manifest:release-pip-gfx120x-all".to_owned(), - }; - let entries = therock_library_path_entries(&env); - - assert!(entries.iter().any(|entry| entry.ends_with("lib"))); - assert!(entries.iter().any(|entry| entry.ends_with("lib64"))); - assert!( - entries - .iter() - .any(|entry| entry.ends_with(Path::new("lib").join("rocm_sysdeps").join("lib"))) - ); - } - - #[test] - fn llama_server_has_hip_backend_detects_sibling_backend() -> Result<()> { - let root = - std::env::temp_dir().join(format!("rocm-llama-hip-backend-{}", current_unix_millis())); - fs::create_dir_all(&root)?; - let server_path = root.join(if cfg!(windows) { - "llama-server.exe" - } else { - "llama-server" - }); - fs::write(&server_path, "exe")?; - let backend_name = if cfg!(windows) { - "ggml-hip.dll" - } else if cfg!(target_os = "macos") { - "libggml-hip.dylib" - } else { - "libggml-hip.so" - }; - fs::write(root.join(backend_name), "hip")?; - let server = LlamaServer { - program: server_path.display().to_string(), - display: "test".to_owned(), - }; - - let has_backend = llama_server_has_hip_backend(&server); - fs::remove_dir_all(root).ok(); - - assert!(has_backend); - Ok(()) - } - - #[test] - fn windows_therock_staging_places_runtime_next_to_llama_server() -> Result<()> { - let root = std::env::temp_dir().join(format!("rocm-llama-stage-{}", current_unix_millis())); - let paths = AppPaths { - config_dir: root.join("config"), - data_dir: root.join("data"), - cache_dir: root.join("cache"), - }; - let server_dir = root.join("server"); - fs::create_dir_all(&server_dir)?; - let server_path = server_dir.join("llama-server.exe"); - fs::write(&server_path, "exe")?; - fs::write(server_dir.join("ggml-hip.dll"), "backend")?; - - let sdk_root = root.join("_rocm_sdk_devel"); - let sdk_bin = sdk_root.join("bin"); - fs::create_dir_all(sdk_bin.join("rocblas").join("library"))?; - fs::create_dir_all(sdk_bin.join("hipblaslt").join("library"))?; - for filename in REQUIRED_WINDOWS_THEROCK_EXACT_DLLS { - fs::write(sdk_bin.join(filename), format!("managed {filename}"))?; - } - for filename in [ - "amd_comgr0713.dll", - "hiprtc-builtins07013.dll", - "hiprtc07013.dll", - ] { - fs::write(sdk_bin.join(filename), format!("managed {filename}"))?; - } - fs::write( - sdk_bin - .join("rocblas") - .join("library") - .join("TensileManifest.txt"), - "rocblas data", - )?; - fs::write( - sdk_bin - .join("hipblaslt") - .join("library") - .join("TensileLibrary.dat"), - "hipblaslt data", - )?; - let server = LlamaServer { - program: server_path.display().to_string(), - display: "test".to_owned(), - }; - let runtime_env = TheRockHipRuntimeEnv { - runtime_id: "therock-release:gfx120X-all".to_owned(), - runtime_key: Some("release-pip-gfx120x-all".to_owned()), - root_path: sdk_root, - bin_path: sdk_bin.clone(), - bin_paths: vec![sdk_bin], - library_paths: Vec::new(), - source: "managed_runtime_manifest:release-pip-gfx120x-all".to_owned(), - }; - - let prepared = - stage_windows_therock_llama_server(&paths, &server, &runtime_env, "svc:one")?; - let staged_dir = prepared.staged_dir.expect("staged dir"); - - assert!(prepared.program.ends_with("llama-server.exe")); - assert!(staged_dir.join("llama-server.exe").is_file()); - assert!(staged_dir.join("ggml-hip.dll").is_file()); - assert!(staged_dir.join("amd_comgr0713.dll").is_file()); - assert!(staged_dir.join("amdhip64_7.dll").is_file()); - assert!(staged_dir.join("hiprtc-builtins07013.dll").is_file()); - assert!(staged_dir.join("hiprtc07013.dll").is_file()); - assert!(staged_dir.join("hipblas.dll").is_file()); - assert!( - staged_dir - .join("rocblas") - .join("library") - .join("TensileManifest.txt") - .is_file() - ); - assert!( - staged_dir - .join("hipblaslt") - .join("library") - .join("TensileLibrary.dat") - .is_file() - ); - fs::remove_dir_all(root).ok(); - Ok(()) - } - - #[test] - fn windows_therock_staging_supports_split_runtime_wheel_bins() -> Result<()> { - let root = - std::env::temp_dir().join(format!("rocm-llama-stage-split-{}", current_unix_millis())); - let paths = AppPaths { - config_dir: root.join("config"), - data_dir: root.join("data"), - cache_dir: root.join("cache"), - }; - let server_dir = root.join("server"); - fs::create_dir_all(&server_dir)?; - let server_path = server_dir.join("llama-server.exe"); - fs::write(&server_path, "exe")?; - fs::write(server_dir.join("ggml-hip.dll"), "backend")?; - - let core_root = root.join("_rocm_sdk_core"); - let core_bin = core_root.join("bin"); - let libraries_root = root.join("_rocm_sdk_libraries_gfx120X_all"); - let libraries_bin = libraries_root.join("bin"); - fs::create_dir_all(&core_bin)?; - fs::create_dir_all(libraries_bin.join("rocblas").join("library"))?; - fs::create_dir_all(libraries_bin.join("hipblaslt").join("library"))?; - for filename in [ - "amd_comgr0713.dll", - "amdhip64_7.dll", - "hiprtc-builtins07013.dll", - "hiprtc07013.dll", - "rocm_kpack.dll", - ] { - fs::write(core_bin.join(filename), format!("core {filename}"))?; - } - for filename in [ - "hipblas.dll", - "libhipblaslt.dll", - "rocblas.dll", - "rocsolver.dll", - ] { - fs::write( - libraries_bin.join(filename), - format!("libraries {filename}"), - )?; - } - fs::write( - libraries_bin - .join("rocblas") - .join("library") - .join("TensileManifest.txt"), - "rocblas data", - )?; - fs::write( - libraries_bin - .join("hipblaslt") - .join("library") - .join("TensileLibrary.dat"), - "hipblaslt data", - )?; - let server = LlamaServer { - program: server_path.display().to_string(), - display: "test".to_owned(), - }; - let runtime_env = TheRockHipRuntimeEnv { - runtime_id: "therock-release:gfx120X-all".to_owned(), - runtime_key: Some("release-pip-gfx120x-all".to_owned()), - root_path: core_root, - bin_path: core_bin.clone(), - bin_paths: vec![core_bin, libraries_bin], - library_paths: Vec::new(), - source: "managed_runtime_manifest:release-pip-gfx120x-all".to_owned(), - }; - - let prepared = - stage_windows_therock_llama_server(&paths, &server, &runtime_env, "svc:split")?; - let staged_dir = prepared.staged_dir.expect("staged dir"); - - assert!(staged_dir.join("amd_comgr0713.dll").is_file()); - assert!(staged_dir.join("amdhip64_7.dll").is_file()); - assert!(staged_dir.join("hiprtc-builtins07013.dll").is_file()); - assert!(staged_dir.join("hiprtc07013.dll").is_file()); - assert!(staged_dir.join("hipblas.dll").is_file()); - assert!(staged_dir.join("rocblas.dll").is_file()); - assert!( - staged_dir - .join("rocblas") - .join("library") - .join("TensileManifest.txt") - .is_file() - ); - assert!( - staged_dir - .join("hipblaslt") - .join("library") - .join("TensileLibrary.dat") - .is_file() - ); - fs::remove_dir_all(root).ok(); - Ok(()) - } - - #[test] - fn windows_therock_staging_fails_when_required_runtime_file_is_missing() -> Result<()> { - let root = std::env::temp_dir().join(format!( - "rocm-llama-stage-missing-{}", - current_unix_millis() - )); - let paths = AppPaths { - config_dir: root.join("config"), - data_dir: root.join("data"), - cache_dir: root.join("cache"), - }; - let server_dir = root.join("server"); - fs::create_dir_all(&server_dir)?; - let server_path = server_dir.join("llama-server.exe"); - fs::write(&server_path, "exe")?; - fs::write(server_dir.join("ggml-hip.dll"), "backend")?; - - let sdk_root = root.join("_rocm_sdk_devel"); - let sdk_bin = sdk_root.join("bin"); - fs::create_dir_all(sdk_bin.join("rocblas").join("library"))?; - fs::create_dir_all(sdk_bin.join("hipblaslt").join("library"))?; - for filename in REQUIRED_WINDOWS_THEROCK_EXACT_DLLS { - if *filename != "amdhip64_7.dll" { - fs::write(sdk_bin.join(filename), format!("managed {filename}"))?; - } - } - for filename in [ - "amd_comgr0713.dll", - "hiprtc-builtins07013.dll", - "hiprtc07013.dll", - ] { - fs::write(sdk_bin.join(filename), format!("managed {filename}"))?; - } - fs::write( - sdk_bin.join("rocblas").join("library").join("manifest.txt"), - "rocblas data", - )?; - fs::write( - sdk_bin - .join("hipblaslt") - .join("library") - .join("manifest.txt"), - "hipblaslt data", - )?; - let server = LlamaServer { - program: server_path.display().to_string(), - display: "test".to_owned(), - }; - let runtime_env = TheRockHipRuntimeEnv { - runtime_id: "therock-release:gfx120X-all".to_owned(), - runtime_key: Some("release-pip-gfx120x-all".to_owned()), - root_path: sdk_root, - bin_path: sdk_bin.clone(), - bin_paths: vec![sdk_bin], - library_paths: Vec::new(), - source: "managed_runtime_manifest:release-pip-gfx120x-all".to_owned(), - }; - - let error = stage_windows_therock_llama_server(&paths, &server, &runtime_env, "svc:one") - .expect_err("missing amdhip64_7.dll must fail before launch"); - fs::remove_dir_all(root).ok(); - - let message = error.to_string(); - assert!(message.contains("amdhip64_7.dll")); - assert!(message.contains("no CPU fallback")); - Ok(()) - } - - #[test] - fn serve_http_cli_accepts_protocol_runtime_args() { - let cli = Cli::try_parse_from([ - "rocm-engine-llama-cpp", - "serve-http", - "svc", - "tiny.gguf", - "--host", - "127.0.0.1", - "--port", - "11435", - "--device-policy", - "gpu_required", - "--runtime-id", - "therock-release:gfx120X-all", - "--env-id", - "external-llama.cpp", - "--state-path", - "state.json", - ]) - .expect("serve-http should accept protocol runtime args"); - - match cli.command { - CommandKind::ServeHttp { - device_policy, - runtime_id, - env_id, - .. - } => { - assert_eq!(device_policy.as_deref(), Some("gpu_required")); - assert_eq!(runtime_id.as_deref(), Some("therock-release:gfx120X-all")); - assert_eq!(env_id.as_deref(), Some("external-llama.cpp")); - } - _ => panic!("expected serve-http command"), - } - } - - #[test] - fn launch_cli_accepts_runtime_selection_args() { - let cli = Cli::try_parse_from([ - "rocm-engine-llama-cpp", - "launch", - "svc", - "tiny.gguf", - "--device-policy", - "gpu_preferred", - "--runtime-id", - "runtime-1", - "--env-id", - "external-llama.cpp", - ]) - .expect("launch should accept protocol runtime args"); - - match cli.command { - CommandKind::Launch { - device_policy, - runtime_id, - env_id, - .. - } => { - assert_eq!(device_policy.as_deref(), Some("gpu_preferred")); - assert_eq!(runtime_id.as_deref(), Some("runtime-1")); - assert_eq!(env_id.as_deref(), Some("external-llama.cpp")); - } - _ => panic!("expected launch command"), - } - } - - #[test] - fn serve_http_request_accepts_gpu_required_for_runtime_validation() -> Result<()> { - let request = normalize_serve_http_request(ServeHttpRequest { - service_id: "svc".to_owned(), - model_ref: "tiny.gguf".to_owned(), - host: "127.0.0.1".to_owned(), - port: 11435, - device_policy: Some("gpu_required".to_owned()), - gpu_indices: Vec::new(), - runtime_id: None, - env_id: None, - state_path: PathBuf::from("state.json"), - log_path: None, - engine_recipe: None, - })?; - - assert_eq!(request.device_policy.as_deref(), Some("gpu_required")); - Ok(()) - } - - #[test] - fn gpu_required_fails_loudly_without_managed_therock_sdk() { - let error = - require_managed_therock_hip_runtime_env(None, Some("therock-release"), "gpu_required") - .expect_err("gpu_required should require managed TheRock SDK paths"); - - let message = error.to_string(); - assert!(message.contains("managed TheRock runtime manifest")); - assert!(message.contains("rocm_sdk.root_path")); - assert!(message.contains("rocm_sdk.bin_path")); - assert!(message.contains("no CPU fallback")); - } - - #[test] - fn gpu_required_rejects_non_manifest_therock_source_without_cpu_fallback() { - let error = require_managed_therock_hip_runtime_env( - Some(TheRockHipRuntimeEnv { - runtime_id: "therock-env".to_owned(), - runtime_key: None, - root_path: PathBuf::from("sdk"), - bin_path: PathBuf::from("sdk").join("bin"), - bin_paths: Vec::new(), - library_paths: Vec::new(), - source: "ROCM_CLI_THEROCK_SDK_ROOT".to_owned(), - }), - None, - "gpu_required", - ) - .expect_err("external TheRock SDK roots should not satisfy gpu_required"); - - let message = error.to_string(); - assert!(message.contains("managed runtime manifest")); - assert!(message.contains("no CPU fallback")); - } - - #[test] - fn parse_device_policy_arg_accepts_cli_aliases() -> Result<()> { - assert_eq!(parse_device_policy_arg(None)?, DevicePolicy::GpuRequired); - assert_eq!( - parse_device_policy_arg(Some("gpu"))?, - DevicePolicy::GpuPreferred - ); - assert_eq!(parse_device_policy_arg(Some("cpu"))?, DevicePolicy::CpuOnly); - assert!( - normalize_llama_device_policy(Some(DevicePolicy::CpuOnly)) - .unwrap_err() - .to_string() - .contains("no CPU fallback is used") - ); - Ok(()) - } - - #[test] - fn launch_serve_http_args_preserve_runtime_selection() { - let request = LaunchRequest { - service_id: "svc".to_owned(), - env_id: Some("external-llama.cpp".to_owned()), - runtime_id: Some("runtime-1".to_owned()), - model_ref: "tiny.gguf".to_owned(), - host: "127.0.0.1".to_owned(), - port: 11435, - device_policy: Some(DevicePolicy::GpuRequired), - endpoint_mode: Some("openai".to_owned()), - engine_recipe: None, - gpu_selection: None, - }; - - let args = serve_http_command_args( - &request, - Path::new("state.json"), - Some(Path::new("log.txt")), - ); - - assert!( - args.windows(2) - .any(|pair| pair[0] == "--device-policy" && pair[1] == "gpu_required") - ); - assert!( - args.windows(2) - .any(|pair| pair[0] == "--runtime-id" && pair[1] == "runtime-1") - ); - assert!( - args.windows(2) - .any(|pair| pair[0] == "--env-id" && pair[1] == "external-llama.cpp") - ); - assert!( - args.windows(2) - .any(|pair| pair[0] == "--state-path" && pair[1] == "state.json") - ); - assert!( - args.windows(2) - .any(|pair| pair[0] == "--log-path" && pair[1] == "log.txt") - ); - } - - #[cfg(windows)] - #[test] - fn windows_background_launcher_uses_hidden_start_process() { - let script = windows_background_launcher_script(); - - assert!(script.contains("Start-Process")); - assert!(script.contains("-WindowStyle Hidden")); - assert!(script.contains("-PassThru")); - assert!(script.contains("$args")); - } - - #[test] - fn llama_server_args_add_gpu_layers_for_gpu_required() { - let request = ServeHttpRequest { - service_id: "svc".to_owned(), - model_ref: "tiny.gguf".to_owned(), - host: "127.0.0.1".to_owned(), - port: 11435, - device_policy: Some("gpu_required".to_owned()), - gpu_indices: Vec::new(), - runtime_id: Some("therock-release".to_owned()), - env_id: None, - state_path: PathBuf::from("state.json"), - log_path: None, - engine_recipe: None, - }; - - let args = llama_server_args(&request, true); - - assert!( - args.windows(2) - .any(|pair| pair[0] == "-m" && pair[1] == "tiny.gguf") - ); - assert!( - args.windows(2) - .any(|pair| pair[0] == "--n-gpu-layers" && pair[1] == LLAMA_GPU_LAYERS_VALUE) - ); - } - - #[test] - fn llama_server_args_forward_engine_recipe_flags() { - let request = ServeHttpRequest { - service_id: "svc".to_owned(), - model_ref: "tiny.gguf".to_owned(), - host: "127.0.0.1".to_owned(), - port: 11435, - device_policy: Some("gpu_required".to_owned()), - gpu_indices: Vec::new(), - runtime_id: Some("therock-release".to_owned()), - env_id: None, - state_path: PathBuf::from("state.json"), - log_path: None, - engine_recipe: Some(test_engine_recipe( - ENGINE_NAME, - ENGINE_RECIPE_CONTRACT_VERSION, - )), - }; - - let args = llama_server_args(&request, true); - - assert!(args.iter().any(|arg| arg == "--jinja")); - } - - #[test] - fn serve_http_request_rejects_cpu_policy_without_fallback() { - let error = normalize_serve_http_request(ServeHttpRequest { - service_id: "svc".to_owned(), - model_ref: "tiny.gguf".to_owned(), - host: "127.0.0.1".to_owned(), - port: 11435, - device_policy: Some("cpu_only".to_owned()), - gpu_indices: Vec::new(), - runtime_id: None, - env_id: None, - state_path: PathBuf::from("state.json"), - log_path: None, - engine_recipe: None, - }) - .expect_err("cpu policy should be rejected before launch"); - - assert!(error.to_string().contains("no CPU fallback is used")); - } - - #[test] - fn healthcheck_device_reflects_gpu_required_state() { - let state = json!({ - "device_policy": "gpu_required", - "therock_runtime_env": { - "runtime_id": "therock-release:gfx120X-all" - } - }); - - assert_eq!(healthcheck_device_from_state(Some(&state)), "rocm_gpu"); - assert_eq!( - healthcheck_device_from_state(Some(&json!({ "device_policy": "cpu_only" }))), - "unknown" - ); - } - - #[test] - fn endpoint_url_falls_back_to_host_and_port() { - let state = json!({ - "host": "127.0.0.1", - "port": 12345 - }); - assert_eq!( - endpoint_url_from_state(&state), - Some("http://127.0.0.1:12345/v1".to_owned()) - ); - } - - #[test] - fn endpoint_url_and_parser_support_ipv6_loopback() { - let state = json!({ - "host": "::1", - "port": 12345 - }); - assert_eq!( - endpoint_url_from_state(&state), - Some("http://[::1]:12345/v1".to_owned()) - ); - assert_eq!( - parse_http_endpoint("http://[::1]:12345/v1"), - Some(("::1".to_owned(), 12345)) - ); - } - - #[test] - fn endpoint_response_errors_without_service_state() { - let error = endpoint_response(EndpointRequest { - service_id: format!("missing-{}", current_unix_millis()), - }) - .expect_err("missing service state should not produce a default endpoint"); - - assert!(error.to_string().contains("service state not found")); - } - - #[test] - fn stdio_protocol_routes_all_methods_without_side_effects() { - let service_id = format!( - "missing-protocol-{}-{}", - std::process::id(), - current_unix_millis() - ); - let success_cases = [ - (EngineMethod::Detect, json!({})), - (EngineMethod::Capabilities, json!({})), - ( - EngineMethod::ResolveModel, - json!({ - "model_ref": "missing.gguf", - "device_policy": "gpu_required" - }), - ), - ( - EngineMethod::Healthcheck, - json!({ - "service_id": service_id.as_str() - }), - ), - ( - EngineMethod::Logs, - json!({ - "service_id": service_id.as_str(), - "tail_lines": 4 - }), - ), - ( - EngineMethod::Stop, - json!({ - "service_id": service_id.as_str(), - "force": false - }), - ), - ]; - - for (method, payload) in success_cases { - let response = handle_envelope(EngineRequestEnvelope { method, payload }); - assert!( - response.ok, - "expected protocol method to return a typed success envelope: {:?}", - response.error - ); - } - - let endpoint = handle_envelope(EngineRequestEnvelope { - method: EngineMethod::Endpoint, - payload: json!({ - "service_id": service_id.as_str() - }), - }); - assert!(!endpoint.ok); - assert_eq!( - endpoint.error.as_ref().map(|error| error.code.as_str()), - Some("request_failed") - ); - - for method in [EngineMethod::Install, EngineMethod::Launch] { - let response = handle_envelope(EngineRequestEnvelope { - method, - payload: json!({}), - }); - assert!(!response.ok); - assert_eq!( - response.error.as_ref().map(|error| error.code.as_str()), - Some("invalid_payload") - ); - } - } - - #[test] - fn stop_prefers_recorded_llama_server_pid_over_wrapper_pid() { - let state = json!({ - "pid": 100, - "wrapper_pid": 100, - "server_pid": 200, - }); - assert_eq!(pid_to_terminate_from_state(&state), Some(200)); - - let wrapper_only = json!({ "pid": 100 }); - assert_eq!(pid_to_terminate_from_state(&wrapper_only), Some(100)); - } - - #[test] - fn tail_lines_returns_suffix() -> Result<()> { - let path = std::env::temp_dir().join(format!( - "rocm-llama-tail-{}-{}.log", - std::process::id(), - current_unix_millis() - )); - fs::write(&path, "a\nb\nc\n")?; - let lines = tail_lines(&path, 2)?; - fs::remove_file(path).ok(); - assert_eq!(lines, vec!["b".to_owned(), "c".to_owned()]); - Ok(()) - } -} diff --git a/engines/llama-cpp/src/main.rs b/engines/llama-cpp/src/main.rs deleted file mode 100644 index 32d834aa..00000000 --- a/engines/llama-cpp/src/main.rs +++ /dev/null @@ -1,7 +0,0 @@ -// Copyright © Advanced Micro Devices, Inc., or its affiliates. -// -// SPDX-License-Identifier: MIT - -fn main() -> anyhow::Result<()> { - rocm_engine_llama_cpp::run_cli() -} diff --git a/engines/pytorch/Cargo.toml b/engines/pytorch/Cargo.toml deleted file mode 100644 index 893d4190..00000000 --- a/engines/pytorch/Cargo.toml +++ /dev/null @@ -1,22 +0,0 @@ -[package] -name = "rocm-engine-pytorch" -version.workspace = true -edition.workspace = true -license.workspace = true -repository.workspace = true -rust-version.workspace = true -publish.workspace = true - -[lints] -workspace = true - -[dependencies] -anyhow.workspace = true -async-stream.workspace = true -axum.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 -tokio.workspace = true diff --git a/engines/pytorch/src/lib.rs b/engines/pytorch/src/lib.rs deleted file mode 100644 index be7d1803..00000000 --- a/engines/pytorch/src/lib.rs +++ /dev/null @@ -1,3809 +0,0 @@ -// Copyright © Advanced Micro Devices, Inc., or its affiliates. -// -// SPDX-License-Identifier: MIT - -use anyhow::{Context, Result, bail}; -use clap::{Parser, Subcommand, ValueEnum}; -use rocm_core::{ - AppPaths, DEFAULT_LOCAL_PORT, ModelRecipeRecord, detect_host_therock_family, ensure_uv_binary, - extract_first_gfx_token, format_host_port, format_http_base_url, interactive_terminal, - normalize_runtime_path_for_host, normalize_runtime_path_text_for_host, - normalize_therock_family, require_nonempty, - resolve_model_recipe as resolve_shared_model_recipe, runtime_is_windows, unix_time_millis, - uv_command_env, uv_pip_freeze_args, uv_pip_install_base, uv_venv_args, -}; -use rocm_engine_protocol::{ - DetectRequest, DetectResponse, DevicePolicy, ENGINE_RECIPE_CONTRACT_VERSION, EndpointRequest, - EndpointResponse, EngineCapabilities, EngineDeviceAvailability, EngineMethod, EngineRecipeHint, - EngineRequestEnvelope, EngineResponseEnvelope, GpuSelection, HealthcheckRequest, - HealthcheckResponse, InstallRequest, InstallResponse, LaunchRequest, LaunchResponse, - LogsRequest, LogsResponse, ResolveModelRequest, ResolveModelResponse, StopRequest, - StopResponse, -}; -use serde::{Deserialize, Serialize}; -use serde_json::{Value, json}; -use std::collections::{BTreeMap, VecDeque, hash_map::DefaultHasher}; -use std::fs; -use std::hash::{Hash, Hasher}; -use std::io::{self, Read, Write}; -use std::net::{TcpStream, ToSocketAddrs}; -use std::path::{Path, PathBuf}; -use std::process::{Command, ExitStatus, Stdio}; -use std::thread; -use std::time::{Duration, SystemTime, UNIX_EPOCH}; - -const ENGINE_NAME: &str = "pytorch"; -const DEFAULT_RUNTIME_ID: &str = "therock-release"; -const THEROCK_SIMPLE_INDEX_BASE: &str = "https://rocm.nightlies.amd.com/v2"; -const PYTHON_WORKER_SOURCE: &str = include_str!("python_worker.py"); -const ENGINE_DEPENDENCIES: &[&str] = &[ - "fastapi", - "uvicorn", - "pydantic", - "transformers<5", - "safetensors", - "tokenizers", - "huggingface_hub<1", - "jinja2", -]; -const THEROCK_TORCH_PACKAGES: &[&str] = &["torch", "torchvision", "torchaudio"]; -const TORCH_STACK_DEPENDENCIES: &[&str] = &["accelerate"]; -const DEFAULT_LOG_TAIL_LINES: usize = 200; -const MAX_LOG_TAIL_LINES: usize = 1000; -const HEALTHCHECK_TIMEOUT_MS: u64 = 500; -const REMOVE_ENV_RETRIES: usize = 8; -const REMOVE_ENV_RETRY_DELAY: Duration = Duration::from_millis(250); -const KNOWN_THEROCK_FAMILIES: &[&str] = &[ - "gfx94X-dcgpu", - "gfx950-dcgpu", - "gfx110X-all", - "gfx1151", - "gfx120X-all", -]; - -#[derive(Parser, Debug)] -#[command( - name = "rocm-engine-pytorch", - about = "rocm-cli PyTorch engine", - version -)] -struct Cli { - #[command(subcommand)] - command: CommandKind, -} - -#[derive(Subcommand, Debug)] -enum CommandKind { - Detect, - Capabilities, - Install { - #[arg(long, default_value = DEFAULT_RUNTIME_ID)] - runtime_id: String, - #[arg(long)] - python_version: Option, - #[arg(long)] - reinstall: bool, - }, - ResolveModel { - model_ref: String, - #[arg(long)] - device_policy: Option, - }, - Launch { - service_id: String, - model_ref: String, - #[arg(long, default_value = "127.0.0.1")] - host: String, - #[arg(long, default_value_t = DEFAULT_LOCAL_PORT)] - port: u16, - #[arg(long)] - device_policy: Option, - #[arg(long)] - runtime_id: Option, - #[arg(long)] - env_id: Option, - #[arg(long)] - gpu: Option, - }, - Stdio, - #[command(hide = true)] - ServeHttp { - service_id: String, - model_ref: String, - #[arg(long, default_value = "127.0.0.1")] - host: String, - #[arg(long, default_value_t = DEFAULT_LOCAL_PORT)] - port: u16, - #[arg(long, default_value = "gpu_required")] - device_policy: String, - #[arg(long)] - env_id: Option, - #[arg(long)] - runtime_id: Option, - #[arg(long)] - state_path: PathBuf, - #[arg(long)] - engine_recipe_json: Option, - #[arg(long)] - gpu: Option, - }, -} - -#[derive(Clone, Copy, Debug, Eq, PartialEq, ValueEnum)] -enum DevicePolicyArg { - GpuRequired, - GpuPreferred, - CpuOnly, -} - -#[derive(Debug, Clone, Serialize, Deserialize)] -struct EngineEnvManifest { - env_id: String, - runtime_id: String, - requested_python_version: Option, - python_launcher: String, - python_executable: String, - env_path: PathBuf, - manifest_path: PathBuf, - lock_path: PathBuf, - installed_packages: Vec, - lock_hash: String, - #[serde(default)] - pip_cache_dir: Option, - #[serde(default)] - therock_channel: Option, - #[serde(default)] - therock_family: Option, - #[serde(default)] - therock_index_url: Option, - #[serde(default)] - therock_packages: Vec, - #[serde(default)] - torch_runtime_probe: Option, - warnings: Vec, -} - -#[derive(Debug, Clone)] -struct PythonLauncher { - program: String, - args: Vec, - display: String, -} - -#[derive(Debug, Clone, Copy)] -enum TheRockChannel { - Release, - Nightly, -} - -impl TheRockChannel { - const fn as_str(self) -> &'static str { - match self { - Self::Release => "release", - Self::Nightly => "nightly", - } - } -} - -#[derive(Debug, Clone)] -struct TheRockRuntimeRequest { - channel: TheRockChannel, - family_override: Option, -} - -#[derive(Debug, Clone)] -struct TheRockTorchResolution { - channel: TheRockChannel, - family: String, - index_url: String, - packages: Vec, - source: String, -} - -#[derive(Debug, Clone, Deserialize)] -struct RuntimeRegistryManifest { - runtime_key: String, - runtime_id: String, - channel: String, - format: String, - family: String, - index_url: Option, - selected_artifact_url: Option, - python_executable: Option, -} - -#[derive(Debug, Clone, Copy)] -enum ModelFamily { - Generic, - Qwen, - Glm, - Llama, - Gpt2, -} - -#[derive(Debug, Clone)] -struct ModelRecipe { - canonical_model_id: String, - task: &'static str, - source: String, - loader: &'static str, - trust_remote_code: bool, - chat_template_mode: &'static str, - preferred_dtype: String, - device_policy: DevicePolicy, - estimated_memory: String, - min_gpu_mem_gb: Option, - warnings: Vec, -} - -#[derive(Debug, Clone, Serialize, Deserialize)] -struct TorchRuntimeProbe { - #[serde(default)] - import_ok: bool, - #[serde(default)] - torch_version: Option, - #[serde(default)] - cuda_available: bool, - #[serde(default)] - device_count: u32, - #[serde(default)] - devices: Vec, - #[serde(default)] - error: Option, - #[serde(default)] - rocm_sdk: Option, - #[serde(default)] - torch_rocm_init: Option, -} - -#[derive(Debug, Clone, Serialize, Deserialize)] -struct RocmSdkProbe { - #[serde(default)] - import_ok: bool, - #[serde(default)] - version: Option, - #[serde(default)] - site_packages: Option, - #[serde(default)] - default_target_family: Option, - #[serde(default)] - available_target_families: Vec, - #[serde(default)] - resolved_target_family: Option, - #[serde(default)] - error: Option, -} - -#[derive(Debug, Clone, Serialize, Deserialize)] -struct TorchRocmInitProbe { - #[serde(default)] - present: bool, - #[serde(default)] - path: Option, - #[serde(default)] - check_version: Option, - #[serde(default)] - preload_shortnames: Vec, - #[serde(default)] - error: Option, -} - -#[derive(Debug, Clone, Deserialize)] -struct ServiceRecordSnapshot { - #[serde(default)] - engine: Option, - #[serde(default)] - status: Option, - #[serde(default)] - supervisor_pid: Option, - #[serde(default)] - engine_pid: Option, - #[serde(default)] - log_path: Option, - #[serde(default)] - engine_state_path: Option, - #[serde(default)] - endpoint_url: Option, -} - -#[derive(Debug)] -struct ServiceFiles { - record_path: PathBuf, - record: Option, - record_matches_engine: bool, - state_path: PathBuf, - log_path: PathBuf, -} - -#[derive(Debug, Clone, Copy, Eq, PartialEq)] -enum PidTermination { - Terminated, - Failed, -} - -#[tokio::main] -pub async fn run_cli() -> Result<()> { - let cli = Cli::parse(); - match cli.command { - CommandKind::Detect => { - print_json(&detect_response())?; - } - CommandKind::Capabilities => { - print_json(&capabilities())?; - } - CommandKind::Install { - runtime_id, - python_version, - reinstall, - } => { - let response = install_response(InstallRequest { - runtime_id, - python_version, - env_root: None, - reinstall, - })?; - print_json(&response)?; - } - CommandKind::ResolveModel { - model_ref, - device_policy, - } => { - let response = resolve_model_response(ResolveModelRequest { - model_ref, - runtime_id: None, - device_policy: device_policy.map(Into::into), - recipe_override: None, - engine_recipe: None, - })?; - print_json(&response)?; - } - CommandKind::Launch { - service_id, - model_ref, - host, - port, - device_policy, - runtime_id, - env_id, - gpu, - } => { - let response = launch_service(LaunchRequest { - service_id, - env_id, - runtime_id, - model_ref, - host, - port, - device_policy: device_policy.map(Into::into), - endpoint_mode: Some("openai".to_owned()), - engine_recipe: None, - gpu_selection: parse_gpu_selection_arg(gpu.as_deref())?, - })?; - print_json(&response)?; - } - CommandKind::Stdio => { - let envelope = read_request()?; - let response = handle_envelope(envelope); - print_json(&response)?; - } - CommandKind::ServeHttp { - service_id, - model_ref, - host, - port, - device_policy, - env_id, - runtime_id, - state_path, - engine_recipe_json, - gpu, - } => { - serve_http( - service_id, - model_ref, - host, - port, - parse_device_policy(&device_policy)?, - parse_gpu_indices_arg(gpu.as_deref())?, - env_id, - runtime_id, - state_path, - parse_engine_recipe_json(engine_recipe_json)?, - )?; - } - } - Ok(()) -} - -pub fn builtin_handle_envelope(envelope: EngineRequestEnvelope) -> EngineResponseEnvelope { - handle_envelope(envelope) -} - -#[allow(clippy::too_many_arguments)] -pub fn builtin_serve_http( - service_id: String, - model_ref: String, - host: String, - port: u16, - device_policy: DevicePolicy, - gpu_indices: Vec, - env_id: Option, - runtime_id: Option, - state_path: PathBuf, - engine_recipe: Option, -) -> Result<()> { - serve_http( - service_id, - model_ref, - host, - port, - device_policy, - gpu_indices, - env_id, - runtime_id, - state_path, - engine_recipe, - ) -} - -fn handle_envelope(envelope: EngineRequestEnvelope) -> EngineResponseEnvelope { - match envelope.method { - EngineMethod::Detect => { - deserialize_and_respond::(envelope.payload, |_| { - Ok(detect_response()) - }) - } - EngineMethod::Capabilities => EngineResponseEnvelope::success(capabilities()), - EngineMethod::Install => { - deserialize_and_respond::(envelope.payload, install_response) - } - EngineMethod::ResolveModel => deserialize_and_respond::( - envelope.payload, - resolve_model_response, - ), - EngineMethod::Launch => { - deserialize_and_respond::(envelope.payload, launch_service) - } - EngineMethod::Healthcheck => deserialize_and_respond::( - envelope.payload, - healthcheck_service, - ), - EngineMethod::Endpoint => { - deserialize_and_respond::(envelope.payload, endpoint_response) - } - EngineMethod::Stop => { - deserialize_and_respond::(envelope.payload, stop_service) - } - EngineMethod::Logs => { - deserialize_and_respond::(envelope.payload, logs_response) - } - } -} - -fn deserialize_and_respond( - payload: serde_json::Value, - handler: F, -) -> EngineResponseEnvelope -where - T: for<'de> Deserialize<'de>, - F: FnOnce(T) -> Result, - U: Serialize, -{ - match serde_json::from_value::(payload) { - Ok(request) => match handler(request) { - Ok(response) => EngineResponseEnvelope::success(response), - Err(error) => EngineResponseEnvelope::failure("request_failed", error.to_string()), - }, - Err(error) => EngineResponseEnvelope::failure("invalid_payload", error.to_string()), - } -} - -fn detect_response() -> DetectResponse { - use std::fmt::Write as _; - let manifest = latest_env_manifest().ok().flatten(); - let installed = manifest.is_some(); - let detected_family = detect_host_therock_family(); - let env_id = manifest.as_ref().map(|value| value.env_id.clone()); - let python_version = manifest - .as_ref() - .and_then(|value| value.requested_python_version.clone()) - .or_else(|| Some(default_python_version().to_owned())); - let transformers_version = manifest - .as_ref() - .and_then(|value| find_installed_package(&value.installed_packages, "transformers")); - - let mut notes = Vec::new(); - let torch_probe = manifest.as_ref().and_then(|manifest| { - match probe_torch_runtime(&manifest.python_executable) { - Ok(probe) => { - if probe.import_ok { - notes.push(format!( - "torch probe: cuda_available={} device_count={}", - probe.cuda_available, probe.device_count - )); - if !probe.devices.is_empty() { - notes.push(format!("torch devices: {}", probe.devices.join(", "))); - } - if let Some(init) = probe.torch_rocm_init.as_ref() - && init.present - { - let version = init.check_version.as_deref().unwrap_or(""); - notes.push(format!( - "torch._rocm_init: check_version={} preload_count={}", - version, - init.preload_shortnames.len() - )); - } - if let Some(sdk) = probe.rocm_sdk.as_ref() - && sdk.import_ok - { - let version = sdk.version.as_deref().unwrap_or(""); - let family = sdk.resolved_target_family.as_deref().unwrap_or(""); - notes.push(format!( - "rocm_sdk: version={version} target_family={family}" - )); - } - } else if let Some(error) = probe.error.as_deref() { - notes.push(format!("torch probe import failed: {error}")); - } - Some(probe) - } - Err(error) => { - notes.push(format!("torch probe failed: {error}")); - None - } - } - }); - let torch_version = torch_probe - .as_ref() - .and_then(|probe| probe.torch_version.clone()) - .or_else(|| { - manifest - .as_ref() - .and_then(|value| find_installed_package(&value.installed_packages, "torch")) - }); - if let Some(manifest) = &manifest { - notes.push(format!( - "managed env detected at {}", - manifest.env_path.display() - )); - if let Some(family) = manifest.therock_family.as_deref() { - notes.push(format!("TheRock family: {family}")); - } - if let Some(channel) = manifest.therock_channel.as_deref() { - notes.push(format!("TheRock channel: {channel}")); - } - if let Some(index_url) = manifest.therock_index_url.as_deref() { - notes.push(format!("TheRock index: {index_url}")); - } - notes.extend(manifest.warnings.iter().cloned()); - } else { - notes.push("no managed PyTorch envs found; run `rocm engines install pytorch`".to_owned()); - } - - let rocm_gpu_available = torch_probe - .as_ref() - .map_or_else(|| detected_family.is_some(), |probe| probe.cuda_available); - let rocm_gpu_reason = match torch_probe.as_ref() { - Some(probe) if probe.cuda_available => { - let mut reason = format!("torch.cuda reports {} device(s)", probe.device_count); - if !probe.devices.is_empty() { - let _ = write!(reason, ": {}", probe.devices.join(", ")); - } - Some(reason) - } - Some(probe) => probe - .error - .clone() - .or_else(|| Some("torch.cuda is not available in the managed env".to_owned())), - None => detected_family - .as_ref() - .map(|family| format!("detected host family {family}")) - .or_else(|| Some("no supported TheRock GPU family detected on this host".to_owned())), - }; - - DetectResponse { - installed, - env_id, - runtime_kind: Some("managed_python".to_owned()), - runtime_executable: manifest - .as_ref() - .map(|manifest| manifest.python_executable.clone()), - managed_env: Some(true), - python_version, - torch_version, - transformers_version, - available_devices: vec![ - EngineDeviceAvailability { - kind: "cpu".to_owned(), - available: false, - reason: Some( - "rocm-cli does not offer PyTorch CPU serving; use ROCm GPU execution" - .to_owned(), - ), - }, - EngineDeviceAvailability { - kind: "rocm_gpu".to_owned(), - available: rocm_gpu_available, - reason: rocm_gpu_reason, - }, - ], - capabilities: capabilities(), - notes, - } -} - -fn capabilities() -> EngineCapabilities { - EngineCapabilities { - cpu: false, - rocm_gpu: true, - openai_compatible: true, - tool_calling: true, - quantized_models: "limited".to_owned(), - reasoning_parser: false, - } -} - -fn install_response(request: InstallRequest) -> Result { - require_nonempty(&request.runtime_id, "runtime_id")?; - let manifest = create_or_update_env_manifest(&request)?; - Ok(InstallResponse { - env_id: manifest.env_id, - env_path: manifest.env_path.display().to_string(), - python_executable: manifest.python_executable.clone(), - runtime_kind: Some("managed_python".to_owned()), - runtime_executable: Some(manifest.python_executable), - managed_env: Some(true), - installed_packages: manifest.installed_packages, - capabilities: capabilities(), - lock_hash: manifest.lock_hash, - warnings: manifest.warnings, - }) -} - -fn resolve_model_response(request: ResolveModelRequest) -> Result { - require_nonempty(&request.model_ref, "model_ref")?; - let engine_recipe = accepted_engine_recipe(request.engine_recipe)?; - let recipe = resolve_model_recipe(&request.model_ref)?; - let device_policy = normalize_pytorch_device_policy( - request - .device_policy - .unwrap_or_else(|| default_device_policy_for_recipe(&recipe)), - )?; - Ok(ResolveModelResponse { - canonical_model_id: recipe.canonical_model_id, - task: recipe.task.to_owned(), - source: recipe.source, - revision: "main".to_owned(), - loader: recipe.loader.to_owned(), - trust_remote_code: recipe.trust_remote_code, - chat_template_mode: recipe.chat_template_mode.to_owned(), - dtype: recipe.preferred_dtype.clone(), - device_policy, - estimated_memory: recipe.estimated_memory, - launch_defaults: json!({ - "host": "127.0.0.1", - "port": DEFAULT_LOCAL_PORT, - "endpoint_mode": "openai" - }), - engine_recipe, - warnings: recipe.warnings, - }) -} - -fn accepted_engine_recipe( - engine_recipe: Option, -) -> Result> { - if let Some(hint) = &engine_recipe { - if hint.engine != ENGINE_NAME { - bail!( - "engine_recipe target `{}` does not match adapter `{}`", - hint.engine, - ENGINE_NAME - ); - } - if hint.contract_version != ENGINE_RECIPE_CONTRACT_VERSION { - bail!( - "engine_recipe contract `{}` is unsupported; expected `{}`", - hint.contract_version, - ENGINE_RECIPE_CONTRACT_VERSION - ); - } - } - Ok(engine_recipe) -} - -fn parse_engine_recipe_json(value: Option) -> Result> { - value - .map(|text| { - serde_json::from_str::(&text) - .context("failed to parse engine recipe JSON") - }) - .transpose() - .and_then(accepted_engine_recipe) -} - -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.clone())?; - let device_policy = normalize_pytorch_device_policy( - request.device_policy.unwrap_or(DevicePolicy::GpuRequired), - )?; - - let paths = AppPaths::discover()?; - paths.ensure()?; - fs::create_dir_all(paths.engine_logs_dir(ENGINE_NAME))?; - fs::create_dir_all(paths.engine_state_dir(ENGINE_NAME))?; - - let log_path = paths - .engine_logs_dir(ENGINE_NAME) - .join(format!("{}.log", request.service_id)); - let state_path = paths - .engine_state_dir(ENGINE_NAME) - .join(format!("{}.json", request.service_id)); - let log_file = fs::File::create(&log_path) - .with_context(|| format!("failed to create {}", log_path.display()))?; - let log_file_err = log_file - .try_clone() - .context("failed to clone log file handle")?; - - let current_exe = - std::env::current_exe().context("failed to discover current engine binary")?; - let child = Command::new(command_path(¤t_exe)) - .arg("serve-http") - .arg(&request.service_id) - .arg(&request.model_ref) - .arg("--host") - .arg(&request.host) - .arg("--port") - .arg(request.port.to_string()) - .arg("--device-policy") - .arg(device_policy_name(&device_policy)) - .args(optional_arg("--env-id", request.env_id.as_deref())) - .args(optional_arg("--runtime-id", request.runtime_id.as_deref())) - .args(engine_recipe_json_arg(engine_recipe.as_ref())?) - .args(gpu_indices_arg(request.gpu_selection.as_ref())) - .arg("--state-path") - .arg(&state_path) - .stdin(Stdio::null()) - .stdout(Stdio::from(log_file)) - .stderr(Stdio::from(log_file_err)) - .spawn() - .context("failed to spawn pytorch serve-http process")?; - - fs::write( - &state_path, - serde_json::to_vec_pretty(&json!({ - "engine": ENGINE_NAME, - "service_id": request.service_id, - "model_ref": request.model_ref, - "host": request.host, - "port": request.port, - "pid": child.id(), - "status": "starting", - "engine_recipe": engine_recipe, - "engine_recipe_required_flags": engine_recipe_launch_args(engine_recipe.as_ref()) - }))?, - ) - .with_context(|| format!("failed to write {}", state_path.display()))?; - - let endpoint_url = format!("{}/v1", format_http_base_url(&request.host, request.port)); - Ok(LaunchResponse { - service_id: request.service_id, - pid: child.id(), - endpoint_url, - log_path: log_path.display().to_string(), - state_path: state_path.display().to_string(), - }) -} - -fn healthcheck_service(request: HealthcheckRequest) -> Result { - require_nonempty(&request.service_id, "service_id")?; - let files = service_files(&request.service_id)?; - let (state, state_modified, state_error) = read_service_state(&files.state_path); - let state_status = state - .as_ref() - .and_then(|value| value_string(value, "status")) - .unwrap_or_else(|| "unknown".to_owned()); - let endpoint_url = state.as_ref().and_then(endpoint_url_from_state); - let (health_payload, probe_error) = if state_status == "ready" || state_status == "running" { - match endpoint_url.as_deref() { - Some(endpoint_url) => match query_health_endpoint(endpoint_url) { - Ok(payload) => (Some(payload), None), - Err(error) => (None, Some(error.to_string())), - }, - None => (None, None), - } - } else { - (None, None) - }; - - Ok(build_healthcheck_response( - state.as_ref(), - state_modified, - state_error, - health_payload.as_ref(), - probe_error, - SystemTime::now(), - )) -} - -fn logs_response(request: LogsRequest) -> Result { - require_nonempty(&request.service_id, "service_id")?; - let files = service_files(&request.service_id)?; - let limit = normalize_tail_limit(request.tail_lines); - let recent_lines = if files.log_path.is_file() { - tail_lines(&files.log_path, limit)? - } else { - Vec::new() - }; - - Ok(LogsResponse { - log_path: files.log_path.display().to_string(), - recent_lines, - }) -} - -fn endpoint_response(request: EndpointRequest) -> Result { - require_nonempty(&request.service_id, "service_id")?; - let files = service_files(&request.service_id)?; - endpoint_response_from_files(&files) -} - -fn endpoint_response_from_files(files: &ServiceFiles) -> Result { - let (state, _, _) = read_service_state(&files.state_path); - let endpoint_url = state - .as_ref() - .and_then(endpoint_url_from_state) - .or_else(|| { - files - .record - .as_ref() - .filter(|_| files.record_matches_engine) - .and_then(|record| record.endpoint_url.clone()) - }) - .unwrap_or_else(|| format!("http://127.0.0.1:{DEFAULT_LOCAL_PORT}/v1")); - Ok(openai_endpoint_response(endpoint_url)) -} - -fn openai_endpoint_response(endpoint_url: String) -> EndpointResponse { - EndpointResponse { - endpoint_url, - api_style: "openai".to_owned(), - supported_routes: vec![ - "/healthz".to_owned(), - "/v1/models".to_owned(), - "/v1/chat/completions".to_owned(), - "/v1/completions".to_owned(), - ], - } -} - -fn stop_service(request: StopRequest) -> Result { - require_nonempty(&request.service_id, "service_id")?; - let files = service_files(&request.service_id)?; - let (state, _, _) = read_service_state(&files.state_path); - let mut pids = Vec::new(); - if let Some(pid) = state.as_ref().and_then(|value| value_u32(value, "pid")) { - pids.push(pid); - } - if files.record_matches_engine - && let Some(record) = files.record.as_ref() - { - if let Some(pid) = record.engine_pid { - pids.push(pid); - } - if request.force - && let Some(pid) = record.supervisor_pid - { - pids.push(pid); - } - } - dedupe_pids(&mut pids); - - let had_pid = !pids.is_empty(); - let mut stopped_any = false; - let mut failed_any = false; - for pid in pids { - match terminate_pid(pid, request.force) { - PidTermination::Terminated => stopped_any = true, - PidTermination::Failed => failed_any = true, - } - } - - let state_status = state - .as_ref() - .and_then(|value| value_string(value, "status")) - .or_else(|| { - files - .record - .as_ref() - .and_then(|record| record.status.clone()) - }); - let already_terminal = matches!( - state_status.as_deref(), - Some("stopped" | "failed" | "exited") - ); - let stopped = stopped_any || (had_pid && !failed_any) || (!had_pid && already_terminal); - if stopped { - mark_json_status(&files.state_path, "stopped")?; - if files.record_matches_engine && files.record_path.is_file() { - mark_json_status(&files.record_path, "stopped")?; - } - } - - Ok(StopResponse { - stopped, - graceful: stopped && !request.force && !failed_any, - }) -} - -fn service_files(service_id: &str) -> Result { - let paths = AppPaths::discover()?; - let record_path = paths.service_manifest_path(service_id); - let record = load_service_record(&record_path)?; - let record_matches_engine = record - .as_ref() - .and_then(|record| record.engine.as_deref()) - .map_or_else(|| record.is_some(), |engine| engine == ENGINE_NAME); - let state_path = record - .as_ref() - .filter(|_| record_matches_engine) - .and_then(|record| record.engine_state_path.clone()) - .unwrap_or_else(|| { - paths - .engine_state_dir(ENGINE_NAME) - .join(format!("{service_id}.json")) - }); - let engine_log_path = paths - .engine_logs_dir(ENGINE_NAME) - .join(format!("{service_id}.log")); - let log_path = record - .as_ref() - .filter(|_| record_matches_engine) - .and_then(|record| record.log_path.clone()) - .or_else(|| { - if record.is_some() { - return None; - } - let service_log_path = paths.service_log_path(service_id); - service_log_path.is_file().then_some(service_log_path) - }) - .unwrap_or(engine_log_path); - - Ok(ServiceFiles { - record_path, - record, - record_matches_engine, - state_path, - log_path, - }) -} - -fn load_service_record(path: &Path) -> Result> { - if !path.is_file() { - return Ok(None); - } - let bytes = fs::read(path).with_context(|| format!("failed to read {}", path.display()))?; - serde_json::from_slice(&bytes) - .with_context(|| format!("failed to parse {}", path.display())) - .map(Some) -} - -fn read_service_state(path: &Path) -> (Option, Option, Option) { - let modified = fs::metadata(path) - .and_then(|metadata| metadata.modified()) - .ok(); - match fs::read_to_string(path) { - Ok(content) => match serde_json::from_str::(&content) { - Ok(value) => (Some(value), modified, None), - Err(error) => ( - None, - modified, - Some(format!("failed to parse {}: {error}", path.display())), - ), - }, - Err(error) if error.kind() == std::io::ErrorKind::NotFound => ( - None, - None, - Some(format!("state file not found: {}", path.display())), - ), - Err(error) => ( - None, - modified, - Some(format!("failed to read {}: {error}", path.display())), - ), - } -} - -fn build_healthcheck_response( - state: Option<&Value>, - state_modified: Option, - state_error: Option, - health_payload: Option<&Value>, - probe_error: Option, - now: SystemTime, -) -> HealthcheckResponse { - let state_status = state - .and_then(|value| value_string(value, "status")) - .unwrap_or_else(|| "unknown".to_owned()); - let health_status = health_payload.and_then(|value| value_string(value, "status")); - let status = if matches!(state_status.as_str(), "ready" | "running") - && health_payload.is_none() - && probe_error.is_some() - { - "unreachable".to_owned() - } else if health_status.as_deref() == Some("ok") { - "ready".to_owned() - } else { - state_status - }; - - let device = health_payload - .and_then(|value| value_string(value, "device")) - .or_else(|| state.and_then(|value| value_string(value, "device"))) - .or_else(|| state.and_then(|value| value_string(value, "input_device"))) - .unwrap_or_else(|| "unknown".to_owned()); - let uptime_sec = health_payload - .and_then(|value| value_f64(value, "loaded_at")) - .and_then(|loaded_at| uptime_from_epoch_secs(loaded_at, now)) - .or_else(|| state_modified.and_then(|modified| uptime_from_modified(modified, now))) - .unwrap_or(0); - let queue_depth = health_payload - .and_then(|value| value_u32(value, "queue_depth")) - .or_else(|| state.and_then(|value| value_u32(value, "queue_depth"))) - .unwrap_or(0); - let last_error = state_error - .or_else(|| state.and_then(|value| value_string(value, "last_error"))) - .or_else(|| state.and_then(|value| value_string(value, "error"))) - .or_else(|| { - if status == "unreachable" { - probe_error - } else { - None - } - }); - let tokens_per_sec = health_payload - .and_then(|value| value_f64(value, "tokens_per_sec")) - .or_else(|| state.and_then(|value| value_f64(value, "tokens_per_sec"))) - .map(|value| value as f32); - let model_loaded = status == "ready"; - - HealthcheckResponse { - status, - model_loaded, - device, - uptime_sec, - queue_depth, - last_error, - tokens_per_sec, - } -} - -fn endpoint_url_from_state(state: &Value) -> Option { - value_string(state, "endpoint_url").or_else(|| { - let host = value_string(state, "host")?; - let port = value_u32(state, "port")?; - let port = u16::try_from(port).ok()?; - Some(format!("{}/v1", format_http_base_url(&host, port))) - }) -} - -fn query_health_endpoint(endpoint_url: &str) -> Result { - let (host, port) = parse_http_endpoint(endpoint_url) - .with_context(|| format!("unsupported endpoint URL `{endpoint_url}`"))?; - let addr = (host.as_str(), port) - .to_socket_addrs() - .with_context(|| format!("failed to resolve {host}:{port}"))? - .next() - .with_context(|| format!("no socket addresses resolved for {host}:{port}"))?; - let timeout = Duration::from_millis(HEALTHCHECK_TIMEOUT_MS); - let mut stream = TcpStream::connect_timeout(&addr, timeout) - .with_context(|| format!("failed to connect to {host}:{port}"))?; - stream.set_read_timeout(Some(timeout)).ok(); - stream.set_write_timeout(Some(timeout)).ok(); - let host_header = format_host_port(&host, port); - write!( - stream, - "GET /healthz HTTP/1.1\r\nHost: {host_header}\r\nConnection: close\r\n\r\n" - ) - .context("failed to send health request")?; - - let mut response = String::new(); - stream - .read_to_string(&mut response) - .context("failed to read health response")?; - let (headers, body) = response - .split_once("\r\n\r\n") - .context("health response was missing HTTP body")?; - let status_line = headers.lines().next().unwrap_or_default(); - if !status_line.contains(" 200 ") { - bail!("health endpoint returned {status_line}"); - } - serde_json::from_str(body.trim()).context("failed to parse health response body") -} - -fn parse_http_endpoint(endpoint_url: &str) -> Option<(String, u16)> { - let without_scheme = endpoint_url.trim().strip_prefix("http://")?; - let authority = without_scheme.split('/').next()?.trim(); - if authority.is_empty() { - return None; - } - if let Some(rest) = authority.strip_prefix('[') { - let end = rest.find(']')?; - let host = rest[..end].to_owned(); - let port = rest[end + 1..].strip_prefix(':')?.parse().ok()?; - return Some((host, port)); - } - let (host, port) = authority.rsplit_once(':')?; - Some((host.to_owned(), port.parse().ok()?)) -} - -fn normalize_tail_limit(tail_lines: Option) -> usize { - tail_lines - .unwrap_or(DEFAULT_LOG_TAIL_LINES) - .min(MAX_LOG_TAIL_LINES) -} - -fn tail_lines(path: &Path, limit: usize) -> Result> { - let content = - fs::read_to_string(path).with_context(|| format!("failed to read {}", path.display()))?; - Ok(tail_lines_from_text(&content, limit)) -} - -fn tail_lines_from_text(content: &str, limit: usize) -> Vec { - if limit == 0 { - return Vec::new(); - } - let mut lines = VecDeque::with_capacity(limit); - for line in content.lines() { - if lines.len() == limit { - lines.pop_front(); - } - lines.push_back(line.to_owned()); - } - lines.into_iter().collect() -} - -fn dedupe_pids(pids: &mut Vec) { - let current_pid = std::process::id(); - let mut unique = Vec::new(); - for pid in pids.drain(..) { - if pid == 0 || pid == current_pid || unique.contains(&pid) { - continue; - } - unique.push(pid); - } - *pids = unique; -} - -fn terminate_pid(pid: u32, _force: bool) -> PidTermination { - match rocm_core::terminate_process(pid) { - Ok(()) => PidTermination::Terminated, - Err(_) => PidTermination::Failed, - } -} - -fn mark_json_status(path: &Path, status: &str) -> Result<()> { - let mut value = if path.is_file() { - let content = fs::read_to_string(path) - .with_context(|| format!("failed to read {}", path.display()))?; - serde_json::from_str::(&content).unwrap_or_else(|_| json!({})) - } else { - json!({}) - }; - if !value.is_object() { - value = json!({}); - } - let object = value.as_object_mut().expect("object checked above"); - object.insert("engine".to_owned(), json!(ENGINE_NAME)); - object.insert("status".to_owned(), json!(status)); - object.insert( - "stopped_at_unix_ms".to_owned(), - json!(current_unix_millis()), - ); - 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).context("failed to serialize status update")?, - ) - .with_context(|| format!("failed to write {}", path.display())) -} - -fn value_string(value: &Value, key: &str) -> Option { - value - .get(key) - .and_then(Value::as_str) - .filter(|value| !value.trim().is_empty()) - .map(ToOwned::to_owned) -} - -fn value_u32(value: &Value, key: &str) -> Option { - value - .get(key) - .and_then(Value::as_u64) - .and_then(|value| u32::try_from(value).ok()) -} - -fn value_f64(value: &Value, key: &str) -> Option { - value.get(key).and_then(Value::as_f64) -} - -fn uptime_from_epoch_secs(loaded_at: f64, now: SystemTime) -> Option { - if !loaded_at.is_finite() || loaded_at < 0.0 { - return None; - } - let now_secs = now.duration_since(UNIX_EPOCH).ok()?.as_secs_f64(); - Some((now_secs - loaded_at).max(0.0) as u64) -} - -fn uptime_from_modified(modified: SystemTime, now: SystemTime) -> Option { - now.duration_since(modified) - .ok() - .map(|value| value.as_secs()) -} - -fn current_unix_millis() -> u128 { - SystemTime::now() - .duration_since(UNIX_EPOCH) - .unwrap_or_default() - .as_millis() -} - -fn remove_managed_env_dir_with_retry(env_path: &Path) -> Result<()> { - if !env_path.exists() { - return Ok(()); - } - - let mut delete_target = env_path.to_owned(); - let mut last_error = None; - for attempt in 0..=REMOVE_ENV_RETRIES { - clear_readonly_recursive(&delete_target).ok(); - match fs::remove_dir_all(&delete_target) { - Ok(()) => return Ok(()), - Err(error) => { - last_error = Some(error); - if attempt == 1 - && delete_target == env_path - && let Some(file_name) = env_path.file_name().and_then(|value| value.to_str()) - && let Some(parent) = env_path.parent() - { - let moved_as = - parent.join(format!("{file_name}.deleting-{}", current_unix_millis())); - if fs::rename(env_path, &moved_as).is_ok() { - eprintln!( - "Managed PyTorch env was locked, moved it aside so a fresh env can be created: {}", - moved_as.display() - ); - delete_target = moved_as; - if !env_path.exists() { - return Ok(()); - } - } - } - if attempt < REMOVE_ENV_RETRIES { - thread::sleep(REMOVE_ENV_RETRY_DELAY); - } - } - } - } - - let error = last_error.map_or_else(|| "unknown error".to_owned(), |error| error.to_string()); - bail!( - "failed to remove managed PyTorch env {} after {} attempts. Close any ROCm, Python, or model server process using this folder and try again. Last error: {error}", - env_path.display(), - REMOVE_ENV_RETRIES + 1 - ) -} - -#[cfg(windows)] -#[allow(clippy::permissions_set_readonly_false)] -fn clear_readonly_recursive(path: &Path) -> Result<()> { - if !path.exists() { - return Ok(()); - } - let metadata = - fs::metadata(path).with_context(|| format!("failed to stat {}", path.display()))?; - let mut permissions = metadata.permissions(); - if permissions.readonly() { - permissions.set_readonly(false); - fs::set_permissions(path, permissions) - .with_context(|| format!("failed to clear readonly attribute on {}", path.display()))?; - } - if metadata.is_dir() { - for entry in - fs::read_dir(path).with_context(|| format!("failed to read {}", path.display()))? - { - clear_readonly_recursive(&entry?.path())?; - } - } - Ok(()) -} - -#[cfg(not(windows))] -const fn clear_readonly_recursive(_path: &Path) -> Result<()> { - Ok(()) -} - -fn create_or_update_env_manifest(request: &InstallRequest) -> Result { - let paths = AppPaths::discover()?; - paths.ensure()?; - let engine_envs_dir = request.env_root.as_ref().map_or_else( - || paths.engine_envs_dir(ENGINE_NAME), - |root| { - normalize_runtime_path_for_host(root) - .join(ENGINE_NAME) - .join("envs") - }, - ); - fs::create_dir_all(&engine_envs_dir)?; - fs::create_dir_all(paths.engine_locks_dir(ENGINE_NAME))?; - fs::create_dir_all(paths.engine_manifests_dir(ENGINE_NAME))?; - - let runtime_python_executable = if request.python_version.is_none() { - runtime_python_executable_for_selector(&paths, &request.runtime_id)? - } else { - None - }; - let effective_python_version = request.python_version.clone().or(runtime_python_executable - .as_deref() - .map(runtime_python_major_minor) - .transpose()?); - - let env_id = managed_env_id(&request.runtime_id, effective_python_version.as_deref()); - let env_path = engine_envs_dir.join(&env_id); - let lock_path = paths - .engine_locks_dir(ENGINE_NAME) - .join(format!("{env_id}.txt")); - let manifest_path = paths - .engine_manifests_dir(ENGINE_NAME) - .join(format!("{env_id}.json")); - let existing_manifest = if manifest_path.is_file() { - Some(load_manifest(&manifest_path)?) - } else { - None - }; - let therock_resolution = if std::env::var("ROCM_CLI_PYTORCH_PACKAGE_SPEC").is_ok() { - None - } else { - resolve_therock_torch_resolution(&paths, &request.runtime_id)? - }; - - if !request.reinstall - && let Some(manifest) = existing_manifest - && manifest.env_path.is_dir() - && (manifest_has_torch(&manifest) || therock_resolution.is_none()) - { - return Ok(manifest); - } - - remove_managed_env_dir_with_retry(&env_path)?; - - let launcher = discover_python_launcher( - effective_python_version.as_deref(), - runtime_python_executable.as_deref(), - )?; - let uv = ensure_uv_binary(&paths)?; - let launcher_path = PathBuf::from(&launcher.program); - let venv_args = uv_venv_args(&launcher_path, &env_path); - run_uv_command( - &uv, - venv_args.iter().map(String::as_str), - "create managed pytorch venv", - )?; - - let python_executable = venv_python_path(&env_path); - - let mut engine_dep_args = uv_pip_install_base(&python_executable); - engine_dep_args.extend(["--only-binary".to_owned(), ":all:".to_owned()]); - engine_dep_args.extend( - ENGINE_DEPENDENCIES - .iter() - .map(std::string::ToString::to_string), - ); - run_uv_progress_command( - &uv, - engine_dep_args.iter().map(String::as_str), - "install managed pytorch engine dependencies", - )?; - - let python_executable_string = python_executable.to_string_lossy().to_string(); - let mut warnings = Vec::new(); - let mut therock_channel = None; - let mut therock_family = None; - let mut therock_index_url = None; - let mut therock_packages = Vec::new(); - let maybe_torch_spec = std::env::var("ROCM_CLI_PYTORCH_PACKAGE_SPEC").ok(); - let maybe_extra_index = std::env::var("ROCM_CLI_PYTORCH_EXTRA_INDEX_URL").ok(); - if let Some(torch_spec) = maybe_torch_spec.as_deref() { - let mut args = uv_pip_install_base(&python_executable); - args.push("--only-binary".to_owned()); - args.push(":all:".to_owned()); - if let Some(extra_index) = maybe_extra_index.as_deref() { - args.push("--extra-index-url".to_owned()); - args.push(extra_index.to_owned()); - } - args.push(torch_spec.to_owned()); - run_uv_progress_command( - &uv, - args.iter().map(String::as_str), - "install torch package into managed pytorch env", - )?; - let mut stack_args = uv_pip_install_base(&python_executable); - stack_args.extend(["--only-binary".to_owned(), ":all:".to_owned()]); - stack_args.extend( - TORCH_STACK_DEPENDENCIES - .iter() - .map(std::string::ToString::to_string), - ); - run_uv_progress_command( - &uv, - stack_args.iter().map(String::as_str), - "install pytorch engine runtime dependencies", - )?; - warnings.push( - "using manual torch package override from ROCM_CLI_PYTORCH_PACKAGE_SPEC".to_owned(), - ); - } else { - match therock_resolution { - Some(resolution) => { - install_therock_torch_packages(&uv, &python_executable, &resolution)?; - let mut stack_args = uv_pip_install_base(&python_executable); - stack_args.extend(["--only-binary".to_owned(), ":all:".to_owned()]); - stack_args.extend( - TORCH_STACK_DEPENDENCIES - .iter() - .map(std::string::ToString::to_string), - ); - run_uv_progress_command( - &uv, - stack_args.iter().map(String::as_str), - "install pytorch engine runtime dependencies", - )?; - therock_channel = Some(resolution.channel.as_str().to_owned()); - therock_family = Some(resolution.family.clone()); - therock_index_url = Some(resolution.index_url.clone()); - therock_packages = resolution.packages.clone(); - warnings.push(format!( - "installed TheRock PyTorch packages from {} ({}, source={})", - resolution.index_url, resolution.family, resolution.source - )); - } - None => { - warnings.push( - "torch installation deferred because no supported TheRock GPU family could be resolved; set ROCM_CLI_THEROCK_FAMILY or use a runtime_id like therock-release:gfx950-dcgpu".to_owned(), - ); - } - } - } - - let installed_packages = { - let freeze_args = uv_pip_freeze_args(&python_executable); - let freeze = capture_uv_command( - &uv, - freeze_args.iter().map(String::as_str), - "capture managed pytorch env lockfile", - )?; - freeze - .lines() - .filter(|line| !line.trim().is_empty()) - .map(ToOwned::to_owned) - .collect::>() - }; - let freeze = installed_packages.join("\n") + "\n"; - fs::write(&lock_path, &freeze) - .with_context(|| format!("failed to write {}", lock_path.display()))?; - let lock_hash = simple_hash(&freeze); - let torch_runtime_probe = { - match probe_torch_runtime(&python_executable_string) { - Ok(probe) => { - if therock_family.is_some() { - match probe.rocm_sdk.as_ref() { - Some(sdk) if sdk.import_ok => { - if let (Some(expected), Some(actual)) = ( - probe - .torch_rocm_init - .as_ref() - .and_then(|init| init.check_version.as_deref()), - sdk.version.as_deref(), - ) && expected != actual - { - warnings.push(format!( - "torch._rocm_init expects ROCm SDK {expected}, but rocm_sdk reports {actual}" - )); - } - } - Some(sdk) => warnings.push(format!( - "rocm_sdk import probe failed in managed PyTorch env: {}", - sdk.error.as_deref().unwrap_or("unknown error") - )), - None => warnings.push( - "rocm_sdk probe was not reported by the managed PyTorch env".to_owned(), - ), - } - if !probe - .torch_rocm_init - .as_ref() - .is_some_and(|init| init.present) - { - warnings.push( - "torch._rocm_init was not found; TheRock library preloading may be unavailable" - .to_owned(), - ); - } - } - Some(probe) - } - Err(error) => { - warnings.push(format!("managed PyTorch runtime probe failed: {error}")); - None - } - } - }; - - let manifest = EngineEnvManifest { - env_id, - runtime_id: request.runtime_id.clone(), - requested_python_version: effective_python_version, - python_launcher: launcher.display, - python_executable: python_executable.display().to_string(), - env_path, - manifest_path, - lock_path, - installed_packages, - lock_hash, - pip_cache_dir: None, - therock_channel, - therock_family, - therock_index_url, - therock_packages, - torch_runtime_probe, - warnings, - }; - write_manifest(&manifest)?; - Ok(manifest) -} - -fn latest_env_manifest() -> Result> { - let paths = AppPaths::discover()?; - let manifests_dir = paths.engine_manifests_dir(ENGINE_NAME); - if !manifests_dir.is_dir() { - return Ok(None); - } - - let mut manifests = Vec::new(); - for entry in fs::read_dir(&manifests_dir) - .with_context(|| format!("failed to read {}", manifests_dir.display()))? - { - let entry = entry?; - let path = entry.path(); - if path.extension().and_then(|value| value.to_str()) != Some("json") { - continue; - } - manifests.push(load_manifest(&path)?); - } - manifests.sort_by(|left, right| left.env_id.cmp(&right.env_id)); - Ok(manifests.pop()) -} - -fn load_manifest(path: &Path) -> Result { - let bytes = fs::read(path) - .with_context(|| format!("failed to read engine manifest {}", path.display()))?; - let mut manifest: EngineEnvManifest = serde_json::from_slice(&bytes) - .with_context(|| format!("failed to parse engine manifest {}", path.display()))?; - normalize_manifest_paths_for_host(&mut manifest); - manifest.python_executable = resolve_manifest_python_executable(&manifest) - .display() - .to_string(); - Ok(manifest) -} - -fn normalize_manifest_paths_for_host(manifest: &mut EngineEnvManifest) { - manifest.python_executable = normalize_runtime_path_text_for_host(&manifest.python_executable); - manifest.env_path = normalize_runtime_path_for_host(&manifest.env_path); - manifest.manifest_path = normalize_runtime_path_for_host(&manifest.manifest_path); - manifest.lock_path = normalize_runtime_path_for_host(&manifest.lock_path); - manifest.pip_cache_dir = manifest - .pip_cache_dir - .as_ref() - .map(|path| normalize_runtime_path_for_host(path)); -} - -fn write_manifest(manifest: &EngineEnvManifest) -> Result<()> { - fs::write( - &manifest.manifest_path, - serde_json::to_vec_pretty(manifest).context("failed to serialize engine manifest")?, - ) - .with_context(|| format!("failed to write {}", manifest.manifest_path.display()))?; - Ok(()) -} - -fn discover_python_launcher( - requested_version: Option<&str>, - preferred_python: Option<&str>, -) -> Result { - let mut candidates = Vec::new(); - if let Some(python) = preferred_python - && !python.trim().is_empty() - { - let normalized = normalize_runtime_path_text_for_host(python); - if runtime_is_windows() { - return Ok(PythonLauncher { - program: normalized.clone(), - args: Vec::new(), - display: normalized, - }); - } - candidates.push(PythonLauncher { - program: normalized.clone(), - args: Vec::new(), - display: normalized, - }); - } - if runtime_is_windows() { - if let Some(version) = requested_version { - candidates.push(PythonLauncher { - program: "py".to_owned(), - args: vec![format!("-{version}")], - display: format!("py -{version}"), - }); - } - candidates.push(PythonLauncher { - program: "py".to_owned(), - args: Vec::new(), - display: "py".to_owned(), - }); - candidates.push(PythonLauncher { - program: "python".to_owned(), - args: Vec::new(), - display: "python".to_owned(), - }); - } else { - if let Some(version) = requested_version { - candidates.push(PythonLauncher { - program: format!("python{version}"), - args: Vec::new(), - display: format!("python{version}"), - }); - } - candidates.push(PythonLauncher { - program: "python3".to_owned(), - args: Vec::new(), - display: "python3".to_owned(), - }); - candidates.push(PythonLauncher { - program: "python".to_owned(), - args: Vec::new(), - display: "python".to_owned(), - }); - } - - let mut attempts = Vec::new(); - for launcher in candidates { - let status = Command::new(command_program(&launcher.program)) - .args(&launcher.args) - .arg("--version") - .stdout(Stdio::null()) - .stderr(Stdio::null()) - .status(); - attempts.push(format!("{} -> {:?}", launcher.display, status)); - if matches!(status, Ok(value) if value.success()) { - return Ok(launcher); - } - } - - bail!( - "unable to locate a usable Python launcher for the pytorch engine (runtime_os={}, requested_version={}, preferred_python={}, attempts={})", - rocm_core::runtime_os_name(), - requested_version.unwrap_or(""), - preferred_python.unwrap_or(""), - attempts.join("; ") - ) -} - -fn runtime_python_executable_for_selector( - paths: &AppPaths, - selector: &str, -) -> Result> { - let Some(manifest) = load_runtime_registry_manifest(paths, selector)? else { - return Ok(None); - }; - if manifest.format != "wheel" { - return Ok(None); - } - Ok(manifest.python_executable) -} - -fn runtime_python_major_minor(python_executable: &str) -> Result { - let output = capture_command( - python_executable, - [ - "-c", - "import sys; print(f'{sys.version_info.major}.{sys.version_info.minor}')", - ], - "inspect managed runtime Python version", - )?; - let version = output.trim(); - if version.is_empty() { - bail!("managed runtime Python did not report a version"); - } - Ok(version.to_owned()) -} - -fn venv_python_path(env_path: &Path) -> PathBuf { - if runtime_is_windows() { - env_path.join("Scripts").join("python.exe") - } else { - env_path.join("bin").join("python") - } -} - -fn resolve_manifest_python_executable(manifest: &EngineEnvManifest) -> PathBuf { - let recorded = PathBuf::from(&manifest.python_executable); - if recorded.is_file() { - return recorded; - } - - for candidate in venv_python_candidates(&manifest.env_path) { - if candidate.is_file() { - return candidate; - } - } - - recorded -} - -fn venv_python_candidates(env_path: &Path) -> Vec { - if runtime_is_windows() { - return vec![ - env_path.join("Scripts").join("python.exe"), - env_path.join("Scripts").join("python3.exe"), - ]; - } - - let bin = env_path.join("bin"); - vec![ - bin.join("python"), - bin.join("python3"), - bin.join("python3.12"), - bin.join("python3.11"), - bin.join("python3.10"), - bin.join("python3.9"), - ] -} - -#[allow(clippy::too_many_arguments)] -fn serve_http( - service_id: String, - model_ref: String, - host: String, - port: u16, - device_policy: DevicePolicy, - gpu_indices: Vec, - env_id: Option, - runtime_id: Option, - state_path: PathBuf, - engine_recipe: Option, -) -> Result<()> { - let manifest = ensure_service_env(runtime_id.as_deref(), env_id.as_deref())?; - let device_policy = normalize_pytorch_device_policy(device_policy)?; - let recipe = apply_pytorch_engine_recipe_overrides( - resolve_model_recipe(&model_ref)?, - engine_recipe.as_ref(), - )?; - let worker_script = materialize_python_worker()?; - let endpoint_url = format!("{}/v1", format_http_base_url(&host, port)); - fs::write( - &state_path, - serde_json::to_vec_pretty(&json!({ - "engine": ENGINE_NAME, - "service_id": service_id, - "model_ref": model_ref, - "status": "starting", - "pid": std::process::id(), - "host": host, - "port": port, - "device_policy": device_policy_name(&device_policy), - "env_id": manifest.env_id, - "runtime_id": manifest.runtime_id, - "python_executable": manifest.python_executable, - "preferred_dtype": &recipe.preferred_dtype, - "estimated_memory": &recipe.estimated_memory, - "trust_remote_code": recipe.trust_remote_code, - "engine_recipe": engine_recipe, - "engine_recipe_required_flags": engine_recipe_launch_args(engine_recipe.as_ref()), - "endpoint_url": endpoint_url - }))?, - )?; - let mut worker_command = Command::new(command_program(&manifest.python_executable)); - worker_command - .arg(&worker_script) - .arg("--service-id") - .arg(&service_id) - .arg("--model-ref") - .arg(&model_ref) - .arg("--host") - .arg(&host) - .arg("--port") - .arg(port.to_string()) - .arg("--device-policy") - .arg(device_policy_name(&device_policy)) - .arg("--state-path") - .arg(&state_path) - .arg("--env-id") - .arg(&manifest.env_id) - .arg("--runtime-id") - .arg(&manifest.runtime_id) - .arg("--preferred-dtype") - .arg(&recipe.preferred_dtype) - .args(optional_arg_owned( - "--min-gpu-mem-gb", - recipe.min_gpu_mem_gb.map(|value| value.to_string()), - )) - .args(flag_arg("--trust-remote-code", recipe.trust_remote_code)) - .env("PYTHONUNBUFFERED", "1") - .env("TOKENIZERS_PARALLELISM", "false") - .stdin(Stdio::null()); - rocm_engine_protocol::apply_gpu_visibility(&mut worker_command, &gpu_indices); - worker_command - .stdout(Stdio::inherit()) - .stderr(Stdio::inherit()); - let mut child = worker_command - .spawn() - .context("failed to start python worker for pytorch engine")?; - - let status = child - .wait() - .context("failed waiting for pytorch python worker")?; - if status.success() { - mark_json_status(&state_path, "stopped")?; - Ok(()) - } else { - let current_status = read_service_state(&state_path) - .0 - .and_then(|value| value_string(&value, "status")); - if matches!(current_status.as_deref(), Some("stopped" | "stopping")) { - mark_json_status(&state_path, "stopped")?; - return Ok(()); - } - mark_json_status(&state_path, "failed")?; - bail!("pytorch worker exited with status {status} for service {service_id}") - } -} - -fn canonical_model_id(model_ref: &str) -> String { - if let Ok(Some(recipe)) = resolve_shared_model_recipe(model_ref) - && shared_recipe_supports_engine(&recipe, ENGINE_NAME) - { - return recipe.canonical_model_id; - } - - match model_ref.to_ascii_lowercase().as_str() { - "qwen" | "qwen2.5" | "qwen2.5-1.5b" => "Qwen/Qwen2.5-1.5B-Instruct".to_owned(), - "qwen-tiny" | "tiny-qwen" | "qwen2.5-0.5b" => "Qwen/Qwen2.5-0.5B-Instruct".to_owned(), - "qwen3.5" => "Qwen/Qwen3.5-4B".to_owned(), - "qwen32b" | "qwen3.5-32b" | "qwen3-32b" => "Qwen/Qwen3-32B-FP8".to_owned(), - "glm5" | "glm-5" => "zai-org/GLM-5-FP8".to_owned(), - "llama3.2" | "llama" => "meta-llama/Llama-3.2-3B-Instruct".to_owned(), - "tiny-gpt2" | "gpt2tiny" => "sshleifer/tiny-gpt2".to_owned(), - other if other.contains('/') => model_ref.to_owned(), - _ => model_ref.to_owned(), - } -} - -fn resolve_model_recipe(model_ref: &str) -> Result { - if let Some(recipe) = known_recipe_for_model(model_ref)? { - return Ok(recipe); - } - - let canonical_model_id = canonical_model_id(model_ref); - ensure_pytorch_model_supported(&canonical_model_id)?; - let source = if model_ref.contains('/') { - "huggingface".to_owned() - } else { - "alias".to_owned() - }; - let family = infer_model_family(&canonical_model_id); - let trust_remote_code = matches!(family, ModelFamily::Glm); - let preferred_dtype = preferred_dtype_for_model(&canonical_model_id, family); - let min_gpu_mem_gb = estimate_gpu_memory_gib(&canonical_model_id); - let estimated_memory = format_estimated_memory(&canonical_model_id, min_gpu_mem_gb); - - let mut warnings = Vec::new(); - if trust_remote_code { - warnings.push( - "this model family is configured with trust_remote_code enabled by recipe".to_owned(), - ); - } - if let Some(min_gpu_mem_gb) = min_gpu_mem_gb { - if min_gpu_mem_gb >= 48 { - warnings.push(format!( - "this model looks large (~{min_gpu_mem_gb} GiB GPU memory recommended); implicit CPU serving is disabled for safety" - )); - } else if min_gpu_mem_gb >= 16 { - warnings.push(format!( - "this model may need about {min_gpu_mem_gb} GiB of GPU memory for comfortable serving" - )); - } - } - - Ok(ModelRecipe { - canonical_model_id, - task: "chat", - source, - loader: "transformers", - trust_remote_code, - chat_template_mode: "auto", - preferred_dtype: preferred_dtype.to_owned(), - device_policy: default_device_policy_for_memory(min_gpu_mem_gb), - estimated_memory, - min_gpu_mem_gb, - warnings, - }) -} - -fn known_recipe_for_model(model_ref: &str) -> Result> { - let Some(recipe) = resolve_shared_model_recipe(model_ref)? else { - return Ok(None); - }; - if !shared_recipe_supports_engine(&recipe, ENGINE_NAME) { - return Ok(None); - } - let device_policy = parse_device_policy(&recipe.device_policy) - .unwrap_or_else(|_| default_device_policy_for_memory(recipe.min_gpu_mem_gb)); - let canonical_model_id = recipe.canonical_model_id; - ensure_pytorch_model_supported(&canonical_model_id)?; - let source = recipe.source; - let trust_remote_code = recipe.trust_remote_code; - let dtype = recipe.dtype; - let min_gpu_mem_gb = recipe.min_gpu_mem_gb; - let warnings = recipe.warnings; - Ok(Some(build_known_recipe( - &canonical_model_id, - source, - trust_remote_code, - &dtype, - device_policy, - min_gpu_mem_gb, - warnings, - ))) -} - -fn shared_recipe_supports_engine(recipe: &ModelRecipeRecord, engine: &str) -> bool { - recipe - .preferred_engines - .iter() - .any(|candidate| candidate.eq_ignore_ascii_case(engine)) - || recipe - .engine_recipes - .iter() - .any(|candidate| candidate.engine.eq_ignore_ascii_case(engine)) -} - -fn ensure_pytorch_model_supported(canonical_model_id: &str) -> Result<()> { - if canonical_model_id.eq_ignore_ascii_case("Qwen/Qwen3.5-4B") { - bail!( - "Qwen/Qwen3.5-4B is not supported by the managed PyTorch engine yet: the current Transformers line reports unknown architecture `qwen3_5`. Use `qwen` for the recommended Qwen2.5 1.5B local assistant recipe, or use an engine/runtime that explicitly supports Qwen3.5." - ); - } - Ok(()) -} - -fn apply_pytorch_engine_recipe_overrides( - mut recipe: ModelRecipe, - engine_recipe: Option<&EngineRecipeHint>, -) -> Result { - let Some(engine_recipe) = engine_recipe else { - return Ok(recipe); - }; - - let flags = &engine_recipe.required_flags; - let mut index = 0; - while index < flags.len() { - let flag = flags[index].trim(); - match flag { - "--trust-remote-code" | "--trust-remote-code=true" => { - recipe.trust_remote_code = true; - } - "--no-trust-remote-code" | "--trust-remote-code=false" => { - recipe.trust_remote_code = false; - } - "--preferred-dtype" => { - index += 1; - let value = engine_recipe_flag_value(flags, index, "--preferred-dtype")?; - recipe.preferred_dtype = value.to_owned(); - } - _ if flag.starts_with("--preferred-dtype=") => { - let value = flag - .split_once('=') - .map(|(_, value)| value.trim()) - .unwrap_or_default(); - require_nonempty(value, "--preferred-dtype value")?; - recipe.preferred_dtype = value.to_owned(); - } - "--min-gpu-mem-gb" => { - index += 1; - let value = engine_recipe_flag_value(flags, index, "--min-gpu-mem-gb")?; - recipe.min_gpu_mem_gb = Some(parse_engine_recipe_memory_gb(value)?); - recipe.estimated_memory = - format_estimated_memory(&recipe.canonical_model_id, recipe.min_gpu_mem_gb); - } - _ if flag.starts_with("--min-gpu-mem-gb=") => { - let value = flag - .split_once('=') - .map(|(_, value)| value.trim()) - .unwrap_or_default(); - require_nonempty(value, "--min-gpu-mem-gb value")?; - recipe.min_gpu_mem_gb = Some(parse_engine_recipe_memory_gb(value)?); - recipe.estimated_memory = - format_estimated_memory(&recipe.canonical_model_id, recipe.min_gpu_mem_gb); - } - unsupported => { - bail!("unsupported PyTorch launch recipe flag `{unsupported}`"); - } - } - index += 1; - } - - Ok(recipe) -} - -fn engine_recipe_flag_value<'a>(flags: &'a [String], index: usize, flag: &str) -> Result<&'a str> { - let Some(value) = flags.get(index).map(|value| value.trim()) else { - bail!("{flag} requires a value"); - }; - require_nonempty(value, &format!("{flag} value"))?; - if value.starts_with("--") { - bail!("{flag} requires a value before `{value}`"); - } - Ok(value) -} - -fn parse_engine_recipe_memory_gb(value: &str) -> Result { - value - .parse::() - .with_context(|| format!("invalid --min-gpu-mem-gb value `{value}`")) -} - -fn build_known_recipe( - canonical_model_id: &str, - source: String, - trust_remote_code: bool, - preferred_dtype: &str, - device_policy: DevicePolicy, - min_gpu_mem_gb: Option, - mut warnings: Vec, -) -> ModelRecipe { - if let Some(min_gpu_mem_gb) = min_gpu_mem_gb { - if min_gpu_mem_gb >= 48 - && !warnings - .iter() - .any(|value| value.contains("implicit CPU serving")) - { - warnings.push(format!( - "this model looks large (~{min_gpu_mem_gb} GiB GPU memory recommended); implicit CPU serving is disabled for safety" - )); - if !warnings.iter().any(|value| value.contains("visible GPUs")) { - warnings.push( - "startup will attempt auto device_map placement across visible GPUs when aggregate memory is sufficient" - .to_owned(), - ); - } - } else if min_gpu_mem_gb >= 16 && !warnings.iter().any(|value| value.contains("GPU memory")) - { - warnings.push(format!( - "this model may need about {min_gpu_mem_gb} GiB of GPU memory for comfortable serving" - )); - } - } - - ModelRecipe { - canonical_model_id: canonical_model_id.to_owned(), - task: "chat", - source, - loader: "transformers", - trust_remote_code, - chat_template_mode: "auto", - preferred_dtype: preferred_dtype.to_owned(), - device_policy, - estimated_memory: format_estimated_memory(canonical_model_id, min_gpu_mem_gb), - min_gpu_mem_gb, - warnings, - } -} - -fn infer_model_family(model_ref: &str) -> ModelFamily { - let lower = model_ref.to_ascii_lowercase(); - if lower.contains("qwen") { - ModelFamily::Qwen - } else if lower.contains("glm") { - ModelFamily::Glm - } else if lower.contains("llama") { - ModelFamily::Llama - } else if lower.contains("gpt2") { - ModelFamily::Gpt2 - } else { - ModelFamily::Generic - } -} - -fn preferred_dtype_for_model(model_ref: &str, family: ModelFamily) -> &'static str { - let lower = model_ref.to_ascii_lowercase(); - if lower.contains("fp8") || lower.contains("gptq") || lower.contains("awq") { - "auto" - } else if matches!( - family, - ModelFamily::Qwen | ModelFamily::Glm | ModelFamily::Llama - ) { - "bfloat16" - } else { - "auto" - } -} - -fn infer_parameter_billions(model_ref: &str) -> Option { - let lower = model_ref.to_ascii_lowercase(); - if lower == "zai-org/glm-5-fp8" { - return Some(754.0); - } - - lower - .split(|ch: char| !(ch.is_ascii_alphanumeric() || ch == '.')) - .find_map(|token| { - token - .strip_suffix('b') - .and_then(|value| value.parse::().ok()) - .or_else(|| { - token - .strip_suffix('m') - .and_then(|value| value.parse::().ok()) - .map(|value| value / 1000.0) - }) - }) -} - -fn infer_weight_bytes_per_param(model_ref: &str) -> f32 { - let lower = model_ref.to_ascii_lowercase(); - if lower.contains("int4") || lower.contains("awq") || lower.contains("gptq") { - 0.5 - } else if lower.contains("fp8") || lower.contains("int8") { - 1.0 - } else { - 2.0 - } -} - -fn estimate_gpu_memory_gib(model_ref: &str) -> Option { - let params = infer_parameter_billions(model_ref)?; - let bytes_per_param = infer_weight_bytes_per_param(model_ref); - let overhead = if bytes_per_param <= 1.0 { 1.20 } else { 1.35 }; - Some((params * bytes_per_param * overhead).ceil().max(2.0) as u32) -} - -fn format_estimated_memory(model_ref: &str, min_gpu_mem_gb: Option) -> String { - match (infer_parameter_billions(model_ref), min_gpu_mem_gb) { - (Some(params), Some(min_gpu_mem_gb)) => { - format!( - "~{min_gpu_mem_gb} GiB GPU memory recommended for ~{}B parameters", - trim_float(params) - ) - } - (Some(params), None) => format!( - "~{}B parameters; memory estimate unavailable", - trim_float(params) - ), - (None, Some(min_gpu_mem_gb)) => format!("~{min_gpu_mem_gb} GiB GPU memory recommended"), - (None, None) => "memory estimate unavailable".to_owned(), - } -} - -fn default_device_policy_for_recipe(recipe: &ModelRecipe) -> DevicePolicy { - recipe.device_policy.clone() -} - -fn default_device_policy_for_memory(min_gpu_mem_gb: Option) -> DevicePolicy { - if min_gpu_mem_gb.unwrap_or_default() >= 48 { - DevicePolicy::GpuRequired - } else { - DevicePolicy::GpuPreferred - } -} - -fn trim_float(value: f32) -> String { - if (value.fract() - 0.0).abs() < f32::EPSILON { - format!("{value:.0}") - } else { - format!("{value:.1}") - } -} - -fn parse_device_policy(value: &str) -> Result { - match value { - "gpu_required" => Ok(DevicePolicy::GpuRequired), - "gpu_preferred" => Ok(DevicePolicy::GpuPreferred), - "cpu_only" => Ok(DevicePolicy::CpuOnly), - _ => bail!("unknown device policy: {value}"), - } -} - -/// Parse a `--gpu` CLI value into an optional `GpuSelection` for `LaunchRequest`. -fn parse_gpu_selection_arg(value: Option<&str>) -> Result> { - value - .map(|raw| GpuSelection::parse_cli_value(raw).map_err(anyhow::Error::msg)) - .transpose() -} - -/// Parse a `--gpu` CLI value into explicit device ordinals (empty for `auto`). -fn parse_gpu_indices_arg(value: Option<&str>) -> Result> { - Ok(rocm_engine_protocol::launch_gpu_indices( - parse_gpu_selection_arg(value)?.as_ref(), - )) -} - -/// Re-exec `--gpu` flag for the detached serve-http worker, derived from the -/// launch request's selection. Empty (auto) selections emit no flag. -fn gpu_indices_arg(selection: Option<&GpuSelection>) -> Vec { - match rocm_engine_protocol::gpu_indices_to_csv(&rocm_engine_protocol::launch_gpu_indices( - selection, - )) { - Some(csv) => vec!["--gpu".to_owned(), csv], - None => Vec::new(), - } -} - -fn normalize_pytorch_device_policy(policy: DevicePolicy) -> Result { - match policy { - DevicePolicy::GpuRequired => Ok(DevicePolicy::GpuRequired), - DevicePolicy::GpuPreferred => Ok(DevicePolicy::GpuRequired), - DevicePolicy::CpuOnly => { - bail!("PyTorch adapter requires ROCm GPU execution; no CPU fallback is used") - } - } -} - -const fn device_policy_name(policy: &DevicePolicy) -> &'static str { - match policy { - DevicePolicy::GpuRequired => "gpu_required", - DevicePolicy::GpuPreferred => "gpu_preferred", - DevicePolicy::CpuOnly => "cpu_only", - } -} - -fn slugify(value: &str) -> String { - value - .chars() - .map(|ch| match ch { - 'a'..='z' | 'A'..='Z' | '0'..='9' => ch.to_ascii_lowercase(), - _ => '-', - }) - .collect() -} - -fn managed_env_id(runtime_id: &str, python_version: Option<&str>) -> String { - let python = python_version.unwrap_or(default_python_version()); - format!( - "{}-{}-{}", - rocm_core::runtime_os_name(), - slugify(runtime_id), - slugify(python) - ) -} - -const fn default_python_version() -> &'static str { - "3.12" -} - -fn parse_therock_runtime_request(runtime_id: &str) -> TheRockRuntimeRequest { - let normalized = runtime_id.trim().to_ascii_lowercase(); - let channel = therock_channel_from_str(&normalized); - - let family_override = KNOWN_THEROCK_FAMILIES - .iter() - .find(|family| normalized.contains(&family.to_ascii_lowercase())) - .map(|family| (*family).to_owned()) - .or_else(|| { - extract_first_gfx_token(&normalized) - .and_then(|target| normalize_therock_family(&target)) - }); - - TheRockRuntimeRequest { - channel, - family_override, - } -} - -fn therock_channel_from_str(value: &str) -> TheRockChannel { - if value.trim().to_ascii_lowercase().contains("nightly") { - TheRockChannel::Nightly - } else { - TheRockChannel::Release - } -} - -fn resolve_therock_torch_resolution( - paths: &AppPaths, - runtime_id: &str, -) -> Result> { - let runtime_request = parse_therock_runtime_request(runtime_id); - - if let Some(manifest) = load_runtime_registry_manifest(paths, runtime_id)? { - if manifest.format != "wheel" { - return Ok(None); - } - let family = normalize_therock_family(&manifest.family) - .or_else(|| parse_therock_runtime_request(&manifest.runtime_id).family_override) - .with_context(|| { - format!( - "managed runtime `{}` did not report a supported TheRock family", - manifest.runtime_key - ) - })?; - let python = manifest.python_executable.as_deref().with_context(|| { - format!( - "managed runtime `{}` did not record a Python executable", - manifest.runtime_key - ) - })?; - let packages = pinned_torch_package_specs_from_runtime(python).with_context(|| { - format!( - "managed runtime `{}` did not expose exact torch package versions", - manifest.runtime_key - ) - })?; - let channel = therock_channel_from_str(&manifest.channel); - return Ok(Some(TheRockTorchResolution { - channel, - family: family.clone(), - index_url: manifest - .index_url - .or(manifest.selected_artifact_url) - .unwrap_or_else(|| therock_index_url(&family)), - packages, - source: format!("managed_runtime_manifest:{}", manifest.runtime_key), - })); - } - - if let Some(family) = runtime_request.family_override { - return Ok(Some(TheRockTorchResolution { - channel: runtime_request.channel, - family: family.clone(), - index_url: therock_index_url(&family), - packages: THEROCK_TORCH_PACKAGES - .iter() - .map(|value| (*value).to_owned()) - .collect(), - source: "runtime_id".to_owned(), - })); - } - - if let Ok(value) = std::env::var("ROCM_CLI_THEROCK_FAMILY") - && let Some(family) = normalize_therock_family(&value) - { - return Ok(Some(TheRockTorchResolution { - channel: runtime_request.channel, - family: family.clone(), - index_url: therock_index_url(&family), - packages: THEROCK_TORCH_PACKAGES - .iter() - .map(|item| (*item).to_owned()) - .collect(), - source: "env".to_owned(), - })); - } - - if let Some(family) = detect_host_therock_family() { - return Ok(Some(TheRockTorchResolution { - channel: runtime_request.channel, - family: family.clone(), - index_url: therock_index_url(&family), - packages: THEROCK_TORCH_PACKAGES - .iter() - .map(|item| (*item).to_owned()) - .collect(), - source: "host".to_owned(), - })); - } - - Ok(None) -} - -fn runtime_registry_dir(paths: &AppPaths) -> PathBuf { - paths.data_dir.join("runtimes").join("registry") -} - -fn load_runtime_registry_manifest( - paths: &AppPaths, - selector: &str, -) -> Result> { - let registry_dir = runtime_registry_dir(paths); - if !registry_dir.is_dir() { - return Ok(None); - } - - let exact_path = registry_dir.join(format!("{selector}.json")); - if exact_path.is_file() { - return load_runtime_registry_manifest_path(&exact_path).map(Some); - } - - let mut matches = Vec::new(); - for entry in fs::read_dir(®istry_dir) - .with_context(|| format!("failed to read {}", registry_dir.display()))? - { - let entry = entry?; - let path = entry.path(); - if path.extension().and_then(|value| value.to_str()) != Some("json") { - continue; - } - let manifest = load_runtime_registry_manifest_path(&path)?; - if manifest.runtime_id.eq_ignore_ascii_case(selector) { - matches.push(manifest); - } - } - if matches.len() == 1 { - return Ok(matches.pop()); - } - Ok(None) -} - -fn load_runtime_registry_manifest_path(path: &Path) -> Result { - serde_json::from_slice( - &fs::read(path).with_context(|| format!("failed to read {}", path.display()))?, - ) - .with_context(|| format!("failed to parse runtime manifest {}", path.display())) -} - -fn pinned_torch_package_specs_from_runtime(python_executable: &str) -> Result> { - if let Ok(specs) = pinned_torch_package_specs_from_runtime_metadata(python_executable) { - return Ok(specs); - } - - let code = format!( - "import importlib.metadata as md, json; names = {THEROCK_TORCH_PACKAGES:?}; print(json.dumps({{name: md.version(name) for name in names}}, sort_keys=True))" - ); - let output = capture_command( - python_executable, - ["-c", code.as_str()], - "inspect managed runtime torch package versions", - )?; - parse_torch_package_version_specs(&output) -} - -fn pinned_torch_package_specs_from_runtime_metadata( - python_executable: &str, -) -> Result> { - let packages = installed_package_specs_from_runtime_metadata(python_executable)?; - let versions = packages - .iter() - .filter_map(|spec| spec.split_once("==")) - .map(|(name, version)| (name.to_ascii_lowercase(), version.to_owned())) - .collect::>(); - let mut specs = Vec::new(); - for name in THEROCK_TORCH_PACKAGES { - let version = versions - .get(*name) - .filter(|value| !value.trim().is_empty()) - .with_context(|| format!("runtime metadata is missing package `{name}`"))?; - specs.push(format!("{name}=={version}")); - } - Ok(specs) -} - -fn installed_package_specs_from_runtime_metadata(python_executable: &str) -> Result> { - let mut errors = Vec::new(); - for site_packages in site_packages_candidates_for_python(python_executable) { - match installed_package_specs_from_dist_info_dir(&site_packages) { - Ok(specs) => return Ok(specs), - Err(error) => errors.push(format!("{}: {error}", site_packages.display())), - } - } - bail!( - "could not find installed package metadata ({})", - errors.join("; ") - ) -} - -fn site_packages_candidates_for_python(python_executable: &str) -> Vec { - let python_path = PathBuf::from(normalize_runtime_path_text_for_host(python_executable)); - let Some(parent) = python_path.parent() else { - return Vec::new(); - }; - let mut candidates = Vec::new(); - if runtime_is_windows() { - if parent - .file_name() - .and_then(|value| value.to_str()) - .is_some_and(|value| value.eq_ignore_ascii_case("scripts")) - && let Some(env_root) = parent.parent() - { - candidates.push(env_root.join("Lib").join("site-packages")); - } - candidates.push(parent.join("Lib").join("site-packages")); - } else if parent - .file_name() - .and_then(|value| value.to_str()) - .is_some_and(|value| value == "bin") - && let Some(env_root) = parent.parent() - { - let lib_dir = env_root.join("lib"); - if let Ok(entries) = fs::read_dir(&lib_dir) { - for entry in entries.flatten() { - let path = entry.path(); - let name = entry.file_name(); - let name = name.to_string_lossy(); - if name.starts_with("python") { - candidates.push(path.join("site-packages")); - } - } - } - candidates.push( - lib_dir - .join(format!("python{}", default_python_version())) - .join("site-packages"), - ); - } - candidates.sort(); - candidates.dedup(); - candidates -} - -fn installed_package_specs_from_dist_info_dir(site_packages: &Path) -> Result> { - let mut versions = BTreeMap::new(); - let entries = fs::read_dir(site_packages) - .with_context(|| format!("failed to read {}", site_packages.display()))?; - for entry in entries { - let entry = entry?; - let path = entry.path(); - let is_dist_info = path - .file_name() - .and_then(|value| value.to_str()) - .is_some_and(|value| value.to_ascii_lowercase().ends_with(".dist-info")); - if !path.is_dir() || !is_dist_info { - continue; - } - let metadata_path = path.join("METADATA"); - if !metadata_path.is_file() { - continue; - } - if let Some((name, version)) = dist_info_name_version(&metadata_path)? - && !name.trim().is_empty() - && !version.trim().is_empty() - { - versions.insert(name, version); - } - } - if versions.is_empty() { - bail!("no installed package metadata found"); - } - Ok(versions - .into_iter() - .map(|(name, version)| format!("{name}=={version}")) - .collect()) -} - -#[cfg(test)] -fn pinned_torch_package_specs_from_dist_info_dir(site_packages: &Path) -> Result> { - let packages = installed_package_specs_from_dist_info_dir(site_packages)?; - let versions = packages - .iter() - .filter_map(|spec| spec.split_once("==")) - .map(|(name, version)| (name.to_ascii_lowercase(), version.to_owned())) - .collect::>(); - let mut specs = Vec::new(); - for name in THEROCK_TORCH_PACKAGES { - let version = versions - .get(*name) - .filter(|value| !value.trim().is_empty()) - .with_context(|| format!("runtime metadata is missing package `{name}`"))?; - specs.push(format!("{name}=={version}")); - } - Ok(specs) -} - -fn dist_info_name_version(metadata_path: &Path) -> Result> { - let metadata = fs::read_to_string(metadata_path) - .with_context(|| format!("failed to read {}", metadata_path.display()))?; - let mut name = None; - let mut version = None; - for line in metadata.lines() { - if let Some(value) = line.strip_prefix("Name:") { - name = Some(value.trim().to_ascii_lowercase()); - } else if let Some(value) = line.strip_prefix("Version:") { - version = Some(value.trim().to_owned()); - } - if name.is_some() && version.is_some() { - break; - } - } - Ok(name.zip(version)) -} - -fn parse_torch_package_version_specs(output: &str) -> Result> { - let versions: BTreeMap = - serde_json::from_str(output.trim()).context("failed to parse torch package versions")?; - let mut specs = Vec::new(); - for name in THEROCK_TORCH_PACKAGES { - let version = versions - .get(*name) - .filter(|value| !value.trim().is_empty()) - .with_context(|| format!("runtime Python is missing package `{name}`"))?; - specs.push(format!("{name}=={version}")); - } - Ok(specs) -} - -fn install_therock_torch_packages( - uv: &Path, - python_executable: &Path, - resolution: &TheRockTorchResolution, -) -> Result<()> { - let mut args = uv_pip_install_base(python_executable); - args.push("--index-url".to_owned()); - args.push(resolution.index_url.clone()); - if matches!(resolution.channel, TheRockChannel::Nightly) { - args.extend(["--prerelease".to_owned(), "allow".to_owned()]); - } - args.extend(resolution.packages.iter().cloned()); - run_uv_progress_command( - uv, - args.iter().map(String::as_str), - "install TheRock torch packages into managed pytorch env", - ) -} - -fn therock_index_url(family: &str) -> String { - format!("{THEROCK_SIMPLE_INDEX_BASE}/{family}/") -} - -fn simple_hash(value: &str) -> String { - let mut hasher = DefaultHasher::new(); - value.hash(&mut hasher); - format!("{:016x}", hasher.finish()) -} - -fn find_installed_package(packages: &[String], name: &str) -> Option { - packages.iter().find_map(|entry| { - entry - .strip_prefix(&format!("{name}==")) - .map(ToOwned::to_owned) - }) -} - -fn manifest_has_torch(manifest: &EngineEnvManifest) -> bool { - find_installed_package(&manifest.installed_packages, "torch").is_some() -} - -fn probe_torch_runtime(python_executable: &str) -> Result { - let output = capture_command( - python_executable, - ["-c", PYTORCH_PROBE_SCRIPT], - "probe managed pytorch runtime", - )?; - parse_torch_runtime_probe(&output) -} - -fn parse_torch_runtime_probe(output: &str) -> Result { - serde_json::from_str(output.trim()).context("failed to parse managed pytorch runtime probe") -} - -const PYTORCH_PROBE_SCRIPT: &str = r#" -import ast -import importlib.util -import json - -def inspect_rocm_sdk(): - out = { - "import_ok": False, - "version": None, - "site_packages": None, - "default_target_family": None, - "available_target_families": [], - "resolved_target_family": None, - "error": None, - } - try: - import sysconfig - import rocm_sdk - from rocm_sdk import _dist_info as di - - out["import_ok"] = True - out["version"] = getattr(rocm_sdk, "__version__", None) - out["site_packages"] = sysconfig.get_paths().get("purelib") - out["default_target_family"] = getattr(di, "DEFAULT_TARGET_FAMILY", None) - out["available_target_families"] = list(getattr(di, "AVAILABLE_TARGET_FAMILIES", [])) - try: - out["resolved_target_family"] = di.determine_target_family() - except Exception as exc: - out["error"] = type(exc).__name__ + ": " + str(exc) - except Exception as exc: - out["error"] = type(exc).__name__ + ": " + str(exc) - return out - -def inspect_torch_rocm_init(): - out = { - "present": False, - "path": None, - "check_version": None, - "preload_shortnames": [], - "error": None, - } - try: - spec = importlib.util.find_spec("torch") - locations = list(spec.submodule_search_locations or []) if spec else [] - if not locations: - return out - from pathlib import Path - path = Path(locations[0]) / "_rocm_init.py" - if not path.exists(): - return out - out["present"] = True - out["path"] = str(path) - tree = ast.parse(path.read_text(encoding="utf-8")) - for node in ast.walk(tree): - if not ( - isinstance(node, ast.Call) - and isinstance(node.func, ast.Attribute) - and node.func.attr == "initialize_process" - ): - continue - for kw in node.keywords: - if kw.arg == "check_version": - out["check_version"] = ast.literal_eval(kw.value) - elif kw.arg == "preload_shortnames": - out["preload_shortnames"] = list(ast.literal_eval(kw.value)) - break - except Exception as exc: - out["error"] = type(exc).__name__ + ": " + str(exc) - return out - -rocm_sdk_probe = inspect_rocm_sdk() -rocm_init_probe = inspect_torch_rocm_init() - -try: - import torch - devices = [] - try: - count = torch.cuda.device_count() - devices = [torch.cuda.get_device_name(i) for i in range(count)] - except Exception: - count = 0 - print(json.dumps({ - "import_ok": True, - "torch_version": getattr(torch, "__version__", None), - "cuda_available": bool(torch.cuda.is_available()), - "device_count": int(count), - "devices": devices, - "rocm_sdk": rocm_sdk_probe, - "torch_rocm_init": rocm_init_probe, - })) -except Exception as exc: - print(json.dumps({ - "import_ok": False, - "cuda_available": False, - "device_count": 0, - "devices": [], - "error": str(exc), - "rocm_sdk": rocm_sdk_probe, - "torch_rocm_init": rocm_init_probe, - })) -"#; - -fn run_uv_command<'a, I>(uv: &Path, args: I, context_text: &str) -> Result<()> -where - I: IntoIterator, -{ - let args: Vec<_> = args.into_iter().collect(); - let output = Command::new(uv) - .args(&args) - .envs(uv_command_env()) - .stdin(Stdio::null()) - .stdout(Stdio::null()) - .stderr(Stdio::piped()) - .output() - .with_context(|| format!("failed to launch uv for {context_text}"))?; - if output.status.success() { - return Ok(()); - } - let stderr = String::from_utf8_lossy(&output.stderr).trim().to_owned(); - bail!("{context_text}: uv exited with {}: {stderr}", output.status) -} - -fn run_uv_progress_command<'a, I>(uv: &Path, args: I, context_text: &str) -> Result<()> -where - I: IntoIterator, -{ - let args: Vec<_> = args.into_iter().collect(); - let status = Command::new(uv) - .args(&args) - .envs(uv_command_env()) - .stdin(Stdio::null()) - .stdout(Stdio::inherit()) - .stderr(Stdio::inherit()) - .status() - .with_context(|| format!("failed to launch uv for {context_text}"))?; - if status.success() { - return Ok(()); - } - bail!("{context_text}: uv exited with {status}") -} - -fn capture_uv_command<'a, I>(uv: &Path, args: I, context_text: &str) -> Result -where - I: IntoIterator, -{ - let args: Vec<_> = args.into_iter().collect(); - let output = Command::new(uv) - .args(&args) - .envs(uv_command_env()) - .stdin(Stdio::null()) - .stdout(Stdio::piped()) - .stderr(Stdio::null()) - .output() - .with_context(|| format!("failed to launch uv for {context_text}"))?; - if output.status.success() { - return String::from_utf8(output.stdout).context("uv output was not valid UTF-8"); - } - bail!("{context_text}: uv exited with {}", output.status) -} - -#[allow(dead_code)] -fn command_succeeds<'a, I>(program: &str, args: I) -> Result -where - I: IntoIterator, -{ - let args = args.into_iter().map(ToOwned::to_owned).collect::>(); - let status = Command::new(command_program(program)) - .args(&args) - .stdin(Stdio::null()) - .stdout(Stdio::null()) - .stderr(Stdio::null()) - .status() - .with_context(|| format!("failed to start {}", args.join(" ")))?; - Ok(status.success()) -} - -fn ensure_service_env(runtime_id: Option<&str>, env_id: Option<&str>) -> Result { - if let Some(env_id) = env_id { - let manifest = load_manifest_by_env_id(env_id)?; - if manifest_has_torch(&manifest) { - return Ok(manifest); - } - return create_or_update_env_manifest(&InstallRequest { - runtime_id: manifest.runtime_id.clone(), - python_version: manifest.requested_python_version, - env_root: None, - reinstall: true, - }); - } - - if let Some(runtime_id) = runtime_id { - return create_or_update_env_manifest(&InstallRequest { - runtime_id: runtime_id.to_owned(), - python_version: None, - env_root: None, - reinstall: false, - }); - } - - if let Some(manifest) = latest_runnable_env_manifest()? { - return Ok(manifest); - } - - create_or_update_env_manifest(&InstallRequest { - runtime_id: DEFAULT_RUNTIME_ID.to_owned(), - python_version: None, - env_root: None, - reinstall: false, - }) -} - -fn latest_runnable_env_manifest() -> Result> { - let paths = AppPaths::discover()?; - let manifests_dir = paths.engine_manifests_dir(ENGINE_NAME); - if !manifests_dir.is_dir() { - return Ok(None); - } - - let mut manifests = Vec::new(); - for entry in fs::read_dir(&manifests_dir) - .with_context(|| format!("failed to read {}", manifests_dir.display()))? - { - let entry = entry?; - let path = entry.path(); - if path.extension().and_then(|value| value.to_str()) != Some("json") { - continue; - } - let manifest = load_manifest(&path)?; - if manifest.env_path.is_dir() && manifest_has_torch(&manifest) { - manifests.push(manifest); - } - } - manifests.sort_by(|left, right| left.env_id.cmp(&right.env_id)); - Ok(manifests.pop()) -} - -fn load_manifest_by_env_id(env_id: &str) -> Result { - let paths = AppPaths::discover()?; - let path = paths - .engine_manifests_dir(ENGINE_NAME) - .join(format!("{env_id}.json")); - load_manifest(&path) -} - -fn materialize_python_worker() -> Result { - let paths = AppPaths::discover()?; - let worker_dir = paths.engine_dir(ENGINE_NAME).join("worker"); - fs::create_dir_all(&worker_dir) - .with_context(|| format!("failed to create {}", worker_dir.display()))?; - - let worker_path = worker_dir.join("python_worker.py"); - let needs_write = match fs::read_to_string(&worker_path) { - Ok(current) => current != PYTHON_WORKER_SOURCE, - Err(_) => true, - }; - if needs_write { - fs::write(&worker_path, PYTHON_WORKER_SOURCE) - .with_context(|| format!("failed to write {}", worker_path.display()))?; - } - Ok(worker_path) -} - -fn optional_arg(flag: &str, value: Option<&str>) -> Vec { - match value { - Some(value) => vec![flag.to_owned(), value.to_owned()], - None => Vec::new(), - } -} - -fn engine_recipe_json_arg(engine_recipe: Option<&EngineRecipeHint>) -> Result> { - match engine_recipe { - Some(engine_recipe) => Ok(vec![ - "--engine-recipe-json".to_owned(), - serde_json::to_string(engine_recipe).context("failed to encode engine recipe hint")?, - ]), - None => Ok(Vec::new()), - } -} - -fn engine_recipe_launch_args(engine_recipe: Option<&EngineRecipeHint>) -> Vec { - engine_recipe - .map(|hint| hint.required_flags.clone()) - .unwrap_or_default() -} - -fn optional_arg_owned(flag: &str, value: Option) -> Vec { - match value { - Some(value) => vec![flag.to_owned(), value], - None => Vec::new(), - } -} - -fn flag_arg(flag: &str, enabled: bool) -> Vec { - if enabled { - vec![flag.to_owned()] - } else { - Vec::new() - } -} - -#[allow(dead_code)] -fn run_command<'a, I>(program: &str, args: I, context_label: &str) -> Result<()> -where - I: IntoIterator, -{ - let args = args.into_iter().map(ToOwned::to_owned).collect::>(); - let output = capture_command_files(program, &args, context_label)?; - if output.status.success() { - Ok(()) - } else { - bail!( - "{} failed (status {}): {}", - context_label, - output.status, - String::from_utf8_lossy(&output.stderr) - ); - } -} - -#[allow(dead_code)] -fn run_progress_command<'a, I>(program: &str, args: I, context_label: &str) -> Result<()> -where - I: IntoIterator, -{ - let args = args.into_iter().map(ToOwned::to_owned).collect::>(); - if interactive_terminal() { - let status = Command::new(command_program(program)) - .args(&args) - .stdin(Stdio::null()) - .stdout(Stdio::inherit()) - .stderr(Stdio::inherit()) - .status() - .with_context(|| format!("failed to start {context_label}"))?; - if status.success() { - return Ok(()); - } - bail!("{context_label} failed (status {status})"); - } - if engine_progress_stderr_enabled() { - return run_progress_command_forwarded(program, &args, context_label); - } - run_command(program, args.iter().map(String::as_str), context_label) -} - -#[allow(dead_code)] -fn engine_progress_stderr_enabled() -> bool { - env_value_truthy(std::env::var("ROCM_ENGINE_PROGRESS_STDERR").ok().as_deref()) -} - -#[allow(dead_code)] -fn env_value_truthy(value: Option<&str>) -> bool { - matches!( - value.map(str::trim).map(str::to_ascii_lowercase).as_deref(), - Some("1" | "true" | "yes" | "on") - ) -} - -#[allow(dead_code)] -fn run_progress_command_forwarded( - program: &str, - args: &[String], - context_label: &str, -) -> Result<()> { - let mut child = Command::new(command_program(program)) - .args(args) - .stdin(Stdio::null()) - .stdout(Stdio::piped()) - .stderr(Stdio::piped()) - .spawn() - .with_context(|| format!("failed to start {context_label}"))?; - let stdout = child - .stdout - .take() - .map(forward_child_output_to_stderr) - .context("progress child did not expose stdout")?; - let stderr = child - .stderr - .take() - .map(forward_child_output_to_stderr) - .context("progress child did not expose stderr")?; - let status = child - .wait() - .with_context(|| format!("failed waiting for {context_label}"))?; - let stdout = stdout.join().unwrap_or_default(); - let stderr = stderr.join().unwrap_or_default(); - if status.success() { - return Ok(()); - } - let mut combined = Vec::with_capacity(stdout.len() + stderr.len()); - combined.extend_from_slice(&stdout); - combined.extend_from_slice(&stderr); - bail!( - "{} failed (status {}): {}", - context_label, - status, - String::from_utf8_lossy(&combined) - ); -} - -#[allow(dead_code)] -fn forward_child_output_to_stderr(mut reader: R) -> thread::JoinHandle> -where - R: Read + Send + 'static, -{ - thread::spawn(move || { - let mut collected = Vec::new(); - let mut buffer = [0_u8; 8192]; - let mut stderr = io::stderr(); - loop { - match reader.read(&mut buffer) { - Ok(0) => break, - Ok(read) => { - let chunk = &buffer[..read]; - let _ = stderr.write_all(chunk); - let _ = stderr.flush(); - collected.extend_from_slice(chunk); - } - Err(error) => { - let message = format!("\nfailed to read progress output: {error}\n"); - let _ = stderr.write_all(message.as_bytes()); - collected.extend_from_slice(message.as_bytes()); - break; - } - } - } - collected - }) -} - -fn capture_command<'a, I>(program: &str, args: I, context_label: &str) -> Result -where - I: IntoIterator, -{ - let args = args.into_iter().map(ToOwned::to_owned).collect::>(); - let output = capture_command_files(program, &args, context_label)?; - if !output.status.success() { - bail!( - "{} failed (status {}): {}", - context_label, - output.status, - String::from_utf8_lossy(&output.stderr) - ); - } - String::from_utf8(output.stdout).context("command output was not valid utf-8") -} - -fn command_program(program: &str) -> String { - normalize_runtime_path_text_for_host(program) -} - -fn command_path(path: &Path) -> String { - normalize_runtime_path_text_for_host(&path.display().to_string()) -} - -struct CapturedCommand { - status: ExitStatus, - stdout: Vec, - stderr: Vec, -} - -fn capture_command_files( - program: &str, - args: &[String], - context_label: &str, -) -> Result { - let temp_dir = std::env::temp_dir(); - let stem = format!( - "rocm-pytorch-command-{}-{}", - std::process::id(), - unix_time_millis() - ); - let stdout_path = temp_dir.join(format!("{stem}.out")); - let stderr_path = temp_dir.join(format!("{stem}.err")); - let stdout_file = fs::File::create(&stdout_path) - .with_context(|| format!("failed to create {}", stdout_path.display()))?; - let stderr_file = fs::File::create(&stderr_path) - .with_context(|| format!("failed to create {}", stderr_path.display()))?; - let status_result = Command::new(command_program(program)) - .args(args) - .stdin(Stdio::null()) - .stdout(Stdio::from(stdout_file)) - .stderr(Stdio::from(stderr_file)) - .status() - .with_context(|| format!("failed to start {context_label}")); - let stdout = fs::read(&stdout_path).unwrap_or_default(); - let stderr = fs::read(&stderr_path).unwrap_or_default(); - let _ = fs::remove_file(&stdout_path); - let _ = fs::remove_file(&stderr_path); - Ok(CapturedCommand { - status: status_result?, - stdout, - stderr, - }) -} - -fn read_request() -> Result { - let mut buffer = String::new(); - std::io::stdin() - .read_to_string(&mut buffer) - .context("failed to read stdin for engine request")?; - serde_json::from_str(&buffer).context("failed to parse engine request envelope") -} - -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!(&mut handle)?; - Ok(()) -} - -impl From for DevicePolicy { - fn from(value: DevicePolicyArg) -> Self { - match value { - DevicePolicyArg::GpuRequired => Self::GpuRequired, - DevicePolicyArg::GpuPreferred => Self::GpuPreferred, - DevicePolicyArg::CpuOnly => Self::CpuOnly, - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - - fn test_engine_recipe(engine: &str, contract_version: &str) -> EngineRecipeHint { - EngineRecipeHint { - contract_version: contract_version.to_owned(), - engine: engine.to_owned(), - required_flags: vec!["--trust-remote-code=false".to_owned()], - parser_settings: BTreeMap::default(), - preferred_endpoint: None, - unsupported_combinations: Vec::new(), - notes: vec!["test recipe".to_owned()], - } - } - - #[test] - fn normalize_therock_family_maps_gfx1103_to_gfx110x_all() { - assert_eq!( - normalize_therock_family("gfx1103"), - Some("gfx110X-all".to_owned()) - ); - } - - #[test] - fn normalize_therock_family_maps_gfx1101_to_gfx110x_all() { - assert_eq!( - normalize_therock_family("gfx1101"), - Some("gfx110X-all".to_owned()) - ); - } - - #[test] - fn tail_lines_from_text_returns_requested_suffix() { - assert_eq!( - tail_lines_from_text("one\ntwo\nthree\nfour\n", 2), - vec!["three".to_owned(), "four".to_owned()] - ); - assert!(tail_lines_from_text("one\ntwo\n", 0).is_empty()); - } - - #[test] - fn healthcheck_response_maps_ready_state() { - let state = json!({ - "status": "ready", - "device": "cuda", - "queue_depth": 3, - "tokens_per_sec": 42.5 - }); - let response = build_healthcheck_response( - Some(&state), - Some(UNIX_EPOCH + Duration::from_secs(5)), - None, - None, - None, - UNIX_EPOCH + Duration::from_secs(15), - ); - - assert_eq!(response.status, "ready"); - assert!(response.model_loaded); - assert_eq!(response.device, "cuda"); - assert_eq!(response.uptime_sec, 10); - assert_eq!(response.queue_depth, 3); - assert_eq!(response.tokens_per_sec, Some(42.5)); - } - - #[test] - fn healthcheck_response_marks_ready_state_unreachable_when_probe_fails() { - let state = json!({ - "status": "ready", - "device": "cpu" - }); - let response = build_healthcheck_response( - Some(&state), - None, - None, - None, - Some("connection refused".to_owned()), - UNIX_EPOCH, - ); - - assert_eq!(response.status, "unreachable"); - assert!(!response.model_loaded); - assert_eq!(response.device, "cpu"); - assert_eq!(response.last_error.as_deref(), Some("connection refused")); - } - - #[test] - fn parse_http_endpoint_supports_v1_urls() { - assert_eq!( - parse_http_endpoint("http://127.0.0.1:11435/v1"), - Some(("127.0.0.1".to_owned(), 11435)) - ); - assert_eq!( - parse_http_endpoint("http://[::1]:11435/v1"), - Some(("::1".to_owned(), 11435)) - ); - } - - #[test] - fn endpoint_url_fallback_brackets_ipv6_loopback() { - let state = json!({ - "host": "::1", - "port": 11435 - }); - - assert_eq!( - endpoint_url_from_state(&state), - Some("http://[::1]:11435/v1".to_owned()) - ); - } - - #[test] - fn serve_http_cli_defaults_to_gpu_required_without_cpu_fallback() { - let cli = Cli::try_parse_from([ - "rocm-engine-pytorch", - "serve-http", - "svc", - "qwen", - "--env-id", - "env-1", - "--runtime-id", - "therock-release:gfx120X-all", - "--state-path", - "state.json", - ]) - .expect("serve-http args should parse"); - - match cli.command { - CommandKind::ServeHttp { device_policy, .. } => { - assert_eq!(device_policy, "gpu_required"); - } - _ => panic!("expected serve-http command"), - } - } - - #[test] - fn launch_cli_accepts_runtime_selection_args() { - let cli = Cli::try_parse_from([ - "rocm-engine-pytorch", - "launch", - "svc", - "qwen", - "--runtime-id", - "therock-release:gfx120X-all", - "--env-id", - "pytorch-env-1", - ]) - .expect("launch should accept protocol runtime args"); - - match cli.command { - CommandKind::Launch { - runtime_id, env_id, .. - } => { - assert_eq!(runtime_id.as_deref(), Some("therock-release:gfx120X-all")); - assert_eq!(env_id.as_deref(), Some("pytorch-env-1")); - } - _ => panic!("expected launch command"), - } - } - - #[test] - fn python_worker_disallows_implicit_cpu_fallback_for_gpu_policies() { - let worker = include_str!("python_worker.py"); - - assert!(!worker.contains("cpu_fallback")); - assert!(worker.contains("policy in {\"gpu_required\", \"gpu_preferred\"}")); - assert!(worker.contains("PyTorch CPU serving is not offered by rocm-cli")); - assert!(worker.contains("no CPU fallback is used")); - } - - #[test] - fn python_worker_bridges_qwen_xml_tool_calls_to_openai_tool_calls() { - let worker = include_str!("python_worker.py"); - - assert!(worker.contains("kwargs[\"tools\"] = tools")); - assert!(worker.contains("TOOL_CALL_PATTERN")); - assert!(worker.contains("\"tool_calls\"")); - assert!(worker.contains("\"finish_reason\": finish_reason")); - } - - #[test] - fn pytorch_capabilities_advertise_tool_calling() { - assert!(capabilities().tool_calling); - } - - #[test] - fn engine_progress_forwarding_env_uses_explicit_truthy_values() { - for value in [Some("1"), Some("true"), Some("YES"), Some(" on ")] { - assert!(env_value_truthy(value)); - } - for value in [None, Some(""), Some("0"), Some("false"), Some("please")] { - assert!(!env_value_truthy(value)); - } - } - - #[test] - fn cpu_policy_is_rejected_without_fallback() { - let error = normalize_pytorch_device_policy(DevicePolicy::CpuOnly) - .expect_err("cpu policy should not be accepted by rocm-cli"); - assert!(error.to_string().contains("no CPU fallback is used")); - } - - #[test] - fn stdio_protocol_routes_all_methods_without_side_effects() { - let service_id = format!( - "missing-protocol-{}-{}", - std::process::id(), - rocm_core::unix_time_millis() - ); - let ok_cases = [ - (EngineMethod::Detect, json!({})), - (EngineMethod::Capabilities, json!({})), - ( - EngineMethod::ResolveModel, - json!({ - "model_ref": "qwen", - "device_policy": "gpu_required" - }), - ), - ( - EngineMethod::Healthcheck, - json!({ - "service_id": service_id.as_str() - }), - ), - ( - EngineMethod::Endpoint, - json!({ - "service_id": service_id.as_str() - }), - ), - ( - EngineMethod::Logs, - json!({ - "service_id": service_id.as_str(), - "tail_lines": 4 - }), - ), - ( - EngineMethod::Stop, - json!({ - "service_id": service_id.as_str(), - "force": false - }), - ), - ]; - - for (method, payload) in ok_cases { - let response = handle_envelope(EngineRequestEnvelope { method, payload }); - assert!( - response.ok, - "expected protocol method to return a typed success envelope: {:?}", - response.error - ); - } - - for method in [EngineMethod::Install, EngineMethod::Launch] { - let response = handle_envelope(EngineRequestEnvelope { - method, - payload: json!({}), - }); - assert!(!response.ok); - assert_eq!( - response.error.as_ref().map(|error| error.code.as_str()), - Some("invalid_payload") - ); - } - } - - #[test] - fn endpoint_response_uses_service_state_endpoint() -> Result<()> { - let root = std::env::temp_dir().join(format!( - "rocm-pytorch-endpoint-{}", - rocm_core::unix_time_millis() - )); - let _ = fs::remove_dir_all(&root); - fs::create_dir_all(&root)?; - let state_path = root.join("state.json"); - fs::write( - &state_path, - serde_json::to_vec(&json!({ - "status": "ready", - "endpoint_url": "http://127.0.0.1:32123/v1" - }))?, - )?; - let files = ServiceFiles { - record_path: root.join("service.json"), - record: None, - record_matches_engine: false, - state_path, - log_path: root.join("service.log"), - }; - - let response = endpoint_response_from_files(&files)?; - fs::remove_dir_all(root).ok(); - - assert_eq!(response.endpoint_url, "http://127.0.0.1:32123/v1"); - assert_eq!(response.api_style, "openai"); - assert!(response.supported_routes.contains(&"/healthz".to_owned())); - assert!( - response - .supported_routes - .contains(&"/v1/chat/completions".to_owned()) - ); - Ok(()) - } - - #[test] - fn manifest_python_path_falls_back_to_existing_venv_interpreter() -> Result<()> { - let root = std::env::temp_dir().join(format!( - "rocm-pytorch-python-fallback-{}", - rocm_core::unix_time_millis() - )); - let _ = fs::remove_dir_all(&root); - let env_path = root.join("env"); - let bin_dir = if cfg!(windows) { - env_path.join("Scripts") - } else { - env_path.join("bin") - }; - fs::create_dir_all(&bin_dir)?; - let missing_recorded = if cfg!(windows) { - bin_dir.join("python.exe") - } else { - bin_dir.join("python") - }; - let existing = if cfg!(windows) { - bin_dir.join("python3.exe") - } else { - bin_dir.join("python3") - }; - fs::write(&existing, "python")?; - - let manifest = EngineEnvManifest { - env_id: "env".to_owned(), - runtime_id: "runtime".to_owned(), - requested_python_version: Some("3.12".to_owned()), - python_launcher: "python3.12".to_owned(), - python_executable: missing_recorded.display().to_string(), - env_path, - manifest_path: root.join("manifest.json"), - lock_path: root.join("lock.txt"), - installed_packages: vec!["torch==2.11.0".to_owned()], - lock_hash: "hash".to_owned(), - pip_cache_dir: None, - therock_channel: None, - therock_family: None, - therock_index_url: None, - therock_packages: Vec::new(), - torch_runtime_probe: None, - warnings: Vec::new(), - }; - - let resolved = resolve_manifest_python_executable(&manifest); - fs::remove_dir_all(root).ok(); - - assert_eq!(resolved, existing); - Ok(()) - } - - #[test] - fn parses_torch_runtime_probe() -> Result<()> { - let probe = parse_torch_runtime_probe( - r#"{"import_ok":true,"torch_version":"2.9.0+rocm","cuda_available":true,"device_count":1,"devices":["AMD Radeon"]}"#, - )?; - assert!(probe.import_ok); - assert_eq!(probe.torch_version.as_deref(), Some("2.9.0+rocm")); - assert!(probe.cuda_available); - assert_eq!(probe.device_count, 1); - assert_eq!(probe.devices, vec!["AMD Radeon".to_owned()]); - Ok(()) - } - - #[test] - fn parses_torch_runtime_probe_with_rocm_init_contract() -> Result<()> { - let probe = parse_torch_runtime_probe( - r#"{"import_ok":true,"torch_version":"2.10.0+rocm7.13.0a20260423","cuda_available":true,"device_count":1,"devices":["AMD Radeon RX 9070 XT"],"rocm_sdk":{"import_ok":true,"version":"7.13.0a20260423","site_packages":"C:\\venv\\Lib\\site-packages","default_target_family":"gfx1151","available_target_families":["gfx1151"],"resolved_target_family":"gfx1151","error":null},"torch_rocm_init":{"present":true,"path":"C:\\venv\\Lib\\site-packages\\torch\\_rocm_init.py","check_version":"7.13.0a20260423","preload_shortnames":["amd_comgr","amdhip64","hipblas"],"error":null}}"#, - )?; - - let sdk = probe.rocm_sdk.as_ref().expect("rocm_sdk probe"); - assert!(sdk.import_ok); - assert_eq!(sdk.version.as_deref(), Some("7.13.0a20260423")); - assert_eq!(sdk.resolved_target_family.as_deref(), Some("gfx1151")); - - let init = probe - .torch_rocm_init - .as_ref() - .expect("torch._rocm_init probe"); - assert!(init.present); - assert_eq!(init.check_version.as_deref(), Some("7.13.0a20260423")); - assert_eq!( - init.preload_shortnames, - vec![ - "amd_comgr".to_owned(), - "amdhip64".to_owned(), - "hipblas".to_owned() - ] - ); - Ok(()) - } - - #[test] - fn uv_install_base_targets_venv_python() { - let python = PathBuf::from("/envs/pytorch/bin/python"); - let args = uv_pip_install_base(&python); - assert_eq!(args[0], "pip"); - assert_eq!(args[1], "install"); - assert!( - args.windows(2) - .any(|pair| pair == ["--python", "/envs/pytorch/bin/python"]) - ); - } - - #[test] - fn therock_torch_package_specs_include_full_torch_stack() { - assert_eq!( - THEROCK_TORCH_PACKAGES, - &["torch", "torchvision", "torchaudio"] - ); - } - - #[test] - fn engine_dependencies_keep_transformers_on_windows_rocm_compatible_line() { - assert!(ENGINE_DEPENDENCIES.contains(&"transformers<5")); - assert!(ENGINE_DEPENDENCIES.contains(&"huggingface_hub<1")); - assert!(!ENGINE_DEPENDENCIES.contains(&"transformers")); - } - - #[test] - fn parses_runtime_torch_versions_as_exact_pip_specs() -> Result<()> { - let specs = parse_torch_package_version_specs( - r#"{"torch":"2.11.0+rocm7.13.0a20260416","torchvision":"0.26.0+rocm7.13.0a20260416","torchaudio":"2.11.0+rocm7.13.0a20260416"}"#, - )?; - assert_eq!( - specs, - vec![ - "torch==2.11.0+rocm7.13.0a20260416", - "torchvision==0.26.0+rocm7.13.0a20260416", - "torchaudio==2.11.0+rocm7.13.0a20260416", - ] - ); - Ok(()) - } - - #[test] - fn reads_runtime_torch_versions_from_dist_info_metadata() -> Result<()> { - let root = std::env::temp_dir().join(format!( - "rocm-pytorch-dist-info-{}", - rocm_core::unix_time_millis() - )); - let _ = fs::remove_dir_all(&root); - fs::create_dir_all(&root)?; - for (name, version) in [ - ("torch", "2.10.0+rocm7.13.0a20260511"), - ("torchvision", "0.25.0+rocm7.13.0a20260511"), - ("torchaudio", "2.10.0+rocm7.13.0a20260511"), - ] { - let dist_info = root.join(format!("{name}-{version}.dist-info")); - fs::create_dir_all(&dist_info)?; - fs::write( - dist_info.join("METADATA"), - format!("Metadata-Version: 2.4\nName: {name}\nVersion: {version}\n"), - )?; - } - - let specs = pinned_torch_package_specs_from_dist_info_dir(&root)?; - fs::remove_dir_all(root).ok(); - - assert_eq!( - specs, - vec![ - "torch==2.10.0+rocm7.13.0a20260511", - "torchvision==0.25.0+rocm7.13.0a20260511", - "torchaudio==2.10.0+rocm7.13.0a20260511", - ] - ); - Ok(()) - } - - #[test] - fn runtime_torch_version_parse_requires_full_stack() { - let error = parse_torch_package_version_specs( - r#"{"torch":"2.11.0+rocm7.13.0a20260416","torchvision":"0.26.0+rocm7.13.0a20260416"}"#, - ) - .expect_err("missing torchaudio must fail"); - assert!(error.to_string().contains("torchaudio")); - } - - #[test] - fn known_model_recipe_comes_from_shared_registry() { - let recipe = resolve_model_recipe("qwen").expect("recipe should resolve"); - assert_eq!(recipe.canonical_model_id, "Qwen/Qwen2.5-1.5B-Instruct"); - assert_eq!(recipe.source, "alias"); - assert_eq!(recipe.preferred_dtype, "bfloat16"); - assert_eq!(recipe.device_policy, DevicePolicy::GpuPreferred); - } - - #[test] - fn qwen35_is_rejected_before_pytorch_launch() { - let error = - resolve_model_recipe("qwen3.5").expect_err("qwen3.5 should be gated before launch"); - - assert!(error.to_string().contains("unknown architecture `qwen3_5`")); - assert!(error.to_string().contains("Use `qwen`")); - } - - #[test] - fn engine_recipe_overrides_supported_pytorch_worker_settings() -> Result<()> { - let mut hint = test_engine_recipe(ENGINE_NAME, ENGINE_RECIPE_CONTRACT_VERSION); - hint.required_flags = vec![ - "--trust-remote-code".to_owned(), - "--preferred-dtype".to_owned(), - "float16".to_owned(), - "--min-gpu-mem-gb=16".to_owned(), - ]; - - let recipe = - apply_pytorch_engine_recipe_overrides(resolve_model_recipe("tiny-gpt2")?, Some(&hint))?; - - assert!(recipe.trust_remote_code); - assert_eq!(recipe.preferred_dtype, "float16"); - assert_eq!(recipe.min_gpu_mem_gb, Some(16)); - assert!(recipe.estimated_memory.contains("16 GiB")); - Ok(()) - } - - #[test] - fn engine_recipe_rejects_unknown_pytorch_launch_flags() { - let mut hint = test_engine_recipe(ENGINE_NAME, ENGINE_RECIPE_CONTRACT_VERSION); - hint.required_flags = vec!["--enable-auto-tool-choice".to_owned()]; - - let error = apply_pytorch_engine_recipe_overrides( - resolve_model_recipe("tiny-gpt2").unwrap(), - Some(&hint), - ) - .expect_err("unknown PyTorch launch flag should fail"); - - assert!( - error - .to_string() - .contains("unsupported PyTorch launch recipe flag") - ); - } - - #[test] - fn resolve_model_echoes_matching_engine_recipe() -> Result<()> { - let hint = test_engine_recipe(ENGINE_NAME, ENGINE_RECIPE_CONTRACT_VERSION); - let response = resolve_model_response(ResolveModelRequest { - model_ref: "qwen".to_owned(), - runtime_id: None, - device_policy: Some(DevicePolicy::GpuRequired), - recipe_override: None, - engine_recipe: Some(hint.clone()), - })?; - - assert_eq!(response.engine_recipe, Some(hint)); - Ok(()) - } - - #[test] - fn resolve_model_rejects_mismatched_engine_recipe() { - let error = resolve_model_response(ResolveModelRequest { - model_ref: "qwen".to_owned(), - runtime_id: None, - device_policy: Some(DevicePolicy::GpuRequired), - recipe_override: None, - engine_recipe: Some(test_engine_recipe("vllm", ENGINE_RECIPE_CONTRACT_VERSION)), - }) - .expect_err("mismatched engine recipe should fail"); - - assert!(error.to_string().contains("does not match adapter")); - } - - #[test] - fn resolve_model_rejects_unsupported_engine_recipe_contract() { - let error = resolve_model_response(ResolveModelRequest { - model_ref: "qwen".to_owned(), - runtime_id: None, - device_policy: Some(DevicePolicy::GpuRequired), - recipe_override: None, - engine_recipe: Some(test_engine_recipe(ENGINE_NAME, "999.0.0")), - }) - .expect_err("unsupported recipe contract should fail"); - - assert!(error.to_string().contains("unsupported")); - } - - #[test] - fn tiny_model_recipe_uses_gpu_policy_from_registry() { - let recipe = resolve_model_recipe("tiny-gpt2").expect("recipe should resolve"); - assert_eq!(recipe.canonical_model_id, "sshleifer/tiny-gpt2"); - assert_eq!(recipe.device_policy, DevicePolicy::GpuRequired); - assert_eq!(recipe.min_gpu_mem_gb, Some(2)); - assert!(!recipe.trust_remote_code); - } -} diff --git a/engines/pytorch/src/main.rs b/engines/pytorch/src/main.rs deleted file mode 100644 index ee82cc96..00000000 --- a/engines/pytorch/src/main.rs +++ /dev/null @@ -1,7 +0,0 @@ -// Copyright © Advanced Micro Devices, Inc., or its affiliates. -// -// SPDX-License-Identifier: MIT - -fn main() -> anyhow::Result<()> { - rocm_engine_pytorch::run_cli() -} diff --git a/engines/pytorch/src/python_worker.py b/engines/pytorch/src/python_worker.py deleted file mode 100644 index d0195b3e..00000000 --- a/engines/pytorch/src/python_worker.py +++ /dev/null @@ -1,636 +0,0 @@ -# Copyright © Advanced Micro Devices, Inc., or its affiliates. -# -# SPDX-License-Identifier: MIT - -import argparse -import json -import os -import re -import sys -import time -import traceback -from collections.abc import Iterator -from threading import Thread -from typing import Any - -import torch -import uvicorn -from fastapi import FastAPI, HTTPException -from fastapi.responses import JSONResponse, StreamingResponse -from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer - - -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description="rocm-cli pytorch worker") - parser.add_argument("--service-id", required=True) - parser.add_argument("--model-ref", required=True) - parser.add_argument("--host", required=True) - parser.add_argument("--port", required=True, type=int) - parser.add_argument("--device-policy", required=True) - parser.add_argument("--state-path", required=True) - parser.add_argument("--env-id", required=True) - parser.add_argument("--runtime-id", required=True) - parser.add_argument("--preferred-dtype", default="auto") - parser.add_argument("--min-gpu-mem-gb", type=float) - parser.add_argument("--trust-remote-code", action="store_true") - return parser.parse_args() - - -def write_state(path: str, payload: dict[str, Any]) -> None: - os.makedirs(os.path.dirname(path), exist_ok=True) - temp_path = f"{path}.tmp" - with open(temp_path, "w", encoding="utf-8") as handle: - json.dump(payload, handle, indent=2) - os.replace(temp_path, path) - - -def collect_gpu_inventory() -> dict[str, Any]: - inventory: dict[str, Any] = { - "cuda_available": False, - "gpu_count": 0, - "per_gpu_mem_gb": [], - "max_single_gpu_mem_gb": None, - "total_gpu_mem_gb": 0.0, - } - try: - if not torch.cuda.is_available(): - return inventory - - per_gpu_mem_gb: list[float] = [] - for index in range(torch.cuda.device_count()): - total_bytes = torch.cuda.get_device_properties(index).total_memory - per_gpu_mem_gb.append(total_bytes / float(1024**3)) - - inventory["cuda_available"] = True - inventory["gpu_count"] = len(per_gpu_mem_gb) - inventory["per_gpu_mem_gb"] = per_gpu_mem_gb - inventory["max_single_gpu_mem_gb"] = ( - max(per_gpu_mem_gb) if per_gpu_mem_gb else None - ) - inventory["total_gpu_mem_gb"] = sum(per_gpu_mem_gb) - return inventory - except Exception: - return inventory - - -def detect_device( - policy: str, - min_gpu_mem_gb: float | None, -) -> tuple[str, dict[str, Any], str | None, str]: - inventory = collect_gpu_inventory() - if policy == "cpu_only": - raise RuntimeError( - "PyTorch CPU serving is not offered by rocm-cli; no CPU fallback is used" - ) - - if not inventory["cuda_available"]: - message = "torch.cuda.is_available() is false" - if policy in {"gpu_required", "gpu_preferred"}: - raise RuntimeError( - f"device policy {policy} does not permit implicit CPU execution because {message}; " - "no CPU fallback is used" - ) - raise RuntimeError(f"unsupported device policy {policy}") - - max_single_gpu_mem_gb = inventory["max_single_gpu_mem_gb"] - if ( - min_gpu_mem_gb is not None - and max_single_gpu_mem_gb is not None - and max_single_gpu_mem_gb + 0.5 < min_gpu_mem_gb - ): - message = ( - f"detected {max_single_gpu_mem_gb:.1f} GiB GPU memory but recipe recommends " - f"{min_gpu_mem_gb:.1f} GiB" - ) - if policy in {"gpu_required", "gpu_preferred"}: - raise RuntimeError(f"{message}; no CPU fallback is used") - raise RuntimeError(f"unsupported device policy {policy}") - - return "cuda", inventory, None, "single_gpu" - - -def normalize_content(content: Any) -> str: - if isinstance(content, list): - parts: list[str] = [] - for item in content: - if isinstance(item, dict): - if item.get("type") == "text": - parts.append(str(item.get("text", ""))) - else: - parts.append(json.dumps(item, ensure_ascii=False)) - else: - parts.append(str(item)) - return "".join(parts) - if content is None: - return "" - return str(content) - - -TOOL_CALL_PATTERN = re.compile(r"\s*(.*?)\s*", re.DOTALL) - - -def parse_tool_calls(text: str) -> tuple[str, list[dict[str, Any]]]: - tool_calls: list[dict[str, Any]] = [] - content_parts: list[str] = [] - cursor = 0 - for match in TOOL_CALL_PATTERN.finditer(text): - content_parts.append(text[cursor : match.start()]) - cursor = match.end() - raw_call = match.group(1).strip() - try: - payload = json.loads(raw_call) - except Exception: - content_parts.append(match.group(0)) - continue - - name = payload.get("name") - arguments = payload.get("arguments", {}) - if isinstance(payload.get("function"), dict): - function = payload["function"] - name = name or function.get("name") - arguments = function.get("arguments", arguments) - if not isinstance(name, str) or not name.strip(): - content_parts.append(match.group(0)) - continue - if isinstance(arguments, str): - try: - arguments = json.loads(arguments) if arguments.strip() else {} - except Exception: - arguments = {} - if not isinstance(arguments, dict): - arguments = {} - - call_index = len(tool_calls) - tool_calls.append( - { - "id": f"call_{call_index}", - "type": "function", - "function": { - "name": name.strip(), - "arguments": json.dumps(arguments, ensure_ascii=False), - }, - } - ) - - content_parts.append(text[cursor:]) - content = "".join(content_parts).strip() - return content, tool_calls - - -def resolve_torch_dtype(preferred_dtype: str) -> tuple[torch.dtype | None, str]: - normalized = preferred_dtype.strip().lower() - if normalized == "float32": - return torch.float32, "float32" - if normalized in {"float16", "fp16"}: - return torch.float16, "float16" - if normalized in {"bfloat16", "bf16"}: - return torch.bfloat16, "bfloat16" - - try: - if torch.cuda.is_available() and torch.cuda.is_bf16_supported(): - return torch.bfloat16, "bfloat16" - except Exception: - pass - return torch.float16, "float16" - - -class Runtime: - def __init__(self, args: argparse.Namespace): - self.args = args - self.device, self.gpu_inventory, self.device_note, self.placement_mode = ( - detect_device(args.device_policy, args.min_gpu_mem_gb) - ) - self.gpu_mem_gb = self.gpu_inventory.get("max_single_gpu_mem_gb") - self.total_gpu_mem_gb = self.gpu_inventory.get("total_gpu_mem_gb") - self.gpu_count = self.gpu_inventory.get("gpu_count", 0) - self.trust_remote_code = bool(args.trust_remote_code) - self.loaded_at = time.time() - self.compute_dtype_label = "float32" - - print( - f"[rocm-engine-pytorch] loading model={args.model_ref} device={self.device} " - f"trust_remote_code={self.trust_remote_code}", - flush=True, - ) - if self.device_note: - print(f"[rocm-engine-pytorch] device note: {self.device_note}", flush=True) - - tokenizer_kwargs: dict[str, Any] = { - "trust_remote_code": self.trust_remote_code, - } - self.tokenizer = AutoTokenizer.from_pretrained( - args.model_ref, **tokenizer_kwargs - ) - if self.tokenizer.pad_token is None and self.tokenizer.eos_token is not None: - self.tokenizer.pad_token = self.tokenizer.eos_token - - model_kwargs: dict[str, Any] = { - "trust_remote_code": self.trust_remote_code, - "low_cpu_mem_usage": True, - } - if self.device == "cuda": - torch_dtype, dtype_label = resolve_torch_dtype(args.preferred_dtype) - self.compute_dtype_label = dtype_label - if torch_dtype is not None: - model_kwargs["torch_dtype"] = torch_dtype - # Place the entire model on the single visible GPU. HIP_VISIBLE_DEVICES - # pins exactly one device, which torch sees as ordinal 0. Serving a - # model across multiple GPUs is not supported. - model_kwargs["device_map"] = {"": 0} - - self.model = AutoModelForCausalLM.from_pretrained( - args.model_ref, **model_kwargs - ) - self.model.eval() - if self.device == "cpu": - self.model = self.model.to("cpu") - self.input_device = self.resolve_input_device() - - write_state( - args.state_path, - { - "engine": "pytorch", - "service_id": args.service_id, - "env_id": args.env_id, - "runtime_id": args.runtime_id, - "model_ref": args.model_ref, - "status": "ready", - "pid": os.getpid(), - "device": self.device, - "device_policy": args.device_policy, - "device_note": self.device_note, - "placement_mode": self.placement_mode, - "input_device": str(self.input_device), - "gpu_count": self.gpu_count, - "per_gpu_mem_gb": self.gpu_inventory.get("per_gpu_mem_gb"), - "gpu_mem_gb": self.gpu_mem_gb, - "total_gpu_mem_gb": self.total_gpu_mem_gb, - "preferred_dtype": args.preferred_dtype, - "compute_dtype": self.compute_dtype_label, - "trust_remote_code": self.trust_remote_code, - "endpoint_url": f"http://{args.host}:{args.port}/v1", - }, - ) - print( - f"[rocm-engine-pytorch] ready service={args.service_id} endpoint=http://{args.host}:{args.port}/v1", - flush=True, - ) - - def model_device(self) -> str: - return str(self.input_device) - - def resolve_input_device(self) -> torch.device: - if self.device == "cpu": - return torch.device("cpu") - - try: - input_embeddings = self.model.get_input_embeddings() - if input_embeddings is not None and hasattr(input_embeddings, "weight"): - return input_embeddings.weight.device - except Exception: - pass - - hf_device_map = getattr(self.model, "hf_device_map", None) - if isinstance(hf_device_map, dict): - for device in hf_device_map.values(): - if isinstance(device, int): - return torch.device(f"cuda:{device}") - if isinstance(device, str) and device.startswith("cuda"): - return torch.device(device) - - return torch.device("cuda:0") - - def build_chat_prompt( - self, - messages: list[dict[str, Any]], - tools: list[dict[str, Any]] | None, - ) -> str: - normalized_messages = [ - { - "role": str(message.get("role", "user")), - "content": normalize_content(message.get("content")), - } - for message in messages - ] - if hasattr(self.tokenizer, "apply_chat_template"): - try: - kwargs: dict[str, Any] = { - "tokenize": False, - "add_generation_prompt": True, - } - if tools: - kwargs["tools"] = tools - return self.tokenizer.apply_chat_template(normalized_messages, **kwargs) - except Exception: - pass - - lines = [ - f"{message['role']}: {message['content']}" - for message in normalized_messages - ] - if tools: - lines.insert( - 0, - "system: Available tools are listed as JSON. When using a tool, return " - '{"name": "", "arguments": {}} inside tags.\n' - f"\n{json.dumps(tools, ensure_ascii=False)}\n", - ) - lines.append("assistant:") - return "\n".join(lines) - - def build_completion_prompt(self, prompt: Any) -> str: - if isinstance(prompt, str): - return prompt - return normalize_content(prompt) - - def encode_prompt(self, prompt: str) -> dict[str, torch.Tensor]: - encoded = self.tokenizer(prompt, return_tensors="pt") - return {key: value.to(self.model_device()) for key, value in encoded.items()} - - def generation_kwargs( - self, - max_tokens: int | None, - temperature: float | None, - ) -> dict[str, Any]: - kwargs: dict[str, Any] = { - "max_new_tokens": max_tokens or 256, - "pad_token_id": self.tokenizer.pad_token_id or self.tokenizer.eos_token_id, - "eos_token_id": self.tokenizer.eos_token_id, - "use_cache": True, - } - if temperature is not None and temperature > 0: - kwargs["do_sample"] = True - kwargs["temperature"] = temperature - else: - kwargs["do_sample"] = False - return kwargs - - def generate( - self, - prompt: str, - max_tokens: int | None, - temperature: float | None, - ) -> str: - encoded = self.encode_prompt(prompt) - generation_kwargs = self.generation_kwargs(max_tokens, temperature) - - with torch.inference_mode(): - output = self.model.generate(**encoded, **generation_kwargs) - - prompt_tokens = encoded["input_ids"].shape[-1] - completion_tokens = output[0][prompt_tokens:] - return self.tokenizer.decode( - completion_tokens, skip_special_tokens=True - ).strip() - - def generate_stream( - self, - prompt: str, - max_tokens: int | None, - temperature: float | None, - ) -> Iterator[str]: - encoded = self.encode_prompt(prompt) - generation_kwargs = self.generation_kwargs(max_tokens, temperature) - streamer = TextIteratorStreamer( - self.tokenizer, - skip_prompt=True, - skip_special_tokens=True, - ) - generation_kwargs["streamer"] = streamer - - errors: list[Exception] = [] - - def run_generation() -> None: - try: - with torch.inference_mode(): - self.model.generate(**encoded, **generation_kwargs) - except Exception as exc: - errors.append(exc) - - worker = Thread(target=run_generation, daemon=True) - worker.start() - for chunk in streamer: - if chunk: - yield chunk - worker.join() - if errors: - raise errors[0] - - -def stream_chat_chunks(service_id: str, model: str, chunks: Iterator[str]): - sent_role = False - for chunk in chunks: - delta: dict[str, Any] = {"content": chunk} - if not sent_role: - delta["role"] = "assistant" - sent_role = True - payload = { - "id": f"chatcmpl-{service_id}", - "object": "chat.completion.chunk", - "model": model, - "choices": [ - { - "index": 0, - "delta": delta, - "finish_reason": None, - } - ], - } - yield f"data: {json.dumps(payload, ensure_ascii=False)}\n\n" - - final_payload = { - "id": f"chatcmpl-{service_id}", - "object": "chat.completion.chunk", - "model": model, - "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], - } - yield f"data: {json.dumps(final_payload, ensure_ascii=False)}\n\n" - yield "data: [DONE]\n\n" - - -def stream_completion_chunks(service_id: str, model: str, chunks: Iterator[str]): - for chunk in chunks: - payload = { - "id": f"cmpl-{service_id}", - "object": "text_completion", - "model": model, - "choices": [{"index": 0, "text": chunk, "finish_reason": None}], - } - yield f"data: {json.dumps(payload, ensure_ascii=False)}\n\n" - - final_payload = { - "id": f"cmpl-{service_id}", - "object": "text_completion", - "model": model, - "choices": [{"index": 0, "text": "", "finish_reason": "stop"}], - } - yield f"data: {json.dumps(final_payload, ensure_ascii=False)}\n\n" - yield "data: [DONE]\n\n" - - -def create_app(runtime: Runtime) -> FastAPI: - app = FastAPI(title="rocm-cli pytorch engine") - - @app.get("/healthz") - async def healthz(): - return JSONResponse( - { - "status": "ok", - "engine": "pytorch", - "service_id": runtime.args.service_id, - "model": runtime.args.model_ref, - "device": runtime.device, - "device_policy": runtime.args.device_policy, - "device_note": runtime.device_note, - "placement_mode": runtime.placement_mode, - "input_device": runtime.model_device(), - "gpu_count": runtime.gpu_count, - "per_gpu_mem_gb": runtime.gpu_inventory.get("per_gpu_mem_gb"), - "gpu_mem_gb": runtime.gpu_mem_gb, - "total_gpu_mem_gb": runtime.total_gpu_mem_gb, - "compute_dtype": runtime.compute_dtype_label, - "loaded_at": runtime.loaded_at, - } - ) - - @app.get("/v1/models") - async def models(): - return JSONResponse( - { - "object": "list", - "data": [ - { - "id": runtime.args.model_ref, - "object": "model", - "owned_by": "rocm-cli", - } - ], - } - ) - - @app.post("/v1/chat/completions") - async def chat_completions(request: dict[str, Any]): - messages = request.get("messages") or [] - if not isinstance(messages, list) or not messages: - raise HTTPException( - status_code=400, detail="messages must be a non-empty list" - ) - - stream = bool(request.get("stream")) - max_tokens = request.get("max_tokens") - temperature = request.get("temperature") - model = request.get("model") or runtime.args.model_ref - tools = request.get("tools") - if not isinstance(tools, list): - tools = None - prompt = runtime.build_chat_prompt(messages, tools) - - if stream: - return StreamingResponse( - stream_chat_chunks( - runtime.args.service_id, - model, - runtime.generate_stream(prompt, max_tokens, temperature), - ), - media_type="text/event-stream", - ) - - text = runtime.generate(prompt, max_tokens, temperature) - content, tool_calls = parse_tool_calls(text) - message: dict[str, Any] = {"role": "assistant", "content": content} - finish_reason = "stop" - if tool_calls: - message["tool_calls"] = tool_calls - finish_reason = "tool_calls" - return JSONResponse( - { - "id": f"chatcmpl-{runtime.args.service_id}", - "object": "chat.completion", - "model": model, - "choices": [ - { - "index": 0, - "message": message, - "finish_reason": finish_reason, - } - ], - } - ) - - @app.post("/v1/completions") - async def completions(request: dict[str, Any]): - prompt = runtime.build_completion_prompt(request.get("prompt")) - if not prompt: - raise HTTPException(status_code=400, detail="prompt must not be empty") - - stream = bool(request.get("stream")) - max_tokens = request.get("max_tokens") - temperature = request.get("temperature") - model = request.get("model") or runtime.args.model_ref - - if stream: - return StreamingResponse( - stream_completion_chunks( - runtime.args.service_id, - model, - runtime.generate_stream(prompt, max_tokens, temperature), - ), - media_type="text/event-stream", - ) - - text = runtime.generate(prompt, max_tokens, temperature) - return JSONResponse( - { - "id": f"cmpl-{runtime.args.service_id}", - "object": "text_completion", - "model": model, - "choices": [ - { - "index": 0, - "text": text, - "finish_reason": "stop", - } - ], - } - ) - - return app - - -def main() -> int: - args = parse_args() - try: - runtime = Runtime(args) - app = create_app(runtime) - uvicorn.run(app, host=args.host, port=args.port, log_level="info") - return 0 - except Exception as exc: - write_state( - args.state_path, - { - "engine": "pytorch", - "service_id": args.service_id, - "env_id": args.env_id, - "runtime_id": args.runtime_id, - "model_ref": args.model_ref, - "status": "failed", - "pid": os.getpid(), - "placement_mode": None, - "input_device": None, - "gpu_count": 0, - "per_gpu_mem_gb": [], - "gpu_mem_gb": None, - "total_gpu_mem_gb": None, - "error": str(exc), - "traceback": traceback.format_exc(), - }, - ) - print( - f"[rocm-engine-pytorch] startup failed: {exc}", file=sys.stderr, flush=True - ) - traceback.print_exc() - return 1 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/engines/sglang/Cargo.toml b/engines/sglang/Cargo.toml deleted file mode 100644 index e8ff59b9..00000000 --- a/engines/sglang/Cargo.toml +++ /dev/null @@ -1,23 +0,0 @@ -[package] -name = "rocm-engine-sglang" -version.workspace = true -edition.workspace = true -license.workspace = true -repository.workspace = true -rust-version.workspace = true -publish.workspace = true - -[lints] -workspace = true - -[[bin]] -name = "rocm-engine-sglang" -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 diff --git a/engines/sglang/src/lib.rs b/engines/sglang/src/lib.rs deleted file mode 100644 index 73162318..00000000 --- a/engines/sglang/src/lib.rs +++ /dev/null @@ -1,1856 +0,0 @@ -// Copyright © Advanced Micro Devices, Inc., or its affiliates. -// -// SPDX-License-Identifier: MIT - -use anyhow::{Context, Result, bail}; -use clap::{Parser, Subcommand}; -use rocm_core::{ - AppPaths, DEFAULT_LOCAL_PORT, format_http_base_url, openai_models_endpoint_has_model, - require_nonempty, -}; -use rocm_engine_protocol::{ - DEFAULT_LOG_TAIL_LINES, DetectRequest, DetectResponse, DevicePolicy, - ENGINE_RECIPE_CONTRACT_VERSION, EndpointRequest, EndpointResponse, EngineCapabilities, - EngineDeviceAvailability, EngineMethod, EngineRecipeHint, EngineRequestEnvelope, - EngineResponseEnvelope, GpuSelection, HealthcheckRequest, HealthcheckResponse, InstallRequest, - InstallResponse, LaunchRequest, LaunchResponse, LogsRequest, LogsResponse, ResolveModelRequest, - ResolveModelResponse, StopRequest, StopResponse, -}; -use serde::de::DeserializeOwned; -use serde::{Deserialize, Serialize}; -use serde_json::{Value, json}; -use std::collections::hash_map::DefaultHasher; -use std::ffi::OsString; -use std::fs; -use std::hash::{Hash, Hasher}; -use std::io::{Read, Write}; -use std::path::{Path, PathBuf}; -use std::process::{Command as ProcessCommand, Stdio}; -use std::time::{Duration, SystemTime, UNIX_EPOCH}; - -const ENGINE_NAME: &str = "sglang"; -const DEFAULT_HOST: &str = "127.0.0.1"; -const HEALTHCHECK_TIMEOUT_MS: u64 = 700; - -#[derive(Parser, Debug)] -#[command(name = "rocm-engine-sglang", about = "rocm-cli SGLang engine adapter")] -struct Cli { - #[command(subcommand)] - command: CommandKind, -} - -#[derive(Subcommand, Debug)] -enum CommandKind { - Detect, - Capabilities, - Install { - #[arg(long, default_value = "external-sglang")] - runtime_id: String, - #[arg(long)] - reinstall: bool, - }, - ResolveModel { - model_ref: String, - #[arg(long)] - device_policy: Option, - }, - Launch { - service_id: String, - model_ref: String, - #[arg(long, default_value = DEFAULT_HOST)] - host: String, - #[arg(long, default_value_t = DEFAULT_LOCAL_PORT)] - port: u16, - #[arg(long)] - device_policy: Option, - #[arg(long)] - runtime_id: Option, - #[arg(long)] - env_id: Option, - #[arg(long)] - gpu: Option, - }, - Stdio, - ServeHttp { - service_id: String, - model_ref: String, - #[arg(long, default_value = DEFAULT_HOST)] - host: String, - #[arg(long, default_value_t = DEFAULT_LOCAL_PORT)] - port: u16, - #[arg(long, default_value = "gpu_required")] - device_policy: String, - #[arg(long)] - runtime_id: Option, - #[arg(long)] - env_id: Option, - #[arg(long)] - state_path: PathBuf, - #[arg(long)] - engine_recipe_json: Option, - #[arg(long)] - gpu: Option, - }, -} - -#[derive(Debug, Clone)] -struct SglangRuntime { - runtime_id: String, - env_id: String, - command: PathBuf, - launcher: SglangLauncher, - python_executable: Option, - version: Option, - source: String, - sdk_root: Option, - sdk_bin: Option, - sdk_bin_paths: Vec, - sdk_library_paths: Vec, -} - -#[derive(Debug, Clone, Copy, Eq, PartialEq)] -enum SglangLauncher { - Command, - PythonModule, -} - -#[derive(Debug, Clone, Deserialize)] -struct TheRockRuntimeManifest { - #[serde(default)] - runtime_key: Option, - #[serde(default)] - runtime_id: Option, - #[serde(default)] - python_executable: Option, - #[serde(default)] - rocm_sdk: Option, - #[serde(default)] - installed_at_unix_ms: Option, -} - -#[derive(Debug, Clone, Deserialize)] -struct RocmSdkRuntimeProbe { - #[serde(default)] - import_ok: bool, - #[serde(default)] - root_path: Option, - #[serde(default)] - bin_path: Option, - #[serde(default)] - bin_paths: Vec, - #[serde(default)] - library_paths: Vec, -} - -#[derive(Debug, Clone)] -struct ServiceFiles { - state_path: PathBuf, - log_path: PathBuf, -} - -pub fn run_cli() -> Result<()> { - let cli = Cli::parse(); - match cli.command { - CommandKind::Detect => print_json(&detect_response())?, - CommandKind::Capabilities => print_json(&capabilities())?, - CommandKind::Install { - runtime_id, - reinstall, - } => print_json(&install_response(InstallRequest { - runtime_id, - python_version: None, - env_root: None, - reinstall, - })?)?, - CommandKind::ResolveModel { - model_ref, - device_policy, - } => print_json(&resolve_model_response(ResolveModelRequest { - model_ref, - runtime_id: None, - device_policy: device_policy - .as_deref() - .map(|value| parse_device_policy_arg(Some(value))) - .transpose()?, - recipe_override: None, - engine_recipe: None, - })?)?, - CommandKind::Launch { - service_id, - model_ref, - host, - port, - device_policy, - runtime_id, - env_id, - gpu, - } => print_json(&launch_service(LaunchRequest { - service_id, - env_id, - runtime_id, - model_ref, - host, - port, - device_policy: Some(parse_device_policy_arg(device_policy.as_deref())?), - endpoint_mode: Some("openai".to_owned()), - engine_recipe: None, - gpu_selection: parse_gpu_selection_arg(gpu.as_deref())?, - })?)?, - CommandKind::Stdio => { - let envelope = read_request()?; - print_json(&handle_envelope(envelope))?; - } - CommandKind::ServeHttp { - service_id, - model_ref, - host, - port, - device_policy, - runtime_id, - env_id, - state_path, - engine_recipe_json, - gpu, - } => serve_http(ServeHttpRequest { - service_id, - model_ref, - host, - port, - device_policy: parse_device_policy_arg(Some(&device_policy))?, - gpu_indices: parse_gpu_indices_arg(gpu.as_deref())?, - runtime_id, - env_id, - state_path, - engine_recipe: parse_engine_recipe_json(engine_recipe_json)?, - })?, - } - Ok(()) -} - -pub fn builtin_handle_envelope(envelope: EngineRequestEnvelope) -> EngineResponseEnvelope { - handle_envelope(envelope) -} - -#[allow(clippy::too_many_arguments)] -pub fn builtin_serve_http( - service_id: String, - model_ref: String, - host: String, - port: u16, - device_policy: DevicePolicy, - gpu_indices: Vec, - runtime_id: Option, - env_id: Option, - state_path: PathBuf, - engine_recipe: Option, -) -> Result<()> { - serve_http(ServeHttpRequest { - service_id, - model_ref, - host, - port, - device_policy, - gpu_indices, - runtime_id, - env_id, - state_path, - engine_recipe, - }) -} - -fn handle_envelope(envelope: EngineRequestEnvelope) -> EngineResponseEnvelope { - match envelope.method { - EngineMethod::Detect => { - deserialize_and_respond::(envelope.payload, |_| { - Ok(detect_response()) - }) - } - EngineMethod::Capabilities => EngineResponseEnvelope::success(capabilities()), - EngineMethod::Install => { - deserialize_and_respond::(envelope.payload, install_response) - } - EngineMethod::ResolveModel => deserialize_and_respond::( - envelope.payload, - resolve_model_response, - ), - EngineMethod::Launch => { - deserialize_and_respond::(envelope.payload, launch_service) - } - EngineMethod::Healthcheck => deserialize_and_respond::( - envelope.payload, - healthcheck_service, - ), - EngineMethod::Endpoint => { - deserialize_and_respond::(envelope.payload, endpoint_response) - } - EngineMethod::Stop => { - deserialize_and_respond::(envelope.payload, stop_service) - } - EngineMethod::Logs => { - deserialize_and_respond::(envelope.payload, logs_response) - } - } -} - -fn deserialize_and_respond(payload: Value, handler: F) -> EngineResponseEnvelope -where - T: DeserializeOwned, - F: FnOnce(T) -> Result, - U: Serialize, -{ - match serde_json::from_value::(payload) { - Ok(request) => match handler(request) { - Ok(response) => EngineResponseEnvelope::success(response), - Err(error) => EngineResponseEnvelope::failure("request_failed", error.to_string()), - }, - Err(error) => EngineResponseEnvelope::failure("invalid_payload", error.to_string()), - } -} - -fn detect_response() -> DetectResponse { - let runtime = resolve_sglang_runtime(None); - let installed = runtime.is_ok(); - let mut notes = Vec::new(); - if cfg!(windows) { - notes.push(windows_unsupported_message().to_owned()); - } else if let Err(error) = runtime.as_ref() { - notes.push(error.to_string()); - } - let runtime = runtime.ok(); - if let Some(runtime) = runtime.as_ref() { - notes.push(format!( - "SGLang launcher resolved from {}; no CPU fallback is used", - runtime.source - )); - } - - DetectResponse { - installed, - env_id: runtime.as_ref().map(|runtime| runtime.env_id.clone()), - runtime_kind: Some("external_sglang".to_owned()), - runtime_executable: runtime - .as_ref() - .map(|runtime| runtime.command.display().to_string()), - managed_env: Some(runtime.as_ref().is_some_and(runtime_is_managed)), - python_version: runtime - .as_ref() - .and_then(|runtime| runtime.version.as_ref()) - .map(|version| format!("SGLang {version}")), - torch_version: None, - transformers_version: None, - available_devices: vec![EngineDeviceAvailability { - kind: "rocm_gpu".to_owned(), - available: installed && !cfg!(windows), - reason: if cfg!(windows) { - Some(windows_unsupported_message().to_owned()) - } else if installed { - None - } else { - Some("SGLang is not installed in a Linux/WSL ROCm Python environment".to_owned()) - }, - }], - capabilities: capabilities(), - notes, - } -} - -fn capabilities() -> EngineCapabilities { - EngineCapabilities { - cpu: false, - rocm_gpu: !cfg!(windows), - openai_compatible: true, - tool_calling: false, - quantized_models: "sglang-supported".to_owned(), - reasoning_parser: false, - } -} - -fn install_response(request: InstallRequest) -> Result { - let runtime = resolve_sglang_runtime(Some(&request.runtime_id))?; - let env_path = runtime - .command - .parent() - .map_or_else(|| PathBuf::from("."), Path::to_path_buf); - Ok(InstallResponse { - env_id: runtime.env_id.clone(), - env_path: env_path.display().to_string(), - python_executable: runtime - .python_executable - .as_ref() - .unwrap_or(&runtime.command) - .display() - .to_string(), - runtime_kind: Some("external_sglang".to_owned()), - runtime_executable: Some(runtime.command.display().to_string()), - managed_env: Some(runtime_is_managed(&runtime)), - installed_packages: vec![format!( - "sglang{}", - runtime - .version - .as_deref() - .map(|version| format!("=={version}")) - .unwrap_or_default() - )], - capabilities: capabilities(), - lock_hash: runtime_lock_hash(&runtime), - warnings: sglang_runtime_warnings(&runtime), - }) -} - -fn runtime_is_managed(runtime: &SglangRuntime) -> bool { - runtime.source.starts_with("managed_runtime_manifest") -} - -fn sglang_runtime_warnings(runtime: &SglangRuntime) -> Vec { - let runtime_scope = if runtime_is_managed(runtime) { - "rocm-cli records this SGLang launcher from a managed TheRock runtime; it does not pip install SGLang automatically" - } else { - "rocm-cli records this as an external SGLang runtime; it does not pip install SGLang automatically" - }; - vec![ - runtime_scope.to_owned(), - "SGLang serving remains ROCm GPU required; no CPU fallback is used".to_owned(), - ] -} - -fn resolve_model_response(request: ResolveModelRequest) -> Result { - let device_policy = normalize_sglang_device_policy(request.device_policy)?; - let engine_recipe = accepted_engine_recipe(request.engine_recipe)?; - Ok(ResolveModelResponse { - canonical_model_id: request.model_ref, - task: "text-generation".to_owned(), - source: "huggingface_or_local".to_owned(), - revision: "main".to_owned(), - loader: "sglang".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": DEFAULT_HOST, - "port": DEFAULT_LOCAL_PORT - }), - engine_recipe, - warnings: vec![ - "SGLang is treated as a ROCm GPU engine in rocm-cli; select another engine explicitly for CPU serving".to_owned(), - ], - }) -} - -fn accepted_engine_recipe( - engine_recipe: Option, -) -> Result> { - if let Some(hint) = &engine_recipe { - if hint.engine != ENGINE_NAME { - bail!( - "engine_recipe target `{}` does not match adapter `{}`", - hint.engine, - ENGINE_NAME - ); - } - if hint.contract_version != ENGINE_RECIPE_CONTRACT_VERSION { - bail!( - "engine_recipe contract `{}` is unsupported; expected `{}`", - hint.contract_version, - ENGINE_RECIPE_CONTRACT_VERSION - ); - } - } - Ok(engine_recipe) -} - -fn parse_engine_recipe_json(value: Option) -> Result> { - value - .map(|text| { - serde_json::from_str::(&text) - .context("failed to parse engine recipe JSON") - }) - .transpose() - .and_then(accepted_engine_recipe) -} - -fn launch_service(request: LaunchRequest) -> Result { - let device_policy = normalize_sglang_device_policy(request.device_policy)?; - let engine_recipe = accepted_engine_recipe(request.engine_recipe)?; - let runtime = resolve_sglang_runtime(request.runtime_id.as_deref())?; - let state_path = AppPaths::discover()? - .engine_state_dir(ENGINE_NAME) - .join(format!("{}.json", request.service_id)); - let serve_request = ServeHttpRequest { - service_id: request.service_id.clone(), - model_ref: request.model_ref.clone(), - host: request.host.clone(), - port: request.port, - device_policy, - gpu_indices: rocm_engine_protocol::launch_gpu_indices(request.gpu_selection.as_ref()), - runtime_id: request.runtime_id.clone(), - env_id: request.env_id.clone(), - state_path: state_path.clone(), - engine_recipe, - }; - let log_path = AppPaths::discover()? - .engine_logs_dir(ENGINE_NAME) - .join(format!("{}.log", request.service_id)); - let child = spawn_sglang_server(&serve_request, &runtime, Some(&log_path))?; - let pid = child.id(); - write_running_state(&serve_request, &runtime, pid)?; - Ok(LaunchResponse { - service_id: request.service_id, - pid, - endpoint_url: endpoint_url(&request.host, request.port), - log_path: log_path.display().to_string(), - state_path: state_path.display().to_string(), - }) -} - -#[derive(Debug, Clone)] -struct ServeHttpRequest { - service_id: String, - model_ref: String, - host: String, - port: u16, - device_policy: DevicePolicy, - gpu_indices: Vec, - runtime_id: Option, - env_id: Option, - state_path: PathBuf, - engine_recipe: Option, -} - -fn serve_http(request: ServeHttpRequest) -> Result<()> { - let runtime = resolve_sglang_runtime(request.runtime_id.as_deref())?; - let mut child = spawn_sglang_server(&request, &runtime, None)?; - write_running_state(&request, &runtime, child.id())?; - let status = child.wait().context("failed waiting for SGLang server")?; - write_terminal_state( - &request.state_path, - if status.success() { - "stopped" - } else { - "failed" - }, - )?; - if status.success() { - Ok(()) - } else { - std::process::exit(status.code().unwrap_or(1)); - } -} - -fn spawn_sglang_server( - request: &ServeHttpRequest, - runtime: &SglangRuntime, - log_path: Option<&Path>, -) -> Result { - require_nonempty(&request.service_id, "service_id")?; - require_nonempty(&request.model_ref, "model_ref")?; - if !matches!(request.device_policy, DevicePolicy::GpuRequired) { - bail!("SGLang launch requires ROCm GPU execution; no CPU fallback is used"); - } - - if let Some(parent) = request.state_path.parent() { - fs::create_dir_all(parent) - .with_context(|| format!("failed to create {}", parent.display()))?; - } - if let Some(log_path) = log_path - && let Some(parent) = log_path.parent() - { - fs::create_dir_all(parent) - .with_context(|| format!("failed to create {}", parent.display()))?; - } - - let mut command = ProcessCommand::new(&runtime.command); - command - .args(sglang_server_args( - runtime.launcher, - &request.model_ref, - &request.host, - request.port, - )) - .args(engine_recipe_launch_args(request.engine_recipe.as_ref())) - .stdin(Stdio::null()); - apply_therock_env(&mut command, runtime)?; - rocm_engine_protocol::apply_gpu_visibility(&mut command, &request.gpu_indices); - if let Some(log_path) = log_path { - let log = fs::File::create(log_path) - .with_context(|| format!("failed to create {}", log_path.display()))?; - command.stdout(Stdio::from( - log.try_clone().context("failed to clone log handle")?, - )); - command.stderr(Stdio::from(log)); - } - - command.spawn().with_context(|| { - format!( - "failed to spawn SGLang command {}", - runtime.command.display() - ) - }) -} - -fn sglang_server_args( - launcher: SglangLauncher, - model_ref: &str, - host: &str, - port: u16, -) -> Vec { - let mut args = Vec::new(); - match launcher { - SglangLauncher::Command => { - args.push("serve".to_owned()); - } - SglangLauncher::PythonModule => { - args.push("-m".to_owned()); - args.push("sglang.launch_server".to_owned()); - } - } - args.extend([ - "--model-path".to_owned(), - model_ref.to_owned(), - "--host".to_owned(), - host.to_owned(), - "--port".to_owned(), - port.to_string(), - "--attention-backend".to_owned(), - "triton".to_owned(), - ]); - args -} - -fn engine_recipe_launch_args(engine_recipe: Option<&EngineRecipeHint>) -> Vec { - engine_recipe - .map(|hint| hint.required_flags.clone()) - .unwrap_or_default() -} - -fn healthcheck_service(request: HealthcheckRequest) -> Result { - require_nonempty(&request.service_id, "service_id")?; - let files = service_files(&request.service_id)?; - let state = read_service_state(&files.state_path).ok(); - let endpoint_url = state.as_ref().and_then(endpoint_url_from_state); - let model_ref = state - .as_ref() - .and_then(|value| value_string(value, "model_ref")); - let ready = endpoint_url - .as_deref() - .map(|endpoint| query_loaded_model_endpoint(endpoint, model_ref.as_deref())) - .transpose() - .unwrap_or(None) - .unwrap_or(false); - let status = if ready { - "ready".to_owned() - } else { - state - .as_ref() - .and_then(|value| value_string(value, "status")) - .unwrap_or_else(|| "unknown".to_owned()) - }; - Ok(HealthcheckResponse { - status, - model_loaded: ready, - device: if state.is_some() { - "rocm_gpu".to_owned() - } else { - "unknown".to_owned() - }, - uptime_sec: 0, - queue_depth: 0, - last_error: None, - tokens_per_sec: None, - }) -} - -fn endpoint_response(request: EndpointRequest) -> Result { - require_nonempty(&request.service_id, "service_id")?; - let files = service_files(&request.service_id)?; - let state = read_service_state(&files.state_path) - .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(), - ], - }) -} - -fn logs_response(request: LogsRequest) -> Result { - require_nonempty(&request.service_id, "service_id")?; - let files = service_files(&request.service_id)?; - let limit = request.tail_lines.unwrap_or(DEFAULT_LOG_TAIL_LINES); - Ok(LogsResponse { - log_path: files.log_path.display().to_string(), - recent_lines: if files.log_path.is_file() { - tail_lines(&files.log_path, limit)? - } else { - Vec::new() - }, - }) -} - -fn stop_service(request: StopRequest) -> Result { - require_nonempty(&request.service_id, "service_id")?; - let files = service_files(&request.service_id)?; - let state = read_service_state(&files.state_path).ok(); - let stopped = match state.as_ref().and_then(pid_from_state) { - Some(pid) => terminate_pid(pid, request.force), - None => false, - }; - if stopped { - write_terminal_state(&files.state_path, "stopped")?; - } - Ok(StopResponse { - stopped, - graceful: stopped && !request.force, - }) -} - -fn resolve_sglang_runtime(runtime_id: Option<&str>) -> Result { - if cfg!(windows) { - bail!("{}", windows_unsupported_message()); - } - - if let Some(command) = std::env::var_os("ROCM_CLI_SGLANG_COMMAND") - .or_else(|| std::env::var_os("SGLANG_COMMAND")) - .map(PathBuf::from) - { - let command = resolve_command_path(&command)?; - return Ok(SglangRuntime { - runtime_id: runtime_id.unwrap_or("external-sglang").to_owned(), - env_id: "external-sglang-command".to_owned(), - command, - launcher: SglangLauncher::Command, - python_executable: None, - version: None, - source: "environment command".to_owned(), - sdk_root: None, - sdk_bin: None, - sdk_bin_paths: Vec::new(), - sdk_library_paths: Vec::new(), - }); - } - - if let Some(python) = std::env::var_os("ROCM_CLI_SGLANG_PYTHON") - .or_else(|| std::env::var_os("SGLANG_PYTHON")) - .map(PathBuf::from) - .filter(|path| path.is_file()) - { - return runtime_from_python( - python, - runtime_id.unwrap_or("external-sglang-python"), - "environment python", - None, - None, - Vec::new(), - Vec::new(), - ); - } - - if let Some(runtime) = resolve_managed_runtime(runtime_id)? { - return Ok(runtime); - } - - if let Some(command) = find_command_on_path("sglang") { - return Ok(SglangRuntime { - runtime_id: runtime_id.unwrap_or("external-sglang-path").to_owned(), - env_id: "external-sglang-path".to_owned(), - command, - launcher: SglangLauncher::Command, - python_executable: None, - version: None, - source: "PATH".to_owned(), - sdk_root: None, - sdk_bin: None, - sdk_bin_paths: Vec::new(), - sdk_library_paths: Vec::new(), - }); - } - - bail!( - "SGLang is not installed in a Linux/WSL ROCm Python environment. Install/build SGLang against a ROCm-capable Python environment, then set ROCM_CLI_SGLANG_COMMAND, set ROCM_CLI_SGLANG_PYTHON, or install it into the active rocm-cli TheRock runtime. Native Windows is skipped; no CPU fallback is used." - ) -} - -fn runtime_from_python( - python: PathBuf, - runtime_id: &str, - source: &str, - sdk_root: Option, - sdk_bin: Option, - sdk_bin_paths: Vec, - sdk_library_paths: Vec, -) -> Result { - let (command, launcher, version) = if let Some(command) = sglang_command_from_python(&python) { - ( - command, - SglangLauncher::Command, - probe_sglang_version(&python).ok().flatten(), - ) - } else { - let version = probe_sglang_version(&python) - .with_context(|| format!("SGLang package not found in {}", python.display()))? - .unwrap_or_else(|| "unknown".to_owned()); - (python.clone(), SglangLauncher::PythonModule, Some(version)) - }; - Ok(SglangRuntime { - runtime_id: runtime_id.to_owned(), - env_id: format!("external-sglang-{}", stable_id_component(runtime_id)), - command, - launcher, - python_executable: Some(python), - version, - source: source.to_owned(), - sdk_root, - sdk_bin, - sdk_bin_paths, - sdk_library_paths, - }) -} - -fn resolve_managed_runtime(runtime_id: Option<&str>) -> Result> { - let paths = AppPaths::discover()?; - let registry = paths.data_dir.join("runtimes").join("registry"); - if !registry.is_dir() { - return Ok(None); - } - let mut manifests = Vec::new(); - for entry in - fs::read_dir(®istry).with_context(|| format!("failed to read {}", registry.display()))? - { - let path = entry?.path(); - if path.extension().and_then(|value| value.to_str()) != Some("json") { - continue; - } - let bytes = - fs::read(&path).with_context(|| format!("failed to read {}", path.display()))?; - let Ok(manifest) = serde_json::from_slice::(&bytes) else { - continue; - }; - if !runtime_matches(&manifest, runtime_id) { - continue; - } - manifests.push((manifest.installed_at_unix_ms.unwrap_or(0), manifest)); - } - manifests.sort_by_key(|(installed_at, _)| std::cmp::Reverse(*installed_at)); - - for (_, manifest) in manifests { - let Some(python) = manifest - .python_executable - .clone() - .filter(|path| path.is_file()) - else { - continue; - }; - let runtime_id = manifest - .runtime_id - .as_deref() - .unwrap_or("therock-sglang-runtime") - .to_owned(); - let source = manifest.runtime_key.as_deref().map_or_else( - || "managed_runtime_manifest".to_owned(), - |key| format!("managed_runtime_manifest:{key}"), - ); - let (sdk_root, sdk_bin, sdk_bin_paths, sdk_library_paths) = manifest - .rocm_sdk - .as_ref() - .filter(|probe| probe.import_ok) - .map_or((None, None, Vec::new(), Vec::new()), |probe| { - ( - probe.root_path.clone(), - probe.bin_path.clone(), - probe.bin_paths.clone(), - probe.library_paths.clone(), - ) - }); - if let Ok(runtime) = runtime_from_python( - python, - &runtime_id, - &source, - sdk_root, - sdk_bin, - sdk_bin_paths, - sdk_library_paths, - ) { - return Ok(Some(runtime)); - } - } - Ok(None) -} - -fn runtime_matches(manifest: &TheRockRuntimeManifest, requested: Option<&str>) -> bool { - let Some(requested) = requested.map(str::trim).filter(|value| !value.is_empty()) else { - return true; - }; - let requested = requested.to_ascii_lowercase(); - if requested == "external" || requested == "external-sglang" { - return false; - } - for candidate in [ - manifest.runtime_id.as_deref(), - manifest.runtime_key.as_deref(), - ] - .into_iter() - .flatten() - { - let candidate = candidate.to_ascii_lowercase(); - if candidate == requested || candidate.starts_with(&requested) { - return true; - } - } - false -} - -fn normalize_sglang_device_policy(policy: Option) -> Result { - match policy.unwrap_or(DevicePolicy::GpuRequired) { - DevicePolicy::GpuRequired => Ok(DevicePolicy::GpuRequired), - DevicePolicy::GpuPreferred => Ok(DevicePolicy::GpuRequired), - DevicePolicy::CpuOnly => { - bail!("SGLang adapter is ROCm GPU-only in rocm-cli; no CPU fallback is used") - } - } -} - -fn parse_device_policy_arg(policy: Option<&str>) -> Result { - match policy.unwrap_or("gpu_required") { - "gpu" | "gpu_required" => Ok(DevicePolicy::GpuRequired), - "gpu_preferred" => Ok(DevicePolicy::GpuPreferred), - "cpu" | "cpu_only" => Ok(DevicePolicy::CpuOnly), - other => bail!("unsupported device policy: {other}"), - } -} - -/// Parse a `--gpu` CLI value into an optional `GpuSelection` for `LaunchRequest`. -fn parse_gpu_selection_arg(value: Option<&str>) -> Result> { - value - .map(|raw| GpuSelection::parse_cli_value(raw).map_err(anyhow::Error::msg)) - .transpose() -} - -/// Parse a `--gpu` CLI value into explicit device ordinals (empty for `auto`). -fn parse_gpu_indices_arg(value: Option<&str>) -> Result> { - Ok(rocm_engine_protocol::launch_gpu_indices( - parse_gpu_selection_arg(value)?.as_ref(), - )) -} - -fn sglang_command_from_python(python: &Path) -> Option { - let dir = python.parent()?; - candidate_command_names("sglang") - .into_iter() - .map(|name| dir.join(name)) - .find(|path| path.is_file()) -} - -fn find_command_on_path(name: &str) -> Option { - let path = std::env::var_os("PATH")?; - for dir in std::env::split_paths(&path) { - for candidate in candidate_command_names(name) { - let path = dir.join(candidate); - if path.is_file() { - return Some(path); - } - } - } - None -} - -fn resolve_command_path(command: &Path) -> Result { - if command.components().count() > 1 || command.is_absolute() { - if command.is_file() { - return Ok(command.to_path_buf()); - } - bail!( - "configured SGLang command is not a file: {}", - command.display() - ); - } - find_command_on_path(&command.display().to_string()).with_context(|| { - format!( - "configured SGLang command `{}` was not found on PATH", - command.display() - ) - }) -} - -fn candidate_command_names(name: &str) -> Vec { - if cfg!(windows) { - vec![ - format!("{name}.exe"), - format!("{name}.cmd"), - name.to_owned(), - ] - } else { - vec![name.to_owned()] - } -} - -fn probe_sglang_version(python: &Path) -> Result> { - let script = r#"import importlib.metadata, importlib.util, json -spec = importlib.util.find_spec("sglang") -version = None -if spec is not None: - try: - version = importlib.metadata.version("sglang") - except importlib.metadata.PackageNotFoundError: - version = "unknown" -print(json.dumps({"present": spec is not None, "version": version})) -"#; - let output = ProcessCommand::new(python) - .arg("-c") - .arg(script) - .output() - .with_context(|| format!("failed to probe SGLang with {}", python.display()))?; - if !output.status.success() { - bail!( - "SGLang probe failed: {}", - String::from_utf8_lossy(&output.stderr).trim() - ); - } - let value: Value = - serde_json::from_slice(&output.stdout).context("invalid SGLang probe JSON")?; - if value - .get("present") - .and_then(Value::as_bool) - .unwrap_or(false) - { - Ok(value - .get("version") - .and_then(Value::as_str) - .map(str::to_owned)) - } else { - bail!("Python environment does not contain the SGLang package") - } -} - -fn apply_therock_env(command: &mut ProcessCommand, runtime: &SglangRuntime) -> Result<()> { - let Some(root) = runtime.sdk_root.as_ref() else { - return Ok(()); - }; - let bin = runtime.sdk_bin.as_ref(); - command - .env("ROCM_SDK_ROOT", root) - .env("ROCM_PATH", root) - .env("ROCM_HOME", root) - .env("HIP_PATH", root) - .env("ROCM_CLI_THEROCK_RUNTIME_ID", &runtime.runtime_id); - if std::env::var_os("GPU_ARCHS").is_none() - && let Some(arch) = sglang_rocm_gpu_arch(runtime) - { - command - .env("GPU_ARCHS", &arch) - .env("AMDGPU_TARGET", &arch) - .env("PYTORCH_ROCM_ARCH", &arch); - } - if let Some(bin) = bin { - command.env("ROCM_CLI_THEROCK_SDK_BIN", bin).env( - "PATH", - prepend_path_entries(&runtime_bin_paths(runtime), std::env::var_os("PATH"))?, - ); - } else if !runtime.sdk_bin_paths.is_empty() { - command.env( - "PATH", - prepend_path_entries(&runtime_bin_paths(runtime), std::env::var_os("PATH"))?, - ); - } - if !cfg!(windows) { - command.env( - "LD_LIBRARY_PATH", - prepend_path_entries( - &therock_library_path_entries(runtime), - std::env::var_os("LD_LIBRARY_PATH"), - )?, - ); - } - Ok(()) -} - -fn sglang_rocm_gpu_arch(runtime: &SglangRuntime) -> Option { - detect_sglang_rocm_gpu_arch_from_sdk(runtime) - .or_else(|| sglang_rocm_gpu_arch_from_text(&runtime.runtime_id)) - .or_else(|| sglang_rocm_gpu_arch_from_text(&runtime.source)) -} - -fn detect_sglang_rocm_gpu_arch_from_sdk(runtime: &SglangRuntime) -> Option { - for bin_dir in runtime_bin_paths(runtime) { - let tool = bin_dir.join(if cfg!(windows) { - "rocm_agent_enumerator.exe" - } else { - "rocm_agent_enumerator" - }); - if !tool.is_file() { - continue; - } - let Ok(output) = ProcessCommand::new(&tool).output() else { - continue; - }; - if !output.status.success() { - continue; - } - let stdout = String::from_utf8_lossy(&output.stdout); - if let Some(arch) = sglang_rocm_gpu_arch_from_text(&stdout) { - return Some(arch); - } - let stderr = String::from_utf8_lossy(&output.stderr); - if let Some(arch) = sglang_rocm_gpu_arch_from_text(&stderr) { - return Some(arch); - } - } - None -} - -fn sglang_rocm_gpu_arch_from_text(text: &str) -> Option { - text.split(|ch: char| !ch.is_ascii_alphanumeric()) - .find_map(normalize_sglang_rocm_gpu_arch) -} - -fn normalize_sglang_rocm_gpu_arch(value: &str) -> Option { - let value = value.trim().to_ascii_lowercase(); - match value.as_str() { - "gfx90a" | "gfx940" | "gfx941" | "gfx942" | "gfx950" => Some(value), - value if value.starts_with("gfx942") || value.starts_with("gfx94") => { - Some("gfx942".to_owned()) - } - value if value.starts_with("gfx950") => Some("gfx950".to_owned()), - _ => None, - } -} - -fn runtime_bin_paths(runtime: &SglangRuntime) -> Vec { - let mut entries = Vec::new(); - if let Some(bin) = runtime.sdk_bin.as_ref() { - entries.push(bin.clone()); - } - entries.extend(runtime.sdk_bin_paths.iter().cloned()); - dedupe_paths(entries) -} - -fn therock_library_path_entries(runtime: &SglangRuntime) -> Vec { - let Some(root) = runtime.sdk_root.as_ref() else { - return dedupe_paths(runtime.sdk_library_paths.clone()); - }; - let mut entries = runtime.sdk_library_paths.clone(); - entries.extend([ - root.join("lib"), - root.join("lib64"), - root.join("lib").join("rocm_sysdeps").join("lib"), - ]); - if cfg!(target_os = "linux") { - let wsl_dxcore_lib = PathBuf::from("/usr/lib/wsl/lib"); - if wsl_dxcore_lib.is_dir() { - entries.push(wsl_dxcore_lib); - } - } - dedupe_paths(entries) -} - -fn dedupe_paths(entries: Vec) -> Vec { - let mut deduped = Vec::new(); - for entry in entries { - if !entry.as_os_str().is_empty() && !deduped.iter().any(|seen| seen == &entry) { - deduped.push(entry); - } - } - deduped -} - -fn prepend_path_entries(entries: &[PathBuf], current: Option) -> Result { - let mut parts = Vec::new(); - for entry in entries { - if !entry.as_os_str().is_empty() && !parts.iter().any(|part: &PathBuf| part == entry) { - parts.push(entry.clone()); - } - } - if let Some(current) = current { - for entry in std::env::split_paths(¤t) { - if !entry.as_os_str().is_empty() && !parts.iter().any(|part| part == &entry) { - parts.push(entry); - } - } - } - std::env::join_paths(parts).context("failed to compose runtime path") -} - -fn service_files(service_id: &str) -> Result { - let paths = AppPaths::discover()?; - Ok(ServiceFiles { - state_path: paths - .engine_state_dir(ENGINE_NAME) - .join(format!("{service_id}.json")), - log_path: paths - .engine_logs_dir(ENGINE_NAME) - .join(format!("{service_id}.log")), - }) -} - -fn write_running_state( - request: &ServeHttpRequest, - runtime: &SglangRuntime, - pid: u32, -) -> Result<()> { - write_state( - &request.state_path, - &json!({ - "service_id": request.service_id, - "engine": ENGINE_NAME, - "status": "running", - "pid": pid, - "model_ref": request.model_ref, - "host": request.host, - "port": request.port, - "endpoint_url": endpoint_url(&request.host, request.port), - "device_policy": "gpu_required", - "runtime_id": runtime.runtime_id, - "requested_runtime_id": request.runtime_id, - "env_id": request.env_id.as_deref().unwrap_or(runtime.env_id.as_str()), - "runtime_executable": runtime.command, - "server_pid": pid, - "engine_recipe": request.engine_recipe, - "engine_recipe_required_flags": engine_recipe_launch_args(request.engine_recipe.as_ref()), - "therock_runtime_env": therock_runtime_env_state(runtime), - "started_at_unix_ms": current_unix_millis() - }), - ) -} - -fn therock_runtime_env_state(runtime: &SglangRuntime) -> Option { - let root = runtime.sdk_root.as_ref()?; - Some(json!({ - "runtime_id": runtime.runtime_id, - "env_id": runtime.env_id, - "root": root.display().to_string(), - "bin": runtime.sdk_bin.as_ref().map(|path| path.display().to_string()), - "bin_paths": runtime_bin_paths(runtime) - .into_iter() - .map(|path| path.display().to_string()) - .collect::>(), - "library_paths": therock_library_path_entries(runtime) - .into_iter() - .map(|path| path.display().to_string()) - .collect::>(), - "source": runtime.source, - })) -} - -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(), Value::String(status.to_owned())); - object.insert( - "stopped_at_unix_ms".to_owned(), - Value::from(current_unix_millis() as u64), - ); - } - write_state(state_path, &state) -} - -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).context("failed to serialize SGLang state")?, - ) - .with_context(|| format!("failed to write {}", path.display())) -} - -fn endpoint_url(host: &str, port: u16) -> String { - format!("{}/v1", format_http_base_url(host, port)) -} - -fn endpoint_url_from_state(state: &Value) -> Option { - value_string(state, "endpoint_url").or_else(|| { - let host = value_string(state, "host")?; - let port = state.get("port")?.as_u64()?; - let port = u16::try_from(port).ok()?; - Some(endpoint_url(&host, port)) - }) -} - -fn query_loaded_model_endpoint(endpoint_url: &str, model_ref: Option<&str>) -> Result { - openai_models_endpoint_has_model( - endpoint_url, - model_ref, - Duration::from_millis(HEALTHCHECK_TIMEOUT_MS), - ) -} - -fn pid_from_state(state: &Value) -> Option { - state - .get("pid")? - .as_u64() - .and_then(|pid| pid.try_into().ok()) -} - -fn terminate_pid(pid: u32, _force: bool) -> bool { - rocm_core::terminate_process(pid).is_ok() -} - -fn tail_lines(path: &Path, limit: usize) -> Result> { - let text = - fs::read_to_string(path).with_context(|| format!("failed to read {}", path.display()))?; - Ok(tail_lines_from_text(&text, limit)) -} - -fn tail_lines_from_text(text: &str, limit: usize) -> Vec { - if limit == 0 { - return Vec::new(); - } - let mut lines = text - .lines() - .rev() - .take(limit) - .map(str::to_owned) - .collect::>(); - lines.reverse(); - lines -} - -fn value_string(value: &Value, key: &str) -> Option { - value.get(key)?.as_str().map(str::to_owned) -} - -fn stable_id_component(value: &str) -> String { - value - .chars() - .map(|ch| { - if ch.is_ascii_alphanumeric() || ch == '-' || ch == '_' { - ch.to_ascii_lowercase() - } else { - '-' - } - }) - .collect() -} - -fn runtime_lock_hash(runtime: &SglangRuntime) -> String { - let mut hasher = DefaultHasher::new(); - runtime.runtime_id.hash(&mut hasher); - runtime.command.hash(&mut hasher); - runtime.version.hash(&mut hasher); - format!("{:016x}", hasher.finish()) -} - -fn current_unix_millis() -> u128 { - SystemTime::now() - .duration_since(UNIX_EPOCH) - .unwrap_or_default() - .as_millis() -} - -const fn windows_unsupported_message() -> &'static str { - "SGLang ROCm serving is supported by rocm-cli only on Linux/WSL; native Windows SGLang is skipped. No CPU fallback is used." -} - -fn read_request() -> Result { - let mut buffer = String::new(); - std::io::stdin() - .read_to_string(&mut buffer) - .context("failed to read stdin for engine request")?; - serde_json::from_str(&buffer).context("failed to parse engine request envelope") -} - -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!(&mut handle)?; - Ok(()) -} - -#[cfg(test)] -mod tests { - use super::*; - - fn test_engine_recipe(engine: &str, contract_version: &str) -> EngineRecipeHint { - EngineRecipeHint { - contract_version: contract_version.to_owned(), - engine: engine.to_owned(), - required_flags: vec!["--reasoning-parser".to_owned(), "qwen3".to_owned()], - parser_settings: std::collections::BTreeMap::default(), - preferred_endpoint: None, - unsupported_combinations: Vec::new(), - notes: vec!["test recipe".to_owned()], - } - } - - #[test] - fn parse_gpu_args_map_to_indices() { - assert_eq!(parse_gpu_indices_arg(None).unwrap(), Vec::::new()); - assert_eq!( - parse_gpu_indices_arg(Some("auto")).unwrap(), - Vec::::new() - ); - assert_eq!(parse_gpu_indices_arg(Some("3")).unwrap(), vec![3]); - assert!(parse_gpu_selection_arg(Some("x")).is_err()); - assert!(parse_gpu_selection_arg(Some("1,3")).is_err()); - } - - #[test] - fn cpu_policy_is_rejected_without_fallback() { - let error = normalize_sglang_device_policy(Some(DevicePolicy::CpuOnly)) - .expect_err("SGLang CPU policy must fail"); - assert!(error.to_string().contains("no CPU fallback is used")); - } - - #[test] - fn gpu_preferred_resolves_to_gpu_required() -> Result<()> { - assert_eq!( - normalize_sglang_device_policy(Some(DevicePolicy::GpuPreferred))?, - DevicePolicy::GpuRequired - ); - Ok(()) - } - - #[test] - fn engine_recipe_launch_args_forward_required_flags() { - let hint = test_engine_recipe(ENGINE_NAME, ENGINE_RECIPE_CONTRACT_VERSION); - - assert_eq!( - engine_recipe_launch_args(Some(&hint)), - vec!["--reasoning-parser".to_owned(), "qwen3".to_owned()] - ); - } - - #[test] - fn server_args_support_command_and_python_module_launchers() { - assert_eq!( - sglang_server_args(SglangLauncher::Command, "qwen", "127.0.0.1", 30000), - vec![ - "serve", - "--model-path", - "qwen", - "--host", - "127.0.0.1", - "--port", - "30000", - "--attention-backend", - "triton" - ] - ); - assert_eq!( - sglang_server_args(SglangLauncher::PythonModule, "qwen", "0.0.0.0", 30001), - vec![ - "-m", - "sglang.launch_server", - "--model-path", - "qwen", - "--host", - "0.0.0.0", - "--port", - "30001", - "--attention-backend", - "triton" - ] - ); - } - - #[test] - fn managed_runtime_allows_sibling_command_without_shared_python_package() -> Result<()> { - let root = std::env::temp_dir().join(format!( - "rocm-sglang-runtime-test-{}", - current_unix_millis() - )); - fs::create_dir_all(&root)?; - let python = root.join(if cfg!(windows) { - "python.exe" - } else { - "python" - }); - let command = root.join( - candidate_command_names("sglang") - .into_iter() - .next() - .expect("candidate command name"), - ); - fs::write(&python, "")?; - fs::write(&command, "")?; - - let runtime = runtime_from_python( - python.clone(), - "therock-release:gfx120X-all", - "managed_runtime_manifest:test", - None, - None, - Vec::new(), - Vec::new(), - )?; - - assert_eq!(runtime.command, command); - assert_eq!(runtime.launcher, SglangLauncher::Command); - assert_eq!(runtime.python_executable.as_deref(), Some(python.as_path())); - assert_eq!(runtime.version, None); - - fs::remove_dir_all(root)?; - Ok(()) - } - - #[test] - fn endpoint_response_errors_without_service_state() { - let error = endpoint_response(EndpointRequest { - service_id: format!("missing-{}", current_unix_millis()), - }) - .expect_err("missing service state should not produce a default endpoint"); - - assert!(error.to_string().contains("service state not found")); - } - - #[test] - fn endpoint_url_falls_back_to_host_and_port() { - let state = json!({ - "host": "127.0.0.1", - "port": 12345 - }); - assert_eq!( - endpoint_url_from_state(&state), - Some("http://127.0.0.1:12345/v1".to_owned()) - ); - let ipv6_state = json!({ - "host": "::1", - "port": 12345 - }); - assert_eq!( - endpoint_url_from_state(&ipv6_state), - Some("http://[::1]:12345/v1".to_owned()) - ); - } - - #[test] - fn resolve_model_echoes_matching_engine_recipe() -> Result<()> { - let hint = test_engine_recipe(ENGINE_NAME, ENGINE_RECIPE_CONTRACT_VERSION); - let response = resolve_model_response(ResolveModelRequest { - model_ref: "Qwen/Qwen3.5-4B".to_owned(), - runtime_id: None, - device_policy: Some(DevicePolicy::GpuRequired), - recipe_override: None, - engine_recipe: Some(hint.clone()), - })?; - - assert_eq!(response.engine_recipe, Some(hint)); - Ok(()) - } - - #[test] - fn resolve_model_rejects_mismatched_engine_recipe() { - let error = resolve_model_response(ResolveModelRequest { - model_ref: "Qwen/Qwen3.5-4B".to_owned(), - runtime_id: None, - device_policy: Some(DevicePolicy::GpuRequired), - recipe_override: None, - engine_recipe: Some(test_engine_recipe( - "pytorch", - ENGINE_RECIPE_CONTRACT_VERSION, - )), - }) - .expect_err("mismatched engine recipe should fail"); - - assert!(error.to_string().contains("does not match adapter")); - } - - #[test] - fn resolve_model_rejects_unsupported_engine_recipe_contract() { - let error = resolve_model_response(ResolveModelRequest { - model_ref: "Qwen/Qwen3.5-4B".to_owned(), - runtime_id: None, - device_policy: Some(DevicePolicy::GpuRequired), - recipe_override: None, - engine_recipe: Some(test_engine_recipe(ENGINE_NAME, "999.0.0")), - }) - .expect_err("unsupported recipe contract should fail"); - - assert!(error.to_string().contains("unsupported")); - } - - #[test] - fn tail_lines_returns_suffix() -> Result<()> { - let path = std::env::temp_dir().join(format!( - "rocm-sglang-tail-{}-{}.log", - std::process::id(), - current_unix_millis() - )); - fs::write(&path, "a\nb\nc\n")?; - let lines = tail_lines(&path, 2)?; - fs::remove_file(path).ok(); - assert_eq!(lines, vec!["b".to_owned(), "c".to_owned()]); - Ok(()) - } - - #[test] - fn stdio_protocol_routes_all_methods_without_side_effects() { - let service_id = format!( - "missing-protocol-{}-{}", - std::process::id(), - current_unix_millis() - ); - let success_cases = [ - (EngineMethod::Detect, json!({})), - (EngineMethod::Capabilities, json!({})), - ( - EngineMethod::ResolveModel, - json!({ - "model_ref": "Qwen/Qwen3.5-4B", - "device_policy": "gpu_required" - }), - ), - ( - EngineMethod::Healthcheck, - json!({ - "service_id": service_id.as_str() - }), - ), - ( - EngineMethod::Logs, - json!({ - "service_id": service_id.as_str(), - "tail_lines": 4 - }), - ), - ( - EngineMethod::Stop, - json!({ - "service_id": service_id.as_str(), - "force": false - }), - ), - ]; - - for (method, payload) in success_cases { - let response = handle_envelope(EngineRequestEnvelope { method, payload }); - assert!( - response.ok, - "expected protocol method to return a typed success envelope: {:?}", - response.error - ); - } - - let endpoint = handle_envelope(EngineRequestEnvelope { - method: EngineMethod::Endpoint, - payload: json!({ - "service_id": service_id.as_str() - }), - }); - assert!(!endpoint.ok); - assert_eq!( - endpoint.error.as_ref().map(|error| error.code.as_str()), - Some("request_failed") - ); - - for method in [EngineMethod::Install, EngineMethod::Launch] { - let response = handle_envelope(EngineRequestEnvelope { - method, - payload: json!({}), - }); - assert!(!response.ok); - assert_eq!( - response.error.as_ref().map(|error| error.code.as_str()), - Some("invalid_payload") - ); - } - } - - #[test] - fn therock_library_path_entries_include_sysdeps_for_hip_apps() { - let root = PathBuf::from(if cfg!(windows) { - r"C:\rocm-sdk" - } else { - "/tmp/rocm-sdk" - }); - let runtime = SglangRuntime { - runtime_id: "therock-release:gfx120X-all".to_owned(), - env_id: "external-sglang-therock".to_owned(), - command: PathBuf::from("sglang"), - launcher: SglangLauncher::Command, - python_executable: None, - version: None, - source: "managed_runtime_manifest:test".to_owned(), - sdk_root: Some(root.clone()), - sdk_bin: Some(root.join("bin")), - sdk_bin_paths: vec![root.join("runtime").join("bin")], - sdk_library_paths: vec![root.join("runtime").join("lib")], - }; - let entries = therock_library_path_entries(&runtime); - assert!(entries.contains(&root.join("runtime").join("lib"))); - assert!(entries.contains(&root.join("lib"))); - assert!( - entries - .iter() - .any(|entry| entry.ends_with(Path::new("lib").join("rocm_sysdeps").join("lib"))) - ); - } - - #[test] - fn sglang_rocm_gpu_arch_maps_runtime_families_for_aiter() { - assert_eq!( - sglang_rocm_gpu_arch_from_text("therock-release:gfx94X-dcgpu"), - Some("gfx942".to_owned()) - ); - assert_eq!( - sglang_rocm_gpu_arch_from_text("gfx942:sramecc+:xnack-"), - Some("gfx942".to_owned()) - ); - assert_eq!( - sglang_rocm_gpu_arch_from_text("release-pip-gfx950-dcgpu"), - Some("gfx950".to_owned()) - ); - assert_eq!(sglang_rocm_gpu_arch_from_text("gfx120X-all"), None); - } - - #[test] - fn managed_env_reflects_managed_runtime_manifest_source() { - let runtime = SglangRuntime { - runtime_id: "therock-release:gfx120X-all".to_owned(), - env_id: "external-sglang-therock".to_owned(), - command: PathBuf::from(if cfg!(windows) { - r"C:\venv\Scripts\sglang.exe" - } else { - "/home/user/.venv/bin/sglang" - }), - launcher: SglangLauncher::Command, - python_executable: None, - version: Some("0.5.12".to_owned()), - source: "managed_runtime_manifest:sglang-source-pip-gfx120x-all".to_owned(), - sdk_root: Some(PathBuf::from(if cfg!(windows) { - r"C:\rocm-sdk" - } else { - "/home/user/.venv/lib/python/site-packages/_rocm_sdk_devel" - })), - sdk_bin: Some(PathBuf::from(if cfg!(windows) { - r"C:\rocm-sdk\bin" - } else { - "/home/user/.venv/lib/python/site-packages/_rocm_sdk_devel/bin" - })), - sdk_bin_paths: Vec::new(), - sdk_library_paths: Vec::new(), - }; - - assert!(runtime_is_managed(&runtime)); - assert!( - sglang_runtime_warnings(&runtime) - .join("\n") - .contains("managed TheRock runtime") - ); - - let external = SglangRuntime { - source: "environment command".to_owned(), - sdk_root: None, - sdk_bin: None, - sdk_bin_paths: Vec::new(), - sdk_library_paths: Vec::new(), - ..runtime - }; - assert!(!runtime_is_managed(&external)); - assert!( - sglang_runtime_warnings(&external) - .join("\n") - .contains("external SGLang runtime") - ); - } - - #[test] - fn running_state_records_managed_therock_env_for_gpu_verification() -> Result<()> { - let state_path = std::env::temp_dir().join(format!( - "rocm-sglang-state-{}-{}.json", - std::process::id(), - current_unix_millis() - )); - let request = ServeHttpRequest { - service_id: "sglang-test".to_owned(), - model_ref: "Qwen/Qwen2.5-0.5B-Instruct".to_owned(), - host: "127.0.0.1".to_owned(), - port: 11435, - device_policy: DevicePolicy::GpuRequired, - gpu_indices: Vec::new(), - runtime_id: Some("runtime-key-gfx120x".to_owned()), - env_id: None, - state_path: state_path.clone(), - engine_recipe: None, - }; - let runtime = SglangRuntime { - runtime_id: "therock-release:gfx120X-all".to_owned(), - env_id: "external-sglang-therock".to_owned(), - command: PathBuf::from(if cfg!(windows) { - r"C:\venv\Scripts\sglang.exe" - } else { - "/home/user/.venv/bin/sglang" - }), - launcher: SglangLauncher::Command, - python_executable: None, - version: Some("0.5.12".to_owned()), - source: "managed_runtime_manifest:test".to_owned(), - sdk_root: Some(PathBuf::from(if cfg!(windows) { - r"C:\rocm-sdk" - } else { - "/home/user/.venv/lib/python/site-packages/_rocm_sdk_devel" - })), - sdk_bin: Some(PathBuf::from(if cfg!(windows) { - r"C:\rocm-sdk\bin" - } else { - "/home/user/.venv/lib/python/site-packages/_rocm_sdk_devel/bin" - })), - sdk_bin_paths: vec![PathBuf::from(if cfg!(windows) { - r"C:\rocm-sdk\extra-bin" - } else { - "/home/user/.venv/lib/python/site-packages/_rocm_sdk_libraries/bin" - })], - sdk_library_paths: vec![PathBuf::from(if cfg!(windows) { - r"C:\rocm-sdk\extra-lib" - } else { - "/home/user/.venv/lib/python/site-packages/_rocm_sdk_libraries/lib" - })], - }; - - write_running_state(&request, &runtime, 12345)?; - let state = read_service_state(&state_path)?; - fs::remove_file(&state_path).ok(); - - assert_eq!(state.get("server_pid").and_then(Value::as_u64), Some(12345)); - assert_eq!( - state.get("runtime_id").and_then(Value::as_str), - Some("therock-release:gfx120X-all") - ); - assert_eq!( - state.get("requested_runtime_id").and_then(Value::as_str), - Some("runtime-key-gfx120x") - ); - let runtime_env = state - .get("therock_runtime_env") - .expect("runtime env should be recorded"); - assert_eq!( - runtime_env.get("runtime_id").and_then(Value::as_str), - Some("therock-release:gfx120X-all") - ); - assert!( - runtime_env - .get("root") - .and_then(Value::as_str) - .is_some_and(|root| root.contains("rocm")) - ); - assert!( - runtime_env - .get("bin_paths") - .and_then(Value::as_array) - .is_some_and(|paths| paths.len() >= 2) - ); - assert!( - runtime_env - .get("library_paths") - .and_then(Value::as_array) - .is_some_and(|paths| !paths.is_empty()) - ); - Ok(()) - } -} diff --git a/engines/sglang/src/main.rs b/engines/sglang/src/main.rs deleted file mode 100644 index a976c29c..00000000 --- a/engines/sglang/src/main.rs +++ /dev/null @@ -1,7 +0,0 @@ -// Copyright © Advanced Micro Devices, Inc., or its affiliates. -// -// SPDX-License-Identifier: MIT - -fn main() -> anyhow::Result<()> { - rocm_engine_sglang::run_cli() -} diff --git a/engines/vllm/src/lib.rs b/engines/vllm/src/lib.rs index 9057873a..4b9d704c 100644 --- a/engines/vllm/src/lib.rs +++ b/engines/vllm/src/lib.rs @@ -1854,7 +1854,7 @@ mod tests { device_policy: Some(DevicePolicy::GpuRequired), recipe_override: None, engine_recipe: Some(test_engine_recipe( - "pytorch", + "lemonade", ENGINE_RECIPE_CONTRACT_VERSION, )), }) diff --git a/install.ps1 b/install.ps1 index 2a2a19e3..f70d5150 100644 --- a/install.ps1 +++ b/install.ps1 @@ -181,7 +181,7 @@ function Write-MinimalConfigIfMissing { New-Item -ItemType Directory -Force -Path $configDir | Out-Null $json = @' { - "default_engine": "pytorch", + "default_engine": "lemonade", "telemetry": { "mode": "local" }, diff --git a/install.sh b/install.sh index 47c27b0f..2358ab49 100755 --- a/install.sh +++ b/install.sh @@ -116,7 +116,7 @@ write_minimal_config_if_missing() { config_tmp="${tmp_dir}/config.json" cat > "${config_tmp}" <<'JSON' { - "default_engine": "pytorch", + "default_engine": "lemonade", "telemetry": { "mode": "local" }, diff --git a/scripts/acceptance-install-upgrade-tui-uninstall.ps1 b/scripts/acceptance-install-upgrade-tui-uninstall.ps1 index 3084ccac..dac9dd26 100644 --- a/scripts/acceptance-install-upgrade-tui-uninstall.ps1 +++ b/scripts/acceptance-install-upgrade-tui-uninstall.ps1 @@ -295,7 +295,7 @@ try { Assert-File (Join-Path $PemInstallDir "rocm.exe") Assert-File (Join-Path $PemInstallDir ".rocm-cli-manifest") Assert-File $ConfigFile - if (-not (Select-String -LiteralPath $ConfigFile -Pattern '"default_engine"\s*:\s*"pytorch"' -Quiet)) { + if (-not (Select-String -LiteralPath $ConfigFile -Pattern '"default_engine"\s*:\s*"lemonade"' -Quiet)) { Fail "installer did not seed minimal config with the expected default engine" } @@ -413,12 +413,12 @@ try { Assert-Missing (Join-Path $missingSignatureInstallDir "rocm.exe") Assert-Missing (Join-Path $missingSignatureInstallDir ".rocm-cli-manifest") - Set-Content -LiteralPath $ConfigFile -Value '{"default_engine":"llama.cpp"}' -Encoding utf8 + Set-Content -LiteralPath $ConfigFile -Value '{"default_engine":"vllm"}' -Encoding utf8 Invoke-Checked "acceptance: first install" $psExe $installArgs $InstallLog1 if (-not (Select-String -LiteralPath $InstallLog1 -Pattern "signature verified" -Quiet)) { Fail "installer did not report signature verification" } - if (-not (Select-String -LiteralPath $ConfigFile -Pattern '"default_engine"\s*:\s*"llama.cpp"' -Quiet)) { + if (-not (Select-String -LiteralPath $ConfigFile -Pattern '"default_engine"\s*:\s*"vllm"' -Quiet)) { Fail "installer overwrote an existing config file" } Assert-File (Join-Path $InstallDir "rocm.exe") diff --git a/scripts/acceptance-install-upgrade-tui-uninstall.sh b/scripts/acceptance-install-upgrade-tui-uninstall.sh index 228ea485..e83631d6 100755 --- a/scripts/acceptance-install-upgrade-tui-uninstall.sh +++ b/scripts/acceptance-install-upgrade-tui-uninstall.sh @@ -256,20 +256,20 @@ assert_file "${INSTALL_DIR}/rocm" assert_file "${INSTALL_DIR}/rocmd" assert_file "${INSTALL_DIR}/.rocm-cli-manifest" assert_file "${INSTALL_CONFIG_FILE}" -grep -q '"default_engine"[[:space:]]*:[[:space:]]*"pytorch"' "${INSTALL_CONFIG_FILE}" \ +grep -q '"default_engine"[[:space:]]*:[[:space:]]*"lemonade"' "${INSTALL_CONFIG_FILE}" \ || fail "installer did not seed minimal config with the expected default engine" assert_file "${BASHRC_FILE}" grep -F "${INSTALL_DIR}" "${BASHRC_FILE}" >/dev/null \ || fail "installer did not add install dir to the shell profile" echo "acceptance: simulate stale prior install entry and reinstall" -printf '%s\n' '{"default_engine":"llama.cpp"}' > "${INSTALL_CONFIG_FILE}" +printf '%s\n' '{"default_engine":"vllm"}' > "${INSTALL_CONFIG_FILE}" echo "stale" > "${INSTALL_DIR}/rocm-engine-stale" echo "${INSTALL_DIR}/rocm-engine-stale" >> "${INSTALL_DIR}/.rocm-cli-manifest" run_installer | tee "${INSTALL_LOG_2}" assert_missing "${INSTALL_DIR}/rocm-engine-stale" assert_file "${INSTALL_DIR}/.rocm-cli-manifest" -grep -q '"default_engine"[[:space:]]*:[[:space:]]*"llama.cpp"' "${INSTALL_CONFIG_FILE}" \ +grep -q '"default_engine"[[:space:]]*:[[:space:]]*"vllm"' "${INSTALL_CONFIG_FILE}" \ || fail "installer overwrote an existing config file" grep -q "removing previous rocm-cli install" "${INSTALL_LOG_2}" \ || fail "installer did not report removal of previous install" @@ -284,7 +284,7 @@ env \ XDG_CONFIG_HOME="${XDG_CONFIG_HOME}" \ XDG_DATA_HOME="${XDG_DATA_HOME}" \ XDG_CACHE_HOME="${XDG_CACHE_HOME}" \ - "${INSTALL_DIR}/rocm" config set-default-engine pytorch >/dev/null + "${INSTALL_DIR}/rocm" config set-default-engine vllm >/dev/null tui_command="$( printf '%q ' \ @@ -326,7 +326,7 @@ if [[ "${tui_status}" -ne 0 ]]; then fi assert_file "${CONFIG_FILE}" -grep -q '"default_engine"[[:space:]]*:[[:space:]]*"pytorch"' "${CONFIG_FILE}" \ +grep -q '"default_engine"[[:space:]]*:[[:space:]]*"vllm"' "${CONFIG_FILE}" \ || fail "config smoke did not persist the expected default engine" assert_file "${TUI_LOG}" [[ -s "${TUI_LOG}" ]] || fail "TUI smoke log was empty" diff --git a/scripts/atom_therock_gpu_test.py b/scripts/atom_therock_gpu_test.py deleted file mode 100644 index dcb48570..00000000 --- a/scripts/atom_therock_gpu_test.py +++ /dev/null @@ -1,481 +0,0 @@ -#!/usr/bin/env python3 -# Copyright © Advanced Micro Devices, Inc., or its affiliates. -# -# SPDX-License-Identifier: MIT - -"""End-to-end ATOM ROCm GPU smoke test for rocm-cli managed TheRock runtimes.""" - -from __future__ import annotations - -import argparse -import json -import os -import platform -import subprocess -import tempfile -import time -from pathlib import Path -from typing import Any - -from vllm_therock_gpu_test import ( - DEFAULT_MATH_MODULES, - completion_text, - first_model_id, - get_json, - http_request, - load_runtime_registry_manifests, - post_json, - resolve_path, - resolve_runtime_selector, - restore_env, - rocm_cli_state_paths, - run_json, - stop_service, - verify_loaded_modules, - verify_managed_env, - write_config, - write_runtime_manifest, -) - - -def main() -> int: - args = parse_args() - if args.self_test: - return run_self_test() - - if platform.system() == "Windows": - raise SystemExit( - "ATOM GPU acceptance is Linux/WSL only; no CPU fallback is allowed" - ) - - repo_root = Path(__file__).resolve().parents[1] - engine = resolve_path(args.engine, repo_root) - runtime_id = resolve_runtime_id(args.runtime_id) - math_modules = args.math_module or DEFAULT_MATH_MODULES - - if not engine.is_file(): - raise SystemExit(f"engine binary not found: {engine}") - reject_external_runtime_env() - - env = os.environ.copy() - detect = run_json([str(engine), "detect"], env=env, timeout=args.timeout) - assert_atom_gpu_detected(detect) - - capabilities = run_json( - [str(engine), "capabilities"], env=env, timeout=args.timeout - ) - if capabilities.get("cpu"): - raise RuntimeError("ATOM capabilities unexpectedly report CPU support") - if not capabilities.get("rocm_gpu"): - raise RuntimeError("ATOM capabilities did not report ROCm GPU support") - if not capabilities.get("openai_compatible"): - raise RuntimeError("ATOM capabilities did not report OpenAI-compatible serving") - assert_atom_cpu_policy_rejected(engine, args.model, env=env, timeout=args.timeout) - - process, state_path, log_path = start_atom( - engine=engine, - model=args.model, - runtime_id=runtime_id, - args=args, - env=env, - repo_root=repo_root, - ) - - try: - health = wait_atom_health(args.host, args.port, args.timeout) - models = get_json(args.host, args.port, "/v1/models", timeout=args.timeout) - served_model = first_model_id(models) - completion = post_json( - args.host, - args.port, - "/v1/completions", - { - "model": served_model, - "prompt": args.prompt, - "max_tokens": args.max_tokens, - "temperature": 0, - }, - timeout=args.timeout, - ) - state = json.loads(state_path.read_text(encoding="utf-8")) - verify_atom_state(state) - env_values = verify_managed_env(state) - module_paths = verify_loaded_modules(state, math_modules) - text = completion_text(completion) - if not text.strip(): - raise RuntimeError("ATOM completion returned empty text") - - summary = { - "ok": True, - "launch_mode": args.launch_mode, - "service_id": args.service_id, - "model": args.model, - "served_model": served_model, - "runtime_id": runtime_id, - "health": health, - "completion_text": text, - "state_path": str(state_path), - "log_path": str(log_path), - "server_pid": state.get("server_pid") or state.get("pid"), - "therock_runtime_env": state.get("therock_runtime_env"), - "verified_env": env_values, - "verified_modules": module_paths, - } - print(json.dumps(summary, indent=2)) - finally: - if not args.keep_running: - stop_service(process, state_path) - - return 0 - - -def parse_args() -> argparse.Namespace: - default_engine = ( - Path("target/debug/rocm-engine-atom.exe") - if platform.system() == "Windows" - else Path("target/debug/rocm-engine-atom") - ) - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--engine", default=str(default_engine)) - parser.add_argument( - "--model", - default=os.environ.get("ROCM_CLI_ATOM_TEST_MODEL", "Qwen/Qwen3-0.6B"), - help="small Hugging Face model id or local model path served by ATOM", - ) - parser.add_argument( - "--runtime-id", - help=( - "managed TheRock runtime key or unambiguous runtime id; defaults to " - "the active rocm-cli runtime" - ), - ) - parser.add_argument("--service-id", default="atom-gpu-e2e") - parser.add_argument("--host", default="127.0.0.1") - parser.add_argument("--port", type=int, default=11443) - parser.add_argument("--timeout", type=int, default=300) - parser.add_argument("--prompt", default="Once upon a time") - parser.add_argument("--max-tokens", type=int, default=8) - parser.add_argument( - "--launch-mode", - choices=("serve-http", "launch"), - default="serve-http", - help="serve-http runs the wrapper directly; launch verifies background launch", - ) - parser.add_argument( - "--math-module", - action="append", - help=( - "math-library basename where at least one must load from the managed " - "TheRock SDK wheel directories; defaults to libhipblas/libhipblaslt/librocblas" - ), - ) - parser.add_argument("--keep-running", action="store_true") - parser.add_argument( - "--self-test", - action="store_true", - help="run offline selection checks for this script and exit", - ) - return parser.parse_args() - - -def reject_external_runtime_env() -> None: - blocked = [ - name - for name in [ - "ROCM_CLI_ATOM_COMMAND", - "ATOM_COMMAND", - "ROCM_CLI_ATOM_PYTHON", - "ATOM_PYTHON", - ] - if os.environ.get(name) - ] - if blocked: - raise RuntimeError( - "ATOM GPU acceptance requires discovery through a rocm-cli managed " - f"TheRock runtime manifest; unset external runtime overrides: {', '.join(blocked)}" - ) - - -def resolve_runtime_id(explicit: str | None) -> str: - config_path, registry_dir = rocm_cli_state_paths() - manifests = load_runtime_registry_manifests(registry_dir) - if explicit: - return resolve_runtime_selector(explicit, manifests, "explicit --runtime-id") - if config_path.is_file(): - config = json.loads(config_path.read_text(encoding="utf-8")) - active_key = config.get("active_runtime_key") - if isinstance(active_key, str) and active_key.strip(): - return resolve_runtime_selector( - active_key.strip(), - manifests, - f"active_runtime_key in {config_path}", - exact_key_only=True, - ) - runtime_id = config.get("default_runtime_id") - if isinstance(runtime_id, str) and runtime_id.strip(): - return resolve_runtime_selector( - runtime_id.strip(), - manifests, - f"default_runtime_id in {config_path}", - ) - raise RuntimeError( - "no active managed TheRock runtime was found; run " - "`rocm runtimes activate ` or pass --runtime-id " - ) - - -def assert_atom_gpu_detected(detect: dict[str, Any]) -> None: - devices = detect.get("available_devices", []) - rocm_gpu = next((item for item in devices if item.get("kind") == "rocm_gpu"), None) - if not detect.get("installed") or not rocm_gpu or not rocm_gpu.get("available"): - raise RuntimeError( - "ATOM ROCm GPU runtime was not detected through a managed TheRock " - "runtime; no CPU fallback is allowed:\n" + json.dumps(detect, indent=2) - ) - if not detect.get("managed_env"): - raise RuntimeError( - "ATOM acceptance requires a rocm-cli managed TheRock runtime manifest" - ) - - -def assert_atom_cpu_policy_rejected( - engine: Path, - model: str, - *, - env: dict[str, str], - timeout: int, -) -> None: - completed = subprocess.run( - [str(engine), "resolve-model", model, "--device-policy", "cpu_only"], - env=env, - text=True, - capture_output=True, - timeout=timeout, - check=False, - ) - output = completed.stdout + completed.stderr - if completed.returncode == 0: - raise RuntimeError( - "ATOM accepted cpu_only during GPU acceptance; no fallback is allowed" - ) - if "no CPU fallback is used" not in output: - raise RuntimeError( - "ATOM cpu_only rejection did not include the no-fallback explanation:\n" - + output - ) - - -def verify_atom_state(state: dict[str, Any]) -> None: - if state.get("device_policy") != "gpu_required": - raise RuntimeError(f"ATOM service state did not record gpu_required: {state!r}") - if not (state.get("server_pid") or state.get("pid")): - raise RuntimeError("ATOM service state is missing server_pid/pid") - runtime_env = state.get("therock_runtime_env") - if not isinstance(runtime_env, dict): - raise RuntimeError("ATOM service state is missing therock_runtime_env") - source = runtime_env.get("source") - if not isinstance(source, str) or not source.startswith("managed_runtime_manifest"): - raise RuntimeError( - "ATOM service did not launch from a managed TheRock runtime manifest; " - f"source={source!r}" - ) - - -def wait_atom_health(host: str, port: int, timeout: int) -> dict[str, Any]: - deadline = time.monotonic() + timeout - last_error: Exception | None = None - while time.monotonic() < deadline: - try: - status, body = http_request(host, port, "GET", "/health", None, timeout=3) - if status < 400: - return { - "status_code": status, - "body": body.decode("utf-8", errors="replace"), - } - except Exception as exc: - last_error = exc - time.sleep(0.5) - raise RuntimeError(f"ATOM server did not become healthy: {last_error}") - - -def start_atom( - *, - engine: Path, - model: str, - runtime_id: str, - args: argparse.Namespace, - env: dict[str, str], - repo_root: Path, -) -> tuple[subprocess.Popen[bytes] | None, Path, Path]: - if args.launch_mode == "launch": - return start_launch(engine, model, runtime_id, args, env) - - state_path = repo_root / "target" / "test-state" / f"{args.service_id}.json" - log_path = repo_root / "target" / "test-logs" / f"{args.service_id}.log" - command = [ - str(engine), - "serve-http", - args.service_id, - model, - "--host", - args.host, - "--port", - str(args.port), - "--device-policy", - "gpu_required", - "--runtime-id", - runtime_id, - "--state-path", - str(state_path), - ] - state_path.parent.mkdir(parents=True, exist_ok=True) - log_path.parent.mkdir(parents=True, exist_ok=True) - if state_path.exists(): - state_path.unlink() - log_file = log_path.open("wb") - try: - process = subprocess.Popen( - command, - env=env, - stdin=subprocess.DEVNULL, - stdout=log_file, - stderr=subprocess.STDOUT, - ) - finally: - log_file.close() - return process, state_path, log_path - - -def start_launch( - engine: Path, - model: str, - runtime_id: str, - args: argparse.Namespace, - env: dict[str, str], -) -> tuple[None, Path, Path]: - command = [ - str(engine), - "launch", - args.service_id, - model, - "--host", - args.host, - "--port", - str(args.port), - "--device-policy", - "gpu_required", - "--runtime-id", - runtime_id, - ] - started = time.monotonic() - response = run_json(command, env=env, timeout=args.timeout) - elapsed = time.monotonic() - started - if elapsed > 10: - raise RuntimeError( - f"launch took {elapsed:.1f}s; background launch should return promptly" - ) - return None, Path(response["state_path"]), Path(response["log_path"]) - - -def run_self_test() -> int: - scratch_root = ( - Path(__file__).resolve().parents[1] / ".rocm-work" / "script-self-tests" - ) - scratch_root.mkdir(parents=True, exist_ok=True) - with tempfile.TemporaryDirectory(prefix="atom-", dir=scratch_root) as temp: - root = Path(temp) - config_dir = root / "config" - data_dir = root / "data" - registry_dir = data_dir / "runtimes" / "registry" - registry_dir.mkdir(parents=True) - config_dir.mkdir(parents=True) - write_runtime_manifest( - registry_dir, "runtime-old", "therock-release:gfx120X-all" - ) - write_runtime_manifest( - registry_dir, "runtime-new", "therock-release:gfx120X-all" - ) - write_runtime_manifest(registry_dir, "runtime-other", "therock-release:gfx1151") - write_config( - config_dir, - { - "active_runtime_key": "runtime-old", - "default_runtime_id": "therock-release:gfx120X-all", - }, - ) - - old_env = { - "ROCM_CLI_CONFIG_DIR": os.environ.get("ROCM_CLI_CONFIG_DIR"), - "ROCM_CLI_DATA_DIR": os.environ.get("ROCM_CLI_DATA_DIR"), - "ROCM_CLI_ATOM_COMMAND": os.environ.get("ROCM_CLI_ATOM_COMMAND"), - "ATOM_COMMAND": os.environ.get("ATOM_COMMAND"), - "ROCM_CLI_ATOM_PYTHON": os.environ.get("ROCM_CLI_ATOM_PYTHON"), - "ATOM_PYTHON": os.environ.get("ATOM_PYTHON"), - } - try: - os.environ["ROCM_CLI_CONFIG_DIR"] = str(config_dir) - os.environ["ROCM_CLI_DATA_DIR"] = str(data_dir) - for key in [ - "ROCM_CLI_ATOM_COMMAND", - "ATOM_COMMAND", - "ROCM_CLI_ATOM_PYTHON", - "ATOM_PYTHON", - ]: - os.environ.pop(key, None) - - assert resolve_runtime_id(None) == "runtime-old" - assert resolve_runtime_id("runtime-new") == "runtime-new" - assert resolve_runtime_id("therock-release:gfx1151") == "runtime-other" - - os.environ["ROCM_CLI_ATOM_COMMAND"] = "/tmp/not-used-atom" - try: - reject_external_runtime_env() - except RuntimeError as exc: - assert "managed TheRock runtime manifest" in str(exc) - else: - raise AssertionError("external ATOM override did not fail") - os.environ.pop("ROCM_CLI_ATOM_COMMAND", None) - - write_config( - config_dir, {"default_runtime_id": "therock-release:gfx120X-all"} - ) - try: - resolve_runtime_id(None) - except RuntimeError as exc: - message = str(exc) - assert "runtime-old" in message and "runtime-new" in message - else: - raise AssertionError("ambiguous default_runtime_id did not fail") - - write_config(config_dir, {"active_runtime_key": "missing-runtime"}) - try: - resolve_runtime_id(None) - except RuntimeError as exc: - assert "missing-runtime" in str(exc) - else: - raise AssertionError("missing active runtime key did not fail") - finally: - restore_env(old_env) - - detect = { - "installed": True, - "managed_env": True, - "available_devices": [{"kind": "rocm_gpu", "available": True}], - } - assert_atom_gpu_detected(detect) - verify_atom_state( - { - "server_pid": 123, - "device_policy": "gpu_required", - "therock_runtime_env": { - "runtime_id": "therock-release:gfx120X-all", - "source": "managed_runtime_manifest:runtime-old", - }, - } - ) - print("ATOM GPU script self-test passed") - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/scripts/build_single_exe_release.py b/scripts/build_single_exe_release.py index 01980912..286acbe7 100644 --- a/scripts/build_single_exe_release.py +++ b/scripts/build_single_exe_release.py @@ -30,12 +30,8 @@ WINDOWS_BINARIES = [ "rocm.exe", "rocmd.exe", - "rocm-engine-pytorch.exe", - "rocm-engine-llama-cpp.exe", "rocm-engine-lemonade.exe", - "rocm-engine-atom.exe", "rocm-engine-vllm.exe", - "rocm-engine-sglang.exe", ] LINUX_BINARIES = [name[:-4] for name in WINDOWS_BINARIES] PLATFORMS = {"windows-amd64", "linux-amd64"} diff --git a/scripts/llama_cpp_therock_gpu_test.py b/scripts/llama_cpp_therock_gpu_test.py deleted file mode 100644 index fa24b4b2..00000000 --- a/scripts/llama_cpp_therock_gpu_test.py +++ /dev/null @@ -1,689 +0,0 @@ -#!/usr/bin/env python3 -# Copyright © Advanced Micro Devices, Inc., or its affiliates. -# -# SPDX-License-Identifier: MIT - -"""End-to-end llama.cpp GPU smoke test for rocm-cli managed TheRock runtimes.""" - -from __future__ import annotations - -import argparse -import contextlib -import http.client -import json -import os -import platform -import subprocess -import tempfile -import time -import urllib.request -from pathlib import Path -from typing import Any - -DEFAULT_MODEL_URL = ( - "https://huggingface.co/ggml-org/tiny-llamas/resolve/main/stories260K.gguf" -) -MODULES_TO_VERIFY = { - "amdhip64_7.dll", - "ggml-hip.dll", - "hipblas.dll", - "libhipblaslt.dll", - "rocblas.dll", -} -LINUX_MODULES_TO_VERIFY = { - "libamdhip64": "libamdhip64", - "libhipblas": "libhipblas", - "libhipblaslt": "libhipblaslt", - "librocblas": "librocblas", -} - - -def main() -> int: - args = parse_args() - if args.self_test: - return run_self_test() - - repo_root = Path(__file__).resolve().parents[1] - engine = resolve_path(args.engine, repo_root) - llama_server = resolve_path(args.llama_server, repo_root) - model_path = resolve_path(args.model_path, repo_root) - runtime_id = resolve_runtime_id(args.runtime_id) - - if not engine.is_file(): - raise SystemExit(f"engine binary not found: {engine}") - if not llama_server.is_file(): - raise SystemExit( - f"HIP llama-server not found: {llama_server}; no CPU fallback is allowed" - ) - ensure_model(model_path, args.model_url) - - env = os.environ.copy() - env["ROCM_CLI_LLAMA_CPP_SERVER"] = str(llama_server) - - detect = run_json([str(engine), "detect"], env=env, timeout=args.timeout) - assert_rocm_gpu_detected(detect) - - process, state_path, log_path = start_serve_http( - engine=engine, - model_path=model_path, - runtime_id=runtime_id, - args=args, - env=env, - repo_root=repo_root, - ) - - try: - health = wait_health(args.host, args.port, args.timeout) - completion = post_json( - args.host, - args.port, - "/v1/completions", - { - "model": model_path.name, - "prompt": "Once upon a time", - "max_tokens": 8, - "temperature": 0, - }, - timeout=args.timeout, - ) - state = json.loads(state_path.read_text(encoding="utf-8")) - log_text = log_path.read_text(encoding="utf-8", errors="replace") - assert_gpu_log(log_text) - module_paths = verify_loaded_modules(state) - - summary = { - "ok": True, - "launch_mode": args.launch_mode, - "service_id": args.service_id, - "health": health, - "completion_text": completion["choices"][0]["text"], - "runtime_id": runtime_id, - "state_path": str(state_path), - "log_path": str(log_path), - "launcher_pid": process.pid if process else state.get("pid"), - "staged_runtime_dir": state.get("staged_runtime_dir"), - "verified_modules": module_paths, - } - print(json.dumps(summary, indent=2)) - finally: - if not args.keep_running: - stop_service(process, state_path) - - return 0 - - -def parse_args() -> argparse.Namespace: - default_engine = cargo_binary_path("debug", "rocm-engine-llama-cpp") - default_server = ( - Path("target/llama.cpp-build-hip/bin/llama-server.exe") - if platform.system() == "Windows" - else Path("target/llama.cpp-build-hip/bin/llama-server") - ) - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--engine", default=str(default_engine)) - parser.add_argument( - "--llama-server", - default=os.environ.get("ROCM_CLI_LLAMA_CPP_SERVER", str(default_server)), - ) - parser.add_argument("--model-path", default="target/models/stories260K.gguf") - parser.add_argument("--model-url", default=DEFAULT_MODEL_URL) - parser.add_argument( - "--runtime-id", - help=( - "managed TheRock runtime key or unambiguous runtime id; defaults to the " - "active rocm-cli runtime" - ), - ) - parser.add_argument("--service-id", default="llama-gpu-e2e") - parser.add_argument("--host", default="127.0.0.1") - parser.add_argument("--port", type=int, default=11439) - parser.add_argument("--timeout", type=int, default=60) - parser.add_argument( - "--launch-mode", - choices=("serve-http", "launch"), - default="serve-http", - help="serve-http runs the wrapper directly; launch verifies captured-output background launch", - ) - parser.add_argument("--keep-running", action="store_true") - parser.add_argument( - "--self-test", - action="store_true", - help="run offline selection checks for this script and exit", - ) - return parser.parse_args() - - -def resolve_path(value: str, repo_root: Path) -> Path: - path = Path(value).expanduser() - if not path.is_absolute(): - path = repo_root / path - return path.resolve() - - -def cargo_binary_path(profile: str, name: str) -> Path: - target_root = Path(os.environ.get("CARGO_TARGET_DIR", "target")).expanduser() - return target_root / profile / exe_name(name) - - -def exe_name(name: str) -> str: - return f"{name}.exe" if platform.system() == "Windows" else name - - -def resolve_runtime_id(explicit: str | None) -> str: - config_path, registry_dir = rocm_cli_state_paths() - manifests = load_runtime_registry_manifests(registry_dir) - if explicit: - return resolve_runtime_selector(explicit, manifests, "explicit --runtime-id") - if config_path.is_file(): - config = json.loads(config_path.read_text(encoding="utf-8")) - active_key = config.get("active_runtime_key") - if isinstance(active_key, str) and active_key.strip(): - return resolve_runtime_selector( - active_key.strip(), - manifests, - f"active_runtime_key in {config_path}", - exact_key_only=True, - ) - runtime_id = config.get("default_runtime_id") - if isinstance(runtime_id, str) and runtime_id.strip(): - return resolve_runtime_selector( - runtime_id.strip(), - manifests, - f"default_runtime_id in {config_path}", - ) - raise RuntimeError( - "no active managed TheRock runtime was found; run " - "`rocm runtimes activate ` or pass --runtime-id " - ) - - -def load_runtime_registry_manifests(registry_dir: Path) -> list[dict[str, str]]: - manifests: list[dict[str, str]] = [] - if not registry_dir.is_dir(): - return manifests - for path in registry_dir.glob("*.json"): - try: - manifest = json.loads(path.read_text(encoding="utf-8")) - except json.JSONDecodeError: - continue - runtime_key = manifest.get("runtime_key") - runtime_id = manifest.get("runtime_id") - if not isinstance(runtime_key, str) or not runtime_key.strip(): - runtime_key = path.stem - if not isinstance(runtime_id, str) or not runtime_id.strip(): - continue - manifests.append( - { - "runtime_key": runtime_key.strip(), - "runtime_id": runtime_id.strip(), - } - ) - return manifests - - -def resolve_runtime_selector( - selector: str, - manifests: list[dict[str, str]], - source: str, - *, - exact_key_only: bool = False, -) -> str: - exact_key_matches = [ - manifest - for manifest in manifests - if manifest["runtime_key"].lower() == selector.lower() - ] - if len(exact_key_matches) == 1: - return exact_key_matches[0]["runtime_key"] - if exact_key_only: - raise RuntimeError( - f"{source} points to `{selector}`, but that exact runtime key is not " - "registered; run `rocm runtimes list` and activate an installed runtime" - ) - - runtime_id_matches = [ - manifest - for manifest in manifests - if manifest["runtime_id"].lower() == selector.lower() - ] - if len(runtime_id_matches) == 1: - return runtime_id_matches[0]["runtime_key"] - if len(runtime_id_matches) > 1: - raise RuntimeError( - f"{source} selector `{selector}` matches more than one installed runtime. " - "Activate one by runtime_key first: " - + runtime_keys_text(runtime_id_matches) - ) - raise RuntimeError( - f"{source} selector `{selector}` did not match any registered runtime; " - "run `rocm runtimes list` and activate an installed runtime" - ) - - -def runtime_keys_text(manifests: list[dict[str, str]]) -> str: - keys = sorted({manifest["runtime_key"] for manifest in manifests}) - return ", ".join(keys) if keys else "" - - -def rocm_cli_state_paths() -> tuple[Path, Path]: - config_dir = os.environ.get("ROCM_CLI_CONFIG_DIR") - data_dir = os.environ.get("ROCM_CLI_DATA_DIR") - if config_dir or data_dir: - config_base = Path(config_dir) if config_dir else default_rocm_cli_config_dir() - data_base = Path(data_dir) if data_dir else default_rocm_cli_data_dir() - return ( - config_base / "config.json", - data_base / "runtimes" / "registry", - ) - return default_rocm_cli_config_dir() / "config.json", ( - default_rocm_cli_data_dir() / "runtimes" / "registry" - ) - - -def default_rocm_cli_config_dir() -> Path: - return Path.home() / ".rocm" - - -def default_rocm_cli_data_dir() -> Path: - return Path.home() / ".rocm" - - -def run_self_test() -> int: - scratch_root = ( - Path(__file__).resolve().parents[1] / ".rocm-work" / "script-self-tests" - ) - scratch_root.mkdir(parents=True, exist_ok=True) - with tempfile.TemporaryDirectory(prefix="llama-cpp-", dir=scratch_root) as temp: - root = Path(temp) - config_dir = root / "config" - data_dir = root / "data" - registry_dir = data_dir / "runtimes" / "registry" - registry_dir.mkdir(parents=True) - config_dir.mkdir(parents=True) - write_runtime_manifest( - registry_dir, "runtime-old", "therock-release:gfx120X-all" - ) - write_runtime_manifest( - registry_dir, "runtime-new", "therock-release:gfx120X-all" - ) - write_runtime_manifest(registry_dir, "runtime-other", "therock-release:gfx1151") - write_config( - config_dir, - { - "active_runtime_key": "runtime-old", - "default_runtime_id": "therock-release:gfx120X-all", - }, - ) - - old_env = { - "ROCM_CLI_CONFIG_DIR": os.environ.get("ROCM_CLI_CONFIG_DIR"), - "ROCM_CLI_DATA_DIR": os.environ.get("ROCM_CLI_DATA_DIR"), - } - try: - os.environ["ROCM_CLI_CONFIG_DIR"] = str(config_dir) - os.environ["ROCM_CLI_DATA_DIR"] = str(data_dir) - assert resolve_runtime_id(None) == "runtime-old" - assert resolve_runtime_id("runtime-new") == "runtime-new" - assert resolve_runtime_id("therock-release:gfx1151") == "runtime-other" - - write_config( - config_dir, {"default_runtime_id": "therock-release:gfx120X-all"} - ) - try: - resolve_runtime_id(None) - except RuntimeError as exc: - message = str(exc) - assert "runtime-old" in message and "runtime-new" in message - else: - raise AssertionError("ambiguous default_runtime_id did not fail") - - write_config(config_dir, {"active_runtime_key": "missing-runtime"}) - try: - resolve_runtime_id(None) - except RuntimeError as exc: - assert "missing-runtime" in str(exc) - else: - raise AssertionError("missing active runtime key did not fail") - finally: - restore_env(old_env) - print("llama.cpp GPU script self-test passed") - return 0 - - -def write_runtime_manifest( - registry_dir: Path, runtime_key: str, runtime_id: str -) -> None: - (registry_dir / f"{runtime_key}.json").write_text( - json.dumps({"runtime_key": runtime_key, "runtime_id": runtime_id}), - encoding="utf-8", - ) - - -def write_config(config_dir: Path, payload: dict[str, str]) -> None: - (config_dir / "config.json").write_text(json.dumps(payload), encoding="utf-8") - - -def restore_env(values: dict[str, str | None]) -> None: - for key, value in values.items(): - if value is None: - os.environ.pop(key, None) - else: - os.environ[key] = value - - -def ensure_model(path: Path, url: str) -> None: - if path.is_file() and path.stat().st_size > 0: - return - path.parent.mkdir(parents=True, exist_ok=True) - with urllib.request.urlopen(url, timeout=60) as response: - path.write_bytes(response.read()) - if path.stat().st_size == 0: - raise RuntimeError(f"downloaded empty GGUF model: {path}") - - -def run_json( - command: list[str], *, env: dict[str, str], timeout: int -) -> dict[str, Any]: - completed = subprocess.run( - command, - env=env, - text=True, - capture_output=True, - timeout=timeout, - check=False, - ) - if completed.returncode != 0: - raise RuntimeError( - f"command failed ({completed.returncode}): {' '.join(command)}\n" - f"stdout:\n{completed.stdout}\nstderr:\n{completed.stderr}" - ) - try: - return json.loads(completed.stdout) - except json.JSONDecodeError as exc: - raise RuntimeError( - f"command did not return JSON: {' '.join(command)}\n{completed.stdout}" - ) from exc - - -def start_serve_http( - *, - engine: Path, - model_path: Path, - runtime_id: str, - args: argparse.Namespace, - env: dict[str, str], - repo_root: Path, -) -> tuple[subprocess.Popen[bytes] | None, Path, Path]: - if args.launch_mode == "launch": - return start_launch(engine, model_path, runtime_id, args, env=env) - - data_root = test_data_root(repo_root) - state_path = data_root / "test-state" / f"{args.service_id}.json" - log_path = data_root / "test-logs" / f"{args.service_id}.log" - command = [ - str(engine), - "serve-http", - args.service_id, - str(model_path), - "--host", - args.host, - "--port", - str(args.port), - "--device-policy", - "gpu_required", - "--runtime-id", - runtime_id, - "--state-path", - str(state_path), - ] - state_path.parent.mkdir(parents=True, exist_ok=True) - log_path.parent.mkdir(parents=True, exist_ok=True) - if state_path.exists(): - state_path.unlink() - log_file = log_path.open("wb") - try: - process = subprocess.Popen( - command, - env=env, - stdin=subprocess.DEVNULL, - stdout=log_file, - stderr=subprocess.STDOUT, - ) - finally: - log_file.close() - return process, state_path, log_path - - -def test_data_root(repo_root: Path) -> Path: - override = os.environ.get("ROCM_CLI_DATA_DIR") - if override: - return Path(override).expanduser() - return repo_root / "target" - - -def start_launch( - engine: Path, - model_path: Path, - runtime_id: str, - args: argparse.Namespace, - *, - env: dict[str, str], -) -> tuple[None, Path, Path]: - command = [ - str(engine), - "launch", - args.service_id, - str(model_path), - "--host", - args.host, - "--port", - str(args.port), - "--device-policy", - "gpu_required", - "--runtime-id", - runtime_id, - ] - started = time.monotonic() - response = run_json(command, env=env, timeout=args.timeout) - elapsed = time.monotonic() - started - if elapsed > 10: - raise RuntimeError( - f"launch took {elapsed:.1f}s; captured-output background launch should return promptly" - ) - state_path = Path(response["state_path"]) - log_path = Path(response["log_path"]) - return None, state_path, log_path - - -def assert_rocm_gpu_detected(detect: dict[str, Any]) -> None: - devices = detect.get("available_devices", []) - rocm_gpu = next((d for d in devices if d.get("kind") == "rocm_gpu"), None) - if not rocm_gpu or not rocm_gpu.get("available"): - raise RuntimeError( - "llama.cpp ROCm GPU was not detected; no CPU fallback is allowed:\n" - + json.dumps(detect, indent=2) - ) - - -def wait_health(host: str, port: int, timeout: int) -> dict[str, Any]: - deadline = time.monotonic() + timeout - last_error: Exception | None = None - while time.monotonic() < deadline: - try: - health = get_json(host, port, "/health", timeout=3) - if health.get("status") == "ok": - return health - except Exception as exc: - last_error = exc - time.sleep(0.5) - raise RuntimeError(f"llama-server did not become healthy: {last_error}") - - -def get_json(host: str, port: int, path: str, *, timeout: int) -> dict[str, Any]: - connection = http.client.HTTPConnection(host, port, timeout=timeout) - try: - connection.request("GET", path) - response = connection.getresponse() - body = response.read() - finally: - connection.close() - if response.status >= 400: - raise RuntimeError(f"GET {path} failed: HTTP {response.status}: {body!r}") - return json.loads(body.decode("utf-8")) - - -def post_json( - host: str, port: int, path: str, payload: dict[str, Any], *, timeout: int -) -> dict[str, Any]: - connection = http.client.HTTPConnection(host, port, timeout=timeout) - body = json.dumps(payload).encode("utf-8") - try: - connection.request( - "POST", path, body=body, headers={"Content-Type": "application/json"} - ) - response = connection.getresponse() - response_body = response.read() - finally: - connection.close() - if response.status >= 400: - raise RuntimeError( - f"POST {path} failed: HTTP {response.status}: {response_body!r}" - ) - return json.loads(response_body.decode("utf-8")) - - -def assert_gpu_log(log_text: str) -> None: - if "ROCm0" not in log_text: - raise RuntimeError("llama.cpp log did not show ROCm0 GPU usage") - if "C:\\WINDOWS\\SYSTEM32\\amdhip64_7.dll".lower() in log_text.lower(): - raise RuntimeError("llama.cpp loaded amdhip64_7.dll from System32") - - -def verify_loaded_modules(state: dict[str, Any]) -> dict[str, str]: - if platform.system() == "Windows": - return verify_windows_modules(state) - return verify_proc_maps(state) - - -def verify_windows_modules(state: dict[str, Any]) -> dict[str, str]: - server_pid = state.get("server_pid") - staged_dir = state.get("staged_runtime_dir") - if not server_pid or not staged_dir: - raise RuntimeError("state is missing server_pid or staged_runtime_dir") - command = ( - "$names=@(" - + ",".join(f"'{name}'" for name in sorted(MODULES_TO_VERIFY)) - + "); " - + f"Get-Process -Id {int(server_pid)} -Module | " - + "Where-Object { $names -contains $_.ModuleName } | " - + "Select-Object ModuleName,FileName | ConvertTo-Json -Depth 4" - ) - completed = subprocess.run( - ["powershell", "-NoProfile", "-Command", command], - text=True, - capture_output=True, - timeout=20, - check=False, - ) - if completed.returncode != 0: - raise RuntimeError(f"module inspection failed:\n{completed.stderr}") - if not completed.stdout.strip(): - raise RuntimeError("module inspection returned no HIP modules") - parsed = json.loads(completed.stdout) - rows = parsed if isinstance(parsed, list) else [parsed] - modules = {row["ModuleName"].lower(): row["FileName"] for row in rows} - missing = sorted(MODULES_TO_VERIFY - set(modules)) - if missing: - raise RuntimeError(f"missing loaded HIP modules: {missing}") - staged_prefix = str(Path(staged_dir)).lower() - for name, filename in modules.items(): - lower = filename.lower() - if "\\system32\\" in lower: - raise RuntimeError(f"{name} loaded from System32: {filename}") - if not lower.startswith(staged_prefix): - raise RuntimeError(f"{name} did not load from staged runtime: {filename}") - return modules - - -def verify_proc_maps(state: dict[str, Any]) -> dict[str, str]: - server_pid = state.get("server_pid") - runtime_env = state.get("therock_runtime_env") or {} - runtime_root = runtime_env.get("root") - if not server_pid or not runtime_root: - raise RuntimeError("state is missing server_pid or therock_runtime_env.root") - maps_path = Path("/proc") / str(int(server_pid)) / "maps" - if not maps_path.is_file(): - raise RuntimeError(f"process maps file was not found: {maps_path}") - module_paths: dict[str, str] = {} - for line in maps_path.read_text(encoding="utf-8", errors="replace").splitlines(): - if "/" not in line: - continue - path = line.split(maxsplit=5)[-1] - name = Path(path).name.lower() - for label, needle in LINUX_MODULES_TO_VERIFY.items(): - if name.startswith(f"{needle}.so"): - module_paths.setdefault(label, path) - missing = sorted(set(LINUX_MODULES_TO_VERIFY) - set(module_paths)) - if missing: - raise RuntimeError(f"missing loaded HIP/BLAS modules: {missing}") - roots = managed_therock_module_roots(runtime_env) - for name, path in module_paths.items(): - lower = str(Path(path)).lower() - if not any(lower.startswith(root) for root in roots): - raise RuntimeError( - f"{name} did not load from managed TheRock SDK wheel directories: {path}" - ) - return module_paths - - -def managed_therock_module_roots(runtime_env: dict[str, Any]) -> set[str]: - roots: set[str] = set() - for key in ["root", "bin"]: - add_module_root(roots, runtime_env.get(key)) - for key in ["bin_paths", "library_paths"]: - values = runtime_env.get(key) - if isinstance(values, list): - for value in values: - add_module_root(roots, value) - - runtime_root = runtime_env.get("root") - if isinstance(runtime_root, str) and runtime_root.strip(): - root_path = Path(runtime_root) - parent = root_path.parent - if root_path.name.startswith("_rocm_sdk_") and parent.is_dir(): - for sibling in parent.glob("_rocm_sdk_*"): - if sibling.is_dir(): - add_module_root(roots, sibling) - return roots - - -def add_module_root(roots: set[str], value: Any) -> None: - if not isinstance(value, (str, Path)) or not str(value).strip(): - return - path = Path(value) - roots.add(str(path).lower()) - with contextlib.suppress(OSError): - roots.add(str(path.resolve()).lower()) - - -def stop_service(process: subprocess.Popen[bytes] | None, state_path: Path) -> None: - state: dict[str, Any] = {} - if state_path.is_file(): - state = json.loads(state_path.read_text(encoding="utf-8")) - process_pid = process.pid if process else None - for pid in [state.get("server_pid"), state.get("pid"), process_pid]: - if not pid: - continue - if platform.system() == "Windows": - subprocess.run( - ["taskkill", "/PID", str(int(pid)), "/T", "/F"], - stdout=subprocess.DEVNULL, - stderr=subprocess.DEVNULL, - check=False, - ) - else: - subprocess.run(["kill", str(int(pid))], check=False) - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/scripts/package-linux-release.sh b/scripts/package-linux-release.sh index 1dd23827..5f9be65b 100755 --- a/scripts/package-linux-release.sh +++ b/scripts/package-linux-release.sh @@ -28,9 +28,9 @@ rm -rf "${ROOT_DIR}" rm -f "${ARCHIVE_PATH}" "${ARCHIVE_PATH}.sha256" "${ARCHIVE_PATH}.sig" "${TAR_PATH}" mkdir -p "${ROOT_DIR}/bin" -# The first-party engines (pytorch/llama.cpp/lemonade/atom/vllm/sglang) are -# compiled into `rocm` and run in-process; the standalone rocm-engine-* binaries -# are only an external third-party plugin fallback, so they are not shipped. +# The first-party engines (lemonade/vllm) are compiled into `rocm` and run +# in-process; the standalone rocm-engine-* binaries are only an external +# third-party plugin fallback, so they are not shipped. cp "${BINARY_DIR}/rocm" "${ROOT_DIR}/bin/" cp "${BINARY_DIR}/rocmd" "${ROOT_DIR}/bin/" cp README.md LICENSE.TXT install.sh "${ROOT_DIR}/" diff --git a/scripts/pytorch_therock_gpu_test.py b/scripts/pytorch_therock_gpu_test.py deleted file mode 100644 index 084eac66..00000000 --- a/scripts/pytorch_therock_gpu_test.py +++ /dev/null @@ -1,880 +0,0 @@ -#!/usr/bin/env python3 -# Copyright © Advanced Micro Devices, Inc., or its affiliates. -# -# SPDX-License-Identifier: MIT - -"""End-to-end PyTorch GPU smoke test for rocm-cli managed TheRock runtimes.""" - -from __future__ import annotations - -import argparse -import http.client -import json -import os -import platform -import subprocess -import sys -import tempfile -import time -from pathlib import Path -from typing import Any - -DEFAULT_MODEL_REF = "hf-internal-testing/tiny-random-gpt2" -WINDOWS_MODULE_PREFIXES = ("amdhip64", "hipblas", "rocblas", "torch_hip") -LINUX_MODULE_PREFIXES = ("libamdhip64", "libhipblas", "librocblas", "libtorch_hip") - - -def main() -> int: - args = parse_args() - if args.self_test: - return run_self_test() - - repo_root = Path(__file__).resolve().parents[1] - engine = resolve_path(args.engine, repo_root) - - if not engine.is_file(): - raise SystemExit( - f"PyTorch engine binary not found: {engine}\n" - "Build it with `cargo build -p rocm-engine-pytorch`, or pass " - "`--engine `." - ) - - env = os.environ.copy() - localize_huggingface_cache(env, repo_root) - print_step("Checking for an AMD GPU and a managed PyTorch folder...") - runtime_id = resolve_runtime_id(args.runtime_id) - detect = run_json([str(engine), "detect"], env=env, timeout=args.timeout) - assert_rocm_gpu_detected(detect) - - if args.env_id: - env_id = args.env_id.strip() - manifest = load_engine_manifest(env_id) - assert_env_matches_runtime(manifest, runtime_id) - else: - env_id, manifest = resolve_engine_env_for_runtime(runtime_id) - assert_managed_env_manifest_ready(manifest) - - process, state_path, log_path = start_serve_http( - engine=engine, - model_ref=args.model_ref, - env_id=env_id, - runtime_id=runtime_id, - args=args, - env=env, - repo_root=repo_root, - ) - print_step("Starting the PyTorch test server in AMD GPU mode.") - print(f"Log file: {log_path}", flush=True) - - try: - health = wait_health( - args.host, args.port, args.timeout, process, state_path, log_path - ) - print_step("PyTorch is running on the AMD GPU. Sending a tiny prompt...") - models = get_json(args.host, args.port, "/v1/models", timeout=args.timeout) - completion = post_json( - args.host, - args.port, - "/v1/completions", - { - "model": args.model_ref, - "prompt": "ROCm is", - "max_tokens": 8, - "temperature": 0, - }, - timeout=args.timeout, - ) - state = json.loads(state_path.read_text(encoding="utf-8")) - assert_gpu_state(state, health) - assert_models(models, args.model_ref) - assert_completion(completion) - module_paths = ( - {} - if args.skip_module_check - else verify_loaded_modules(state, manifest.get("env_path")) - ) - - summary = { - "ok": True, - "message": "Success: PyTorch ran on your AMD GPU using TheRock ROCm.", - "service_id": args.service_id, - "model_ref": args.model_ref, - "health": health, - "completion_text": completion["choices"][0]["text"], - "env_id": env_id, - "runtime_id": state.get("runtime_id"), - "state_path": str(state_path), - "log_path": str(log_path), - "worker_pid": state.get("pid"), - "verified_modules": module_paths, - } - print_step("Success: PyTorch ran on your AMD GPU using TheRock ROCm.") - print(json.dumps(summary, indent=2)) - finally: - if not args.keep_running: - stop_service(process, state_path) - - return 0 - - -def print_step(message: str) -> None: - print(f"[pytorch-gpu-test] {message}", flush=True) - - -def parse_args() -> argparse.Namespace: - default_engine = cargo_binary_path("debug", "rocm-engine-pytorch") - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--engine", default=str(default_engine)) - parser.add_argument("--model-ref", default=DEFAULT_MODEL_REF) - parser.add_argument("--service-id", default="pytorch-gpu-e2e") - parser.add_argument("--host", default="127.0.0.1") - parser.add_argument("--port", type=int, default=11441) - parser.add_argument("--timeout", type=int, default=180) - parser.add_argument( - "--env-id", help="managed PyTorch env id; defaults to engine detect" - ) - parser.add_argument( - "--runtime-id", - help=( - "managed TheRock runtime key or unambiguous runtime id; defaults to the " - "active rocm-cli runtime" - ), - ) - parser.add_argument( - "--skip-module-check", - action="store_true", - help="skip OS loaded-module inspection; AMD GPU state is still required", - ) - parser.add_argument("--keep-running", action="store_true") - parser.add_argument( - "--self-test", - action="store_true", - help="run offline selection checks for this script and exit", - ) - return parser.parse_args() - - -def resolve_path(value: str, repo_root: Path) -> Path: - path = Path(value).expanduser() - if not path.is_absolute(): - path = repo_root / path - return path.resolve() - - -def cargo_binary_path(profile: str, name: str) -> Path: - target_root = Path(os.environ.get("CARGO_TARGET_DIR", "target")).expanduser() - return target_root / profile / exe_name(name) - - -def exe_name(name: str) -> str: - return f"{name}.exe" if platform.system() == "Windows" else name - - -def run_json( - command: list[str], *, env: dict[str, str], timeout: int -) -> dict[str, Any]: - completed = subprocess.run( - command, - env=env, - text=True, - capture_output=True, - timeout=timeout, - check=False, - ) - if completed.returncode != 0: - raise RuntimeError( - f"command failed ({completed.returncode}): {' '.join(command)}\n" - f"stdout:\n{completed.stdout}\nstderr:\n{completed.stderr}" - ) - try: - return json.loads(completed.stdout) - except json.JSONDecodeError as exc: - raise RuntimeError( - f"command did not return JSON: {' '.join(command)}\n{completed.stdout}" - ) from exc - - -def localize_huggingface_cache(env: dict[str, str], repo_root: Path) -> None: - cache_root = test_cache_root(repo_root) / "huggingface" - env.setdefault("HF_HOME", str(cache_root)) - env.setdefault("HUGGINGFACE_HUB_CACHE", str(cache_root / "hub")) - env.setdefault("TRANSFORMERS_CACHE", str(cache_root / "transformers")) - - -def assert_rocm_gpu_detected(detect: dict[str, Any]) -> None: - devices = detect.get("available_devices", []) - rocm_gpu = next( - (device for device in devices if device.get("kind") == "rocm_gpu"), - None, - ) - if not rocm_gpu or not rocm_gpu.get("available"): - raise RuntimeError( - "PyTorch could not see an AMD GPU. No CPU fallback is allowed.\n" - "Run `rocm examine`, fix AMD driver/GPU detection there, then retry this test.\n" - + json.dumps(detect, indent=2) - ) - - -def assert_managed_env_manifest_ready(manifest: dict[str, Any]) -> None: - env_path = manifest.get("env_path") - if not isinstance(env_path, str) or not env_path.strip(): - raise RuntimeError("managed PyTorch env manifest is missing env_path") - if not Path(env_path).is_dir(): - raise RuntimeError( - f"managed PyTorch env path from manifest does not exist: {env_path}" - ) - - -def assert_env_matches_runtime(manifest: dict[str, Any], runtime_id: str) -> None: - manifest_runtime = manifest.get("runtime_id") - if not isinstance(manifest_runtime, str) or not manifest_runtime.strip(): - raise RuntimeError("managed PyTorch env manifest is missing runtime_id") - if manifest_runtime.lower() != runtime_id.lower(): - raise RuntimeError( - "managed PyTorch env does not belong to the selected ROCm runtime: " - f"env_id={manifest.get('env_id')}, env_runtime_id={manifest_runtime}, " - f"selected_runtime_key={runtime_id}" - ) - - -def start_serve_http( - *, - engine: Path, - model_ref: str, - env_id: str, - runtime_id: str | None, - args: argparse.Namespace, - env: dict[str, str], - repo_root: Path, -) -> tuple[subprocess.Popen[bytes], Path, Path]: - data_root = test_data_root(repo_root) - state_path = data_root / "test-state" / f"{args.service_id}.json" - log_path = data_root / "test-logs" / f"{args.service_id}.log" - command = [ - str(engine), - "serve-http", - args.service_id, - model_ref, - "--host", - args.host, - "--port", - str(args.port), - "--device-policy", - "gpu_required", - "--env-id", - env_id, - "--state-path", - str(state_path), - ] - if runtime_id: - command.extend(["--runtime-id", runtime_id]) - - state_path.parent.mkdir(parents=True, exist_ok=True) - log_path.parent.mkdir(parents=True, exist_ok=True) - if state_path.exists(): - state_path.unlink() - log_file = log_path.open("wb") - try: - process = subprocess.Popen( - command, - env=env, - stdin=subprocess.DEVNULL, - stdout=log_file, - stderr=subprocess.STDOUT, - ) - finally: - log_file.close() - return process, state_path, log_path - - -def test_cache_root(repo_root: Path) -> Path: - override = os.environ.get("ROCM_CLI_CACHE_DIR") - if override: - return Path(override).expanduser() / "test-cache" - return repo_root / "target" / "test-cache" - - -def test_data_root(repo_root: Path) -> Path: - override = os.environ.get("ROCM_CLI_DATA_DIR") - if override: - return Path(override).expanduser() - return repo_root / "target" - - -def wait_health( - host: str, - port: int, - timeout: int, - process: subprocess.Popen[bytes], - state_path: Path, - log_path: Path, -) -> dict[str, Any]: - deadline = time.monotonic() + timeout - last_error: Exception | None = None - last_log_line_count = 0 - last_progress = 0.0 - while time.monotonic() < deadline: - try: - health = get_json(host, port, "/healthz", timeout=3) - if health.get("status") == "ok": - return health - except Exception as exc: - last_error = exc - exit_code = process.poll() - if exit_code is not None: - raise RuntimeError( - f"PyTorch worker exited before it became healthy (exit {exit_code}).\n" - + failure_context(state_path, log_path) - ) - now = time.monotonic() - if now - last_progress >= 5: - last_log_line_count = print_new_log_lines(log_path, last_log_line_count) - print_state_progress(state_path) - last_progress = now - time.sleep(0.5) - raise RuntimeError( - f"PyTorch worker did not become healthy: {last_error}\n" - + failure_context(state_path, log_path) - ) - - -def print_new_log_lines(log_path: Path, previous_count: int) -> int: - if not log_path.is_file(): - print_step("Waiting for PyTorch to start. The log file is not ready yet.") - return previous_count - lines = log_path.read_text(encoding="utf-8", errors="replace").splitlines() - if len(lines) <= previous_count: - print_step("Still waiting for PyTorch to load the tiny model...") - return len(lines) - for line in lines[previous_count:][-8:]: - print(f"[pytorch-gpu-test][log] {line}", flush=True) - return len(lines) - - -def print_state_progress(state_path: Path) -> None: - if not state_path.is_file(): - return - try: - state = json.loads(state_path.read_text(encoding="utf-8")) - except Exception: - return - status = state.get("status") - if status == "starting": - print_step( - "The test server is starting. This can take a little while on first run." - ) - elif status: - print_step(f"Current status: {status}") - - -def failure_context(state_path: Path, log_path: Path) -> str: - parts = [f"state file: {state_path}", f"log file: {log_path}"] - if state_path.is_file(): - try: - state = json.loads(state_path.read_text(encoding="utf-8")) - visible_state = { - key: state.get(key) - for key in [ - "status", - "error", - "device", - "device_policy", - "runtime_id", - "env_id", - ] - if key in state - } - parts.append("state: " + json.dumps(visible_state, indent=2)) - except Exception as exc: - parts.append(f"state could not be read: {exc}") - if log_path.is_file(): - log_lines = log_path.read_text(encoding="utf-8", errors="replace").splitlines() - if log_lines: - parts.append("last log lines:\n" + "\n".join(log_lines[-40:])) - return "\n".join(parts) - - -def get_json(host: str, port: int, path: str, *, timeout: int) -> dict[str, Any]: - connection = http.client.HTTPConnection(host, port, timeout=timeout) - try: - connection.request("GET", path) - response = connection.getresponse() - body = response.read() - finally: - connection.close() - if response.status >= 400: - raise RuntimeError(f"GET {path} failed: HTTP {response.status}: {body!r}") - return json.loads(body.decode("utf-8")) - - -def post_json( - host: str, port: int, path: str, payload: dict[str, Any], *, timeout: int -) -> dict[str, Any]: - connection = http.client.HTTPConnection(host, port, timeout=timeout) - body = json.dumps(payload).encode("utf-8") - try: - connection.request( - "POST", path, body=body, headers={"Content-Type": "application/json"} - ) - response = connection.getresponse() - response_body = response.read() - finally: - connection.close() - if response.status >= 400: - raise RuntimeError( - f"POST {path} failed: HTTP {response.status}: {response_body!r}" - ) - return json.loads(response_body.decode("utf-8")) - - -def assert_gpu_state(state: dict[str, Any], health: dict[str, Any]) -> None: - for label, payload in [("state", state), ("health", health)]: - if payload.get("device") != "cuda": - raise RuntimeError( - f"PyTorch {label} reported {payload.get('device')!r}, not cuda; " - "no CPU fallback is allowed" - ) - if payload.get("device_policy") != "gpu_required": - raise RuntimeError( - f"PyTorch {label} reported device_policy={payload.get('device_policy')!r}, " - "expected gpu_required" - ) - gpu_count = payload.get("gpu_count") - if not isinstance(gpu_count, int) or gpu_count < 1: - raise RuntimeError(f"PyTorch {label} did not report a visible GPU") - if state.get("status") != "ready": - raise RuntimeError(f"PyTorch state is not ready: {state}") - - -def assert_models(models: dict[str, Any], model_ref: str) -> None: - rows = models.get("data") - if not isinstance(rows, list) or not rows: - raise RuntimeError("/v1/models did not return any models") - ids = {row.get("id") for row in rows if isinstance(row, dict)} - if model_ref not in ids: - raise RuntimeError(f"/v1/models did not include {model_ref!r}: {models}") - - -def assert_completion(completion: dict[str, Any]) -> None: - choices = completion.get("choices") - if not isinstance(choices, list) or not choices: - raise RuntimeError(f"/v1/completions returned no choices: {completion}") - text = choices[0].get("text") if isinstance(choices[0], dict) else None - if not isinstance(text, str): - raise RuntimeError(f"/v1/completions did not return text: {completion}") - - -def resolve_runtime_id(explicit: str | None) -> str: - config_path, registry_dir = rocm_cli_state_paths() - manifests = load_runtime_registry_manifests(registry_dir) - if explicit: - return resolve_runtime_selector(explicit, manifests, "explicit --runtime-id") - if config_path.is_file(): - config = json.loads(config_path.read_text(encoding="utf-8")) - active_key = config.get("active_runtime_key") - if isinstance(active_key, str) and active_key.strip(): - return resolve_runtime_selector( - active_key.strip(), - manifests, - f"active_runtime_key in {config_path}", - exact_key_only=True, - ) - runtime_id = config.get("default_runtime_id") - if isinstance(runtime_id, str) and runtime_id.strip(): - return resolve_runtime_selector( - runtime_id.strip(), - manifests, - f"default_runtime_id in {config_path}", - ) - raise RuntimeError( - "no active managed TheRock runtime was found; run " - "`rocm runtimes activate ` or pass --runtime-id " - ) - - -def load_runtime_registry_manifests(registry_dir: Path) -> list[dict[str, str]]: - manifests: list[dict[str, str]] = [] - if not registry_dir.is_dir(): - return manifests - for path in registry_dir.glob("*.json"): - try: - manifest = json.loads(path.read_text(encoding="utf-8")) - except json.JSONDecodeError: - continue - runtime_key = manifest.get("runtime_key") - runtime_id = manifest.get("runtime_id") - if not isinstance(runtime_key, str) or not runtime_key.strip(): - runtime_key = path.stem - if not isinstance(runtime_id, str) or not runtime_id.strip(): - continue - manifests.append( - { - "runtime_key": runtime_key.strip(), - "runtime_id": runtime_id.strip(), - } - ) - return manifests - - -def resolve_runtime_selector( - selector: str, - manifests: list[dict[str, str]], - source: str, - *, - exact_key_only: bool = False, -) -> str: - exact_key_matches = [ - manifest - for manifest in manifests - if manifest["runtime_key"].lower() == selector.lower() - ] - if len(exact_key_matches) == 1: - return exact_key_matches[0]["runtime_key"] - if exact_key_only: - raise RuntimeError( - f"{source} points to `{selector}`, but that exact runtime key is not " - "registered; run `rocm runtimes list` and activate an installed runtime" - ) - - runtime_id_matches = [ - manifest - for manifest in manifests - if manifest["runtime_id"].lower() == selector.lower() - ] - if len(runtime_id_matches) == 1: - return runtime_id_matches[0]["runtime_key"] - if len(runtime_id_matches) > 1: - raise RuntimeError( - f"{source} selector `{selector}` matches more than one installed runtime. " - "Activate one by runtime_key first: " - + runtime_keys_text(runtime_id_matches) - ) - raise RuntimeError( - f"{source} selector `{selector}` did not match any registered runtime; " - "run `rocm runtimes list` and activate an installed runtime" - ) - - -def runtime_keys_text(manifests: list[dict[str, str]]) -> str: - keys = sorted({manifest["runtime_key"] for manifest in manifests}) - return ", ".join(keys) if keys else "" - - -def rocm_cli_state_paths() -> tuple[Path, Path]: - config_dir = os.environ.get("ROCM_CLI_CONFIG_DIR") - data_dir = os.environ.get("ROCM_CLI_DATA_DIR") - config_base = ( - Path(config_dir).expanduser() if config_dir else default_rocm_cli_dir() - ) - data_base = ( - Path(data_dir).expanduser() - if data_dir - else default_rocm_cli_data_dir(config_base) - ) - return config_base / "config.json", data_base / "runtimes" / "registry" - - -def default_rocm_cli_dir() -> Path: - return Path.home() / ".rocm" - - -def rocm_cli_data_dir() -> Path: - override = os.environ.get("ROCM_CLI_DATA_DIR") - if override: - return Path(override).expanduser() - config_dir = os.environ.get("ROCM_CLI_CONFIG_DIR") - config_base = ( - Path(config_dir).expanduser() if config_dir else default_rocm_cli_dir() - ) - return default_rocm_cli_data_dir(config_base) - - -def default_rocm_cli_data_dir(config_base: Path) -> Path: - config_path = config_base / "config.json" - try: - config = json.loads(config_path.read_text(encoding="utf-8")) - except (OSError, json.JSONDecodeError): - return default_rocm_cli_dir() - setup = config.get("setup") - if isinstance(setup, dict): - therock_venv = setup.get("therock_venv") - if isinstance(therock_venv, str) and therock_venv.strip(): - return Path(therock_venv).expanduser() - return default_rocm_cli_dir() - - -def load_engine_manifest(env_id: str) -> dict[str, Any]: - manifest_path = ( - rocm_cli_data_dir() / "engines" / "pytorch" / "manifests" / f"{env_id}.json" - ) - if not manifest_path.is_file(): - raise RuntimeError( - f"managed PyTorch env manifest was not found: {manifest_path}" - ) - return json.loads(manifest_path.read_text(encoding="utf-8")) - - -def resolve_engine_env_for_runtime(runtime_id: str) -> tuple[str, dict[str, Any]]: - manifest_dir = rocm_cli_data_dir() / "engines" / "pytorch" / "manifests" - if not manifest_dir.is_dir(): - raise RuntimeError( - "no managed PyTorch environments were found. Install one with " - f"`rocm engines install pytorch --runtime-id {runtime_id}`" - ) - matches: list[tuple[int, str, dict[str, Any]]] = [] - for path in manifest_dir.glob("*.json"): - try: - manifest = json.loads(path.read_text(encoding="utf-8")) - except json.JSONDecodeError: - continue - manifest_runtime = manifest.get("runtime_id") - env_id = manifest.get("env_id") - if ( - isinstance(manifest_runtime, str) - and manifest_runtime.lower() == runtime_id.lower() - and isinstance(env_id, str) - and env_id.strip() - ): - try: - modified = path.stat().st_mtime_ns - except OSError: - modified = 0 - matches.append((modified, env_id.strip(), manifest)) - if not matches: - raise RuntimeError( - "no managed PyTorch environment matches the selected ROCm runtime " - f"`{runtime_id}`. Install one with " - f"`rocm engines install pytorch --runtime-id {runtime_id}`" - ) - matches.sort(reverse=True, key=lambda row: (row[0], row[1])) - _, env_id, manifest = matches[0] - return env_id, manifest - - -def run_self_test() -> int: - scratch_root = ( - Path(__file__).resolve().parents[1] / ".rocm-work" / "script-self-tests" - ) - scratch_root.mkdir(parents=True, exist_ok=True) - with tempfile.TemporaryDirectory(prefix="pytorch-", dir=scratch_root) as temp: - root = Path(temp) - config_dir = root / "config" - data_dir = root / "data" - registry_dir = data_dir / "runtimes" / "registry" - manifest_dir = data_dir / "engines" / "pytorch" / "manifests" - env_root = data_dir / "engines" / "pytorch" / "envs" - registry_dir.mkdir(parents=True) - manifest_dir.mkdir(parents=True) - env_root.mkdir(parents=True) - config_dir.mkdir(parents=True) - - write_runtime_manifest( - registry_dir, "runtime-old", "therock-release:gfx120X-all" - ) - write_runtime_manifest( - registry_dir, "runtime-new", "therock-release:gfx120X-all" - ) - write_runtime_manifest(registry_dir, "runtime-other", "therock-release:gfx1151") - write_engine_manifest(manifest_dir, env_root, "env-old", "runtime-old") - write_engine_manifest(manifest_dir, env_root, "env-new", "runtime-new") - write_engine_manifest(manifest_dir, env_root, "env-other", "runtime-other") - write_config( - config_dir, - { - "active_runtime_key": "runtime-old", - "default_runtime_id": "therock-release:gfx120X-all", - }, - ) - - old_env = { - "ROCM_CLI_CONFIG_DIR": os.environ.get("ROCM_CLI_CONFIG_DIR"), - "ROCM_CLI_DATA_DIR": os.environ.get("ROCM_CLI_DATA_DIR"), - } - try: - os.environ["ROCM_CLI_CONFIG_DIR"] = str(config_dir) - os.environ["ROCM_CLI_DATA_DIR"] = str(data_dir) - assert resolve_runtime_id(None) == "runtime-old" - env_id, manifest = resolve_engine_env_for_runtime("runtime-old") - assert env_id == "env-old" - assert_env_matches_runtime(manifest, "runtime-old") - assert_managed_env_manifest_ready(manifest) - - assert resolve_runtime_id("runtime-new") == "runtime-new" - assert resolve_runtime_id("therock-release:gfx1151") == "runtime-other" - try: - assert_env_matches_runtime( - load_engine_manifest("env-new"), "runtime-old" - ) - except RuntimeError as exc: - assert "selected ROCm runtime" in str(exc) - else: - raise AssertionError("mismatched env/runtime pair did not fail") - - write_config( - config_dir, {"default_runtime_id": "therock-release:gfx120X-all"} - ) - try: - resolve_runtime_id(None) - except RuntimeError as exc: - message = str(exc) - assert "runtime-old" in message and "runtime-new" in message - else: - raise AssertionError("ambiguous default_runtime_id did not fail") - - write_config(config_dir, {"active_runtime_key": "missing-runtime"}) - try: - resolve_runtime_id(None) - except RuntimeError as exc: - assert "missing-runtime" in str(exc) - else: - raise AssertionError("missing active runtime key did not fail") - finally: - restore_env(old_env) - print("PyTorch GPU script self-test passed") - return 0 - - -def write_runtime_manifest( - registry_dir: Path, runtime_key: str, runtime_id: str -) -> None: - (registry_dir / f"{runtime_key}.json").write_text( - json.dumps({"runtime_key": runtime_key, "runtime_id": runtime_id}), - encoding="utf-8", - ) - - -def write_engine_manifest( - manifest_dir: Path, - env_root: Path, - env_id: str, - runtime_id: str, -) -> None: - env_path = env_root / env_id - env_path.mkdir(parents=True) - payload = { - "env_id": env_id, - "runtime_id": runtime_id, - "env_path": str(env_path), - } - (manifest_dir / f"{env_id}.json").write_text(json.dumps(payload), encoding="utf-8") - - -def write_config(config_dir: Path, payload: dict[str, str]) -> None: - (config_dir / "config.json").write_text(json.dumps(payload), encoding="utf-8") - - -def restore_env(values: dict[str, str | None]) -> None: - for key, value in values.items(): - if value is None: - os.environ.pop(key, None) - else: - os.environ[key] = value - - -def verify_loaded_modules(state: dict[str, Any], env_path: Any) -> dict[str, str]: - if not isinstance(env_path, str) or not env_path.strip(): - raise RuntimeError("managed PyTorch env manifest is missing env_path") - if platform.system() == "Windows": - return verify_windows_modules(state, Path(env_path)) - return verify_proc_maps(state, Path(env_path)) - - -def verify_windows_modules(state: dict[str, Any], env_path: Path) -> dict[str, str]: - worker_pid = state.get("pid") - if not worker_pid: - raise RuntimeError("state is missing the PyTorch worker pid") - command = ( - f"Get-Process -Id {int(worker_pid)} -Module | " - "Select-Object ModuleName,FileName | ConvertTo-Json -Depth 4" - ) - completed = subprocess.run( - ["powershell", "-NoProfile", "-Command", command], - text=True, - capture_output=True, - timeout=20, - check=False, - ) - if completed.returncode != 0: - raise RuntimeError(f"module inspection failed:\n{completed.stderr}") - if not completed.stdout.strip(): - raise RuntimeError("module inspection returned no loaded modules") - parsed = json.loads(completed.stdout) - rows = parsed if isinstance(parsed, list) else [parsed] - modules: dict[str, str] = {} - for row in rows: - name = str(row.get("ModuleName", "")).lower() - filename = str(row.get("FileName", "")) - if any(name.startswith(prefix) for prefix in WINDOWS_MODULE_PREFIXES): - modules[name] = filename - ensure_expected_modules(modules, WINDOWS_MODULE_PREFIXES, env_path) - return modules - - -def verify_proc_maps(state: dict[str, Any], env_path: Path) -> dict[str, str]: - worker_pid = state.get("pid") - if not worker_pid: - raise RuntimeError("state is missing the PyTorch worker pid") - maps_path = Path("/proc") / str(int(worker_pid)) / "maps" - if not maps_path.is_file(): - raise RuntimeError(f"process maps file was not found: {maps_path}") - modules: dict[str, str] = {} - for line in maps_path.read_text(encoding="utf-8", errors="replace").splitlines(): - if "/" not in line: - continue - path = line.split(maxsplit=5)[-1] - if not path.startswith("/"): - continue - name = Path(path).name.lower() - if any(name.startswith(prefix) for prefix in LINUX_MODULE_PREFIXES): - modules.setdefault(name, path) - ensure_expected_modules(modules, LINUX_MODULE_PREFIXES, env_path) - return modules - - -def ensure_expected_modules( - modules: dict[str, str], - prefixes: tuple[str, ...], - env_path: Path, -) -> None: - if not modules: - raise RuntimeError("no loaded PyTorch HIP modules were found") - if not any(name.startswith(prefixes[0]) for name in modules): - raise RuntimeError(f"missing loaded HIP runtime module matching {prefixes[0]}") - - env_prefix = str(env_path.resolve()).lower() - for name, filename in modules.items(): - lower_path = str(Path(filename).resolve()).lower() - if platform.system() == "Windows" and "\\system32\\" in lower_path: - raise RuntimeError(f"{name} loaded from System32: {filename}") - if not lower_path.startswith(env_prefix): - raise RuntimeError( - f"{name} did not load from the managed PyTorch env: {filename}" - ) - - -def stop_service(process: subprocess.Popen[bytes], state_path: Path) -> None: - state: dict[str, Any] = {} - if state_path.is_file(): - state = json.loads(state_path.read_text(encoding="utf-8")) - for pid in [state.get("pid"), process.pid]: - if not pid: - continue - if platform.system() == "Windows": - subprocess.run( - ["taskkill", "/PID", str(int(pid)), "/T", "/F"], - stdout=subprocess.DEVNULL, - stderr=subprocess.DEVNULL, - check=False, - ) - else: - subprocess.run(["kill", str(int(pid))], check=False) - - -if __name__ == "__main__": - try: - raise SystemExit(main()) - except KeyboardInterrupt: - raise SystemExit(130) from None - except Exception as exc: - print(f"error: {exc}", file=sys.stderr) - raise SystemExit(1) from exc diff --git a/scripts/sglang_therock_gpu_test.py b/scripts/sglang_therock_gpu_test.py deleted file mode 100644 index 3de342b1..00000000 --- a/scripts/sglang_therock_gpu_test.py +++ /dev/null @@ -1,560 +0,0 @@ -#!/usr/bin/env python3 -# Copyright © Advanced Micro Devices, Inc., or its affiliates. -# -# SPDX-License-Identifier: MIT - -"""End-to-end SGLang ROCm GPU smoke test for rocm-cli managed TheRock runtimes.""" - -from __future__ import annotations - -import argparse -import json -import os -import platform -import subprocess -import tempfile -import time -from pathlib import Path -from typing import Any - -from vllm_therock_gpu_test import ( - DEFAULT_MATH_MODULES, - HIP_MODULE, - completion_text, - first_model_id, - get_json, - managed_therock_module_roots, - post_json, - resolve_path, - resolve_runtime_selector, - restore_env, - rocm_cli_state_paths, - run_json, - write_config, - write_runtime_manifest, -) - - -def main() -> int: - args = parse_args() - if args.self_test: - return run_self_test() - - if platform.system() == "Windows": - print( - json.dumps( - { - "ok": True, - "skipped": True, - "reason": ( - "SGLang GPU acceptance is skipped on native Windows; " - "use WSL/Linux for ROCm GPU serving. No CPU fallback is allowed." - ), - }, - indent=2, - ) - ) - return 0 - - repo_root = Path(__file__).resolve().parents[1] - engine = resolve_path(args.engine, repo_root) - runtime_id = resolve_runtime_id(args.runtime_id) - math_modules = args.math_module or DEFAULT_MATH_MODULES - - if not engine.is_file(): - raise SystemExit(f"engine binary not found: {engine}") - reject_external_runtime_env() - - env = os.environ.copy() - detect = run_json([str(engine), "detect"], env=env, timeout=args.timeout) - assert_sglang_gpu_detected(detect) - - capabilities = run_json( - [str(engine), "capabilities"], env=env, timeout=args.timeout - ) - if capabilities.get("cpu"): - raise RuntimeError("SGLang capabilities unexpectedly report CPU support") - if not capabilities.get("rocm_gpu"): - raise RuntimeError("SGLang capabilities did not report ROCm GPU support") - - process, state_path, log_path = start_sglang( - engine=engine, - model=args.model, - runtime_id=runtime_id, - args=args, - env=env, - repo_root=repo_root, - ) - - try: - health = wait_sglang_openai_ready(args.host, args.port, args.timeout) - models = get_json(args.host, args.port, "/v1/models", timeout=args.timeout) - served_model = first_model_id(models) - completion = post_json( - args.host, - args.port, - "/v1/completions", - { - "model": served_model, - "prompt": args.prompt, - "max_tokens": args.max_tokens, - "temperature": 0, - }, - timeout=args.timeout, - ) - state = json.loads(state_path.read_text(encoding="utf-8")) - env_values = verify_managed_env(state) - module_paths = verify_loaded_modules(state, math_modules) - text = completion_text(completion) - if not text.strip(): - raise RuntimeError("SGLang completion returned empty text") - - summary = { - "ok": True, - "launch_mode": args.launch_mode, - "service_id": args.service_id, - "model": args.model, - "served_model": served_model, - "runtime_id": runtime_id, - "health": health, - "completion_text": text, - "state_path": str(state_path), - "log_path": str(log_path), - "server_pid": state.get("server_pid") or state.get("pid"), - "therock_runtime_env": state.get("therock_runtime_env"), - "verified_env": env_values, - "verified_modules": module_paths, - } - print(json.dumps(summary, indent=2)) - finally: - if not args.keep_running: - stop_service(process, state_path) - - return 0 - - -def parse_args() -> argparse.Namespace: - default_engine = ( - Path("target/debug/rocm-engine-sglang.exe") - if platform.system() == "Windows" - else Path("target/debug/rocm-engine-sglang") - ) - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--engine", default=str(default_engine)) - parser.add_argument( - "--model", - default=os.environ.get( - "ROCM_CLI_SGLANG_TEST_MODEL", "Qwen/Qwen2.5-1.5B-Instruct" - ), - help="small Hugging Face model id or local model path served by SGLang", - ) - parser.add_argument( - "--runtime-id", - help=( - "managed TheRock runtime key or unambiguous runtime id; defaults to " - "the active rocm-cli runtime" - ), - ) - parser.add_argument("--service-id", default="sglang-gpu-e2e") - parser.add_argument("--host", default="127.0.0.1") - parser.add_argument("--port", type=int, default=11442) - parser.add_argument("--timeout", type=int, default=300) - parser.add_argument("--prompt", default="Once upon a time") - parser.add_argument("--max-tokens", type=int, default=8) - parser.add_argument( - "--launch-mode", - choices=("serve-http", "launch"), - default="serve-http", - help="serve-http runs the wrapper directly; launch verifies background launch", - ) - parser.add_argument( - "--math-module", - action="append", - help=( - "math-library basename where at least one must load from the managed " - "TheRock SDK wheel directories; defaults to libhipblas/libhipblaslt/librocblas" - ), - ) - parser.add_argument("--keep-running", action="store_true") - parser.add_argument( - "--self-test", - action="store_true", - help="run offline selection checks for this script and exit", - ) - return parser.parse_args() - - -def reject_external_runtime_env() -> None: - blocked = [ - name - for name in [ - "ROCM_CLI_SGLANG_COMMAND", - "SGLANG_COMMAND", - "ROCM_CLI_SGLANG_PYTHON", - "SGLANG_PYTHON", - ] - if os.environ.get(name) - ] - if blocked: - raise RuntimeError( - "SGLang GPU acceptance requires discovery through a rocm-cli managed " - f"TheRock runtime manifest; unset external runtime overrides: {', '.join(blocked)}" - ) - - -def wait_sglang_openai_ready(host: str, port: int, timeout: int) -> dict[str, Any]: - deadline = time.monotonic() + timeout - last_error: Exception | None = None - while time.monotonic() < deadline: - try: - models = get_json(host, port, "/v1/models", timeout=3) - first_model_id(models) - return {"status_code": 200, "body": "v1/models ready"} - except Exception as exc: - last_error = exc - time.sleep(0.5) - raise RuntimeError( - f"SGLang OpenAI model endpoint did not become ready: {last_error}" - ) - - -def resolve_runtime_id(explicit: str | None) -> str: - config_path, registry_dir = rocm_cli_state_paths() - manifests = load_runtime_registry_manifests(registry_dir) - if explicit: - return resolve_runtime_selector(explicit, manifests, "explicit --runtime-id") - if config_path.is_file(): - config = json.loads(config_path.read_text(encoding="utf-8")) - active_key = config.get("active_runtime_key") - if isinstance(active_key, str) and active_key.strip(): - return resolve_runtime_selector( - active_key.strip(), - manifests, - f"active_runtime_key in {config_path}", - exact_key_only=True, - ) - runtime_id = config.get("default_runtime_id") - if isinstance(runtime_id, str) and runtime_id.strip(): - return resolve_runtime_selector( - runtime_id.strip(), - manifests, - f"default_runtime_id in {config_path}", - ) - raise RuntimeError( - "no active managed TheRock runtime was found; run " - "`rocm runtimes activate ` or pass --runtime-id " - ) - - -def load_runtime_registry_manifests(registry_dir: Path) -> list[dict[str, str]]: - manifests: list[dict[str, str]] = [] - if not registry_dir.is_dir(): - return manifests - for path in registry_dir.glob("*.json"): - try: - manifest = json.loads(path.read_text(encoding="utf-8")) - except json.JSONDecodeError: - continue - runtime_key = manifest.get("runtime_key") - runtime_id = manifest.get("runtime_id") - if not isinstance(runtime_key, str) or not runtime_key.strip(): - runtime_key = path.stem - if not isinstance(runtime_id, str) or not runtime_id.strip(): - continue - manifests.append( - { - "runtime_key": runtime_key.strip(), - "runtime_id": runtime_id.strip(), - } - ) - return manifests - - -def run_self_test() -> int: - scratch_root = ( - Path(__file__).resolve().parents[1] / ".rocm-work" / "script-self-tests" - ) - scratch_root.mkdir(parents=True, exist_ok=True) - with tempfile.TemporaryDirectory(prefix="sglang-", dir=scratch_root) as temp: - root = Path(temp) - config_dir = root / "config" - data_dir = root / "data" - registry_dir = data_dir / "runtimes" / "registry" - registry_dir.mkdir(parents=True) - config_dir.mkdir(parents=True) - write_runtime_manifest( - registry_dir, "runtime-old", "therock-release:gfx120X-all" - ) - write_runtime_manifest( - registry_dir, "runtime-new", "therock-release:gfx120X-all" - ) - write_runtime_manifest(registry_dir, "runtime-other", "therock-release:gfx1151") - write_config( - config_dir, - { - "active_runtime_key": "runtime-old", - "default_runtime_id": "therock-release:gfx120X-all", - }, - ) - - old_env = { - "ROCM_CLI_CONFIG_DIR": os.environ.get("ROCM_CLI_CONFIG_DIR"), - "ROCM_CLI_DATA_DIR": os.environ.get("ROCM_CLI_DATA_DIR"), - "ROCM_CLI_SGLANG_COMMAND": os.environ.get("ROCM_CLI_SGLANG_COMMAND"), - "SGLANG_COMMAND": os.environ.get("SGLANG_COMMAND"), - "ROCM_CLI_SGLANG_PYTHON": os.environ.get("ROCM_CLI_SGLANG_PYTHON"), - "SGLANG_PYTHON": os.environ.get("SGLANG_PYTHON"), - } - try: - os.environ["ROCM_CLI_CONFIG_DIR"] = str(config_dir) - os.environ["ROCM_CLI_DATA_DIR"] = str(data_dir) - for key in [ - "ROCM_CLI_SGLANG_COMMAND", - "SGLANG_COMMAND", - "ROCM_CLI_SGLANG_PYTHON", - "SGLANG_PYTHON", - ]: - os.environ.pop(key, None) - - assert resolve_runtime_id(None) == "runtime-old" - assert resolve_runtime_id("runtime-new") == "runtime-new" - assert resolve_runtime_id("therock-release:gfx1151") == "runtime-other" - - os.environ["ROCM_CLI_SGLANG_COMMAND"] = "/tmp/not-used-sglang" - try: - reject_external_runtime_env() - except RuntimeError as exc: - assert "managed TheRock runtime manifest" in str(exc) - else: - raise AssertionError("external SGLang override did not fail") - os.environ.pop("ROCM_CLI_SGLANG_COMMAND", None) - - write_config( - config_dir, {"default_runtime_id": "therock-release:gfx120X-all"} - ) - try: - resolve_runtime_id(None) - except RuntimeError as exc: - message = str(exc) - assert "runtime-old" in message and "runtime-new" in message - else: - raise AssertionError("ambiguous default_runtime_id did not fail") - - write_config(config_dir, {"active_runtime_key": "missing-runtime"}) - try: - resolve_runtime_id(None) - except RuntimeError as exc: - assert "missing-runtime" in str(exc) - else: - raise AssertionError("missing active runtime key did not fail") - finally: - restore_env(old_env) - print("SGLang GPU script self-test passed") - return 0 - - -def assert_sglang_gpu_detected(detect: dict[str, Any]) -> None: - devices = detect.get("available_devices", []) - rocm_gpu = next((item for item in devices if item.get("kind") == "rocm_gpu"), None) - if not detect.get("installed") or not rocm_gpu or not rocm_gpu.get("available"): - raise RuntimeError( - "SGLang ROCm GPU runtime was not detected through a managed TheRock " - "runtime; no CPU fallback is allowed:\n" + json.dumps(detect, indent=2) - ) - if not detect.get("managed_env"): - raise RuntimeError( - "SGLang acceptance requires a rocm-cli managed TheRock runtime manifest" - ) - - -def start_sglang( - *, - engine: Path, - model: str, - runtime_id: str, - args: argparse.Namespace, - env: dict[str, str], - repo_root: Path, -) -> tuple[subprocess.Popen[bytes] | None, Path, Path]: - if args.launch_mode == "launch": - return start_launch(engine, model, runtime_id, args, env) - - state_path = repo_root / "target" / "test-state" / f"{args.service_id}.json" - log_path = repo_root / "target" / "test-logs" / f"{args.service_id}.log" - command = [ - str(engine), - "serve-http", - args.service_id, - model, - "--host", - args.host, - "--port", - str(args.port), - "--device-policy", - "gpu_required", - "--runtime-id", - runtime_id, - "--state-path", - str(state_path), - ] - state_path.parent.mkdir(parents=True, exist_ok=True) - log_path.parent.mkdir(parents=True, exist_ok=True) - if state_path.exists(): - state_path.unlink() - log_file = log_path.open("wb") - try: - process = subprocess.Popen( - command, - env=env, - stdin=subprocess.DEVNULL, - stdout=log_file, - stderr=subprocess.STDOUT, - ) - finally: - log_file.close() - return process, state_path, log_path - - -def start_launch( - engine: Path, - model: str, - runtime_id: str, - args: argparse.Namespace, - env: dict[str, str], -) -> tuple[None, Path, Path]: - command = [ - str(engine), - "launch", - args.service_id, - model, - "--host", - args.host, - "--port", - str(args.port), - "--device-policy", - "gpu_required", - "--runtime-id", - runtime_id, - ] - started = time.monotonic() - response = run_json(command, env=env, timeout=args.timeout) - elapsed = time.monotonic() - started - if elapsed > 10: - raise RuntimeError( - f"launch took {elapsed:.1f}s; background launch should return promptly" - ) - return None, Path(response["state_path"]), Path(response["log_path"]) - - -def verify_managed_env(state: dict[str, Any]) -> dict[str, str]: - pid = state.get("server_pid") or state.get("pid") - runtime_env = state.get("therock_runtime_env") or {} - runtime_root = runtime_env.get("root") - runtime_bin = runtime_env.get("bin") - if not pid or not runtime_root: - raise RuntimeError( - "state is missing server_pid/pid or therock_runtime_env.root" - ) - environ_path = Path("/proc") / str(int(pid)) / "environ" - if not environ_path.is_file(): - raise RuntimeError(f"process environ file was not found: {environ_path}") - entries = environ_path.read_bytes().split(b"\0") - env: dict[str, str] = {} - for entry in entries: - if not entry or b"=" not in entry: - continue - key, value = entry.split(b"=", 1) - env[key.decode("utf-8", errors="replace")] = value.decode( - "utf-8", errors="replace" - ) - - expected = { - "ROCM_SDK_ROOT": runtime_root, - "ROCM_PATH": runtime_root, - "ROCM_HOME": runtime_root, - "HIP_PATH": runtime_root, - } - if runtime_bin: - expected["ROCM_CLI_THEROCK_SDK_BIN"] = runtime_bin - for key, value in expected.items(): - if env.get(key) != value: - raise RuntimeError( - f"managed TheRock env mismatch for {key}: expected {value!r}, got {env.get(key)!r}" - ) - expected_runtime_id = runtime_env.get("runtime_id") - if env.get("ROCM_CLI_THEROCK_RUNTIME_ID") != expected_runtime_id: - raise RuntimeError( - "ROCM_CLI_THEROCK_RUNTIME_ID does not match managed runtime state" - ) - return {key: env[key] for key in expected if key in env} - - -def verify_loaded_modules( - state: dict[str, Any], math_modules: list[str] -) -> dict[str, str]: - pid = state.get("server_pid") or state.get("pid") - runtime_env = state.get("therock_runtime_env") or {} - runtime_root = runtime_env.get("root") - if not pid or not runtime_root: - raise RuntimeError( - "state is missing server_pid/pid or therock_runtime_env.root" - ) - maps_path = Path("/proc") / str(int(pid)) / "maps" - if not maps_path.is_file(): - raise RuntimeError(f"process maps file was not found: {maps_path}") - - module_paths: dict[str, str] = {} - for line in maps_path.read_text(encoding="utf-8", errors="replace").splitlines(): - if "/" not in line: - continue - path = line.split(maxsplit=5)[-1] - name = Path(path).name.lower() - for module in [HIP_MODULE, *math_modules]: - if name.startswith(f"{module.lower()}.so"): - module_paths.setdefault(module, path) - if HIP_MODULE not in module_paths: - raise RuntimeError(f"missing loaded HIP module: {HIP_MODULE}") - math_loaded = sorted(set(math_modules) & set(module_paths)) - if not math_loaded: - raise RuntimeError( - f"missing loaded ROCm math module; expected one of {math_modules}" - ) - - roots = managed_therock_module_roots(Path(runtime_root)) - for module, path in module_paths.items(): - lower = str(Path(path)).lower() - if not any(lower.startswith(root) for root in roots): - raise RuntimeError( - f"{module} did not load from managed TheRock SDK wheel directories: {path}" - ) - return module_paths - - -def stop_service(process: subprocess.Popen[bytes] | None, state_path: Path) -> None: - state: dict[str, Any] = {} - if state_path.is_file(): - state = json.loads(state_path.read_text(encoding="utf-8")) - process_pid = process.pid if process else None - for pid in [state.get("server_pid"), state.get("pid"), process_pid]: - if not pid: - continue - subprocess.run( - ["kill", str(int(pid))], - stdout=subprocess.DEVNULL, - stderr=subprocess.DEVNULL, - check=False, - ) - time.sleep(0.5) - for pid in [state.get("server_pid"), state.get("pid"), process_pid]: - if not pid: - continue - subprocess.run( - ["kill", "-9", str(int(pid))], - stdout=subprocess.DEVNULL, - stderr=subprocess.DEVNULL, - check=False, - ) - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/scripts/smoke_local.py b/scripts/smoke_local.py index 0ab099fb..a603fdbb 100644 --- a/scripts/smoke_local.py +++ b/scripts/smoke_local.py @@ -127,96 +127,11 @@ def binary_paths( return { "rocm": binary_dir / exe_name("rocm"), "rocmd": binary_dir / exe_name("rocmd"), - "pytorch": binary_dir / exe_name("rocm-engine-pytorch"), - "llama": binary_dir / exe_name("rocm-engine-llama-cpp"), "lemonade": binary_dir / exe_name("rocm-engine-lemonade"), - "atom": binary_dir / exe_name("rocm-engine-atom"), "vllm": binary_dir / exe_name("rocm-engine-vllm"), - "sglang": binary_dir / exe_name("rocm-engine-sglang"), } -def create_fake_llama_server(smoke_root: Path) -> Path: - fake_dir = smoke_root / "fake-bin" - fake_dir.mkdir(parents=True, exist_ok=True) - server_py = fake_dir / "fake_llama_server.py" - server_py.write_text( - """ -from __future__ import annotations - -import http.server -import json -import os -import socketserver -import sys -import time - - -def arg_value(name: str, default: str) -> str: - if name not in sys.argv: - return default - index = sys.argv.index(name) - if index + 1 >= len(sys.argv): - return default - return sys.argv[index + 1] - - -host = arg_value("--host", "127.0.0.1") -port = int(arg_value("--port", "11435")) -ready_port = int(os.environ.get("ROCM_CLI_FAKE_LLAMA_READY_PORT", "11437")) -print("fake llama-server " + " ".join(sys.argv[1:]), flush=True) - -if port != ready_port: - raise SystemExit(0) - - -class Handler(http.server.BaseHTTPRequestHandler): - def log_message(self, format: str, *args: object) -> None: - return - - def do_GET(self) -> None: - if self.path == "/health": - self.send_response(200) - self.end_headers() - self.wfile.write(b"OK") - return - if self.path == "/v1/models": - payload = json.dumps({"data": [{"id": "tiny.gguf"}]}).encode("utf-8") - self.send_response(200) - self.send_header("content-type", "application/json") - self.send_header("content-length", str(len(payload))) - self.end_headers() - self.wfile.write(payload) - return - self.send_response(404) - self.end_headers() - - -socketserver.TCPServer.allow_reuse_address = True -with socketserver.TCPServer((host, port), Handler) as httpd: - httpd.timeout = 0.2 - deadline = time.time() + 10 - while time.time() < deadline: - httpd.handle_request() -""".lstrip(), - encoding="utf-8", - ) - if platform.system() == "Windows": - path = fake_dir / "llama-server.cmd" - path.write_text( - f'@echo off\r\n"{sys.executable}" "{server_py}" %*\r\n', - encoding="utf-8", - ) - else: - path = fake_dir / "llama-server" - path.write_text( - f'#!/usr/bin/env sh\nexec "{sys.executable}" "{server_py}" "$@"\n', - encoding="utf-8", - ) - path.chmod(0o755) - return path - - def main() -> int: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--profile", choices=["debug", "release"], default="debug") @@ -249,15 +164,10 @@ def main() -> int: shutil.rmtree(smoke_root) smoke_root.mkdir(parents=True, exist_ok=True) reject_port = free_tcp_port() - env["ROCM_CLI_LLAMA_CPP_SERVER"] = str(create_fake_llama_server(smoke_root)) rocm = str(paths["rocm"]) rocmd = str(paths["rocmd"]) - pytorch = str(paths["pytorch"]) - llama = str(paths["llama"]) - atom = str(paths["atom"]) vllm = str(paths["vllm"]) - sglang = str(paths["sglang"]) version = run("rocm version", [rocm, "version"], env=env) assert_contains(version, "rocm ", "rocm version") @@ -269,11 +179,8 @@ def main() -> int: assert_contains(examine, "managed_services: 0", "rocm examine first-run state") engines = run("rocm engines list", [rocm, "engines", "list"], env=env) - assert_contains(engines, "llama.cpp", "rocm engines list") - assert_contains(engines, "pytorch", "rocm engines list") - assert_contains(engines, "atom", "rocm engines list") + assert_contains(engines, "lemonade", "rocm engines list") assert_contains(engines, "vllm", "rocm engines list") - assert_contains(engines, "sglang", "rocm engines list") telemetry = run( "rocm config set telemetry off", @@ -286,25 +193,9 @@ def main() -> int: assert_contains(config_show, "telemetry_mode: off", "config telemetry off") assert_contains(config_show, "telemetry_policy: disabled", "config telemetry off") - llama_install = parse_json( - run( - "llama.cpp direct external install probe", - [llama, "install", "--runtime-id", "external"], - env=env, - ), - "llama.cpp direct external install probe", - ) - if ( - not isinstance(llama_install, dict) - or llama_install.get("runtime_kind") != "external_llama_server" - or llama_install.get("managed_env") is not False - or "python_executable:" in json.dumps(llama_install) - ): - fail(f"unexpected llama.cpp external install probe: {llama_install}") - engine_install = run( "rocm engines install requires exact runtime", - [rocm, "engines", "install", "llama.cpp"], + [rocm, "engines", "install", "vllm"], env=env, expect_failure=True, ) @@ -376,11 +267,11 @@ def main() -> int: ) plan = run( - "rocm freeform llama plan", - [rocm, "serve qwen with llama.cpp"], + "rocm freeform vllm plan", + [rocm, "serve qwen with vllm"], env=env, ) - assert_contains(plan, "engine: llama.cpp", "rocm freeform plan") + assert_contains(plan, "engine: vllm", "rocm freeform plan") assert_contains( plan, "no CPU fallback is implied", @@ -463,101 +354,6 @@ def main() -> int: prefetch_failure, "prefetch_artifact requires", "sandbox prefetch validation" ) - parse_json(run("pytorch detect", [pytorch, "detect"], env=env), "pytorch detect") - pytorch_capabilities = parse_json( - run("pytorch capabilities", [pytorch, "capabilities"], env=env), - "pytorch capabilities", - ) - if not isinstance(pytorch_capabilities, dict) or not pytorch_capabilities.get( - "openai_compatible" - ): - fail( - f"pytorch capabilities did not report OpenAI-compatible serving: {pytorch_capabilities}" - ) - - pytorch_model = run( - "pytorch resolve qwen", [pytorch, "resolve-model", "qwen"], env=env - ) - assert_contains(pytorch_model, "Qwen/Qwen2.5-1.5B-Instruct", "pytorch resolve qwen") - qwen35_failure = run( - "pytorch reject qwen3.5", - [pytorch, "resolve-model", "qwen3.5"], - env=env, - expect_failure=True, - ) - assert_contains( - qwen35_failure, - "not supported by the managed PyTorch engine", - "pytorch reject qwen3.5", - ) - tiny_pytorch = parse_json( - run( - "pytorch resolve tiny gpu recipe", - [pytorch, "resolve-model", "tiny-gpt2"], - env=env, - ), - "pytorch resolve tiny-gpt2", - ) - if ( - not isinstance(tiny_pytorch, dict) - or tiny_pytorch.get("canonical_model_id") != "sshleifer/tiny-gpt2" - or tiny_pytorch.get("device_policy") != "gpu_required" - or tiny_pytorch.get("dtype") != "float16" - ): - fail( - f"pytorch tiny-gpt2 did not resolve as GPU-required recipe: {tiny_pytorch}" - ) - - parse_json(run("llama.cpp detect", [llama, "detect"], env=env), "llama.cpp detect") - llama_capabilities = parse_json( - run("llama.cpp capabilities", [llama, "capabilities"], env=env), - "llama.cpp capabilities", - ) - if ( - not isinstance(llama_capabilities, dict) - or not llama_capabilities.get("openai_compatible") - or llama_capabilities.get("quantized_models") != "gguf" - ): - fail( - f"llama.cpp capabilities did not report expected GGUF/OpenAI support: {llama_capabilities}" - ) - - llama_model = run( - "llama.cpp resolve gguf", [llama, "resolve-model", "tiny.gguf"], env=env - ) - assert_contains(llama_model, "tiny.gguf", "llama.cpp resolve gguf") - llama_cpu = run( - "llama.cpp reject cpu", - [llama, "launch", "smoke-cpu", "tiny.gguf", "--device-policy", "cpu_only"], - env=env, - expect_failure=True, - ) - assert_contains(llama_cpu, "no CPU fallback is used", "llama.cpp reject cpu") - - parse_json(run("atom detect", [atom, "detect"], env=env), "atom detect") - atom_capabilities = parse_json( - run("atom capabilities", [atom, "capabilities"], env=env), - "atom capabilities", - ) - if ( - not isinstance(atom_capabilities, dict) - or not atom_capabilities.get("openai_compatible") - or atom_capabilities.get("cpu") - ): - fail( - f"atom capabilities did not report GPU-only OpenAI serving: {atom_capabilities}" - ) - - atom_model = run("atom resolve qwen", [atom, "resolve-model", "qwen"], env=env) - assert_contains(atom_model, "qwen", "atom resolve qwen") - atom_cpu = run( - "atom reject cpu", - [atom, "resolve-model", "qwen", "--device-policy", "cpu_only"], - env=env, - expect_failure=True, - ) - assert_contains(atom_cpu, "no CPU fallback is used", "atom reject cpu") - parse_json(run("vllm detect", [vllm, "detect"], env=env), "vllm detect") vllm_capabilities = parse_json( run("vllm capabilities", [vllm, "capabilities"], env=env), @@ -582,40 +378,14 @@ def main() -> int: ) assert_contains(vllm_cpu, "no CPU fallback is used", "vllm reject cpu") - parse_json(run("sglang detect", [sglang, "detect"], env=env), "sglang detect") - sglang_capabilities = parse_json( - run("sglang capabilities", [sglang, "capabilities"], env=env), - "sglang capabilities", - ) - if ( - not isinstance(sglang_capabilities, dict) - or not sglang_capabilities.get("openai_compatible") - or sglang_capabilities.get("cpu") - ): - fail( - f"sglang capabilities did not report GPU-only OpenAI serving: {sglang_capabilities}" - ) - - sglang_model = run( - "sglang resolve qwen", [sglang, "resolve-model", "qwen"], env=env - ) - assert_contains(sglang_model, "qwen", "sglang resolve qwen") - sglang_cpu = run( - "sglang reject cpu", - [sglang, "resolve-model", "qwen", "--device-policy", "cpu_only"], - env=env, - expect_failure=True, - ) - assert_contains(sglang_cpu, "no CPU fallback is used", "sglang reject cpu") - - llama_gpu_required = run( - "llama.cpp reject required gpu", + vllm_gpu_required = run( + "vllm reject required gpu", [ rocm, "serve", - "tiny.gguf", + "qwen", "--engine", - "llama.cpp", + "vllm", "--device", "gpu_required", "--foreground", @@ -625,21 +395,21 @@ def main() -> int: env=env, expect_failure=True, ) - assert_contains(llama_gpu_required, "gpu_required", "llama.cpp reject required gpu") + assert_contains(vllm_gpu_required, "gpu_required", "vllm reject required gpu") assert_not_contains( - llama_gpu_required, + vllm_gpu_required, "CPU fallback", - "llama.cpp reject required gpu", + "vllm reject required gpu", ) - llama_cpu_serve = run( - "rocm llama.cpp reject cpu serve", + vllm_cpu_serve = run( + "rocm vllm reject cpu serve", [ rocm, "serve", - "tiny.gguf", + "qwen", "--engine", - "llama.cpp", + "vllm", "--device", "cpu", "--foreground", @@ -650,9 +420,9 @@ def main() -> int: expect_failure=True, ) assert_contains( - llama_cpu_serve, + vllm_cpu_serve, "CPU mode is not a fallback path", - "rocm llama.cpp reject cpu serve", + "rocm vllm reject cpu serve", ) assert_path_missing( diff --git a/scripts/therock_sdk_install_test.py b/scripts/therock_sdk_install_test.py index 50b38092..5e0db569 100644 --- a/scripts/therock_sdk_install_test.py +++ b/scripts/therock_sdk_install_test.py @@ -444,18 +444,18 @@ def main() -> int: if not args.skip_build: cargo = resolve_cargo() run( - "build rocm and llama.cpp adapter", - [cargo, "build", "-p", "rocm", "-p", "rocm-engine-llama-cpp"], + "build rocm and vLLM adapter", + [cargo, "build", "-p", "rocm", "-p", "rocm-engine-vllm"], env=os.environ.copy(), timeout=1200, ) rocm = binary_path(args.profile, args.target_dir, "rocm") - llama = binary_path(args.profile, args.target_dir, "rocm-engine-llama-cpp") + vllm = binary_path(args.profile, args.target_dir, "rocm-engine-vllm") if not rocm.is_file(): fail(f"missing rocm binary: {rocm}") - if not llama.is_file(): - fail(f"missing llama.cpp adapter binary: {llama}") + if not vllm.is_file(): + fail(f"missing vLLM adapter binary: {vllm}") bootstrap_python = ensure_bootstrap_python(test_root, args.python) env = isolated_env(test_root, bootstrap_python, args.family) @@ -561,16 +561,16 @@ def main() -> int: ) verify_rocm_manifest_packages(manifest) - llama_detect = run( - "llama.cpp adapter detects TheRock HIP env", - [str(llama), "detect"], + vllm_detect = run( + "vLLM adapter detects the managed TheRock runtime", + [str(vllm), "detect"], env=env, timeout=120, ) assert_contains( - llama_detect, - "TheRock HIP runtime env available", - "llama.cpp TheRock HIP env detection", + vllm_detect, + "external_vllm", + "vLLM TheRock runtime detection", ) print("therock-sdk-install: ok") From a6da53d8f01a67c00be7763392795d5a368872cd Mon Sep 17 00:00:00 2001 From: Roman Inflianskas Date: Mon, 6 Jul 2026 14:57:49 +0000 Subject: [PATCH 2/2] Fix stale engine references from review Align the Qwen2.5-1.5B recipe's loader/dtype/artifact_hint with the GGUF model Lemonade actually resolves it to, and drop stale mentions of removed serving engines (SGLang, PyTorch, atom, llama.cpp) from a recipe warning string, the serve-summary module doc, and the CLI assistant skill doc. Signed-off-by: Roman Inflianskas --- apps/rocm/src/serve_summary.rs | 8 ++++---- crates/rocm-core/src/lib.rs | 8 ++++---- skills/rocm-cli-assistant/SKILL.md | 8 ++++---- 3 files changed, 12 insertions(+), 12 deletions(-) diff --git a/apps/rocm/src/serve_summary.rs b/apps/rocm/src/serve_summary.rs index de50d757..d1d176fb 100644 --- a/apps/rocm/src/serve_summary.rs +++ b/apps/rocm/src/serve_summary.rs @@ -12,10 +12,10 @@ //! //! Everything here operates at the CLI/HTTP layer *above* the per-engine //! adapters, so the summary shape is identical for every serving engine -//! (lemonade, vLLM, SGLang, PyTorch, llama.cpp, atom). Only two things vary by -//! engine, and both are already normalized elsewhere: the health path used for -//! readiness, and whether the server reports token usage (which only affects -//! whether throughput is exact, approximated, or `n/a`). +//! (lemonade, vLLM). Only two things vary by engine, and both are already +//! normalized elsewhere: the health path used for readiness, and whether the +//! server reports token usage (which only affects whether throughput is +//! exact, approximated, or `n/a`). use std::fmt::Write as _; use std::io::{IsTerminal, Write}; diff --git a/crates/rocm-core/src/lib.rs b/crates/rocm-core/src/lib.rs index 3c931359..9ac79076 100644 --- a/crates/rocm-core/src/lib.rs +++ b/crates/rocm-core/src/lib.rs @@ -5327,15 +5327,15 @@ pub fn builtin_model_recipes() -> Vec { task: "chat".to_owned(), source: "recipe_index".to_owned(), revision: "main".to_owned(), - loader: "transformers".to_owned(), + loader: "llamacpp".to_owned(), trust_remote_code: false, - dtype: "float16".to_owned(), + dtype: "gguf".to_owned(), device_policy: "gpu_required".to_owned(), min_gpu_mem_gb: Some(6), recommended_system_ram_gb: Some(8), quantization: Some("none; recommended small assistant recipe".to_owned()), artifact_hint: Some( - "Hugging Face model id; selected as the built-in low-VRAM assistant path" + "Lemonade model id; resolved internally to a GGUF build of this model" .to_owned(), ), artifacts: Vec::new(), @@ -5470,7 +5470,7 @@ pub fn builtin_model_recipes() -> Vec { chat_template_mode: "auto".to_owned(), preferred_engines: vec!["vllm".to_owned()], warnings: vec![ - "not a verified PyTorch smoke path: Transformers 4.57.6 reports unknown architecture qwen3_5" + "not a verified vLLM smoke path: Transformers 4.57.6 reports unknown architecture qwen3_5" .to_owned(), ], }, diff --git a/skills/rocm-cli-assistant/SKILL.md b/skills/rocm-cli-assistant/SKILL.md index 1d3b532b..a68b86bd 100644 --- a/skills/rocm-cli-assistant/SKILL.md +++ b/skills/rocm-cli-assistant/SKILL.md @@ -11,7 +11,7 @@ Use this skill when answering ROCm CLI local assistant questions. ## Status And Running Questions - Inspect before answering. For "is X running", "what is running", status, or port questions, call a read-only tool first. -- Use `services list --all` for vLLM, SGLang, PyTorch, Lemonade, llama.cpp, qwen, and general local model servers. +- Use `services list --all` for vLLM, Lemonade, qwen, and general local model servers. - Use `comfyui status` or `port_status` for ComfyUI and port 8188. - Interpret `running_state=running` as running, `running_state=starting` as starting, `running_state=not_running` as not running, and no matching row as unknown or not managed by ROCm CLI. - Treat `localhost` and `127.0.0.1` as the same loopback endpoint. @@ -24,11 +24,11 @@ Use this skill when answering ROCm CLI local assistant questions. ## Engines And Assistant -- vLLM, SGLang, PyTorch, Lemonade, and llama.cpp are serving engines. -- The built-in assistant is fixed to qwen served by Lemonade with GPU required. Do not switch the built-in assistant to vLLM or SGLang. +- vLLM and Lemonade are serving engines. +- The built-in assistant is fixed to qwen served by Lemonade with GPU required. Do not switch the built-in assistant to vLLM. - Installing an engine and running a model server are different states. Answer each separately when the user asks. - `rocm serve` accepts `--gpu auto|` to pick the AMD GPU. `auto` (default) prefers a GPU that looks idle from `amd-smi` VRAM telemetry 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. -- On native Windows, vLLM and SGLang serving/install live checks are skipped; tell the user to use WSL/Linux for those ROCm GPU engines and do not suggest CPU fallback. +- On native Windows, vLLM serving/install live checks are skipped; tell the user to use WSL/Linux for that ROCm GPU engine and do not suggest CPU fallback. ## ComfyUI