Accept wheel metadata version ranges in strict guard rails.

Infer the CUDA Toolkit line from both wildcard-pinned and range-based cuda-toolkit requirements so strict process-wide guard rails keep working for editable wheel installs used by nvrtc and nvJitLink.

Made-with: Cursor

Author: rwgk

cuda_pathfinder/cuda/pathfinder/_compatibility_guard_rails.py

@@ -53,9 +53,8 @@
 ConstraintArg: TypeAlias = int | str | tuple[str, int] | None
 
 _CTK_VERSION_RE = re.compile(r"^(?P<major>\d+)\.(?P<minor>\d+)")
-_REQUIRES_DIST_RE = re.compile(
-    r"^\s*(?P<name>[A-Za-z0-9_.-]+)\s*==\s*(?P<version>[0-9][A-Za-z0-9.+-]*?)(?:\.\*)?(?:\s*;|$)"
-)
+_REQUIRES_DIST_RE = re.compile(r"^\s*(?P<name>[A-Za-z0-9_.-]+)\s*(?P<specifier_text>[^;]*)(?:\s*;|$)")
+_VERSION_SPECIFIER_RE = re.compile(r"^\s*(?P<operator>==|<=|>=|<|>)\s*(?P<version>[0-9][A-Za-z0-9.+-]*?(?:\.\*)?)\s*$")
 
 _STATIC_LIBS_PACKAGED_WITH: dict[str, PackagedWith] = {
     "cudadevrt": "ctk",
@@ -113,6 +112,12 @@ def __str__(self) -> str:
         return f"{self.operator}{self.value}"
 
 
+@dataclass(frozen=True, slots=True)
+class VersionSpecifier:
+    operator: ConstraintOperator
+    version: str
+
+
 @dataclass(frozen=True, slots=True)
 class ResolvedItem:
     name: str
@@ -185,6 +190,63 @@ def _distribution_name(dist: importlib.metadata.Distribution) -> str | None:
     return metadata.get("Name")
 
 
+def _release_version_parts(version: str) -> tuple[int, ...] | None:
+    match = re.match(r"^\d+(?:\.\d+)*", version)
+    if match is None:
+        return None
+    return tuple(int(part) for part in match.group(0).split("."))
+
+
+def _compare_release_versions(lhs: tuple[int, ...], rhs: tuple[int, ...]) -> int:
+    max_len = max(len(lhs), len(rhs))
+    lhs_padded = lhs + (0,) * (max_len - len(lhs))
+    rhs_padded = rhs + (0,) * (max_len - len(rhs))
+    if lhs_padded < rhs_padded:
+        return -1
+    if lhs_padded > rhs_padded:
+        return 1
+    return 0
+
+
+def _parse_version_specifiers(specifier_text: str) -> tuple[VersionSpecifier, ...]:
+    stripped = specifier_text.strip()
+    if not stripped:
+        return ()
+    parsed: list[VersionSpecifier] = []
+    for raw_clause in stripped.split(","):
+        match = _VERSION_SPECIFIER_RE.match(raw_clause)
+        if match is None:
+            return ()
+        parsed.append(VersionSpecifier(operator=match.group("operator"), version=match.group("version")))
+    return tuple(parsed)
+
+
+def _version_satisfies_specifiers(version: str, specifiers: tuple[VersionSpecifier, ...]) -> bool:
+    if not specifiers:
+        return False
+    for specifier in specifiers:
+        if specifier.operator == "==":
+            prefix = specifier.version.removesuffix(".*")
+            if version == prefix or version.startswith(prefix + "."):
+                continue
+            return False
+        candidate_parts = _release_version_parts(version)
+        required_parts = _release_version_parts(specifier.version)
+        if candidate_parts is None or required_parts is None:
+            return False
+        comparison = _compare_release_versions(candidate_parts, required_parts)
+        if specifier.operator == "<" and comparison < 0:
+            continue
+        if specifier.operator == "<=" and comparison <= 0:
+            continue
+        if specifier.operator == ">" and comparison > 0:
+            continue
+        if specifier.operator == ">=" and comparison >= 0:
+            continue
+        return False
+    return True
+
+
 @functools.cache
 def _owned_distribution_candidates(abs_path: str) -> tuple[tuple[str, str], ...]:
     normalized_abs_path = os.path.normpath(os.path.abspath(abs_path))
@@ -201,27 +263,42 @@ def _owned_distribution_candidates(abs_path: str) -> tuple[tuple[str, str], ...]
 
 
 @functools.cache
-def _cuda_toolkit_requirement_maps() -> tuple[tuple[str, CtkVersion, dict[str, tuple[str, ...]]], ...]:
-    results: list[tuple[str, CtkVersion, dict[str, tuple[str, ...]]]] = []
+def _cuda_toolkit_requirement_maps() -> tuple[
+    tuple[str, CtkVersion, dict[str, tuple[tuple[VersionSpecifier, ...], ...]]], ...
+]:
+    results: list[tuple[str, CtkVersion, dict[str, tuple[tuple[VersionSpecifier, ...], ...]]]] = []
     for dist in importlib.metadata.distributions():
         dist_name = _distribution_name(dist)
         if _normalize_distribution_name(dist_name or "") != "cuda-toolkit":
             continue
         ctk_version = _parse_ctk_version(dist.version)
         if ctk_version is None:
             continue
-        requirement_map: dict[str, set[str]] = {}
+        requirement_map: dict[str, set[tuple[VersionSpecifier, ...]]] = {}
         for requirement in dist.requires or ():
             match = _REQUIRES_DIST_RE.match(requirement)
             if match is None:
                 continue
             req_name = _normalize_distribution_name(match.group("name"))
-            requirement_map.setdefault(req_name, set()).add(match.group("version"))
+            parsed_specifiers = _parse_version_specifiers(match.group("specifier_text"))
+            if not parsed_specifiers:
+                continue
+            requirement_map.setdefault(req_name, set()).add(parsed_specifiers)
         results.append(
             (
                 dist.version,
                 ctk_version,
-                {name: tuple(sorted(prefixes)) for name, prefixes in requirement_map.items()},
+                {
+                    name: tuple(
+                        sorted(
+                            specifier_sets,
+                            key=lambda specifiers: tuple(
+                                (specifier.operator, specifier.version) for specifier in specifiers
+                            ),
+                        )
+                    )
+                    for name, specifier_sets in requirement_map.items()
+                },
             )
         )
     return tuple(results)
@@ -232,9 +309,9 @@ def _wheel_metadata_for_abs_path(abs_path: str) -> CtkMetadata | None:
     for owner_name, owner_version in _owned_distribution_candidates(abs_path):
         normalized_owner_name = _normalize_distribution_name(owner_name)
         for toolkit_dist_version, ctk_version, requirement_map in _cuda_toolkit_requirement_maps():
-            requirement_prefixes = requirement_map.get(normalized_owner_name, ())
+            requirement_specifier_sets = requirement_map.get(normalized_owner_name, ())
             if not any(
-                owner_version == prefix or owner_version.startswith(prefix + ".") for prefix in requirement_prefixes
+                _version_satisfies_specifiers(owner_version, specifiers) for specifiers in requirement_specifier_sets
             ):
                 continue
             matched_versions[ctk_version] = (

cuda_pathfinder/tests/test_compatibility_guard_rails.py

@@ -83,6 +83,26 @@ def _driver_cuda_version(encoded: int) -> DriverCudaVersion:
     )
 
 
+class _FakeDistribution:
+    def __init__(
+        self,
+        *,
+        name: str,
+        version: str,
+        root: Path,
+        files: tuple[str, ...] = (),
+        requires: tuple[str, ...] = (),
+    ) -> None:
+        self.metadata = {"Name": name}
+        self.version = version
+        self.files = tuple(Path(file) for file in files)
+        self.requires = list(requires)
+        self._root = root
+
+    def locate_file(self, file: Path) -> Path:
+        return self._root / file
+
+
 def _assert_real_ctk_backed_path(path: str) -> None:
     norm_path = os.path.normpath(os.path.abspath(path))
     if "site-packages" in Path(norm_path).parts:
@@ -441,6 +461,49 @@ def fake_load_nvidia_dynamic_lib(libname: str) -> LoadedDL:
         guard_rails.find_nvidia_header_directory("nvrtc")
 
 
+@pytest.mark.parametrize(
+    "requirement",
+    (
+        "nvidia-nvjitlink == 13.2.78.*; extra == 'nvjitlink'",
+        "nvidia-nvjitlink<14,>=13.2.78; extra == 'nvjitlink'",
+    ),
+)
+def test_wheel_metadata_accepts_exact_and_range_requirements(monkeypatch, tmp_path, requirement):
+    site_packages = tmp_path / "site-packages"
+    lib_path = _touch(site_packages / "nvidia" / "cu13" / "lib" / "libnvJitLink.so.13")
+    owner_dist = _FakeDistribution(
+        name="nvidia-nvjitlink",
+        version="13.2.78",
+        root=site_packages,
+        files=("nvidia/cu13/lib/libnvJitLink.so.13",),
+    )
+    cuda_toolkit_dist = _FakeDistribution(
+        name="cuda-toolkit",
+        version="13.2.1",
+        root=site_packages,
+        requires=(requirement,),
+    )
+
+    compatibility_module._owned_distribution_candidates.cache_clear()
+    compatibility_module._cuda_toolkit_requirement_maps.cache_clear()
+    try:
+        monkeypatch.setattr(
+            compatibility_module.importlib.metadata,
+            "distributions",
+            lambda: (owner_dist, cuda_toolkit_dist),
+        )
+
+        metadata = compatibility_module._wheel_metadata_for_abs_path(lib_path)
+    finally:
+        compatibility_module._owned_distribution_candidates.cache_clear()
+        compatibility_module._cuda_toolkit_requirement_maps.cache_clear()
+
+    assert metadata is not None
+    assert metadata.ctk_version.major == 13
+    assert metadata.ctk_version.minor == 2
+    assert metadata.source == "wheel metadata via nvidia-nvjitlink==13.2.78 pinned by cuda-toolkit==13.2.1"
+
+
 def test_constraints_accept_string_and_tuple_forms(monkeypatch, tmp_path):
     ctk_root = tmp_path / "cuda-12.9"
     _write_version_json(ctk_root, "12.9.20250531")