Refactor ML worker Dockerfiles: mamba + shared base images

REGISTRY.md

@@ -117,3 +117,13 @@
 | Worker | Path | Description |
 |--------|------|-------------|
 | test_worker | `workers/test_worker/` | Internal test worker for CI infrastructure |
+
+---
+
+## Shared Base Images (ML Workers)
+
+| Image | Dockerfile | CUDA | Used by |
+|-------|-----------|------|---------|
+| `nimbusimage/sam2-worker-base` | `workers/base_docker_images/Dockerfile.sam2_worker_base` | 12.1 | All 5 SAM2 workers |
+| `nimbusimage/sam2-worker-base-m1` | `workers/base_docker_images/Dockerfile.sam2_worker_base_M1` | 11.8 | All 5 SAM2 workers (M1) |
+| `nimbusimage/cuda-ml-worker-base` | `workers/base_docker_images/Dockerfile.cuda_ml_worker_base` | 11.8 | cellpose, cellpose_train, cellposesam, stardist, SAM1 workers |

build_machine_learning_workers.sh

@@ -19,51 +19,83 @@ else
     DOCKERFILE="Dockerfile"
 fi
 
+# ============================================================
+# Build shared base images first
+# ============================================================
+
+echo "============================================================"
+echo "Building shared base images..."
+echo "============================================================"
+
+if [ "$ARCH" == "arm64" ]; then
+    echo "Building SAM2 worker base image (M1)"
+    docker build . -f ./workers/base_docker_images/Dockerfile.sam2_worker_base_M1 -t nimbusimage/sam2-worker-base-m1:latest $NO_CACHE
+else
+    echo "Building SAM2 worker base image"
+    docker build . -f ./workers/base_docker_images/Dockerfile.sam2_worker_base -t nimbusimage/sam2-worker-base:latest $NO_CACHE
+
+    echo "Building CUDA ML worker base image"
+    docker build . -f ./workers/base_docker_images/Dockerfile.cuda_ml_worker_base -t nimbusimage/cuda-ml-worker-base:latest $NO_CACHE
+fi
+
+echo "============================================================"
+echo "Building individual workers..."
+echo "============================================================"
+
+# ============================================================
+# Piscis (uses docker-compose, standalone CUDA 12.4, x86_64 only)
+# ============================================================
+
 echo "Building Piscis worker"
 docker compose -f ./workers/annotations/piscis/docker-compose.yaml build $NO_CACHE
 
-echo "Building Cellpose worker"
-docker build ./workers/annotations/cellpose/ -t annotations/cellpose_worker:latest $NO_CACHE
+# ============================================================
+# Cellpose workers (standalone, own conda env for GPU compatibility)
+# ============================================================
+
+if [ "$ARCH" != "arm64" ]; then
+    echo "Building Cellpose worker"
+    docker build ./workers/annotations/cellpose/ -t annotations/cellpose_worker:latest $NO_CACHE
 
-echo "Building Cellpose train worker"
-docker build ./workers/annotations/cellpose_train/ -t annotations/cellpose_train_worker:latest $NO_CACHE
+    echo "Building Cellpose train worker"
+    docker build ./workers/annotations/cellpose_train/ -t annotations/cellpose_train_worker:latest $NO_CACHE
 
-echo "Building Cellpose-SAM worker"
-docker build ./workers/annotations/cellposesam/ -t annotations/cellposesam_worker:latest $NO_CACHE
+    echo "Building Cellpose-SAM worker"
+    docker build ./workers/annotations/cellposesam/ -t annotations/cellposesam_worker:latest $NO_CACHE
 
