andygrove opened a new pull request, #1951:
URL: https://github.com/apache/datafusion-ballista/pull/1951
# Which issue does this PR close?
Closes #1943.
# Rationale for this change
The shuffle reader fetches each remote partition as an open-ended `do_get`
stream with **no global in-flight-bytes / in-flight-requests governor**. That
forces a lose-lose choice between three failure corners:
| Corner | Symptom |
|---|---|
| Multiplex many streams over one h2 connection | h2 64 KB connection-window
deadlock |
| One exclusive connection per stream | connection churn / **ephemeral-port
exhaustion** at high `target_partitions` (#1943) |
| Broken pipe mid-body | fatal — a partition is an unbounded, non-retriable
stream |
Spark solves the same problem with a **reduce-side in-flight governor**
(`maxBytesInFlight` / `maxReqsInFlight` / `maxBlocksInFlightPerAddress`, see
`ShuffleBlockFetcherIterator`). That governor is what makes *multiplexing over
very few connections* safe: it bounds how much data is in flight at the
application layer, so the transport is never flooded and a small, reused
connection set is sufficient. This PR ports that model onto Ballista's
gRPC/Arrow-Flight transport.
# What changes are included in this PR?
1. **Reduce-side in-flight governor.** Each remote fetch acquires three
tokio semaphore permits — an in-flight-**bytes** budget (acquiring
`min(block_size, max_bytes)`, charged from `PartitionStats.num_bytes`), an
in-flight-**request** count, and a **per-address** slot — and holds them (via a
drop-guard on the returned stream) until the body is fully consumed. This
replaces the previous count-only semaphore whose permit was released as soon as
the stream handle was obtained (before the body was read), so it bounds
concurrent *in-flight data*, not just connection establishment. By bounding
concurrency it sharply reduces how many connections are opened per peer, which
is what resolves the ephemeral-port exhaustion in #1943.
2. **h2 flow-control window sizing.** The shuffle data-plane client's HTTP/2
initial connection and stream windows are made configurable and default `>=`
the governor byte budget, so the application-level governor — not the 64 KB
transport window — is the binding backpressure (the property Spark gets for
free from Netty). Control-plane clients are unchanged.
3. **Retriable, buffered block fetch.** Each partition body is buffered into
memory as a discrete, idempotent unit and the whole fetch is retried on a
transport error (up to `io_retries_times`), so a mid-body broken pipe refetches
the block instead of failing the task. The outer retry owns retry policy
(transport-error-gated) and disables the client's inner establish-retry to
avoid multiplying attempts.
4. **Configuration** (all with Spark-parity defaults):
- `ballista.shuffle.reader.max_bytes_in_flight` (48 MiB)
- `ballista.shuffle.reader.max_blocks_in_flight_per_address` (128)
- `ballista.shuffle.reader.default_block_size_bytes` (1 MiB, used when
partition stats carry no byte count)
- `ballista.client.initial_connection_window_size` (64 MiB) /
`ballista.client.initial_stream_window_size` (16 MiB)
- reuses the existing `ballista.shuffle.max_concurrent_read_requests`
(64) as the request-count cap and `io_retries_times` (3) for fetch retry.
**Testing.** Unit tests cover the permit helpers, the `GovernedStream`
release-on-drop guard, and the transport-gated retry (retriable vs
non-retriable, attempt counts). A standalone integration test runs a shuffle
query under a deliberately tiny in-flight budget (64 KiB, 2 blocks/address) to
prove it completes without deadlock. The full `sort_shuffle` standalone suite
passes.
**Preliminary cluster result (SF100, 2×8 = 16 task slots):** with this
change, `target_partitions=128` — which previously died very early on
shuffle-fetch port exhaustion — runs the full TPC-H set, and
`target_partitions=32` completes ~10% faster than before. (Larger-partition
runs also exposed an unrelated scheduler memory ceiling that is a separate
concern.)
**Deferred (follow-ups, not in this PR):** spilling oversized partition
bodies to disk (`spark.maxRemoteBlockSizeFetchToMem` analog — until then a
single partition larger than the budget is buffered in full, and the byte
budget bounds compressed *wire* bytes, so decoded RAM is larger by the
compression ratio), and a hard `connections_per_peer` cap (which would require
a shared-connection pool model).
# Are there any user-facing changes?
New configuration keys are added (listed above), all with defaults, so
existing deployments behave sensibly with no changes. The shuffle read path now
bounds in-flight bytes/requests and buffers each partition body into memory
before yielding it downstream. No public API breakage.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
