rusty_stack/installers/common/
nvidia_blocklist.rs

1//! NVIDIA/CUDA exclusion — the SINGLE source of truth for the hard prime
2//! requirement that NO NVIDIA/CUDA dependency is ever installed.
3//!
4//! # Tenet (v0.3.0 remediation)
5//!
6//! "Nvidia CUDA deps are NEVER installed by ANY component — ESPECIALLY Rusty
7//! Llama. THIS IS A HARD PRIME REQUIREMENT. NOT A SINGLE NVIDIA DEP MAY EVER BE
8//! INSTALLED. Rusty is for AMD GPUs."
9//!
10//! # Why this module exists
11//!
12//! The crate previously had THREE divergent CUDA/nvidia blocklists:
13//! - `installer.rs::filter_cuda_requirements` (requirements-file filter, narrow)
14//! - `megatron.rs::is_safe_package` (single-package check, broadest)
15//! - `textgen.rs::EXCLUDED_PATTERNS` (regex form)
16//!
17//! They disagreed (e.g. the first lacks `cudnn`/`cublas`/`nvjitlink`/`nccl`
18//! that the second catches), and `comfyui` used NONE of them. Everything now
19//! routes through here so the rule is defined once and applied everywhere.
20//!
21//! # Enforcement layers
22//!
23//! This module is the *detection + filtering* layer. It feeds two enforcement
24//! mechanisms used together:
25//! 1. **Requirements filtering** — [`filter_requirements`] strips banned lines
26//!    from `requirements.txt` files before `pip install -r`.
27//! 2. **`--no-deps` + explicit deps** (Stage 3) — for torch-adjacent packages
28//!    whose transitive deps could pull `nvidia-*` wheels; installed with
29//!    `--no-deps` and a curated dependency list. pip *constraints* cannot
30//!    cleanly "ban" a package, so `--no-deps` is the transitive guarantee.
31//! 3. **Post-install contamination check** — [`contaminated_packages`] scans
32//!    `pip list` output so verification (Stage 4) can raise
33//!    `InstallerError::NvidiaContamination`.
34
35/// Package-name prefixes that mark an NVIDIA/CUDA dependency.
36///
37/// Matched case-insensitively against the normalized package name (prefix match).
38/// `triton` is a prefix so it also catches `triton-windows`, `triton[...]`, etc.
39pub const BLOCKED_PREFIXES: &[&str] = &[
40    "nvidia",
41    "cuda",
42    "cudnn",
43    "cublas",
44    "cufft",
45    "curand",
46    "cusolver",
47    "cusparse",
48    "cusparselt",
49    "nccl",
50    "nvtx",
51    "nvjitlink",
52    "nv-pic",
53    "tensorrt",
54    "triton",
55    "pytorch-cuda",
56    "torch-cuda",
57];
58
59/// Exact package names (after normalization) that are NVIDIA/CUDA/ROCm-conflicting.
60///
61/// These must come from the ROCm build managed by Rusty, never from a CUDA
62/// default wheel.
63pub const BLOCKED_EXACT: &[&str] = &[
64    "torch",
65    "torchvision",
66    "torchaudio",
67    "xformers",
68    "flash-attn",
69    "flash_attn",
70];
71
72/// CUDA wheel URL markers (in URLs like `...+cu124...whl`).
73pub const CUDA_URL_MARKERS: &[&str] = &[
74    "+cu1", "+cu2", "+cu3", "+cu4", "cu118", "cu121", "cu124", "cu126", "cu128", "cu129",
75];
76
77/// Normalize a requirement spec to its bare package name.
78///
79/// Strips environment markers (`;`), extras (`[...]`), version specifiers
80/// (`==`, `>=`, etc.), and trailing whitespace. `"nvidia-cudnn-cu12==8.9..."`
81/// → `"nvidia-cudnn-cu12"`; `"triton-windows [test]>=3.5 ; sys_platform==..."`
82/// → `"triton-windows"`.
83pub fn extract_pkg_name(requirement: &str) -> &str {
84    let req = requirement.trim();
85    let after_marker = req.split(';').next().unwrap_or(req);
86    let after_extras = after_marker.split('[').next().unwrap_or(after_marker);
87    after_extras
88        .split(&['=', '<', '>', '~', '!', ' ', '\t'][..])
89        .next()
90        .unwrap_or(after_extras)
91        .trim()
92}
93
94/// Is this package name an NVIDIA/CUDA/ROCm-conflicting package?
95///
96/// Use on a normalized name (see [`extract_pkg_name`]) or a raw requirement —
97/// this normalizes first. Case-insensitive.
98pub fn is_cuda_nvidia_package(requirement: &str) -> bool {
99    let pkg = extract_pkg_name(requirement).to_lowercase();
100    if pkg.is_empty() {
101        return false;
102    }
103    if BLOCKED_EXACT
104        .iter()
105        .any(|exact| exact.eq_ignore_ascii_case(&pkg))
106    {
107        return true;
108    }
109    if BLOCKED_PREFIXES
110        .iter()
111        .any(|prefix| pkg.starts_with(prefix))
112    {
113        return true;
114    }
115    // cupy-cuda*, any name containing "cuda" (e.g. cupy-cuda12, torch-cuda).
116    if pkg.contains("cuda") {
117        return true;
118    }
119    false
120}
121
122/// Does this wheel URL carry a CUDA build marker? (e.g. `...+cu124...whl`)
123///
124/// Also rejects direct wheel URLs whose **filename** embeds a blocked NVIDIA/
125/// CUDA runtime package name (e.g. `.../nvidia_cublas_cu12-...whl` with no
126/// `+cu` marker) so they cannot bypass the blocklist via a bare URL.
127pub fn is_cuda_wheel_url(url: &str) -> bool {
128    let lower = url.to_lowercase();
129    if CUDA_URL_MARKERS.iter().any(|marker| lower.contains(marker))
130        || (lower.contains("+cu") && lower.contains(".whl"))
131    {
132        return true;
133    }
134    // Direct-wheel bypass guard: derive the package name from the URL filename
135    // and reject if it is a blocked NVIDIA/CUDA runtime package. A wheel
136    // filename is `name-version-python-abi-tag.whl`; the package name is the
137    // dash-separated segment(s) before the version, with underscores normalized
138    // to dashes (`nvidia_cublas_cu12` → `nvidia-cublas-cu12`).
139    if lower.contains("http://") || lower.contains("https://") || lower.contains(".whl") {
140        if let Some(pkg) = extract_pkg_name_from_url_filename(&lower) {
141            if is_nvidia_cuda_runtime_package(&pkg) {
142                return true;
143            }
144        }
145    }
146    false
147}
148
149/// Derive the package name from the filename portion of a wheel URL.
150///
151/// `https://.../nvidia_cublas_cu12-12.1.3.1-cp310-...whl` →
152/// `nvidia-cublas-cu12`. Returns `None` when no version-like segment can be
153/// found (so a non-wheel or unparseable URL is not falsely matched).
154fn extract_pkg_name_from_url_filename(url_lower: &str) -> Option<String> {
155    // Take the last path segment (the filename).
156    let filename = url_lower.rsplit(['/', '\\']).next()?;
157    let filename = filename.trim();
158    if filename.is_empty() {
159        return None;
160    }
161    // Strip a trailing .whl if present.
162    let stem = filename.strip_suffix(".whl").unwrap_or(filename);
163    // The package name is everything before the first segment that looks like a
164    // version (starts with a digit). Wheel filenames are
165    // `name-version-python-abi-platform`, so splitting on '-' and taking the
166    // leading non-version segments reconstructs the normalized package name.
167    let mut name_parts: Vec<&str> = Vec::new();
168    for seg in stem.split('-') {
169        if seg.is_empty() {
170            continue;
171        }
172        // A version segment starts with a digit (e.g. "12.1.3.1", "2.4.0").
173        if seg.as_bytes().first().is_some_and(|b| b.is_ascii_digit()) {
174            break;
175        }
176        name_parts.push(seg);
177    }
178    if name_parts.is_empty() {
179        return None;
180    }
181    Some(name_parts.join("-"))
182}
183
184/// NVIDIA/CUDA *runtime-library* prefixes (the `nvidia-*`/`cuda*` wheels), for
185/// the **execution chokepoint**. Deliberately EXCLUDES the torch/triton family
186/// — those are legitimately installed from the ROCm index by the core
187/// pytorch/triton installers, and are instead kept out of third-party installs
188/// by [`is_cuda_nvidia_package`] (requirements filtering) + `--no-deps`.
189pub const NVIDIA_RUNTIME_PREFIXES: &[&str] = &[
190    "nvidia",
191    "cuda",
192    "cudnn",
193    "cublas",
194    "cufft",
195    "curand",
196    "cusolver",
197    "cusparse",
198    "cusparselt",
199    "nccl",
200    "nvtx",
201    "nvjitlink",
202    "nv-pic",
203    "tensorrt",
204    "pytorch-cuda",
205    "torch-cuda",
206];
207
208/// Is this an NVIDIA/CUDA *runtime* package (a `nvidia-*`/`cuda*` lib)?
209///
210/// Narrower than [`is_cuda_nvidia_package`]: returns `false` for
211/// torch/torchvision/torchaudio/triton/xformers/flash-attn (core, ROCm-managed).
212/// Used by the pip execution chokepoint to block any nvidia/cuda runtime wheel
213/// from *any* installer, while still permitting the core pytorch/triton install.
214pub fn is_nvidia_cuda_runtime_package(requirement: &str) -> bool {
215    let pkg = extract_pkg_name(requirement).to_lowercase();
216    if pkg.is_empty() {
217        return false;
218    }
219    if NVIDIA_RUNTIME_PREFIXES
220        .iter()
221        .any(|prefix| pkg.starts_with(prefix))
222    {
223        return true;
224    }
225    // cupy-cuda*, any name containing "cuda".
226    if pkg.contains("cuda") {
227        return true;
228    }
229    false
230}
231
232/// Is this a CUDA-only source line (e.g. `ik_llama` builds)?
233fn is_cuda_only_source(line_lower: &str) -> bool {
234    // ik_llama ships CUDA-only wheels.
235    if line_lower.contains("ik_llama") {
236        return true;
237    }
238    // exllamav3 / flash_attn wheels with CUDA markers.
239    if (line_lower.contains("exllamav3") || line_lower.contains("flash_attn"))
240        && (line_lower.contains("+cu") || line_lower.contains("cu1"))
241    {
242        return true;
243    }
244    false
245}
246
247/// Filter a requirements FILE, returning the content with all
248/// NVIDIA/CUDA/ROCm-conflicting lines removed.
249///
250/// Propagates read failures as `Err` so callers can warn / treat the component
251/// as fatal rather than silently installing an empty dependency set (which a
252/// missing/unreadable requirements file previously masqueraded as "all
253/// filtered"). Keeps ordering; drops comments/blank lines. This is the unified
254/// replacement for the former `installer.rs::filter_cuda_requirements`.
255pub fn filter_requirements_file(path: &str) -> std::io::Result<String> {
256    std::fs::read_to_string(path).map(|c| filter_requirements(&c))
257}
258
259/// Filter requirements TEXT, returning lines that are safe to install on a
260/// ROCm/AMD stack (all NVIDIA/CUDA/ROCm-conflicting lines removed).
261pub fn filter_requirements(content: &str) -> String {
262    let mut kept = Vec::new();
263    for line in content.lines() {
264        let trimmed = line.trim();
265        if trimmed.is_empty() || trimmed.starts_with('#') {
266            continue;
267        }
268        if is_cuda_nvidia_package(trimmed) {
269            continue;
270        }
271        let line_lower = trimmed.to_lowercase();
272        if line_lower.contains("http://") || line_lower.contains("https://") {
273            if is_cuda_wheel_url(trimmed) || is_cuda_only_source(&line_lower) {
274                continue;
275            }
276        } else if is_cuda_only_source(&line_lower) {
277            continue;
278        }
279        kept.push(line.to_string());
280    }
281    kept.join("\n")
282}
283
284/// Scan `pip list` output for any installed NVIDIA/CUDA **runtime** packages.
285///
286/// Uses the NARROW runtime check ([`is_nvidia_cuda_runtime_package`]) plus the
287/// CUDA-wheel-URL check so valid ROCm-managed `torch`/`torchvision`/`triton`
288/// are NOT flagged as contamination (they are legitimately installed from the
289/// ROCm index). Only `nvidia-*`/`cuda*` runtime libs and CUDA wheels are
290/// reported. Returns the offending package display lines (for the
291/// `InstallerError::NvidiaContamination` payload and verification reports).
292pub fn contaminated_packages(pip_list_output: &str) -> Vec<String> {
293    let mut found = Vec::new();
294    for line in pip_list_output.lines() {
295        let name = line.split_whitespace().next().unwrap_or("");
296        if is_nvidia_cuda_runtime_package(name) || is_cuda_wheel_url(line) {
297            found.push(line.trim().to_string());
298        }
299    }
300    found
301}
302
303// ===========================================================================
304// Tests
305// ===========================================================================
306
307#[cfg(test)]
308mod tests {
309    use super::*;
310
311    #[test]
312    fn nvidia_wheels_are_blocked() {
313        assert!(is_cuda_nvidia_package("nvidia-cudnn-cu12"));
314        assert!(is_cuda_nvidia_package("nvidia-cublas-cu12==12.1.3.1"));
315        assert!(is_cuda_nvidia_package("nccl"));
316        assert!(is_cuda_nvidia_package("nvjitlink"));
317        assert!(is_cuda_nvidia_package("cuda-python"));
318        assert!(is_cuda_nvidia_package("cupy-cuda12x"));
319        assert!(is_cuda_nvidia_package("tensorrt"));
320    }
321
322    #[test]
323    fn torch_family_is_blocked() {
324        assert!(is_cuda_nvidia_package("torch"));
325        assert!(is_cuda_nvidia_package("torchvision>=0.17"));
326        assert!(is_cuda_nvidia_package("torchaudio"));
327        assert!(is_cuda_nvidia_package("triton"));
328        assert!(is_cuda_nvidia_package("triton-windows==3.5.1"));
329        assert!(is_cuda_nvidia_package("xformers"));
330        assert!(is_cuda_nvidia_package("flash-attn"));
331    }
332
333    #[test]
334    fn safe_amd_packages_pass() {
335        assert!(!is_cuda_nvidia_package("deepspeed"));
336        assert!(!is_cuda_nvidia_package("transformers"));
337        assert!(!is_cuda_nvidia_package("accelerate"));
338        assert!(!is_cuda_nvidia_package("einops"));
339        assert!(!is_cuda_nvidia_package("numpy"));
340        assert!(!is_cuda_nvidia_package("bitsandbytes"));
341        assert!(!is_cuda_nvidia_package("vllm"));
342        assert!(!is_cuda_nvidia_package("wandb"));
343    }
344
345    #[test]
346    fn cuda_wheel_urls_detected() {
347        assert!(is_cuda_wheel_url(
348            "https://x/llama_cpp_binaries-0.124.0+cu124-py3.whl"
349        ));
350        assert!(is_cuda_wheel_url("https://x/torch-2.4.0+cu118.whl"));
351        assert!(is_cuda_wheel_url("https://x/flash_attn-2.6+cu128.whl"));
352        assert!(!is_cuda_wheel_url("https://x/torch-2.4.0+rocm624.whl"));
353        assert!(!is_cuda_wheel_url("https://x/numpy-1.26.whl"));
354    }
355
356    #[test]
357    fn direct_nvidia_wheel_url_blocked_without_cu_marker() {
358        // B1: a direct wheel URL embedding a blocked nvidia/cuda runtime package
359        // name with NO +cu marker must still be rejected.
360        assert!(is_cuda_wheel_url(
361            "https://download.pytorch.org/whl/cu121/nvidia_cublas_cu12-12.1.3.1-cp310-cp310-linux_x86_64.whl"
362        ));
363        assert!(is_cuda_wheel_url(
364            "https://example.com/nvidia-cudnn-cu12-8.9.2-cp310-none-manylinux1_x86_64.whl"
365        ));
366        assert!(is_cuda_wheel_url(
367            "https://example.com/cuda_runtime-12.1.0-py3-none-any.whl"
368        ));
369        // ROCm torch wheel (underscore form) must NOT be blocked by the filename
370        // check (no nvidia/cuda runtime name).
371        assert!(!is_cuda_wheel_url(
372            "https://download.pytorch.org/whl/rocm6.2/torch-2.4.0-cp310-cp310-linux_x86_64.whl"
373        ));
374        assert!(!is_cuda_wheel_url(
375            "https://example.com/numpy-1.26.4-cp310-none-any.whl"
376        ));
377    }
378
379    #[test]
380    fn filter_requirements_strips_banned_lines() {
381        let reqs = "\
382# comment
383torch==2.4.0
384nvidia-cudnn-cu12
385numpy>=1.26
386transformers
387https://download.pytorch.org/whl/cu121/xformers-0.0.27%2Bcu121.whl
388ik_llama_cpp_cuda-0.1.whl
389deepspeed
390";
391        let filtered = filter_requirements(reqs);
392        assert!(!filtered.contains("torch"));
393        assert!(!filtered.contains("nvidia-cudnn"));
394        assert!(!filtered.contains("xformers"));
395        assert!(!filtered.contains("ik_llama"));
396        assert!(filtered.contains("numpy>=1.26"));
397        assert!(filtered.contains("transformers"));
398        assert!(filtered.contains("deepspeed"));
399        assert!(!filtered.contains("# comment"));
400    }
401
402    #[test]
403    fn extract_pkg_name_handles_specs() {
404        assert_eq!(
405            extract_pkg_name("nvidia-cudnn-cu12==8.9.2.26"),
406            "nvidia-cudnn-cu12"
407        );
408        assert_eq!(
409            extract_pkg_name("triton-windows [test]>=3.5.1 ; sys_platform == 'win32'"),
410            "triton-windows"
411        );
412        assert_eq!(extract_pkg_name("  torch  "), "torch");
413    }
414
415    #[test]
416    fn contamination_scan_finds_nvidia_runtime_only() {
417        // B3: torch/torchvision ROCm packages are NOT contamination — only the
418        // nvidia-* runtime wheels are flagged.
419        let pip_list = "\
420Package        Version
421-------------  -------
422deepspeed      0.14.5
423nvidia-cublas-cu12  12.1.3.1
424torch          2.4.0
425torchvision    0.19.0
426triton         3.1.0
427numpy          1.26.4
428";
429        let found = contaminated_packages(pip_list);
430        assert_eq!(found.len(), 1, "only nvidia-cublas flagged: {found:?}");
431        assert!(found.iter().any(|l| l.starts_with("nvidia-cublas")));
432        // ROCm-managed torch family must NOT be flagged.
433        assert!(!found.iter().any(|l| l.contains("torch")));
434        assert!(!found.iter().any(|l| l.contains("triton")));
435        assert!(!found.iter().any(|l| l.contains("deepspeed")));
436    }
437}
rusty_stack/installers/common/nvidia_blocklist.rs

rusty_stack/installers/common/
nvidia_blocklist.rs