-echo "Building Stardist worker"
-docker build ./workers/annotations/stardist/ -t annotations/stardist_worker:latest $NO_CACHE
+    echo "Building Stardist worker"
+    docker build . -f ./workers/annotations/stardist/Dockerfile -t annotations/stardist_worker:latest $NO_CACHE
+fi
+
+# SAM1 workers (have M1 variants)
+echo "Building SAM few-shot segmentation worker"
+docker build . -f ./workers/annotations/sam_fewshot_segmentation/$DOCKERFILE -t annotations/sam_fewshot_segmentation:latest $NO_CACHE
+
+echo "Building SAM automatic mask generator worker"
+docker build . -f ./workers/annotations/sam_automatic_mask_generator/$DOCKERFILE -t annotations/sam_automatic_mask_generator:latest $NO_CACHE
+
+# ============================================================
+# SAM2 workers (standalone, own conda env for GPU compatibility)
+# ============================================================
 
 echo "Building SAM2 automatic mask generator worker"
 docker build . -f ./workers/annotations/sam2_automatic_mask_generator/Dockerfile -t annotations/sam2_automatic_mask_generator:latest $NO_CACHE
-# Command for M1:
-# docker build . -f ./workers/annotations/sam2_automatic_mask_generator/Dockerfile_M1 -t annotations/sam2_automatic_mask_generator:latest $NO_CACHE
 
 echo "Building SAM2 few-shot segmentation worker"
 docker build . -f ./workers/annotations/sam2_fewshot_segmentation/Dockerfile -t annotations/sam2_fewshot_segmentation:latest $NO_CACHE
-# Command for M1:
-# docker build . -f ./workers/annotations/sam2_fewshot_segmentation/Dockerfile_M1 -t annotations/sam2_fewshot_segmentation:latest $NO_CACHE
-
-echo "Building SAM few-shot segmentation worker"
-docker build . -f ./workers/annotations/sam_fewshot_segmentation/Dockerfile -t annotations/sam_fewshot_segmentation:latest $NO_CACHE
-# Command for M1:
-# docker build . -f ./workers/annotations/sam_fewshot_segmentation/Dockerfile_M1 -t annotations/sam_fewshot_segmentation:latest $NO_CACHE
 
 echo "Building SAM2 propagate worker"
-docker build . -f ./workers/annotations/sam2_propagate/$DOCKERFILE -t annotations/sam2_propagate_worker:latest $NO_CACHE
+docker build . -f ./workers/annotations/sam2_propagate/Dockerfile -t annotations/sam2_propagate_worker:latest $NO_CACHE
 
 echo "Building SAM2 refine worker"
-docker build . -f ./workers/annotations/sam2_refine/$DOCKERFILE -t annotations/sam2_refine_worker:latest $NO_CACHE
+docker build . -f ./workers/annotations/sam2_refine/Dockerfile -t annotations/sam2_refine_worker:latest $NO_CACHE
 
 echo "Building SAM2 video worker"
-docker build . -f ./workers/annotations/sam2_video/$DOCKERFILE -t annotations/sam2_video_worker:latest $NO_CACHE
+docker build . -f ./workers/annotations/sam2_video/Dockerfile -t annotations/sam2_video_worker:latest $NO_CACHE
+
+# ============================================================
+# Standalone workers (CUDA 12.1, unique dependencies)
+# ============================================================
 
 echo "Building CondensateNet worker"
 docker build . -f ./workers/annotations/condensatenet/$DOCKERFILE -t annotations/condensatenet:latest $NO_CACHE
-
-# These are some legacy workers that are no longer used.
-#docker build ./workers/annotations/cellori_segmentation/ -t annotations/cellori_segmentation_worker:latest --label isUPennContrastWorker --label isAnnotationWorker --label "interfaceName=Cellori" --label "interfaceCategory=Cellori"
-#docker build ./workers/annotations/deepcell/ -t annotations/deepcell_worker:latest --label isUPennContrastWorker --label isAnnotationWorker --label "interfaceName=DeepCell" --label "interfaceCategory=Deepcell"
-#docker build ./workers/test_worker/ -t both/test_worker:latest  --label isUPennContrastWorker --label isAnnotationWorker --label isPropertyWorker --label "annotationShape=point" --label "interfaceName=Test worker" --label "interfaceCategory=Test"
-#docker build ./workers/annotations/test_multiple_annotation/ -t annotations/test_multiple_annotation:latest --label isUPennContrastWorker --label isAnnotationWorker --label "interfaceName=Random square" --label "interfaceCategory=random"
-

workers/annotations/cellpose/Dockerfile

