Pipeline & top-level API¶

class Pipeline(work_source)¶

Bases: object

Fluent builder that compiles user ops into an executable pipeline.

add_op(op: BaseOp, /, *, name: str | None = None, placement: str = 'auto') → Pipeline[source]¶

add_op(name: str, /, *, process_many: Callable[[list[Any]], list[Any]], preserves_cursor_order: bool, accumulator: Callable[[...], Accumulator[Any]] | None = None, process_one: Callable[[Any], list[Any]] | None = None, validation_samples: Callable[[], list[SampleRecord]] | None = None, parallelism: int = 1, placement: str = 'auto', indexable: bool = False, batch_shape_sensitive: bool = False, requires_serial_state: bool = False) → Pipeline

Append a custom operator — two forms.

Instance form: pipe.add_op(op, name=..., placement=...). Pass a prebuilt BaseOp subclass instance. Use this when the op needs full lifecycle control — __init__ for configuration, setup for per-worker resource construction (tokenizers, model handles) and OpContext services, traits for any combination of OpTraits fields. The framework deep-copies the instance per parallel worker, so self.* attributes set in setup are isolated per worker on every runner (including the thread runner). This matches the lifecycle built-in operators use. Traits come from op.traits() rather than kwargs.

Kwargs form: pipe.add_op(name, *, process_many=..., preserves_cursor_order=..., ...). Convenience for stateless transforms. The framework builds an internal BaseOp subclass from the supplied callables + accumulator factory + traits. process_many must not carry state between calls; any per-lane or cross-invocation state belongs in the accumulator. Read-only shared resources (codec tables, thresholds) can be captured in a closure over the callables. See the “Accumulators and Operators” page in the documentation for the operator/accumulator split.

Parameters:

op – Instance form. The BaseOp instance to attach.
name (str | None) – Operator name used in plan graphs, metrics, and logs. The instance form defaults to the class name of op.
placement (str) – Placement hint passed to the planner ("auto", "local", or a runner-specific tag).
process_many (Callable[[list[Any]], list[Any]] | None) – Kwargs form, required. Callable applied to each ready batch from the accumulator. Receives a list of upstream items, returns the list of downstream items. Runs in parallel workers when parallelism > 1.
accumulator (Callable[[...], Accumulator[Any]] | None) –
Kwargs form. Optional factory returning a fresh Accumulator instance. The framework calls the factory once at runner setup and again on reset_buffers between runs, so a factory (rather than an instance) is required for correctness. Two signatures are supported and auto-detected:
- Callable[[], Accumulator] — simplest form. Use lambda: CountingAccumulator(max_batch=N) for size-based per-lane batching (it always groups by lane).
- Callable[*, deterministic, ctx], Accumulator] — for accumulators whose construction depends on the deterministic mode (e.g. enabling latency-based flush only when deterministic is False) or runtime context. Example:
```
accumulator=lambda *, deterministic, ctx: (
    CountingAccumulator(
        max_batch=N,
        max_latency_ms=None if deterministic else 3,
    )
)
```
Defaults to a PassthroughAccumulator factory — each upstream micro-batch becomes one ready batch as-is.
process_one (Callable[[Any], list[Any]] | None) – Kwargs form. Optional fast path for single-element processing. If omitted, the framework wraps each element in a list and routes it through process_many.
validation_samples (Callable[[], list[SampleRecord]] | None) – Kwargs form. Optional factory returning a list of SampleRecord instances for the validation harness. Override when process_many requires payload fields beyond the validator’s synthetic {'text': str, 'value': int} — letting the harness run the full op-level + state-diff suite instead of degrading to OP_REJECTS_GENERIC_PAYLOAD. Buggy factories surface OP_VALIDATION_SAMPLES_FACTORY_FAILED and the validator falls back to synthetic records. Include at least two distinct lane_id values so the cross-call-state probe stays meaningful.
parallelism (int | None) – Kwargs form. Number of worker invocations to run in parallel for this op. Default 1 (serial). Increase when process_many is CPU/GPU-bound.
preserves_cursor_order (bool | None) – Kwargs form, required. True when process_many emits records whose chunk_id order matches their inputs (1:1 maps, payload transforms, non-reordering filters). False when the op reorders, shuffles, or packs. The planner picks a different eviction strategy based on this trait — getting it wrong corrupts checkpoint semantics silently. See the Checkpointing page in the documentation.
indexable (bool | None) – Kwargs form. Whether the op preserves indexability through the plan. Default False; set True only if the transform is 1:1 and deterministic.
batch_shape_sensitive (bool | None) – Kwargs form. Set True when process_many’s output can depend on how inputs are grouped into micro-batches (per-batch RNG, statistics, etc.). In deterministic mode this disables latency-based accumulator flushing for stages containing this op, preserving strong determinism at the cost of some throughput.
requires_serial_state (bool | None) – Kwargs form. Set True when the op’s accumulator holds cross-invocation state that cannot be sharded across parallel worker instances. In deterministic mode the planner pins parallelism=1 for ops with this trait.

