Rebuild the tensorflow base image on python:3.12

base_images/CHANGELOG.md

@@ -1,7 +1,23 @@
 # OpenProblems Base Images v1.2.0
 
+## MAJOR CHANGES
+
+* Rebuild `base_tensorflow_nvidia` on `python:3.12` with `tensorflow[and-cuda]` (PR #49).
+  nvcr.io stopped publishing TensorFlow images after 25.02, which left the image stuck on
+  TensorFlow 2.17, NumPy 1.26 and Scanpy 1.10. Taking CUDA from the `and-cuda` extra pins it
+  to whatever the TensorFlow wheel was built against, so the two can no longer drift apart
+  and silently fall back to the CPU.
+
+* Bump TensorFlow from 2.17 to 2.21, NumPy from 1.26 to 2.x and Scanpy from 1.10 to 1.12 in
+  `base_tensorflow_nvidia` (PR #49).
+
 ## BUG FIXES
 
+* `base_tensorflow_nvidia`: put the CUDA libraries that the `and-cuda` extra installs under
+  `site-packages/nvidia/*/lib` on the linker path with `ldconfig` (PR #49). TensorFlow does
+  not look there by itself, so without this it reports `Cannot dlopen some GPU libraries`,
+  finds no GPU and quietly trains on the CPU.
+
 * `base_pytorch_nvidia`: drop the pip-installed `cmake` in favour of apt's (PR #48). The pip
   shims in `/usr/local/bin` shadow `/usr/bin/cmake`, and they fail with `ModuleNotFoundError:
   No module named 'cmake'` when a build calls them from inside a pip build isolation
@@ -12,6 +28,10 @@
 * Check that `cmake`, `cpack` and `ctest` resolve to the apt-provided binaries in
   `base_pytorch_nvidia` (PR #48).
 
+* Check that the CUDA libraries are on the linker path in `base_tensorflow_nvidia` (PR #49).
+  There is no GPU in CI, but the linker can be asked whether it would find them, which is
+  enough to catch the failure above.
+
 # OpenProblems Base Images v1.1.0
 
 ## MAJOR CHANGES

base_images/src/tensorflow_nvidia/config.vsh.yaml

@@ -1,25 +1,41 @@
 name: base_tensorflow_nvidia
-description: An nvcr.io tensorflow with anndata preinstalled.
+description: |
+  A Python 3.12 image with GPU-enabled TensorFlow and anndata preinstalled.
+
+  nvcr.io stopped publishing TensorFlow images after 25.02, so this image is
+  built on the same `python:3.12` base as `base_python`. The CUDA and cuDNN
+  runtimes come from the `and-cuda` extra, which pins them to whatever the
+  TensorFlow wheel was built against -- installing TensorFlow on a base image
+  that ships its own CUDA lets the two drift apart and silently fall back to
+  the CPU.
 test_resources:
   - type: python_script
     path: test_tensorflow.py
   - type: python_script
     path: ../python/test_anndata.py
 engines:
   - type: docker
-    image: nvcr.io/nvidia/tensorflow:25.02-tf2-py3
+    image: python:3.12
     setup:
       - type: apt
         packages:
           - procps
           - git
       - type: python
         packages:
-          - numpy~=1.26.0
+          - tensorflow[and-cuda]~=2.21.0
           - anndata~=0.12.0
-          - scanpy~=1.10.0
+          - scanpy~=1.12
           - pyyaml
           - requests
           - jsonschema
         github:
           - openproblems-bio/core#subdirectory=packages/python/openproblems
+      # TensorFlow does not find the CUDA libraries that the `and-cuda` extra
+      # installs under `site-packages/nvidia/*/lib`, so put them on the linker
+      # path. Without this it silently falls back to the CPU.
+      - type: docker
+        run: |
+          find "$(python -c 'import nvidia, os; print(os.path.dirname(nvidia.__file__))')" -maxdepth 2 -name lib -type d > /etc/ld.so.conf.d/nvidia-pip.conf && \
+            test -s /etc/ld.so.conf.d/nvidia-pip.conf && \
+            ldconfig

base_images/src/tensorflow_nvidia/test_tensorflow.py

@@ -1,5 +1,6 @@
 import tensorflow as tf
 import numpy as np
+import ctypes
 import tempfile
 import os
 
@@ -14,6 +15,14 @@
 
 print(f"Using TensorFlow version: {tf.__version__}", flush=True)
 
+# The CUDA libraries live in site-packages/nvidia/*/lib, which is not on the
+# linker path by default. There is no GPU in CI to catch that, but the linker
+# can be asked whether it would find them.
+print("\n--- Checking that the CUDA libraries are on the linker path ---", flush=True)
+for soname in ["libcudnn.so.9", "libcublas.so.12"]:
+    ctypes.CDLL(soname)
+    print(f"  - {soname} loaded", flush=True)
+
 # Check for and list available physical devices (CPU/GPU)
 gpus = tf.config.list_physical_devices('GPU')
 if gpus: