Distributed processing¶

Split one iterator run across jobs that never talk to each other.

On a cluster, an iterator run is split into jobs — SLURM array tasks, Fractal parallel tasks — that share a filesystem but cannot talk to each other. ngio's partition model is built for exactly that: each job restricts the same iterator to its own share with for_job, and no locks or coordination are ever needed. This tutorial makes the recipe run on a laptop by standing a plain Python loop in for the cluster — one loop iteration per job, where a job is whatever one scheduler task would run.

Step 1: set up¶

from pathlib import Path

from ngio import open_ome_zarr_container
from ngio.utils import download_ome_zarr_dataset

# Download the dataset
download_dir = Path("./data").absolute()
hcs_path = download_ome_zarr_dataset("CardiomyocyteTinyMip", download_dir=download_dir)

# Open the OME-Zarr container. The store is on disk — a distributed run cannot
# use an in-memory store, since each process would write its own private copy.
ome_zarr = open_ome_zarr_container(hcs_path / "B" / "03" / "0")
image = ome_zarr.get_image()

The segmentation function is the watershed pipeline from the stitching tutorial:

import numpy as np
from scipy import ndimage as ndi
from skimage.feature import peak_local_max
from skimage.filters import threshold_otsu
from skimage.morphology import remove_small_objects
from skimage.segmentation import watershed


# Same watershed pipeline as docs/snippets/tutorials/stitching.py; each snippet
# script stands alone, so the function is repeated rather than imported.
def segment(patch: np.ndarray) -> np.ndarray:
    # Smooth → Otsu threshold → distance transform → seeded watershed → cleanup
    smooth = ndi.gaussian_filter(patch, sigma=4)
    mask = smooth > threshold_otsu(smooth)
    distance = ndi.distance_transform_edt(mask)
    coords = peak_local_max(distance, min_distance=20, labels=mask)
    markers = np.zeros(distance.shape, dtype=np.int32)
    markers[tuple(coords.T)] = np.arange(1, len(coords) + 1)
    seg = watershed(-distance, markers, mask=mask).astype(np.uint32)
    return remove_small_objects(seg, max_size=500)

Step 2: check the partition layout¶

A chunk is one atomic write object, so tiles that share an output chunk must travel in the same job: effective parallelism follows the output's chunking, not the tiling. Derive the output label with the defaults and the constraint shows itself:

from ngio.iterators import SegmentationIterator

# Deriving with no arguments inherits the image's chunking: two fat chunks.
label = ome_zarr.derive_label("nuclei_distributed", overwrite=True)
print(f"write granularity: {label.write_granularity}")

iterator = SegmentationIterator(
    image, label, channel_selection="DAPI", axes_order=["y", "x"]
).by_grid(size_x=512, size_y=512)

layout = [iterator.for_job(i, n_jobs=4).partition_indices for i in range(4)]
print(f"tiles per job: {[len(part) for part in layout]}")

write granularity: (1, 2160, 2560) tiles per job: [25, 25, 0, 0]

Fifty tiles, but two fat chunks: two working jobs, and two empty no-ops. This partition_indices listing is the pre-flight check worth running before anything is submitted. Chunk the output to match the work instead:

# Chunk the output like the tiling (chunks are given in the reference image's
# axes, here czyx; the channel axis is squeezed away on a label).
label = ome_zarr.derive_label(
    "nuclei_distributed", chunks=(1, 1, 512, 512), overwrite=True
)
print(f"write granularity: {label.write_granularity}")

# `by_write_units` tiles by exactly that granularity, so every tile is its own
# independent write unit.
iterator = SegmentationIterator(
    image, label, channel_selection="DAPI", axes_order=["y", "x"]
).by_write_units()

layout = [iterator.for_job(i, n_jobs=4).partition_indices for i in range(4)]
print(f"tiles per job: {[len(part) for part in layout]}")

write granularity: (1, 512, 512) tiles per job: [13, 13, 12, 12]

from matplotlib.colors import ListedColormap

image_data = image.get_as_numpy(c=0, axes_order=["y", "x"])
job_map = np.full(image_data.shape, np.nan)
for job, indices in enumerate(layout):
    for index in indices:
        box = iterator.rois[index].to_pixel(pixel_size=image.pixel_size)
        y, x = box["y"], box["x"]
        job_map[
            int(y.start) : int(y.start + y.length),
            int(x.start) : int(x.start + x.length),
        ] = job

fig, ax = plt.subplots(figsize=(8, 3.4))
show_image(ax, image_data, title="which job owns which write unit (n_jobs=4)")
job_colors = ListedColormap(["#2e6fd6", "#22a699", "#f4a63a", "#c2185b"])
ax.imshow(job_map, cmap=job_colors, alpha=0.35, interpolation="nearest")
for edge in range(512, image_data.shape[1], 512):
    ax.axvline(edge - 0.5, color="white", lw=0.4)
for edge in range(512, image_data.shape[0], 512):
    ax.axhline(edge - 0.5, color="white", lw=0.4)
print(
    figure_html(
        fig,
        alt="The image tiled into 50 write units, each tinted by the job that "
        "owns it; the four jobs interleave freely because no two of them "
        "ever share a write unit.",
    )
)