Returns:

Self for method chaining.

Return type:

Pipeline

Examples:

# Instance form — class with setup/traits/accumulator overrides.
class Tokenize(BaseOp):
    def __init__(self, name):
        super().__init__()
        self._name = name
        self._tokenizer = None
    def traits(self):
        return OpTraits(preserves_cursor_order=True, parallelism=4)
    def setup(self, ctx):
        super().setup(ctx)
        self._tokenizer = load_tokenizer(self._name)
    def process_many(self, elems):
        return [self._tokenizer.encode(e) for e in elems]

pipeline.add_op(Tokenize("gpt2"))

# Kwargs form — stateless windowed transform.
from zephon.ops import CountingAccumulator
pipeline.add_op(
    "windowed_transform",
    process_many=lambda elems: [...],
    accumulator=lambda: CountingAccumulator(max_batch=64),
    parallelism=4,
    preserves_cursor_order=True,
)

batch(microbatch_size, *, drop_last=True, placement='auto')[source]¶

Return type:: Pipeline

checkpoint()[source]¶

Return type:: dict[str, Any]

decode_text(parallelism=None, **kwargs)[source]¶

Return type:: Pipeline

enable_observability(tracking=ExecutionTrackingMode.NODES, *, sink=None)[source]¶

Enable runtime metrics collection for the compiled pipeline.

Return type:: Pipeline

ensure_mixture(*, max_buffer_size=1000, drain_target_ratio=0.8, obsolete_drain_rate=0.1, weight_by='auto', warn_tolerance=None, mixture=None, placement='auto', parallelism=None)[source]¶

Enforce mixture ratios using adaptive Smooth Weighted Round Robin.

Place this operator after tokenization (for token-level) or after filtering (for sample-level). The operator automatically detects token fields from common names (input_ids, tokens, token_ids, ids).

Uses adaptive buffering: emit immediately when SWRR’s ideal component is available, buffer when the desired component isn’t present yet. Falls back to emitting best-available when max_buffer_size is reached.

The SWRR algorithm tracks deficit (target - actual) and always picks the component most “owed” samples, ensuring smooth, deterministic convergence.

Parameters:

max_buffer_size (int | None) – Samples to hold while waiting for the component the target mixture needs next. Default 1000; when the buffer fills, the operator emits what it has, so the output mixture may drift if a component stays scarce. Pass None to instead drop the surplus it cannot place on-target — at every epoch boundary and at end of stream. This makes the mixture exact but discards data, so use it only when an exact mixture matters more than keeping every sample. How much is dropped depends on the stream’s skew and on flush_every_k_chunks (smaller → more dropped). ensure_mixture is always non-monotonic, so that cadence defaults to 8; in strict mode the buffer is then discarded on roughly every flush. If drops are too frequent, raise flush_every_k_chunks above the default to give scarce components more time to absorb the buffer (at the cost of holding more in memory between flushes).
drain_target_ratio (float) – When forced to emit (buffer hits max_buffer_size), drain the buffer down to this fraction of max_buffer_size before stopping. Default is 0.8 (drain to 80% of max). Ignored when max_buffer_size=None.
obsolete_drain_rate (float) – Fraction of emissions reserved for draining obsolete components (those no longer in the current mixture target). Default is 0.1 (10%), meaning 1 in every 10 emissions drains an obsolete sample.
weight_by (Callable[[SampleRecord], float] | Literal['samples', 'auto'] | str) –
How to compute sample weights. Options:
- auto (default): Auto-detect token field from common names (input_ids, tokens, token_ids, ids). Raises if not found.
- samples: Each sample has weight 1.
- Explicit field name (e.g., input_ids): Use that field’s length.
- Callable: Custom function taking SampleRecord, returning float.
warn_tolerance (float | None) – If set, warn when mixture drift exceeds this value (0.05 = ±5%). If None (default), no warnings are emitted.
mixture (dict[str, float] | None) – Explicit mixture target {component_name: float}. If None, derived from chunk mixture via engine context.
placement (str) – Placement hint for this operator.
parallelism (int | None) – Override default parallelism for this operator.

Returns:

Self for method chaining.

Return type:

Pipeline

Examples

# Token-level (default) - place after tokenize
pipeline.fetch().tokenize(tokenizer_id="gpt2", field="text").ensure_mixture()

