Accelerator images, and a trimmed multi-stage build

.github/workflows/images.yml

@@ -67,7 +67,7 @@ jobs:
 
   images:
     needs: version
-    runs-on: ubuntu-latest
+    runs-on: ${{ matrix.runner || 'ubuntu-latest' }}
     strategy:
       fail-fast: false
       matrix:
@@ -123,9 +123,54 @@ jobs:
             desktop: 0
             python: "3.12"
             gpu: true
+          # The accelerator add-on (C kernels, and the CUDA library on the cuda one),
+          # fetched from its private release with the ACCELERATOR_TOKEN secret.
+          # arm64 builds natively on an Arm runner rather than under QEMU.
+          - variant: accel-amd64
+            flavour: web
+            extras: gui
+            torch_index: ""
+            desktop: 0
+            python: "3.14"
+            gpu: false
+            accel: 1
+            accel_slug: linux-x86_64
+            accel_cuda: 0
+          - variant: accel-arm64
+            flavour: web
+            extras: gui
+            torch_index: ""
+            desktop: 0
+            python: "3.14"
+            gpu: false
+            accel: 1
+            accel_slug: linux-aarch64
+            accel_cuda: 0
+            runner: ubuntu-24.04-arm
+            platform: linux/arm64
+          - variant: cuda-accel-amd64
+            flavour: web
+            extras: gui,cellpose
+            torch_index: https://download.pytorch.org/whl/cu126
+            desktop: 0
+            python: "3.14"
+            gpu: true
+            accel: 1
+            accel_slug: linux-x86_64
+            accel_cuda: 1
     steps:
       - uses: actions/checkout@v4
 
+      - name: Check the accelerator token
+        if: matrix.accel == 1
+        env:
+          ACCEL_TOKEN: ${{ secrets.ACCELERATOR_TOKEN }}
+        run: |
+          if [ -z "$ACCEL_TOKEN" ]; then
+            echo "::error::ACCELERATOR_TOKEN is not set. It needs read access to contents of FLIMKit/FLIMKit-accelerator-plugin."
+            exit 1
+          fi
+
       - name: Free disk space
         if: matrix.gpu
         run: |
@@ -139,7 +184,7 @@ jobs:
         uses: docker/build-push-action@v6
         with:
           context: .
-          platforms: linux/amd64
+          platforms: ${{ matrix.platform || 'linux/amd64' }}
           load: true
           tags: flimkit-smoke:${{ matrix.variant }}
           build-args: |
@@ -149,6 +194,11 @@ jobs:
             TORCH_INDEX=${{ matrix.torch_index }}
             INCLUDE_DESKTOP=${{ matrix.desktop }}
             FLAVOUR=${{ matrix.flavour }}
+            ACCELERATOR=${{ matrix.accel || 0 }}
+            ACCEL_SLUG=${{ matrix.accel_slug || 'linux-x86_64' }}
+            ACCEL_CUDA=${{ matrix.accel_cuda || 0 }}
+          secrets: |
+            accel_token=${{ secrets.ACCELERATOR_TOKEN }}
           cache-from: type=gha,scope=${{ matrix.variant }}
           cache-to: type=gha,mode=max,scope=${{ matrix.variant }}
 
@@ -166,6 +216,18 @@ jobs:
             | python3 -c "import json, sys; s = json.load(sys.stdin); print('web UI up:', s['form'], len(s['fields']), 'fields')"
           docker exec smoke python -c "import flimkit; print('flimkit', flimkit.__file__)"
 
+      - name: Smoke test the accelerator
+        if: matrix.accel == 1
+        run: |
+          docker exec smoke python - <<'PY'
+          import ctypes, os
+          d = os.environ['FLIMKIT_PLUGIN_PATH'] + '/flimkit_accelerator_plugin/'
+          lib = ctypes.CDLL(d + 'libflimkit_accelerator.so')
+          print('C kernel abi', lib.flimkit_c_abi_version(), 'threads', lib.flimkit_c_threads())
+          assert lib.flimkit_c_abi_version() == 3
+          assert os.path.exists(d + 'libflimkit_accelerator_cuda.so') == bool(${{ matrix.accel_cuda }})
+          PY
+
       - name: Smoke test the bridge
         if: matrix.flavour == 'bridge'
         run: |
@@ -197,7 +259,7 @@ jobs:
         uses: docker/build-push-action@v6
         with:
           context: .
