perf(workflow): ⚡ segment a batch of tiles per job to amortize model load

Every segment job paid CUDA init plus a Cellpose weight load -- tens of
seconds -- before touching its first voxel, and at ~1250 one-tile jobs that
startup dominated the cheap tiles. The plugin already caches the model per
process (_model_cache), but a one-tile job never got a second tile to serve
it.

prepare now groups the occupied tiles into batches of `tiles_per_job` and the
segment rule scatters over batches instead of tiles, building the segmentation
function once and looping. Tiles in a batch run sequentially, so they share
the cached model and never contend for the GPU; batches still write disjoint
chunks and run in parallel across GPUs.

Snakemake's `group:` directive was the obvious route but is wrong here:
GroupResources.basic_layered sums integer resources across jobs in one
toposort layer, and all segment tiles sit in a single layer, so a group would
have requested N x mem_mb / N x cpus_per_task and run N Cellpose evaluations
concurrently on the one GPU it was granted.

Defaults to 1, which reproduces the previous one-job-per-tile behaviour. The
markers move from seg/<tile>.done to seg/<batch>.done and now carry their tile
list as JSON.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

Author: lguerard

workflow/config/config.yaml

@@ -36,6 +36,13 @@ overlap: 30 # halo (≈ one object diameter); scalar, or per-axis [z, y, x]
 #             all of it wasted z. patchworks.auto_overlap(d, voxel_size=...)
 #             computes the per-axis values for you.
 skip_empty: true # skip background tiles
+# Tiles per SLURM job. Each job pays CUDA init + a Cellpose weight load (tens
+# of seconds) once, then reuses the cached model for the rest of its batch, so
+# at 1 that startup can dominate a fast tile. Tiles in a batch run
+# SEQUENTIALLY, so the job's wall time is roughly N x per-tile time: size N
+# from a measured tile (`seff <jobid>`) and keep it inside the profile's
+# runtime/QOS ceiling. A failed job retries the whole batch.
+tiles_per_job: 1
 empty_threshold: null # null → Otsu; or a number
 
 # ---- segmentation -----------------------------------------------------------

workflow/config/config_cilia.yaml

@@ -28,6 +28,9 @@ gpu_memory_gb: null
 # than pushing the halo up against the tile depth.
 overlap: [8, 30, 30]
 skip_empty: true
+# Tiles per SLURM job; they run sequentially and share one decon/GPU context.
+# Raise once you know a tile's runtime -- job wall time is ~N x per-tile time.
+tiles_per_job: 1
 empty_threshold: null
 
 method: "custom"

workflow/config/config_cyto.yaml

@@ -20,6 +20,9 @@ gpu_memory_gb: null
 # z-halo only needs to cover one cell in z, not one cell in x/y.
 overlap: [4, 30, 30]
 skip_empty: true
+# Tiles per SLURM job; they run sequentially and share one model load. Raise
+# once you know a tile's runtime -- job wall time is ~N x per-tile time.
+tiles_per_job: 1
 empty_threshold: null
 
 method: "cellpose"

workflow/config/config_nuclei.yaml

@@ -20,6 +20,9 @@ gpu_memory_gb: null
 # z-halo only needs to cover one nucleus in z, not one nucleus in x/y.
 overlap: [4, 30, 30]
 skip_empty: true
+# Tiles per SLURM job; they run sequentially and share one model load. Raise
+# once you know a tile's runtime -- job wall time is ~N x per-tile time.
+tiles_per_job: 1
 empty_threshold: null
 
 method: "cellpose"

workflow/rules/common.smk

@@ -36,8 +36,13 @@ STEPLOG = f"{LOGS}/steps.log"
 MODEL_OK = f"{RUN}/model.ready"
 
 
-def occupied_done(wildcards):
-    """Per-tile markers for the occupied tiles (resolved after the checkpoint)."""
+def batch_done(wildcards):
+    """Per-batch markers for the segment jobs (resolved after the checkpoint).
+
+    prepare groups the occupied tiles into batches of ``tiles_per_job``; one
+    marker is produced per batch, not per tile, so the fan-in shrinks with the
+    batch size.
+    """
     tiles = checkpoints.prepare.get().output.tiles
-    occupied = json.loads(Path(tiles).read_text())["occupied"]
-    return [f"{RUN}/seg/{i}.done" for i in occupied]
+    manifest = json.loads(Path(tiles).read_text())
+    return [f"{RUN}/seg/{i}.done" for i in range(len(manifest["batches"]))]

workflow/rules/merge.smk

@@ -2,7 +2,7 @@
 
 rule merge:
     input:
-        occupied_done,
+        batch_done,
     output:
         touch(f"{RUN}/labels.done"),
     log:

workflow/rules/segment.smk

@@ -28,15 +28,21 @@ checkpoint prepare:
 
 
 rule segment:
-    """Segment one tile on a GPU and write it into the stage store."""
+    """Segment one batch of tiles on a GPU and write them into the stage store.
+
+    A batch is `tiles_per_job` tiles (see config), processed sequentially in
+    one process so they share a single CUDA init and a single Cellpose model
+    load. Batches write disjoint chunks, so any number of them run in
+    parallel across GPUs.
+    """
     input:
         tiles=TILES,
         stage=STAGE_OK,
         image=IMAGE_OK,
         model=MODEL_OK,
     output:
-        f"{RUN}/seg/{{index}}.done",
+        f"{RUN}/seg/{{batch}}.done",
     log:
-        f"{LOGS}/segment/{{index}}.log",
+        f"{LOGS}/segment/{{batch}}.log",
     script:
         "../scripts/segment_tile.py"

workflow/scripts/prepare_tiles.py

@@ -102,6 +102,16 @@
     occ = info["occupancy"].ravel()  # row-major, matches spatial_tiles
     occupied = [i for i in range(len(tiles)) if occ[i]]
 
+# One SLURM job per tile means every tile pays CUDA init + a model load before
+# its first voxel. Batching amortizes that over `tiles_per_job` tiles, which
+# run sequentially in one process so they share the cached model and never
+# contend for the GPU. 1 = the old one-job-per-tile behaviour.
+tiles_per_job = max(1, int(cfg.get("tiles_per_job", 1)))
+batches = [
+    occupied[i : i + tiles_per_job]
+    for i in range(0, len(occupied), tiles_per_job)
+]
+
 create_stage(stage_path(work_dir, label_name), image.shape, tile_shape)
 
 Path(work_dir, label_name, "tiles.json").write_text(
@@ -111,8 +121,13 @@
             "overlap": list(overlap),
             "n_tiles": len(tiles),
             "occupied": occupied,
+            "tiles_per_job": tiles_per_job,
+            "batches": batches,
         },
         indent=2,
     )
 )
-print(f"[patchworks] {len(occupied)}/{len(tiles)} tiles to segment")
+print(
+    f"[patchworks] {len(occupied)}/{len(tiles)} tiles to segment "
+    f"in {len(batches)} job(s) of up to {tiles_per_job}"
+)

workflow/scripts/segment_tile.py

@@ -1,31 +1,45 @@
-"""Snakemake script: segment ONE tile and write it to the shared stage.
+"""Snakemake script: segment one BATCH of tiles into the shared stage.
 
-Scattered over tile indices, so each tile is its own SLURM job and many GPUs
-run in parallel. Each job writes a disjoint chunk of the stage store.
+Scattered over batches, so several batches run as separate SLURM jobs and many
+GPUs work in parallel. Within a batch the tiles run sequentially in this one
+process: CUDA is initialised once and the Cellpose model is loaded once (see
+``patchworks.plugins.cellpose._model_cache``), instead of once per tile. Each
+batch writes disjoint chunks of the stage store, so batches never collide.
 """
 
+import json
+
 from patchworks import stage_tile
 
 from _pw import build_fn, load_tiles_json, open_image, stage_path, start_log
 
 start_log(snakemake.log[0])  # noqa: F821
 cfg = snakemake.config  # noqa: F821
-index = int(snakemake.wildcards.index)  # noqa: F821
+batch = int(snakemake.wildcards.batch)  # noqa: F821
 work_dir = cfg["work_dir"]
 label_name = cfg.get("label_name", "labels")
 
 manifest = load_tiles_json(snakemake.input.tiles)  # noqa: F821
 image = open_image(work_dir, cfg["channel"], cfg["level"])
+indices = manifest["batches"][batch]
+
+# Built once for the whole batch: this is what makes the model load amortize.
+fn = build_fn(cfg)
+stage = stage_path(work_dir, label_name)
+tile_shape = tuple(manifest["tile_shape"])
+
+for index in indices:
+    stage_tile(
+        image,
+        fn,
+        stage,
+        index,
+        tile_shape=tile_shape,
+        # Scalar (older manifests) or per-axis list; stage_tile normalizes both.
+        overlap=manifest["overlap"],
+    )
+    print(f"[patchworks] segmented tile {index}")
 
-stage_tile(
-    image,
-    build_fn(cfg),
-    stage_path(work_dir, label_name),
-    index,
-    tile_shape=tuple(manifest["tile_shape"]),
-    # Scalar (older manifests) or per-axis list; stage_tile normalizes both.
-    overlap=manifest["overlap"],
-)
-
-open(snakemake.output[0], "w").close()  # noqa: F821
-print(f"[patchworks] segmented tile {index}")
+with open(snakemake.output[0], "w") as fh:  # noqa: F821
+    json.dump({"batch": batch, "tiles": indices}, fh)
+print(f"[patchworks] batch {batch}: {len(indices)} tile(s) done")