# Sample-level enforcement (after filter)
pipeline.fetch().filter(...).ensure_mixture(weight_by="samples")

# With explicit token field
pipeline.fetch().tokenize(tokenizer_id="gpt2", field="text").ensure_mixture(
    weight_by="input_ids"
)

# With warnings for drift (warn if >5% deviation)
pipeline.fetch().tokenize(tokenizer_id="gpt2", field="text").ensure_mixture(
    warn_tolerance=0.05
)

# Explicit mixture target (override chunk mixture)
pipeline.fetch().tokenize(tokenizer_id="gpt2", field="text").ensure_mixture(
    mixture={"English": 0.7, "German": 0.3}
)

explain()[source]¶

Return type:: str

fetch(parallelism=None, max_batch=None)[source]¶

Configure fetch operator parallelism and batch size.

Parameters:

parallelism (int | None) – Number of parallel fetch workers. If None, uses default.
max_batch (int | None) – Maximum batch size for fetch accumulator. If None, uses default (64).

Return type:

Pipeline

fetch_parallelism(parallelism, max_batch=None)[source]¶

Override the implicit FetchOp parallelism and batch size.

Parameters:

parallelism (int | None) – Number of parallel fetch workers. If None, uses default from traits.
max_batch (int | None) – Maximum batch size for fetch accumulator. If None, uses default (64).

Return type:

Pipeline

fetch_timing_snapshot()[source]¶

Return fetch timings for the inline engine, or None if absent.

Data is collected only in ExecutionTrackingMode.NODES.

Return type:: FetchTimingSummary | None

inflight_summary()[source]¶

Return per-lane inflight chunk counts: {lane_id: count}.

In MTP mode this reads from shared memory (non-blocking, zero IPC). In inline mode it reads the engine dict directly. Returns {} if the engine is not running.

Return type:: dict[int, int]

property is_indexable: bool¶

map_batch(transform_fn, *, drop_none=True, placement='auto', parallelism=None)[source]¶

Add a batch-level map transformation operator.

Applies a user-provided transformation function to each SampleBatch. This operator must be placed after a batch() call in the pipeline. For per-sample transforms (before batching), use map_transform() instead.

Parameters:

transform_fn (Callable[[SampleBatch], SampleBatch | None]) – Callable that transforms a SampleBatch. If it returns None and drop_none=True, the batch is filtered out. Supports lambdas, closures, and nested functions.
drop_none (bool) – If True, drop batches where transform_fn returns None.
placement (str) – Placement hint for this operator.
parallelism (int | None) – Override default parallelism for this operator.

Returns:

Self for method chaining.

Return type:

Pipeline

Example

>>> pipeline.batch(32).map_batch(lambda b: process(b))

map_transform(transform_fn, *, drop_none=True, placement='auto', parallelism=None)[source]¶

Add a map-style transformation operator.

Applies a user-provided transformation function to each sample’s payload. Supports lambda functions, closures, and nested functions seamlessly with multiprocessing-based runners (uses cloudpickle for serialization).

Parameters:

transform_fn (Callable[[bytes | memoryview | str | int | float | complex | Any | list[bytes | memoryview | str | int | float | complex | Any | list[SamplePayload] | dict[Any, SamplePayload]] | dict[Any, bytes | memoryview | str | int | float | complex | Any | list[SamplePayload] | dict[Any, SamplePayload]]], bytes | memoryview | str | int | float | complex | Any | list[bytes | memoryview | str | int | float | complex | Any | list[SamplePayload] | dict[Any, SamplePayload]] | dict[Any, bytes | memoryview | str | int | float | complex | Any | list[SamplePayload] | dict[Any, SamplePayload]] | None]) – Callable that transforms the payload. If it returns None and drop_none=True, the sample is filtered out. Supports lambdas, closures, and nested functions.
drop_none (bool) – If True, drop samples where transform_fn returns None.
placement (str) – Placement hint for this operator.
parallelism (int | None) – Override default parallelism for this operator.

Returns:

Self for method chaining.

Return type:

Pipeline

Example

>>> pipeline.map_transform(lambda x: {"value": x["value"] * 2})

materialize(placement='auto', parallelism=None)[source]¶

Return type:: Pipeline

metrics_snapshot()[source]¶

Return node metrics, or None unless inline tracking is active.

Return type:: PipelineSummary | None

mtp_queue_stats()[source]¶

Return occupancy of the MTP hand-off queue (see MTPQueueStats).

Returns None when no MTP subprocess is live: mtp_mode is off, iteration has not started or has finished, or the daemon fallback chose inline mode. Cheap to call once per training step.

Return type:: MTPQueueStats | None

options(**hints)[source]¶

