diff --git a/tools/Polygraphy/CHANGELOG.md b/tools/Polygraphy/CHANGELOG.md index cea3567ee8..f9c9a95d86 100644 --- a/tools/Polygraphy/CHANGELOG.md +++ b/tools/Polygraphy/CHANGELOG.md @@ -3,6 +3,13 @@ Dates are in YYYY-MM-DD format. +## Unreleased +### Fixed +- Polygraphy now finds the CUDA runtime installed by the CUDA 13+ `nvidia-cuda-runtime` pip wheel, which + places it under `nvidia/cu13` instead of `nvidia/cuda_runtime`. Previously, without a CUDA Toolkit or + `LD_LIBRARY_PATH` entry, TensorRT runners failed with `OSError: libcudart.so: cannot open shared object file`. + + ## v0.53.6 (2026-09-22) ### Added - `Comparator.run()` now accepts a `save_input_blob_path` parameter, exposed on the CLI as diff --git a/tools/Polygraphy/polygraphy/cuda/cuda.py b/tools/Polygraphy/polygraphy/cuda/cuda.py index 9344753246..dbe1bbc32c 100644 --- a/tools/Polygraphy/polygraphy/cuda/cuda.py +++ b/tools/Polygraphy/polygraphy/cuda/cuda.py @@ -16,6 +16,7 @@ # import contextlib import ctypes +import glob import importlib.util import os import platform @@ -32,24 +33,27 @@ def void_ptr(val=None): return ctypes.c_void_p(val) -def _cuda_runtime_wheel_dirs(subdir): - # The pip `nvidia-cuda-runtime-cuXX` wheel installs the runtime under - # `/nvidia/cuda_runtime/{bin (Windows), lib (Linux)}`. - dirs = [ - # Some `sys.path` entries are files (a console-script `.exe`, a zipapp), not dirs; skip those. - os.path.join(p, "nvidia", "cuda_runtime", subdir) - for p in sys.path - if p and os.path.isdir(p) +def _cuda_runtime_wheel_dirs(*subdirs): + # The pip CUDA runtime wheel installs the runtime under `/nvidia/cuda_runtime` + # for CUDA 12 (`nvidia-cuda-runtime-cu12`) and `/nvidia/cu` for CUDA 13+ + # (`nvidia-cuda-runtime`). + # Some `sys.path` entries are files (a console-script `.exe`, a zipapp), not dirs; skip those. + nvidia_dirs = [ + os.path.join(p, "nvidia") for p in sys.path if p and os.path.isdir(p) ] + pkg_dirs = [os.path.join(d, "cuda_runtime") for d in nvidia_dirs] # find_spec also locates the package when it is import-reachable (e.g. via a `.pth`) but its # directory is not itself on `sys.path`. Best-effort: an uninstalled wheel just raises. with contextlib.suppress(Exception): spec = importlib.util.find_spec("nvidia.cuda_runtime") if spec is not None and spec.submodule_search_locations: - dirs += [ - os.path.join(loc, subdir) for loc in spec.submodule_search_locations - ] - return dirs + pkg_dirs += list(spec.submodule_search_locations) + pkg_dirs += [ + pkg_dir + for d in nvidia_dirs + for pkg_dir in sorted(glob.glob(os.path.join(d, "cu[0-9]*"))) + ] + return [os.path.join(d, subdir) for d in pkg_dirs for subdir in subdirs] def _find_cuda_lib_dirs(): @@ -79,7 +83,7 @@ def _find_cuda_lib_dirs(): if root for path in [root, *[os.path.join(root, s) for s in bin_subdirs]] ] - cuda_paths += _cuda_runtime_wheel_dirs("bin") + cuda_paths += _cuda_runtime_wheel_dirs("bin", os.path.join("bin", "x86_64")) cuda_paths += os.environ.get("PATH", "").split(os.path.pathsep) else: cuda_paths = [ @@ -139,13 +143,15 @@ def __init__(self): hint = ( "Ensure the CUDA Toolkit is installed and that either `CUDA_PATH` points to its " "root or the directory containing `cudart64_*.dll` is on `PATH`. " - "Alternatively, `pip install nvidia-cuda-runtime-cuXX` provides the runtime." + "Alternatively, `pip install nvidia-cuda-runtime` (CUDA 13+) or " + "`nvidia-cuda-runtime-cu12` provides the runtime." ) else: hint = ( "Ensure the CUDA Toolkit is installed and that the directory containing " "`libcudart.so*` is on `LD_LIBRARY_PATH`. " - "Alternatively, `pip install nvidia-cuda-runtime-cuXX` provides the runtime." + "Alternatively, `pip install nvidia-cuda-runtime` (CUDA 13+) or " + "`nvidia-cuda-runtime-cu12` provides the runtime." ) log_func( f"Could not find the CUDA runtime library.\n{hint}\n" diff --git a/tools/Polygraphy/tests/cuda/test_cuda.py b/tools/Polygraphy/tests/cuda/test_cuda.py index ad957ff8f4..4866412b38 100644 --- a/tools/Polygraphy/tests/cuda/test_cuda.py +++ b/tools/Polygraphy/tests/cuda/test_cuda.py @@ -127,6 +127,28 @@ class FakeSpec: dirs = _find_cuda_lib_dirs() assert os.path.join(pkg_dir, "bin") in dirs + @pytest.mark.parametrize( + "platform_name,lib_subdir", + [("linux", "lib"), ("win32", os.path.join("bin", "x86_64"))], + ) + def test_cuda13_runtime_wheel_dirs_included( + self, monkeypatch, tmp_path, platform_name, lib_subdir + ): + # CUDA 13+ runtime wheels install under `nvidia/cu` instead of `nvidia/cuda_runtime`; + # other `nvidia/cu*` packages (e.g. cuBLAS) must not be picked up. + (tmp_path / "nvidia" / "cu13").mkdir(parents=True) + (tmp_path / "nvidia" / "cublas").mkdir() + monkeypatch.setattr("sys.platform", platform_name) + monkeypatch.setattr("sys.path", [str(tmp_path)]) + monkeypatch.setattr("platform.machine", lambda: "AMD64") + monkeypatch.setattr(importlib.util, "find_spec", lambda name: None) + monkeypatch.setattr(os, "environ", {}) + + dirs = _find_cuda_lib_dirs() + + assert os.path.join(str(tmp_path), "nvidia", "cu13", lib_subdir) in dirs + assert not any("cublas" in d for d in dirs) + def test_windows_deduplicates_preserving_order(self, monkeypatch): # CUDA_PATH and a versioned variable can point at the same root; the candidate list must # not contain duplicates, and first-seen order is preserved.