reduce_merge

Purpose

Fold partial outputs (often after multi-worker) into final datasets per the DRTML reduce_merge plan.

Layers

Layer	Role
processing/reduce_merge/run.py	lifecycle, work-unit plan, metrics ticks
processing/reduce_merge/bucket_plan.py	shuffle buckets (or flat group) → work units
processing/reduce_merge/parallel_jobs.py	workers via job_map + open_job_worker_session; write lanes + manifest merge
processing/job_map.py	shared offline map / round-robin assign
framework/reduce/build.py	one-pass multi-output sum; optional bucket_ids
storage/run/reduce_merge_*	I/O / SQL merge adapters

Storage must not import processing metrics. Processing passes an on_input_ticks callback.

Pattern

pattern: reduce_merge
execution:
  entry: drtoller.framework.processing.reduce_merge.run:execute_reduce_merge
reduce_merge:
  entry: drtoller.framework.processing.reduce_merge.run:execute_reduce_merge
  upstream_runs_param: reduce.upstream_runs
  jobs:
    - input: partial_counts
      output: merged_counts
      group_by: [kind, token_id]
      sum_columns: [count]

Dispatch → run_reduce_merge.

jobs[] is the preferred form and permits multiple independent reductions in one run. The legacy single-job inputs / outputs / group_by / sum_columns fields are still compiled when jobs is empty.

Jobs that share the same input key are one input group: the build path scans that input once and accumulates all group jobs (one-pass multi-output).

Parallel work units

With early-shuffle upstream (bucket-NNN/), each shuffle bucket is a work unit. Flat layouts keep one unit per input group.

When work_units > 1 and runtime.parallelism.backend is processes or ray with workers > 1:

• units are round-robin assigned via job_map.assign_round_robin (same idea as partition_loop shard assign);
• each worker uses a private write lane (seed_output_shard_id);
• driver merges output manifests (manifest_merge / persist_merged_output_manifests).

So workers=8 can keep eight CPUs busy on one large dataset (64 buckets → 8×8), not just one task per dataset.

inprocess / workers=1 / a single work unit stays sequential on the driver session.

Ray/processes also call ensure_raycluster_shape before jobs, same as partition_loop.

Live progress (Ray / processes)

Workers do not scrape Prometheus. They put incremental WorkerMetricDelta on the shared metric channel (processing/reduce_merge/live_metrics.py). A driver drain thread applies ticks to RunMetrics and heartbeats about every 5 s so Grafana stays live while map_jobs blocks. inprocess ticks the driver session directly.

On Ray, the metric collector actor is pinned to the head (num_cpus=0). If that actor dies during autoscaler idle-down, the run continues — it does not fail the job.

Early shuffle (parquet)

If upstream parquet outputs declared DRTML shuffle and all shards are bucketed (no flat mix), reduce_sum aggregates per shuffle bucket across upstream runs, then flushes. Peak RAM tracks roughly one bucket’s keyspace instead of the full combiner.

Producer shuffle parts grow to the part-size limit while writing. There is no driver compact after merge.

Flat (non-shuffled) layouts still use an in-memory partition combiner for that path.

Prefer fixing OOM with producer shuffle: + streaming reduce — not by forcing inprocess / workers=1.

Backend routes

• PostgreSQL input → PostgreSQL output: one SQL/procedure aggregation path.
• Parquet input → PostgreSQL output: bounded streaming reduction and bulk writes.
• Parquet input → parquet output: Python reduce_sum (with early-shuffle when available).
• Qdrant and Neo4j outputs are rejected.

framework/reduce/dispatch.py selects the backend route; storage and DB layers do not import processing metrics.