Return type:: Pipeline

pack_flat(max_length, *, num_bins=None, max_sequences_per_bin=None, algorithm='first_fit', tokens_field='auto', pad_token_id=None, emit_positions=True, drop_oversized=None, min_sequence_length=1, shuffle_strategy=None, shuffle_seed=None, flush_strategy='fifo', candidate_pool_size=None, max_candidate_age=None, homogeneity='none', groups=None, placement='auto', parallelism=None)[source]¶

Add a sequence-packing operator that emits flat training records.

Each bin is emitted flat as {tokens_field: concat[+pad], "positions"?} — no packed_samples — so SampleBatch.to_training consumes it directly (positions is surfaced automatically). Uncapped wrap emits only full bins. Capped wrap and the other algorithms pad partial bins, including their final remainder, to max_length with pad_token_id. Buffered algorithms may reorder segments, and best_fit_wrap may emit a suffix before its remaining prefix. When emit_positions=True, positions still reset at every emitted segment; they do not restore source order. Only the token field and length-aligned sliceable fields survive; scalar and non-aligned payload is dropped (use pack_sequences() to keep it).

Parameters:

max_length (int) – Fixed length of every emitted record.
max_sequences_per_bin (int | None) – Optional positive bound on real segments per output bin, excluding padding. Split fragments count once per bin.
num_bins (int | None) – Number of bins to maintain per packing group. Required for first_fit/best_fit and not allowed for wrapping algorithms.
algorithm (Literal['first_fit', 'best_fit', 'wrap', 'best_fit_wrap']) – "first_fit" (default), "best_fit", "wrap", or "best_fit_wrap" (candidate-based buffered wrapping with length-based selection and at most one split per bin).
tokens_field (str) – Token field to concatenate ("auto" or explicit name).
pad_token_id (int | None) – Fill value for the token field when padding partial bins (aligned fields pad with 0). Required for first_fit, best_fit, best_fit_wrap, and capped wrap; unused for uncapped wrap. Any embeddable id works — to_training masks the pad tail from the loss by position (meta.padding_length), not by id.
emit_positions (bool) – Include the positions array marking document boundaries (cumsum(positions == 0) - 1 → doc ids). Defaults to True; set False for classic concatenated blocks.
drop_oversized (bool | None) – Whether first/best should drop records longer than max_length. Defaults to True for first/best and False for wrap/best_fit_wrap, which split instead.
min_sequence_length (int) – Remaining capacity below which a bin is emitted.
shuffle_strategy (Literal['random', 'length', None]) – Strategy for ordering sequences before packing.
shuffle_seed (int | None) – Seed for random shuffling when shuffle_strategy=”random”.
flush_strategy (Literal['fifo', 'fullest']) – “fifo” (default) or “fullest” when num_bins is reached.
candidate_pool_size (int | None) – best_fit_wrap only — candidate lookahead per packing group. May be exceeded until the pool contains max_length tokens or reaches the segment cap. Defaults to 1024.
max_candidate_age (int | None) – best_fit_wrap only — candidate arrivals before a still-buffered record is force-placed. Defaults to 8 * candidate_pool_size.
homogeneity (Literal['none', 'group', 'full']) – If "full", each packed sample stays within a single mixing domain (mixture component); "group" keeps it within a single groups group of domains; "none" (default) mixes freely. See PackSequences.__init__().
groups (DomainGroups | Mapping[str, Sequence[str]] | None) – Required for homogeneity="group" — a zephon.ops.DomainGroups or {group: [component, ...]} mapping naming which mixing domains may share a packed sample. Members are mixture-component names (the source’s component_ids), validated against them at build time.
placement (str) – Placement strategy for this operator.
parallelism (int | None) – Worker count for materializing packed-bin payloads. Bin assignment remains serial, so output is unchanged.

Return type:

Pipeline

pack_sequences(max_length, *, num_bins=None, max_sequences_per_bin=None, algorithm='first_fit', tokens_field='auto', length_fn=None, drop_oversized=None, min_sequence_length=1, shuffle_strategy=None, shuffle_seed=None, flush_strategy='fifo', pack_payloads='keep_list', candidate_pool_size=None, max_candidate_age=None, homogeneity='none', groups=None, placement='auto', parallelism=None)[source]¶

Add a sequence-packing operator that preserves segment boundaries.

Each bin is emitted as {"packed_samples": [seg0, seg1, ...]} — an ordered list of constituent records (first_fit/best_fit) or record slices (wrap/best_fit_wrap). Each list element remains a distinct segment; buffered algorithms may reorder segments, and best_fit_wrap may emit a suffix before its remaining prefix. pack_payloads merges that list (default keeps it as-is). For flat, tensor-ready training records with positions, use pack_flat() instead.

