Remove CUDA_VERSION env variable and deprecated apt-key

Dockerfile.tmpl

@@ -78,17 +78,9 @@ RUN apt-get install -y ocl-icd-libopencl1 clinfo && \
     uv pip install --system /tmp/lightgbm/*.whl && \
     rm -rf /tmp/lightgbm && \
     /tmp/clean-layer.sh
-
-# Remove CUDA_VERSION from non-GPU image.
-{{ else }}
-ENV CUDA_VERSION=""
 {{ end }}
 
 
-# Update GPG key per documentation at https://cloud.google.com/compute/docs/troubleshooting/known-issues
-RUN curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | sudo apt-key add -
-RUN curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | sudo apt-key --keyring /usr/share/keyrings/cloud.google.gpg add -
-
 # Use a fixed apt-get repo to stop intermittent failures due to flaky httpredir connections,
 # as described by Lionel Chan at http://stackoverflow.com/a/37426929/5881346
 RUN sed -i "s/httpredir.debian.org/debian.uchicago.edu/" /etc/apt/sources.list && \

tests/common.py

@@ -11,7 +11,10 @@ def getAcceleratorName():
     except FileNotFoundError:
         return("nvidia-smi not found.")
 
-gpu_test = unittest.skipIf(len(os.environ.get('CUDA_VERSION', '')) == 0, 'Not running GPU tests')
+def isGPU():
+    return os.path.isfile('/proc/driver/nvidia/version')
+
+gpu_test = unittest.skipIf(not isGPU(), 'Not running GPU tests')
 # b/342143152 P100s are slowly being unsupported in new release of popular ml tools such as RAPIDS. 
 p100_exempt = unittest.skipIf(getAcceleratorName() == "Tesla P100-PCIE-16GB", 'Not running p100 exempt tests')
 tpu_test = unittest.skipIf(len(os.environ.get('ISTPUVM', '')) == 0, 'Not running TPU tests')

tests/test_jax.py

@@ -6,7 +6,7 @@
 import jax
 import jax.numpy as np
 
-from common import gpu_test
+from common import gpu_test, isGPU
 from jax import grad, jit
 
 
@@ -21,5 +21,5 @@ def test_grad(self):
         self.assertEqual(0.4199743, ag)
 
     def test_backend(self):
-        expected_backend = 'cpu' if len(os.environ.get('CUDA_VERSION', '')) == 0 else 'gpu'
+        expected_backend = 'cpu' if not isGPU() else 'gpu'
         self.assertEqual(expected_backend, jax.default_backend())