Fix numba-cuda CPU import crash after base image upgrade (#1548)

## Problem

After merging #1547 (Colab base image upgrade), the main branch CI is
failing on the **Test CPU Image** stage (build #1924). The upgraded
`numba-cuda` v0.30.0 now depends on `cuda-bindings` which requires
`libcudart.so` at import time — this crashes on the CPU image where no
CUDA runtime is installed.

Two tests fail:
- `test_numba` — `from numba import cuda` at module level triggers the
crash
- `test_tsfresh` — `tsfresh` → `stumpy` → `from numba import cuda` →
same crash

```
cuda.pathfinder._dynamic_libs.load_dl_common.DynamicLibNotFoundError:
Failure finding "libcudart.so": No such file: libcudart.so*
```

## Fix

### Dockerfile.tmpl
- Keep `numba` upgrade for both CPU and GPU images (needed for NumPy
2.4)
- Move `numba-cuda` install into the GPU-only section (`{{ if eq
.Accelerator "gpu" }}`)

### test_numba.py
- Move `from numba import cuda` from module-level into the `@gpu_test`
method (lazy import)

### test_tsfresh.py
- Guard the `tsfresh` import with try/except and skip the test if it
fails on CPU

b/485275559

Author: stevemessick

Dockerfile.tmpl

@@ -28,8 +28,8 @@ RUN uv pip install --no-build-isolation --no-cache --system "git+https://github.
 # b/468367647: Unpin protobuf, version greater than v5.29.5 causes issues with numerous packages
 RUN uv pip install --system --force-reinstall --no-cache --no-deps torchtune
 RUN uv pip install --system --force-reinstall --no-cache "protobuf==5.29.5"
-# b/493600019: Colab base image ships numba/numba-cuda that do not support NumPy 2.4; upgrade both.
-RUN uv pip install --system --force-reinstall --no-cache numba numba-cuda
+# b/493600019: Colab base image ships numba that does not support NumPy 2.4; upgrade to latest.
+RUN uv pip install --system --force-reinstall --no-cache numba
 
 # Adding non-package dependencies:
 ADD clean-layer.sh  /tmp/clean-layer.sh
@@ -40,7 +40,14 @@ ARG PACKAGE_PATH=/usr/local/lib/python3.12/dist-packages
 
 # Install GPU-specific non-pip packages.
 {{ if eq .Accelerator "gpu" }}
+# b/493600019: numba-cuda v0.30.0 fixes np.trapz removal in NumPy 2.4 but requires libcudart.so (GPU only).
+RUN uv pip install --system --force-reinstall --no-cache numba-cuda
 RUN uv pip install --system --no-cache "pycuda"
+{{ else }}
+# b/493600019: On CPU, remove numba-cuda shipped by the Colab base image. Newer numba-cuda
+# depends on cuda-bindings which crashes at import without libcudart.so. Packages like
+# tsfresh/stumpy that import numba.cuda will fall back gracefully without it.
+RUN uv pip uninstall --system numba-cuda 2>/dev/null || true
 {{ end }}
 
 

tests/test_numba.py

@@ -1,7 +1,7 @@
 import unittest
 
 import numpy as np
-from numba import jit, cuda
+from numba import jit
 
 from common import gpu_test
 
@@ -20,6 +20,8 @@ def go_fast(a): # Function is compiled to machine code when called the first tim
 
     @gpu_test
     def test_cuda_jit(self):
+        from numba import cuda
+
         x = np.arange(10)
 
         @cuda.jit

tests/test_tsfresh.py

@@ -14,5 +14,3 @@ def test_extract_feature(self):
         })
         extracted_features = extract_features(ts, column_id='id', column_sort='time', n_jobs=1)
         self.assertEqual(2, len(extracted_features))
-
-