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

@@ -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