Merge pull request #24 from imcf/feat/snakemake-workflow

feat: Snakemake workflow — per-tile GPU SLURM jobs

Author: lguerard

docs/guide/snakemake.md

@@ -0,0 +1,213 @@
+# Cluster workflow (Snakemake + SLURM)
+
+`tile_process` runs every tile **serially on one GPU**. For a large 3-D image
+that can be days. The bundled Snakemake workflow instead submits **one GPU job
+per tile**, so with *N* GPUs the segmentation is ~*N*× faster. This page walks
+through running it from scratch.
+
+```text
+convert ──▶ prepare (checkpoint) ──▶ segment {tile}  ──▶ merge
+                                     one GPU SLURM job per tile
+```
+
+## 1. Get the workflow
+
+The workflow lives in the `workflow/` directory of the patchworks repository
+(it is not shipped inside the pip package — it is a set of Snakemake files you
+run):
+
+```bash
+git clone https://github.com/imcf/patchworks
+cd patchworks/workflow
+```
+
+## 2. Install the dependencies
+
+You need patchworks with the workflow + reader + segmentation extras, in the
+environment Snakemake will use:
+
+```bash
+pip install "patchworks[workflow,cellpose,imaris,bioio]"
+```
+
+- `workflow` → Snakemake + the SLURM executor plugin
+- `cellpose` → the segmentation model
+- `imaris` / `bioio` → read your input format (`.ims`, `.czi`, `.lif`, …)
+
+On a cluster, do this inside a conda/venv that the compute nodes can see, or let
+each rule activate a conda env (see *Conda*, below).
+
+## 3. Configure the run
+
+Copy and edit `config/config.yaml`. Every field:
+
+```yaml
+# input / output
+input: "/data/scan.ims"        # .ims/.czi/.lif/.nd2/ome-tiff/.zarr
+work_dir: "/scratch/results"   # everything is written here
+
+# conversion (input → pyramidal OME-ZARR)
+reuse_pyramid: true            # .ims: copy its own pyramid (fast)
+convert_chunks: null           # null → bounded auto chunks; or [c,z,y,x]
+shard: false                   # true → pack chunks into shards (fewer files)
+
+# tiling
+channel: 0                     # channel to segment (null = keep all)
+level: 0                       # pyramid level (0 = full resolution)
+tile_shape: "auto"             # "auto", or e.g. [16, 1024, 1024] (zyx)
+overlap: 30                    # halo ≈ one object diameter
+skip_empty: true               # skip background tiles
+empty_threshold: null          # null → Otsu
+
+# segmentation
+method: "cellpose"             # "cellpose" (GPU) or "threshold" (no GPU)
+label_name: "cellpose"         # name under image.zarr/labels/
+cellpose:
+  model: "cyto3"
+  diameter: 30
+  do_3D: true
+  gpu: true
+  # extra model.eval() kwargs, e.g. flow_threshold: 0.4
+
+# label pyramid
+pyramid_levels: 5
+pyramid_downscale: 2
+sequential_labels: true        # renumber labels to a contiguous 1..N
+```
+
+!!! tip "Tile size vs runtime"
+    `tile_shape: "auto"` sizes each tile to your GPU's VRAM. Smaller tiles =
+    more (faster) jobs; very large 3-D tiles are slow. Keep `do_3D: false` (2-D
+    per slice) if your objects segment fine per slice — it is much faster.
+
+## 4. Dry-run (always do this first)
+
+Check the plan without running anything:
+
+```bash
+python -m snakemake -s Snakefile --configfile config/config.yaml -n -p
+```
+
+You should see `convert`, `prepare`, and a note that the **checkpoint** will add
+the `segment` jobs after `prepare` runs. (The number of segment jobs is only
+known after `prepare` decides which tiles are non-empty.)
+
+## 5a. Run locally (single machine)
+
+```bash
+python -m snakemake -s Snakefile --configfile config/config.yaml --cores 8
+```
+
+Tiles run on the local machine (one at a time on the GPU). Good for a small
+image or a smoke test.
+
+## 5b. Run on SLURM (one GPU job per tile)
+
+Edit `profile/slurm/config.yaml` for **your** cluster — partitions, account,
+and the GPU request:
+
+```yaml
+executor: slurm
+jobs: 64                       # max concurrent SLURM jobs ≈ GPUs you can grab
+default-resources:
+  slurm_partition: "cpu"       # your CPU partition
+  # slurm_account: "my_account"
+  mem_mb: 16000
+  cpus_per_task: 4
+  runtime: 60
+set-resources:
+  segment:                     # the GPU step
+    slurm_partition: "gpu"     # your GPU partition
+    slurm_extra: "'--gres=gpu:1'"
+    mem_mb: 32000
+    runtime: 120
+  merge:
+    mem_mb: 128000
+    runtime: 240
+```
+
+Then launch (from a login node — Snakemake submits and watches the jobs):
+
+```bash
+python -m snakemake --workflow-profile profile/slurm \
+                    --configfile config/config.yaml
+```
+
+Snakemake submits `convert`, then `prepare`, then **one `segment` job per
+non-empty tile** (up to `jobs:` at once → that many GPUs in parallel), then
+`merge`. Raise `jobs:` to use more GPUs.
+
+!!! note "GPU request flag"
+    Clusters differ. `--gres=gpu:1` is common; some need `--gpus=1` or a
+    specific gres name (`--gres=gpu:a100:1`). Put whatever `sbatch` flag your
+    cluster needs in `slurm_extra`.
+
+## 6. Monitor
+
+- **Snakemake** prints each job as it submits/finishes and a `X of Y steps`
+  counter.
+- **SLURM**: `squeue --me` shows your queued/running jobs (`smk-segment`, …);
+  logs land where your profile/cluster sends them.
+- **patchworks** logs (`processing tile k/N`, ETA) are inside each job's stdout.
+
+## 7. Output
+
+Everything is under `work_dir`:
+
+```text
+results/
+  image.zarr/                 # converted, pyramidal OME-ZARR
+  image.zarr/labels/<name>/   # the segmentation (multi-scale, calibrated)
+```
+
+The labels live **inside** the image store. View image + labels together:
+
+```python
+from patchworks.plugins.napari import view_in_napari
+view_in_napari("/scratch/results/image.zarr")   # auto-loads the labels
+```
+
+## 8. Re-running and resuming
+
+Snakemake is resumable — if jobs fail or you cancel, just relaunch the same
+command and it picks up only the missing tiles. To force a clean rerun, delete
+`work_dir` (or the relevant outputs).
+
+## Conda (optional)
+
+To have each rule run in a named conda env instead of the active one, add
+`--use-conda` and point the rules at an env; or activate your env in a SLURM
+prologue. The simplest path is a single shared env that the compute nodes see.
+
+## Troubleshooting
+
+| Symptom | Fix |
+|---------|-----|
+| `snakemake: command not found` | use `python -m snakemake` |
+| Segment jobs pend forever | wrong `slurm_partition`/`slurm_extra` GPU flag for your cluster |
+| `cellpose is not installed` in a job | the job's env lacks `patchworks[cellpose]` |
+| Reading the input fails | install the matching reader (`patchworks[imaris]`/`[bioio]` + a `bioio-*`) |
+| Out of GPU memory | smaller `tile_shape`, or `do_3D: false` |
+| Very slow | confirm GPU is used (`nvidia-smi`); try 2-D or a lower `level` |
+
+## How it works (for the curious)
+
+The rule scripts are thin wrappers over patchworks' public API, so you can build
+the same per-tile distribution yourself:
+
+```python
+from patchworks import (
+    load_ome_zarr, spatial_tiles, create_stage, stage_tile, merge_tile_labels
+)
+from patchworks.plugins.ome_zarr import write_labels
+
+img = load_ome_zarr("image.zarr", channel=0)
+tiles = spatial_tiles(img.shape, tile_shape=(16, 1024, 1024))
+create_stage("stage.zarr", img.shape, (16, 1024, 1024))
+# (distribute these across jobs:)
+for i in range(len(tiles)):
+    stage_tile(img, my_fn, "stage.zarr", i, tile_shape=(16, 1024, 1024), overlap=30)
+merged = merge_tile_labels("stage.zarr", input_component="staged",
+                           write_to="merged.zarr", sequential_labels=True)
+write_labels("image.zarr", merged, name="cells")
+```

