[CI] Improve build time dependency resolution

[fheinecke]

Description

There are a few issues with how build-time dependencies are resolved for native code build paths today:

• Installed python packages are not checked when looking up some dependencies
• The CUDA toolkit root is not auto-discovered when it's built via python deps
• Paths to libraries and headers installed via python packages are not plumbed to cmake

The goal here is to allow all build-time requirements to be installed via wheel. This fixes many of the issues blocking TE wheel build isolation.

This change aims to be 100% backwards compatible.

Fixes # (issue)

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

Please list the changes introduced in this PR:

• Shared libraries and headers installed via wheel are now provided to cmake
• CUDA toolkit root built via python wheels is now provided to cmake
• NVCC installed via wheel is now used when compiling the JAX extension

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[fheinecke]

/te-ci

[fheinecke]

/te-ci

[greptile-apps]

Greptile Summary

This PR expands native build dependency discovery so CUDA and NCCL components installed through Python packages can be passed into the CMake and JAX build paths.

• Discovers NVCC and the CUDA toolkit through environment, Python-package, PATH, and standard-install locations.
• Passes wheel-provided CUDA headers and libraries to CMake.
• Uses discovered NCCL headers and libraries when building the JAX extension.
• Propagates CUDA compiler paths into the NCCL EP submodule build.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains within the scope of the previous review thread.

Important Files Changed

Filename	Overview
build_tools/utils.py	Centralizes CUDA, NVCC, and NCCL discovery across explicit, Python-package, PATH, and standard installation locations.
build_tools/build_ext.py	Supplies discovered CUDA compiler, toolkit root, and wheel-layout libraries to CMake and reports missing NVCC explicitly.
build_tools/jax.py	Adds discovered NCCL headers and direct library input to the JAX native extension build while preserving system-library fallback.
setup.py	Propagates discovered CUDA and NVCC paths into the NCCL EP submodule environment without replacing explicit user settings.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A[Python package build] --> B[Discover NVCC and CUDA root]
  B --> C[Configure common CMake extension]
  B --> D[Discover NCCL headers and library]
  D --> E[Build JAX extension]
  B --> F[Configure NCCL EP submodule]
  C --> G[Package native artifacts]
  E --> G
  F --> G

Loading

_{Reviews (2): Last reviewed commit: "Standardize CUDA home path and NVCC look..." | Re-trigger Greptile}

[timmoon10]

We have somewhat inconsistent search logic:

• cuda_home_path: prioritize CUDA_HOME, otherwise look in PYTHONPATH for nvidia/*-cu13/bin (assumed to hold NVCC).
• cuda_toolkit_include_path: prioritize CUDA_HOME, then look in PATH for NVCC, then in /usr/local/cuda/include.
• nvcc_path: look in cuda_home_path(), then PATH, then Python libs, then /usr/local/cuda

It seems cuda_home_path and cuda_toolkit_include_path are primarily looking for NVCC, so could we consolidate logic?

• nvcc_path: prioritize CUDA_HOME, then look in PYTHONPATH, then Python libs, then PATH, then /usr/local/cuda
• cuda_home_path: look in nvcc_path()
• cuda_toolkit_include_path: look in cuda_home_path()

[timmoon10]

Overall LGTM. I have a suggestion to consolidate logic, but it's not required if this PR is blocking more important improvements.

[fheinecke]

/te-ci

[ksivaman]

LGTM