cellmap_flow.serving.chunkmirage_server

The inference server, built on chunkmirage.

ChunkmirageServer serves one model’s predictions over one dataset as zarr v2, for Neuroglancer. chunkmirage does the serving: reading the raw data around each chunk, caching the model’s output, queueing chunks for the device (finest level first, dropping those nobody waits for any more) and the zarr format. cellmap-flow keeps the model, its chains (serving.chunkmirage_ops) and the routes the dashboard reads.

A layer is zarr://<server>/<model><ARGS_KEY><blob><ARGS_KEY>/zarr: a name, then chunkmirage’s format. The name is resolved into a pipeline on its first request (pipeline_for), and the blob’s chain is the layer’s own, so layers sharing a server never share a chain. The array has one level, s0, with any channel axis first.

/__control__/model_info (and its older name /output_probe) report the model’s geometry and output activation; /__control__/restart asks a finetune job to train again, and /__control__/stop to stop training after the epoch it is in. model_info says "engine": "chunkmirage": the release before this one had two engines, and its dashboard reads the field to tell their URLs apart.

CELLMAP_FLOW_PREDICTION_CACHE_BYTES (default 16 GiB; 0 for none) bounds the memory kept for computed chunks, the model’s output and the raw data read for it; CELLMAP_FLOW_RAW_CACHE_BYTES (default 1 GiB) tensorstore’s decoded raw chunks.

Attributes

`logger`
`ENGINE`
`PREDICTION_CACHE_BYTES_ENV`
`PREDICTION_CACHE_BYTES_DEFAULT`
`RAW_CACHE_BYTES_ENV`
`RAW_CACHE_BYTES_DEFAULT`
`HALF_PRECISION_ENV`
`INFERENCE_PREFETCH`

Classes

ChunkmirageServer

One model's predictions over dataset_name, served by chunkmirage.

Module Contents

cellmap_flow.serving.chunkmirage_server.logger

cellmap_flow.serving.chunkmirage_server.ENGINE = 'chunkmirage'

cellmap_flow.serving.chunkmirage_server.PREDICTION_CACHE_BYTES_ENV = 'CELLMAP_FLOW_PREDICTION_CACHE_BYTES'

cellmap_flow.serving.chunkmirage_server.PREDICTION_CACHE_BYTES_DEFAULT = 17179869184

cellmap_flow.serving.chunkmirage_server.RAW_CACHE_BYTES_ENV = 'CELLMAP_FLOW_RAW_CACHE_BYTES'

cellmap_flow.serving.chunkmirage_server.RAW_CACHE_BYTES_DEFAULT = 1073741824

cellmap_flow.serving.chunkmirage_server.HALF_PRECISION_ENV = 'CELLMAP_FLOW_HALF_PRECISION'

cellmap_flow.serving.chunkmirage_server.INFERENCE_PREFETCH = 2

class cellmap_flow.serving.chunkmirage_server.ChunkmirageServer(dataset_name, model_config, restart_callback=None, stop_callback=None, restart_token=None, resample=False)

One model’s predictions over dataset_name, served by chunkmirage.

resample reads a level resampled to the model’s input voxel size when the data has none at it (open_raw(..., resample=True), see cellmap_flow.io.raw), instead of reading the nearest level as if it were at that size; cellmap_flow serve --resample. restart_callback enables POST /__control__/restart, and stop_callback POST /__control__/stop, which then only accept requests carrying restart_token in the X-Restart-Token header.

CELLMAP_FLOW_GPU_SLOTS (default 1) is how many chunks use the device at once (inference.runner.DeviceSlots); CELLMAP_FLOW_HALF_PRECISION=1 serves the model under half-precision autocast, as half_precision = True in its config does, if it agrees with fp32 at warmup (see Inferencer). The Inferencer is built before anything else, so a model whose output contradicts its declared shapes fails here, on the warmup forward, before the server announces itself.

restart_callback = None

stop_callback = None

restart_token = None

resample = False

inferencer

geometry

axes

has_channel

served

registry

source

app

pipeline_for(name): The pipeline for layer name: the raw data, the model on the layer’s normalization, then the layer’s postprocessing. Only this server’s model and data; the name only chooses the chains.

weights_changed(): The model’s weights changed in place: layers resolved from now on are computed anew (the finetune loop’s new iteration, under a new name).

model_info() → dict

The served model’s geometry and output activation, and the engine serving it.

Both halves exist so the dashboard does not have to build the model itself: it runs on whatever node launched the jobs, often without a usable GPU. The geometry comes from the validated config, so it is reported even when the warmup probe failed; script-defined models expose nothing else to the dashboard, which makes this the only place it can learn their shapes and channel names. The activation half is read once, as the model is loaded (inference/runner.py), and lets the dashboard propose a postprocessing chain, or flag one that contradicts the model: output_class, read off the network’s graph (output_class_from “graph”, with output_op, the op that ends it) or else from the warmup probe’s range (output_min, output_max) on noise (“probe”, a guess); None, from neither, when the graph does not say and the probe’s output was flat. A resampled input really is at input_voxel_size; a relabelled one is at the level’s, which moves the output (see to_model_info).

read_chunk(index=(0, 0, 0), name='check'): Chunk index of layer name (no chain: the process’s), computed here without HTTP: what cellmap_flow infer --server-check runs.

run(port=None, certfile=None, keyfile=None)

Serve until stopped, on port (any free one for None or 0).

The address is announced, as the IP_PATTERN marker and in the ready file a launcher asked for (jobs/ready.py), once the port is bound and the server is taking requests.

stop(): Ask a running run to return.