perf(workers): slim the remaining ML images via real multi-stage builds

[arjunrajlab]

Applies the pattern from #160 (SAM/SAM2) to the eight other images built by build_machine_learning_workers.sh: cellpose, cellpose_train, cellposesam, cellposesam_train, stardist, condensatenet, piscis/predict, piscis/train.

Note

Built, measured and exercised on x86_64 + RTX 3060 against a live NimbusImage stack. The size table below is measured, not estimated. One checklist item — an actual retrain — remains unverified; see Verification.

The problem

All eight carried the exact lineage #160 diagnosed:

• a *-devel CUDA base whose CUDA stack is shipped twice (PyTorch loads the cuDNN/cuBLAS/cuFFT/NCCL inside its own nvidia-* pip wheels, never the image's copy)
• a no-op FROM base as build — nothing was ever copied between stages, so the whole build environment shipped
• no PIP_NO_CACHE_DIR (a full second copy of the wheels) and no conda clean (a second copy of the env)
• full-history git clones
• a copy-pasted apt block installing r-base twice, plus software-properties-common and python3 — none of which any of these workers use

What changed

Per worker: a real two-stage build (devel base for the build stage, *-runtime for the final stage, copying only the finished conda env, the editable-install trees and the baked-in model cache), ENV PIP_NO_CACHE_DIR=1, conda clean --all --yes, shallow clones, and the dead apt packages dropped. Worker .py files now land in the final stage, so editing one no longer invalidates the annotation_utilities / worker_client installs.

Model caches are carried across the stage boundary so first-run downloads are still avoided — /root/.cellpose (both cellpose 3.x and 4.x resolve from ~/.cellpose/models), /root/.keras (csbdeep's from_pretrained → keras.utils.get_file), /models (condensatenet — real files, snapshot_download uses local_dir_use_symlinks=False) and /root/.piscis.

No environment.yml was touched. Unlike the SAM workers, none of these had a dependency that was provably unused.

stardist is the exception on the cuDNN tag

Its runtime stage keeps cudnn8; the other seven drop it. TensorFlow 2.11 predates the tensorflow[and-cuda] extra (2.14), so its wheel declares no nvidia-* dependencies and dlopens libcudnn.so.8 from the image. A plain -runtime tag would not fail the build — TF logs Could not load dynamic library libcudnn.so.8 and silently falls back to CPU.

Confirmed against the built images: pip list shows nvidia-cudnn-cu12/cu13 wheels in all seven torch workers and none in stardist.

Two piscis bugs fixed along the way

• Both images ran git clone https://github.com/arjunrajlaboratory/ImageAnalysisProject/ and installed annotation_utilities / worker_client from that clone, so the build used whatever was on the default branch rather than the tree being built — the same bug Reduce ML worker image sizes via multi-stage builds and cache cleanup #160 fixed in sam_automatic_mask_generator. They now COPY the local directories.
• Miniconda → Miniforge, matching every other worker. Drops the defaults channel and both conda tos accept calls; piscis/environment.yml is only python=3.11 + pip.

The r-base trap (second commit)

Dropping r-base is not purely cosmetic. r-base-core depends on libglib2.0-0, libx11-6, libxt6 (→ libsm6, libice6), libtk8.6 (→ libxext6) and libcairo2 (→ libxcb1), so every worker carrying it had been getting those shared libraries by accident.

That breaks exactly one worker: piscis depends on opencv-python, not -headless, and piscis/transforms.py does import cv2 as cv at module scope. Without those libs import piscis raises ImportError — in the build stage too, since download_models.py imports the package. Both piscis stages now install libgl1 libglib2.0-0 libsm6 libice6 libx11-6 libxext6 explicitly.

The rest of the fleet was cleared by reading DT_NEEDED off every wheel's .so files against what each wheel vendors:

Wheel	Needs from the image
opencv-python (piscis)	libGL, libglib/libgthread, libX11, libSM, libICE, libXext, libxcb
opencv-python-headless (cellpose ×4)	nothing beyond glibc/libstdc++
torch	nothing — vendors its own libgomp and the whole CUDA stack
tensorflow==2.11.0 (stardist)	nothing linked; dlopens libcudnn.so.8 by name
stardist==0.9.1, rtree, shapely	nothing beyond glibc/libstdc++

ptxas / nvcc are gone from the runtime stage (third commit)

The *-runtime images have no /usr/local/cuda/bin, so the toolkit binaries disappear. Visible in one place: stardist logs Couldn't invoke ptxas --version / Relying on driver to perform ptx compilation on every run. The fallback is correct — a real segmentation returned the right object count with cuDNN 8906 loaded on the GPU — so this is accepted rather than fixed; adding cuda-nvcc back would cost a few hundred MB to silence log noise.

The seven torch workers are unaffected, and not because they never JIT: the triton wheel bundles its own ptxas under site-packages/triton/backends/nvidia/bin/. Documented in todo/ml-worker-image-size.md and inline in the stardist Dockerfile.

Measured impact

Shipped image size, unpacked, from docker image inspect .Size (base-10).

Worker	Before	After	Saved
cellpose	25.23 G	8.87 G	−16.36 G (65%)
cellpose_train	25.07 G	8.88 G	−16.20 G (65%)
cellposesam	23.21 G	11.21 G	−12.00 G (52%)
cellposesam_train	23.05 G	11.22 G	−11.83 G (51%)
stardist	20.44 G	6.69 G	−13.74 G (67%)
condensatenet	24.09 G	8.28 G	−15.81 G (66%)
piscis/predict	18.53 G	8.71 G	−9.82 G (53%)
piscis/train	18.87 G	8.82 G	−10.05 G (53%)
Fleet	178.5 G	72.7 G	−105.8 G (59%)

An earlier revision of this description estimated a 215 G → 96 G fleet (−55%). That was written without a Docker daemon and overstated the baseline: the real saving is smaller in absolute terms (−105.8 G, not −119 G) but a larger fraction (59%, not 55%).

Build times were 136–293 s per worker. Compressed pull sizes were not measured — every figure above is unpacked. The earlier claim that cold build time is "roughly a wash" is unverified and withdrawn: the pre-change images were not rebuilt, so there is no before-time to compare against.

Verification

1. Model caches survived the stage copy — /root/.cellpose/models (cyto*/nuclei* for cellpose 3, cpsam+cpsam_v2 for cellpose 4), /root/.keras/models/StarDist2D, /models/condensatenet-v1, /root/.piscis/models (four dated models + the rajlab collection). No first-run download.
2. GPU genuinely used from the runtime base — torch.cuda.is_available() true for the seven torch workers; stardist reports [PhysicalDevice('/physical_device:GPU:0')] with Loaded cuDNN version 8906 and no Could not load dynamic library libcudnn line.
3. import cv2 succeeds in both piscis images — the check that would have caught the r-base trap.
4. Real jobs against a live NimbusImage stack, all writing annotations back to Girder: cellposesam (17 objects), stardist (17), cellpose (12), condensatenet (7976 condensates), piscis predict (spots). cellposesam was additionally run through the NimbusImage UI, whose job log shows runtime: nvidia — so the isGPUWorker label still routes to the GPU queue through girder_worker.
5. Training workers — GPU visible, checkpoints present, the copied model directory is writable, and from cellpose import train imports.

Not verified

An actual retrain. Item 5 covers the preconditions but no training run was executed, so nothing has yet written a real checkpoint into a directory that arrives via COPY --from. This is the weakest point in the verification.

Out of scope

deconwolf is the last GPU worker on this lineage. It is an image-processing worker rather than an ML one and the only one compiling a native binary, so it needs its own pass — todo/ml-worker-image-size.md records what that involves, including using ldd /usr/bin/dw to enumerate its runtime libs rather than assuming the base image carries them.

---

Generated by Claude Code