From e10c649c89deb2b94b1b1adfec9f4e6936e1c471 Mon Sep 17 00:00:00 2001 From: Roman Inflianskas Date: Mon, 3 Aug 2026 12:34:09 +0000 Subject: [PATCH] fix(uv): colocate the uv cache with the ROCm CLI data dir uv was never told where to cache, so its cache stayed at $HOME/.cache/uv while managed environments are created under the data dir. When the two are on different filesystems (split /home, container overlay, a data dir moved via ROCM_CLI_DATA_DIR or --prefix), uv cannot hardlink and silently copies every file, so each environment carries a full duplicate of the ROCm SDK and torch stack. Add managed_uv_cache_dir() alongside the other managed_*_dir helpers and set UV_CACHE_DIR from it in uv_command_env(), which now takes &AppPaths so the cache always follows the data dir. An UV_CACHE_DIR already present in the environment is a deliberate choice (the e2e harness sets one to share a cache across scenarios) and is left untouched. Closes #160 Signed-off-by: Roman Inflianskas --- apps/rocm/src/comfyui.rs | 4 +- apps/rocm/src/therock.rs | 25 ++++++--- crates/rocm-core/src/lib.rs | 27 +++++----- crates/rocm-core/src/runtime.rs | 6 +++ crates/rocm-core/src/uv.rs | 92 ++++++++++++++++++++++++++++++--- engines/vllm/src/lib.rs | 2 +- 6 files changed, 128 insertions(+), 28 deletions(-) diff --git a/apps/rocm/src/comfyui.rs b/apps/rocm/src/comfyui.rs index 0d21a491..fdc1d6ea 100644 --- a/apps/rocm/src/comfyui.rs +++ b/apps/rocm/src/comfyui.rs @@ -348,6 +348,7 @@ pub(crate) fn install( let uv = ensure_uv_binary(paths) .context("failed to acquire uv binary for ComfyUI dependency install")?; run_uv_logged_command( + paths, &uv, uv_install_args(&runtime.python, &packages), Some(&runtime_env), @@ -1400,6 +1401,7 @@ fn uv_install_args(venv_python: &Path, packages: &[String]) -> Vec { } fn run_uv_logged_command( + paths: &AppPaths, uv: &Path, args: Vec, runtime_env: Option<&ComfyUiRuntimeEnvironment>, @@ -1421,7 +1423,7 @@ fn run_uv_logged_command( .stdin(Stdio::null()) .stdout(Stdio::piped()) .stderr(Stdio::piped()); - for (key, value) in uv_command_env() { + for (key, value) in uv_command_env(paths) { command.env(key, value); } if let Some(runtime_env) = runtime_env { diff --git a/apps/rocm/src/therock.rs b/apps/rocm/src/therock.rs index 204a15ae..56db7406 100644 --- a/apps/rocm/src/therock.rs +++ b/apps/rocm/src/therock.rs @@ -905,7 +905,7 @@ fn install_wheel_runtime( "Creating Python environment at {}.", install_root.display() )); - ensure_uv_venv(&uv, &python_launcher.executable, &install_root)?; + ensure_uv_venv(paths, &uv, &python_launcher.executable, &install_root)?; let env_python = venv_python_path(&install_root); progress_line(format!( @@ -920,6 +920,7 @@ fn install_wheel_runtime( } install_args.extend(therock_pip_package_specs(&resolution.package_versions)); run_uv_progress_command( + paths, &uv, install_args .iter() @@ -2214,7 +2215,12 @@ fn extract_tarball(archive_path: &Path, target_dir: &Path) -> Result<()> { ) } -fn ensure_uv_venv(uv: &Path, python_launcher: &Path, install_root: &Path) -> Result<()> { +fn ensure_uv_venv( + paths: &AppPaths, + uv: &Path, + python_launcher: &Path, + install_root: &Path, +) -> Result<()> { let env_python = venv_python_path(install_root); if env_python.is_file() { if run_command( @@ -2241,7 +2247,7 @@ fn ensure_uv_venv(uv: &Path, python_launcher: &Path, install_root: &Path) -> Res .map(String::as_str) .collect::>() .as_slice(), - &uv_command_env(), + &uv_command_env(paths), "create managed TheRock runtime virtual environment", )?; if !env_python.is_file() { @@ -2626,10 +2632,15 @@ fn run_command_with_env( bail!("{context_text}: {detail}") } -fn run_uv_progress_command(uv: &Path, args: &[&str], context_text: &str) -> Result<()> { +fn run_uv_progress_command( + paths: &AppPaths, + uv: &Path, + args: &[&str], + context_text: &str, +) -> Result<()> { let mut command = Command::new(uv); command.args(args); - for (key, value) in &uv_command_env() { + for (key, value) in &uv_command_env(paths) { command.env(key, value); } let status = command @@ -2733,7 +2744,7 @@ fn ensure_managed_python(paths: &AppPaths) -> Result { progress_line(format!("Installing Python {version} via uv...")); let status = Command::new(&uv) .args(["python", "install", &version]) - .envs(uv_command_env()) + .envs(uv_command_env(paths)) .stdin(Stdio::null()) .stdout(Stdio::inherit()) .stderr(Stdio::inherit()) @@ -2746,7 +2757,7 @@ fn ensure_managed_python(paths: &AppPaths) -> Result { progress_line(format!("Finding Python {version}...")); let output = Command::new(&uv) .args(["python", "find", &version]) - .envs(uv_command_env()) + .envs(uv_command_env(paths)) .stdin(Stdio::null()) .stdout(Stdio::piped()) .stderr(Stdio::null()) diff --git a/crates/rocm-core/src/lib.rs b/crates/rocm-core/src/lib.rs index 92e46e08..0615799c 100644 --- a/crates/rocm-core/src/lib.rs +++ b/crates/rocm-core/src/lib.rs @@ -49,21 +49,22 @@ use runtime::home_rocm_dir; pub use runtime::{ RuntimeHost, RuntimePlatform, current_executable_path, default_cache_dir, default_config_dir, default_data_dir, default_interactive_shell_program, managed_logs_dir, managed_pip_cache_dir, - managed_runtime_cache_dir, managed_tools_dir, normalize_runtime_path_for_host, - normalize_runtime_path_for_storage, normalize_runtime_path_text_for_host, - normalize_runtime_path_text_for_platform, normalize_runtime_path_text_for_storage, - platform_binary_name, prepend_runtime_path, runtime_directory_label, - runtime_drive_root_for_key, runtime_drive_roots, runtime_exe_suffix, runtime_home_dir, - runtime_install_root_is_protected, runtime_is_linux, runtime_is_windows, runtime_os_name, - runtime_path_for_child, runtime_path_for_windows_child, runtime_path_is_same_or_inside, - runtime_path_list_join, runtime_path_list_split, runtime_path_sort_key, - runtime_path_text_is_absolute_for_host, runtime_path_text_is_absolute_for_platform, - runtime_paths_equivalent, runtime_python_activation_hint, runtime_python_activation_script, - runtime_python_bin_dir_name, runtime_python_env_bin_dir, runtime_python_executable_in_env, - runtime_python_executable_name, runtime_rocm_library_filename, shell_command_for_host, + managed_runtime_cache_dir, managed_tools_dir, managed_uv_cache_dir, + normalize_runtime_path_for_host, normalize_runtime_path_for_storage, + normalize_runtime_path_text_for_host, normalize_runtime_path_text_for_platform, + normalize_runtime_path_text_for_storage, platform_binary_name, prepend_runtime_path, + runtime_directory_label, runtime_drive_root_for_key, runtime_drive_roots, runtime_exe_suffix, + runtime_home_dir, runtime_install_root_is_protected, runtime_is_linux, runtime_is_windows, + runtime_os_name, runtime_path_for_child, runtime_path_for_windows_child, + runtime_path_is_same_or_inside, runtime_path_list_join, runtime_path_list_split, + runtime_path_sort_key, runtime_path_text_is_absolute_for_host, + runtime_path_text_is_absolute_for_platform, runtime_paths_equivalent, + runtime_python_activation_hint, runtime_python_activation_script, runtime_python_bin_dir_name, + runtime_python_env_bin_dir, runtime_python_executable_in_env, runtime_python_executable_name, + runtime_rocm_library_filename, shell_command_for_host, }; pub use uv::{ - DEFAULT_UV_TIMEOUT_SECS, ensure_uv_binary, uv_binary_name, uv_command_env, + DEFAULT_UV_TIMEOUT_SECS, UV_CACHE_DIR_ENV, ensure_uv_binary, uv_binary_name, uv_command_env, uv_http_timeout_secs, uv_pip_freeze_args, uv_pip_install_base, uv_venv_args, }; diff --git a/crates/rocm-core/src/runtime.rs b/crates/rocm-core/src/runtime.rs index 940f9e97..e66668cb 100644 --- a/crates/rocm-core/src/runtime.rs +++ b/crates/rocm-core/src/runtime.rs @@ -226,6 +226,12 @@ pub fn managed_pip_cache_dir(root: &Path) -> PathBuf { normalize_runtime_path_for_host(root).join("pip-cache") } +/// `uv`'s content-addressed cache, kept under the managed root so it shares a filesystem +/// with the environments `uv` populates and hardlinking keeps working (see issue #160). +pub fn managed_uv_cache_dir(root: &Path) -> PathBuf { + normalize_runtime_path_for_host(root).join("uv-cache") +} + pub fn managed_logs_dir(root: &Path) -> PathBuf { normalize_runtime_path_for_host(root).join("logs") } diff --git a/crates/rocm-core/src/uv.rs b/crates/rocm-core/src/uv.rs index 7d6dbd3f..b6d432e2 100644 --- a/crates/rocm-core/src/uv.rs +++ b/crates/rocm-core/src/uv.rs @@ -13,11 +13,14 @@ use anyhow::{Context, Result, bail}; use serde::{Deserialize, Serialize}; +use std::ffi::OsStr; use std::path::{Path, PathBuf}; use std::process::Command; use std::time::Duration; -use crate::runtime::{managed_tools_dir, runtime_is_windows, runtime_os_name}; +use crate::runtime::{ + managed_tools_dir, managed_uv_cache_dir, runtime_is_windows, runtime_os_name, +}; use crate::{AppPaths, download_file_to_path, unix_time_millis}; /// Default network timeout, in seconds, applied to `uv` HTTP operations. @@ -58,13 +61,42 @@ pub fn uv_http_timeout_secs() -> u64 { .unwrap_or(DEFAULT_UV_TIMEOUT_SECS) } -/// Environment pairs to apply when spawning `uv` so network behavior is configured -/// consistently (uv reads `UV_HTTP_TIMEOUT` rather than accepting a `--timeout` flag). -pub fn uv_command_env() -> Vec<(String, String)> { - vec![( +/// Environment variable `uv` reads to locate its content-addressed cache. +pub const UV_CACHE_DIR_ENV: &str = "UV_CACHE_DIR"; + +/// Environment pairs to apply when spawning `uv`. +/// +/// Network behavior is configured consistently (uv reads `UV_HTTP_TIMEOUT` rather than +/// accepting a `--timeout` flag) and the cache lives beside the managed environments it +/// populates. +/// +/// Without a cache inside the managed root, `uv` caches under `$HOME/.cache/uv`; when +/// that is on a different filesystem from the data directory, `uv` cannot hardlink and +/// silently copies every file, so each environment carries a full duplicate of the SDK +/// and torch stack. An `UV_CACHE_DIR` already present in the environment is a deliberate +/// choice (the e2e harness shares one cache across scenarios) and is left alone. +pub fn uv_command_env(paths: &AppPaths) -> Vec<(String, String)> { + let inherited = std::env::var_os(UV_CACHE_DIR_ENV).filter(|value| !value.is_empty()); + uv_command_env_with_inherited_cache(paths, inherited.as_deref()) +} + +fn uv_command_env_with_inherited_cache( + paths: &AppPaths, + inherited_cache_dir: Option<&OsStr>, +) -> Vec<(String, String)> { + let mut env = vec![( "UV_HTTP_TIMEOUT".to_owned(), uv_http_timeout_secs().to_string(), - )] + )]; + if inherited_cache_dir.is_none() { + env.push(( + UV_CACHE_DIR_ENV.to_owned(), + managed_uv_cache_dir(&paths.data_dir) + .to_string_lossy() + .into_owned(), + )); + } + env } /// Arguments for `uv venv`, creating an environment at `env_root` using `python`. @@ -367,6 +399,54 @@ mod tests { ); } + fn test_paths(root: &str) -> AppPaths { + AppPaths { + config_dir: PathBuf::from(root).join("config"), + data_dir: PathBuf::from(root), + cache_dir: PathBuf::from(root).join("cache"), + } + } + + fn cache_dir_in(env: &[(String, String)]) -> Option { + env.iter() + .find(|(key, _)| key == UV_CACHE_DIR_ENV) + .map(|(_, value)| value.clone()) + } + + #[test] + fn command_env_cache_dir_is_derived_from_data_dir() { + let paths = test_paths("/managed/root"); + let env = uv_command_env_with_inherited_cache(&paths, None); + assert_eq!( + cache_dir_in(&env), + Some( + managed_uv_cache_dir(&paths.data_dir) + .to_string_lossy() + .into_owned() + ) + ); + assert!(env.iter().any(|(key, _)| key == "UV_HTTP_TIMEOUT")); + } + + #[test] + fn command_env_keeps_an_explicit_cache_dir() { + // The e2e harness sets UV_CACHE_DIR to share one cache across scenarios; a user can + // do the same. Such a choice must be inherited, not overridden. + let paths = test_paths("/managed/root"); + let env = uv_command_env_with_inherited_cache(&paths, Some(OsStr::new("/shared/uv-cache"))); + assert_eq!(cache_dir_in(&env), None); + } + + #[test] + fn command_env_cache_dir_follows_a_relocated_data_dir() { + let default_root = test_paths("/home/user/.rocm"); + let moved_root = test_paths("/mnt/big/rocm"); + assert_ne!( + cache_dir_in(&uv_command_env_with_inherited_cache(&default_root, None)), + cache_dir_in(&uv_command_env_with_inherited_cache(&moved_root, None)) + ); + } + #[test] fn slug_sanitizes_unexpected_characters() { assert_eq!(slug("0.8.4"), "0.8.4"); diff --git a/engines/vllm/src/lib.rs b/engines/vllm/src/lib.rs index ff083318..25faae88 100644 --- a/engines/vllm/src/lib.rs +++ b/engines/vllm/src/lib.rs @@ -939,7 +939,7 @@ fn install_vllm_with_uv(python: &Path, reinstall: bool) -> Result<()> { args.push(index_url.clone()); let output = ProcessCommand::new(&uv) .args(args) - .envs(uv_command_env()) + .envs(uv_command_env(&paths)) .output() .context("failed to launch uv pip install for vLLM")?; if output.status.success() {