Parameters:

max_length (int) – Maximum length for packed bins.
max_sequences_per_bin (int | None) – Optional positive bound on segments per bin; split fragments count once per bin. Capped wrap emits partial bins, including the final remainder.
num_bins (int | None) – Number of bins to maintain per packing group. Required for first_fit/best_fit and not allowed for wrapping algorithms.
algorithm (Literal['first_fit', 'best_fit', 'wrap', 'best_fit_wrap']) – "first_fit" (default), "best_fit", "wrap", or "best_fit_wrap" (candidate-based buffered wrapping with length-based selection and at most one split per bin).
tokens_field (str) – Token field to slice ("auto" or an explicit name); consumed by wrap/best_fit_wrap. first/best keep whole payloads.
length_fn (Callable[[SampleRecord], int] | None) – Optional callable measuring packing length, for first_fit/best_fit only (e.g. a precomputed length field with no token field to slice). None measures len(payload[tokens_field]). Not allowed with wrap or best_fit_wrap (length is the sliced field’s length).
drop_oversized (bool | None) – Whether first/best should drop records longer than max_length. Defaults to True for first/best and False for wrap/best_fit_wrap, which split instead.
min_sequence_length (int) – Remaining capacity below which a bin is emitted.
shuffle_strategy (Literal['random', 'length', None]) – Strategy for ordering sequences before packing (“random”, “length”, or None).
shuffle_seed (int | None) – Seed for random shuffling when shuffle_strategy=”random”.
flush_strategy (Literal['fifo', 'fullest']) – “fifo” (default) flushes oldest bins first; “fullest” flushes bins with the smallest remaining capacity first.
pack_payloads (str | Callable[[list[Any]], Any]) – How to merge the segment list. “keep_list” (default), “torch_tensor”, “numpy_array”, or a custom callable taking list[Any].
candidate_pool_size (int | None) – best_fit_wrap only — candidate lookahead per packing group. May be exceeded until the pool contains max_length tokens or reaches the segment cap. Defaults to 1024.
max_candidate_age (int | None) – best_fit_wrap only — candidate arrivals before a still-buffered record is force-placed. Defaults to 8 * candidate_pool_size.
homogeneity (Literal['none', 'group', 'full']) – If "full", each packed sample stays within a single mixing domain (mixture component); "group" keeps it within a single groups group of domains; "none" (default) mixes freely. See PackSequences.__init__().
groups (DomainGroups | Mapping[str, Sequence[str]] | None) – Required for homogeneity="group" — a zephon.ops.DomainGroups or {group: [component, ...]} mapping naming which mixing domains may share a packed sample. Members are mixture-component names (the source’s component_ids), validated against them at build time.
placement (str) – Placement strategy for this operator.
parallelism (int | None) – Worker count for materializing packed-bin payloads. Bin assignment remains serial, so output is unchanged.

Return type:

Pipeline

prefetch(buffer_size=1024, parallelism=None, *, placement='local')[source]¶

Add a prefetch operator to warm the cache before fetching.

The prefetch operator looks ahead in the sample stream and downloads shards to the local cache before they’re needed by FetchOp. This significantly reduces fetch latency when loading from remote storage (S3, GCS).

The prefetch node is inserted before the fetch node in the pipeline.

Parameters:

buffer_size (int) – Number of samples to buffer for lookahead (default: 1024). Larger values provide more prefetch opportunities but use more memory.
parallelism (int | None) – Number of concurrent worker threads for downloads (default: 4). Higher values increase download parallelism.
placement (str) – Placement hint for the prefetch operator (default: “local”).

Returns:

Self for method chaining.

Return type:

Pipeline

Example

>>> pipeline = (
...     Pipeline(work_source)
...     .prefetch(buffer_size=2048, parallelism=8)
...     .decode_text()
...     .tokenize(field="text")
...     .batch(32)
... )

Note: The fetch operator is automatically added by Pipeline.__init__, so you don’t need to call .fetch() explicitly. The prefetch operator is inserted before the implicit fetch operator.

Note

Prefetch is most effective with: - Remote storage (S3, GCS) with high download latency - Sequential or predictable shard access patterns - Large shards where download time is significant

For local storage or already-cached data, prefetch has minimal benefit.

prefetch_timing_snapshot()[source]¶

Return prefetch timings for the inline engine, or None if absent.

Data is collected only in ExecutionTrackingMode.NODES.

Return type:: PrefetchTimingSummary | None

preflight_tokenizers(*, strict=True)[source]¶

Load and validate every tokenizer configured on the pipeline.

Each tokenizer operator is probed through a deep copy, leaving the pipeline’s operator state unchanged. This may download tokenizer files.

