feat(bootstrap): opt-in AMD ROCm torch install (#124) [needs AMD-hardware validation] - #154
Conversation
Detection already routes ROCm through torch.cuda (get_best_device + _configure_rocm_if_needed), but the default install ships the CUDA torch build, so AMD-only machines fall back to CPU. Add an opt-in post-sync step: when OMNIVOICE_TORCH_VARIANT=rocm is set, the bootstrap reinstalls torch/torchaudio from the ROCm wheel index (default https://download.pytorch.org/whl/rocm6.2, overridable via OMNIVOICE_TORCH_INDEX). Strictly gated — default (unset) leaves the CUDA/CPU path untouched — and non-fatal: a failed ROCm reinstall keeps the working default build and points the user at docs/install/linux.md. rocm_opt_in() + rocm_torch_reinstall_args() are pure and unit-tested (gating + index override + arg shape). cargo test green. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
📝 WalkthroughWalkthroughThe PR adds opt-in AMD/ROCm PyTorch installation support to the Tauri bootstrap module. It introduces environment variable detection, a helper to build reinstall arguments targeting a ROCm wheel index, and conditional post-sync reinstall logic that handles failures non-fatally. Test coverage validates environment matching, index resolution, and argument construction. ChangesROCm PyTorch Variant Installation
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Possibly related PRs
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
| Filename | Overview |
|---|---|
| frontend/src-tauri/src/bootstrap.rs | Adds opt-in ROCm torch reinstall after first-run uv sync, but the step is unreachable from the repair path and the existing-healthy-venv path, so a venv repair silently reverts a ROCm user to the CUDA/CPU build; tests correctly cover the gate logic but use non-thread-safe env mutation. |
Flowchart
%%{init: {'theme': 'neutral'}}%%
flowchart TD
A([ensure_venv_ready]) --> B{venv_py exists\nand backend_dir exists?}
B -- Yes --> C{uvicorn\nimportable?}
C -- Yes --> D[Sync source dirs\nfrom bundle]
D --> E([return early — ROCm block\nnever reached])
C -- No --> F[uv sync frozen\nrepair path]
F -- success --> G([return early — ROCm block\nnever reached])
F -- fail --> H([fail + return None])
B -- No --> I[First-run bootstrap:\ncopy bundle resources]
I --> J[uv venv\nmirror cascade]
J --> K[uv sync --frozen]
K -- fail --> L([fail + return None])
K -- success --> M{OMNIVOICE_TORCH_VARIANT\n== rocm?}
M -- No / unset --> N([return Some venv_py, backend_dir])
M -- Yes --> O[uv pip install --reinstall\ntorch torchaudio\n--index-url rocm_url]
O -- success --> P([return Some venv_py, backend_dir])
O -- fail --> Q[log warn + emit hint\nkeep default build]
Q --> P
style E fill:#f9a,stroke:#c33
style G fill:#f9a,stroke:#c33
style O fill:#9cf,stroke:#06c
Comments Outside Diff (1)
-
frontend/src-tauri/src/bootstrap.rs, line 391-409 (link)Repair path silently downgrades ROCm torch back to CUDA
The venv-repair branch (triggered when
uvicornis no longer importable) runsuv sync [--frozen]which re-installstorchfrom the lockfile — the CUDA-flavoured wheel — and returns early at line 406 without ever reaching the ROCm reinstall block. Any user withOMNIVOICE_TORCH_VARIANT=rocmwhose venv is repaired will silently regress to CPU/CUDA inference with no log warning and no user-visible hint. The non-fatal fallback message that exists in the fresh-bootstrap path is entirely absent here.
Reviews (1): Last reviewed commit: "feat(bootstrap): opt-in AMD ROCm torch i..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| fn rocm_opt_in_gates_strictly_on_the_env_var() { | ||
| // This test owns OMNIVOICE_TORCH_VARIANT / _INDEX for its duration; no | ||
| // other test reads them. | ||
| std::env::remove_var("OMNIVOICE_TORCH_VARIANT"); | ||
| std::env::remove_var("OMNIVOICE_TORCH_INDEX"); | ||
| assert!(rocm_opt_in().is_none(), "unset → no ROCm (default CUDA/CPU path)"); | ||
|
|
||
| std::env::set_var("OMNIVOICE_TORCH_VARIANT", "cuda"); | ||
| assert!(rocm_opt_in().is_none(), "non-rocm value → no ROCm"); | ||
|
|
||
| std::env::set_var("OMNIVOICE_TORCH_VARIANT", "ROCm"); | ||
| assert_eq!(rocm_opt_in().as_deref(), Some(ROCM_TORCH_INDEX), "case-insensitive opt-in → default index"); | ||
|
|
||
| std::env::set_var("OMNIVOICE_TORCH_INDEX", "https://example.test/rocm6.3"); | ||
| assert_eq!(rocm_opt_in().as_deref(), Some("https://example.test/rocm6.3"), "index override honored"); | ||
|
|
||
| std::env::remove_var("OMNIVOICE_TORCH_VARIANT"); | ||
| std::env::remove_var("OMNIVOICE_TORCH_INDEX"); | ||
| } |
There was a problem hiding this comment.
std::env::set_var / remove_var are not thread-safe in multi-threaded test runs
cargo test runs unit tests on multiple OS threads by default. POSIX setenv/unsetenv (which back these calls) are explicitly documented as not async-signal-safe and can corrupt the environment block when called concurrently with any env::var or env::vars read in another thread. The comment "no other test reads them" mitigates the most obvious race, but a global-env read anywhere in the process (e.g., from a dependency, the Tauri runtime, or a future test) can still race. Consider guarding the test with a Mutex over a module-level OnceLock, using the serial_test crate's #[serial] attribute, or running this test isolated with -- --test-threads=1 in CI.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@frontend/src-tauri/src/bootstrap.rs`:
- Around line 557-577: Existing venvs skip the ROCm reinstall because the
rocm_opt_in() block only runs on the fresh-first-run path; add an idempotent
helper (e.g., ensure_rocm_reinstall_if_needed) that detects whether the venv's
installed torch is already a ROCm build (inspect torch.__version__ /
torch.version.hip or run a small Python probe inside the venv) and only when
ROCm is requested by rocm_opt_in() and the probe shows a non-ROCm torch, invoke
the same reinstall flow used in the fresh path (reuse Command setup,
apply_uv_http_env, rocm_torch_reinstall_args and run_streaming). Call this
helper from the existing-venv branches in ensure_venv_ready (and the repair
path) before returning so users who set OMNIVOICE_TORCH_VARIANT=rocm get the
reinstall only when necessary.
- Around line 614-633: The test rocm_opt_in_gates_strictly_on_the_env_var
mutates process-global env vars via std::env::set_var/remove_var (used by
rocm_opt_in and ROCM_TORCH_INDEX) which is unsafe under parallel tests; fix by
either marking the test serial (e.g., use a serial_test attribute or run with
--test-threads=1 so it runs alone) or refactor rocm_opt_in to accept an injected
environment/config map (add a rocm_opt_in_from_env(&dyn Env) or similar and call
that from the test with a controlled HashMap) and update the test to use the
injected variant instead of touching std::env.
- Around line 562-577: Wrap the ROCm reinstall logic in a Linux-only guard so
the rocm_opt_in() path cannot run on macOS/Windows: surround the existing if let
Some(rocm_url) = rocm_opt_in() { ... } block with a platform check (e.g. if
cfg!(target_os = "linux") { ... } or use #[cfg(target_os = "linux")] on the
block) so the calls to rocm_opt_in(), creation of rocm_cmd, apply_uv_http_env,
rocm_torch_reinstall_args, run_streaming and emit_log only execute on Linux.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 9fe5512a-5a56-42af-96b5-e70dcc381399
📒 Files selected for processing (1)
frontend/src-tauri/src/bootstrap.rs
| // Opt-in AMD ROCm (#124): the default install ships the CUDA torch build, | ||
| // so AMD-only machines fall back to CPU. If the user set | ||
| // OMNIVOICE_TORCH_VARIANT=rocm, reinstall torch/torchaudio from the ROCm | ||
| // wheel index. Non-fatal: a failure keeps the working CUDA/CPU build rather | ||
| // than breaking first-run. Default (unset) leaves everything unchanged. | ||
| if let Some(rocm_url) = rocm_opt_in() { | ||
| log::info!("OMNIVOICE_TORCH_VARIANT=rocm → reinstalling torch from {}", rocm_url); | ||
| let mut rocm_cmd = Command::new(&uv_path); | ||
| rocm_cmd.env_remove("PYTHONHOME").env_remove("PYTHONPATH").env_remove("LD_LIBRARY_PATH"); | ||
| apply_uv_http_env(&mut rocm_cmd); | ||
| rocm_cmd.args(rocm_torch_reinstall_args(&rocm_url)).current_dir(&project_dir); | ||
| let rocm_status = run_streaming(app, "installing_deps", &mut rocm_cmd); | ||
| if !matches!(rocm_status, Ok(ref s) if s.success()) { | ||
| log::warn!("ROCm torch reinstall failed ({:?}); keeping default torch build", rocm_status); | ||
| emit_log( | ||
| app, "installing_deps", | ||
| "ROCm torch reinstall failed — keeping the default torch build. \ | ||
| See docs/install/linux.md (AMD GPU) to install the ROCm wheel manually.", | ||
| ); | ||
| } | ||
| } |
There was a problem hiding this comment.
Opt-in is skipped on already-bootstrapped installs.
This block is only reached on the first-run/fresh-sync path. When a venv already exists and uvicorn imports, ensure_venv_ready returns at Line 382 (and the repair path at Line 406) before reaching here. So a user who installed earlier and then sets OMNIVOICE_TORCH_VARIANT=rocm gets no reinstall and still sees CPU — they'd have to "Clean & Retry". This also affects the PR's own validation steps (set the var, relaunch, confirm GPU). Surfacing the reinstall on the existing-venv path needs idempotency (detect a ROCm torch is already installed) so it doesn't re-run on every launch.
Want me to sketch an idempotent helper that runs the ROCm reinstall on the existing-venv path only when the installed torch isn't already a ROCm build?
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@frontend/src-tauri/src/bootstrap.rs` around lines 557 - 577, Existing venvs
skip the ROCm reinstall because the rocm_opt_in() block only runs on the
fresh-first-run path; add an idempotent helper (e.g.,
ensure_rocm_reinstall_if_needed) that detects whether the venv's installed torch
is already a ROCm build (inspect torch.__version__ / torch.version.hip or run a
small Python probe inside the venv) and only when ROCm is requested by
rocm_opt_in() and the probe shows a non-ROCm torch, invoke the same reinstall
flow used in the fresh path (reuse Command setup, apply_uv_http_env,
rocm_torch_reinstall_args and run_streaming). Call this helper from the
existing-venv branches in ensure_venv_ready (and the repair path) before
returning so users who set OMNIVOICE_TORCH_VARIANT=rocm get the reinstall only
when necessary.
| if let Some(rocm_url) = rocm_opt_in() { | ||
| log::info!("OMNIVOICE_TORCH_VARIANT=rocm → reinstalling torch from {}", rocm_url); | ||
| let mut rocm_cmd = Command::new(&uv_path); | ||
| rocm_cmd.env_remove("PYTHONHOME").env_remove("PYTHONPATH").env_remove("LD_LIBRARY_PATH"); | ||
| apply_uv_http_env(&mut rocm_cmd); | ||
| rocm_cmd.args(rocm_torch_reinstall_args(&rocm_url)).current_dir(&project_dir); | ||
| let rocm_status = run_streaming(app, "installing_deps", &mut rocm_cmd); | ||
| if !matches!(rocm_status, Ok(ref s) if s.success()) { | ||
| log::warn!("ROCm torch reinstall failed ({:?}); keeping default torch build", rocm_status); | ||
| emit_log( | ||
| app, "installing_deps", | ||
| "ROCm torch reinstall failed — keeping the default torch build. \ | ||
| See docs/install/linux.md (AMD GPU) to install the ROCm wheel manually.", | ||
| ); | ||
| } | ||
| } |
There was a problem hiding this comment.
❓ Verification inconclusive
🏁 Script executed:
#!/bin/bash
# ROCm opt-in usage + any target_os gate in bootstrap.rs and the spawn_backend caller
rg -nP -C3 '(rocm_opt_in|OMNIVOICE_TORCH_VARIANT|target_os\s*=)' --type=rust
echo '--- ensure_venv_ready call sites ---'
rg -nP -C3 '\bensure_venv_ready\s*\(' --type=rustRepository: debpalash/OmniVoice-Studio
Length of output: 21831
Gate the ROCm torch reinstall on Linux
frontend/src-tauri/src/bootstrap.rs’s ROCm reinstall block is only controlled by rocm_opt_in() and has no target_os = "linux" guard around the reinstall logic, so an opt-in path could run on macOS/Windows where ROCm wheels won’t apply. Gate the reinstall on Linux (or ensure rocm_opt_in()/caller returns None off-Linux).
🛡️ Proposed Linux-only guard
- if let Some(rocm_url) = rocm_opt_in() {
+ // ROCm wheels are Linux-only; skip the reinstall on macOS/Windows.
+ if let Some(rocm_url) = rocm_opt_in().filter(|_| cfg!(target_os = "linux")) {🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@frontend/src-tauri/src/bootstrap.rs` around lines 562 - 577, Wrap the ROCm
reinstall logic in a Linux-only guard so the rocm_opt_in() path cannot run on
macOS/Windows: surround the existing if let Some(rocm_url) = rocm_opt_in() { ...
} block with a platform check (e.g. if cfg!(target_os = "linux") { ... } or use
#[cfg(target_os = "linux")] on the block) so the calls to rocm_opt_in(),
creation of rocm_cmd, apply_uv_http_env, rocm_torch_reinstall_args,
run_streaming and emit_log only execute on Linux.
| #[test] | ||
| fn rocm_opt_in_gates_strictly_on_the_env_var() { | ||
| // This test owns OMNIVOICE_TORCH_VARIANT / _INDEX for its duration; no | ||
| // other test reads them. | ||
| std::env::remove_var("OMNIVOICE_TORCH_VARIANT"); | ||
| std::env::remove_var("OMNIVOICE_TORCH_INDEX"); | ||
| assert!(rocm_opt_in().is_none(), "unset → no ROCm (default CUDA/CPU path)"); | ||
|
|
||
| std::env::set_var("OMNIVOICE_TORCH_VARIANT", "cuda"); | ||
| assert!(rocm_opt_in().is_none(), "non-rocm value → no ROCm"); | ||
|
|
||
| std::env::set_var("OMNIVOICE_TORCH_VARIANT", "ROCm"); | ||
| assert_eq!(rocm_opt_in().as_deref(), Some(ROCM_TORCH_INDEX), "case-insensitive opt-in → default index"); | ||
|
|
||
| std::env::set_var("OMNIVOICE_TORCH_INDEX", "https://example.test/rocm6.3"); | ||
| assert_eq!(rocm_opt_in().as_deref(), Some("https://example.test/rocm6.3"), "index override honored"); | ||
|
|
||
| std::env::remove_var("OMNIVOICE_TORCH_VARIANT"); | ||
| std::env::remove_var("OMNIVOICE_TORCH_INDEX"); | ||
| } |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
fd -t f '^Cargo.toml$' | xargs rg -nP 'edition\s*=|name\s*='Repository: debpalash/OmniVoice-Studio
Length of output: 137
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# locate the test and rocm_opt_in implementation
FILE="frontend/src-tauri/src/bootstrap.rs"
# show around the provided lines for exact context
sed -n '560,690p' "$FILE" | cat -n
# find rocm_opt_in definition and any related helpers
rg -n "fn rocm_opt_in|rocm_opt_in\s*\(" "$FILE"Repository: debpalash/OmniVoice-Studio
Length of output: 4523
🌐 Web query:
Rust std::env::set_var remove_var thread safety data race undefined behavior concurrent tests
💡 Result:
In Rust, the functions std::env::set_var and std::env::remove_var are marked as unsafe because they are inherently thread-unsafe on many operating systems (particularly Unix-like systems) [1][2]. Using these functions in a multi-threaded program can lead to undefined behavior (UB), which is why they now require an unsafe block [1][3][4]. The core issue is that operating systems do not provide a thread-safe mechanism to modify the process environment while other threads are reading it [1]. Because many C libraries and even parts of the Rust standard library (such as DNS lookups via std::net::ToSocketAddrs) may call libc functions like getenv internally without notice, it is impossible to guarantee that no other thread is reading the environment when set_var or remove_var is called [1]. Consequently, a data race can occur between a write (set_var/remove_var) and an implicit read in another thread, leading to memory corruption or crashes [5][1][3]. Regarding concurrent tests, Rust runs tests in parallel by default [6]. If your tests call set_var or remove_var, they can easily trigger this undefined behavior because multiple threads will be active [1][6]. Furthermore, even if the threads were somehow synchronized, environment variables are process-wide state; changing them in one test will affect all other tests running in the same process, leading to non-deterministic test failures [6]. To safely handle environment variables in tests, you should avoid set_var and remove_var entirely in multi-threaded contexts [1]. Recommended alternatives include: 1. Running tests sequentially by using the --test-threads=1 argument with cargo test [6]. 2. Designing your code to accept environment-like configurations as arguments rather than reading them directly from the process environment [1]. 3. For integration tests that require specific environment variables, run them as separate processes rather than as threads within the same test binary [1][7].
Citations:
- 1: https://doc.rust-lang.org/std/env/fn.set_var.html
- 2: https://doc.rust-jp.rs/edition-guide/rust-2024/newly-unsafe-functions.html
- 3: https://users.rust-lang.org/t/unsafe-std-set-var-change/112704
- 4: https://users.rust-lang.org/t/why-env-set-var-unsafe-now/134233
- 5: Consider deprecating and/or modifying behavior of std::env::set_var rust-lang/rust#90308
- 6: https://dev.to/someb1oody/rust-guide-116-controlling-test-execution-parallel-and-sequential-tests-1gji
- 7:
std::env::{set_var, remove_var}is called in tests without safety documentation rust-lang/rust#148432
Avoid process-global env mutation in parallel tests (rocm_opt_in_gates_strictly_on_the_env_var)
std::env::set_var / std::env::remove_var mutate process-wide state; with cargo test running tests in parallel, a concurrent std::env::var / getenv in any thread (not just these two vars) can lead to undefined behavior and flaky failures. Rust crate edition here is 2021, so the main issue is the parallelism/UB risk, not edition gating.
Serialize this test (e.g., serial_test / --test-threads=1) or refactor the gating logic to accept an env/config value as an input instead of mutating the real process environment.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@frontend/src-tauri/src/bootstrap.rs` around lines 614 - 633, The test
rocm_opt_in_gates_strictly_on_the_env_var mutates process-global env vars via
std::env::set_var/remove_var (used by rocm_opt_in and ROCM_TORCH_INDEX) which is
unsafe under parallel tests; fix by either marking the test serial (e.g., use a
serial_test attribute or run with --test-threads=1 so it runs alone) or refactor
rocm_opt_in to accept an injected environment/config map (add a
rocm_opt_in_from_env(&dyn Env) or similar and call that from the test with a
controlled HashMap) and update the test to use the injected variant instead of
touching std::env.
1 participant
Completes the installer side of #124 (the docs side landed in #151).
What
An opt-in post-sync step in
ensure_venv_ready: whenOMNIVOICE_TORCH_VARIANT=rocmis set, the bootstrap reinstalls torch/torchaudio from the ROCm wheel index (defaulthttps://download.pytorch.org/whl/rocm6.2, overridable viaOMNIVOICE_TORCH_INDEX), replacing the default CUDA build so AMD GPUs are used instead of CPU.The detection side already works (
get_best_device()routes ROCm throughtorch.cuda;_configure_rocm_if_needed()setsHSA_OVERRIDE_GFX_VERSION) — this just installs the right wheel.Safety
rocmvalue) → no-op, the CUDA/CPU first-run path is byte-for-byte unchanged.Tests
cargo testgreen —rocm_torch_reinstall_args(arg shape + index) androcm_opt_in(strict gating: unset → None, non-rocm → None, case-insensitiverocm→ default index,OMNIVOICE_TORCH_INDEXoverride honored).I can't runtime-verify the actual ROCm install (no AMD GPU here). The logic compiles and the gate/builder are unit-tested, but please validate on an AMD/Linux box (
OMNIVOICE_TORCH_VARIANT=rocm, confirm Settings → System reports the GPU) before merging. Holding for that.No default-behavior change, no version bump.
🤖 Generated with Claude Code
Summary by CodeRabbit
New Features
Tests