Step 3: run the three-phase recipe¶

Schedulers like Fractal run distributed work as init → parallel tasks → consolidate, and the iterator verbs map onto those phases one to one: prepare_jobs performs any setup the run needs (wiping stale scratch state from earlier runs first) and returns one JSON-ready argument set per non-empty partition; each parallel task rebuilds the identical iterator and runs its own share; the consolidate task's finalize() is the one global step.

Every phase rebuilds the identical iterator from scratch — cheap, because construction is metadata-only — and derives its share on its own:

def build_iterator() -> SegmentationIterator:
    # Rebuilt identically in every phase; construction is metadata-only.
    return (
        SegmentationIterator(
            image,
            ome_zarr.get_label("nuclei_distributed"),
            channel_selection="DAPI",
            axes_order=["y", "x"],
            consolidation_mode="auto",
        )
        .with_stitch()
        .by_write_units()
        .with_halo(x=32, y=32)
    )

This run stitches, so prepare_jobs is required — the stitch scratch has to be created once, before any job writes, and the init step is that moment. (For a plain, unstitched writer it is optional: for_job and finalize alone are enough.) It returns the parallelization list, with empty partitions already dropped:

args_list = build_iterator().prepare_jobs(n_jobs=4)
print(args_list)

[{'job_index': 0, 'n_jobs': 4}, {'job_index': 1, 'n_jobs': 4}, {'job_index': 2, 'n_jobs': 4}, {'job_index': 3, 'n_jobs': 4}]

# On a real cluster each iteration is its own scheduler task.
for args in args_list:
    build_iterator().for_job(**args).segment(segment)

A slice's segment deliberately does not finalize: until the gather runs, only the iterated level is up to date, and the banked tiles are not yet reconciled. The consolidate task runs the one global step — it verifies every expected bank exists (a half-finished run errors, naming the tiles that never banked), resolves the seams, and rebuilds the pyramid:

build_iterator().finalize()

final = ome_zarr.get_label("nuclei_distributed")
print(f"objects: {len(np.unique(final.get_as_numpy())) - 1}")

objects: 1412

The result is bit-identical to the serial with_stitch() run of the stitching tutorial. Three properties are worth knowing:

A failed job never destroys the others' banks — re-run just that job (banking is idempotent) and gather as planned.
Every step validates a plan fingerprint stamped at init: change the tiling, halo, stitch config, or n_jobs between phases and the run fails loudly.
Every job must use the same n_jobs and the same iterator construction — the fingerprint catches most drift, but a custom seam matcher or the function itself cannot be fingerprinted, so declare the identical chain in every phase.

Step 4: measure across jobs¶

The read-only iterators end in a global join that per-job runs cannot reproduce piecewise, so their topic verbs are partition-aware: on a for_job slice, measure banks the job's raw pre-join records as a partial and returns None, and finalize() runs the single global join and returns the table — the same three-phase recipe, verb for verb:

import pandas as pd
from skimage.measure import regionprops_table

from ngio import Roi


def measure(image_data: np.ndarray, label_data: np.ndarray, roi: Roi) -> pd.DataFrame:
    props = regionprops_table(
        label_image=label_data.squeeze(-1),  # remove the channel axis
        intensity_image=image_data,
        properties=["label", "area", "mean_intensity"],
    )
    return pd.DataFrame(props)

from ngio.iterators import FeatureExtractorIterator


def build_measure_iterator() -> FeatureExtractorIterator:
    return FeatureExtractorIterator(
        image,
        ome_zarr.get_label("nuclei_distributed"),
        axes_order=["y", "x", "c"],
    ).by_blocks(num_x=2, num_y=2)


# init task
measure_args = build_measure_iterator().prepare_jobs(n_jobs=2)

# parallel tasks: on a slice, `measure` banks a partial and returns None
for args in measure_args:
    build_measure_iterator().for_job(**args).measure(measure)

# consolidate task: the one global join returns the table; storing it is yours
table = build_measure_iterator().finalize()
ome_zarr.add_table("nuclei_features_distributed", table, overwrite=True)
print(f"rows: {len(table.dataframe)}")

rows: 1477

Sanity check: read the table back¶

print(table_html(ome_zarr.get_table("nuclei_features_distributed").dataframe.head()))

label	area	mean_intensity-0	roi_name
1	6018.00	278.41	t0_z0_y0_x0
2	3738.00	253.19	t0_z0_y0_x0
3	5631.00	243.17	t0_z0_y0_x0
4	4838.00	249.84	t0_z0_y0_x0
5	7374.00	249.39	t0_z0_y0_x0

There are slightly more rows than objects: the 2×2 blocks split some objects, and a split object is measured by both sides. Every row is stamped with the roi_index / roi_name it came from, and reconciling the duplicates is a declared join away — the feature extraction tutorial shows one.

Next steps¶

Iterators guide — the partition model in full: why jobs need no locks, and what finalize refuses.
Stitching — the serial version of this run.
Feature extraction — joins, halos and duplicate-row reconciliation for measure.