zephon.ops

Operator authoring kit: the contracts you subclass or pass to add_op().

Built-in operators are configured via the Pipeline builder methods (.shuffle(), .batch(), .tokenize(), …).

class Accumulator

Bases: ABC, Generic[T]

Base class for accumulators that run on the pump thread.

Accumulators consume inputs in deterministic order, maintain any cross-invocation state (buffers/bins), and emit ready invocation batches.

Contract for deterministic parallelism: - All cross-invocation, output-producing state must live in the accumulator - Worker process_many must be deterministic given its input batch - Accumulators are invoked only on the pump thread (no concurrent calls)

abstractmethod push_many(elems)

Consume elements in stream order and return 0..N ready invocation batches.

Parameters:

elems (Sequence[T]) – Input elements to accumulate.

Returns:

List of ready batches to dispatch to workers. May be empty if more elements are needed before a batch is ready.

Return type:

list[tuple[list[T], int]]

abstractmethod flush(*, reset=False, lane_id=None)

Emit any remaining ready batches.

Called both at upstream close (default, reset=False) and mid-stream by flush sentinels (reset=True). Mid-stream flushes must fully reset internal state so the accumulator is indistinguishable from a freshly constructed instance.

Flush sentinels are injected per lane, so a mid-stream flush is scoped to lane_id: only that lane’s state is emitted and reset, leaving other lanes’ epochs untouched. Flushing every lane at one lane’s epoch boundary corrupts the others and breaks deterministic replay across checkpoint/restore. lane_id=None flushes all lanes (used at upstream close, when every lane is done).

Parameters:

• reset (bool) – True when triggered by a flush sentinel mid-stream (the accumulator must reset to fresh state), False when called at upstream close.

• lane_id (int | None) – When set, flush only this lane. When None, flush all lanes.

Returns:

List of remaining batches. These are typically partial batches that were waiting for more elements.

Return type:

list[tuple[list[T], int]]

property reads_payload: bool

Whether push_many() may access element payloads.

When True, the runner resolves lazy SHM payloads before calling push_many(). Default False — most accumulators only read metadata (e.g. lane_id).

has_pending_data(lane_id=None)

Return True if buffered data would be emitted on flush().

Default implementation returns False. Subclasses with internal buffers should override this to return True when they have pending data.

Parameters:

lane_id (int | None) – When set, report only this lane’s pending state. When None, report whether any lane has pending data.

Return type:

bool

try_epoch_reset(boundary_chunk_id, lane_id=None)

Attempt a delayed epoch reset for stalled flush sentinels.

Called when a flush sentinel was stalled (held behind buffered data) and the runner wants to check whether the accumulator can now be reset.

Returns True if and only if future outputs are equivalent to a fresh accumulator on the retained suffix — i.e., the accumulator’s behavior going forward is indistinguishable from one that was freshly constructed and fed only the records still in the buffer. When returning True, the method must also perform any necessary state reset (e.g., clearing SWRR emission history).

The default implementation is conservative: it returns True only when there is no pending data at all. Accumulators that buffer SampleRecord elements can override this to check chunk_id directly and return True as soon as no pre-boundary records remain, even if post-boundary records are still buffered.

Parameters:

• boundary_chunk_id (int) – The _boundary_cid from the stalled flush sentinel. Records with chunk_id < boundary_chunk_id are pre-boundary (old epoch).

• lane_id (int | None) – The stalled sentinel’s lane. When set, only this lane’s buffered records gate the reset; other lanes are independent.

Returns:

True if the reset was performed and the sentinel can be released. False if pre-boundary records still exist in the buffer.

Return type:

bool

class BaseOp

Bases: ABC

Base class for operators — both built-in and user-authored.

Attach a subclass to a pipeline via the instance form of zephon.Pipeline.add_op():

pipeline.add_op(MyOp(...))

Built-in operators subclass BaseOp and are attached internally; the same class hierarchy is what user code subclasses for ops that need full lifecycle control (__init__ / setup / traits / accumulator / process_many). The kwargs form of add_op is the lighter alternative for stateless transforms — it wraps user callables in an internal BaseOp subclass, so the same contract documented here applies either way.

Operators are stateless. Instance attributes set in __init__ are read-only configuration or shared resources (a codec table, a threshold, a tokenizer name). process_many must not depend on or mutate state carried across calls — any cross-invocation or per-lane state (buffers, counters, open bins) lives in the operator’s accumulator, not on the op instance. This is the invariant that lets the runtime fan process_many out across parallel workers deterministically.

Fan-out operators must preserve SampleMeta.sample_id and derive child lineage with helpers from zephon.ops.children. This preserves deterministic ordering and replay.

Lifecycle: __init__ vs setup

__init__ runs once on the user’s main process: store config and lightweight references, and leave runtime fields defaulted to None. The op instance is then deep-copied per parallel worker; on the process and Ray runners it is also cloudpickled across the process boundary, so anything assigned in __init__ must survive cloudpickle on those runners.

