cellmap_flow.serving.chunkmirage_ops

A served model and its chains as chunkmirage ops.

chunkmirage (chunkmirage.ops.Op) reads the raw data around each chunk, runs the ops on it, caches what the cached ones make and serves the result. A layer’s pipeline is layer_ops:

InferenceOp: the model, on its input voxel size, with the layer’s input normalization. Its output is cached, so a layer revisiting a region, or another layer with the same normalization, does not run the model again.
PostprocessOp, one per step of the layer’s postprocessing chain: the same PostProcessor steps the dashboard’s Output tab builds, run as they always were. Changing them leaves the model’s cached output valid.

The model itself, loaded and warmed up once per process, is a ServedModel registered by name (serve_model): ops are frozen pydantic models that chunkmirage rebuilds whenever a pipeline is built, so they only name it.

Importing this module is cheap (no torch), since chunkmirage loads the ops through the chunkmirage.ops entry point in any process that serves.

Attributes

`logger`
`EQUIVALENCES_TIMEOUT_SECONDS`
`EQUIVALENCES_EVERY_SECONDS`

Classes

`ServedModel`	A model a process serves, and the grid its chunks lie on.
`BlockInput`	What a model reads its input from: the block chunkmirage read for a chunk.
`InferenceOp`	A served model's raw output, `(channels, z, y, x)`, from the layer's
`PostprocessOp`	One step of the layer's postprocessing chain on the model's output.
`DevicePostprocessOp`	A postprocessing step that runs on the GPU (its `uses_device`), at

Functions

`serve_model`(served)
`served_model`(name) → ServedModel
`chunk_roi`(index, spatial_block, output_voxel_size, origin) → funlib.geometry.Roi	The region (nm) chunk `index` covers: `spatial_block` output voxels
`layer_ops`(model, input_norm, postprocess[, ...])	The ops of a layer's pipeline: the model on `input_norm`, then one op

Module Contents

cellmap_flow.serving.chunkmirage_ops.logger

cellmap_flow.serving.chunkmirage_ops.EQUIVALENCES_TIMEOUT_SECONDS = 10

cellmap_flow.serving.chunkmirage_ops.EQUIVALENCES_EVERY_SECONDS = 5

class cellmap_flow.serving.chunkmirage_ops.ServedModel

A model a process serves, and the grid its chunks lie on.

runner: the warmed-up Inferencer (with no device slots: the InferenceOp’s slots bound the device). translation: the centre of input voxel 0 in nm (worldarray’s translation) of the grid the model’s reads are mapped onto, whose voxel size is input_voxel_size: the level open_raw (cellmap_flow.io.raw) opened for the model, at that voxel size, resampled to it, or read as if it were at it (its corner rescaled too). level_shape: that level’s voxels (its resampled ones when resampled). origin: output voxel 0’s lower corner in nm, where the chunk grid starts. block: output voxels per chunk, spatial. halo: input voxels the model reads beyond each side of what it writes. ratio: input voxels per output voxel.

name: str

runner: Any

translation: tuple

level_shape: tuple

origin: numpy.ndarray

block: tuple

input_voxel_size: tuple

output_voxel_size: tuple

halo: tuple

ratio: tuple

output_channels: int

has_channel: bool

output_dtype: Any

stage_dtype: numpy.dtype

weights_version: int = 0

weights_changed(): Say the model’s weights changed in place (a finetune iteration): pipelines built from now on get new cache keys, so nothing computed before is served.

cellmap_flow.serving.chunkmirage_ops.serve_model(served: ServedModel)

cellmap_flow.serving.chunkmirage_ops.served_model(name: str) → ServedModel

class cellmap_flow.serving.chunkmirage_ops.BlockInput(block, block_start, voxel_size, translation, level_shape, input_norms)

What a model reads its input from: the block chunkmirage read for a chunk.

It stands in for the RawInput (cellmap_flow.io.raw) a ModelRunner reads with. Every config’s own process_chunk or predict, like the default forward, only calls to_ndarray_ts(roi) with the chunk grown by its context. The level’s level_shape voxels are voxel_size apart (nm) with voxel 0’s centre at translation, the grid RawInput reads on; the block starts at voxel block_start of it.

block

block_start

voxel_size

translation

level_shape

input_norms