@@ -1,4 +1,4 @@
-FROM nvidia/cuda:11.8.0-cudnn8-devel-ubuntu22.04 as base
+FROM nvidia/cuda:11.8.0-cudnn8-devel-ubuntu22.04 AS base
 LABEL isUPennContrastWorker=True
 LABEL com.nvidia.volumes.needed="nvidia_driver"
 
@@ -30,7 +30,7 @@ RUN wget \
     && bash Miniforge3-Linux-x86_64.sh -b \
     && rm -f Miniforge3-Linux-x86_64.sh
 
-FROM base as build
+FROM base AS build
 
 COPY ./environment.yml /
 RUN conda env create --file /environment.yml
@@ -40,10 +40,10 @@ RUN git clone https://github.com/arjunrajlaboratory/DeepTile
 RUN pip install -e /DeepTile
 RUN pip install rtree shapely
 
-RUN git clone https://github.com/Kitware/UPennContrast/
+RUN git clone https://github.com/arjunrajlaboratory/NimbusImage/
 
-RUN pip install -r /UPennContrast/devops/girder/annotation_client/requirements.txt
-RUN pip install -e /UPennContrast/devops/girder/annotation_client/
+RUN pip install -r /NimbusImage/devops/girder/annotation_client/requirements.txt
+RUN pip install -e /NimbusImage/devops/girder/annotation_client/
 
 COPY ./download_models.py /
 RUN python /download_models.py
@@ -64,4 +64,4 @@ LABEL isUPennContrastWorker=True \
       description="Uses Cellpose to find cells and nuclei" \
       defaultToolName="Cellpose"
 
-ENTRYPOINT ["conda", "run", "--no-capture-output", "-n", "worker", "python", "/entrypoint.py"]
+ENTRYPOINT ["conda", "run", "--no-capture-output", "-n", "worker", "python", "/entrypoint.py"]

workers/annotations/cellpose/environment.yml

@@ -7,4 +7,4 @@ dependencies:
 - shapely
 - pip:
   - cellpose==3.1.1.2
-  - deeptile
+  - deeptile

workers/annotations/cellpose_train/Dockerfile

@@ -1,4 +1,4 @@
-FROM nvidia/cuda:11.8.0-cudnn8-devel-ubuntu22.04 as base
+FROM nvidia/cuda:11.8.0-cudnn8-devel-ubuntu22.04 AS base
 LABEL isUPennContrastWorker=True
 LABEL com.nvidia.volumes.needed="nvidia_driver"
 
@@ -30,18 +30,18 @@ RUN wget \
     && bash Miniforge3-Linux-x86_64.sh -b \
     && rm -f Miniforge3-Linux-x86_64.sh
 
-FROM base as build
+FROM base AS build
 
 COPY ./environment.yml /
 RUN conda env create --file /environment.yml
 SHELL ["conda", "run", "-n", "worker", "/bin/bash", "-c"]
 
 RUN pip install rtree shapely
 
-RUN git clone https://github.com/Kitware/UPennContrast/
+RUN git clone https://github.com/arjunrajlaboratory/NimbusImage/
 
-RUN pip install -r /UPennContrast/devops/girder/annotation_client/requirements.txt
-RUN pip install -e /UPennContrast/devops/girder/annotation_client/
+RUN pip install -r /NimbusImage/devops/girder/annotation_client/requirements.txt
+RUN pip install -e /NimbusImage/devops/girder/annotation_client/
 
 COPY ./download_models.py /
 RUN python /download_models.py
@@ -64,4 +64,4 @@ LABEL isUPennContrastWorker="" \
       defaultToolName="Cellpose retrain" \
       annotationShape="polygon"
 
-ENTRYPOINT ["conda", "run", "--no-capture-output", "-n", "worker", "python", "/entrypoint.py"]
+ENTRYPOINT ["conda", "run", "--no-capture-output", "-n", "worker", "python", "/entrypoint.py"]

workers/annotations/cellpose_train/environment.yml

@@ -7,4 +7,4 @@ dependencies:
 - shapely
 - pip:
   - cellpose==3.1.1.2
