Multi-step pipeline with resource-aware dispatch#
A three-step image pipeline on the astronaut photo, set up so the resource-aware dispatcher actually has something to balance:
edges — per-block Sobel edge magnitude. CPU-bound.
segment — per-block watershed using local-minima markers. Tagged as “GPU” (pseudo-resource — we don’t actually use a GPU, but treating the step as scarce demonstrates how
requires=and a globalresources=budget interact).recolor — for every segment found in step 2, replace its pixels with the average colour from the original image. CPU.
Tasks 1 and 3 share a CPU pool of 4. Task 2 is the bottleneck — only one “GPU” worker ever runs concurrently. Once the edges task drains, its CPU workers exit and the recolor task can grow into the freed budget. This is the cluster pattern: GPU-bound stage in the middle, CPU-bound stages on either side, sharing budgets across the pipeline.
import tempfile
import time
from pathlib import Path
import daisy.logging as gl
import daisy.v2 as daisy
import matplotlib.pyplot as plt
import numpy as np
import zarr
from scipy.ndimage import label as scipy_label
from skimage.data import astronaut
from skimage.filters import sobel
from skimage.segmentation import watershed
_TMP = Path(tempfile.mkdtemp(prefix="daisy_astronaut_pipeline_"))
gl.set_log_basedir(_TMP / "logs")
print(f"output paths under: {_TMP}")
output paths under: /tmp/daisy_astronaut_pipeline_c6muwe5g
Load the astronaut#
INPUT = astronaut() # (512, 512, 3) uint8
H, W = INPUT.shape[:2]
BLOCK = 32
print(
f"input: {INPUT.shape}, dtype={INPUT.dtype}, "
f"block grid: {H // BLOCK}×{W // BLOCK} = {(H // BLOCK) * (W // BLOCK)} blocks"
)
plt.figure(figsize=(5, 5))
plt.imshow(INPUT, interpolation="nearest")
plt.title("input")
plt.axis("off")
input: (512, 512, 3), dtype=uint8, block grid: 16×16 = 256 blocks
(np.float64(-0.5), np.float64(511.5), np.float64(511.5), np.float64(-0.5))
Allocate the intermediate arrays#
Each step writes into its own array and reads from the previous step’s output. Every daisy worker is its own process, so the stages exchange data through zarr arrays on disk, chunked to the block size so concurrent workers never touch the same chunk. The process functions open the group themselves: a path pickles, an open handle does not.
ARRAY_PATH = str(_TMP / "arrays.zarr")
_arrays = zarr.open_group(ARRAY_PATH, mode="w")
_arrays.create_array("edges", shape=(H, W), chunks=(BLOCK, BLOCK), dtype=np.uint8)
_arrays.create_array("labels", shape=(H, W), chunks=(BLOCK, BLOCK), dtype=np.uint32)
_arrays.create_array(
"output", shape=(H, W, 3), chunks=(BLOCK, BLOCK, 3), dtype=np.uint8
)
# Per-task concurrency tracking, so we can plot how the dispatcher
# distributes workers across the pipeline. Workers append one line per
# enter/exit to a shared log (small O_APPEND writes are atomic); the
# driver rebuilds the timeline from it after the run.
EVENTS = str(_TMP / "events.log")
def _record(task_id, kind):
with open(EVENTS, "a") as f:
f.write(f"{time.time():.6f} {task_id} {kind}\n")
Step 1 — edges (CPU, max 4 workers, 1 CPU per worker)#
Read the input slice, convert to grayscale, take the Sobel magnitude,
and write to EDGES. No halo — Sobel uses reflection padding at
block boundaries, which is fine for this demo (the resulting block
seams are visible in the output, on purpose).
def step_edges(block):
_record("edges", "enter")
try:
time.sleep(0.002) # simulate a more expensive operation
off = list(block.write_roi.offset)
shape = list(block.write_roi.shape)
r0, c0 = off[0], off[1]
rs, cs = shape[0], shape[1]
gray = INPUT[r0 : r0 + rs, c0 : c0 + cs].mean(axis=-1) / 255.0
arrays = zarr.open_group(ARRAY_PATH, mode="r+")
arrays["edges"][r0 : r0 + rs, c0 : c0 + cs] = (
(sobel(gray) * 255).clip(0, 255).astype(np.uint8)
)
finally:
_record("edges", "exit")
Step 2 — watershed (pseudo-“GPU”, max 1 worker)#
Read the edge map, build markers from low-edge pixels (interior of
objects), run a watershed, offset the local labels into a global
range so blocks don’t collide, and write to LABELS. We tag this as
requires={"gpu": 1} to demonstrate the resource model — even
though it’s pure CPU work in reality, the dispatcher will only
schedule one watershed worker at a time because the global budget
only has "gpu": 1.
def step_segment(block):
_record("segment", "enter")
try:
time.sleep(0.0025) # simulate the expensive "GPU" step
off = list(block.write_roi.offset)
shape = list(block.write_roi.shape)
r0, c0 = off[0], off[1]
rs, cs = shape[0], shape[1]
arrays = zarr.open_group(ARRAY_PATH, mode="r+")
edges_tile = (
arrays["edges"][r0 : r0 + rs, c0 : c0 + cs].astype(np.float32) / 255.0
)
# Markers: where the edge response is low (object interiors).
markers, _ = scipy_label(edges_tile < 0.04)
seg = watershed(edges_tile, markers=markers).astype(np.uint32)
# Make labels globally unique by offsetting per block.
block_offset = block.block_id[1] * BLOCK * BLOCK
arrays["labels"][r0 : r0 + rs, c0 : c0 + cs] = np.where(
seg > 0, seg + block_offset, 0
)
finally:
_record("segment", "exit")
Step 3 — recolor (CPU, max 4 workers, 1 CPU per worker)#
For each segment in this block, replace its pixels with the average RGB colour of the original image over that segment’s mask. Output is a low-poly-style abstraction of the astronaut.
def step_recolor(block):
_record("recolor", "enter")
try:
time.sleep(0.004)
off = list(block.write_roi.offset)
shape = list(block.write_roi.shape)
r0, c0 = off[0], off[1]
rs, cs = shape[0], shape[1]
arrays = zarr.open_group(ARRAY_PATH, mode="r+")
labels_tile = arrays["labels"][r0 : r0 + rs, c0 : c0 + cs]
rgb_tile = INPUT[r0 : r0 + rs, c0 : c0 + cs]
out_tile = np.zeros_like(rgb_tile)
for label_id in np.unique(labels_tile):
mask = labels_tile == label_id
out_tile[mask] = rgb_tile[mask].mean(axis=0).astype(np.uint8)
arrays["output"][r0 : r0 + rs, c0 : c0 + cs] = out_tile
finally:
_record("recolor", "exit")
Define the pipeline#
Three tasks chained via + composition: edges → segment → recolor.
requires declares each task’s per-worker resource cost; max_workers
is the hard cap. resources={"cpu": 4, "gpu": 1} is the global budget
— enough CPU headroom that the pipeline isn’t CPU-starved, with the
watershed stage gated by a single “GPU” slot.
edges_task = daisy.Task(
task_id="edges",
total_roi=daisy.Roi([0, 0], [H, W]),
read_roi=daisy.Roi([0, 0], [BLOCK, BLOCK]),
write_roi=daisy.Roi([0, 0], [BLOCK, BLOCK]),
process_function=step_edges,
read_write_conflict=False,
max_workers=3,
max_retries=0,
requires={"cpu": 1},
)
segment_task = daisy.Task(
task_id="segment",
total_roi=daisy.Roi([0, 0], [H, W]),
read_roi=daisy.Roi([0, 0], [BLOCK, BLOCK]),
write_roi=daisy.Roi([0, 0], [BLOCK, BLOCK]),
process_function=step_segment,
read_write_conflict=False,
max_workers=4, # cap is 4, but global budget caps at 1
max_retries=0,
requires={"gpu": 1},
)
recolor_task = daisy.Task(
task_id="recolor",
total_roi=daisy.Roi([0, 0], [H, W]),
read_roi=daisy.Roi([0, 0], [BLOCK, BLOCK]),
write_roi=daisy.Roi([0, 0], [BLOCK, BLOCK]),
process_function=step_recolor,
read_write_conflict=False,
max_workers=4,
max_retries=0,
requires={"cpu": 1},
)
pipeline = edges_task + segment_task + recolor_task
Run the pipeline#
The Pipeline’s edge structure tells the scheduler that each stage
depends on its upstream. Block 0 of recolor can start as soon as
block 0 of segment completes; it doesn’t have to wait for segment
to fully drain.
t_start = time.perf_counter()
ok = pipeline.run_blockwise(resources={"cpu": 4, "gpu": 1})
elapsed = time.perf_counter() - t_start
assert ok
print(f"\npipeline elapsed: {elapsed * 1e3:.0f} ms")
Execution Summary
-----------------
task blocks completed skipped failed orphaned
───────── ─────── ────────── ─────── ────── ────────
edges ✔ 256 256 0 0 0
segment ✔ 256 256 0 0 0
recolor ✔ 256 256 0 0 0
Resource Utilization
--------------------
edges 256 blocks, 3 workers, 3.93 s wall
blocks round trip 7.2 ms ± 2.1 ms
workers startup 787.0 ms ± 156.8 ms idle 47.2 µs ± 62.1 µs teardown 147.6 ms ± 50.4 ms
exits 3 clean
id startup blocks idle idle / block compute compute / block teardown exit
0 892.7 ms 73 3.7 ms 50.2 µs ± 74.2 µs 516.1 ms 7.1 ms ± 1.5 ms 138.6 ms clean
1 861.5 ms 77 3.7 ms 48.4 µs ± 63.3 µs 547.2 ms 7.1 ms ± 1.5 ms 201.9 ms clean
2 606.8 ms 106 4.7 ms 44.2 µs ± 51.7 µs 791.6 ms 7.5 ms ± 2.8 ms 102.3 ms clean
segment 256 blocks, 1 worker, 3.93 s wall
blocks round trip 9.3 ms ± 3.0 ms
workers startup 695.8 ms idle 53.1 µs ± 33.7 µs teardown 113.2 ms
exits 1 clean
id startup blocks idle idle / block compute compute / block teardown exit
3 695.8 ms 256 13.6 ms 53.1 µs ± 33.7 µs 2.37 s 9.3 ms ± 3.0 ms 113.2 ms clean
recolor 256 blocks, 4 workers, 3.93 s wall
blocks round trip 10.6 ms ± 2.9 ms
workers startup 519.0 ms ± 21.5 ms idle 15.9 ms ± 9.9 ms teardown 107.2 ms ± 9.8 ms
exits 4 clean
id startup blocks idle idle / block compute compute / block teardown exit
4 490.5 ms 75 1.02 s 13.6 ms ± 10.8 ms 864.3 ms 11.5 ms ± 4.9 ms 109.6 ms clean
5 515.4 ms 65 1.01 s 15.6 ms ± 10.1 ms 674.5 ms 10.4 ms ± 1.4 ms 113.2 ms clean
6 539.9 ms 60 1.00 s 16.7 ms ± 9.4 ms 604.4 ms 10.1 ms ± 0.9 ms 113.2 ms clean
7 530.1 ms 56 1.02 s 18.3 ms ± 8.3 ms 565.0 ms 10.1 ms ± 1.5 ms 92.8 ms clean
pipeline elapsed: 3954 ms
Visualise each stage of the pipeline#
EDGES = _arrays["edges"][:]
LABELS = _arrays["labels"][:]
OUTPUT = _arrays["output"][:]
fig, axes = plt.subplots(1, 4, figsize=(16, 4))
axes[0].imshow(INPUT, interpolation="nearest")
axes[0].set_title("1. input")
axes[0].axis("off")
axes[1].imshow(EDGES, cmap="gray")
axes[1].set_title("2. Sobel edges")
axes[1].axis("off")
# Hash-coloured labels.
def label_to_rgb(seg, bg_colour=(0, 0, 0)):
seg = np.asarray(seg).astype(np.int64)
ids, inverse = np.unique(seg, return_inverse=True)
rng = np.random.default_rng(42)
colours = 0.2 + 0.8 * rng.random((ids.size, 3))
for i, lid in enumerate(ids):
if lid == 0:
colours[i] = bg_colour
return (colours[inverse].reshape(*seg.shape, 3) * 255).astype(np.uint8)
axes[2].imshow(label_to_rgb(LABELS), interpolation="nearest")
axes[2].set_title(f"3. watershed labels ({len(np.unique(LABELS))} segments)")
axes[2].axis("off")
axes[3].imshow(OUTPUT, interpolation="nearest")
axes[3].set_title("4. recoloured")
axes[3].axis("off")
fig.tight_layout()
Concurrency timeline — was the budget honoured?#
Plot, for every event in the run, how many workers of each task were alive simultaneously. Two things to confirm visually:
segment(the “GPU” step) never exceeds 1 alive worker — the global budget caps it.edgesandrecolortogether never exceed 4 alive workers simultaneously, even though each task’s ownmax_workers=4cap would individually allow 4.As
edgesdrains, its alive count falls;recolor’s alive count rises into the freed CPU budget — no respawning, just the dispatcher spawning new workers for the next ready task.
_timeline = []
_alive = {"edges": 0, "segment": 0, "recolor": 0}
with open(EVENTS) as f:
for line in sorted(f):
t, task_id, kind = line.split()
_alive[task_id] += 1 if kind == "enter" else -1
_timeline.append(
{"t": float(t), "task": task_id, "kind": kind, "alive": dict(_alive)}
)
t0 = _timeline[0]["t"]
times = [(e["t"] - t0) * 1e3 for e in _timeline]
edges_alive = [e["alive"]["edges"] for e in _timeline]
segment_alive = [e["alive"]["segment"] for e in _timeline]
recolor_alive = [e["alive"]["recolor"] for e in _timeline]
panels = [
(
"edges (CPU)",
edges_alive,
"#4878CF",
4,
"max_workers=3 — shares cpu=4 with recolor",
),
(
"segment ('GPU')",
segment_alive,
"#956CB4",
1,
"max_workers=4 — capped by gpu budget = 1",
),
(
"recolor (CPU)",
recolor_alive,
"#6ACC65",
4,
"max_workers=4 — shares cpu=4 with edges",
),
]
fig, axes = plt.subplots(3, 1, figsize=(11, 6.5), sharex=True)
for ax, (label, data, colour, cap, subtitle) in zip(axes, panels, strict=False):
ax.step(times, data, where="post", color=colour, linewidth=1.5)
ax.fill_between(times, data, step="post", alpha=0.25, color=colour)
ax.axhline(cap, color="black", linestyle=":", alpha=0.5, linewidth=0.9)
ax.text(times[-1], cap, f" cap={cap}", va="center", fontsize=8, alpha=0.7)
ax.set_ylabel("alive workers")
ax.set_ylim(-0.2, max(cap, max(data)) + 0.6)
ax.set_title(f"{label} — {subtitle}", fontsize=10, loc="left")
axes[-1].set_xlabel("time since pipeline start (ms)")
fig.tight_layout()
# Combined CPU concurrency check: edges + recolor.
combined = [e["alive"]["edges"] + e["alive"]["recolor"] for e in _timeline]
print(f"peak (edges + recolor) alive = {max(combined)} (budget cpu=4)")
print(f"peak segment alive = {max(segment_alive)} (budget gpu=1)")
plt.show()
peak (edges + recolor) alive = 4 (budget cpu=4)
peak segment alive = 1 (budget gpu=1)
