does nvrtc work as deps?

EnzymeAD · Oct 15, 2024 · 8475fc3 · 8475fc3
1 parent 221c7d8
commit 8475fc3
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/.buildkite/gpu_pipeline.yml b/.buildkite/gpu_pipeline.yml
@@ -51,6 +51,7 @@ steps:
       export LD_LIBRARY_PATH="`pwd`/bazel-bin/test/test.runfiles/pypi_nvidia_cublas_cu12/site-packages/nvidia/cublas/lib:\$LD_LIBRARY_PATH"
       export LD_LIBRARY_PATH="`pwd`/bazel-bin/test/llama.runfiles/pypi_nvidia_cuda_cupti_cu12/site-packages/nvidia/cuda_cupti/lib:\$LD_LIBRARY_PATH"
       export LD_LIBRARY_PATH="`pwd`/bazel-bin/test/llama.runfiles/pypi_nvidia_cuda_runtime_cu12/site-packages/nvidia/cuda_runtime/lib:\$LD_LIBRARY_PATH"
+      export LD_LIBRARY_PATH="`pwd`/bazel-bin/test/llama.runfiles/pypi_nvidia_cuda_nvrtc_cu12/site-packages/nvidia/cuda_nvrtc/lib:\$LD_LIBRARY_PATH"
       export PATH="`pwd`/bazel-bin/test/llama.runfiles/pypi_nvidia_cuda_nvcc_cu12/site-packages/nvidia/cuda_nvcc/bin:\$PATH"
       export TF_CPP_MIN_LOG_LEVEL=0
 

diff --git a/test/BUILD b/test/BUILD
@@ -63,7 +63,7 @@ TEST_DEPS = [
         "@pypi_absl_py//:pkg",
     ] + select({
         ":use_tpu": ["@pypi_libtpu_nightly//:pkg", "@pypi_requests//:pkg"],
-        "@bazel_tools//src/conditions:linux_x86_64": ["@pypi_jax_cuda12_plugin//:pkg"],
+        "@bazel_tools//src/conditions:linux_x86_64": ["@pypi_jax_cuda12_plugin//:pkg", "@pypi_nvidia_cuda_nvrtc_cu12//:pkg"],
         "//conditions:default": []
     })