patchworks 2.6.3__tar.gz → 2.6.5__tar.gz

{patchworks-2.6.3 → patchworks-2.6.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: patchworks
-Version: 2.6.3
+Version: 2.6.5
 Summary: Tiled processing of arbitrarily large images with globally consistent labels
 Project-URL: Homepage, https://github.com/imcf/patchworks
 Project-URL: Issues, https://github.com/imcf/patchworks/issues

{patchworks-2.6.3 → patchworks-2.6.5}/docs/guide/snakemake.md

@@ -398,12 +398,21 @@ results/image.zarr/labels/nuclei_labels/
 results/image.zarr/labels/cyto_labels/
 ```
 
-!!! tip "Keep `tile_shape` (and `level`) identical across configs"
-    Different segmentations of the same image can use different `channel` and
-    `cellpose:` settings freely, but keep `tile_shape`/`level` the same across
+!!! tip "Keep `level` identical across configs; `tile_shape` usually takes care of itself"
+    Different segmentations of the same image can use different `channel`
+    and `cellpose:` settings freely, but keep `level` the same across
     configs — the label arrays then share the exact same chunk layout, which
     [`label_relations()`](label_relations.md) requires.
 
+    `tile_shape` is a little more forgiving under `run_multi.py`: leaving it
+    on `"auto"` (the default) would normally size a `nuclei_channel` config's
+    tile smaller than a single-channel one (the sizer charges per channel),
+    breaking that same requirement — `run_multi.py` detects this and pins
+    every config to the smallest computed tile automatically, so mixed
+    single-/two-channel configs under `"auto"` just work. Setting an
+    explicit `tile_shape` yourself still requires it to be identical across
+    configs, same as before.
+
 See [Relating labels across segmentations](label_relations.md) for what
 `label_relations()` returns and how to save it yourself — the cluster
 workflow's own automation is below.
@@ -451,6 +460,18 @@ the GPU partition stays busy instead of idling through every config's
 `prepare` and multi-hour `merge` in turn. A config that fails does **not**
 abort the others; you get a per-config status and a non-zero exit.
 
+!!! tip "The relate step runs on the cluster too, under `multi-slurm`"
+    `label_relations()` streams every chunk of two full-resolution label
+    volumes — real CPU/IO work, not orchestration. Under `multi-slurm` it is
+    submitted as its own `srun` job (`scripts/relate.py`) instead of running
+    in the driver process on the login node, the same fix already applied to
+    the occupancy map. Tune its allocation with `--relate-partition`,
+    `--relate-mem`, `--relate-cpus` and `--relate-time` (defaults: `scicore`,
+    `32G`, `8`, `180` minutes) — these are wide-margin guesses, not measured
+    numbers, so raise them for a very large or very object-dense pair. Under
+    plain `multi` (no `--profile`), it still runs locally, in-process, as
+    before.
+
 !!! tip "After a killed run"
     Snakemake only releases its lock on a clean exit, so a run that was killed
     (Ctrl-C, an SSH drop, an OOM) leaves the directory locked. Each phase has

{patchworks-2.6.3 → patchworks-2.6.5}/tests/test_run_multi.py

@@ -160,28 +160,114 @@ def test_occupancy_is_not_rebuilt_by_the_driver():
     assert "occupancy.zarr" in src
 
 
-def test_auto_tile_shape_with_a_lone_nuclei_channel_is_refused():
-    """Matching `tile_shape` *values* are not enough when one config is 2-ch.
+def test_relate_is_submitted_via_slurm_under_profile():
+    """The relate step must never run in-process on the submit host.
 
-    "auto" == "auto" passes the plain equality check, but the sizer charges
-    per channel, so the nuclei_channel config gets a smaller tile. The label
-    groups then disagree on chunk layout and label_relations raises -- after
-    every segmentation has already run, which is the expensive way to find out.
+    Same failure mode as the occupancy map: label_relations() streams every
+    chunk of two full-resolution label volumes. A prior fix moved the map
+    build off the login node; the relate step made the identical mistake and
+    hung there for the same reason until this fix.
+    """
+    src = (_workflow_dir() / "scripts" / "run_multi.py").read_text()
+    assert "from relate import run_relations" in src
+    assert '"srun"' in src
+    assert "label_relations(" not in src
+
+
+def test_relate_script_has_the_real_bookkeeping():
+    """relate.py must be the actual implementation, not a stub.
+
+    Submitting the wrong (or a trimmed-down) script would silently produce a
+    workbook missing the unmatched-label rows the docstring promises.
+    """
+    src = (_workflow_dir() / "scripts" / "relate.py").read_text()
+    assert "def run_relations(" in src
+    assert "label_relations" in src
+    assert "openpyxl" in src
+
+
+def test_mixed_nuclei_channel_auto_passes_validation():
+    """A channel-count mismatch under `tile_shape: "auto"` is no longer
+
+    refused at validation time -- it's resolved automatically instead (see
+    `_resolve_shared_tile_shape`), which needs the converted image's real
+    shape/dtype and so can only run after phase A, not from
+    `_validate_configs()`. This used to `sys.exit` here; asserting that
+    would now be testing the wrong layer.
     """
     paths = [Path("a.yaml"), Path("b.yaml")]
     base = {"work_dir": "/w", "tile_shape": "auto", "level": 0}
 
-    bad = [
+    mixed = [
         {**base, "label_name": "a", "channel": 0, "nuclei_channel": 1},
         {**base, "label_name": "b", "channel": 2},
     ]
-    with pytest.raises(SystemExit):
-        _validate_configs(paths, bad)
+    assert _validate_configs(paths, mixed) == "/w"
 
     # Same pair with one explicit shape is fine: both get that tile.
-    pinned = [{**c, "tile_shape": [16, 512, 512]} for c in bad]
+    pinned = [{**c, "tile_shape": [16, 512, 512]} for c in mixed]
     assert _validate_configs(paths, pinned) == "/w"
 
     # And "auto" is fine when every config carries the same channel count.
-    both = [{**bad[0]}, {**bad[1], "nuclei_channel": 3}]
+    both = [{**mixed[0]}, {**mixed[1], "nuclei_channel": 3}]
     assert _validate_configs(paths, both) == "/w"
+
+
+def test_resolve_shared_tile_shape_pins_the_smallest_candidate(
+    tmp_path, monkeypatch
+):
+    """The shared tile must be the tightest of every config's own budget.
+
+    A larger tile than some config's own sizer output would ask that config
+    for more memory than its settings were judged to need -- only the
+    smallest candidate is safe for every config at once.
+    """
+    import run_multi
+
+    class _FakeImage:
+        shape = (10, 100, 100)
+        dtype = "uint16"
+
+    calls = []
+
+    def _fake_load_ome_zarr(store, *, channel, level):
+        calls.append((store, channel, level))
+        return _FakeImage()
+
+    # One tile per config, matched up by call order (channel 0 then 1).
+    fake_tiles = [(8, 64, 64), (4, 32, 32)]
+
+    def _fake_sizer_cellpose(shape, dtype, **kwargs):
+        return fake_tiles[len(calls) - 1]
+
+    monkeypatch.setattr(
+        "patchworks.load_ome_zarr", _fake_load_ome_zarr, raising=False
+    )
+    monkeypatch.setattr(
+        "patchworks.auto_tile_shape_cellpose",
+        _fake_sizer_cellpose,
+        raising=False,
+    )
+
+    seg_cfgs = [
+        {
+            "channel": 0,
+            "level": 0,
+            "method": "cellpose",
+            "cellpose": {"do_3D": True, "gpu": True},
+        },
+        {
+            "channel": 1,
+            "level": 0,
+            "nuclei_channel": 2,
+            "method": "cellpose",
+            "cellpose": {"do_3D": True, "gpu": True},
+        },
+    ]
+    out = run_multi._resolve_shared_tile_shape(
+        seg_cfgs, "/w/image.zarr", str(tmp_path)
+    )
+    assert out == tmp_path / ".multi_tile_shape.generated.yaml"
+    written = yaml.safe_load(out.read_text())
+    assert written == {"tile_shape": [4, 32, 32]}  # the smaller candidate
+    assert len(calls) == 2

patchworks-2.6.5/workflow/scripts/relate.py

@@ -0,0 +1,149 @@
+"""Compute label_relations for configured pairs and write .xlsx.
+
+Split out of run_multi.py so this step can be submitted as its own SLURM job
+instead of running in-process on the login node. It streams every chunk of
+two full-resolution label volumes -- real CPU/IO work, not orchestration --
+same reasoning as the occupancy-map fix (see run_multi.py's phase A comment).
+
+Usage (called by run_multi.py under --profile, but also runnable standalone,
+e.g. under srun):
+    python scripts/relate.py --work-dir /path/to/work_dir \
+        --image-store /path/to/work_dir/image.zarr \
+        --relations '[{"a": "nuclei_labels", "b": "cyto_labels", "output": "nuclei_to_cyto.xlsx"}]'
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+
+def _label_ids(image_store: str, name: str) -> list[int]:
+    """Ids present in a label image, without scanning the volume.
+
+    The merge writes n_objects/sequential_labels into the label group's
+    attrs precisely so consumers don't have to re-derive the id set; the ids
+    are 1..n_objects by construction. Fall back to a full scan only for a
+    label group written before those attrs existed.
+    """
+    import dask.array as da
+    import zarr
+
+    attrs = dict(zarr.open_group(f"{image_store}/labels/{name}").attrs)
+    if attrs.get("sequential_labels") and attrs.get("n_objects") is not None:
+        return list(range(1, int(attrs["n_objects"]) + 1))
+    print(
+        f"[relate] {name}: no n_objects attr, falling back to a full scan "
+        "for its id set",
+        flush=True,
+    )
+    arr = da.from_zarr(image_store, component=f"labels/{name}/0")
+    return sorted(int(x) for x in da.unique(arr[arr > 0]).compute())
+
+
+def run_relations(
+    work_dir: str, image_store: str, relations: list[dict]
+) -> None:
+    """Compute and write every configured relation pair as an .xlsx workbook.
+
+    Parameters
+    ----------
+    work_dir : str
+        Directory relation workbooks are written into (a relation's
+        ``output``, when relative, resolves against this).
+    image_store : str
+        The shared ``image.zarr`` holding every config's ``labels/<name>``.
+    relations : list of dict
+        Each ``{"a": ..., "b": ..., "output": ...}`` (``output`` optional,
+        defaults to ``<a>_to_<b>.xlsx``), matching ``multi.yaml``'s
+        ``relations:`` list.
+    """
+    import dask.array as da
+    import openpyxl
+
+    from patchworks import label_relations
+
+    for rel in relations:
+        a_name, b_name = rel["a"], rel["b"]
+        out_path = Path(work_dir) / rel.get(
+            "output", f"{a_name}_to_{b_name}.xlsx"
+        )
+        print(f"[relate] relating {a_name} -> {b_name} …", flush=True)
+        a = da.from_zarr(image_store, component=f"labels/{a_name}/0")
+        b = da.from_zarr(image_store, component=f"labels/{b_name}/0")
+        table = label_relations(a, b)
+
+        # label_relations() only returns a-objects that touch a b-object.
+        # Pull the full id sets so unmatched a-objects (zero overlap) and
+        # b-objects with no matches at all still get a row -- otherwise
+        # they'd silently vanish instead of counting as zero.
+        a_ids = _label_ids(image_store, a_name)
+        b_ids = _label_ids(image_store, b_name)
+
+        per_b = {b_id: {"count": 0, "overlap_voxels": 0} for b_id in b_ids}
+        for m in table.values():
+            agg = per_b.get(m["match"])
+            if agg is not None:
+                agg["count"] += 1
+                agg["overlap_voxels"] += m["overlap_voxels"]
+
+        wb = openpyxl.Workbook()
+        ws_a = wb.active
+        ws_a.title = a_name[:31]  # Excel sheet-name length limit
+        ws_a.append(
+            [
+                f"{a_name}_id",
+                f"{b_name}_id",
+                "overlap_voxels",
+                "overlap_fraction",
+            ]
+        )
+        for a_id in a_ids:
+            m = table.get(a_id)
+            if m is None:
+                ws_a.append([a_id, None, 0, 0])  # no overlap -- still counted
+            else:
+                ws_a.append(
+                    [
+                        a_id,
+                        m["match"],
+                        m["overlap_voxels"],
+                        m["overlap_fraction"],
+                    ]
+                )
+
+        ws_b = wb.create_sheet(title=b_name[:31])
+        ws_b.append(
+            [f"{b_name}_id", f"{a_name}_count", "total_overlap_voxels"]
+        )
+        for b_id in b_ids:
+            agg = per_b[b_id]
+            ws_b.append([b_id, agg["count"], agg["overlap_voxels"]])
+
+        wb.save(out_path)
+        print(
+            f"[relate] wrote {out_path} "
+            f"({len(a_ids)} {a_name}, {len(b_ids)} {b_name})",
+            flush=True,
+        )
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--work-dir", required=True)
+    parser.add_argument("--image-store", required=True)
+    parser.add_argument(
+        "--relations",
+        required=True,
+        help=(
+            "JSON list of {a, b, output} dicts, matching multi.yaml's "
+            "relations:"
+        ),
+    )
+    args = parser.parse_args()
+    run_relations(args.work_dir, args.image_store, json.loads(args.relations))
+
+
+if __name__ == "__main__":
+    main()

{patchworks-2.6.3 → patchworks-2.6.5}/workflow/scripts/run_multi.py

@@ -20,12 +20,16 @@ Once all segmentations succeed, each configured relation pair is computed via
 patchworks.label_relations and written as an Excel workbook in work_dir,
 with two sheets: one row per a-object (unmatched ones included, with an
 empty b-id and zeros) and one row per b-object (a-object count + total
-overlap, including b-objects with zero matches).
+overlap, including b-objects with zero matches). Under --profile, this runs
+as a submitted SLURM job (see scripts/relate.py) rather than in-process here
+-- same reasoning as the occupancy-map fix: it streams entire label volumes,
+which is real work, not orchestration, and does not belong on the login node.
 """
 
 from __future__ import annotations
 
 import argparse
+import json
 import re
 import subprocess
 import sys
@@ -50,6 +54,7 @@ def _snakemake_cmd(
     extra: list[str] | None = None,
     jobname_prefix: str | None = None,
     common: Path | None = None,
+    extra_configfiles: list[Path] | None = None,
 ) -> list[str]:
     """Build one snakemake invocation.
 
@@ -62,10 +67,17 @@ def _snakemake_cmd(
     settings every config shares -- the input, the work_dir, everything
     ``convert`` reads -- live in one file and the per-config file carries only
     what actually differs.
+
+    *extra_configfiles*, when given, are appended after *configfile* and so
+    win over both it and *common* -- used to pin a driver-computed value
+    (e.g. a resolved ``tile_shape``) across every config without editing any
+    config file on disk.
     """
     configfiles = [str(configfile.resolve())]
     if common is not None:
         configfiles.insert(0, str(common.resolve()))
+    if extra_configfiles:
+        configfiles += [str(p.resolve()) for p in extra_configfiles]
     cmd = [
         "snakemake",
         "-s",
@@ -230,21 +242,12 @@ def _validate_configs(paths: list[Path], cfgs: list[dict]) -> str:
 
     # `tile_shape: "auto"` is identical as a *value* across configs while
     # producing different tiles: the sizer charges per channel, so a config
-    # with nuclei_channel gets a smaller one. The label groups then disagree on
-    # chunk layout and label_relations raises -- after every segmentation has
-    # run. Matching values are not enough here, so check the inputs that feed
-    # the sizer instead.
-    if {repr(cfg.get("tile_shape", "auto")) for cfg in cfgs} == {repr("auto")}:
-        if len({cfg.get("nuclei_channel") is not None for cfg in cfgs}) != 1:
-            problems.append(
-                'tile_shape: "auto" sizes a nuclei_channel config smaller '
-                "(a tile carries two channels), so the label groups would "
-                "not share a chunk layout and label_relations would fail "
-                f"after every segmentation had run; got "
-                f"{_spread('nuclei_channel')}. Set one explicit tile_shape in "
-                "the file `common:` points at, sized for the two-channel "
-                "config."
-            )
+    # with nuclei_channel gets a smaller one. The label groups would then
+    # disagree on chunk layout and label_relations would raise -- after every
+    # segmentation had run. That used to be a hard error asking for a manual
+    # explicit tile_shape; main() now resolves and pins one automatically
+    # (see _resolve_shared_tile_shape), once the converted image exists to
+    # size against, so there is nothing to check here anymore.
 
     # Phase A converts once, from the first config. Anything `convert` reads
     # out of a later config is therefore silently ignored -- someone setting
@@ -288,6 +291,92 @@ def _validate_configs(paths: list[Path], cfgs: list[dict]) -> str:
     return work_dirs.pop()
 
 
+def _resolve_shared_tile_shape(
+    seg_cfgs: list[dict], image_store: str, work_dir: str
+) -> Path:
+    """Auto-size ``tile_shape`` once, shared across every config.
+
+    ``tile_shape: "auto"`` resolves differently per config when
+    ``nuclei_channel`` differs -- the sizer charges per channel, so a
+    two-channel config gets a smaller tile. Left alone, the label groups
+    would end up with different chunk layouts and ``label_relations`` would
+    raise, after every segmentation had already run.
+
+    Computes what each config's own settings (channel count, ``do_3D``,
+    ``diameter``, GPU budget) would actually produce, then pins every config
+    to the *smallest* of them by voxel count -- the most memory-constrained
+    case, and safe for every config since a smaller tile only ever asks for
+    less memory than that config's own budget allows, never more. Configs
+    can therefore use different ``do_3D``/``diameter``/channel-count
+    settings and still end up with one shared, valid ``tile_shape``.
+
+    Only called once the converted image exists (the sizer needs its real
+    shape/dtype), so this runs from ``main()`` after phase A, not from
+    ``_validate_configs()``.
+
+    Returns
+    -------
+    Path
+        A generated one-key YAML file (``tile_shape: [...]``), meant to be
+        passed as an ``extra_configfiles`` entry to ``_snakemake_cmd`` so it
+        overrides every config's own (or common's) ``tile_shape`` value.
+    """
+    from functools import partial
+
+    import numpy as np
+    from patchworks import (
+        auto_tile_shape,
+        auto_tile_shape_cellpose,
+        load_ome_zarr,
+    )
+
+    candidates = []
+    for cfg in seg_cfgs:
+        image = load_ome_zarr(
+            image_store, channel=cfg["channel"], level=int(cfg.get("level", 0))
+        )
+        gpu_gb = cfg.get("gpu_memory_gb")
+        gpu_bytes = int(gpu_gb * 1024**3) if gpu_gb else None
+        n_channels = 2 if cfg.get("nuclei_channel") is not None else 1
+        method = cfg.get("method", "cellpose")
+        if method == "cellpose":
+            cp = cfg["cellpose"]
+            sizer = partial(
+                auto_tile_shape_cellpose,
+                do_3D=cp.get("do_3D", False),
+                use_gpu=cp.get("gpu", True),
+                diameter=cp.get("diameter"),
+                gpu_memory=gpu_bytes,
+                n_channels=n_channels,
+            )
+        else:
+            sizer = partial(
+                auto_tile_shape,
+                use_gpu=gpu_bytes is not None,
+                gpu_memory=gpu_bytes,
+                n_channels=n_channels,
+            )
+        candidates.append(tuple(int(x) for x in sizer(image.shape, image.dtype)))
+
+    tile_shape = min(candidates, key=lambda t: int(np.prod(t)))
+    print(
+        f'[run_multi] tile_shape: "auto" resolves differently across '
+        f"configs (nuclei_channel differs); pinning every config to the "
+        f"smallest computed tile {list(tile_shape)} so the label arrays "
+        f"share a chunk layout. Candidates were {[list(c) for c in candidates]}.",
+        flush=True,
+    )
+
+    override_path = Path(work_dir) / ".multi_tile_shape.generated.yaml"
+    override_path.write_text(
+        "# Generated by run_multi.py -- pins tile_shape across configs so\n"
+        "# label_relations sees matching chunk layouts. Safe to delete; it\n"
+        "# is regenerated on the next multi-config run.\n"
+        f"tile_shape: {list(tile_shape)}\n"
+    )
+    return override_path
+
+
 def _resolve(workflow_dir: Path, path_str: str) -> Path:
     path = Path(path_str)
     return path if path.is_absolute() else workflow_dir / path
@@ -337,6 +426,28 @@ def main() -> None:
             "otherwise look identical: no email either way."
         ),
     )
+    parser.add_argument(
+        "--relate-partition",
+        default="scicore",
+        help="SLURM partition for the relate step under --profile (default: scicore)",
+    )
+    parser.add_argument(
+        "--relate-mem",
+        default="32G",
+        help="srun --mem for the relate step under --profile (default: 32G)",
+    )
+    parser.add_argument(
+        "--relate-cpus",
+        type=int,
+        default=8,
+        help="srun --cpus-per-task for the relate step under --profile (default: 8)",
+    )
+    parser.add_argument(
+        "--relate-time",
+        type=int,
+        default=180,
+        help="srun --time in minutes for the relate step under --profile (default: 180)",
+    )
     args = parser.parse_args()
 
     workflow_dir = Path(__file__).resolve().parent.parent
@@ -427,6 +538,24 @@ def main() -> None:
         )
         sys.exit(rc)
 
+    # tile_shape: "auto" resolves differently per config when nuclei_channel
+    # differs (the sizer charges per channel) -- pin every config to one
+    # shared, computed value so the label arrays end up with matching chunk
+    # layouts. Needs the just-converted image, so this can only happen here,
+    # not in _validate_configs(). Dry runs never reach a real image.zarr.
+    tile_override = None
+    if not args.dry_run and {
+        cfg.get("tile_shape", "auto") for cfg in seg_cfgs
+    } == {"auto"}:
+        channel_counts = {
+            2 if cfg.get("nuclei_channel") is not None else 1
+            for cfg in seg_cfgs
+        }
+        if len(channel_counts) > 1:
+            tile_override = _resolve_shared_tile_shape(
+                seg_cfgs, image_store, work_dir
+            )
+
     # Phase B: the segmentations touch disjoint files under
     # work_dir/<label_name>/, so run them together and let the GPU partition
     # stay busy instead of idling through each config's prepare and merge.
@@ -444,6 +573,7 @@ def main() -> None:
             # Names the config in squeue, so concurrent runs are tellable apart.
             jobname_prefix=slurm_jobname_prefix(cfg["label_name"]),
             common=common_path,
+            extra_configfiles=[tile_override] if tile_override else None,
         )
         print(f"[run_multi] $ {' '.join(cmd)}", flush=True)
         procs.append((cfg_path.name, subprocess.Popen(cmd, cwd=workflow_dir)))
@@ -467,96 +597,47 @@ def main() -> None:
     if args.dry_run or not relations:
         return
 
-    import dask.array as da
-    import openpyxl
-    import zarr
-
-    from patchworks import label_relations
-
-    def _label_ids(name: str) -> list[int]:
-        """Ids present in a label image, without scanning the volume.
-
-        The merge writes n_objects/sequential_labels into the label group's
-        attrs precisely so consumers don't have to re-derive the id set; the
-        ids are 1..n_objects by construction. Fall back to the full scan only
-        for a label group written before those attrs existed -- that scan runs
-        here on the login node, so it is worth avoiding.
-        """
-        attrs = dict(zarr.open_group(f"{image_store}/labels/{name}").attrs)
-        if (
-            attrs.get("sequential_labels")
-            and attrs.get("n_objects") is not None
-        ):
-            return list(range(1, int(attrs["n_objects"]) + 1))
-        print(
-            f"[run_multi] {name}: no n_objects attr, falling back to a full "
-            "scan for its id set",
-            flush=True,
-        )
-        arr = da.from_zarr(image_store, component=f"labels/{name}/0")
-        return sorted(int(x) for x in da.unique(arr[arr > 0]).compute())
+    if args.profile:
+        # Real CPU/IO work -- tens of thousands of zarr chunk reads for a
+        # full-resolution label volume -- not orchestration, so (like the
+        # occupancy map) it does not belong in this driver process on the
+        # login node. Submit it as its own job instead.
+        cmd = [
+            "srun",
+            "--partition",
+            args.relate_partition,
+            "--mem",
+            args.relate_mem,
+            "--cpus-per-task",
+            str(args.relate_cpus),
+            "--time",
+            str(args.relate_time),
+            "--job-name",
+            "pw-relate",
+            sys.executable,
+            str(workflow_dir / "scripts" / "relate.py"),
+            "--work-dir",
+            work_dir,
+            "--image-store",
+            image_store,
+            "--relations",
+            json.dumps(relations),
+        ]
+        rc = _run(cmd, workflow_dir)
+        if rc != 0:
+            print(
+                f"[run_multi] ERROR: relate step failed (exit {rc}). "
+                "Segmentations already succeeded -- only the relation "
+                "workbook(s) are missing. Re-run with the same --config to "
+                "retry just this step.",
+                file=sys.stderr,
+            )
+            sys.exit(rc)
+        return
 
-    for rel in relations:
-        a_name, b_name = rel["a"], rel["b"]
-        out_path = Path(work_dir) / rel.get(
-            "output", f"{a_name}_to_{b_name}.xlsx"
-        )
-        print(f"[run_multi] relating {a_name} -> {b_name} …", flush=True)
-        a = da.from_zarr(image_store, component=f"labels/{a_name}/0")
-        b = da.from_zarr(image_store, component=f"labels/{b_name}/0")
-        table = label_relations(a, b)
-
-        # label_relations() only returns a-objects that touch a b-object.
-        # Pull the full id sets so unmatched a-objects (zero overlap) and
-        # b-objects with no matches at all still get a row -- otherwise
-        # they'd silently vanish instead of counting as zero.
-        a_ids = _label_ids(a_name)
-        b_ids = _label_ids(b_name)
-
-        per_b = {b_id: {"count": 0, "overlap_voxels": 0} for b_id in b_ids}
-        for m in table.values():
-            agg = per_b.get(m["match"])
-            if agg is not None:
-                agg["count"] += 1
-                agg["overlap_voxels"] += m["overlap_voxels"]
-
-        wb = openpyxl.Workbook()
-        ws_a = wb.active
-        ws_a.title = a_name[:31]  # Excel sheet-name length limit
-        ws_a.append(
-            [
-                f"{a_name}_id",
-                f"{b_name}_id",
-                "overlap_voxels",
-                "overlap_fraction",
-            ]
-        )
-        for a_id in a_ids:
-            m = table.get(a_id)
-            if m is None:
-                ws_a.append([a_id, None, 0, 0])  # no overlap -- still counted
-            else:
-                ws_a.append(
-                    [
-                        a_id,
-                        m["match"],
-                        m["overlap_voxels"],
-                        m["overlap_fraction"],
-                    ]
-                )
-
-        ws_b = wb.create_sheet(title=b_name[:31])
-        ws_b.append([f"{b_name}_id", f"{a_name}_count", "total_overlap_voxels"])
-        for b_id in b_ids:
-            agg = per_b[b_id]
-            ws_b.append([b_id, agg["count"], agg["overlap_voxels"]])
-
-        wb.save(out_path)
-        print(
-            f"[run_multi] wrote {out_path} "
-            f"({len(a_ids)} {a_name}, {len(b_ids)} {b_name})",
-            flush=True,
-        )
+    from relate import run_relations
+
+    run_relations(work_dir, image_store, relations)
 
 
 if __name__ == "__main__":