rusty_stack/installers/components/
rccl.rs

1//! Targeted workaround for the ROCm 7.2.1+ RCCL compiler regression.
2//!
3//! The distro RCCL package remains authoritative. This module only builds and
4//! activates a sealed overlay when a two-GPU collective reproduces a known
5//! ROCm/RCCL multi-GPU failure fingerprint.
6
7use anyhow::{bail, Context, Result};
8use sha2::{Digest, Sha256};
9use std::collections::BTreeSet;
10use std::fs;
11use std::io::{Read, Write};
12#[cfg(unix)]
13use std::os::unix::fs::PermissionsExt;
14use std::path::{Path, PathBuf};
15use std::process::{Command, Output};
16
17pub const RCCL_COMMIT: &str = "96a25b5fd6f73fba58c7d83eb57cf19a50230aa4";
18const RCCL_REPO: &str = "https://github.com/ROCm/rccl.git";
19const LIBXML_URL: &str = "https://archive.ubuntu.com/ubuntu/pool/main/libx/libxml2/libxml2_2.9.13+dfsg-1ubuntu0.12_amd64.deb";
20const LIBXML_SHA256: &str = "b3678e6e4b166bc0e4226fb118d489ab51802c772914421397c9dcb2dd0e0d2b";
21const ICU_URL: &str = "https://archive.ubuntu.com/ubuntu/pool/main/i/icu/libicu70_70.1-2_amd64.deb";
22const ICU_SHA256: &str = "58a154f6307289813da2276f900498ef536ae7c0522d2cf31a3c3c5cf62dfd9a";
23const ISSUE_URL: &str = "https://github.com/ROCm/ROCm/issues/6074";
24
25#[derive(Debug, Clone, Copy)]
26struct RcclRepairProfile {
27    id: &'static str,
28    compiler_id: &'static str,
29    compiler_url: &'static str,
30    compiler_sha256: &'static str,
31    compiler_size: u64,
32    hipify_id: &'static str,
33    hipify_url: &'static str,
34    hipify_sha256: &'static str,
35    hipify_size: u64,
36}
37
38const ROCM_72_REPAIR_PROFILE: RcclRepairProfile = RcclRepairProfile {
39    id: "rocm-7.2.x",
40    compiler_id: "rocm-llvm-7.2.0",
41    compiler_url: "https://repo.radeon.com/rocm/apt/7.2/pool/main/r/rocm-llvm/rocm-llvm_22.0.0.26014.70200-43~22.04_amd64.deb",
42    compiler_sha256: "88d604fe0ab4a2502f0c7e0f8d3b7b49e829ae3057971de681a68b2fdb2505d2",
43    compiler_size: 549_321_042,
44    hipify_id: "hipify-rocm-7.2.0",
45    hipify_url: "https://raw.githubusercontent.com/ROCm/HIPIFY/refs/tags/rocm-7.2.0/bin/hipify-perl",
46    hipify_sha256: "036b3b07a4fe9a921d220f21a1fa5a3cfc22343bb0417d729015b87ae0463291",
47    hipify_size: 948_312,
48};
49
50#[derive(Debug, Clone, PartialEq, Eq)]
51pub struct ShellCommand {
52    pub program: String,
53    pub args: Vec<String>,
54    pub env: Vec<(String, String)>,
55}
56
57#[derive(Debug, Clone)]
58pub struct RcclInstaller {
59    home: PathBuf,
60    python_bin: PathBuf,
61    rocm_path: PathBuf,
62    rocm_version: String,
63    rocm_channel: String,
64    gpu_archs: Vec<String>,
65    gpu_count: usize,
66    rocr_tokens: Vec<String>,
67}
68
69impl RcclInstaller {
70    pub fn new(
71        home: impl Into<PathBuf>,
72        python_bin: impl Into<PathBuf>,
73        rocm_path: impl Into<PathBuf>,
74        rocm_version: impl Into<String>,
75        rocm_channel: impl Into<String>,
76        rocr_visible_devices: Option<String>,
77    ) -> Self {
78        let gpus = crate::gpu::detect_discrete_amd_gpus();
79        let gpu_archs = gpus
80            .iter()
81            .filter_map(|gpu| gpu.gfx_arch.clone())
82            .collect::<BTreeSet<_>>()
83            .into_iter()
84            .collect();
85        let rocr_tokens = rocr_visible_devices
86            .as_deref()
87            .map(parse_visibility_tokens)
88            .filter(|tokens| !tokens.is_empty())
89            .unwrap_or_else(|| (0..gpus.len()).map(|idx| idx.to_string()).collect());
90        Self {
91            home: home.into(),
92            python_bin: python_bin.into(),
93            rocm_path: rocm_path.into(),
94            rocm_version: rocm_version.into(),
95            rocm_channel: rocm_channel.into(),
96            gpu_count: gpus.len(),
97            gpu_archs,
98            rocr_tokens,
99        }
100    }
101
102    /// Probe the installed RCCL and return a build command only for ROCm#6074.
103    pub fn prepare_remediation(&self) -> Result<Option<ShellCommand>> {
104        if self.gpu_count < 2 {
105            return Ok(None);
106        }
107        if self.gpu_archs.is_empty() {
108            bail!(
109                "RCCL remediation requires detected gfx architectures; refusing a guessed target"
110            );
111        }
112        if self.rocr_tokens.len() < 2 {
113            bail!("RCCL remediation requires at least two ROCr visibility tokens");
114        }
115
116        match self.active_overlay_state()? {
117            OverlayState::Valid { version_dir, sha } => {
118                let probe = self.run_probe(Some((&version_dir, &sha)))?;
119                if overlay_probe_succeeded(&probe) {
120                    return Ok(None);
121                }
122                if !unseal_requested() {
123                    bail!(
124                        "sealed RCCL overlay failed functional revalidation; set MLSTACK_UNSEAL_CORE=1 to rebuild it:\n{}",
125                        output_text(&probe)
126                    );
127                }
128            }
129            OverlayState::Invalid(reason) if !unseal_requested() => {
130                bail!(
131                    "sealed RCCL overlay is invalid ({reason}); set MLSTACK_UNSEAL_CORE=1 to rebuild it"
132                )
133            }
134            OverlayState::Invalid(_) | OverlayState::Missing => {}
135        }
136
137        let probe = self.run_probe(None)?;
138        if probe.status.success() {
139            return Ok(None);
140        }
141        let output = output_text(&probe);
142        if !is_compiler_regression(&output) {
143            bail!("RCCL multi-GPU probe failed outside the known ROCm/RCCL collective failure fingerprints:\n{output}");
144        }
145
146        let profile = self.repair_profile()?;
147        let compiler = self.ensure_toolchain(profile)?;
148        let hipify_perl = self.ensure_hipify_perl(profile)?;
149        Ok(Some(self.build_command(&compiler, &hipify_perl)?))
150    }
151
152    /// Publish the completed build as a read-only overlay and prove both ranks use it.
153    pub fn activate_and_verify(&self) -> Result<()> {
154        let profile = self.repair_profile()?;
155        let source = self.source_dir();
156        let built = find_named(&source.join("build"), "librccl.so.1.0")
157            .or_else(|| find_named(&source.join("build"), "librccl.so.1"))
158            .context("RCCL build completed without librccl.so.1.0")?;
159        let sha = sha256_file(&built)?;
160        let version_name = format!("{}-{}", &RCCL_COMMIT[..12], profile.compiler_id);
161        let root = self.overlay_root();
162        let version_dir = root.join(&version_name);
163        let staging = root.join(format!(".{version_name}.tmp-{}", std::process::id()));
164
165        fs::create_dir_all(&root)?;
166        if staging.exists() {
167            fs::remove_dir_all(&staging)?;
168        }
169        if version_dir.exists() {
170            if !unseal_requested() {
171                bail!(
172                    "RCCL overlay version already exists but was not reusable; refusing overwrite"
173                );
174            }
175            make_tree_writable(&version_dir)?;
176            fs::remove_dir_all(&version_dir)?;
177        }
178
179        let lib_dir = staging.join("lib");
180        let py_dir = staging.join("python");
181        fs::create_dir_all(&lib_dir)?;
182        fs::create_dir_all(&py_dir)?;
183        fs::copy(&built, lib_dir.join("librccl.so.1.0"))?;
184        make_symlink("librccl.so.1.0", &lib_dir.join("librccl.so.1"))?;
185        make_symlink("librccl.so.1", &lib_dir.join("librccl.so"))?;
186        fs::write(py_dir.join("sitecustomize.py"), SITECUSTOMIZE_SOURCE)?;
187
188        let archs = self.gpu_archs.join(";");
189        fs::write(
190            staging.join("manifest"),
191            format!(
192                "source={RCCL_REPO}\ncommit={RCCL_COMMIT}\nprofile={}\ncompiler={}\nhipify={}\narchitectures={archs}\nsha256={sha}\nissue={ISSUE_URL}\n",
193                profile.id, profile.compiler_id, profile.hipify_id
194            ),
195        )?;
196        // NOTE: RCCL no longer writes its own env.sh/env.fish. The persistent-env
197        // generator (bootstrap::env_setup::rccl_overlay_*_block) is the SINGLE
198        // source of truth for env vars — it reads the `active` symlink (which
199        // points at this `version_dir`) plus the manifest above and emits an
200        // idempotent, guarded block. The old component-written env did an
201        // un-guarded `export PYTHONPATH=X${PYTHONPATH:+:$PYTHONPATH}` that
202        // accumulated a duplicate entry on every `source ~/.mlstack_env`.
203
204        set_tree_read_only(&staging)?;
205        fs::rename(&staging, &version_dir)?;
206
207        let probe = self.run_probe(Some((&version_dir, &sha)))?;
208        let output = output_text(&probe);
209        if !overlay_probe_succeeded(&probe) {
210            make_tree_writable(&version_dir)?;
211            fs::remove_dir_all(&version_dir)?;
212            bail!("corrected RCCL overlay failed final two-GPU validation:\n{output}");
213        }
214
215        let active_tmp = root.join(".active.tmp");
216        if active_tmp.exists() {
217            fs::remove_file(&active_tmp)?;
218        }
219        make_symlink(&version_name, &active_tmp)?;
220        fs::rename(&active_tmp, root.join("active"))?;
221        Ok(())
222    }
223
224    fn overlay_root(&self) -> PathBuf {
225        self.home.join(".mlstack/components/rccl")
226    }
227
228    fn source_dir(&self) -> PathBuf {
229        self.home
230            .join(".mlstack/src")
231            .join(format!("rccl-{}", &RCCL_COMMIT[..12]))
232    }
233
234    fn active_overlay_state(&self) -> Result<OverlayState> {
235        let active = self.overlay_root().join("active");
236        let metadata = match fs::symlink_metadata(&active) {
237            Ok(metadata) => metadata,
238            Err(err) if err.kind() == std::io::ErrorKind::NotFound => {
239                return Ok(OverlayState::Missing);
240            }
241            Err(err) => return Err(err.into()),
242        };
243        if !metadata.file_type().is_symlink() {
244            return Ok(OverlayState::Invalid(
245                "active marker is not a symlink".into(),
246            ));
247        }
248        if !active.exists() {
249            return Ok(OverlayState::Invalid(
250                "active symlink target is missing".into(),
251            ));
252        }
253        let manifest = active.join("manifest");
254        let library = active.join("lib/librccl.so.1.0");
255        if !manifest.is_file() || !library.is_file() {
256            return Ok(OverlayState::Invalid("missing manifest or library".into()));
257        }
258        let text = fs::read_to_string(manifest)?;
259        let expected = text
260            .lines()
261            .find_map(|line| line.strip_prefix("sha256="))
262            .unwrap_or_default();
263        if expected.is_empty() || sha256_file(&library)? != expected {
264            return Ok(OverlayState::Invalid("SHA-256 mismatch".into()));
265        }
266        if !active.join("python/sitecustomize.py").is_file() {
267            return Ok(OverlayState::Invalid("activation files missing".into()));
268        }
269        // NOTE: do NOT require env.sh / env.fish here. This component no longer
270        // writes them — the persistent-env generator
271        // (bootstrap::env_setup::rccl_overlay_*_block) is the single source of
272        // truth for env vars and emits the RCCL overlay block inline in
273        // ~/.mlstack_env / activate-global.fish. Requiring the old component-
274        // written files would mark every newly-installed overlay Invalid and
275        // never report Valid.
276        Ok(OverlayState::Valid {
277            version_dir: active,
278            sha: expected.to_string(),
279        })
280    }
281
282    fn probe_path(&self) -> PathBuf {
283        self.home.join(".mlstack/cache/rccl-two-rank-probe.py")
284    }
285
286    fn run_probe(&self, overlay: Option<(&Path, &str)>) -> Result<Output> {
287        let path = self.probe_path();
288        fs::create_dir_all(path.parent().context("probe cache has no parent")?)?;
289        fs::write(&path, PROBE_SOURCE)?;
290
291        let mut command = Command::new(&self.python_bin);
292        command.arg(&path);
293        command.env_remove("HIP_VISIBLE_DEVICES");
294        command.env_remove("CUDA_VISIBLE_DEVICES");
295        command.env_remove("GPU_DEVICE_ORDINAL");
296        command.env_remove("MLSTACK_RCCL_OVERLAY_LIB");
297        command.env_remove("MLSTACK_RCCL_OVERLAY_SHA256");
298        command.env_remove("MLSTACK_RCCL_LOADER_PID");
299        command.env("ROCR_VISIBLE_DEVICES", self.rocr_tokens.join(","));
300        if let Some((version_dir, sha)) = overlay {
301            command.env(
302                "MLSTACK_RCCL_OVERLAY_LIB",
303                version_dir.join("lib/librccl.so.1.0"),
304            );
305            command.env("MLSTACK_RCCL_OVERLAY_SHA256", sha);
306            command.env("NCCL_P2P_DISABLE", "1");
307            command.env("RCCL_P2P_DISABLE", "1");
308            let existing = std::env::var_os("PYTHONPATH").unwrap_or_default();
309            let mut paths = vec![version_dir.join("python")];
310            paths.extend(std::env::split_paths(&existing));
311            command.env("PYTHONPATH", std::env::join_paths(paths)?);
312        }
313        command.output().with_context(|| {
314            format!(
315                "failed to launch RCCL probe with {}",
316                self.python_bin.display()
317            )
318        })
319    }
320
321    fn repair_profile(&self) -> Result<&'static RcclRepairProfile> {
322        if self.rocm_channel.eq_ignore_ascii_case("latest") && self.rocm_version.starts_with("7.2")
323        {
324            return Ok(&ROCM_72_REPAIR_PROFILE);
325        }
326        bail!(
327            "RCCL repair has no pinned profile for ROCm channel '{}' version '{}'",
328            self.rocm_channel,
329            self.rocm_version
330        )
331    }
332
333    fn ensure_toolchain(&self, profile: &RcclRepairProfile) -> Result<PathBuf> {
334        let root = self
335            .home
336            .join(format!(".mlstack/toolchains/{}", profile.compiler_id));
337        if let Some(clang) = find_named(&root, "clang++") {
338            if compiler_works(&clang, &compat_library_path(&root)) {
339                return Ok(clang);
340            }
341        }
342
343        let cache = self.home.join(".mlstack/cache/downloads");
344        fs::create_dir_all(&cache)?;
345        let compiler_deb = cache.join(format!("{}.deb", profile.compiler_id));
346        download_pinned(
347            profile.compiler_url,
348            profile.compiler_sha256,
349            Some(profile.compiler_size),
350            &compiler_deb,
351        )?;
352        fs::create_dir_all(&root)?;
353        extract_deb(&compiler_deb, &root)?;
354
355        let clang = find_named(&root, "clang++").context("pinned ROCm compiler lacks clang++")?;
356        if !compiler_works(&clang, "") {
357            let compat = root.join("compat");
358            fs::create_dir_all(&compat)?;
359            let xml = cache.join("libxml2-2.9.13.deb");
360            let icu = cache.join("libicu70.deb");
361            download_pinned(LIBXML_URL, LIBXML_SHA256, None, &xml)?;
362            download_pinned(ICU_URL, ICU_SHA256, None, &icu)?;
363            extract_deb(&xml, &compat)?;
364            extract_deb(&icu, &compat)?;
365        }
366        if !compiler_works(&clang, &compat_library_path(&root)) {
367            bail!(
368                "pinned ROCm compiler profile {} cannot run, even with isolated compatibility DSOs",
369                profile.id
370            );
371        }
372        Ok(clang)
373    }
374
375    fn ensure_hipify_perl(&self, profile: &RcclRepairProfile) -> Result<PathBuf> {
376        let tool = self.home.join(format!(
377            ".mlstack/toolchains/{}/bin/hipify-perl",
378            profile.hipify_id
379        ));
380        if tool.is_file() && sha256_file(&tool)?.eq_ignore_ascii_case(profile.hipify_sha256) {
381            make_executable(&tool)?;
382            return Ok(tool);
383        }
384
385        let cache = self.home.join(".mlstack/cache/downloads");
386        fs::create_dir_all(&cache)?;
387        let cached = cache.join(profile.hipify_id);
388        download_pinned(
389            profile.hipify_url,
390            profile.hipify_sha256,
391            Some(profile.hipify_size),
392            &cached,
393        )?;
394        fs::create_dir_all(tool.parent().context("hipify tool path has no parent")?)?;
395        fs::copy(&cached, &tool)?;
396        make_executable(&tool)?;
397        if !sha256_file(&tool)?.eq_ignore_ascii_case(profile.hipify_sha256) {
398            bail!("cached hipify-perl SHA-256 changed during installation");
399        }
400        Ok(tool)
401    }
402
403    fn build_command(&self, compiler: &Path, hipify_perl: &Path) -> Result<ShellCommand> {
404        let source = self.source_dir();
405        fs::create_dir_all(source.parent().context("RCCL source has no parent")?)?;
406        let archs = self.gpu_archs.join(";");
407        let profile = self.repair_profile()?;
408        let compat = compat_library_path(
409            &self
410                .home
411                .join(format!(".mlstack/toolchains/{}", profile.compiler_id)),
412        );
413        let script = format!(
414            r#"set -eu
415if [ ! -d {src}/.git ]; then git clone --no-checkout {repo} {src}; fi
416git -C {src} fetch --depth 1 origin {commit}
417git -C {src} checkout --detach --force {commit}
418test "$(git -C {src} rev-parse HEAD)" = {commit}
419rm -rf {src}/build
420git -C {src} diff --quiet
421test -z "$(git -C {src} status --porcelain --untracked-files=all)"
422CC={cc} CXX={cxx} CFLAGS=--rocm-path={rocm} CXXFLAGS=--rocm-path={rocm} \
423cmake -S {src} -B {src}/build -DCMAKE_BUILD_TYPE=Release -DBUILD_TESTS=OFF \
424-DROCM_PATH={rocm} -DCMAKE_PREFIX_PATH={rocm} -DCMAKE_C_COMPILER={cc} \
425-DCMAKE_CXX_COMPILER={cxx} -DCMAKE_HIP_COMPILER={cxx} \
426-Dhipify-perl_executable={hipify} \
427-DGPU_TARGETS={archs_q} -DAMDGPU_TARGETS={archs_q}
428cmake --build {src}/build --parallel 2
429"#,
430            src = shell_quote_path(&source),
431            repo = shell_quote(RCCL_REPO),
432            commit = shell_quote(RCCL_COMMIT),
433            cc = shell_quote_path(&compiler.with_file_name("clang")),
434            cxx = shell_quote_path(compiler),
435            hipify = shell_quote_path(hipify_perl),
436            rocm = shell_quote_path(&self.rocm_path),
437            archs_q = shell_quote(&archs),
438        );
439        let mut env = Vec::new();
440        if !compat.is_empty() {
441            let inherited = std::env::var("LD_LIBRARY_PATH").unwrap_or_default();
442            env.push((
443                "LD_LIBRARY_PATH".into(),
444                if inherited.is_empty() {
445                    compat
446                } else {
447                    format!("{compat}:{inherited}")
448                },
449            ));
450        }
451        Ok(ShellCommand {
452            program: "sh".into(),
453            args: vec!["-c".into(), script],
454            env,
455        })
456    }
457}
458
459#[derive(Debug, Clone, PartialEq, Eq)]
460enum OverlayState {
461    Missing,
462    Valid { version_dir: PathBuf, sha: String },
463    Invalid(String),
464}
465
466pub fn is_compiler_regression(output: &str) -> bool {
467    let lower = output.to_ascii_lowercase();
468    let collective_context =
469        (lower.contains("rccl") || lower.contains("nccl") || lower.contains("processgroup"))
470            && (lower.contains("distributed")
471                || lower.contains("all_reduce")
472                || lower.contains("allreduce"));
473    let present_state_error = lower.contains("operation cannot be performed in present state")
474        || lower.contains("the operation cannot be performed in the present state");
475    let invalid_pointer_error = lower.contains("invalid device pointer")
476        && (lower.contains("ncclunhandledcudaerror")
477            || lower.contains("hip failure")
478            || lower.contains("unhandled cuda error"));
479    collective_context && (present_state_error || invalid_pointer_error)
480}
481
482fn output_text(output: &Output) -> String {
483    format!(
484        "{}{}",
485        String::from_utf8_lossy(&output.stdout),
486        String::from_utf8_lossy(&output.stderr)
487    )
488}
489
490fn parse_visibility_tokens(value: &str) -> Vec<String> {
491    value
492        .split(',')
493        .map(str::trim)
494        .filter(|token| !token.is_empty())
495        .map(str::to_string)
496        .collect()
497}
498
499fn overlay_probe_succeeded(output: &Output) -> bool {
500    if !output.status.success() {
501        return false;
502    }
503    let text = output_text(output);
504    text.matches("RCCL_OVERLAY_OK").count() >= 2 && text.matches("sum=3.0").count() >= 2
505}
506
507fn unseal_requested() -> bool {
508    std::env::var("MLSTACK_UNSEAL_CORE").as_deref() == Ok("1")
509}
510
511fn download_pinned(url: &str, expected: &str, size: Option<u64>, dest: &Path) -> Result<()> {
512    if dest.is_file()
513        && sha256_file(dest)?.eq_ignore_ascii_case(expected)
514        && size.is_none_or(|n| dest.metadata().map(|m| m.len() == n).unwrap_or(false))
515    {
516        return Ok(());
517    }
518    let part = dest.with_extension("part");
519    let _ = fs::remove_file(&part);
520    let response = ureq::get(url)
521        .header("User-Agent", "rusty-stack-rccl-remediation")
522        .call()
523        .with_context(|| format!("failed to download pinned artifact {url}"))?;
524    let mut reader = response.into_body().into_reader();
525    let mut file = fs::File::create(&part)?;
526    std::io::copy(&mut reader, &mut file)?;
527    file.flush()?;
528    file.sync_all()?;
529    if size.is_some_and(|n| part.metadata().map(|m| m.len() != n).unwrap_or(true)) {
530        let _ = fs::remove_file(&part);
531        bail!("pinned artifact size mismatch for {url}");
532    }
533    let actual = sha256_file(&part)?;
534    if !actual.eq_ignore_ascii_case(expected) {
535        let _ = fs::remove_file(&part);
536        bail!("pinned artifact SHA-256 mismatch for {url}: expected {expected}, got {actual}");
537    }
538    fs::rename(part, dest)?;
539    Ok(())
540}
541
542fn extract_deb(deb: &Path, dest: &Path) -> Result<()> {
543    let temp = tempfile::tempdir()?;
544    let status = Command::new("ar")
545        .arg("x")
546        .arg(deb)
547        .current_dir(temp.path())
548        .status()
549        .context("ar is required to extract the pinned ROCm toolchain")?;
550    if !status.success() {
551        bail!("failed to extract {}", deb.display());
552    }
553    let data = fs::read_dir(temp.path())?
554        .filter_map(|entry| entry.ok().map(|e| e.path()))
555        .find(|path| {
556            path.file_name()
557                .and_then(|n| n.to_str())
558                .is_some_and(|n| n.starts_with("data.tar"))
559        })
560        .context("Debian artifact has no data archive")?;
561    let status = Command::new("tar")
562        .args(["-xf"])
563        .arg(data)
564        .arg("-C")
565        .arg(dest)
566        .status()
567        .context("tar is required to extract the pinned ROCm toolchain")?;
568    if !status.success() {
569        bail!("failed to unpack {}", deb.display());
570    }
571    Ok(())
572}
573
574fn compiler_works(clang: &Path, library_path: &str) -> bool {
575    let lld = clang.with_file_name("ld.lld");
576    let mut cmd = Command::new(if lld.is_file() { &lld } else { clang });
577    cmd.arg("--version");
578    if !library_path.is_empty() {
579        cmd.env("LD_LIBRARY_PATH", library_path);
580    }
581    cmd.output().is_ok_and(|output| output.status.success())
582}
583
584fn compat_library_path(root: &Path) -> String {
585    let compat = root.join("compat");
586    let mut dirs = BTreeSet::new();
587    collect_library_dirs(&compat, &mut dirs);
588    std::env::join_paths(dirs)
589        .ok()
590        .and_then(|p| p.into_string().ok())
591        .unwrap_or_default()
592}
593
594fn collect_library_dirs(path: &Path, dirs: &mut BTreeSet<PathBuf>) {
595    let Ok(entries) = fs::read_dir(path) else {
596        return;
597    };
598    for entry in entries.flatten() {
599        let child = entry.path();
600        if child.is_dir() {
601            collect_library_dirs(&child, dirs);
602        } else if child
603            .file_name()
604            .and_then(|n| n.to_str())
605            .is_some_and(|n| n.starts_with("libxml2.so") || n.starts_with("libicu"))
606        {
607            if let Some(parent) = child.parent() {
608                dirs.insert(parent.to_path_buf());
609            }
610        }
611    }
612}
613
614fn find_named(root: &Path, name: &str) -> Option<PathBuf> {
615    let entries = fs::read_dir(root).ok()?;
616    for entry in entries.flatten() {
617        let path = entry.path();
618        if path.is_dir() {
619            if let Some(found) = find_named(&path, name) {
620                return Some(found);
621            }
622        } else if path.file_name().and_then(|n| n.to_str()) == Some(name) {
623            return Some(path);
624        }
625    }
626    None
627}
628
629fn sha256_file(path: &Path) -> Result<String> {
630    let mut file = fs::File::open(path)?;
631    let mut hash = Sha256::new();
632    let mut buffer = [0_u8; 1024 * 1024];
633    loop {
634        let read = file.read(&mut buffer)?;
635        if read == 0 {
636            break;
637        }
638        hash.update(&buffer[..read]);
639    }
640    Ok(hash
641        .finalize()
642        .iter()
643        .map(|byte| format!("{byte:02x}"))
644        .collect())
645}
646
647#[cfg(unix)]
648fn make_symlink(target: impl AsRef<Path>, link: &Path) -> Result<()> {
649    std::os::unix::fs::symlink(target, link)?;
650    Ok(())
651}
652
653#[cfg(not(unix))]
654fn make_symlink(_target: impl AsRef<Path>, _link: &Path) -> Result<()> {
655    bail!("RCCL overlay activation requires Unix symbolic links")
656}
657
658#[cfg(unix)]
659fn make_executable(path: &Path) -> Result<()> {
660    fs::set_permissions(path, fs::Permissions::from_mode(0o755))?;
661    Ok(())
662}
663
664#[cfg(not(unix))]
665fn make_executable(_path: &Path) -> Result<()> {
666    bail!("RCCL hipify tool requires Unix executable permissions")
667}
668
669#[cfg(unix)]
670fn set_tree_read_only(root: &Path) -> Result<()> {
671    for entry in fs::read_dir(root)? {
672        let path = entry?.path();
673        if path.is_dir() {
674            set_tree_read_only(&path)?;
675            fs::set_permissions(&path, fs::Permissions::from_mode(0o555))?;
676        } else if !path.is_symlink() {
677            fs::set_permissions(&path, fs::Permissions::from_mode(0o444))?;
678        }
679    }
680    fs::set_permissions(root, fs::Permissions::from_mode(0o555))?;
681    Ok(())
682}
683
684#[cfg(not(unix))]
685fn set_tree_read_only(_root: &Path) -> Result<()> {
686    Ok(())
687}
688
689#[cfg(unix)]
690fn make_tree_writable(root: &Path) -> Result<()> {
691    fs::set_permissions(root, fs::Permissions::from_mode(0o755))?;
692    for entry in fs::read_dir(root)? {
693        let path = entry?.path();
694        if path.is_dir() {
695            make_tree_writable(&path)?;
696        } else if !path.is_symlink() {
697            fs::set_permissions(&path, fs::Permissions::from_mode(0o644))?;
698        }
699    }
700    Ok(())
701}
702
703#[cfg(not(unix))]
704fn make_tree_writable(_root: &Path) -> Result<()> {
705    Ok(())
706}
707
708fn shell_quote(value: &str) -> String {
709    format!("'{}'", value.replace('\'', "'\"'\"'"))
710}
711
712fn shell_quote_path(path: &Path) -> String {
713    shell_quote(&path.to_string_lossy())
714}
715
716const SITECUSTOMIZE_SOURCE: &str = r#"import hashlib
717import importlib.machinery
718import os
719from pathlib import Path
720import sys
721import sysconfig
722
723def _activate_rccl_overlay():
724    lib = os.environ.get("MLSTACK_RCCL_OVERLAY_LIB")
725    expected = os.environ.get("MLSTACK_RCCL_OVERLAY_SHA256")
726    if not (lib and expected):
727        return
728    path = Path(lib).resolve()
729    try:
730        digest = hashlib.sha256()
731        with path.open("rb") as stream:
732            for chunk in iter(lambda: stream.read(1024 * 1024), b""):
733                digest.update(chunk)
734        actual = digest.hexdigest()
735    except OSError as exc:
736        sys.stderr.write(f"rusty-stack: sealed RCCL overlay unreadable: {exc}\n")
737        os._exit(126)
738    if actual != expected:
739        sys.stderr.write("rusty-stack: sealed RCCL overlay SHA-256 mismatch\n")
740        os._exit(126)
741
742    spec = importlib.machinery.PathFinder.find_spec("torch", sys.path)
743    if spec is not None and spec.origin:
744        torch_dir = Path(spec.origin).parent
745        torch_lib = torch_dir / "lib"
746        targets = list(torch_dir.glob("_C*.so"))
747        targets += [torch_lib / name for name in (
748            "libtorch_python.so", "libtorch.so", "libtorch_hip.so"
749        ) if (torch_lib / name).exists()]
750        loaders = (
751            Path("/lib64/ld-linux-x86-64.so.2"),
752            Path("/lib/x86_64-linux-gnu/ld-linux-x86-64.so.2"),
753        )
754        loader = next((candidate for candidate in loaders if candidate.exists()), None)
755        if loader is None or not targets:
756            sys.stderr.write("rusty-stack: cannot activate sealed RCCL overlay safely\n")
757            os._exit(126)
758        library_path = [str(path.parent), str(torch_lib)]
759        libdir = sysconfig.get_config_var("LIBDIR")
760        if libdir:
761            library_path.append(libdir)
762        if os.environ.get("LD_LIBRARY_PATH"):
763            library_path.append(os.environ["LD_LIBRARY_PATH"])
764        env = os.environ.copy()
765        env["MLSTACK_RCCL_LOADER_PID"] = str(os.getpid())
766        argv = list(getattr(sys, "orig_argv", [sys.executable, *sys.argv]))
767        os.execve(str(loader), [
768            str(loader), "--inhibit-rpath", ":".join(map(str, targets)),
769            "--library-path", ":".join(library_path),
770            sys.executable, *argv[1:]
771        ], env)
772
773if os.environ.get("MLSTACK_RCCL_LOADER_PID") == str(os.getpid()):
774    os.environ.pop("MLSTACK_RCCL_LOADER_PID", None)
775elif os.environ.get("MLSTACK_RCCL_OVERLAY_LIB") and os.environ.get("MLSTACK_RCCL_OVERLAY_SHA256"):
776    class _RustyRcclImportHook:
777        def find_spec(self, fullname, path=None, target=None):
778            if fullname == "torch" or fullname.startswith("torch."):
779                _activate_rccl_overlay()
780            return None
781    sys.meta_path.insert(0, _RustyRcclImportHook())
782"#;
783
784const PROBE_SOURCE: &str = r#"import os
785from pathlib import Path
786import subprocess
787import sys
788import tempfile
789
790rank_source = r'''
791import os
792from pathlib import Path
793import sys
794
795rank = int(sys.argv[1])
796rendezvous = sys.argv[2]
797token = sys.argv[3]
798os.environ["ROCR_VISIBLE_DEVICES"] = token
799os.environ.pop("HIP_VISIBLE_DEVICES", None)
800os.environ.pop("CUDA_VISIBLE_DEVICES", None)
801os.environ.pop("GPU_DEVICE_ORDINAL", None)
802
803import torch
804import torch.distributed as dist
805
806device_count = torch.cuda.device_count()
807if device_count != 1:
808    raise RuntimeError(f"rank {rank} expected one isolated GPU for ROCR token {token}, found {device_count}")
809name = torch.cuda.get_device_name(0)
810
811initialized = False
812try:
813    dist.init_process_group("nccl", init_method=f"file://{rendezvous}", rank=rank, world_size=2)
814    initialized = True
815    value = torch.tensor([float(rank + 1)], device="cuda:0")
816    dist.all_reduce(value)
817    torch.cuda.synchronize()
818    loaded = [line.split()[-1] for line in Path("/proc/self/maps").read_text().splitlines()
819              if "librccl.so" in line and "/" in line]
820    expected = os.environ.get("MLSTACK_RCCL_OVERLAY_LIB")
821    if expected and not any(Path(item).resolve() == Path(expected).resolve() for item in loaded):
822        raise RuntimeError(f"RCCL overlay not loaded: {loaded}")
823    if value.item() != 3.0:
824        raise RuntimeError(f"wrong collective sum {value.item()}")
825    print(f"RCCL_OVERLAY_OK rank={rank} token={token} name={name} sum={value.item():.1f} lib={loaded}", flush=True)
826finally:
827    if initialized:
828        dist.destroy_process_group()
829'''
830
831def rank_env(token):
832    env = os.environ.copy()
833    env["ROCR_VISIBLE_DEVICES"] = token
834    env.pop("HIP_VISIBLE_DEVICES", None)
835    env.pop("CUDA_VISIBLE_DEVICES", None)
836    env.pop("GPU_DEVICE_ORDINAL", None)
837    env.pop("MLSTACK_RCCL_LOADER_PID", None)
838    return env
839
840if __name__ == "__main__":
841    devices = [item.strip() for item in os.environ.get("ROCR_VISIBLE_DEVICES", "").split(",") if item.strip()]
842    if len(devices) < 2:
843        raise SystemExit("RCCL probe requires at least two ROCr visibility tokens")
844    fd, rendezvous = tempfile.mkstemp(prefix="rusty-rccl-")
845    os.close(fd)
846    os.unlink(rendezvous)
847    procs = []
848    try:
849        for rank, token in enumerate(devices[:2]):
850            procs.append((
851                rank,
852                token,
853                subprocess.Popen(
854                    [sys.executable, "-c", rank_source, str(rank), rendezvous, token],
855                    env=rank_env(token),
856                    stdout=subprocess.PIPE,
857                    stderr=subprocess.PIPE,
858                    text=True,
859                ),
860            ))
861        output = []
862        failed = False
863        for rank, token, proc in procs:
864            try:
865                stdout, stderr = proc.communicate(timeout=300)
866            except subprocess.TimeoutExpired:
867                proc.kill()
868                stdout, stderr = proc.communicate()
869                failed = True
870                stderr += f"\nrank {rank} token {token} timed out\n"
871            output.append(stdout)
872            output.append(stderr)
873            if proc.returncode != 0:
874                failed = True
875        text = "".join(output)
876        print(text, end="")
877        if failed:
878            raise SystemExit(1)
879        for rank, token in enumerate(devices[:2]):
880            if f"RCCL_OVERLAY_OK rank={rank} token={token} " not in text:
881                raise SystemExit(f"rank {rank} token {token} did not report success")
882    finally:
883        for _, _, proc in procs:
884            if proc.poll() is None:
885                proc.kill()
886        Path(rendezvous).unlink(missing_ok=True)
887"#;
888
889#[cfg(test)]
890mod tests {
891    use super::*;
892
893    #[test]
894    fn classifies_only_known_rocm_rccl_collective_failures() {
895        assert!(is_compiler_regression(
896            "torch.distributed ProcessGroupNCCL all_reduce: operation cannot be performed in present state"
897        ));
898        assert!(is_compiler_regression(
899            "torch.distributed all_reduce NCCL error ncclUnhandledCudaError HIP failure: invalid device pointer"
900        ));
901        assert!(!is_compiler_regression(
902            "operation cannot be performed in present state"
903        ));
904        assert!(!is_compiler_regression(
905            "HIP failure: invalid device pointer"
906        ));
907        assert!(!is_compiler_regression("distributed nccl timeout"));
908    }
909
910    #[test]
911    fn remediation_is_pinned_amd_only_and_fail_closed() {
912        assert!(ROCM_72_REPAIR_PROFILE
913            .compiler_url
914            .starts_with("https://repo.radeon.com/rocm/"));
915        assert!(ROCM_72_REPAIR_PROFILE
916            .hipify_url
917            .starts_with("https://raw.githubusercontent.com/ROCm/HIPIFY/"));
918        assert!(LIBXML_URL.starts_with("https://archive.ubuntu.com/ubuntu/"));
919        assert!(ICU_URL.starts_with("https://archive.ubuntu.com/ubuntu/"));
920        assert!(!SITECUSTOMIZE_SOURCE.contains("LD_PRELOAD"));
921        assert!(SITECUSTOMIZE_SOURCE.contains("--inhibit-rpath"));
922        assert!(SITECUSTOMIZE_SOURCE.contains("sha256"));
923        assert!(SITECUSTOMIZE_SOURCE.contains("MLSTACK_RCCL_LOADER_PID"));
924        assert!(SITECUSTOMIZE_SOURCE.contains("os.getpid()"));
925        assert!(SITECUSTOMIZE_SOURCE.contains("sys.meta_path.insert"));
926        assert!(PROBE_SOURCE.contains("ROCR_VISIBLE_DEVICES"));
927        assert!(PROBE_SOURCE.contains("RCCL_OVERLAY_OK"));
928        assert!(PROBE_SOURCE.contains("pop(\"HIP_VISIBLE_DEVICES\""));
929        assert!(PROBE_SOURCE.contains("pop(\"CUDA_VISIBLE_DEVICES\""));
930        assert!(PROBE_SOURCE.contains("pop(\"GPU_DEVICE_ORDINAL\""));
931        assert!(PROBE_SOURCE.contains("subprocess.Popen"));
932        assert!(!PROBE_SOURCE.contains("torch.multiprocessing"));
933    }
934
935    #[test]
936    fn overlay_lives_outside_python_environment() {
937        let temp = tempfile::tempdir().unwrap();
938        let home = temp.path().to_path_buf();
939        let installer = RcclInstaller {
940            python_bin: home.join(".mlstack/global/bin/python"),
941            home: home.clone(),
942            rocm_path: PathBuf::from("/opt/rocm"),
943            rocm_version: "7.2.4".into(),
944            rocm_channel: "latest".into(),
945            gpu_archs: vec!["gfx1100".into(), "gfx1101".into()],
946            gpu_count: 2,
947            rocr_tokens: vec!["2".into(), "5".into()],
948        };
949        assert_eq!(
950            installer.overlay_root(),
951            home.join(".mlstack/components/rccl")
952        );
953        let command = installer
954            .build_command(
955                Path::new("/toolchain/bin/clang++"),
956                Path::new("/toolchain/bin/hipify-perl"),
957            )
958            .unwrap();
959        let script = command.args.last().unwrap();
960        assert!(script.contains("BUILD_TESTS=OFF"));
961        assert!(script.contains("--parallel 2"));
962        assert!(script.contains("gfx1100;gfx1101"));
963        assert!(script.contains("--rocm-path="));
964        assert!(script.contains("-Dhipify-perl_executable="));
965    }
966
967    #[test]
968    fn overlay_exports_safe_consumer_rdna_transport() {
969        assert!(SITECUSTOMIZE_SOURCE.contains("MLSTACK_RCCL_OVERLAY_LIB"));
970        let env_text = format!(
971            "export MLSTACK_RCCL_OVERLAY_LIB={}\nexport MLSTACK_RCCL_OVERLAY_SHA256=sha\nexport NCCL_P2P_DISABLE=1\nexport RCCL_P2P_DISABLE=1\n",
972            shell_quote_path(Path::new("/overlay/lib/librccl.so.1.0"))
973        );
974        assert!(env_text.contains("NCCL_P2P_DISABLE=1"));
975        assert!(env_text.contains("RCCL_P2P_DISABLE=1"));
976    }
977
978    #[test]
979    fn rccl_repair_profile_is_version_aware() {
980        let installer = RcclInstaller {
981            home: PathBuf::from("/home/test"),
982            python_bin: PathBuf::from("python3"),
983            rocm_path: PathBuf::from("/opt/rocm"),
984            rocm_version: "7.2.4".into(),
985            rocm_channel: "latest".into(),
986            gpu_archs: vec!["gfx1100".into(), "gfx1101".into()],
987            gpu_count: 2,
988            rocr_tokens: vec!["0".into(), "1".into()],
989        };
990        assert_eq!(installer.repair_profile().unwrap().id, "rocm-7.2.x");
991
992        let legacy = RcclInstaller {
993            rocm_version: "6.4.3".into(),
994            rocm_channel: "legacy".into(),
995            ..installer
996        };
997        assert!(legacy.repair_profile().is_err());
998    }
999
1000    #[test]
1001    fn visibility_tokens_are_trimmed_without_guessing() {
1002        assert_eq!(
1003            parse_visibility_tokens(" 2, 5 ,,"),
1004            vec!["2".to_string(), "5".to_string()]
1005        );
1006    }
1007
1008    #[cfg(unix)]
1009    #[test]
1010    fn dangling_active_symlink_fails_closed() {
1011        let temp = tempfile::tempdir().unwrap();
1012        let installer = RcclInstaller {
1013            home: temp.path().to_path_buf(),
1014            python_bin: PathBuf::from("python3"),
1015            rocm_path: PathBuf::from("/opt/rocm"),
1016            rocm_version: "7.2.4".into(),
1017            rocm_channel: "latest".into(),
1018            gpu_archs: vec!["gfx1100".into(), "gfx1101".into()],
1019            gpu_count: 2,
1020            rocr_tokens: vec!["0".into(), "1".into()],
1021        };
1022        fs::create_dir_all(installer.overlay_root()).unwrap();
1023        make_symlink("missing-version", &installer.overlay_root().join("active")).unwrap();
1024
1025        assert!(matches!(
1026            installer.active_overlay_state().unwrap(),
1027            OverlayState::Invalid(reason) if reason.contains("missing")
1028        ));
1029    }
1030}
rusty_stack/installers/components/rccl.rs

rusty_stack/installers/components/
rccl.rs
