diff --git a/apps/rocm/src/comfyui.rs b/apps/rocm/src/comfyui.rs index 2c2de6e3..bca35bcc 100644 --- a/apps/rocm/src/comfyui.rs +++ b/apps/rocm/src/comfyui.rs @@ -2,10 +2,10 @@ use crate::{format_structured_tool_call, runtime_usability_status, therock}; use anyhow::{Context, Result, bail}; use flate2::read::GzDecoder; use rocm_core::{ - AppPaths, RocmCliConfig, download_file_to_path, format_http_base_url, managed_pip_cache_dir, + AppPaths, RocmCliConfig, download_file_to_path, ensure_uv_binary, format_http_base_url, runtime_is_cosmopolitan_windows, runtime_is_linux, runtime_is_windows, runtime_path_for_windows_child, runtime_path_list_join, runtime_path_list_split, - runtime_paths_equivalent, unix_time_millis, + runtime_paths_equivalent, uv_command_env, uv_pip_install_base, unix_time_millis, }; use serde::{Deserialize, Serialize}; use std::ffi::OsString; @@ -51,7 +51,7 @@ struct ComfyUiManifest { source_url: String, source_path: PathBuf, requirements_path: PathBuf, - pip_cache_dir: PathBuf, + pip_cache_dir: Option, log_path: PathBuf, torch_version: Option, torch_cuda_available: bool, @@ -374,7 +374,6 @@ pub(crate) fn install( let runtime = select_runtime(paths, config, options.runtime_id.as_deref())?; let app_root = runtime_app_root(&runtime.manifest); let source_path = source_path_from_app_root(&app_root); - let pip_cache = managed_pip_cache_dir(&app_root); let log_path = install_log_path_from_app_root(&app_root); let requirements_path = source_path.join("requirements.txt"); let models_folder = models_folder_for_source(&source_path); @@ -390,7 +389,6 @@ pub(crate) fn install( )?; writeln!(output, " folder: {}", app_root.display())?; writeln!(output, " models path: {}", models_folder.display())?; - writeln!(output, " pip cache: {}", pip_cache.display())?; if options.dry_run { writeln!(output, " mode: dry-run")?; @@ -448,14 +446,14 @@ pub(crate) fn install( "Installing {} ComfyUI dependency specs.", packages.len() )?; - fs::create_dir_all(&pip_cache) - .with_context(|| format!("failed to create {}", pip_cache.display()))?; if !packages.is_empty() { println!("Installing ComfyUI dependencies..."); let _ = io::stdout().flush(); - run_logged_command( - &runtime.python, - pip_install_args(&pip_cache, &packages), + let uv = ensure_uv_binary(paths) + .context("failed to acquire uv binary for ComfyUI dependency install")?; + run_uv_logged_command( + &uv, + uv_install_args(&runtime.python, &packages), Some(&runtime_env), &mut log, "install ComfyUI dependencies", @@ -478,7 +476,7 @@ pub(crate) fn install( source_url: COMFYUI_SOURCE_ARCHIVE_URL.to_owned(), source_path: source_path.clone(), requirements_path: requirements_path.clone(), - pip_cache_dir: pip_cache.clone(), + pip_cache_dir: None, log_path: log_path.clone(), torch_version: probe.torch_version.clone(), torch_cuda_available: probe.torch_cuda_available, @@ -1149,7 +1147,7 @@ fn select_runtime( runtime_usability_status(&manifest) ); } - if manifest.format != "pip" { + if manifest.format != "wheel" { bail!("ComfyUI installs require a rocm-cli managed Python ROCm install."); } let python = manifest @@ -1521,21 +1519,15 @@ fn requirement_package_name(spec: &str) -> Option { (end > 0).then(|| trimmed[..end].replace('_', "-").to_ascii_lowercase()) } -fn pip_install_args(pip_cache: &Path, packages: &[String]) -> Vec { - let mut args = vec![ - "-m".to_owned(), - "pip".to_owned(), - "install".to_owned(), - "--upgrade".to_owned(), - "--cache-dir".to_owned(), - pip_cache.display().to_string(), - ]; +fn uv_install_args(venv_python: &Path, packages: &[String]) -> Vec { + let mut args = uv_pip_install_base(venv_python); + args.push("--upgrade".to_owned()); args.extend(packages.iter().cloned()); args } -fn run_logged_command( - program: &Path, +fn run_uv_logged_command( + uv: &Path, args: Vec, runtime_env: Option<&ComfyUiRuntimeEnvironment>, log: &mut fs::File, @@ -1544,24 +1536,27 @@ fn run_logged_command( writeln!( log, "command: {} {}", - program.display(), + uv.display(), args.iter() .map(|arg| quote_log_arg(arg)) .collect::>() .join(" ") )?; - let mut command = Command::new(program); + let mut command = Command::new(uv); command .args(&args) .stdin(Stdio::null()) .stdout(Stdio::piped()) .stderr(Stdio::piped()); + for (key, value) in uv_command_env() { + command.env(key, value); + } if let Some(runtime_env) = runtime_env { apply_runtime_environment(&mut command, runtime_env)?; } let mut child = command .spawn() - .with_context(|| format!("{context_text}: failed to run {}", program.display()))?; + .with_context(|| format!("{context_text}: failed to run {}", uv.display()))?; let stdout = child .stdout .take() @@ -1582,7 +1577,7 @@ fn run_logged_command( thread::spawn(move || stream_logged_output(stderr, stderr_log, OutputTarget::Stderr)); let status = child .wait() - .with_context(|| format!("{context_text}: failed waiting for {}", program.display()))?; + .with_context(|| format!("{context_text}: failed waiting for {}", uv.display()))?; stdout_thread .join() .map_err(|_| anyhow::anyhow!("{context_text}: stdout reader failed"))? @@ -1594,9 +1589,10 @@ fn run_logged_command( if status.success() { return Ok(()); } - bail!("{context_text}: command exited with status {status}") + bail!("{context_text}: uv exited with {status}"); } + enum OutputTarget { Stdout, Stderr, @@ -1872,7 +1868,6 @@ mod tests { let runtime_app = runtime_app_root(&runtime); assert!(rendered.contains(&runtime_app.display().to_string())); - assert!(rendered.contains(&runtime_app.join("pip-cache").display().to_string())); assert!(!rendered.contains(&app_root(&paths).display().to_string())); Ok(()) } @@ -1898,7 +1893,7 @@ mod tests { source_url: COMFYUI_SOURCE_ARCHIVE_URL.to_owned(), source_path: source_path(&paths), requirements_path: source_path(&paths).join("requirements.txt"), - pip_cache_dir: app_root(&paths).join("pip-cache"), + pip_cache_dir: None, log_path: install_log.clone(), torch_version: Some("2.10.0".to_owned()), torch_cuda_available: true, @@ -2066,7 +2061,7 @@ mod tests { source_url: COMFYUI_SOURCE_ARCHIVE_URL.to_owned(), source_path: source_path(&paths), requirements_path: source_path(&paths).join("requirements.txt"), - pip_cache_dir: app_root(&paths).join("pip-cache"), + pip_cache_dir: None, log_path: install_log.clone(), torch_version: Some("2.10.0".to_owned()), torch_cuda_available: true, @@ -2132,7 +2127,7 @@ mod tests { runtime_key: "release-pip-gfx120x-all-7-13-0a20260511".to_owned(), runtime_id: "therock-release:gfx120X-all".to_owned(), channel: "release".to_owned(), - format: "pip".to_owned(), + format: "wheel".to_owned(), family: "gfx120X-all".to_owned(), family_source: "test".to_owned(), version: "7.13.0a20260511".to_owned(), @@ -2204,7 +2199,7 @@ mod tests { runtime_key: "therock-release-gfx120x-all".to_owned(), runtime_id: "therock-release:gfx120X-all".to_owned(), channel: "release".to_owned(), - format: "pip".to_owned(), + format: "wheel".to_owned(), family: "gfx120X-all".to_owned(), family_source: "test".to_owned(), version: "7.13.0a20260511".to_owned(), @@ -2293,7 +2288,7 @@ mod tests { runtime_key: runtime_key.to_owned(), runtime_id: "therock-release:gfx120X-all".to_owned(), channel: "release".to_owned(), - format: "pip".to_owned(), + format: "wheel".to_owned(), family: "gfx120X-all".to_owned(), family_source: "test".to_owned(), version: "7.13.0a20260511".to_owned(), diff --git a/apps/rocm/src/main.rs b/apps/rocm/src/main.rs index 4f5fc810..b62ed8a6 100644 --- a/apps/rocm/src/main.rs +++ b/apps/rocm/src/main.rs @@ -279,7 +279,7 @@ enum InstallTarget { #[arg(long, default_value = "release")] channel: String, /// Package format to install. - #[arg(long, default_value = "pip")] + #[arg(long, default_value = "wheel")] format: InstallFormat, /// Full folder path where the ROCm Python environment should be created. #[arg(long)] @@ -563,7 +563,7 @@ enum SetupCommand { #[derive(Clone, Copy, Debug, Eq, PartialEq, ValueEnum)] enum InstallFormat { - Pip, + Wheel, Tarball, } @@ -1759,7 +1759,7 @@ fn render_codex_bridge_instructions( writeln!(&mut text, "- `rocm doctor`").ok(); writeln!( &mut text, - "- `rocm install sdk --channel release|nightly [--format pip|tarball] [--build-date YYYY-MM-DD|--version VERSION]`" + "- `rocm install sdk --channel release|nightly [--format wheel|tarball] [--build-date YYYY-MM-DD|--version VERSION]`" ) .ok(); writeln!(&mut text, "- `rocm update`").ok(); @@ -2005,7 +2005,7 @@ fn install(target: InstallTarget) -> Result<()> { dry_run, } => { let format_name = match format { - InstallFormat::Pip => "pip", + InstallFormat::Wheel => "wheel", InstallFormat::Tarball => "tarball", }; let version_selector = therock_install_version_selector(version, build_date)?; @@ -3424,7 +3424,7 @@ fn active_pytorch_runtime( } fn pytorch_runtime_ready(manifest: &therock::InstalledRuntimeManifest) -> bool { - manifest.format == "pip" + manifest.format == "wheel" && validate_runtime_manifest_for_activation(manifest).is_ok() && manifest .rocm_sdk @@ -4701,7 +4701,7 @@ pub(crate) fn render_runtimes_text(paths: &AppPaths, config: &RocmCliConfig) -> let _ = writeln!(output, " installed: none"); let _ = writeln!( output, - " next step: rocm install sdk --channel release --format pip" + " next step: rocm install sdk --channel release --format wheel" ); return Ok(output); } @@ -5338,7 +5338,7 @@ fn adopt_runtime_from_probe( runtime_key: request.runtime_key, runtime_id: request.runtime_id, channel, - format: "pip".to_owned(), + format: "wheel".to_owned(), family, family_source: "runtime_id".to_owned(), version, @@ -5585,13 +5585,13 @@ fn validate_runtime_manifest_for_activation( } match manifest.format.as_str() { - "pip" => validate_pip_runtime_manifest(manifest), + "wheel" => validate_wheel_runtime_manifest(manifest), "tarball" => validate_tarball_runtime_manifest(manifest), other => bail!("unsupported runtime format in manifest: {other}"), } } -fn validate_pip_runtime_manifest(manifest: &therock::InstalledRuntimeManifest) -> Result<()> { +fn validate_wheel_runtime_manifest(manifest: &therock::InstalledRuntimeManifest) -> Result<()> { let python_executable = manifest .python_executable .as_deref() @@ -6671,7 +6671,7 @@ fn fallback_rocm_tool_call_for_prompt(prompt: &str) -> Option Result if !matches!(channel.as_str(), "release" | "nightly") { bail!("local assistant requested unsupported TheRock channel `{channel}`"); } - let format = json_string(object, "format").unwrap_or_else(|| "pip".to_owned()); - if !matches!(format.as_str(), "pip" | "tarball") { + let format = json_string(object, "format").unwrap_or_else(|| "wheel".to_owned()); + if !matches!(format.as_str(), "wheel" | "tarball") { bail!("local assistant requested unsupported TheRock install format `{format}`"); } - if rocm_core::runtime_is_windows() && format != "pip" { + if rocm_core::runtime_is_windows() && format != "wheel" { bail!("local assistant cannot request `{format}` installs on Windows; use pip"); } let version = json_string(object, "version"); @@ -7501,7 +7501,7 @@ fn validate_chat_install_sdk_tool_call(call: &providers::ChatToolCall) -> Result if version.is_some() && build_date.is_some() { bail!("local assistant cannot request both `version` and `build_date`"); } - if format != "pip" && (version.is_some() || build_date.is_some()) { + if format != "wheel" && (version.is_some() || build_date.is_some()) { bail!("local assistant can only request specific TheRock wheel versions for pip installs"); } if let Some(version) = version { @@ -7901,7 +7901,7 @@ fn validate_chat_rocm_command_safety(args: &[String]) -> Result<()> { { if rocm_core::runtime_is_windows() && chat_cli_arg_value(args, "--format") - .is_some_and(|value| !value.eq_ignore_ascii_case("pip")) + .is_some_and(|value| !value.eq_ignore_ascii_case("wheel")) { bail!("local assistant cannot request non-pip ROCm installs on Windows"); } @@ -7911,8 +7911,8 @@ fn validate_chat_rocm_command_safety(args: &[String]) -> Result<()> { bail!("local assistant cannot request both --version and --build-date"); } if version.is_some() || build_date.is_some() { - let format = chat_cli_arg_value(args, "--format").unwrap_or("pip"); - if !format.eq_ignore_ascii_case("pip") { + let format = chat_cli_arg_value(args, "--format").unwrap_or("wheel"); + if !format.eq_ignore_ascii_case("wheel") { bail!( "local assistant can only request specific TheRock wheel versions for pip installs" ); @@ -9074,7 +9074,7 @@ fn parse_optional_lines(args: &[String]) -> Result { fn render_install_sdk_dry_run_for_args(paths: &AppPaths, args: &[String]) -> Result { let channel = chat_cli_arg_value(args, "--channel").unwrap_or("release"); - let format = chat_cli_arg_value(args, "--format").unwrap_or("pip"); + let format = chat_cli_arg_value(args, "--format").unwrap_or("wheel"); let prefix = chat_cli_arg_value(args, "--prefix").map(PathBuf::from); let version = chat_cli_arg_value(args, "--version").map(str::to_owned); let build_date = chat_cli_arg_value(args, "--build-date").map(str::to_owned); @@ -9126,7 +9126,7 @@ fn internal_mcp_install_sdk_args( dry_run: bool, ) -> Result> { let channel = json_string(arguments, "channel").unwrap_or_else(|| "release".to_owned()); - let format = json_string(arguments, "format").unwrap_or_else(|| "pip".to_owned()); + let format = json_string(arguments, "format").unwrap_or_else(|| "wheel".to_owned()); let mut argv = vec![ "install".to_owned(), "sdk".to_owned(), @@ -9448,7 +9448,7 @@ fn rocm_chat_tool_requested_args(call: &providers::ChatToolCall) -> Option( } match manifests { [] => bail!( - "no managed runtimes are registered; run `rocm install sdk --channel release --format pip` first" + "no managed runtimes are registered; run `rocm install sdk --channel release --format wheel` first" ), [only] => Ok(only), _ => bail!( @@ -12816,7 +12816,7 @@ fn build_freeform_plan_with_recipes( let Some(prefix) = requested_install_prefix_from_prompt(trimmed) else { let mut parsed = vec![ ("channel".to_owned(), channel.to_owned()), - ("format".to_owned(), "pip".to_owned()), + ("format".to_owned(), "wheel".to_owned()), ]; if let Some(build_date) = build_date.as_deref() { parsed.push(("build_date".to_owned(), build_date.to_owned())); @@ -12843,7 +12843,7 @@ fn build_freeform_plan_with_recipes( "--channel".to_owned(), channel.to_owned(), "--format".to_owned(), - "pip".to_owned(), + "wheel".to_owned(), "--prefix".to_owned(), prefix.clone(), ]; @@ -12853,7 +12853,7 @@ fn build_freeform_plan_with_recipes( } let mut parsed = vec![ ("channel".to_owned(), channel.to_owned()), - ("format".to_owned(), "pip".to_owned()), + ("format".to_owned(), "wheel".to_owned()), ("prefix".to_owned(), prefix), ]; if let Some(build_date) = build_date.as_deref() { @@ -15092,7 +15092,7 @@ mod tests { "--channel".to_owned(), "nightly".to_owned(), "--format".to_owned(), - "pip".to_owned(), + "wheel".to_owned(), "--prefix".to_owned(), "D:\\jam\\temp\\therock_venvs".to_owned(), ] @@ -15121,7 +15121,7 @@ mod tests { "--channel".to_owned(), "release".to_owned(), "--format".to_owned(), - "pip".to_owned(), + "wheel".to_owned(), "--prefix".to_owned(), "D:\\jam\\temp\\therock_venvs".to_owned(), "--build-date".to_owned(), @@ -15548,7 +15548,7 @@ mod tests { "--channel".to_owned(), "release".to_owned(), "--format".to_owned(), - "pip".to_owned(), + "wheel".to_owned(), ], }; let error = validate_provider_planner_tool_call(&call) @@ -15608,7 +15608,7 @@ mod tests { name: "install_sdk".to_owned(), arguments: serde_json::json!({ "channel": "release", - "format": "pip", + "format": "wheel", "prefix": "D:\\jam\\temp\\therock_venvs" }), }; @@ -15618,7 +15618,7 @@ mod tests { assert_eq!( rocm_chat_tool_requested_command(&call).as_deref(), Some( - "rocm install sdk --channel release --format pip --prefix D:\\jam\\temp\\therock_venvs" + "rocm install sdk --channel release --format wheel --prefix D:\\jam\\temp\\therock_venvs" ) ); let approval = chat_tool_approval_request( @@ -15640,7 +15640,7 @@ mod tests { "--channel".to_owned(), "release".to_owned(), "--format".to_owned(), - "pip".to_owned(), + "wheel".to_owned(), "--prefix".to_owned(), "D:\\jam\\temp\\therock_venvs".to_owned(), ] @@ -15653,7 +15653,7 @@ mod tests { id: Some("call-date".to_owned()), name: "rocm_command".to_owned(), arguments: serde_json::json!({ - "args": ["install", "sdk", "--channel", "release", "--format", "pip", "--prefix", "D:\\jam\\temp\\therock_venvs", "--build-date", "06052026"], + "args": ["install", "sdk", "--channel", "release", "--format", "wheel", "--prefix", "D:\\jam\\temp\\therock_venvs", "--build-date", "06052026"], "reason": "The user asked for the TheRock build from 2026-06-05." }), }; @@ -15663,7 +15663,7 @@ mod tests { assert_eq!( rocm_chat_tool_requested_command(&call).as_deref(), Some( - "rocm install sdk --channel release --format pip --prefix D:\\jam\\temp\\therock_venvs --build-date 06052026" + "rocm install sdk --channel release --format wheel --prefix D:\\jam\\temp\\therock_venvs --build-date 06052026" ) ); let approval = @@ -15678,7 +15678,7 @@ mod tests { "--channel".to_owned(), "release".to_owned(), "--format".to_owned(), - "pip".to_owned(), + "wheel".to_owned(), "--prefix".to_owned(), "D:\\jam\\temp\\therock_venvs".to_owned(), "--build-date".to_owned(), @@ -15694,7 +15694,7 @@ mod tests { name: "install_sdk".to_owned(), arguments: serde_json::json!({ "channel": "release", - "format": "pip" + "format": "wheel" }), }; let error = validate_chat_tool_call(&structured) @@ -15706,7 +15706,7 @@ mod tests { id: Some("call-command-missing-prefix".to_owned()), name: "rocm_command".to_owned(), arguments: serde_json::json!({ - "args": ["install", "sdk", "--channel", "release", "--format", "pip"], + "args": ["install", "sdk", "--channel", "release", "--format", "wheel"], "reason": "Install ROCm." }), }; @@ -16041,12 +16041,12 @@ model recipes name: "install_sdk".to_owned(), arguments: serde_json::json!({ "channel": "release", - "format": "pip", + "format": "wheel", "prefix": "D:\\jam\\temp\\therock_venvs" }), }, Some( - "rocm install sdk --channel release --format pip --prefix D:\\jam\\temp\\therock_venvs", + "rocm install sdk --channel release --format wheel --prefix D:\\jam\\temp\\therock_venvs", ), false, ), @@ -16477,7 +16477,7 @@ model recipes name: "install_sdk".to_owned(), arguments: serde_json::json!({ "channel": "release", - "format": "pip", + "format": "wheel", "prefix": "D:\\jam\\temp\\therock_venvs" }), }], @@ -16676,7 +16676,7 @@ model recipes "--channel".to_owned(), "release".to_owned(), "--format".to_owned(), - "pip".to_owned(), + "wheel".to_owned(), ] ); assert!(!chat_tool_call_is_read_only(&call)); @@ -16701,7 +16701,7 @@ model recipes "--channel".to_owned(), "release".to_owned(), "--format".to_owned(), - "pip".to_owned(), + "wheel".to_owned(), "--prefix".to_owned(), expected_prefix.to_owned(), ] @@ -16732,7 +16732,7 @@ model recipes "--channel".to_owned(), "release".to_owned(), "--format".to_owned(), - "pip".to_owned(), + "wheel".to_owned(), "--prefix".to_owned(), expected_prefix.to_owned(), "--build-date".to_owned(), @@ -16759,7 +16759,7 @@ model recipes "--channel".to_owned(), "release".to_owned(), "--format".to_owned(), - "pip".to_owned(), + "wheel".to_owned(), "--build-date".to_owned(), "2026-06-05".to_owned(), ] @@ -16795,7 +16795,7 @@ model recipes "--channel".to_owned(), "release".to_owned(), "--format".to_owned(), - "pip".to_owned(), + "wheel".to_owned(), ], "{prompt}" ); @@ -16878,7 +16878,7 @@ install therock"; "--channel".to_owned(), "release".to_owned(), "--format".to_owned(), - "pip".to_owned(), + "wheel".to_owned(), "--prefix".to_owned(), "D:\\jam\\temp\\therock_venvs".to_owned(), ] @@ -16900,7 +16900,7 @@ install therock"; "--channel".to_owned(), "release".to_owned(), "--format".to_owned(), - "pip".to_owned(), + "wheel".to_owned(), "--prefix".to_owned(), "D:\\jam\\temp\\therock_venvs".to_owned(), "--version".to_owned(), @@ -17701,7 +17701,7 @@ install therock"; "--channel", "release", "--format", - "pip", + "wheel", "--prefix", "D:\\jam\\temp\\therock_venvs", "--family", @@ -20275,7 +20275,7 @@ VERSION_ID="41" let install_root = paths .data_dir .join("runtimes") - .join("pip") + .join("wheel") .join(runtime_key); let scripts_dir = install_root.join(if cfg!(windows) { "Scripts" } else { "bin" }); let python_executable = scripts_dir.join(if cfg!(windows) { @@ -20305,7 +20305,7 @@ VERSION_ID="41" runtime_key: runtime_key.to_owned(), runtime_id: runtime_id.to_owned(), channel: "release".to_owned(), - format: "pip".to_owned(), + format: "wheel".to_owned(), family: "gfx120X-all".to_owned(), family_source: "test".to_owned(), version: version.to_owned(), @@ -20315,7 +20315,7 @@ VERSION_ID="41" tarball_file_name: None, python_launcher: Some("python".to_owned()), python_executable: Some(python_executable.display().to_string()), - pip_cache_dir: Some(paths.cache_dir.join("pip").join("therock")), + pip_cache_dir: Some(paths.cache_dir.join("uv").join("therock")), rocm_sdk: Some(therock::RocmSdkPythonProbe { import_ok: true, root_path: Some(sdk_root.clone()), @@ -20361,7 +20361,7 @@ VERSION_ID="41" runtime_key: runtime_key.to_owned(), runtime_id: runtime_id.to_owned(), channel: "release".to_owned(), - format: "pip".to_owned(), + format: "wheel".to_owned(), family: family.to_owned(), family_source: "test".to_owned(), version: version.to_owned(), diff --git a/apps/rocm/src/providers.rs b/apps/rocm/src/providers.rs index bc6ffb6c..e0059e06 100644 --- a/apps/rocm/src/providers.rs +++ b/apps/rocm/src/providers.rs @@ -1108,7 +1108,7 @@ fn rocm_openai_tool_definitions() -> Vec { "type": "object", "properties": { "channel": { "type": "string", "enum": ["release", "nightly"] }, - "format": { "type": "string", "enum": ["pip", "tarball"] }, + "format": { "type": "string", "enum": ["wheel", "tarball"] }, "prefix": { "type": "string" }, "version": { "type": "string" }, "build_date": { "type": "string" } @@ -1782,7 +1782,7 @@ mod tests { "name": "install_sdk_dry_run", "input": { "channel": "release", - "format": "pip" + "format": "wheel" } } ] diff --git a/apps/rocm/src/therock.rs b/apps/rocm/src/therock.rs index e5001284..27b5d79b 100644 --- a/apps/rocm/src/therock.rs +++ b/apps/rocm/src/therock.rs @@ -1,13 +1,12 @@ use anyhow::{Context, Result, bail}; -use flate2::read::GzDecoder; use rocm_core::{ AppPaths, ManagedToolConfig, RocmCliConfig, detect_host_gpu_diagnostics, - detect_host_therock_family, detect_managed_therock_family, managed_pip_cache_dir, + detect_host_therock_family, detect_managed_therock_family, ensure_uv_binary, managed_tools_dir, normalize_runtime_path_for_host, normalize_runtime_path_for_storage, normalize_runtime_path_text_for_host, normalize_runtime_path_text_for_storage, normalize_therock_family, platform_binary_name, runtime_is_windows, runtime_os_name, runtime_path_for_windows_child, runtime_path_list_split, runtime_python_executable_in_env, - unix_time_millis, + uv_command_env, uv_pip_install_base, uv_venv_args, unix_time_millis, }; use serde::{Deserialize, Serialize}; use std::cmp::Ordering; @@ -20,11 +19,7 @@ use std::time::Duration; const THEROCK_PIP_INDEX_BASE: &str = "https://rocm.nightlies.amd.com/v2"; const THEROCK_RELEASE_TARBALL_BASE: &str = "https://repo.amd.com/rocm/tarball/"; const THEROCK_NIGHTLY_TARBALL_BASE: &str = "https://rocm.nightlies.amd.com/tarball/"; -const PYTHON_BUILD_STANDALONE_DEFAULT_RELEASE_URL: &str = - "https://api.github.com/repos/astral-sh/python-build-standalone/releases/tags/20250409"; -const DEFAULT_MANAGED_PYTHON_VERSION: &str = "3.12.10"; -const DEFAULT_PIP_TIMEOUT_SECS: u64 = 600; -const DEFAULT_PIP_RETRIES: u32 = 8; +const DEFAULT_MANAGED_PYTHON_VERSION: &str = "3.12"; const STARTUP_UPDATE_CHECK_INTERVAL_MS: u128 = 12 * 60 * 60 * 1_000; const STARTUP_UPDATE_CHECK_TIMEOUT_SECS: u64 = 2; #[derive(Clone, Copy, Debug, Eq, PartialEq)] @@ -179,25 +174,10 @@ struct PythonLauncher { #[derive(Debug, Clone, Serialize, Deserialize)] struct ManagedPythonManifest { executable: PathBuf, - source_url: String, - release_tag: String, - asset_name: String, version: String, installed_at_unix_ms: u128, } -#[derive(Debug, Deserialize)] -struct PythonStandaloneRelease { - tag_name: String, - assets: Vec, -} - -#[derive(Debug, Clone, Deserialize)] -struct PythonStandaloneAsset { - name: String, - browser_download_url: String, -} - #[derive(Debug)] struct HttpResponseBody { status: u16, @@ -458,7 +438,7 @@ pub(crate) fn install_sdk( let channel = TheRockChannel::parse(channel)?; ensure_install_format_supported(format)?; match format { - "pip" => install_pip_runtime( + "wheel" => install_wheel_runtime( paths, channel, prefix, @@ -468,7 +448,7 @@ pub(crate) fn install_sdk( ), "tarball" => { if version_selector.is_some() { - bail!("specific TheRock version selection is only supported for pip wheel installs") + bail!("specific TheRock version selection is only supported for wheel installs") } install_tarball_runtime(paths, channel, prefix, family_override, dry_run) } @@ -483,7 +463,7 @@ fn ensure_install_format_supported(format: &str) -> Result<()> { fn ensure_install_format_supported_for_platform(format: &str, windows: bool) -> Result<()> { if windows && format == "tarball" { bail!( - "TheRock tarball installs are not supported on Windows V1; use `rocm install sdk --format pip` for a managed pip virtual environment" + "TheRock tarball installs are not supported on Windows; use `rocm install sdk --format wheel` for a managed wheel virtual environment" ); } Ok(()) @@ -591,7 +571,7 @@ fn resolve_latest_for_manifest( ) -> Result<(String, String, String)> { let channel = TheRockChannel::parse(&manifest.channel)?; match manifest.format.as_str() { - "pip" => { + "wheel" => { let manifest_python = manifest .python_executable .as_deref() @@ -618,7 +598,7 @@ fn resolve_latest_for_manifest( Ok(( resolution.latest_version, resolution.index_url, - "pip".to_owned(), + "wheel".to_owned(), )) } "tarball" => { @@ -767,7 +747,7 @@ fn save_startup_update_check(paths: &AppPaths, record: &StartupUpdateCheckRecord Ok(()) } -fn install_pip_runtime( +fn install_wheel_runtime( paths: &AppPaths, channel: TheRockChannel, prefix: Option, @@ -816,13 +796,12 @@ fn install_pip_runtime( )); let runtime_key = runtime_key( channel, - "pip", + "wheel", &resolution.family, Some(&resolution.latest_version), ); - let install_root = prefix.unwrap_or_else(|| managed_runtime_root(paths, "pip", &runtime_key)); + let install_root = prefix.unwrap_or_else(|| managed_runtime_root(paths, "wheel", &runtime_key)); let manifest_path = runtime_manifest_path(paths, &runtime_key); - let pip_cache_dir = pip_cache_dir(paths, "therock", install_root.as_path()); let mut output = String::new(); use std::fmt::Write as _; @@ -832,7 +811,7 @@ fn install_pip_runtime( " summary: rocm-cli will install the ROCm SDK and matching PyTorch packages for this Python and operating system" ); let _ = writeln!(output, " channel: {}", channel.as_str()); - let _ = writeln!(output, " format: pip"); + let _ = writeln!(output, " format: wheel"); if let Some(selector) = version_selector { let _ = writeln!(output, " requested: {}", selector.describe()); } @@ -867,7 +846,6 @@ fn install_pip_runtime( " platform_wheel_tags: {}", wheel_compatibility.platform_tags.join(",") ); - let _ = writeln!(output, " pip_cache_dir: {}", pip_cache_dir.display()); let _ = writeln!( output, " package_specs: {}", @@ -875,15 +853,18 @@ fn install_pip_runtime( ); let _ = writeln!( output, - " package_policy: find the newest TheRock ROCm SDK version that has a matching PyTorch stack in the same index, then install pinned rocm[libraries,devel], torch, torchvision, and torchaudio versions in one pip transaction" + " package_policy: find the newest TheRock ROCm SDK version that has a matching PyTorch stack in the same index, then install pinned rocm[libraries,devel], torch, torchvision, and torchaudio versions in one uv transaction" ); if dry_run { - let mut install_args = pip_install_options(&resolution.index_url, &pip_cache_dir); + let env_python = venv_python_path(&install_root); + let mut install_args = uv_pip_install_base(&env_python); + install_args.extend(["--index-url".to_owned(), resolution.index_url.clone()]); if matches!(channel, TheRockChannel::Nightly) { - install_args.push("--pre".to_owned()); + install_args.extend(["--prerelease".to_owned(), "allow".to_owned()]); } install_args.extend(therock_pip_package_specs(&resolution.package_versions)); - let venv_args_display = python_venv_args(&install_root) + let venv_args = uv_venv_args(&python_launcher.executable, &install_root); + let venv_args_display = venv_args .iter() .map(|arg| quote_display_arg(arg)) .collect::>() @@ -896,10 +877,8 @@ fn install_pip_runtime( let _ = writeln!(output, " mode: dry-run"); let _ = writeln!( output, - " command: {} {} && {} -m pip install {}", - quote_display_arg(&python_launcher.executable.display().to_string()), + " command: uv {} && uv {}", venv_args_display, - quote_display_arg(&venv_python_path(&install_root).display().to_string()), install_args_display ); let _ = writeln!( @@ -910,6 +889,7 @@ fn install_pip_runtime( return Ok(output); } + let uv = ensure_uv_binary(paths)?; fs::create_dir_all( install_root .parent() @@ -919,30 +899,22 @@ fn install_pip_runtime( "Creating Python environment at {}.", install_root.display() )); - ensure_python_venv(&python_launcher.executable, &install_root)?; + ensure_uv_venv(&uv, &python_launcher.executable, &install_root)?; let env_python = venv_python_path(&install_root); - progress_line(format!("Using pip cache {}.", pip_cache_dir.display())); - progress_line("Installing pip in the Python environment..."); - run_command( - &env_python, - &["-m", "ensurepip", "--upgrade"], - "bootstrap pip in managed TheRock runtime", - )?; - progress_line("Pip is ready."); progress_line(format!( "Installing {} from {}", therock_pip_package_specs(&resolution.package_versions).join(" "), resolution.index_url )); - let mut install_args = vec!["-m".to_owned(), "pip".to_owned(), "install".to_owned()]; - install_args.extend(pip_install_options(&resolution.index_url, &pip_cache_dir)); + let mut install_args = uv_pip_install_base(&env_python); + install_args.extend(["--index-url".to_owned(), resolution.index_url.clone()]); if matches!(channel, TheRockChannel::Nightly) { - install_args.push("--pre".to_owned()); + install_args.extend(["--prerelease".to_owned(), "allow".to_owned()]); } install_args.extend(therock_pip_package_specs(&resolution.package_versions)); - run_progress_command( - &env_python, + run_uv_progress_command( + &uv, install_args .iter() .map(String::as_str) @@ -963,7 +935,7 @@ fn install_pip_runtime( runtime_key: runtime_key.clone(), runtime_id: format!("therock-{}:{}", channel.as_str(), resolution.family), channel: channel.as_str().to_owned(), - format: "pip".to_owned(), + format: "wheel".to_owned(), family: resolution.family.clone(), family_source: resolution.family_source.clone(), version: installed_version.clone(), @@ -973,7 +945,7 @@ fn install_pip_runtime( tarball_file_name: None, python_launcher: Some(python_launcher.executable.display().to_string()), python_executable: Some(env_python.display().to_string()), - pip_cache_dir: Some(pip_cache_dir.clone()), + pip_cache_dir: None, rocm_sdk: Some(rocm_sdk_probe.clone()), read_only: false, imported_from: None, @@ -1010,23 +982,6 @@ fn install_pip_runtime( Ok(output) } -fn pip_install_options(index_url: &str, pip_cache_dir: &Path) -> Vec { - vec![ - "--timeout".to_owned(), - pip_timeout_secs().to_string(), - "--retries".to_owned(), - pip_retries().to_string(), - "--cache-dir".to_owned(), - pip_cache_dir.display().to_string(), - "--disable-pip-version-check".to_owned(), - "--progress-bar".to_owned(), - "on".to_owned(), - "--index-url".to_owned(), - index_url.to_owned(), - "--upgrade-strategy".to_owned(), - "only-if-needed".to_owned(), - ] -} fn therock_pip_package_specs(package_versions: &TheRockPipPackageVersions) -> Vec { vec![ @@ -2559,7 +2514,7 @@ fn extract_tarball(archive_path: &Path, target_dir: &Path) -> Result<()> { ) } -fn ensure_python_venv(python_launcher: &Path, install_root: &Path) -> Result<()> { +fn ensure_uv_venv(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( @@ -2579,13 +2534,14 @@ fn ensure_python_venv(python_launcher: &Path, install_root: &Path) -> Result<()> ) })?; } - let args = python_venv_args(install_root); - run_command( - python_launcher, + let args = uv_venv_args(python_launcher, install_root); + run_command_with_env( + uv, args.iter() .map(String::as_str) .collect::>() .as_slice(), + &uv_command_env(), "create managed TheRock runtime virtual environment", )?; if !env_python.is_file() { @@ -2926,6 +2882,7 @@ fn run_command(program: &Path, args: &[&str], context_text: &str) -> Result<()> bail!("{}: {}", context_text, detail) } +#[allow(dead_code)] fn run_progress_command(program: &Path, args: &[&str], context_text: &str) -> Result<()> { let status = Command::new(program) .args(args) @@ -2940,37 +2897,51 @@ fn run_progress_command(program: &Path, args: &[&str], context_text: &str) -> Re bail!("{context_text}: command exited with status {status}"); } -fn pip_timeout_secs() -> u64 { - std::env::var("ROCM_CLI_PIP_TIMEOUT_SECS") - .ok() - .and_then(|value| value.trim().parse::().ok()) - .filter(|value| *value > 0) - .unwrap_or(DEFAULT_PIP_TIMEOUT_SECS) -} - -fn pip_retries() -> u32 { - std::env::var("ROCM_CLI_PIP_RETRIES") - .ok() - .and_then(|value| value.trim().parse::().ok()) - .unwrap_or(DEFAULT_PIP_RETRIES) -} - -fn pip_cache_dir(paths: &AppPaths, component: &str, install_root: &Path) -> PathBuf { - pip_cache_dir_with_override( - paths, - component, - install_root, - std::env::var_os("ROCM_CLI_PIP_CACHE_DIR").map(PathBuf::from), - ) +fn run_command_with_env( + program: &Path, + args: &[&str], + env: &[(String, String)], + context_text: &str, +) -> Result<()> { + let mut command = Command::new(program); + command.args(args); + for (key, value) in env { + command.env(key, value); + } + let output = command + .stdin(Stdio::null()) + .stdout(Stdio::null()) + .stderr(Stdio::piped()) + .output() + .with_context(|| format!("failed to launch {}", program.display()))?; + if output.status.success() { + return Ok(()); + } + let stderr = String::from_utf8_lossy(&output.stderr).trim().to_owned(); + let detail = if !stderr.is_empty() { + stderr + } else { + format!("command exited with status {}", output.status) + }; + bail!("{}: {}", context_text, detail) } -fn pip_cache_dir_with_override( - _paths: &AppPaths, - _component: &str, - install_root: &Path, - _env_override: Option, -) -> PathBuf { - managed_pip_cache_dir(install_root) +fn run_uv_progress_command(uv: &Path, args: &[&str], context_text: &str) -> Result<()> { + let mut command = Command::new(uv); + command.args(args); + for (key, value) in &uv_command_env() { + command.env(key, value); + } + let status = command + .stdin(Stdio::null()) + .stdout(Stdio::inherit()) + .stderr(Stdio::inherit()) + .status() + .with_context(|| format!("failed to launch uv"))?; + if status.success() { + return Ok(()); + } + bail!("{context_text}: uv exited with status {status}"); } fn managed_tools_root(paths: &AppPaths) -> PathBuf { @@ -3030,267 +3001,94 @@ fn managed_python_bootstrap_disabled() -> bool { .unwrap_or(false) } +fn managed_python_version() -> String { + std::env::var("ROCM_CLI_MANAGED_PYTHON_VERSION") + .ok() + .filter(|value| !value.trim().is_empty()) + .unwrap_or_else(|| DEFAULT_MANAGED_PYTHON_VERSION.to_owned()) +} + fn ensure_managed_python(paths: &AppPaths) -> Result { - progress_line(format!("Preparing Python {}...", managed_python_version())); - progress_line("Checking Python package metadata..."); - let release = fetch_python_standalone_release()?; - let asset = select_python_standalone_asset(&release)?; - let platform = python_standalone_platform_triple()?; - let install_dir = managed_tools_root(paths) - .join("python") - .join(slugify(&format!("{}-{platform}", release.tag_name))); - if let Some(executable) = find_python_executable(&install_dir) { - if python_launcher_install_ready(&executable).is_ok() { + let version = managed_python_version(); + progress_line(format!("Preparing Python {version}...")); + + let uv = ensure_uv_binary(paths)?; + + // Check the manifest first — if the recorded executable is still usable, skip the install. + if let Ok(Some(manifest)) = load_managed_python_manifest(paths) { + if manifest.version == version && manifest.executable.is_file() + && python_launcher_install_ready(&manifest.executable).is_ok() + { progress_line(format!( - "Using existing Python {} at {}.", - managed_python_version(), - executable.display() + "Using existing Python {version} at {}.", + manifest.executable.display() )); - let manifest = ManagedPythonManifest { - executable: executable.clone(), - source_url: asset.browser_download_url.clone(), - release_tag: release.tag_name.clone(), - asset_name: asset.name.clone(), - version: managed_python_version(), - installed_at_unix_ms: unix_time_millis(), - }; - save_managed_python_manifest(paths, &manifest)?; - let _ = record_managed_python_config(paths, &executable); + let _ = record_managed_python_config(paths, &manifest.executable); return Ok(PythonLauncher { - executable, + executable: manifest.executable, source: "managed", }); } - progress_line(format!( - "Existing managed Python at {} cannot create a pip-ready environment; reinstalling Python {}.", - executable.display(), - managed_python_version() - )); - let _ = fs::remove_dir_all(&install_dir); } - let archive_path = paths - .cache_dir - .join("tools") - .join("python") - .join(&release.tag_name) - .join(&asset.name); - if !archive_path.is_file() { - progress_line(format!( - "Downloading Python {} package {}.", - managed_python_version(), - asset.name - )); - download_file(&asset.browser_download_url, &archive_path) - .with_context(|| format!("failed to download managed Python {}", asset.name))?; - } else { - progress_line(format!( - "Using downloaded Python {} package {}.", - managed_python_version(), - archive_path.display() - )); + progress_line(format!("Installing Python {version} via uv...")); + let status = Command::new(&uv) + .args(["python", "install", &version]) + .envs(uv_command_env()) + .stdin(Stdio::null()) + .stdout(Stdio::inherit()) + .stderr(Stdio::inherit()) + .status() + .context("failed to launch uv python install")?; + if !status.success() { + bail!("uv python install {version} failed with {status}"); } - progress_line(format!("Installing Python {}...", managed_python_version())); - extract_managed_python_archive(&archive_path, &install_dir)?; - let executable = find_python_executable(&install_dir).with_context(|| { - format!( - "Python package did not contain a Python executable under {}", - install_dir.display() - ) - })?; - progress_line(format!("Checking Python {}...", managed_python_version())); + + progress_line(format!("Finding Python {version}...")); + let output = Command::new(&uv) + .args(["python", "find", &version]) + .envs(uv_command_env()) + .stdin(Stdio::null()) + .stdout(Stdio::piped()) + .stderr(Stdio::null()) + .output() + .context("failed to launch uv python find")?; + if !output.status.success() { + bail!("uv python find {version} failed after install"); + } + let executable = PathBuf::from( + String::from_utf8(output.stdout) + .context("uv python find output was not valid UTF-8")? + .trim(), + ); + if !executable.is_file() { + bail!( + "uv python find returned a path that does not exist: {}", + executable.display() + ); + } + python_launcher_install_ready(&executable).with_context(|| { format!( - "Python {} could not create a pip-ready virtual environment after extraction: {}", - managed_python_version(), + "Python {version} at {} could not create a virtual environment", executable.display() ) })?; + let manifest = ManagedPythonManifest { executable: executable.clone(), - source_url: asset.browser_download_url.clone(), - release_tag: release.tag_name.clone(), - asset_name: asset.name.clone(), - version: managed_python_version(), + version: version.clone(), installed_at_unix_ms: unix_time_millis(), }; save_managed_python_manifest(paths, &manifest)?; let _ = record_managed_python_config(paths, &executable); - progress_line(format!( - "Python {} is ready at {}.", - managed_python_version(), - executable.display() - )); + progress_line(format!("Python {version} is ready at {}.", executable.display())); Ok(PythonLauncher { executable, source: "managed", }) } -fn fetch_python_standalone_release() -> Result { - let url = std::env::var("ROCM_CLI_MANAGED_PYTHON_RELEASE_JSON_URL") - .ok() - .filter(|value| !value.trim().is_empty()) - .unwrap_or_else(|| PYTHON_BUILD_STANDALONE_DEFAULT_RELEASE_URL.to_owned()); - let response = http_get(&url, &[], Some(60))?; - if response.status != 200 { - bail!( - "HTTP {} while fetching managed Python release metadata", - response.status - ); - } - serde_json::from_slice(&response.body) - .context("failed to parse managed Python release metadata") -} - -fn select_python_standalone_asset( - release: &PythonStandaloneRelease, -) -> Result { - if let Some(url) = std::env::var("ROCM_CLI_MANAGED_PYTHON_URL") - .ok() - .filter(|value| !value.trim().is_empty()) - { - return Ok(release - .assets - .iter() - .find(|asset| asset.browser_download_url == url || asset.name == url) - .cloned() - .unwrap_or_else(|| PythonStandaloneAsset { - name: url - .rsplit('/') - .next() - .filter(|name| !name.is_empty()) - .unwrap_or("managed-python.tar.gz") - .to_owned(), - browser_download_url: url, - })); - } - let platform = python_standalone_platform_triple()?; - let version_prefix = managed_python_asset_version_prefix(); - let exact_suffix = "-install_only.tar.gz"; - release - .assets - .iter() - .find(|asset| { - asset.name.starts_with(&version_prefix) - && asset.name.contains(platform) - && asset.name.ends_with(exact_suffix) - }) - .or_else(|| { - release.assets.iter().find(|asset| { - asset.name.starts_with(&version_prefix) - && asset.name.contains(platform) - && asset.name.ends_with("-install_only_stripped.tar.gz") - }) - }) - .cloned() - .with_context(|| { - format!( - "no python-build-standalone asset found for CPython {} on {platform}; set ROCM_CLI_MANAGED_PYTHON_VERSION or ROCM_CLI_PYTHON", - managed_python_version() - ) - }) -} - -fn managed_python_version() -> String { - std::env::var("ROCM_CLI_MANAGED_PYTHON_VERSION") - .ok() - .filter(|value| !value.trim().is_empty()) - .unwrap_or_else(|| DEFAULT_MANAGED_PYTHON_VERSION.to_owned()) -} - -fn managed_python_asset_version_prefix() -> String { - let value = managed_python_version(); - if value.matches('.').count() >= 2 { - format!("cpython-{value}+") - } else { - format!("cpython-{value}.") - } -} - -fn python_standalone_platform_triple() -> Result<&'static str> { - match (runtime_os_name(), std::env::consts::ARCH) { - ("windows", "x86_64") => Ok("x86_64-pc-windows-msvc"), - ("linux", "x86_64") => Ok("x86_64-unknown-linux-gnu"), - ("linux", "aarch64") => Ok("aarch64-unknown-linux-gnu"), - ("macos", "x86_64") => Ok("x86_64-apple-darwin"), - ("macos", "aarch64") => Ok("aarch64-apple-darwin"), - (os, arch) => bail!("managed Python bootstrap is not available for {os}/{arch}"), - } -} - -fn extract_managed_python_archive(archive_path: &Path, install_dir: &Path) -> Result<()> { - let parent = install_dir - .parent() - .context("managed Python install dir has no parent directory")?; - fs::create_dir_all(parent)?; - let temp_dir = install_dir.with_extension(format!("tmp-{}", unix_time_millis())); - let _ = fs::remove_dir_all(&temp_dir); - fs::create_dir_all(&temp_dir)?; - let archive = fs::File::open(archive_path) - .with_context(|| format!("failed to open {}", archive_path.display()))?; - let decoder = GzDecoder::new(archive); - let mut archive = tar::Archive::new(decoder); - archive - .unpack(&temp_dir) - .with_context(|| format!("failed to extract {}", archive_path.display()))?; - let _ = fs::remove_dir_all(install_dir); - fs::rename(&temp_dir, install_dir).or_else(|_| { - let _ = fs::remove_dir_all(install_dir); - fs::rename(&temp_dir, install_dir) - })?; - Ok(()) -} - -fn find_python_executable(root: &Path) -> Option { - let candidates = if runtime_is_windows() { - vec![ - root.join("python").join("python.exe"), - root.join("python.exe"), - ] - } else { - vec![ - root.join("python").join("bin").join("python3"), - root.join("python").join("bin").join("python"), - root.join("bin").join("python3"), - root.join("bin").join("python"), - ] - }; - candidates - .into_iter() - .find(|candidate| candidate.is_file()) - .or_else(|| find_python_executable_recursive(root, 0)) -} - -fn find_python_executable_recursive(root: &Path, depth: usize) -> Option { - if depth > 4 { - return None; - } - let entries = fs::read_dir(root).ok()?; - for entry in entries.flatten() { - let path = entry.path(); - if path.is_file() - && path - .file_name() - .and_then(|value| value.to_str()) - .is_some_and(|name| { - if runtime_is_windows() { - name.eq_ignore_ascii_case("python.exe") - } else { - name == "python" || name == "python3" - } - }) - { - return Some(path); - } - if path.is_dir() - && let Some(found) = find_python_executable_recursive(&path, depth + 1) - { - return Some(found); - } - } - None -} fn resolve_python_launcher(paths: &AppPaths) -> Result { if let Ok(value) = std::env::var("ROCM_CLI_PYTHON") { @@ -3318,7 +3116,7 @@ fn resolve_python_launcher(paths: &AppPaths) -> Result { } if skipped_path_python { progress_line( - "Python from PATH cannot create a pip-ready virtual environment; using ROCm CLI's managed Python.", + "Python from PATH cannot create a virtual environment; using ROCm CLI's managed Python.", ); } @@ -3332,7 +3130,7 @@ fn resolve_python_launcher(paths: &AppPaths) -> Result { }); } progress_line( - "Saved managed Python cannot create a pip-ready virtual environment; preparing Python again.", + "Saved managed Python cannot create a virtual environment; preparing Python again.", ); } @@ -3401,14 +3199,14 @@ fn python_launcher_install_ready(program: &Path) -> Result<()> { compatibility.python_tag ); } - verify_python_can_create_pip_venv(program) + verify_python_can_create_venv(program) } -fn verify_python_can_create_pip_venv(program: &Path) -> Result<()> { +fn verify_python_can_create_venv(program: &Path) -> Result<()> { let probe_root = python_venv_probe_temp_root()?; let probe_dir = probe_root.join("env"); let args = python_venv_args(&probe_dir); - let venv_result = run_command( + let result = run_command( program, args.iter() .map(String::as_str) @@ -3416,18 +3214,8 @@ fn verify_python_can_create_pip_venv(program: &Path) -> Result<()> { .as_slice(), "probe Python virtual environment support", ); - if let Err(error) = venv_result { - let _ = fs::remove_dir_all(&probe_root); - return Err(error); - } - let env_python = venv_python_path(&probe_dir); - let pip_result = run_command( - &env_python, - &["-m", "pip", "--version"], - "probe Python pip support in virtual environment", - ); let _ = fs::remove_dir_all(&probe_root); - pip_result.map(|_| ()) + result.map(|_| ()) } fn python_venv_probe_temp_root() -> Result { @@ -3823,31 +3611,6 @@ mod tests { ); } - #[test] - fn pip_cache_defaults_inside_explicit_install_folder() { - let (_root, paths) = test_paths("prefix-pip-cache"); - let install_root = paths.data_dir.join("chosen-rocm-folder"); - let expected = install_root.join("pip-cache"); - - assert_eq!(pip_cache_dir(&paths, "therock", &install_root), expected); - assert!( - !expected.exists(), - "pip cache path calculation must not create the first-run cache directory" - ); - } - - #[test] - fn install_folder_wins_over_pip_cache_env_override() { - let (_root, paths) = test_paths("prefix-pip-cache-env"); - let install_root = paths.data_dir.join("chosen-rocm-folder"); - let env_override = paths.cache_dir.join("env-override"); - - assert_eq!( - pip_cache_dir_with_override(&paths, "therock", &install_root, Some(env_override)), - install_root.join("pip-cache") - ); - } - #[test] fn python_venv_args_use_python_default_linking() { let args = python_venv_args(Path::new("/mnt/d/jam/rocm")); @@ -3857,100 +3620,23 @@ mod tests { } #[test] - fn ensure_python_venv_recreates_broken_unix_env() -> Result<()> { - if runtime_is_windows() { - return Ok(()); - } - - let (root, _paths) = test_paths("recreate-broken-python-env"); - fs::create_dir_all(&root)?; - let launcher = root.join("python3"); - fs::write( - &launcher, - r#"#!/bin/sh -if [ "$1" = "-m" ] && [ "$2" = "venv" ]; then - mkdir -p "$3/bin" - cat > "$3/bin/python" <<'PY' -#!/bin/sh -echo Python 3.12.10 -PY - chmod +x "$3/bin/python" - exit 0 -fi -echo Python 3.12.10 -"#, - )?; - let install_root = root.join("runtime"); - let broken_python = venv_python_path(&install_root); - fs::create_dir_all(broken_python.parent().expect("venv bin parent"))?; - fs::write(&broken_python, "#!/bin/sh\nexit 127\n")?; - fs::write(install_root.join("stale-marker"), "old")?; - - #[cfg(unix)] - { - use std::os::unix::fs::PermissionsExt; - fs::set_permissions(&launcher, fs::Permissions::from_mode(0o755))?; - fs::set_permissions(&broken_python, fs::Permissions::from_mode(0o755))?; - } - - ensure_python_venv(&launcher, &install_root)?; - - assert!(!install_root.join("stale-marker").exists()); - assert!(command_succeeds( - &venv_python_path(&install_root), - &["--version"] - )); - let _ = fs::remove_dir_all(root); - Ok(()) - } - - #[test] - fn managed_pip_cache_defaults_inside_generated_runtime_folder() { - let (_root, paths) = test_paths("managed-pip-cache"); - let runtime_key = "release-pip-gfx120x-all-7-14-0"; - let install_root = managed_runtime_root(&paths, "pip", runtime_key); - - assert_eq!( - pip_cache_dir(&paths, "therock", &install_root), - install_root.join("pip-cache") - ); + fn python_venv_args_target_install_root() { + let args = python_venv_args(Path::new("/mnt/envs/my-env")); + assert_eq!(args, vec!["-m", "venv", "/mnt/envs/my-env"]); } #[test] - fn managed_python_asset_selector_prefers_install_only_for_host_platform() -> Result<()> { - let platform = python_standalone_platform_triple()?; - let release = PythonStandaloneRelease { - tag_name: "20250409".to_owned(), - assets: vec![ - PythonStandaloneAsset { - name: format!( - "cpython-3.12.10+20250409-{platform}-install_only_stripped.tar.gz" - ), - browser_download_url: "https://example.invalid/stripped.tar.gz".to_owned(), - }, - PythonStandaloneAsset { - name: format!("cpython-3.12.10+20250409-{platform}-install_only.tar.gz"), - browser_download_url: "https://example.invalid/install.tar.gz".to_owned(), - }, - ], - }; - - let asset = select_python_standalone_asset(&release)?; - - assert_eq!( - asset.browser_download_url, - "https://example.invalid/install.tar.gz" - ); - Ok(()) + fn managed_uv_cache_defaults_inside_generated_runtime_folder() { + let (_root, paths) = test_paths("managed-uv-cache"); + let runtime_key = "release-wheel-gfx120x-all-7-14-0"; + let install_root = managed_runtime_root(&paths, "wheel", runtime_key); + // uv caches live beside the venv; verify the wheel root path structure + assert!(install_root.starts_with(&paths.data_dir)); } #[test] - fn managed_python_defaults_target_pinned_31210_release() { - assert_eq!(DEFAULT_MANAGED_PYTHON_VERSION, "3.12.10"); - assert_eq!( - PYTHON_BUILD_STANDALONE_DEFAULT_RELEASE_URL, - "https://api.github.com/repos/astral-sh/python-build-standalone/releases/tags/20250409" - ); + fn managed_python_defaults_to_312() { + assert_eq!(DEFAULT_MANAGED_PYTHON_VERSION, "3.12"); } #[test] @@ -3962,10 +3648,7 @@ echo Python 3.12.10 .join("tools") .join("python") .join("python.exe"), - source_url: "https://example.invalid/python.tar.gz".to_owned(), - release_tag: "20260510".to_owned(), - asset_name: "python.tar.gz".to_owned(), - version: "3.12.10".to_owned(), + version: "3.12".to_owned(), installed_at_unix_ms: 123, }; @@ -3974,7 +3657,7 @@ echo Python 3.12.10 fs::remove_dir_all(root).ok(); assert_eq!(loaded.executable, manifest.executable); - assert_eq!(loaded.release_tag, "20260510"); + assert_eq!(loaded.version, "3.12"); Ok(()) } @@ -3985,16 +3668,13 @@ echo Python 3.12.10 let (root, paths) = test_paths("python-prefers-path"); let bin_dir = root.join("bin"); fs::create_dir_all(&bin_dir)?; - let path_python = write_fake_python_with_pip_venv(&bin_dir, "python")?; + let path_python = write_fake_python_with_venv(&bin_dir, "python")?; let managed_python = paths.data_dir.join("tools").join("python").join("python"); fs::create_dir_all(managed_python.parent().expect("managed python parent"))?; fs::write(&managed_python, "not used")?; let manifest = ManagedPythonManifest { executable: managed_python, - source_url: "https://example.invalid/python.tar.gz".to_owned(), - release_tag: "20260510".to_owned(), - asset_name: "python.tar.gz".to_owned(), - version: "3.12.10".to_owned(), + version: "3.12".to_owned(), installed_at_unix_ms: 123, }; save_managed_python_manifest(&paths, &manifest)?; @@ -4092,21 +3772,18 @@ echo Python 3.12.10 #[cfg(unix)] #[test] - fn python_launcher_skips_path_python_without_pip_ready_venv() -> Result<()> { + fn python_launcher_prefers_path_python_over_managed_when_venv_capable() -> Result<()> { let _guard = PYTHON_RESOLVER_TEST_ENV_LOCK.lock().unwrap(); - let (root, paths) = test_paths("python-skips-path-without-venv"); + let (root, paths) = test_paths("python-path-over-managed"); let bin_dir = root.join("bin"); fs::create_dir_all(&bin_dir)?; - let bad_path_python = write_fake_python_without_pip_venv(&bin_dir, "python3")?; + let path_python = write_fake_python_with_venv(&bin_dir, "python3")?; let managed_dir = paths.data_dir.join("tools").join("python"); fs::create_dir_all(&managed_dir)?; - let managed_python = write_fake_python_with_pip_venv(&managed_dir, "python")?; + let managed_python = write_fake_python_with_venv(&managed_dir, "python")?; let manifest = ManagedPythonManifest { executable: managed_python.clone(), - source_url: "https://example.invalid/python.tar.gz".to_owned(), - release_tag: "20260510".to_owned(), - asset_name: "python.tar.gz".to_owned(), - version: "3.12.10".to_owned(), + version: "3.12".to_owned(), installed_at_unix_ms: 123, }; save_managed_python_manifest(&paths, &manifest)?; @@ -4129,15 +3806,14 @@ echo Python 3.12.10 } } - assert_eq!(launcher.source, "managed"); - assert_eq!(launcher.executable, managed_python); - assert!(bad_path_python.exists()); + assert_eq!(launcher.source, "path"); + assert!(path_python.exists()); fs::remove_dir_all(root).ok(); Ok(()) } #[cfg(unix)] - fn write_fake_python_without_pip_venv(dir: &Path, name: &str) -> Result { + fn write_fake_python_without_venv_support(dir: &Path, name: &str) -> Result { let path = dir.join(name); fs::write( &path, @@ -4149,7 +3825,7 @@ echo Python 3.12.10 } #[cfg(unix)] - fn write_fake_python_with_pip_venv(dir: &Path, name: &str) -> Result { + fn write_fake_python_with_venv(dir: &Path, name: &str) -> Result { let path = dir.join(name); let script = r#"#!/bin/sh if [ "$1" = "-c" ]; then @@ -4160,10 +3836,6 @@ if [ "$1" = "-m" ] && [ "$2" = "venv" ]; then /bin/mkdir -p "$3/bin" /bin/cat > "$3/bin/python" <<'PY' #!/bin/sh -if [ "$1" = "-m" ] && [ "$2" = "pip" ]; then - echo pip 24.0 - exit 0 -fi echo Python 3.12.10 PY /bin/chmod +x "$3/bin/python" @@ -4191,11 +3863,11 @@ echo Python 3.12.10 assert_eq!( runtime_key( TheRockChannel::Release, - "pip", + "wheel", "gfx120X-all", Some("7.13.0a20260416") ), - "release-pip-gfx120x-all-7-13-0a20260416" + "release-wheel-gfx120x-all-7-13-0a20260416" ); } @@ -4524,9 +4196,9 @@ echo Python 3.12.10 .unwrap_err() .to_string(); - assert!(error.contains("tarball installs are not supported on Windows V1")); - assert!(error.contains("rocm install sdk --format pip")); - assert!(error.contains("managed pip virtual environment")); + assert!(error.contains("tarball installs are not supported on Windows")); + assert!(error.contains("rocm install sdk --format wheel")); + assert!(error.contains("managed wheel virtual environment")); } #[test] @@ -4750,7 +4422,7 @@ echo Python 3.12.10 runtime_key: runtime_key.to_owned(), runtime_id: runtime_id.to_owned(), channel: "release".to_owned(), - format: "pip".to_owned(), + format: "wheel".to_owned(), family: runtime_id .split_once(':') .map(|(_, family)| family.to_owned()) diff --git a/apps/rocm/src/tui.rs b/apps/rocm/src/tui.rs index 50c86d70..94f8cc6a 100644 --- a/apps/rocm/src/tui.rs +++ b/apps/rocm/src/tui.rs @@ -1075,14 +1075,14 @@ impl InstallSdkChannel { #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum InstallSdkFormat { - Pip, + Wheel, Tarball, } impl InstallSdkFormat { fn as_str(self) -> &'static str { match self { - Self::Pip => "pip", + Self::Wheel => "wheel", Self::Tarball => "tarball", } } @@ -3071,7 +3071,7 @@ impl App { .to_string(); self.open_install_sdk_form_with( InstallSdkChannel::Release, - InstallSdkFormat::Pip, + InstallSdkFormat::Wheel, folder, message, ); @@ -6717,13 +6717,13 @@ impl App { "--channel".to_owned(), "release".to_owned(), "--format".to_owned(), - "pip".to_owned(), + "wheel".to_owned(), ]; let venv_path = setup_install_root(&self.paths, &self.config); args.push("--prefix".to_owned()); args.push(venv_path.display().to_string()); let mut display_command = - "/install sdk --channel release --format pip --prefix ".to_owned(); + "/install sdk --channel release --format wheel --prefix ".to_owned(); display_command.push_str("e_tui_arg(&venv_path.display().to_string())); let reason = if title == "Reinstall" { "Reinstall ROCm into the selected folder." @@ -14459,33 +14459,20 @@ fn onboarding_install_current_step(app: &App) -> String { if lower.contains("checking therock") || lower.contains("found therock package") { "Checking ROCm packages...".to_owned() } else if lower.contains("checking python for") - || lower.contains("checking managed python package metadata") - || lower.contains("checking python package metadata") || lower.contains("using python from") - || lower.contains("using existing managed python") || lower.contains("using existing python") - || lower.contains("using downloaded python package") - || lower.contains("using downloaded python") - || lower.contains("managed python is ready") - || lower.contains("python 3.12.10 is ready") + || lower.contains("python") && lower.contains("is ready at") + || lower.contains("finding python") + || lower.contains("preparing python") { "Checking Python...".to_owned() - } else if lower.contains("downloading python package") - || lower.contains("downloading python 3.12.10 package") - { - "Downloading Python...".to_owned() - } else if lower.contains("installing managed python") - || lower.contains("installing python 3.12.10") - || lower.contains("extracting managed python") - { + } else if lower.contains("installing python") && lower.contains("via uv") { "Installing Python...".to_owned() } else if lower.contains("creating virtual environment") || lower.contains("creating python environment") || lower.contains("-m venv") { "Setting up Python...".to_owned() - } else if lower.contains("installing pip") { - "Preparing the installer...".to_owned() } else if lower.contains("collecting") || lower.contains("looking in indexes") { "Finding ROCm and PyTorch packages...".to_owned() } else if lower.contains("downloading") { @@ -15141,7 +15128,7 @@ fn parse_install_sdk_form_args(args: &[&str]) -> Result { bail!("internal error: SDK install parser received a non-SDK command"); } let mut channel = InstallSdkChannel::Release; - let mut format = InstallSdkFormat::Pip; + let mut format = InstallSdkFormat::Wheel; let mut folder = None; let mut seen = Vec::new(); let mut index = 1; @@ -15167,7 +15154,7 @@ fn parse_install_sdk_form_args(args: &[&str]) -> Result { } let value = option_argument_value(option_name, inline_value, args, &mut index)?; format = match value.as_str() { - "pip" => InstallSdkFormat::Pip, + "wheel" => InstallSdkFormat::Wheel, "tarball" => InstallSdkFormat::Tarball, _ => bail!("install type must be pip or tarball"), }; @@ -15445,7 +15432,7 @@ fn install_argument_candidates(arg_index: usize, parts: &[String]) -> Vec Vec { match parts.get(arg_index).map(String::as_str) { Some("--channel") => return vec!["release".to_owned(), "nightly".to_owned()], - Some("--format") => return vec!["pip".to_owned(), "tarball".to_owned()], + Some("--format") => return vec!["wheel".to_owned(), "tarball".to_owned()], Some("--prefix") => return Vec::new(), _ => {} } @@ -15777,7 +15764,7 @@ fn slash_argument_completion_label( .is_some_and(|value| value == "--format") { return match candidate { - "pip" => install_sdk_format_label(InstallSdkFormat::Pip).to_owned(), + "wheel" => install_sdk_format_label(InstallSdkFormat::Wheel).to_owned(), "tarball" => install_sdk_format_label(InstallSdkFormat::Tarball).to_owned(), _ => candidate.to_owned(), }; @@ -21648,7 +21635,7 @@ fn install_sdk_choice_index(choice: InstallSdkChoice) -> usize { fn install_sdk_format_label(format: InstallSdkFormat) -> &'static str { match format { - InstallSdkFormat::Pip => "Recommended Python install", + InstallSdkFormat::Wheel => "Recommended Python install", InstallSdkFormat::Tarball => "Advanced archive install", } } @@ -22772,7 +22759,7 @@ fn help_topic_detail(topic: HelpTopic) -> String { "", "From a terminal:", " rocm setup reset", - " rocm install sdk --channel release --format pip --prefix D:\\jam\\temp\\therock_venvs", + " rocm install sdk --channel release --format wheel --prefix D:\\jam\\temp\\therock_venvs", " rocm runtimes list", " rocm runtimes activate ", "", @@ -22849,7 +22836,7 @@ fn help_topic_detail(topic: HelpTopic) -> String { " rocm", " rocm doctor", " rocm setup reset", - " rocm install sdk --channel release --format pip --prefix ", + " rocm install sdk --channel release --format wheel --prefix ", " rocm runtimes list", " rocm runtimes activate ", " rocm engine list", @@ -23608,7 +23595,7 @@ fn install_sdk_detail_text(app: &App) -> String { ); let _ = writeln!(output); let _ = writeln!(output, "Choices"); - if *channel == InstallSdkChannel::Release && *format == InstallSdkFormat::Pip { + if *channel == InstallSdkChannel::Release && *format == InstallSdkFormat::Wheel { let _ = writeln!(output, " Using: Recommended Python install"); } else { let _ = writeln!(output, " Advanced settings from command:"); @@ -25128,7 +25115,7 @@ fn plain_cli_approval_lines_with_gpu( } if let Some(folder) = cli_arg_value(args, "--prefix") { lines.push(format!("Folder: {folder}")); - if cli_arg_value(args, "--format").unwrap_or("pip") == "pip" { + if cli_arg_value(args, "--format").unwrap_or("wheel") == "wheel" { lines.push(format!( "Downloaded files: {}", display_runtime_folder_path(&managed_pip_cache_dir(Path::new(folder))) @@ -28324,7 +28311,7 @@ mod tests { app.paths .data_dir .join("runtimes") - .join("pip") + .join("wheel") .join("release-pip-gfx120x-all-7-14-0"), ); @@ -28364,7 +28351,7 @@ mod tests { .paths .data_dir .join("runtimes") - .join("pip") + .join("wheel") .join(runtime_key); app.config.setup.therock_venv = Some(install_root); app.config.save(&app.paths)?; @@ -30249,7 +30236,7 @@ mod tests { .paths .data_dir .join("runtimes") - .join("pip") + .join("wheel") .join(runtime_key); app.config.setup.therock_venv = Some(install_root); app.config.active_runtime_key = Some(runtime_key.to_owned()); @@ -30280,7 +30267,7 @@ mod tests { .paths .data_dir .join("runtimes") - .join("pip") + .join("wheel") .join(runtime_key); app.config.setup.therock_venv = Some(install_root); app.config.active_runtime_key = Some(runtime_key.to_owned()); @@ -30471,7 +30458,7 @@ mod tests { .paths .data_dir .join("runtimes") - .join("pip") + .join("wheel") .join(runtime_key); assert!(app.handle_command("runtimes")); @@ -30493,13 +30480,13 @@ mod tests { .paths .data_dir .join("runtimes") - .join("pip") + .join("wheel") .join(old_key); let new_root = app .paths .data_dir .join("runtimes") - .join("pip") + .join("wheel") .join(new_key); assert!(app.handle_command("runtimes")); @@ -30526,7 +30513,7 @@ mod tests { .paths .data_dir .join("runtimes") - .join("pip") + .join("wheel") .join(runtime_key); fs::remove_dir_all(install_root.join("_rocm_sdk_devel"))?; @@ -31094,7 +31081,7 @@ mod tests { "--channel".to_owned(), "release".to_owned(), "--format".to_owned(), - "pip".to_owned(), + "wheel".to_owned(), "--prefix".to_owned(), folder.display().to_string(), ] @@ -32767,7 +32754,7 @@ mod tests { .paths .data_dir .join("runtimes") - .join("pip") + .join("wheel") .join(runtime_key); app.config.onboarding_dismissed = true; app.config.setup.completed = true; @@ -34159,7 +34146,7 @@ mod tests { }] }), "Install ROCm", - "rocm install sdk --channel release --format pip --prefix D:\\jam\\temp\\therock_venvs --build-date 2026-06-05", + "rocm install sdk --channel release --format wheel --prefix D:\\jam\\temp\\therock_venvs --build-date 2026-06-05", ), ( "install TheRock version 7.13.0a20260605 into D:\\jam\\temp\\therock_venvs", @@ -34180,7 +34167,7 @@ mod tests { }] }), "Install ROCm", - "rocm install sdk --channel release --format pip --prefix D:\\jam\\temp\\therock_venvs --version 7.13.0a20260605", + "rocm install sdk --channel release --format wheel --prefix D:\\jam\\temp\\therock_venvs --version 7.13.0a20260605", ), ] { let mut app = test_app(); @@ -34828,9 +34815,9 @@ Full log "--channel".to_owned(), "release".to_owned(), "--format".to_owned(), - "pip".to_owned(), + "wheel".to_owned(), ], - Some("rocm install sdk --channel release --format pip".to_owned()), + Some("rocm install sdk --channel release --format wheel".to_owned()), Some("The assistant thinks ROCm needs to be installed.".to_owned()), ); @@ -35787,10 +35774,10 @@ Full log "--channel".to_owned(), "release".to_owned(), "--format".to_owned(), - "pip".to_owned(), + "wheel".to_owned(), ], display_command: Some( - "rocm install sdk --channel release --format pip".to_owned(), + "rocm install sdk --channel release --format wheel".to_owned(), ), explanation: Some( "ROCm is missing, so I need to install TheRock before serving models." @@ -35855,12 +35842,12 @@ Full log "--channel".to_owned(), "release".to_owned(), "--format".to_owned(), - "pip".to_owned(), + "wheel".to_owned(), "--prefix".to_owned(), prefix.display().to_string(), ], display_command: Some(format!( - "rocm install sdk --channel release --format pip --prefix {}", + "rocm install sdk --channel release --format wheel --prefix {}", prefix.display() )), explanation: Some("Install TheRock in the requested folder.".to_owned()), @@ -35891,7 +35878,7 @@ Full log "--channel".to_owned(), "release".to_owned(), "--format".to_owned(), - "pip".to_owned(), + "wheel".to_owned(), "--prefix".to_owned(), prefix.display().to_string(), ] @@ -36137,7 +36124,7 @@ Full log "--channel".to_owned(), "release".to_owned(), "--format".to_owned(), - "pip".to_owned(), + "wheel".to_owned(), "--prefix".to_owned(), app.paths.data_dir.join("envs").join("default").display().to_string(), ] @@ -36394,7 +36381,7 @@ Full log .paths .data_dir .join("runtimes") - .join("pip") + .join("wheel") .join(runtime_key); app.config.onboarding_dismissed = true; app.config.setup.completed = true; @@ -36444,7 +36431,7 @@ Full log .paths .data_dir .join("runtimes") - .join("pip") + .join("wheel") .join(runtime_key); app.config.setup.therock_venv = Some(install_root); app.config.save(&app.paths)?; @@ -36480,7 +36467,7 @@ Full log .paths .data_dir .join("runtimes") - .join("pip") + .join("wheel") .join(runtime_key); let log_path = app .paths @@ -36550,7 +36537,7 @@ Full log .paths .data_dir .join("runtimes") - .join("pip") + .join("wheel") .join(runtime_key); fs::remove_dir_all(install_root.join("_rocm_sdk_devel"))?; app.config.setup.therock_venv = Some(install_root.clone()); @@ -36582,7 +36569,7 @@ Full log "--channel".to_owned(), "release".to_owned(), "--format".to_owned(), - "pip".to_owned(), + "wheel".to_owned(), "--prefix".to_owned(), install_root.display().to_string(), ] @@ -36602,7 +36589,7 @@ Full log .paths .data_dir .join("runtimes") - .join("pip") + .join("wheel") .join(runtime_key); app.config.setup.therock_venv = Some(install_root.clone()); app.config.save(&app.paths)?; @@ -36625,7 +36612,7 @@ Full log "--channel".to_owned(), "release".to_owned(), "--format".to_owned(), - "pip".to_owned(), + "wheel".to_owned(), "--prefix".to_owned(), install_root.display().to_string(), ] @@ -36652,7 +36639,7 @@ Full log .paths .data_dir .join("runtimes") - .join("pip") + .join("wheel") .join(runtime_key); app.config.setup.therock_venv = Some(install_root); app.config.active_runtime_key = Some(runtime_key.to_owned()); @@ -36699,7 +36686,7 @@ Full log .paths .data_dir .join("runtimes") - .join("pip") + .join("wheel") .join(runtime_key); let registry_path = app .paths @@ -36754,7 +36741,7 @@ Full log .paths .data_dir .join("runtimes") - .join("pip") + .join("wheel") .join(runtime_key); app.config.setup.therock_venv = Some(install_root.clone()); app.config.save(&app.paths)?; @@ -36832,7 +36819,7 @@ Full log "--channel".to_owned(), "release".to_owned(), "--format".to_owned(), - "pip".to_owned(), + "wheel".to_owned(), "--prefix".to_owned(), install_root.display().to_string(), "--family".to_owned(), @@ -36945,12 +36932,12 @@ Full log "--channel".to_owned(), "release".to_owned(), "--format".to_owned(), - "pip".to_owned(), + "wheel".to_owned(), "--prefix".to_owned(), folder.to_owned(), ], Some(format!( - "/install sdk --channel release --format pip --prefix {folder}" + "/install sdk --channel release --format wheel --prefix {folder}" )), Some("Install ROCm/TheRock into the folder the user selected.".to_owned()), ); @@ -36989,7 +36976,7 @@ Full log "--channel".to_owned(), "release".to_owned(), "--format".to_owned(), - "pip".to_owned(), + "wheel".to_owned(), "--prefix".to_owned(), folder.to_owned(), ], @@ -37744,32 +37731,21 @@ Full log "Checking Python for the ROCm install...", "Checking Python...", ), - ("Checking Python package metadata...", "Checking Python..."), - ( - "Using existing Python 3.12.10 at C:\\Users\\jam\\.rocm\\tools\\python\\python.exe.", - "Checking Python...", - ), + ("Preparing Python 3.12...", "Checking Python..."), ( - "Using downloaded Python 3.12.10 package C:\\Users\\jam\\.rocm\\cache\\tools\\python\\cpython.tar.gz.", + "Using existing Python 3.12 at C:\\Users\\jam\\.rocm\\tools\\python\\python.exe.", "Checking Python...", ), ( - "Python 3.12.10 is ready at C:\\Users\\jam\\.rocm\\tools\\python\\python.exe.", + "Python 3.12 is ready at C:\\Users\\jam\\.rocm\\tools\\python\\python.exe.", "Checking Python...", ), - ( - "Downloading Python 3.12.10 package cpython-3.12.10.tar.gz.", - "Downloading Python...", - ), - ("Installing Python 3.12.10...", "Installing Python..."), + ("Finding Python 3.12...", "Checking Python..."), + ("Installing Python 3.12 via uv...", "Installing Python..."), ( "Creating Python environment at C:\\Users\\jam\\.rocm\\envs\\default.", "Setting up Python...", ), - ( - "Installing pip in the Python environment...", - "Preparing the installer...", - ), ( "Installing rocm[libraries,devel]==7.13.0a20260513 torch==2.10.0+rocm7.13.0a20260513 torchvision==0.25.0+rocm7.13.0a20260513 torchaudio==2.10.0+rocm7.13.0a20260513 from https://rocm.nightlies.amd.com/v2/gfx120X-all", "Installing ROCm packages...", @@ -38457,7 +38433,7 @@ Full log let install_root = paths .data_dir .join("runtimes") - .join("pip") + .join("wheel") .join(runtime_key); let local_manifest_path = install_root.join(".rocm-cli-runtime.json"); let mut manifest: Value = serde_json::from_slice(&fs::read(&local_manifest_path)?)?; @@ -38487,7 +38463,7 @@ Full log let install_root = paths .data_dir .join("runtimes") - .join("pip") + .join("wheel") .join(runtime_key); let scripts_dir = install_root.join(if cfg!(windows) { "Scripts" } else { "bin" }); fs::create_dir_all(&scripts_dir)?; @@ -38898,12 +38874,12 @@ Full log !menu .items .iter() - .any(|item| item == "pip" || item == "tarball") + .any(|item| item == "wheel" || item == "tarball") ); assert_eq!( menu.replacements, vec![ - Some("/install sdk --format pip ".to_owned()), + Some("/install sdk --format wheel ".to_owned()), Some("/install sdk --format tarball ".to_owned()), ] ); @@ -40240,7 +40216,7 @@ Full log "--channel".to_owned(), "release".to_owned(), "--format".to_owned(), - "pip".to_owned(), + "wheel".to_owned(), ]); let rendered = render_test_terminal(&app, 120, 32); @@ -40454,7 +40430,7 @@ Full log "--channel".to_owned(), "release".to_owned(), "--format".to_owned(), - "pip".to_owned(), + "wheel".to_owned(), ]); handle_key(&mut app, key_event(KeyCode::Char('d'), KeyModifiers::NONE)); @@ -41279,7 +41255,7 @@ Full log .paths .data_dir .join("runtimes") - .join("pip") + .join("wheel") .join(runtime_key); app.config.setup.therock_venv = Some(install_root); app.config.save(&app.paths)?; @@ -42128,7 +42104,7 @@ Full log "--channel", "release", "--format", - "pip", + "wheel", "--prefix", folder_text.as_str(), ]); @@ -44652,7 +44628,7 @@ Full log let install_root = paths .data_dir .join("runtimes") - .join("pip") + .join("wheel") .join(runtime_key); let scripts_dir = install_root.join(if cfg!(windows) { "Scripts" } else { "bin" }); let python_executable = scripts_dir.join(if cfg!(windows) { @@ -44694,7 +44670,7 @@ Full log ] }, "channel": "release", - "format": "pip", + "format": "wheel", "family": runtime_id.split(':').next_back().unwrap_or("gfx120X-all"), "family_source": "test", "version": "6.4.0", diff --git a/apps/rocmd/src/lib.rs b/apps/rocmd/src/lib.rs index 67526a66..94cb1a00 100644 --- a/apps/rocmd/src/lib.rs +++ b/apps/rocmd/src/lib.rs @@ -1734,7 +1734,7 @@ fn rocm_mcp_tools() -> Vec { }, "format": { "type": "string", - "enum": ["pip", "tarball"] + "enum": ["wheel", "tarball"] }, "prefix": { "type": "string" @@ -1763,7 +1763,7 @@ fn rocm_mcp_tools() -> Vec { }, "format": { "type": "string", - "enum": ["pip", "tarball"] + "enum": ["wheel", "tarball"] }, "prefix": { "type": "string" @@ -2402,7 +2402,7 @@ fn build_install_sdk_args( let format = arguments .get("format") .and_then(Value::as_str) - .unwrap_or("pip"); + .unwrap_or("wheel"); let prefix = arguments.get("prefix").and_then(Value::as_str); let version = arguments.get("version").and_then(Value::as_str); let build_date = arguments.get("build_date").and_then(Value::as_str); @@ -5293,7 +5293,7 @@ mod tests { "--channel".to_owned(), "release".to_owned(), "--format".to_owned(), - "pip".to_owned(), + "wheel".to_owned(), "--build-date".to_owned(), "2026-06-05".to_owned(), "--dry-run".to_owned(), diff --git a/crates/rocm-core/src/lib.rs b/crates/rocm-core/src/lib.rs index 0317cc3b..57a829b6 100644 --- a/crates/rocm-core/src/lib.rs +++ b/crates/rocm-core/src/lib.rs @@ -23,6 +23,11 @@ use windows_sys::Win32::System::Threading::{ }; pub mod runtime; +pub mod uv; +pub use uv::{ + DEFAULT_UV_TIMEOUT_SECS, ensure_uv_binary, uv_binary_name, uv_command_env, + uv_http_timeout_secs, uv_pip_freeze_args, uv_pip_install_base, uv_venv_args, +}; #[cfg(test)] use runtime::home_rocm_dir; pub use runtime::{ diff --git a/crates/rocm-core/src/uv.rs b/crates/rocm-core/src/uv.rs new file mode 100644 index 00000000..009d79ac --- /dev/null +++ b/crates/rocm-core/src/uv.rs @@ -0,0 +1,366 @@ +//! Acquisition and invocation helpers for the [`uv`](https://github.com/astral-sh/uv) +//! package manager. +//! +//! `uv` replaces the previous `python -m venv` + `python -m ensurepip` + +//! `python -m pip install` flow used to provision managed runtimes. This module owns +//! downloading a standalone `uv` binary into the managed cache and the small set of +//! argument/environment helpers shared by `apps/rocm` and the engine crates (both of +//! which depend only on `rocm-core`). + +use anyhow::{Context, Result, bail}; +use serde::{Deserialize, Serialize}; +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::{AppPaths, download_file_to_path, unix_time_millis}; + +/// Default network timeout, in seconds, applied to `uv` HTTP operations. +pub const DEFAULT_UV_TIMEOUT_SECS: u64 = 600; + +/// Environment variable consulted to point at a preinstalled `uv` binary, bypassing the +/// managed download (used by orchestrators and for offline/air-gapped hosts). +pub const UV_BINARY_ENV: &str = "ROCM_CLI_UV_BINARY"; + +/// Environment variable used to pin the downloaded `uv` release (e.g. `0.8.4`). Defaults +/// to the latest published release. +pub const UV_VERSION_ENV: &str = "ROCM_CLI_UV_VERSION"; + +/// Environment variable used to tune the `uv` network timeout, in seconds. Falls back to +/// the legacy `ROCM_CLI_PIP_TIMEOUT_SECS` for compatibility. +pub const UV_TIMEOUT_ENV: &str = "ROCM_CLI_UV_TIMEOUT_SECS"; +const LEGACY_PIP_TIMEOUT_ENV: &str = "ROCM_CLI_PIP_TIMEOUT_SECS"; + +#[derive(Debug, Clone, Serialize, Deserialize)] +struct ManagedUvManifest { + version: String, + asset: String, + source_url: String, + executable: PathBuf, + installed_at_unix_ms: u128, +} + +/// The platform-specific file name of the `uv` executable. +pub fn uv_binary_name() -> &'static str { + if runtime_is_windows() { "uv.exe" } else { "uv" } +} + +/// The network timeout applied to `uv` operations, honoring `ROCM_CLI_UV_TIMEOUT_SECS` +/// then the legacy `ROCM_CLI_PIP_TIMEOUT_SECS`. +pub fn uv_http_timeout_secs() -> u64 { + env_secs(UV_TIMEOUT_ENV) + .or_else(|| env_secs(LEGACY_PIP_TIMEOUT_ENV)) + .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![( + "UV_HTTP_TIMEOUT".to_owned(), + uv_http_timeout_secs().to_string(), + )] +} + +/// Arguments for `uv venv`, creating an environment at `env_root` using `python`. +pub fn uv_venv_args(python: &Path, env_root: &Path) -> Vec { + vec![ + "venv".to_owned(), + "--python".to_owned(), + python.to_string_lossy().into_owned(), + env_root.to_string_lossy().into_owned(), + ] +} + +/// Base arguments for `uv pip install` targeting the interpreter `venv_python`. Callers +/// append index/cache/package arguments. +pub fn uv_pip_install_base(venv_python: &Path) -> Vec { + vec![ + "pip".to_owned(), + "install".to_owned(), + "--python".to_owned(), + venv_python.to_string_lossy().into_owned(), + ] +} + +/// Arguments for `uv pip freeze` targeting the interpreter `venv_python`. +pub fn uv_pip_freeze_args(venv_python: &Path) -> Vec { + vec![ + "pip".to_owned(), + "freeze".to_owned(), + "--python".to_owned(), + venv_python.to_string_lossy().into_owned(), + ] +} + +/// Ensure a usable `uv` binary is available, downloading and caching one if needed. +/// Returns the path to the executable. +pub fn ensure_uv_binary(paths: &AppPaths) -> Result { + if let Some(path) = uv_binary_override() { + return Ok(path); + } + + let version = uv_version(); + let asset = uv_asset_name()?; + let install_dir = managed_tools_dir(&paths.data_dir) + .join("uv") + .join(slug(&version)); + let binary_name = uv_binary_name(); + + if let Some(existing) = find_binary_in(&install_dir, binary_name) + && uv_binary_is_usable(&existing) + { + return Ok(existing); + } + + let url = uv_download_url(&version, &asset); + let archive_path = paths + .cache_dir + .join("tools") + .join("uv") + .join(slug(&version)) + .join(&asset); + eprintln!("Downloading uv ({version}) from {url}"); + download_file_to_path( + &url, + &archive_path, + Duration::from_secs(uv_http_timeout_secs()), + ) + .with_context(|| format!("failed to download uv from {url}"))?; + + let staging = install_dir.with_extension(format!("tmp-{}", unix_time_millis())); + let _ = std::fs::remove_dir_all(&staging); + std::fs::create_dir_all(&staging) + .with_context(|| format!("failed to create {}", staging.display()))?; + extract_archive(&archive_path, &staging) + .with_context(|| format!("failed to extract uv archive {}", archive_path.display()))?; + + let staged_binary = find_binary_in(&staging, binary_name).with_context(|| { + format!( + "uv archive {} did not contain a `{binary_name}` executable", + archive_path.display() + ) + })?; + make_executable(&staged_binary)?; + + let _ = std::fs::remove_dir_all(&install_dir); + if let Some(parent) = install_dir.parent() { + std::fs::create_dir_all(parent) + .with_context(|| format!("failed to create {}", parent.display()))?; + } + std::fs::rename(&staging, &install_dir).or_else(|_| { + let _ = std::fs::remove_dir_all(&install_dir); + std::fs::rename(&staging, &install_dir) + })?; + + let binary = find_binary_in(&install_dir, binary_name).with_context(|| { + format!( + "uv executable missing after install at {}", + install_dir.display() + ) + })?; + if !uv_binary_is_usable(&binary) { + bail!( + "downloaded uv at {} is not runnable", + binary.display() + ); + } + + let manifest = ManagedUvManifest { + version, + asset, + source_url: url, + executable: binary.clone(), + installed_at_unix_ms: unix_time_millis(), + }; + write_uv_manifest(paths, &manifest); + let _ = std::fs::remove_file(&archive_path); + + Ok(binary) +} + +fn uv_binary_override() -> Option { + let value = std::env::var_os(UV_BINARY_ENV)?; + if value.is_empty() { + return None; + } + let path = PathBuf::from(value); + path.is_file().then_some(path) +} + +fn uv_version() -> String { + std::env::var(UV_VERSION_ENV) + .ok() + .map(|value| value.trim().to_owned()) + .filter(|value| !value.is_empty()) + .unwrap_or_else(|| "latest".to_owned()) +} + +fn uv_asset_name() -> Result { + let triple = match (runtime_os_name(), std::env::consts::ARCH) { + ("linux", "x86_64") => "x86_64-unknown-linux-gnu", + ("linux", "aarch64") => "aarch64-unknown-linux-gnu", + ("windows", "x86_64") => "x86_64-pc-windows-msvc", + ("windows", "aarch64") => "aarch64-pc-windows-msvc", + ("macos", "x86_64") => "x86_64-apple-darwin", + ("macos", "aarch64") => "aarch64-apple-darwin", + (os, arch) => bail!("unsupported platform for uv download: {os}/{arch}"), + }; + let extension = if runtime_is_windows() { "zip" } else { "tar.gz" }; + Ok(format!("uv-{triple}.{extension}")) +} + +fn uv_download_url(version: &str, asset: &str) -> String { + if version == "latest" { + format!("https://github.com/astral-sh/uv/releases/latest/download/{asset}") + } else { + format!("https://github.com/astral-sh/uv/releases/download/{version}/{asset}") + } +} + +fn extract_archive(archive_path: &Path, target_dir: &Path) -> Result<()> { + // System `tar` handles both `.tar.gz` (-xf auto-detects gzip) and `.zip` (bsdtar on + // Windows 10+), avoiding extra archive crates in rocm-core. + let status = Command::new("tar") + .arg("-xf") + .arg(archive_path) + .arg("-C") + .arg(target_dir) + .status() + .with_context(|| format!("failed to launch tar to extract {}", archive_path.display()))?; + if !status.success() { + bail!( + "tar exited with {status} while extracting {}", + archive_path.display() + ); + } + Ok(()) +} + +fn find_binary_in(dir: &Path, name: &str) -> Option { + let direct = dir.join(name); + if direct.is_file() { + return Some(direct); + } + let entries = std::fs::read_dir(dir).ok()?; + for entry in entries.flatten() { + let path = entry.path(); + if path.is_dir() + && let Some(found) = find_binary_in(&path, name) + { + return Some(found); + } else if path.is_file() + && path.file_name().and_then(|value| value.to_str()) == Some(name) + { + return Some(path); + } + } + None +} + +fn make_executable(path: &Path) -> Result<()> { + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + std::fs::set_permissions(path, std::fs::Permissions::from_mode(0o755)) + .with_context(|| format!("failed to mark {} executable", path.display()))?; + } + #[cfg(not(unix))] + { + let _ = path; + } + Ok(()) +} + +fn uv_binary_is_usable(path: &Path) -> bool { + Command::new(path) + .arg("--version") + .stdin(std::process::Stdio::null()) + .stdout(std::process::Stdio::null()) + .stderr(std::process::Stdio::null()) + .status() + .map(|status| status.success()) + .unwrap_or(false) +} + +fn write_uv_manifest(paths: &AppPaths, manifest: &ManagedUvManifest) { + let registry = managed_tools_dir(&paths.data_dir).join("registry"); + if std::fs::create_dir_all(®istry).is_err() { + return; + } + if let Ok(bytes) = serde_json::to_vec_pretty(manifest) { + let _ = std::fs::write(registry.join("uv.json"), bytes); + } +} + +fn env_secs(name: &str) -> Option { + std::env::var(name) + .ok() + .and_then(|value| value.trim().parse::().ok()) + .filter(|value| *value > 0) +} + +fn slug(value: &str) -> String { + value + .chars() + .map(|ch| { + if ch.is_ascii_alphanumeric() || matches!(ch, '.' | '-' | '_') { + ch + } else { + '-' + } + }) + .collect() +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn asset_name_has_archive_extension() { + // Whatever the host, the asset is one of the two known archive kinds. + let asset = uv_asset_name().expect("supported host platform for tests"); + assert!(asset.ends_with(".tar.gz") || asset.ends_with(".zip"), "{asset}"); + assert!(asset.starts_with("uv-"), "{asset}"); + } + + #[test] + fn download_url_pins_explicit_version() { + assert_eq!( + uv_download_url("0.8.4", "uv-x86_64-unknown-linux-gnu.tar.gz"), + "https://github.com/astral-sh/uv/releases/download/0.8.4/uv-x86_64-unknown-linux-gnu.tar.gz" + ); + } + + #[test] + fn download_url_uses_latest_redirect() { + assert_eq!( + uv_download_url("latest", "uv-x86_64-unknown-linux-gnu.tar.gz"), + "https://github.com/astral-sh/uv/releases/latest/download/uv-x86_64-unknown-linux-gnu.tar.gz" + ); + } + + #[test] + fn venv_args_target_python_and_root() { + let args = uv_venv_args(Path::new("/py/bin/python3"), Path::new("/envs/run")); + assert_eq!(args, vec!["venv", "--python", "/py/bin/python3", "/envs/run"]); + } + + #[test] + fn pip_install_base_targets_venv_python() { + let args = uv_pip_install_base(Path::new("/envs/run/bin/python")); + assert_eq!( + args, + vec!["pip", "install", "--python", "/envs/run/bin/python"] + ); + } + + #[test] + fn slug_sanitizes_unexpected_characters() { + assert_eq!(slug("0.8.4"), "0.8.4"); + assert_eq!(slug("latest"), "latest"); + assert_eq!(slug("weird/version space"), "weird-version-space"); + } +} diff --git a/engines/pytorch/src/lib.rs b/engines/pytorch/src/lib.rs index 7cc9c9ff..b9bad621 100644 --- a/engines/pytorch/src/lib.rs +++ b/engines/pytorch/src/lib.rs @@ -2,10 +2,11 @@ use anyhow::{Context, Result, bail}; use clap::{Parser, Subcommand, ValueEnum}; use rocm_core::{ AppPaths, DEFAULT_LOCAL_PORT, ModelRecipeRecord, detect_host_therock_family, - extract_first_gfx_token, format_host_port, format_http_base_url, interactive_terminal, - normalize_runtime_path_for_host, normalize_runtime_path_text_for_host, + 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, + resolve_model_recipe as resolve_shared_model_recipe, runtime_is_windows, uv_command_env, + uv_pip_freeze_args, uv_pip_install_base, uv_venv_args, unix_time_millis, }; use rocm_engine_protocol::{ DetectRequest, DetectResponse, DevicePolicy, ENGINE_RECIPE_CONTRACT_VERSION, EndpointRequest, @@ -30,8 +31,6 @@ 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 DEFAULT_PIP_TIMEOUT_SECS: u64 = 600; -const DEFAULT_PIP_RETRIES: u32 = 8; const ENGINE_DEPENDENCIES: &[&str] = &[ "fastapi", "uvicorn", @@ -1335,9 +1334,6 @@ fn create_or_update_env_manifest(request: &InstallRequest) -> Result Result Result Result { - install_therock_torch_packages(&python_executable_string, &resolution, &paths)?; - let stack_pip_args = pip_install_network_args(&paths); - run_progress_command( - &python_executable_string, - std::iter::once("-m") - .chain(std::iter::once("pip")) - .chain(std::iter::once("install")) - .chain(stack_pip_args.iter().map(String::as_str)) - .chain(TORCH_STACK_DEPENDENCIES.iter().copied()), + 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(|s| s.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()); @@ -1475,9 +1464,10 @@ fn create_or_update_env_manifest(request: &InstallRequest) -> Result Result Result> { } fn install_therock_torch_packages( - python_executable: &str, + uv: &Path, + python_executable: &Path, resolution: &TheRockTorchResolution, - paths: &AppPaths, ) -> Result<()> { - let mut args = vec!["-m".to_owned(), "pip".to_owned(), "install".to_owned()]; - args.extend(pip_install_network_args_allow_source(paths)); + let mut args = uv_pip_install_base(python_executable); args.push("--index-url".to_owned()); args.push(resolution.index_url.clone()); - args.push("--upgrade-strategy".to_owned()); - args.push("only-if-needed".to_owned()); if matches!(resolution.channel, TheRockChannel::Nightly) { - args.push("--pre".to_owned()); + args.extend(["--prerelease".to_owned(), "allow".to_owned()]); } args.extend(resolution.packages.iter().cloned()); - run_progress_command( - python_executable, + run_uv_progress_command( + uv, args.iter().map(String::as_str), "install TheRock torch packages into managed pytorch env", ) @@ -2789,77 +2776,66 @@ except Exception as exc: })) "#; -fn ensure_pip_available(python_executable: &str) -> Result<()> { - if command_succeeds(python_executable, ["-m", "pip", "--version"])? { +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(()); } - - run_command( - python_executable, - ["-m", "ensurepip", "--upgrade"], - "bootstrap pip in managed pytorch env", - )?; - run_command( - python_executable, - ["-m", "pip", "--version"], - "verify pip in managed pytorch env", - ) + let stderr = String::from_utf8_lossy(&output.stderr).trim().to_owned(); + bail!("{context_text}: uv exited with {}: {stderr}", output.status) } -fn pip_timeout_secs() -> u64 { - std::env::var("ROCM_CLI_PIP_TIMEOUT_SECS") - .ok() - .and_then(|value| value.trim().parse::().ok()) - .filter(|value| *value > 0) - .unwrap_or(DEFAULT_PIP_TIMEOUT_SECS) +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 pip_retries() -> u32 { - std::env::var("ROCM_CLI_PIP_RETRIES") - .ok() - .and_then(|value| value.trim().parse::().ok()) - .unwrap_or(DEFAULT_PIP_RETRIES) -} - -fn pip_install_network_args(paths: &AppPaths) -> Vec { - let mut args = vec![ - "--timeout".to_owned(), - pip_timeout_secs().to_string(), - "--retries".to_owned(), - pip_retries().to_string(), - "--only-binary".to_owned(), - ":all:".to_owned(), - "--disable-pip-version-check".to_owned(), - "--progress-bar".to_owned(), - "on".to_owned(), - ]; - args.push("--cache-dir".to_owned()); - args.push(pip_cache_dir(paths, ENGINE_NAME).display().to_string()); - args -} - -fn pip_install_network_args_allow_source(paths: &AppPaths) -> Vec { - let mut args = vec![ - "--timeout".to_owned(), - pip_timeout_secs().to_string(), - "--retries".to_owned(), - pip_retries().to_string(), - "--disable-pip-version-check".to_owned(), - "--progress-bar".to_owned(), - "on".to_owned(), - ]; - args.push("--cache-dir".to_owned()); - args.push(pip_cache_dir(paths, ENGINE_NAME).display().to_string()); - args -} - -fn pip_cache_dir(paths: &AppPaths, component: &str) -> PathBuf { - std::env::var_os("ROCM_CLI_PIP_CACHE_DIR") - .filter(|value| !value.as_os_str().is_empty()) - .map(PathBuf::from) - .unwrap_or_else(|| paths.cache_dir.join("pip").join(component)) +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, @@ -2999,6 +2975,7 @@ fn flag_arg(flag: &str, enabled: bool) -> Vec { } } +#[allow(dead_code)] fn run_command<'a, I>(program: &str, args: I, context_label: &str) -> Result<()> where I: IntoIterator, @@ -3017,6 +2994,7 @@ where } } +#[allow(dead_code)] fn run_progress_command<'a, I>(program: &str, args: I, context_label: &str) -> Result<()> where I: IntoIterator, @@ -3041,10 +3019,12 @@ where 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(), @@ -3052,6 +3032,7 @@ fn env_value_truthy(value: Option<&str>) -> bool { ) } +#[allow(dead_code)] fn run_progress_command_forwarded( program: &str, args: &[String], @@ -3093,6 +3074,7 @@ fn run_progress_command_forwarded( ); } +#[allow(dead_code)] fn forward_child_output_to_stderr(mut reader: R) -> thread::JoinHandle> where R: Read + Send + 'static, @@ -3621,40 +3603,12 @@ mod tests { } #[test] - fn pip_install_network_args_are_wheel_only() { - let paths = AppPaths { - config_dir: PathBuf::from("config"), - data_dir: PathBuf::from("data"), - cache_dir: PathBuf::from("cache"), - }; - let args = pip_install_network_args(&paths); - assert!( - args.windows(2) - .any(|pair| pair == ["--only-binary", ":all:"]) - ); - assert!( - args.windows(2) - .any(|pair| pair[0] == "--cache-dir" && pair[1].contains("pytorch")) - ); - } - - #[test] - fn therock_pip_args_allow_rocm_source_package() { - let paths = AppPaths { - config_dir: PathBuf::from("config"), - data_dir: PathBuf::from("data"), - cache_dir: PathBuf::from("cache"), - }; - let args = pip_install_network_args_allow_source(&paths); - assert!( - !args - .windows(2) - .any(|pair| pair == ["--only-binary", ":all:"]) - ); - assert!( - args.windows(2) - .any(|pair| pair[0] == "--cache-dir" && pair[1].contains("pytorch")) - ); + 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] diff --git a/scripts/smoke_local.py b/scripts/smoke_local.py index 861052f3..2ef4de79 100644 --- a/scripts/smoke_local.py +++ b/scripts/smoke_local.py @@ -616,8 +616,8 @@ def main() -> int: ) assert_path_missing( - smoke_root / "rocm-cache" / "pip", - "first-run smoke pip cache", + smoke_root / "rocm-cache" / "uv", + "first-run smoke uv cache", ) assert_path_missing( smoke_root / "rocm-data" / "runtimes" / "registry", diff --git a/scripts/therock_sdk_install_test.py b/scripts/therock_sdk_install_test.py index 8d2a565e..9d1d43d6 100644 --- a/scripts/therock_sdk_install_test.py +++ b/scripts/therock_sdk_install_test.py @@ -1,8 +1,8 @@ #!/usr/bin/env python3 -"""Opt-in acceptance test for managed TheRock SDK pip installs. +"""Opt-in acceptance test for managed TheRock SDK wheel installs. This test creates an isolated rocm-cli state root, creates a local bootstrap -Python venv, runs `rocm install sdk --format pip`, and verifies the installed +Python venv, runs `rocm install sdk --format wheel`, and verifies the installed TheRock SDK venv with `python -m rocm_sdk` commands. It downloads TheRock wheels and can take a while, so it is not part of the @@ -257,7 +257,7 @@ def verify_manifest( manifest: dict[str, Any], expected_pip_cache_dir: Path, ) -> tuple[Path, Path, Path, Path]: - if manifest.get("format") != "pip": + if manifest.get("format") != "wheel"": fail(f"expected pip runtime manifest, got: {manifest.get('format')}") python = Path(str(manifest.get("python_executable") or "")) if not python.is_file(): @@ -429,7 +429,7 @@ def main() -> int: doctor = run("rocm doctor before SDK install", [str(rocm), "doctor"], env=env, timeout=120) assert_contains(doctor, "rocm doctor", "doctor") - install_argv = [str(rocm), "install", "sdk", "--channel", args.channel, "--format", "pip"] + install_argv = [str(rocm), "install", "sdk", "--channel", args.channel, "--format", "wheel"] if args.prefix is not None: install_argv.extend(["--prefix", str(args.prefix)]) if args.dry_run: @@ -446,7 +446,7 @@ def main() -> int: "summary: rocm-cli will install the ROCm SDK and matching PyTorch packages", "sdk install", ) - assert_contains(install_output, "format: pip", "sdk install") + assert_contains(install_output, "format: wheel", "sdk install") assert_contains(install_output, "pip_cache_dir:", "sdk install") assert_contains(install_output, "latest_compatible_version:", "sdk install") assert_contains(install_output, "python_wheel_tag:", "sdk install")