Parameters:: strict (bool) – If True, raise ValidationError when any tokenizer fails to initialize.
Returns:: A ValidationReport containing one TOKENIZER_PREFLIGHT_FAILED error per failure.
Return type:: ValidationReport

restore(ckpt)[source]¶

shuffle(buffer_size=None, *, seed=0, algorithm='streaming', warmup_growth=1.5, placement='auto', parallelism=None)[source]¶

Insert a deterministic shuffle buffer.

Parameters:

buffer_size (int | None) – Number of samples to buffer for shuffling. If None, a default of 8192 is used.
seed (int) – RNG seed for the shuffle.
algorithm (Literal['streaming', 'block', 'block_warmup']) – "streaming" uses a deterministic reservoir (default). "block" uses legacy tumbling blocks. "block_warmup" uses tumbling blocks with a post-flush block-size ramp.
warmup_growth (float) – Multiplicative block-size growth per step for algorithm="block_warmup" (ignored otherwise). Default 1.5; lower values ramp more smoothly and reach the full window later.
placement (str) – Placement hint for this operator.
parallelism (int | None) – Override default parallelism for this operator.

Returns:

Self for method chaining.

Return type:

Pipeline

stateful_transform(name, *, init_state, push, flush=None, should_flush=None, transform=None, placement='auto', parallelism=1, indexable=False, preserves_cursor_order=True)[source]¶

Add a stateful transformation with custom accumulation logic.

This is a higher-level alternative to subclassing BaseOp directly. State management runs on the pump thread (serial); the optional transform runs in parallel workers for expensive computation.

Execution model:

push/flush: Run on pump thread (serial) for state management
transform: Runs in parallel workers for expensive per-item processing

This split enables patterns like “deduplicate (serial) then encode (parallel)”.

State is partitioned per lane — each lane keeps its own state instance and your callbacks see one lane at a time, so an epoch-boundary flush resets only that lane (required for deterministic replay when one engine owns several lanes). The per-lane lifecycle:

A lane’s state is lazily initialized on its first record via init_state()
push(state, items) is called with one lane’s records -> (new_state, outputs)
If should_flush returns True, that lane’s flush runs and its state resets
On stream end (all lanes) or a lane’s epoch boundary, flush() emits that lane’s remaining buffered items
Each output item goes through transform (if provided) in parallel

Parameters:

name (str) – Operator name for debugging/metrics.
init_state (Callable[[], _S]) – Factory that creates initial state (called lazily per lane).
push (Callable[[_S, list[SampleRecord]], tuple[_S, list[SampleRecord]]]) – Called with (lane_state, one lane’s records) -> (new_state, outputs). Outputs are emitted immediately; state carries forward for that lane.
flush (Callable[[_S], list[SampleRecord]] | None) – Optional. Emits a lane’s remaining buffered items at end-of-stream. In non-monotonic pipelines, also called per lane mid-stream at that lane’s epoch boundary; the lane’s state is re-initialized afterward.
should_flush (Callable[[_S], bool] | None) – Optional. If returns True, triggers early flush and state reset.
transform (Callable[[list[SampleRecord]], list[SampleRecord]] | None) – Optional. Batch-level transform that runs in parallel workers. Receives the full batch from the accumulator, preserving batch structure for efficient GPU processing, vectorized ops, etc.
placement (str) – Placement hint for this operator.
parallelism (int) – Worker parallelism. Use >1 when transform is expensive.
indexable (bool) – Whether this operator preserves indexability (default False). Set True only if the transform is 1:1 and deterministic.
preserves_cursor_order (bool) – Whether outputs maintain monotone cursor order. Set False when the transform reorders or packs items (e.g., shuffle, bin-packing). This affects which chunk-eviction path the engine uses for the entire plan.

Returns:

Self for method chaining.

Return type:

Pipeline

Example - Custom batching by token count:

>>> def accumulate_by_tokens(state, items, max_tokens=4096):
...     buffer, token_count = state["buffer"], state["tokens"]
...     outputs = []
...     for item in items:
...         item_tokens = len(item.payload["token_ids"])
...         if token_count + item_tokens > max_tokens and buffer:
...             outputs.extend(buffer)
...             buffer, token_count = [], 0
...         buffer.append(item)
...         token_count += item_tokens
...     return {"buffer": buffer, "tokens": token_count}, outputs
...
>>> pipeline.stateful_transform(
...     "batch_by_tokens",
...     init_state=lambda: {"buffer": [], "tokens": 0},
...     push=lambda s, items: accumulate_by_tokens(s, items),
...     flush=lambda s: s["buffer"] if s["buffer"] else [],
... )