mkdocs.yml

@@ -39,6 +39,7 @@ nav:
   - Empty tile skipping: guide/skip_empty.md
   - GPU & distributed: guide/gpu_distributed.md
   - Performance & memory: guide/performance.md
+  - Cluster workflow (Snakemake): guide/snakemake.md
   - OME-ZARR & napari: guide/ome_zarr_napari.md
   - Pitfalls: guide/pitfalls.md
 - Examples:

pyproject.toml

@@ -58,6 +58,8 @@ bioio = [
 imaris = ["imaris-ims-file-reader"]
 # napari enables the interactive viewer plugin.
 napari = ["napari[all]"]
+# workflow runs the Snakemake pipeline (per-tile SLURM jobs across GPUs).
+workflow = ["snakemake>=8", "snakemake-executor-plugin-slurm"]
 dev = ["pytest", "pytest-cov", "scikit-image", "psutil", "tqdm"]
 docs = ["mkdocs-material>=9.0", "mkdocstrings[python]>=0.24"]
 all = [

src/patchworks/__init__.py

@@ -32,6 +32,7 @@
 from ._chunks import auto_overlap, auto_tile_shape, auto_tile_shape_cellpose
 from ._cluster import make_local_cluster
 from ._core import tile_process
+from ._distributed import create_stage, spatial_tiles, stage_tile
 from ._io import estimate_empty_tiles, load_ome_zarr
 from ._merge import merge_tile_labels
 from ._relabel import relabel_sequential_array, relabel_sequential_zarr
@@ -51,4 +52,7 @@
     "make_local_cluster",
     "relabel_sequential_array",
     "relabel_sequential_zarr",
+    "spatial_tiles",
+    "create_stage",
+    "stage_tile",
 ]

src/patchworks/_distributed.py

@@ -0,0 +1,140 @@
+"""Per-tile building blocks for distributed processing.
+
+``tile_process`` runs every tile and merges in one process. To spread tiles
+across separate jobs (e.g. one SLURM GPU job per tile) you need to process a
+*single* tile independently and merge later. These helpers expose exactly that:
+:func:`spatial_tiles` enumerates the tiles, :func:`create_stage` makes the
+shared output store, and :func:`stage_tile` runs ``fn`` on one tile and writes
+it into that store. Stitch the result with
+:func:`patchworks.merge_tile_labels` (or ``zarr_native_merge``).
+"""
+
+from __future__ import annotations
+
+import itertools
+from pathlib import Path
+from typing import Callable, Union
+
+import numpy as np
+import zarr
+
+
+def spatial_tiles(
+    shape: tuple[int, ...], tile_shape: tuple[int, ...]
+) -> list[tuple[slice, ...]]:
+    """Enumerate the tiles covering *shape*, in row-major order.
+
+    Parameters
+    ----------
+    shape : tuple of int
+        Spatial array shape.
+    tile_shape : tuple of int
+        Tile shape.
+
+    Returns
+    -------
+    list of tuple of slice
+        One slice tuple per tile (the same order ``estimate_empty_tiles``'s
+        ``occupancy`` grid uses when ravelled).
+    """
+    grids = [range(0, s, t) for s, t in zip(shape, tile_shape)]
+    return [
+        tuple(
+            slice(o, min(o + t, s))
+            for o, t, s in zip(starts, tile_shape, shape)
+        )
+        for starts in itertools.product(*grids)
+    ]
+
+
+def create_stage(
+    stage_path: Union[str, Path],
+    shape: tuple[int, ...],
+    tile_shape: tuple[int, ...],
+    *,
+    component: str = "staged",
+    dtype=np.int32,
+) -> str:
+    """Create the empty (zero-filled) shared stage store for tiled writes.
+
+    Parameters
+    ----------
+    stage_path : str or Path
+        Destination ``.zarr`` store.
+    shape : tuple of int
+        Full (spatial) array shape.
+    tile_shape : tuple of int
+        Chunk = tile shape (one chunk per tile, so jobs write disjoint files).
+    component : str, optional
+        Array name inside the store (default ``"staged"``).
+    dtype : data-type, optional
+        Label dtype (default ``int32``).
+
+    Returns
+    -------
+    str
+        The stage store path.
+    """
+    root = zarr.open_group(str(stage_path), mode="w")
+    root.create_array(
+        name=component, shape=shape, chunks=tile_shape, dtype=dtype
+    )
+    return str(stage_path)
+
+
+def stage_tile(
+    image,
+    fn: Callable[[np.ndarray], np.ndarray],
+    stage_path: Union[str, Path],
+    index: int,
+    *,
+    tile_shape: tuple[int, ...],
+    overlap: int = 0,
+    component: str = "staged",
+) -> int:
+    """Run *fn* on a single tile and write it into the shared stage store.
+
+    Reads the tile (expanded by *overlap* on every side for boundary context),
+    runs *fn*, trims the halo back off, and writes the result to the tile's
+    disjoint chunk of ``stage_path/component`` — so many of these can run
+    concurrently (one per job) without conflicts.
+
+    Parameters
+    ----------
+    image : array-like
+        The full image (dask/zarr/NumPy), indexable by slices.
+    fn : callable
+        ``(ndarray) -> ndarray`` returning integer labels of the same shape.
+    stage_path : str or Path
+        Stage store created by :func:`create_stage`.
+    index : int
+        Tile index into :func:`spatial_tiles`.
+    tile_shape : tuple of int
+        Tile shape (must match the stage store's chunks).
+    overlap : int, optional
+        Halo added on every side before calling *fn*.
+    component : str, optional
+        Array name inside the stage store.
+
+    Returns
+    -------
+    int
+        The processed tile *index*.
+    """
+    shape = image.shape
+    sl = spatial_tiles(shape, tile_shape)[index]
+    expanded, trims = [], []
+    for s, dim in zip(sl, shape):
+        lo = max(0, s.start - overlap)
+        hi = min(dim, s.stop + overlap)
+        expanded.append(slice(lo, hi))
+        trims.append((s.start - lo, hi - s.stop))
+    block = np.asarray(image[tuple(expanded)])
+    out = np.asarray(fn(block))
+    sel = tuple(
+        slice(left, out.shape[i] - right)
+        for i, (left, right) in enumerate(trims)
+    )
+    dst = zarr.open_group(str(stage_path), mode="r+")[component]
+    dst[sl] = out[sel].astype(dst.dtype)
+    return index