-  - scikit-image
+  - scikit-image

workers/annotations/cellposesam/Dockerfile

@@ -1,4 +1,4 @@
-FROM nvidia/cuda:11.8.0-cudnn8-devel-ubuntu22.04 as base
+FROM nvidia/cuda:11.8.0-cudnn8-devel-ubuntu22.04 AS base
 LABEL isUPennContrastWorker=True
 LABEL com.nvidia.volumes.needed="nvidia_driver"
 
@@ -30,7 +30,7 @@ RUN wget \
     && bash Miniforge3-Linux-x86_64.sh -b \
     && rm -f Miniforge3-Linux-x86_64.sh
 
-FROM base as build
+FROM base AS build
 
 COPY ./environment.yml /
 RUN conda env create --file /environment.yml
@@ -40,10 +40,10 @@ RUN git clone https://github.com/arjunrajlaboratory/DeepTile
 RUN pip install -e /DeepTile
 RUN pip install rtree shapely
 
-RUN git clone https://github.com/Kitware/UPennContrast/
+RUN git clone https://github.com/arjunrajlaboratory/NimbusImage/
 
-RUN pip install -r /UPennContrast/devops/girder/annotation_client/requirements.txt
-RUN pip install -e /UPennContrast/devops/girder/annotation_client/
+RUN pip install -r /NimbusImage/devops/girder/annotation_client/requirements.txt
+RUN pip install -e /NimbusImage/devops/girder/annotation_client/
 
 COPY ./download_models.py /
 RUN python /download_models.py

workers/annotations/cellposesam/environment.yml

@@ -7,4 +7,4 @@ dependencies:
 - shapely
 - pip:
   - cellpose==4.0.1
-  - deeptile
+  - deeptile

workers/annotations/condensatenet/Dockerfile

@@ -32,7 +32,7 @@ FROM base as build
 
 # Copy and create conda environment
 COPY ./workers/annotations/condensatenet/environment.yml /
-RUN conda env create --file /environment.yml
+RUN mamba env create --file /environment.yml
 SHELL ["conda", "run", "-n", "worker", "/bin/bash", "-c"]
 
 # Install NimbusImage annotation client

workers/annotations/deconwolf/Dockerfile

@@ -55,7 +55,7 @@ FROM base AS build
 
 # Copy and create conda environment
 COPY ./workers/annotations/deconwolf/environment.yml /
-RUN conda env create --file /environment.yml
+RUN mamba env create --file /environment.yml
 SHELL ["conda", "run", "-n", "worker", "/bin/bash", "-c"]
 
 # Install NimbusImage annotation client

workers/annotations/piscis/predict/Dockerfile

@@ -1,14 +1,9 @@
-FROM nvidia/cuda:12.4.1-devel-ubuntu22.04 as base
+FROM nvidia/cuda:12.4.1-devel-ubuntu22.04 AS base
 LABEL com.nvidia.volumes.needed="nvidia_driver"
 
 ENV NVIDIA_VISIBLE_DEVICES all
 ENV NVIDIA_DRIVER_CAPABILITIES compute,utility
 
-LABEL isUPennContrastWorker=True
-LABEL isAnnotationWorker=True
-LABEL interfaceName="Piscis (Predict)"
-LABEL interfaceCategory="Piscis"
-
 RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -qy tzdata && \
   apt-get install -qy software-properties-common python3-software-properties && \
   apt-get update && apt-get install -qy \
@@ -20,33 +15,28 @@ RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -qy tzdata
   libssl-dev \
   libjpeg-dev \
   zlib1g-dev \
