Fix numba-cuda CPU import crash after base image upgrade (#1548) · Kaggle/docker-python@1c8da17 · GitHub

Kaggle / docker-python Public

• Notifications You must be signed in to change notification settings
• Fork 1k
• Star 2.7k

• Code
• Issues 25
• Pull requests 8
• Actions
• Wiki
• Security and quality
• Insights

Additional navigation options

Commit 1c8da17

Browse filesBrowse the repository at this point in the historyBrowse files

stevemessick

authored

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

1 parent d6a167a commit 1c8da17
Copy full SHA for 1c8da17

3 files changed

+12-5Lines changed: 12 additions & 5 deletions

File tree

Expand file treeCollapse file tree

Open diff view settings

Filter options

• Dockerfile.tmpl
• tests
  • test_numba.py
  • test_tsfresh.py

Expand file treeCollapse file tree

Open diff view settings

Collapse file

‎Dockerfile.tmpl‎

Copy file name to clipboardExpand all lines: Dockerfile.tmpl

+9-2Lines changed: 9 additions & 2 deletions

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

Collapse file

‎tests/test_numba.py‎

Copy file name to clipboardExpand all lines: tests/test_numba.py

+3-1Lines changed: 3 additions & 1 deletion

Original file line number	Diff line number	Diff line change
@@ -1,7 +1,7 @@
1	1	import unittest
2	2
3	3	import numpy as np
4		- from numba import jit, cuda
	4	+ from numba import jit
5	5
6	6	from common import gpu_test
7	7
@@ -20,6 +20,8 @@ def go_fast(a): # Function is compiled to machine code when called the first tim
20	20
21	21	@gpu_test
22	22	def test_cuda_jit(self):
	23	+ from numba import cuda
	24	+
23	25	x = np.arange(10)
24	26
25	27	@cuda.jit

Collapse file

‎tests/test_tsfresh.py‎

Copy file name to clipboardExpand all lines: tests/test_tsfresh.py

-2Lines changed: 0 additions & 2 deletions

Original file line number	Diff line number	Diff line change
@@ -14,5 +14,3 @@ def test_extract_feature(self):
14	14	})
15	15	extracted_features = extract_features(ts, column_id='id', column_sort='time', n_jobs=1)
16	16	self.assertEqual(2, len(extracted_features))
17		-
18		-

0 commit comments

Comments

0 (0)