to_ndarray_ts(roi): roi (nm) through the input chain, the data’s outermost voxels repeated past it, after the chain, as RawInput.to_ndarray_ts reads it with padding="edge": the voxels whose centres lie in roi (worldarray’s rule). A flat fill there is a wall the model never saw in training: on jrc_jurkat-1 setup_56 found 28% of the nuclei the volume’s edge cuts with zeros past it, 60% with the edge repeated. roi wholly outside the data reads zeros.

cellmap_flow.serving.chunkmirage_ops.chunk_roi(index, spatial_block, output_voxel_size, origin) → funlib.geometry.Roi: The region (nm) chunk index covers: spatial_block output voxels per chunk, the grid starting at origin. The whole chunk, as the model always computed it, also where the data ends inside it.

class cellmap_flow.serving.chunkmirage_ops.InferenceOp(/, **data: Any)

A served model’s raw output, (channels, z, y, x), from the layer’s input normalization of the raw data.

One level only (input_voxel_size): the model’s. Cached, as running the model is the expensive part; slots (set by the server from CELLMAP_FLOW_GPU_SLOTS) bounds how many chunks use the device at once.

name: ClassVar[str] = 'cellmap_flow.inference'

cache: ClassVar[bool] = True

slots: ClassVar[int | None] = 1

model: str

input_norm: list = []

model_post_init(__context): Override this method to perform additional initialization after __init__ and model_construct. This is useful if you want to do some validation that requires the entire model to be initialized.

property served: ServedModel

property halo

input_voxel_size(): The voxel size this op must read, along the data’s last axes and in the source’s units: a model trained at one resolution. The pipeline then runs it on one level (the source’s at that size, else one resampled to it), caches what it makes, and makes the coarser levels by downsampling that. None (the default): it runs on every level.

cache_token(): What this op’s output depends on outside its parameters, as a string that changes when that does: a weights file’s modification time, a model version. None (the default): nothing. It is part of digest, so a pipeline built after it changes has new stage keys and new links, and viewers refetch; DatasetRegistry.refresh rebuilds a served dataset to read it again.

output_info(info: chunkmirage.core.ArrayInfo) → chunkmirage.core.ArrayInfo

apply_at(block, box): Like apply but told where block sits (box = its halo-padded extent in voxels of this scale level). Override when the result must depend on position, e.g. to make per-chunk labels globally unique. Default delegates to apply.

class cellmap_flow.serving.chunkmirage_ops.PostprocessOp(/, **data: Any)

One step of the layer’s postprocessing chain on the model’s output.

step is the step as the Output tab sends it ({"name": ..., **params}); the op runs that PostProcessor, built once per layer so it keeps its state (the blockwise merger’s equivalences, sent to the dashboard at _dashboard_url). A chain is one of these per step, which chunkmirage runs back to back on a chunk. chunk_corner is the chunk’s index on the served grid, as Morton and the affinity segmenter key their ids on.

name: ClassVar[str] = 'cellmap_flow.postprocess'

model: str

step: dict

model_post_init(__context): Override this method to perform additional initialization after __init__ and model_construct. This is useful if you want to do some validation that requires the entire model to be initialized.

for_dashboard(dashboard_url, dataset): Send merged ids to dashboard_url as those of layer dataset.

property served: ServedModel

output_info(info: chunkmirage.core.ArrayInfo) → chunkmirage.core.ArrayInfo

apply_at(block, box): Like apply but told where block sits (box = its halo-padded extent in voxels of this scale level). Override when the result must depend on position, e.g. to make per-chunk labels globally unique. Default delegates to apply.

class cellmap_flow.serving.chunkmirage_ops.DevicePostprocessOp(/, **data: Any)

A postprocessing step that runs on the GPU (its uses_device), at most slots chunks at once (set by the server, as the model’s), the nearest first: several at once slow each other down, and the first chunks of a view come no sooner than the rest.

name: ClassVar[str] = 'cellmap_flow.postprocess_on_device'

slots: ClassVar[int | None] = 1

cellmap_flow.serving.chunkmirage_ops.layer_ops(model, input_norm, postprocess, dashboard_url=None, dataset=''): The ops of a layer’s pipeline: the model on input_norm, then one op per step of postprocess, then, if the chain leaves another dtype than the one served, chunkmirage’s cast to it (as astype).