setup runs once per worker after the deep copy (and, on cross-process runners, after the cloudpickle round trip). It is the only place to read services from OpContext, construct heavy or non-picklable resources (tokenizers, model handles, open file descriptors), and reset per-worker state. Resources built in setup never need to survive pickling.

The convention in this repo is to construct tokenizers (and similar heavy resources) in setup, not __init__:

class MyOp(BaseOp):
    def __init__(self, tokenizer_name: str):
        super().__init__()
        self._tokenizer_name = tokenizer_name  # picklable config
        self._tokenizer = None  # built per-worker in setup()

    def traits(self) -> OpTraits:
        return OpTraits(preserves_cursor_order=True)

    def setup(self, ctx):
        super().setup(ctx)
        self._tokenizer = load_tokenizer(self._tokenizer_name)

    def process_many(self, elems):
        # Records in, records out: rebuild payload, preserve meta.
        for e in elems:
            e.payload = {**e.payload, "ids": self._tokenizer.encode(e.payload["text"])}
        return elems

pipeline.add_op(MyOp("gpt2"))

Defaults supplied:

• setup records ctx.stage_info as self.stage_info. Overrides should call super().setup(ctx) before initializing worker resources.

• accumulator — returns PassthroughAccumulator so each upstream micro-batch is forwarded as one ready batch. Override to enable size-based per-lane batching via CountingAccumulator or any custom Accumulator.

• process_one — wraps the element in a single-item list and routes it through process_many. Override only when the single-element fast path differs (rare).

Abstract — every subclass must implement:

• process_many — the workhorse transform.

• traits — must return an OpTraits with preserves_cursor_order declared; the subclass must declare True (1:1 maps, payload transforms, non-reordering filters) or False (reorders, shuffles, packs).

Instantiating a subclass that doesn’t override both process_many and traits raises TypeError at construction time.

classmethod plan_identity()

Return the stable operator identity used in plan fingerprints.

Return type:

str

setup(ctx)

Record plan metadata before worker processing begins.

The runner calls this once per worker after copying or deserializing the operator. Overrides should call super().setup(ctx).

abstractmethod traits()

Return type:

OpTraits

accumulator(*, deterministic, ctx)

Create the accumulator that defines worker batch boundaries.

The accumulator runs serially and owns state shared across process_many calls.

Parameters:

• deterministic (bool) – Whether to disable behavior such as timed flushing.

• ctx (dict[str, Any]) – Runtime services available to the accumulator.

Returns:

A new accumulator.

Return type:

Accumulator[Any]

process_one(elem)

Return type:

list[Any]

abstractmethod process_many(elems)

Return type:

list[Any]

validation_samples()

Optional records for the validation harness.

Override to opt this op into the validator’s runtime check suite. Returning a small batch (~12 records across a couple of lanes) lets the harness run determinism, cross-call-state, self-mutation, state-diff, and sample-identity probes against process_many.

Why instance-form ops are gated this way: the validator never invokes setup, so probing a BaseOp subclass whose process_many depends on setup-built resources (the typical pattern — tokenizers, model handles, etc.) would raise spuriously. Returning None (the default) tells the validator to skip the runtime checks for this op and surface a single OP_INSTANCE_RUNTIME_CHECKS_SKIPPED warning instead of crashing the pipeline under auto_validation="strict". Static-AST checks (self-writes, non-deterministic stdlib calls) run unconditionally either way.

Returning an empty list, a list of the wrong type, or raising surfaces an OP_VALIDATION_SAMPLES_FACTORY_FAILED warning and falls back to the skip path — the validator never dies on a buggy factory.

Tip: include at least two distinct lane_id values so the cross-call-state probe (which derives an A-vs-B sample-id variant from your records) and the lane-purity accumulator checks remain meaningful. Records must work without setup having run — pre-tokenize / pre-encode your payload rather than relying on the resources the real setup would build.

Return type:

list[SampleRecord] | None

class CountingAccumulator(max_batch, max_latency_ms=None, *, drop_last=False)

Bases: Accumulator[T]

Accumulator that batches elements by count, optionally with time-based flushing.

Elements are routed to per-lane buffers keyed by lane_of(), so each emitted batch is lane-pure. lane_of extracts the lane from whichever element type the accumulator carries (a SampleRecord, or a raw EngineSample tuple at the fetch stage), so the buffer key is always the lane id — which is what lets a per-lane flush sentinel reset one lane alone.

When max_latency_ms is None, this is a pure count-based accumulator suitable for deterministic execution. When set, batches are also flushed after the timeout expires (per lane).

When drop_last is True, flush() discards partial batches instead of emitting them.

has_pending_data(lane_id=None)

Return True if buffered elements remain (in lane_id if given).

Return type:

bool

push_many(elems)

Accumulate elements and emit based on count or time thresholds.

Return type:

list[tuple[list[T], int]]

flush(*, reset=False, lane_id=None)

Emit any remaining buffered elements (or discard if drop_last).

Buffers are keyed by lane, so lane_id is the buffer key for a per-lane flush.

Return type:

list[tuple[list[T], int]]

class DomainGroups(groups)

