[CI] Improve build time dependency resolution

build_tools/build_ext.py

@@ -19,6 +19,8 @@
 
 from .utils import (
     cmake_bin,
+    cuda_home_path,
+    cuda_version,
     debug_build_enabled,
     found_ninja,
     get_frameworks,
@@ -61,6 +63,31 @@ def _build_cmake(self, build_dir: Path, install_dir: Path) -> None:
             f"-DCMAKE_BUILD_TYPE={build_type}",
             f"-DCMAKE_INSTALL_PREFIX={install_dir}",
         ]
+
+        discovered_cuda_home = cuda_home_path()
+        if discovered_cuda_home is not None:
+            configure_command += [
+                f"-DCUDAToolkit_ROOT={discovered_cuda_home}",
+                # CUDA wheels use `lib`, while nvcc's host linker expects `lib64`.
+                f"-DCMAKE_CUDA_FLAGS=-L{discovered_cuda_home}/lib",
+            ]
+
+            if cuda_full_version := cuda_version():
+                cuda_major_version = cuda_full_version[0]
+                cuda_lib_dir = discovered_cuda_home / "lib"
+                configure_command += [
+                    f"-DCUDA_CUDART={cuda_lib_dir / f'libcudart.so.{cuda_major_version}'}",
+                    f"-DCUDA_cudart_LIBRARY={cuda_lib_dir / f'libcudart.so.{cuda_major_version}'}",
+                    f"-DCUDA_cublas_LIBRARY={cuda_lib_dir / f'libcublas.so.{cuda_major_version}'}",
+                    (
+                        f"-DCUDA_cublasLt_LIBRARY={cuda_lib_dir / f'libcublasLt.so.{cuda_major_version}'}"
+                    ),
+                ]
+
+        discovered_nvcc_path = nvcc_path()
+        if discovered_nvcc_path is not None:
+            configure_command.append(f"-DCMAKE_CUDA_COMPILER={discovered_nvcc_path}")
+
         if bool(int(os.getenv("NVTE_USE_CCACHE", "0"))):
             ccache_bin = os.getenv("NVTE_CCACHE_BIN", "ccache")
             configure_command += [
@@ -185,6 +212,11 @@ def _compile_fn(obj, src, ext, cc_args, extra_postargs, pp_opts) -> None:
                             and not framework_extension_only
                         ):
                             nvcc_bin = nvcc_path()
+                            if nvcc_bin is None:
+                                raise RuntimeError(
+                                    f"NVCC not found and is required for building CUDA source {src}"
+                                )
+
                             self.compiler.set_executable("compiler_so", str(nvcc_bin))
                             if isinstance(cflags, dict):
                                 cflags = cflags["nvcc"]

build_tools/jax.py

@@ -16,6 +16,8 @@
     cudnn_frontend_include_path,
     debug_build_enabled,
     setup_mpi_flags,
+    nccl_include_path,
+    nccl_lib_path,
     nccl_ep_enabled,
 )
 from typing import List