-  r-base \
   git \
   libpython3-dev && \
   apt-get clean && rm -rf /var/lib/apt/lists/*
 
-ENV PATH="/root/miniconda3/bin:$PATH"
-ARG PATH="/root/miniconda3/bin:$PATH"
+ENV PATH="/root/miniforge3/bin:$PATH"
+ARG PATH="/root/miniforge3/bin:$PATH"
 
 RUN wget \
-    https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh \
+    https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-Linux-x86_64.sh \
     && mkdir /root/.conda \
-    && bash Miniconda3-latest-Linux-x86_64.sh -b \
-    && rm -f Miniconda3-latest-Linux-x86_64.sh
+    && bash Miniforge3-Linux-x86_64.sh -b \
+    && rm -f Miniforge3-Linux-x86_64.sh
 
-FROM base as build
-
-RUN conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/main && \
-    conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/r
+FROM base AS build
 
 COPY ./environment.yml /
-RUN conda env create --file /environment.yml
+RUN mamba env create --file /environment.yml
 SHELL ["conda", "run", "-n", "worker", "/bin/bash", "-c"]
 
-RUN git clone https://github.com/Kitware/UPennContrast/
-
-RUN pip install -r /UPennContrast/devops/girder/annotation_client/requirements.txt
-RUN pip install -e /UPennContrast/devops/girder/annotation_client/
+RUN git clone https://github.com/arjunrajlaboratory/NimbusImage/
+RUN pip install -r /NimbusImage/devops/girder/annotation_client/requirements.txt
+RUN pip install -e /NimbusImage/devops/girder/annotation_client/
 
 RUN pip install piscis==1.0.0
 RUN pip install flax
@@ -69,4 +59,4 @@ LABEL isUPennContrastWorker=True \
       description="Uses Piscis to find spots" \
       defaultToolName="Spots"
 
-ENTRYPOINT ["conda", "run", "--no-capture-output", "-n", "worker", "python", "/entrypoint.py"]
+ENTRYPOINT ["conda", "run", "--no-capture-output", "-n", "worker", "python", "/entrypoint.py"]

workers/annotations/piscis/train/Dockerfile

@@ -1,14 +1,9 @@
-FROM nvidia/cuda:12.4.1-devel-ubuntu22.04 as base
+FROM nvidia/cuda:12.4.1-devel-ubuntu22.04 AS base
 LABEL com.nvidia.volumes.needed="nvidia_driver"
 
 ENV NVIDIA_VISIBLE_DEVICES all
 ENV NVIDIA_DRIVER_CAPABILITIES compute,utility
 
-LABEL isUPennContrastWorker=True
-LABEL isAnnotationWorker=True
-LABEL interfaceName="Piscis (Train)"
-LABEL interfaceCategory="Piscis"
-
 RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -qy tzdata && \
   apt-get install -qy software-properties-common python3-software-properties && \
   apt-get update && apt-get install -qy \
@@ -20,35 +15,29 @@ RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -qy tzdata
   libssl-dev \
   libjpeg-dev \
   zlib1g-dev \
-  r-base \
   git \
   libpython3-dev && \
   apt-get clean && rm -rf /var/lib/apt/lists/*
 
-ENV PATH="/root/miniconda3/bin:$PATH"
-ARG PATH="/root/miniconda3/bin:$PATH"
+ENV PATH="/root/miniforge3/bin:$PATH"
+ARG PATH="/root/miniforge3/bin:$PATH"
 
 RUN wget \
-    https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh \
+    https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-Linux-x86_64.sh \
     && mkdir /root/.conda \
-    && bash Miniconda3-latest-Linux-x86_64.sh -b \
-    && rm -f Miniconda3-latest-Linux-x86_64.sh
-
-FROM base as build
+    && bash Miniforge3-Linux-x86_64.sh -b \
+    && rm -f Miniforge3-Linux-x86_64.sh
 
-RUN conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/main && \
-    conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/r
+FROM base AS build
 
 COPY ./environment.yml /
-RUN conda env create --file /environment.yml
+RUN mamba env create --file /environment.yml
 SHELL ["conda", "run", "-n", "worker", "/bin/bash", "-c"]
 
-RUN git clone https://github.com/arjunrajlaboratory/NimbusImage
-
+RUN git clone https://github.com/arjunrajlaboratory/NimbusImage/
 RUN pip install -r /NimbusImage/devops/girder/annotation_client/requirements.txt
 RUN pip install -e /NimbusImage/devops/girder/annotation_client/
 
-
 RUN git clone https://github.com/zjniu/Piscis.git
 WORKDIR /Piscis
 RUN pip install .