Bases: object

Groups of mixing-domain (component) names that may share a packed sample.

Members are mixture-component names (the source’s component_ids) — dataset names in StaticMixtureWorkSource, but not necessarily so (a hierarchical source could map several datasets to one component). Validated at construction into an owned map, so mutating the caller’s input afterward can’t change it.

Parameters:

groups (Mapping[str, Sequence[str]]) – Maps a group name to the component names it contains. A component may appear in at most one group. Components absent from every group form their own singleton group.

Example

>>> groups = DomainGroups({"code": ["python", "java"], "web": ["c4"]})
>>> groups.to_member_map()
{'python': 'code', 'java': 'code', 'c4': 'web'}

to_member_map()

Flat {component: group} map (grouped only) — the projection packing consumes.

Return type:

dict[str, str]

validate_against(known_components)

Raise if any grouped component is absent from known_components.

class OpContext(services, stage_info=None)

Bases: object

Container exposing runner-provided services to operator instances.

Plan metadata lives on the context so setup overrides keep a stable signature.

get(key, default=None)

Fetch a service by name, returning default when unavailable.

Return type:

Any

class OpTraits(*, indexable=True, preserves_cursor_order, parallelism=1, batch_shape_sensitive=False, requires_serial_state=False)

Bases: object

Static capabilities an operator advertises to the planner.

indexable — whether the operator preserves indexability through the plan.

preserves_cursor_order — whether the operator preserves per-lane cursor order (no reordering across chunk/offset/lineage). Required — every operator author must make an explicit True (1:1 maps, payload transforms, non-reordering filters) or False (reorders, shuffles, packs) call.

parallelism — suggested parallelism for the operator when not overridden.

batch_shape_sensitive — if True, the operator’s outputs can depend on how inputs are grouped into micro-batches (e.g., per-batch RNG or statistics). In deterministic mode, stages that contain at least one such operator will have time-based flush disabled to preserve strong determinism. When False, ordering determinism suffices and latency flush may be kept for performance.

requires_serial_state — if True, the operator maintains cross-invocation state (e.g., buffers) that must be confined to a single operator instance for determinism. In deterministic mode, operators with this trait will automatically run with parallelism=1. This prevents nondeterministic behavior when multiple worker instances would each maintain separate buffers.

indexable: bool = True

preserves_cursor_order: bool

parallelism: int = 1

batch_shape_sensitive: bool = False

requires_serial_state: bool = False

class PassthroughAccumulator

Bases: Accumulator[T]

Accumulator that emits exactly the incoming microbatch as one invocation batch.

This is the default accumulator for stateless operators. It preserves the upstream batching without any additional buffering.

push_many(elems)

Pass through the input microbatch as-is.

Return type:

list[tuple[list[T], int]]

flush(*, reset=False, lane_id=None)

No buffered state to flush.

Return type:

list[tuple[list[T], int]]

class StageInfo(stage_index=-1, stage_name='', op_index=-1, collect_stats=False)

Bases: object

Plan metadata passed through OpContext.stage_info.

Default values indicate that setup has not run. stage_index and stage_name identify the stage, op_index identifies the operator within it, and collect_stats enables per-operator metrics.

stage_index: int = -1

stage_name: str = ''

op_index: int = -1

collect_stats: bool = False

pack_meta(primary_cursor, contributors, *, lane_id, component_sample_counts, component_token_counts=None, tags=None)

Build metadata for a packed record that merges multiple contributors.

primary_cursor is the replay identity for the packed record and must be unique per lane. contributors lists all contributors included in the pack; any contributor that completes a base offset must set is_last_child=True.

component_sample_counts aggregates how many original samples from each component are included in this pack. For example, if packing 3 samples from component 0 and 2 from component 1, this would be {0: 3, 1: 2}.

component_token_counts optionally provides token counts per component, computed when packing happens after tokenization. If packing before tokenize, pass None and ensure_mixture will fall back to distributing by sample counts.

Return type:

SampleMeta

spawn_child(parent, child_idx, *, is_last_child=False, tags=None)

Create metadata for a child derived from parent.

• Assigns deterministic child lineage for identity/replay.

• Contributors:

  • If the parent has no contributors (single-base record), emit one contributor tied to the child cursor and honor is_last_child.

  • If the parent already has contributors (e.g., packed input), propagate them to the child; if is_last_child=True, mark all propagated contributors as closing.

• Raises if you try to mark a child as closing a base offset that is already marked closed by the parent metadata (e.g., parent contributors already contain an is_last_child=True ref or the parent is a tombstone).

Return type:

SampleMeta

tombstone_meta(ref, lane_id)

Emit a tombstone that marks a base offset complete without payload.

Return type:

SampleMeta

tombstones_for_record(record)

Emit tombstone records for every closing contributor in record.

Each record tracks which chunk offsets it contributes to via contribution_refs(). Only refs with is_last_child=True need a tombstone — intermediate children (from split/spawn) don’t close the base offset, so the engine doesn’t need a signal for them.

Return type:

list[SampleRecord]