@@ -90,6 +92,8 @@ def setup_jax_extension(
 
     # Header files
     include_dirs = get_cuda_include_dirs()
+    if (discovered_nccl_include_path := nccl_include_path()) is not None:
+        include_dirs.append(discovered_nccl_include_path)
     include_dirs.append(cudnn_frontend_include_path())
     include_dirs.extend(
         [
@@ -117,6 +121,12 @@ def setup_jax_extension(
     if nccl_ep_enabled():
         cxx_flags.append("-DNVTE_WITH_NCCL_EP")
 
+    kwargs = {}
+    if (discovered_nccl_lib_path := nccl_lib_path()) is not None:
+        kwargs["extra_objects"] = [str(discovered_nccl_lib_path)]
+    else:
+        kwargs["libraries"] = ["nccl"]
+
     # Define TE/JAX as a Pybind11Extension
     from pybind11.setup_helpers import Pybind11Extension
 
@@ -125,5 +135,5 @@ def setup_jax_extension(
         sources=[str(path) for path in sources],
         include_dirs=[str(path) for path in include_dirs],
         extra_compile_args=cxx_flags,
-        libraries=["nccl"],
+        **kwargs,
     )

build_tools/utils.py

@@ -16,7 +16,7 @@
 from pathlib import Path
 from importlib.metadata import PackageNotFoundError, distribution, version as get_version
 from subprocess import CalledProcessError
-from typing import List, Optional, Tuple, Union
+from typing import Callable, List, Optional, Tuple, Union
 
 
 # Needs to stay consistent with .pre-commit-config.yaml config.
@@ -175,6 +175,90 @@ def found_pybind11() -> bool:
     return False
 
 
+@functools.lru_cache(maxsize=None)
+def cuda_home_path() -> Optional[Path]:
+    """Returns the CUDA home path. This path should contain binaries (e.g. nvcc), headers, and libraries.
+
+    Returns `None` if CUDA is not found."""
+    if cuda_home := os.getenv("CUDA_HOME"):
+        return Path(cuda_home)
+
+    # Check if site-packages contains an `nvidia` directory
+    for site_package in sys.path:
+        if not Path(site_package).is_dir():
+            continue
+
+        nvidia_dir = Path(site_package) / "nvidia"
+        if not nvidia_dir.is_dir():
+            continue
+
+        if Path(nvidia_dir / "bin").is_dir():
+            return nvidia_dir
+
+        # In this case there must be a "CUDA version directory" like `cu12` or `cu13`
+        # that contains the binaries, headers, and libraries
+        for cuda_version_dir in nvidia_dir.iterdir():
+            if not cuda_version_dir.is_dir():
+                continue
+
+            # Verify that the directory name matches the expected `cu##` format
+            if not re.match(r"cu\d+", cuda_version_dir.name):
+                continue
+
+            if not (cuda_version_dir / "bin").is_dir():
+                continue
+
+            return cuda_version_dir
+
+    return None
+
+
+@functools.lru_cache(maxsize=None)
+def nccl_root_path() -> Optional[Path]:
+    """Return the NCCL installation root.
+
+    Returns `None` if NCCL is not found."""
+    if (cuda_home := cuda_home_path()) is not None:
+        nccl_root = cuda_home.parent / "nccl"
+        if nccl_root.is_dir():
+            return nccl_root
+
+    # Check NCCL Python packages
+    for package_name in ["nvidia-nccl-cu13", "nvidia-nccl-cu12"]:
+        try:
+            nccl_distribution = distribution(package_name)
+        except PackageNotFoundError:
+            continue
+
+        nccl_root = Path(nccl_distribution.locate_file("nvidia/nccl"))
+        if nccl_root.is_dir():
+            return nccl_root
+
+    return None
+
+
+@functools.lru_cache(maxsize=None)
+def nccl_include_path() -> Optional[Path]:
+    """Return the NCCL include directory."""
+    if (nccl_root := nccl_root_path()) is not None:
+        include_path = nccl_root / "include"
+        if include_path.is_dir():
+            return include_path
+
+    return None
+
+
+@functools.lru_cache(maxsize=None)
+def nccl_lib_path() -> Optional[Path]:
+    """Return the NCCL shared library path."""
+    if (nccl_root := nccl_root_path()) is not None:
+        lib_path = nccl_root / "lib" / "libnccl.so.2"
+        if lib_path.is_file():
+            return lib_path
+
+    return None
+
+
 @functools.lru_cache(maxsize=None)
 def cuda_toolkit_include_path() -> Tuple[str, str]:
     """Returns root path for cuda toolkit includes.
@@ -198,30 +282,49 @@ def cuda_toolkit_include_path() -> Tuple[str, str]:
 
 
 @functools.lru_cache(maxsize=None)
-def nvcc_path() -> Tuple[str, str]:
-    """Returns the NVCC binary path.
+def nvcc_path() -> Optional[Path]:
+    """Get the NVCC binary path.
 
-    Throws FileNotFoundError if NVCC is not found."""
-    # Try finding NVCC
-    nvcc_bin: Optional[Path] = None
-    if nvcc_bin is None and os.getenv("CUDA_HOME"):
-        # Check in CUDA_HOME
-        cuda_home = Path(os.getenv("CUDA_HOME"))
-        nvcc_bin = cuda_home / "bin" / "nvcc"
-    if nvcc_bin is None:
-        # Check if nvcc is in path
-        nvcc_bin = shutil.which("nvcc")
-        if nvcc_bin is not None:
-            cuda_home = Path(nvcc_bin.rstrip("/bin/nvcc"))
-            nvcc_bin = Path(nvcc_bin)
-    if nvcc_bin is None:
-        # Last-ditch guess in /usr/local/cuda
-        cuda_home = Path("/usr/local/cuda")
-        nvcc_bin = cuda_home / "bin" / "nvcc"
-    if not nvcc_bin.is_file():
-        raise FileNotFoundError(f"Could not find NVCC at {nvcc_bin}")
+    Returns `None` if NVCC is not found.
+    """
+
+    def lookup_via_cuda_home() -> Optional[str]:
+        if (cuda_home := cuda_home_path()) is not None:
+            return cuda_home / "bin" / "nvcc"
+        return None
+
+    def lookup_via_path() -> Optional[str]:
+        if (nvcc_bin := shutil.which("nvcc")) is not None:
+            return nvcc_bin
+        return None
+
+    def lookup_via_distribution() -> Optional[str]:
+        try:
+            return distribution("nvidia-cuda-nvcc").locate_file("bin/nvcc")
+        except PackageNotFoundError:
+            return None
+
+    def lookup_via_local_cuda() -> Optional[str]:
+        return "/usr/local/cuda/bin/nvcc"
+
+    nvcc_lookup_funcs: List[Callable[[], Optional[str]]] = [
+        lookup_via_cuda_home,
+        lookup_via_path,
+        lookup_via_distribution,
+        lookup_via_local_cuda,
+    ]
+
+    for nvcc_lookup_func in nvcc_lookup_funcs:
+        nvcc_bin = nvcc_lookup_func()
 
-    return nvcc_bin
+        if nvcc_bin is None:
+            continue
+
+        nvcc_bin_path = Path(nvcc_bin)
+        if nvcc_bin_path.is_file():
+            return nvcc_bin_path
+
+    return None
 
 
 @functools.lru_cache(maxsize=None)
@@ -324,13 +427,9 @@ def cuda_version() -> Tuple[int, ...]:
     and check pip version.
     """
 
-    try:
-        nvcc_bin = nvcc_path()
-    except FileNotFoundError as e:
-        pass
-    else:
+    if (nvcc_bin := nvcc_path()) is not None:
         output = subprocess.run(
-            [nvcc_bin, "-V"],
+            [str(nvcc_bin), "-V"],
             capture_output=True,
             check=True,
             universal_newlines=True,
@@ -339,12 +438,20 @@ def cuda_version() -> Tuple[int, ...]:
         version = match.group(1).split(".")
         return tuple(int(v) for v in version)
 
-    try:
-        version_str = get_version("nvidia-cuda-runtime-cu12")
-        version_tuple = tuple(int(part) for part in version_str.split(".") if part.isdigit())
-        return version_tuple
-    except importlib.metadata.PackageNotFoundError:
-        raise RuntimeError("Could neither find NVCC executable nor CUDA runtime Python package.")
+    version_str: Optional[str] = None
+    package_names = ["nvidia-cuda-runtime", "nvidia-cuda-runtime-cu13", "nvidia-cuda-runtime-cu12"]
+
+    for package_name in package_names:
+        try:
+            version_str = get_version(package_name)
+        except PackageNotFoundError:
+            pass
+        else:
+            return tuple(int(part) for part in version_str.split(".") if part.isdigit())
+
+    raise RuntimeError(
+        f"Could neither find NVCC executable nor CUDA runtime Python package for {package_names}."
+    )
 
 
 def get_frameworks() -> List[str]: