op_builder: make JIT compile robust to non-curated conda envs

[Andrewxu313]

Summary

The default torch.utils.cpp_extension.load() invocation in OpBuilder.jit_load() fails out-of-the-box on any conda env that wasn't pre-blessed for BatchGen. We hit six distinct failure modes installing v1.0.9.post5 on a Gemini-cluster H20 node (Tencent IDC):

#	Error	Root cause
1	libc10.so: cannot open shared object file	conda activate doesn't add <torch>/lib to LD_LIBRARY_PATH
2	__nvJitLinkCreate_12_8 undefined symbol	torch's cusparse.so links against newer nvJitLink than the system one
3	cuda.h: No such file or directory	conda compiler's sysroot has no CUDA headers
4	numa.h: No such file or directory	conda compiler's sysroot has no libnuma headers (numa.h is in $CONDA_PREFIX/include, just not on path)
5	unsupported GNU version! gcc later than 12 are not supported	conda env defaults to gcc-14, CUDA 12.x nvcc caps at gcc-12
6	cannot find -lnuma	hardcoded -L<env>/lib/stubs doesn't include targets/x86_64-linux/lib where the libnuma symlink lives in conda CUDA bundles

This patch makes jit_load() resilient to all six without changing behaviour on the curated H20 envs (every change is a prepend, idempotent, and skipped if the dir doesn't exist — so it's a no-op when the right paths are already wired up).

What changed

op_builder/builder.py — new OpBuilder._augment_jit_env(), called from jit_load() just before torch.utils.cpp_extension.load:

• Prepends <torch>/lib + every <site-packages>/nvidia/*/lib to LD_LIBRARY_PATH (fixes 1, 2)
• Prepends $CUDA_HOME/include and $CONDA_PREFIX/include to CPATH (fixes 3, 4)
• Prepends $CONDA_PREFIX/{targets/x86_64-linux/lib,lib} to LIBRARY_PATH (helps 6)
• If host gcc major > nvcc's supported max, searches peer conda envs for a compatible compiler and uses it via -ccbin + -allow-unsupported-compiler in NVCC_PREPEND_FLAGS (fixes 5)

op_builder/core_engine.py — extra_ldflags() now adds $CONDA_PREFIX/targets/x86_64-linux/lib and $CONDA_PREFIX/lib in addition to lib/stubs (fixes 6 — the linker invocation itself).

Validated on Gemini node-1 (8× H20, conda env batchgen with gcc-14, peer ray env with gcc-11)

Before this patch: from batchgen import launch_http_server fails through six rounds of env patching, each surfacing a new failure mode.

After this patch: same import succeeds in ~84 s of JIT compile, core_engine.so is cached, subsequent invocations are instant.

Test plan

• Confirm H20-Node1 (curated env) launch is unaffected — re-run bash scripts/run_test.sh glm5 L1 and verify same outcome
• Confirm Gemini node-1 launch now works without an env-wrapper script — python -m batchgen.launch_http_server ... should succeed after just pip install of the v1.0.9.post5 wheels
• Confirm on a CPU-only / non-Hopper conda env that _augment_jit_env is still safe (no AttributeError on missing CUDA paths)
• Run tests/test_deepseek_v4_flash_reference.py after JIT compile to confirm we haven't regressed parity

Notes

• All path additions check os.path.isdir before adding, so this is safe to run on machines with partial / non-conda envs.
• The _find_peer_compatible_cc lookup walks $CONDA_PREFIX/../*/bin/x86_64-conda-linux-gnu-cc — this is the standard conda envs layout, so should hit on most setups. If no peer is found, we still set -allow-unsupported-compiler as a fallback.
• This is draft because (a) we still want to confirm the H20 path is unaffected and (b) the peer-env compiler search is a heuristic that could be tightened.

🤖 Generated with Claude Code