Skip to content

feat: add pipelined exchange operators - #2670

Draft
andygrove wants to merge 6 commits into
apache:mainfrom
andygrove:pipelined-exchange
Draft

andygrove wants to merge 6 commits into
apache:mainfrom
andygrove:pipelined-exchange

Conversation

@andygrove

Copy link
Copy Markdown
Member

Which issue does this PR close?

Part of #2652. This is the second of a series of smaller PRs split out of #2657 (pipelined execution). It is independent of #2669.

Rationale for this change

In pipelined execution, every task of a job runs at once and stages stream record batches to each other instead of exchanging shuffle files. This PR adds the exchange itself: the operators at either end of a pipelined stage boundary, and the per-executor registry that connects them. Nothing plans these operators yet, so this PR can't change how existing jobs run.

#2657 describes the whole feature, its design notes and SF1000 results.

What changes are included in this PR?

  • execution_plans/pipelined/
    • PipelinedRegistry holds a bounded queue for each output partition of each running producer task, keyed by job, stage and task.
      • Broadcast outputs, which every consumer task reads in full, are kept in replay buffers instead.
      • It remembers recently removed jobs and stages, so a leftover task fails at once instead of waiting out its timeout.
    • PipelinedShuffleWriterExec is a stage root. It drives every input partition and routes batches into the queues: hash-partitioned, passthrough or broadcast.
      • Each queue ends with an explicit end marker, so a consumer can tell a finished partition from a failed producer.
      • The writer stops reading its input once every consumer has gone away.
      • It fails if no consumer opens a full queue within ballista.pipelined.registration_timeout_seconds.
    • PipelinedShuffleReaderExec is a stage leaf. It drains all of its producers concurrently: directly from the registry when the producer is on the same executor, and over Flight otherwise. Reading concurrently is what keeps bounded queues free of deadlock for tree-shaped plans.
    • tests.rs runs writers against readers in-process.
  • Protobuf and serde: PipelinedShuffleWriterExecNode, PipelinedShuffleReaderExecNode, PipelinedPartitionLocation, and the FetchPipelinedPartition action, plus BallistaClient::fetch_pipelined_partition. The Flight proxy in serving.rs forwards the new action.
  • Settings: ballista.execution.mode (blocking or pipelined), ballista.pipelined.queue_capacity and ballista.pipelined.registration_timeout_seconds. They are documented in a new section of the config reference.
  • Guards until the rest lands:
    • The scheduler rejects a job with ballista.execution.mode = pipelined ("not supported yet").
    • The executor's Flight service answers FetchPipelinedPartition with unimplemented.
    • Later PRs in the series replace both guards.

How are these changes tested?

  • Unit tests for the registry, covering open and claim, timeouts, replay, cancellation and removed jobs.
  • Unit tests for the writer and reader, run against each other in-process, covering hash, passthrough and broadcast outputs, end markers, producer failure, and early consumer drop.
  • Serde round trips and config parsing.
  • A scheduler test checks that a pipelined job is rejected and no graph is created.

Are there any user-facing changes?

There are three new session settings, documented in the config reference. Setting ballista.execution.mode = pipelined currently makes the job fail with a "not supported yet" error. The default behavior is unchanged.

Add PipelinedShuffleWriterExec, PipelinedShuffleReaderExec and the
per-executor PipelinedRegistry that connects them, with their protobuf
messages, serde, the FetchPipelinedPartition Flight action and the
ballista.execution.mode and ballista.pipelined.* settings.

Nothing plans these operators yet. The scheduler rejects jobs with
ballista.execution.mode = pipelined, and executors answer pipelined
fetches with unimplemented, until later changes wire them in.
…pers

- Lease pooled connections for the life of a remote pipelined stream, and
  build unpooled clients with new_ballista_client, as the shuffle reader
  does. Encode Flight tickets with encode_protobuf.
- Compute the writer's routing once, drop unused builder and getter
  methods, and let JoinSet's drop abort the remaining inputs.
- Track unclaimed queues with a counter instead of rescanning every
  partition on each open, share one bounded set for closed jobs and
  stages, and stop cloning per batch in the queue stream.
- Keep a broadcast reader whole when restricting partitions, forget a
  consumer's waiting slot when it times out, and document the minimums of
  the pipelined settings.
@andygrove

Copy link
Copy Markdown
Member Author

cc @phillipleblanc

Each remote stream checked out its own connection, so a wide pipelined job
opened one per producer and consumer task pair and overflowed the executors'
listen queues. A reader partition now multiplexes its streams to an
executor onto one connection, capped at 64 streams per connection.
tonic's `serve` binds through the standard library, which listens with a
backlog of 128. A pipelined job connects from every reader task on every
executor when it starts, so an executor can get thousands of connections
within a second and overflow both its accept and SYN queues. Overflowed
connections wait for a SYN retransmit, and once the kernel falls back to
SYN cookies, a client whose first segments were dropped is reset as soon
as it sends more. The reader then sees `ConnectionReset` on a stream it
never opened, with nothing logged on either side.

Bind the Flight server through `create_grpc_server_incoming`, which the
executor's gRPC server already uses, and have that helper listen with
`GrpcServerConfig::listen_backlog`, 4096 by default. The kernel caps the
value at `net.core.somaxconn`, which is 4096 on Linux 5.4 and later.
… server answered

A wide pipelined job opens its remote streams in one burst, and a
connection the producer's executor could not accept in time can be reset
by the kernel as the reader starts sending. The reader saw that as a
`ConnectionReset` on a stream it had not opened, and since pipelined
jobs do not retry tasks, one such reset failed the job.

`fetch_pipelined_partition` now reports a failure of the connection
itself, before the server answered, as `GrpcConnectionError`, and keeps
`GrpcActionError` for anything the server answered. The reader retries
the former on a new connection, `io_retries_times` times with
`io_retry_wait_time_ms` between attempts, the settings the blocking
shuffle reader already uses. A failed connection is dropped from its
slot, so the other streams of the partition that share it connect
again rather than fail on it. An error the server answered, including
the producer's own failure, is still final: the producer may already
have handed the partition to the failed stream.

Tests cover a stream open that is reset on its first connection and
served on its second, and a server-side failure that is not retried.
A pipelined reader opened a connection per reader partition to each
producer executor, so a job opened executors x tasks x stage boundaries
connections to every executor at once, enough to overflow its listen
queue however large the queue is.

The executor's client pool now keeps one shared connection per
(host, port, config), which the streams of every task on the executor
multiplex over, so an executor holds one connection to each other
executor. A connection that fails before the server answers is
discarded from the pool, and the open is retried on a new one.

Sharing is safe because of how the windows are set. A stream nobody is
reading holds up to its stream window of data on the connection, and if
stalled streams could use up the connection window they would starve
the streams something is waiting on. The shared connection therefore
advertises the largest connection window HTTP/2 allows and a small
stream window, `ballista.pipelined.stream_window_bytes` (512 KiB by
default), so stalled streams hold at most one window each and a
connection carries over 4000 streams before any could starve the rest.
The executor's server is explicitly left without a limit on concurrent
streams, since a stream the client had to queue behind a limit would
leave its producer blocked on a full queue.

`BallistaClientPool` gains a `shared` method with a default
implementation, so existing pools keep compiling.
@andygrove
andygrove marked this pull request as draft October 27, 2026 19:39
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

documentation Improvements or additions to documentation

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant