Limit maximum data resident in the DAG at any moment (credit-based admission control)

[brian-arnold]

Overview

There is currently no way to bound the total amount of data resident in a DAG at any given
moment. For pipelines where each in-flight dataset holds a large intermediate footprint on
disk — and where that footprint is only released by a terminal cleanup node — the async
orchestrator reliably runs "top-heavy": nearly every dataset completes the first stage before
later stages meaningfully start, so peak storage scales with the total number of submitted
datasets rather than with cluster width.

Motivating case

A 10-step pipeline, 300 datasets, each holding ~1TB of cache while in flight, released by a
terminal cleanup node. Ray autoscales to 25 workers. Submitting all 300 at once overruns
storage capacity, because ~300 datasets accumulate stage-1 cache before the cleanup node runs
for any of them.

Why it happens

AsyncPipelineOrchestrator launches every node as a concurrent asyncio.TaskGroup task
(src/orcapod/pipeline/async_orchestrator.py:181). With no concurrency limit, node 1 drains
its input channel as fast as it fills and submits essentially all work to Ray immediately.
Ray schedules roughly FIFO, so all node-1 tasks occupy the worker slots ahead of any node-2
task. The terminal cleanup node is last in the queue by construction.

The SyncPipelineOrchestrator is maximally top-heavy by design — it fully materializes each
node before starting the next (src/orcapod/pipeline/sync_orchestrator.py:84-118) — so this
feature is only meaningful for the async orchestrator.

Why existing knobs don't solve it

• Bounded channels (Channel(buffer_size=...), src/orcapod/channels.py:157) bound
  queued items on a single edge. They say nothing about how many items are simultaneously
  being worked on across all nodes.
• PodConfig.max_concurrency (src/orcapod/core/function_pod.py:544-578) is per-node.
  Setting N on each of 10 nodes still permits 10 × N datasets in flight, and node 1 keeps
  admitting a new dataset the instant it retires one — which is precisely the failure mode.
• PipelineConfig.default_max_concurrency is not plumbed through the pipeline path at all.
  AsyncPipelineOrchestrator never accepts or forwards a PipelineConfig (it takes its own
  buffer_size/error_policy, async_orchestrator.py:50-56), and
  FunctionJobNode.async_execute calls execution_pod.async_execute(...) without one
  (src/orcapod/core/nodes/function_node.py:2499), so it silently falls back to
  PipelineConfig() → unlimited. src/orcapod/pipeline/execution_context.py:20 still marks
  this integration as deferred. Today only PodConfig.max_concurrency set directly on the pod
  has any effect on the pipeline path.

Proposed approach

Credit-based admission control (the same idea as Flink's credit-based flow control or Reactive
Streams request(n)): bound the number of items resident in the DAG, and release a credit when
an item exits.

Concretely, an asyncio.Semaphore(N) owned by the orchestrator:

1. Acquire a credit before a source node emits a dataset.
2. Release the credit when that dataset exits the DAG at a terminal node.

Both hook points are wrapper-shaped, so nodes need no changes — the existing _CollectingWriter
(async_orchestrator.py:253) is the exact pattern to follow. A new PipelineConfig field
(e.g. max_datasets_in_flight: int | None = None) would expose it, which also gives a reason
to finally plumb PipelineConfig into the orchestrator.

This stays fully async — there is no barrier. The moment dataset 1 finishes cleanup, dataset
N+1 is admitted, so pipelining across stages is preserved and workers stay saturated. The
synchronous alternative (submit in fixed chunks, wait for each chunk to fully drain) is
strictly worse because every chunk pays for its slowest dataset.

Goals & Success Criteria

• A pipeline-level configuration option caps the number of datasets simultaneously resident
  in the DAG.
• With the cap set to N, peak concurrent in-flight datasets never exceeds N, regardless of
  how many datasets are submitted.
• Throughput is not barrier-limited: a new dataset is admitted as soon as any in-flight dataset
  exits, not only when a whole batch completes.
• Credits are released on every exit path, not just the happy one (see Risks).
• Existing pipelines with the cap unset behave exactly as they do today.

Scope & Boundaries

In scope:

• Async orchestrator only.
• Credit accounting for the linear, single-terminal-node case.
• Plumbing PipelineConfig into AsyncPipelineOrchestrator (currently unwired).

Out of scope (candidates for follow-up issues):

• Sync orchestrator — top-heavy by construction; the gate is not meaningful there.
• Byte-aware admission (capping on actual bytes-on-disk rather than dataset count). Dataset
  count is a reasonable first proxy.
• Fan-out/fan-in lineage refcounting, if it proves large enough to split out — see below.

Dependencies & Risks

• Leaked credits deadlock the pipeline. This is the main correctness hazard. A row can
  leave the DAG without reaching a terminal node: a data function returning None filters the
  item (handled at src/orcapod/core/nodes/function_node.py:2442-2476), and a crashed item is
  swallowed and never forwarded (src/orcapod/core/function_pod.py:561-566). Release must be
  tied to "left the DAG by any means," not "reached the terminal node." Get this wrong and the
  pipeline hangs permanently at N in-flight. This deserves more design care than the
  semaphore itself, and explicit tests for the filter and crash paths.
• Fan-out/fan-in needs more than a counter. When the DAG branches, one source dataset can
  become several downstream items across several terminal nodes, so a plain counter is wrong.
  This needs a lineage id per source row, refcounted, released at zero. The system-tag
  provenance machinery already carries per-row lineage, so the identifier exists — but this is
  meaningfully more work than the linear case.

Estimated effort

• Linear DAG, single terminal node: ~1 day (config field, semaphore, gating writer,
  terminal-drain release, tests including filter/crash paths).
• General fan-out/fan-in with lineage refcounting: ~3-5 days on top.

Workarounds available today

• Set pod_config=PodConfig(max_concurrency=N) on the first node only (supported via the
  @function_pod decorator, src/orcapod/core/function_pod.py:1227). Throttling admission at
  the head does not strictly bound total in-flight data, but it is the highest-leverage single
  knob currently available.
• Ray-side equivalent: give the first node a custom resource
  (resources={"admission": 1}) with limited cluster-wide units.
• Chunked submission: run the pipeline over fixed-size batches. Blunt but correct; results are
  DB-cached so reruns are cheap. Costs a straggler barrier per chunk.