-          platforms: linux/amd64
+          platforms: ${{ matrix.platform || 'linux/amd64' }}
           push: true
           tags: |
             ${{ env.IMAGE }}:${{ matrix.variant }}
@@ -209,4 +271,9 @@ jobs:
             TORCH_INDEX=${{ matrix.torch_index }}
             INCLUDE_DESKTOP=${{ matrix.desktop }}
             FLAVOUR=${{ matrix.flavour }}
+            ACCELERATOR=${{ matrix.accel || 0 }}
+            ACCEL_SLUG=${{ matrix.accel_slug || 'linux-x86_64' }}
+            ACCEL_CUDA=${{ matrix.accel_cuda || 0 }}
+          secrets: |
+            accel_token=${{ secrets.ACCELERATOR_TOKEN }}
           cache-from: type=gha,scope=${{ matrix.variant }}

Dockerfile

@@ -1,24 +1,36 @@
 ARG PYTHON_VERSION=3.14
-FROM python:${PYTHON_VERSION}-slim-bookworm
-
-RUN apt-get update && apt-get install -y --no-install-recommends \
-        tzdata \
-        libglib2.0-0 \
-    && rm -rf /var/lib/apt/lists/*
 
-# web serves the browser UI and needs Tk, so it needs an X server to draw on.
-# bridge serves the HTTP API that QuPath and Fiji speak, and needs neither.
-ARG FLAVOUR=web
-RUN if [ "${FLAVOUR}" = "web" ]; then \
-        apt-get update && apt-get install -y --no-install-recommends \
-            xvfb x11-utils python3-tk \
-            libx11-6 libxext6 libxcomposite1 libxdamage1 libxfixes3 libxrandr2 \
-            fonts-dejavu-core \
-        && rm -rf /var/lib/apt/lists/*; \
+# Throwaway stage that fetches the accelerator add-on from its GitHub release.
+# The repo is private, so the token comes in as a build secret and never lands in a layer.
+# With ACCELERATOR=0 it only produces an empty folder.
+FROM debian:bookworm-slim AS accelerator
+ARG ACCELERATOR=0
+ARG ACCEL_SLUG=linux-x86_64
+ARG ACCEL_CUDA=0
+ARG ACCEL_TAG=latest
+RUN mkdir -p /plugins
+RUN --mount=type=secret,id=accel_token \
+    if [ "${ACCELERATOR}" = "1" ]; then \
+        apt-get update && apt-get install -y --no-install-recommends curl ca-certificates unzip jq \
+        && token=$(cat /run/secrets/accel_token) \
+        && repo=https://api.github.com/repos/FLIMKit/FLIMKit-accelerator-plugin/releases \
+        && if [ "${ACCEL_TAG}" = "latest" ]; then url=${repo}/latest; else url=${repo}/tags/${ACCEL_TAG}; fi \
+        && id=$(curl -fsS -H "Authorization: Bearer ${token}" "${url}" \
+               | jq -r ".assets[] | select(.name == \"flimkit-accelerator-${ACCEL_SLUG}.zip\") | .id") \
+        && test -n "${id}" \
+        && curl -fsSL -H "Authorization: Bearer ${token}" -H "Accept: application/octet-stream" \
+               -o /tmp/accel.zip "${repo}/assets/${id}" \
+        && unzip -q /tmp/accel.zip 'flimkit_accelerator_plugin/*' -d /plugins \
+        && if [ "${ACCEL_CUDA}" != "1" ]; then rm -f /plugins/flimkit_accelerator_plugin/*cuda*; fi \
+        && rm -f /plugins/flimkit_accelerator_plugin/kernels.c /plugins/flimkit_accelerator_plugin/kernels.cu \
+        && chmod -R a+rX /plugins; \
     fi
 
-ENV TZ=Etc/UTC \
-    MPLCONFIGDIR=/tmp/mpl-cache
+# Throwaway stage that installs the Python packages into a venv and trims it.
+FROM python:${PYTHON_VERSION}-slim-bookworm AS build
+ARG FLAVOUR=web
+RUN python -m venv /opt/venv
+ENV PATH=/opt/venv/bin:${PATH}
 
 # Empty TORCH_INDEX leaves torch out, which is most of the image size.
 ARG TORCH_INDEX=
@@ -39,6 +51,38 @@ RUN if [ "${FLAVOUR}" = "web" ]; then \
         pip install --no-cache-dir "${FLIMKIT_BRIDGE}"; \
     fi
 
+# Drop what a running install never reads: the packages' own test suites and type
+# stubs. __pycache__ stays, since a non-root user cannot write it back and every
+# start would pay to recompile scipy and friends.
+RUN find /opt/venv -type d -name tests -prune -exec rm -rf {} + \
+    && find /opt/venv -name '*.pyi' -delete
+
+FROM python:${PYTHON_VERSION}-slim-bookworm
+
+RUN apt-get update && apt-get install -y --no-install-recommends \
+        tzdata \
+        libglib2.0-0 \
+    && rm -rf /var/lib/apt/lists/*
+
+# web serves the browser UI and needs Tk, so it needs an X server to draw on.
+# bridge serves the HTTP API that QuPath and Fiji speak, and needs neither.
+ARG FLAVOUR=web
+RUN if [ "${FLAVOUR}" = "web" ]; then \
+        apt-get update && apt-get install -y --no-install-recommends \
+            xvfb x11-utils python3-tk \
+            libx11-6 libxext6 libxcomposite1 libxdamage1 libxfixes3 libxrandr2 \
+            fonts-dejavu-core \
+        && rm -rf /var/lib/apt/lists/*; \
+    fi
+
+ENV TZ=Etc/UTC \
+    MPLCONFIGDIR=/tmp/mpl-cache
+
+# Everything pip installs is built in the stage above and copied across as one venv,
+# so pip's own scratch space and the strip step below never reach this image.
+COPY --from=build /opt/venv /opt/venv
+ENV PATH=/opt/venv/bin:${PATH}
+
 RUN mkdir -p /tmp/mpl-cache && chmod 777 /tmp/mpl-cache
 
 # TrueNAS and other hosts run containers as an arbitrary non-root user. Xvfb
@@ -61,6 +105,11 @@ RUN if [ "${INCLUDE_DESKTOP}" = "1" ]; then \
         && for ic in /usr/share/novnc/app/images/icons/novnc-*.png; do [ -e "$ic" ] && cp -f "$icon" "$ic"; done; \
     fi
 
+# Empty unless ACCELERATOR=1. Folders named in FLIMKIT_PLUGIN_PATH load without the
+# "Load from ~/.flimkit/plugins" switch, which nobody can tick inside a container.
+COPY --from=accelerator /plugins /opt/flimkit-plugins
+ENV FLIMKIT_PLUGIN_PATH=/opt/flimkit-plugins
+
 COPY docker-entrypoint.sh /docker-entrypoint.sh
 RUN chmod +x /docker-entrypoint.sh
 

README.md

@@ -25,9 +25,14 @@ Then open **http://localhost:14500** and log in as `flimkit` with that password.
 | `cellpose` | adds Cellpose and CPU torch, for "Apply cell mask" | ~2.4 GB |
 | `cuda` | adds Cellpose and CUDA torch, `docker run --gpus all ...` | larger again |
 | `rocm` | adds Cellpose and ROCm torch, `docker run --device /dev/kfd --device /dev/dri ...` | larger again |
+| `accel-amd64` | `latest` plus the [accelerator add-on](https://github.com/FLIMKit/flimkit-accelerator-plugin)'s C kernels, for x86_64 CPUs | `latest` + ~0.1 MB |
+| `accel-arm64` | the same on aarch64 (Graviton, Ampere, Apple silicon under Linux), built natively | `latest` + ~0.1 MB |
+| `cuda-accel-amd64` | `cuda` plus the add-on's C kernels and its CUDA library | `cuda` + ~2.3 MB |
 | `cuda-headless` | the bridge API instead of the web UI, with Cellpose and CUDA torch | ~270 MB less than `cuda` |
 | `rocm-headless` | the same for AMD | ~270 MB less than `rocm` |
 
+The three `accel` tags load the add-on from `/opt/flimkit-plugins` through `FLIMKIT_PLUGIN_PATH`, so free-τ fits use its C (and, on `cuda-accel-amd64` with `--gpus all`, CUDA) kernels without any setup. `FLIMKIT_FIT_BACKEND=cpu` turns it off. The CUDA library needs the host driver only.
+
 GPU fitting in FLIMKit is torch, so every GPU image carries it whether or not Cellpose is installed.
 
 Every tag is also published with its FLIMKit version, for example `0.13.5-latest` and `0.13.5-cuda`. Sizes are the unpacked filesystem measured inside the container; the download is smaller.