From d657e6b6cdf54d68ae40dbf5ad9ad43c49f565d2 Mon Sep 17 00:00:00 2001 From: iefgnoix Date: Mon, 20 Nov 2023 11:57:15 -0800 Subject: [PATCH] Configure PATH and LD_LIBRARY_PATH to include cuda (#5761) * configure path and LD_LIBRARY_PATH to include cuda. * also consider 32 bit system. * add the path on gke container --- infra/ansible/config/env.yaml | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/infra/ansible/config/env.yaml b/infra/ansible/config/env.yaml index 83d4ccdd03a..bd4d631ea69 100644 --- a/infra/ansible/config/env.yaml +++ b/infra/ansible/config/env.yaml @@ -16,6 +16,8 @@ release_env: cuda: TF_CUDA_COMPUTE_CAPABILITIES: 7.0,7.5,8.0 XLA_CUDA: 1 + PATH: /usr/local/cuda-{{ cuda_version }}/bin:/usr/local/nvidia/bin${PATH:+:${PATH}} + LD_LIBRARY_PATH: /usr/local/cuda-{{ cuda_version }}/lib64:/usr/local/cuda-{{ cuda_version }}/lib:/usr/local/nvidia/lib64:/usr/local/nvidia/lib${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}:/usr/local/lib # Variables that will be passed to shell environment only for building PyTorch and XLA libs. build_env: @@ -44,6 +46,8 @@ build_env: cuda: TF_CUDA_COMPUTE_CAPABILITIES: 7.0,7.5,8.0 XLA_CUDA: 1 + PATH: /usr/local/cuda-{{ cuda_version }}/bin:/usr/local/nvidia/bin${PATH:+:${PATH}} + LD_LIBRARY_PATH: /usr/local/cuda-{{ cuda_version }}/lib64:/usr/local/cuda-{{ cuda_version }}/lib:/usr/local/nvidia/lib64:/usr/local/nvidia/lib${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}:/usr/local/lib tpu: ACCELERATOR: tpu