Example - Deduplicate (serial) then batch encode (parallel):

>>> def batch_encode(records):
...     # Process whole batch efficiently (e.g., GPU batching)
...     for r in records:
...         r.payload["encoded"] = encode(r.payload["text"])
...     return records
...
>>> pipeline.stateful_transform(
...     "dedupe_and_encode",
...     init_state=lambda: set(),
...     push=lambda seen, items: (
...         seen | {i.payload["id"] for i in items},
...         [i for i in items if i.payload["id"] not in seen]
...     ),
...     transform=batch_encode,  # processes whole batch in parallel
...     parallelism=8,
... )

to_indexable_torch_dataset()[source]¶

Return type:: _DatasetProto

to_torch_dataset(stateful=True)[source]¶

Return type:: _IterableDatasetProto

tokenize(tokenizer=None, tokenizer_id=None, *, field, missing_field='error', add_attention_mask=True, max_length=None, padding=False, truncation=False, return_tensors=None, split_long_samples=False, use_fast=True, preserve_upstream_payload=False, special_tokens='bos_eos', bos_token_id=None, eos_token_id=None, placement='auto', parallelism=None)[source]¶

Return type:: Pipeline

tokenize_chat(tokenizer=None, tokenizer_id=None, *, eos_token=None, field='messages', max_length=None, chat_template=None, apply_chat_template=True, span_source='auto', loss_on_last_turn_only=False, chat_template_kwargs=None, tools_field='tools', enable_thinking_field='enable_thinking', mask_field_out='loss_mask', preserve_upstream_payload=False, placement='auto', parallelism=None)[source]¶

Tokenize chat conversations into ids plus a loss mask.

Full contract: one-shot render+tokenize with assistant spans from {% generation %} tags, a guarded prefix-diff fallback for untagged templates, or exact spans in the no-template path (apply_chat_template=False).

Return type:: Pipeline

validate(*, strict=False)[source]¶

Run the custom-op validation harness against the current graph.

Walks every user-supplied operator (those built by add_op()) and runs a smoke test that checks: process_many is stateless across calls, the accumulator conserves samples on push+flush, flush(reset=True) leaves the accumulator in a fresh state, and has_pending_data() mirrors reality. See the zephon.validation module for the full check list and the “Accumulators and Operators” page in the published documentation for the contracts being checked.

Built-in ops are skipped (framework code, tested separately); a graph with no user ops returns an empty report immediately.

Parameters:: strict (bool) – If True, raise ValidationError when any error is found. Default False — caller inspects report.ok.
Returns:: A ValidationReport with structured issues.
Return type:: ValidationReport

Validation¶

Validation harness for user-supplied operators and accumulators.

The framework guarantees deterministic parallel execution only when the callables passed to Pipeline.add_op honor several contracts that the type system cannot enforce: process_many must be stateless across calls, the accumulator must drain on flush() and reset on flush(reset=True), and so on. This module runs a synthetic-sample smoke test that catches the most common violations before the engine starts. See the “Accumulators and Operators” page in the published documentation for the contracts being checked.

The harness is intentionally a smoke test, not a proof: it cannot catch violations that only surface after many iterations or under specific data shapes. Errors here are real bugs; passing this harness is necessary but not sufficient for correctness.

Unlike the rest of zephon (custom exception subclasses + warnings.warn), this module aggregates findings into a ValidationReport of string-coded Issue records. Aggregation lets the user see every problem in their op at once; the tradeoff is identification by string equality on issue.code rather than by isinstance.

class Issue(severity, code, op_name, message, doc_link='https://datologyai.github.io/zephon/understanding/accumulators_operators.html')[source]¶

Bases: object

A single validation finding.

severity: str¶

code: str¶

op_name: str¶

message: str¶

doc_link: str = 'https://datologyai.github.io/zephon/understanding/accumulators_operators.html'¶

exception ValidationError(report)[source]¶

Bases: RuntimeError

Raised when a strict-mode validation entry point finds at least one error.

Raisers: Pipeline.validate(strict=True), strict auto-validation, and Pipeline.preflight_tokenizers(strict=True).

class ValidationReport(issues=<factory>)[source]¶

Bases: object

Result of running the validation harness against a pipeline.

issues: list[Issue]¶

property ok: bool¶: True when no errors were reported (warnings are allowed).

property errors: list[Issue]¶: Subset of issues with severity error.

property warnings: list[Issue]¶: Subset of issues with severity warning.

format()[source]¶

Human-readable multi-line summary suitable for raising.

Return type:: str

preflight_tokenizers(pipeline)[source]¶

Return tokenizer initialization failures for a pipeline.

