Skip to content

fix(therock): resolve ROCm releases from the current multi-arch pip index #272

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
juhovainio wants to merge 4 commits into from
Closed

fix(therock): resolve ROCm releases from the current multi-arch pip index #272

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
db42fd2
fix(therock): resolve ROCm releases from the current multi-arch pip i…
juhovainio Aug 17, 2026
8c62a8b
test(e2e): guard against the broken multi-arch pip index URL shape
juhovainio Aug 17, 2026
50ecdbb
fix: address PR #272 review feedback
juhovainio Aug 31, 2026
68934fa
style: run cargo fmt on therock.rs import list
juhovainio Aug 31, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
112 changes: 97 additions & 15 deletions apps/rocm/src/therock.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,12 +6,12 @@ use anyhow::{Context, Result, bail};
use rocm_core::{
AppPaths, ManagedToolConfig, RocmCliConfig, detect_host_gpu_diagnostics,
detect_host_therock_family, detect_managed_therock_family, disk_space, ensure_uv_binary,
known_therock_families, 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, 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, verify_rsa_pkcs1_sha256_signature,
known_therock_families, known_therock_family_device_chips, 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, 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, verify_rsa_pkcs1_sha256_signature,
};
#[cfg(test)]
use rocm_core::{
Expand Down Expand Up @@ -826,6 +826,7 @@ fn install_wheel_runtime(
"Found TheRock package family {} version {} with a matching PyTorch stack.",
resolution.family, resolution.latest_version
));
let rocm_extras = therock_rocm_extras(&resolution.family, &resolution.index_url);
let runtime_key = runtime_key(
channel,
"wheel",
Expand Down Expand Up @@ -880,11 +881,11 @@ fn install_wheel_runtime(
let _ = writeln!(
output,
" package_specs: {}",
therock_pip_package_specs(&resolution.package_versions).join(" ")
therock_pip_package_specs(&resolution.package_versions, &rocm_extras).join(" ")
);
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 uv transaction"
" package_policy: find the newest TheRock ROCm SDK version that has a matching PyTorch stack in the same index, then install pinned rocm (with any resolved device extras), torch, torchvision, and torchaudio versions in one uv transaction"
);
if dry_run {
let env_python = venv_python_path(&install_root);
Expand All @@ -893,7 +894,10 @@ fn install_wheel_runtime(
if matches!(channel, TheRockChannel::Nightly) {
install_args.extend(["--prerelease".to_owned(), "allow".to_owned()]);
}
install_args.extend(therock_pip_package_specs(&resolution.package_versions));
install_args.extend(therock_pip_package_specs(
&resolution.package_versions,
&rocm_extras,
));
let venv_args = uv_venv_args(&python_launcher.executable, &install_root);
let venv_args_display = venv_args
.iter()
Expand Down Expand Up @@ -933,15 +937,18 @@ fn install_wheel_runtime(

progress_line(format!(
"Installing {} from {}",
therock_pip_package_specs(&resolution.package_versions).join(" "),
therock_pip_package_specs(&resolution.package_versions, &rocm_extras).join(" "),
resolution.index_url
));
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.extend(["--prerelease".to_owned(), "allow".to_owned()]);
}
install_args.extend(therock_pip_package_specs(&resolution.package_versions));
install_args.extend(therock_pip_package_specs(
&resolution.package_versions,
&rocm_extras,
));
run_uv_progress_command(
paths,
&uv,
Expand Down Expand Up @@ -1012,15 +1019,42 @@ fn install_wheel_runtime(
Ok(output)
}

fn therock_pip_package_specs(package_versions: &TheRockPipPackageVersions) -> Vec<String> {
fn therock_pip_package_specs(
package_versions: &TheRockPipPackageVersions,
rocm_extras: &str,
) -> Vec<String> {
vec![
format!("rocm[libraries,devel]=={}", package_versions.rocm),
format!("rocm[{rocm_extras}]=={}", package_versions.rocm),

@rominf rominf Sep 2, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This shape change breaks a documented acceptance test that CI can't see.

scripts/therock_sdk_install_test.py — unchanged by this PR — asserts:

THEROCK_SDK_PACKAGE_SPEC = "rocm[libraries,devel]"   # :29
...
assert_contains(install_output, f"{THEROCK_SDK_PACKAGE_SPEC}==", "sdk install")   # :499

After this change the release-channel output is rocm[libraries,devel,device-gfx1151]==7.14.0, and "rocm[libraries,devel]==" is not a substring of that. The script's default is --channel release (:412), so every documented invocation in docs/testing.md:180/186/193/200 and docs/manual-testing.md:140 fails on the first release run. It's manual-only — not in ci.yml — so nothing catches it, and AGENTS.md §8 names docs/testing.md checks as part of the gate for touched behavior.

The fix is small: make the constant a prefix, "rocm[libraries,devel", and assert on assert_contains(install_output, THEROCK_SDK_PACKAGE_SPEC). The existing assert_not_contains(install_output, "rocm[devel]") negative still does its job.

Same class of staleness in prose at docs/testing.md:167, which lists the expected plan as "pinned rocm[libraries,devel], torch, torchvision, and torchaudio versions" — worth a sentence noting the release channel now adds device-*.

format!("torch=={}", package_versions.torch),
format!("torchvision=={}", package_versions.torchvision),
format!("torchaudio=={}", package_versions.torchaudio),
]
}

fn is_multi_arch_pip_index(index_url: &str) -> bool {
index_url.trim_end_matches('/') == THEROCK_RELEASE_PIP_MULTI_ARCH_INDEX_BASE
}

/// The `rocm[...]` extras to request for a resolved pip index. The classic
/// per-family index needs only the base `libraries,devel` extras; the flat
/// multi-arch index additionally needs an explicit `device-*` extra or no GPU
/// backend gets installed at all.
fn therock_rocm_extras(family: &str, index_url: &str) -> String {
let mut extras = "libraries,devel".to_owned();
if !is_multi_arch_pip_index(index_url) {
return extras;
}
match known_therock_family_device_chips(family) {
Some(chips) => {
for chip in chips {
let _ = write!(extras, ",device-{chip}");
}
}
None => extras.push_str(",device-all"),
}
extras

@rominf rominf Sep 2, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This fallback fires on every Instinct/datacenter part, and your own CI shows what it costs.

known_therock_family_device_chips returns None for six of the sixteen known families — gfx90X-dgpu, gfx90X-dcgpu, gfx94X-dcgpu, gfx950-dcgpu, gfx101X-dgpu, gfx103X-dgpu. The two gfx94X/gfx950 entries are MI300/MI350; that's the datacenter audience taking the device-all path by default.

From job 99541673879 (MI300X, detected_gfx_target: gfx943 → gfx94X-dcgpu):

Installing rocm[libraries,devel,device-all]==7.14.0 ... from https://repo.amd.com/rocm/whl-multi-arch
Downloading rocm-sdk-device-gfx1030 (458.5MiB)
Downloading rocm-sdk-device-gfx90a  (450.7MiB)
Downloading rocm-sdk-device-gfx950  (650.9MiB)
... 24 device wheels, 4451 MiB total

and then the post-install probe:

rocm_sdk_target_family: gfx1010

So on a gfx943 host we download ~4.3 GiB to use one wheel, and the probe reports the wrong target family — gfx1010 is just the first of the 24. That value is printed to the user and stored in the manifest. The exact-chip lanes in the same run resolve correctly (gfx1151→gfx1151, gfx120X-all→gfx1201), which localises it to device-all.

There's no size guard either: preflight_tarball_space is called only from install_tarball_runtime (:1119), so the wheel path has no disk preflight at all. A constrained host gets a multi-GiB surprise with no warning.

Four of the six are enumerable right now, from the wheel names the same CI log lists:

"gfx94X-dcgpu"  => Some(&["gfx942"]),
"gfx950-dcgpu"  => Some(&["gfx950"]),
"gfx101X-dgpu"  => Some(&["gfx1010", "gfx1011", "gfx1012"]),
"gfx103X-dgpu"  => Some(&["gfx1030", "gfx1031", "gfx1032", "gfx1033",
                          "gfx1034", "gfx1035", "gfx1036"]),

I checked these against normalize_therock_family — gfx942→gfx94X-dcgpu (the starts_with("gfx94") arm), gfx950→gfx950-dcgpu, gfx101*→gfx101X-dgpu, gfx103*→gfx103X-dgpu — so known_therock_family_device_chips_round_trip_to_their_family stays green with all four added. The remaining two (gfx90X-dgpu/gfx90X-dcgpu) are genuine catch-alls that overlap the exact gfx900/gfx906/gfx908/gfx90a families, and device-all is a defensible answer there.

The better answer, if you're reconciling with #329 anyway, is #329's shape: derive the extra from the detected raw arch (device-{raw_arch}) and keep the family table only as the fallback when no raw arch is known. Either way, when device-all is chosen, say so in the dry-run output — right now the user sees it in package_specs with no indication it means "every GPU ROCm supports."

}

fn quote_display_arg(value: &str) -> String {
if value.is_empty()
|| value
Expand Down Expand Up @@ -3368,9 +3402,12 @@ fn parse_version(value: &str) -> Option<ParsedVersion> {

fn therock_index_urls(channel: TheRockChannel, family: &str) -> Vec<String> {

@rominf rominf Sep 2, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

On the #329 collision you flagged — one detail worth pinning down before you agree a merge order, because it's not just a textual conflict.

#329's replacement for this function still contains the bug:

// pr/329 apps/rocm/src/therock.rs:3632-3638
TheRockChannel::Release => vec![
    format!("{}/{family}", release_pip_index_base()),
    format!("{}/{family}", release_pip_multi_arch_index_base()),
],

Classic first, then the per-family multi-arch shape — byte-for-byte main's behavior. #329's new Next candidate doesn't cover this: it points at a different host (whl-next) and returns None outright when raw_arch is unknown, so it never reaches repo.amd.com/rocm/whl-multi-arch. If #329's therock.rs wins the reconciliation, #271 regresses with no test failing — including the new scenario above, for the reason in that comment.

The other direction is the good news: #329's therock_pip_package_specs already derives rocm[libraries,devel,device-{raw_arch}] from the detected arch (and extends it to torch[device-*]/torchvision[device-*], which this PR doesn't), which is the better answer to the device-all problem. So the reconciled version wants this PR's flat candidate plus #329's arch-driven extras — a three-variant generation enum (Legacy / MultiArch / Next) rather than either side winning wholesale.

Whoever merges second should re-verify against #271's actual repro, not just resolve conflicts and trust the suite.

match channel {
// Multi-arch is flat (no per-family path segment) and is where AMD
// publishes current releases; try it first. The classic per-family
// index is kept as a fallback so older releases stay installable.
TheRockChannel::Release => vec![
THEROCK_RELEASE_PIP_MULTI_ARCH_INDEX_BASE.to_owned(),
format!("{THEROCK_RELEASE_PIP_INDEX_BASE}/{family}"),
format!("{THEROCK_RELEASE_PIP_MULTI_ARCH_INDEX_BASE}/{family}"),
],
TheRockChannel::Nightly => vec![format!("{THEROCK_NIGHTLY_PIP_INDEX_BASE}/{family}")],
}
Expand Down Expand Up @@ -4065,7 +4102,7 @@ mod tests {
torchaudio: "2.10.0+rocm7.13.0a20260513".to_owned(),
compatibility_key: "7.13.0a20260513".to_owned(),
};
let package_specs = therock_pip_package_specs(&package_versions);
let package_specs = therock_pip_package_specs(&package_versions, "libraries,devel");

assert_eq!(
package_specs,
Expand All @@ -4078,6 +4115,51 @@ mod tests {
);
}

#[test]
fn therock_index_urls_prefers_multi_arch_then_classic_for_release() {
let urls = therock_index_urls(TheRockChannel::Release, "gfx110X-all");
assert_eq!(
urls,
vec![
"https://repo.amd.com/rocm/whl-multi-arch".to_owned(),
"https://repo.amd.com/rocm/whl/gfx110X-all".to_owned(),
]
);
}

#[test]
fn therock_index_urls_nightly_unchanged() {
let urls = therock_index_urls(TheRockChannel::Nightly, "gfx110X-all");
assert_eq!(
urls,
vec!["https://rocm.nightlies.amd.com/v2/gfx110X-all".to_owned()]
);
}

#[test]
fn therock_rocm_extras_classic_index_is_unchanged() {
assert_eq!(
therock_rocm_extras("gfx110X-all", "https://repo.amd.com/rocm/whl/gfx110X-all"),
"libraries,devel"
);
}

#[test]
fn therock_rocm_extras_multi_arch_adds_exact_device_chips() {
assert_eq!(
therock_rocm_extras("gfx110X-all", "https://repo.amd.com/rocm/whl-multi-arch"),
"libraries,devel,device-gfx1100,device-gfx1101,device-gfx1102,device-gfx1103"
);
}

#[test]
fn therock_rocm_extras_multi_arch_falls_back_to_device_all_for_ambiguous_bucket() {
assert_eq!(
therock_rocm_extras("gfx90X-dcgpu", "https://repo.amd.com/rocm/whl-multi-arch"),
"libraries,devel,device-all"
);
}

/// The downloaded archive is removed once it has been unpacked; keeping it
/// would double the disk cost of every installed SDK version.
#[test]
Expand Down
35 changes: 35 additions & 0 deletions crates/rocm-core/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3969,6 +3969,26 @@ pub const fn known_therock_families() -> &'static [&'static str] {
]
}

/// Exact chip IDs covered by a [`known_therock_families`] bucket.
///
/// `None` for the prefix catch-all buckets (e.g. `gfx90X-dgpu`) whose exact
/// chip membership isn't enumerable from [`normalize_therock_family`] alone.
pub fn known_therock_family_device_chips(family: &str) -> Option<&'static [&'static str]> {
match family {
"gfx900" => Some(&["gfx900"]),
"gfx906" => Some(&["gfx906"]),
"gfx908" => Some(&["gfx908"]),
"gfx90a" => Some(&["gfx90a"]),
"gfx110X-all" => Some(&["gfx1100", "gfx1101", "gfx1102", "gfx1103"]),
"gfx1150" => Some(&["gfx1150"]),
"gfx1151" => Some(&["gfx1151"]),
"gfx1152" => Some(&["gfx1152"]),
"gfx1153" => Some(&["gfx1153"]),
"gfx120X-all" => Some(&["gfx1200", "gfx1201"]),
_ => None,
}
}

fn capture_optional_command(program: &str, args: &[&str]) -> Option<String> {
capture_optional_command_with_timeout(program, args, OPTIONAL_COMMAND_TIMEOUT)
}
Expand Down Expand Up @@ -8897,6 +8917,21 @@ mod tests {
assert!(!known_therock_families().is_empty());
}

#[test]
fn known_therock_family_device_chips_round_trip_to_their_family() {
for family in known_therock_families() {
if let Some(chips) = known_therock_family_device_chips(family) {
for chip in chips {
assert_eq!(
normalize_therock_family(chip).as_deref(),
Some(*family),
"chip `{chip}` must normalize back to family `{family}`"
);
}
}
}
}

#[test]
fn preferred_serve_engine_uses_vllm_for_supported_therock_families() {
assert_eq!(
Expand Down
20 changes: 20 additions & 0 deletions tests/e2e-cucumber/features/runtime_setup.feature
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -84,3 +84,23 @@ Feature: Runtime configuration
When the user tries to adopt the existing install
Then the adoption is refused
And the error explains which install types can be adopted

# Regression guard for the bug where the release-channel multi-arch pip index
# was queried with a per-family path segment (.../whl-multi-arch/{family}/),
# which 403s because that index is flat and 404/403-ed straight into the
# stale classic index every time, silently pinning every install to
# whatever version predated the migration. `--family` bypasses GPU
# auto-detection so this needs no GPU, and `--dry-run` resolves the real
# index without installing anything.
#
# `@nightly` for the same reason as scenario 8 above: this dry-run still
# resolves the real channel index over the network (dry-run only skips the
# venv/download, not index resolution), and the no-GPU mock lane's 64-way
# concurrency from that extra network work is what pushes
# `eai-7960-gen-tps-held-after-scrape-failure` and
# `eai-7960-gen-tps-expiry-boundary` past their validity window. Runs on the
# nightly lanes instead, where scenarios are serialized.
@id:runtime-install-sdk-release-index-shape @nightly
Scenario: 4 - Resolving the SDK from the release channel never uses the broken multi-arch URL shape

@rominf rominf Aug 28, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This scenario issues rocm install sdk --family gfx110X-all --dry-run, which resolves the real repo.amd.com pip index over the network before printing the plan (dry-run skips the install, not the index resolution). That's the same call the runtime-install-records-the-real-folder scenario just above makes, and that one is tagged @nightly specifically because the extra concurrent network work on the no-GPU mock lane (64 scenarios at once) was measured to push two other scenarios past their timing budget 3/3 times. This scenario has no @nightly/@requires-gpu tag, so it will run at that same concurrency on every PR and can reintroduce the same flakiness the sibling scenario's comment describes. Consider tagging it @nightly (or otherwise excluding it from the unserialized mock lane) to match the sibling.

When the user dry-runs installing the SDK for a known family

@rominf rominf Sep 2, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Scenario number 4 is already taken — :30, "Reinstalling the SDK leaves the installed engine's requirements satisfied."

Nothing fails: the harness's uniqueness assert (tests/e2e-cucumber/tests/e2e.rs:1179) covers @id, not the numeric prefix. But the numbers are how the file and the HTML report are navigated, and two "Scenario: 4"s in one feature make them indistinguishable there. Next free in this file is 9.

Then the resolved package index is not the broken per-family multi-arch path
18 changes: 18 additions & 0 deletions tests/e2e-cucumber/tests/e2e/runtime_steps.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -202,6 +202,15 @@ fn assert_engine_ready(world: &mut E2eWorld) {
);
}

#[when("the user dry-runs installing the SDK for a known family")]
async fn user_dry_runs_install_sdk_for_family(world: &mut E2eWorld) {
let stdout = crate::run_rocm_ok(
world,
&["install", "sdk", "--family", "gfx110X-all", "--dry-run"],
);
world.cli_output = Some(stdout);
}

#[when("the user tries to adopt the existing install")]
async fn user_tries_adopt(world: &mut E2eWorld) {
let (stdout, stderr, rc) = crate::run_rocm(
Expand Down Expand Up @@ -365,6 +374,15 @@ async fn assert_update_reports_freshness(world: &mut E2eWorld) {
}
}

#[then("the resolved package index is not the broken per-family multi-arch path")]
async fn assert_index_not_broken_multi_arch(world: &mut E2eWorld) {
let stdout = world.cli_output.as_deref().expect("no dry-run output");
assert!(
!stdout.contains("whl-multi-arch/gfx110X-all"),

@rominf rominf Sep 2, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This assertion is satisfied on unpatched main, so the scenario doesn't guard the fix.

On main, therock_index_urls(Release, family) returns [whl/{family}, whl-multi-arch/{family}] — classic first. resolve_pip_runtime_with_timeout loops the candidates and returns on the first success. The classic index does resolve (that's the premise of #271: it works, it's just pinned at 7.13.0), so the loop returns immediately and the dry-run prints:

index_url: https://repo.amd.com/rocm/whl/gfx110X-all

The broken whl-multi-arch/gfx110X-all URL only ever reaches stdout through the bail! at the bottom of that loop, which is unreachable when a candidate succeeds. So !stdout.contains("whl-multi-arch/gfx110X-all") holds with the bug present, and holds after the fix, and would keep holding if someone reverted therock_index_urls tomorrow. AGENTS.md §3: "test fails before fix and passes after fix."

Asserting the positive fixes it, and it's a one-line change:

assert!(
    stdout.contains("index_url: https://repo.amd.com/rocm/whl-multi-arch\n"),
    "release install did not resolve the flat multi-arch index:\n{stdout}"
);

That fails on main (which prints the classic URL) and passes here. Worth also asserting package_specs: carries a device- extra, since that's the other half of the fix and currently has only unit coverage.

One edge worth naming: if the classic index ever stops resolving entirely, main would bail! and run_rocm_ok's rc == 0 check would panic first — the scenario would fail, but on the command failing, not on the assertion. The positive form is correct in both worlds.

"dry-run resolved the broken per-family multi-arch index shape:\n{stdout}"
);
}

#[then("the adoption is refused")]
async fn assert_adoption_refused(world: &mut E2eWorld) {
let rc = world.cli_rc.expect("no command was run");
Expand Down
Loading