Skip to content

partition_loop

partition_loop is the default “walk the corpus” mode. Framework feeds you batches of logical rows; your processor returns logical results; framework writes datasets, flushes registries, checkpoints, and emits metrics. You never open parquet or call Prometheus.

Mental model:

WorkUnit → ProcessorBatch → batch_processor → ProcessResult → write lanes / checkpoint

Dispatch diagrams: How a run executes. Parallel workers: Parallelism.

  1. Compile manifest and merge params.
  2. Open StorageSession and set up metrics if enabled.
  3. Build the feed plan and enumerate work units.
  4. Load or resume registry state when declared.
  5. Run optional processing.init → in-memory processor_ctx.
  6. Materialize each work unit as a ProcessorBatch.
  7. Call the domain batch_processor.
  8. Drain ProcessResult rows onto write lanes.
  9. Checkpoint on the cadence declared in DRTML.
  10. Finalize manifests, metrics, and run metadata.
def process_batch(
raws: list[Any],
*,
ctx: RunContext,
params: Mapping[str, Any],
processor_ctx: Mapping[str, Any],
partition_id: str,
) -> list[ProcessResult]:
...

Each item is normally a ProcessorBatch (partition id, source run, primary rows, attach, shard / work-unit ids). Return logical rows only:

ProcessResult(
ok=True,
rows={"logical_output": [{"id": "1", "value": "x"}]},
stats={"items_total": 1},
shard_id="source-shard",
)

processing.outputs maps those logical names to physical datasets. The processor never calls storage.

Optional processing.init(ctx=…, params=…, resume=…) must return a dict. Framework may inject prefetched mapping indexes or embed-infer state. Context is reusable in-memory state — not a place to hide I/O.

When a registry is declared, framework loads prior rows on resume, lets domain update RAM state, drains via registry_flush, and checkpoints from dataset declarations. The generic runner currently expects processing.registry_flush.vocab plus processing.registry_ingest (rows derived from one logical ProcessResult.rows output). Legacy registry[] plugins still work but are deprecated.

PartitionRunProfile reads generic runtime.* (batch sizes, progress, resume, metrics, pipeline queue caps). Those keys must exist in params_defaults. Framework never interprets domain.*.

runtime.parallelism.backend selects inprocess, processes, or ray. Multi-worker means primary shards and independent worker sessions. runtime.parallelism.schedule is static (fixed bags) or pull (shared queue). Pull is for processes and ray; it does not enable checkpoint/resume. Checkpoint/resume and run-global mutable registries still require inprocess today — otherwise emit partials and reduce_merge.

On Ray, partition_loop pins KubeRay replicas/minReplicas to packed pods for the run and restores idle on exit. With schedule: pull and deflate_idle: true, replicas shrink as workers drain instead of holding the full shape until the last straggler. Postgres outputs go through one shared writer actor (compute workers enqueue batches). Worker/driver logs include phase timings (pipeline vs domain, PG vs MinIO upload, driver wait p50/p95). Local offset checkpoints are skipped so KubeRay pods do not mkdir host artifact paths. See Ray and Parallelism.

Concern Package
WorkUnit enum / materialize storage/feed
Writes, registry load/drain, checkpoint policy storage/run
Lifecycle glue processing/partition_loop
Worker transport processing/partition_loop/parallel
Prometheus emit telemetry