Each tokenizer operator is deep-copied before its lazy initialization runs, leaving pipeline-owned operators unchanged. This may download tokenizer files.

Parameters:: pipeline (Pipeline) – Pipeline to inspect.
Returns:: A ValidationReport containing one TOKENIZER_PREFLIGHT_FAILED error per failure.
Return type:: ValidationReport

Runtime options¶

Runtime options for Pipeline.options() (RuntimeOptions, IpcTransport).

class RuntimeOptions(runner=None, run_id='default_run_id', per_stage_runner=<factory>, mp_context=<multiprocessing.context.SpawnContext object>, worker_allocation='fit_to_ops', max_workers=None, deterministic=True, prefetch_batches=None, default_stage_prefetch=None, per_stage_prefetch=<factory>, op_queue_capacity=None, ipc_transport='socketpair', ipc_buffer_bytes=524288, mixture_config=None, io_options=<factory>, allow_latency_flush_in_deterministic=True, flush_every_k_chunks=None, shutdown_mode='graceful', auto_validation='strict', mtp_mode=False, mtp_buffer=None, mtp_buffer_bytes=8388608, mtp_prefetch=None, mtp_auto_checkpoint=True, world_size=1, global_rank=0, dp_degree=None, dp_group_id=None, canonical_replicas=None, mapping_strategy=None, coalesce_tensors=True, shm_min_size=4096, max_worker_retries=3, aggregate_dir=None, aggregate_timeout_s=180.0, aggregate_serializer='msgpack', aggregate_compressor='zstd', execution_tracking=ExecutionTrackingMode.OFF, metrics_sink_config=None)[source]¶

Bases: object

User-tunable knobs that influence how the engine constructs runners.

Numeric tuning knobs default to None, resulting in sensible defaults via resolve_* helpers that form a dependency cascade. Each resolver calls the ones above it when its own field is unset:

max_workers              ← os.cpu_count(), clamped [4, 16]
    └── prefetch_batches             = max(8, 2 × max_workers)
            ├── op_queue_capacity        = max(8, prefetch_batches)
            ├── mtp_buffer               = max(4, prefetch_batches // 2)
            │       └── mtp_prefetch         = max(4, mtp_buffer // 10)
            └── default_stage_prefetch   = max(4, prefetch_batches // 4)  [process]
                                           max(8, prefetch_batches // 2)  [threads/inline]

For example, max_workers=24 naturally bumps the prefetch / queue / MTP buffers without the user having to set them too, while prefetch_batches=64 overrides the tail buffer alone and lets the rest auto-derive from it.

By default, max_workers is derived from the host’s CPU count, but is clamped to a range of [4, 16]. For any large deployment this will hit the upper limit, but the clamp allows Zephon to run reasonably on small machines as well.

runner: str | None = None¶

run_id: str = 'default_run_id'¶

per_stage_runner: dict[int, str]¶

mp_context: Any = <multiprocessing.context.SpawnContext object>¶

worker_allocation: Literal['fit_to_ops', 'per_stage_fixed', 'global', 'autotune'] = 'fit_to_ops'¶

max_workers: int | None = None¶

deterministic: bool = True¶

prefetch_batches: int | None = None¶

default_stage_prefetch: int | None = None¶

per_stage_prefetch: dict[int, int]¶

op_queue_capacity: int | None = None¶

ipc_transport: Literal['socketpair', 'pipe'] = 'socketpair'¶

ipc_buffer_bytes: int = 524288¶

mixture_config: MixtureReadConfig | None = None¶

io_options: StoreOptions¶

allow_latency_flush_in_deterministic: bool = True¶

flush_every_k_chunks: int | None = None¶

shutdown_mode: Literal['graceful', 'hard'] = 'graceful'¶

auto_validation: Literal['strict', 'warn', 'off'] = 'strict'¶

mtp_mode: bool = False¶

mtp_buffer: int | None = None¶

mtp_buffer_bytes: int = 8388608¶

mtp_prefetch: int | None = None¶

mtp_auto_checkpoint: bool = True¶

world_size: int = 1¶

global_rank: int = 0¶

dp_degree: int | None = None¶

dp_group_id: int | None = None¶

canonical_replicas: int | None = None¶

mapping_strategy: Literal['contiguous', 'interleaved'] | None = None¶

coalesce_tensors: bool = True¶

shm_min_size: int = 4096¶

max_worker_retries: int = 3¶

aggregate_dir: str | None = None¶

aggregate_timeout_s: float = 180.0¶

aggregate_serializer: str = 'msgpack'¶

aggregate_compressor: str = 'zstd'¶

execution_tracking: ExecutionTrackingMode = 'off'¶

metrics_sink_config: MetricsSinkConfig | None = None¶