NoahKusaba opened a new pull request, #2430:
URL: https://github.com/apache/datafusion-ballista/pull/2430
# Which issue does this PR close?
Closes #2319
# Rationale for this change
`bias` and `round-robin` pick an executor by free vcores alone, so a reduce
task routinely fetches its input over Arrow Flight from an executor that had it
on local disk. `ShuffleReaderExec` already carries the `PartitionLocation`s
naming who wrote each input partition and how many bytes, so the scheduler has
what it needs to place the task where its bytes are. It just doesn't use it.
This adds a policy that does. It is off by default, and it reports the share
of input it kept local, because the gain is real only when data placement and
free capacity disagree.
# What changes are included in this PR?
**The policy** (`ballista/scheduler/src/cluster/affinity.rs`, new) is
selected with `--task-distribution shuffle-affinity`. Per stage, binding runs
two passes:
1. **Affinity.** Every `(pending partition, holder)` pair is ranked by bytes
and taken greedily while holders have vcores free, so a scarce slot goes to the
partition with most to gain, and a partition whose best holder is full settles
for its second best.
2. **Fallback.** Whatever is left is bound `bias`-style onto the remaining
vcores, so no partition idles waiting for a full holder.
An executor must hold 20% of a partition to count as a home for it. Below
that the preference is a guess: in an even shuffle each of `E` executors holds
`1/E`, so past five nothing qualifies and the partition is bound like any
other. Spark uses the same 0.2 in
`MapOutputTracker.getLocationsWithLargestOutputs`.
Attribution follows what one task actually reads. A collapse stage (one
task, every partition) goes whole to the executor holding most of the stage.
`UnionExec` children occupy disjoint partition ranges. Co-partitioned fan-ins
sum into the same partition. Broadcast input counts toward the stage total but
steers no individual partition.
**Plumbing.** `bind_one_where` is a `bind_one` variant that pulls a filtered
subset of the pending queue instead of the front slice, supported by
`PendingPartitions::{queued, next_slice_where}`. Dispatch arms were added to
`InMemoryClusterState::bind_schedulable_tasks` and `poll_work`. Locality scans
are memoized per `(job, stage)` and dropped when a stage succeeds or a job
stops running.
**Observability.** Six `shuffle_locality_*` Prometheus counters,
`ShuffleAffinityPolicy::stats()` for embedded schedulers, and a per-round
`debug!` line. The number to compare between policies is
`shuffle_locality_local_bytes_total / shuffle_locality_input_bytes_total`.
`shuffle_locality_imputed_bytes` flags when a producer reported no size and the
byte counts are padded.
**Docs.** A "Choosing a task distribution" section in the config guide, the
new metrics in the metrics guide, a note in the shuffle contributor guide, and
a fix to the `scheduler-policy` default, which was documented as `pull-staged`
but is `push-staged`.
Two limitations are documented rather than papered over:
- The policy needs `--scheduler-policy push-staged`. Under `pull-staged` the
scheduler sees one executor's capacity at a time and has no placement to
choose, so it warns at startup.
- Locality is executor-level, not host-level. The local-read check tests for
the file in the reader's own work dir, so co-located executors cannot read each
other's output however tasks are placed. Spark's preferences are host-level;
matching that needs a change to the read path, not the scheduler.
## Benchmark
Measured with a harness that drives a throwaway cluster once per policy. It
is kept out of this PR to hold the diff to the scheduler; it lives on the
`shuffle-affinity-benchmark` branch if anyone wants to reproduce these numbers.
4 executors on one machine, 16 shuffle partitions, 15 runs per policy.
`local_share` is the fraction of shuffle reads served from a local file.
**`collapse`**, a global aggregate whose single task reads every partition.
This is the case the policy is for:
| vcores/executor | bias | round-robin | shuffle-affinity |
|---|---|---|---|
| 3 | 25.0% | 18.8% | **37.5%** |
| 4 | 25.0% | 25.0% | 25.0% |
At 3 vcores the 8 map tasks land unevenly and one executor holds 6 of 16
partitions. Affinity finds it every run, the others reach it by luck, and 37.5%
is the ceiling. At 4 vcores the map tasks spread 2/2/2/2, 25% is the ceiling,
and all three tie because there is nothing to win.
**`aggregate`**, a plain hash shuffle. This is the control, showing the
policy costs nothing where it cannot help. Two invocations of the same
benchmark:
| policy | median A | median B | local_share |
|---|---|---|---|
| bias | 1.425s | 1.664s | 24.0% to 27.1% |
| round-robin | 1.363s | 1.467s | 26.0% |
| shuffle-affinity | 1.338s | 1.397s | 24.0% to 26.0% |
Every producer writes every partition at roughly equal size, so all three
read about `1/E`. Both invocations are shown on purpose: the medians move more
between repeat runs of one policy than they do between policies. On a single
machine a "remote" read is a loopback gRPC hop, so wall clock is not the signal
here. The local share is.
## The open questions from #2319
**How much locality is actually available?** Usually none, and the
exceptions are identifiable. On a plain hash shuffle every partition ranks the
executors identically, so affinity degenerates into "pack onto the biggest
holder", which is what `bias` already does, since capacity is what placed the
producers to begin with. That is the tie at `1/E` above. The issue's guess that
the best candidate may hold only `1/N` is correct, and it is now encoded rather
than left implicit: past five holders the policy declines to have an opinion
instead of ranking interchangeable candidates. The 4-vcore collapse row is
worth as much as the 3-vcore one, since it is the same query on the same
cluster, differing only in whether map placement happened to be even.
**What is the cost of the slot idling?** None. This policy never waits: if a
task's preferred executor is full, the task runs elsewhere in the same round.
I built the waiting version (hold a task back until its preferred executor
frees up) and dropped it. It adds a deadline per partition and a rule for how
much locality is worth waiting for, and nothing measured here showed that
paying off. If you revisit it: Spark rechecks held tasks every second, while
Ballista only rechecks when something happens, such as a task finishing.
Usually that is enough, but a wait can run long on an idle cluster.
**Interaction with stragglers and `max_bytes_in_flight`?** Local reads
acquire none of the governor's three semaphores, so a partition moved local
drops its permit demand outright. That much is structural, but it is not
visible in these measurements. `permit_wait_time` came in at 348µs for affinity
against 384µs for `bias` and 297µs for `round-robin`, under 0.03% of a 1.4s
query and well inside the noise. The 48 MiB default is simply not contended by
this workload at this scale, so the question is worth revisiting on a workload
that does contend it.
Ballista has no speculative execution, so there is no duplicate-task
mechanism for affinity to interact with. Retry is the real interaction, and the
cache is tagged with `stage_attempt_num` so a retried stage rescans rather than
placing tasks from a stale view. The residual risk is that affinity
concentrates work on the biggest data holder and turns it into a straggler. The
greedy assignment bounds this by never giving an executor more partitions than
it has free vcores, but this PR does not measure it.
# Tests
36 unit and integration tests.
**Unit.** 26 in `affinity.rs` cover the scan (holder ranking, the 20%
threshold and its exact boundary, union offsets, co-partitioned sums, broadcast
attribution, unsized producers), assignment (scarce slot to the strongest
locality, full first choice falling to the second best), the cache (reused per
stage attempt, invalidated on retry, pruned per job and per stage), and
measurement. Plus 2 for the new `PendingPartitions` methods, 2 for the config
accessor, and 2 for the Prometheus counters.
**Integration.** The push path through `bind_schedulable_tasks` (placement,
that `bias` on the same input differs, and that stats reach the configured
policy instance), the pull path through `poll_work`, and an end-to-end run
asserting locality reaches the metrics collector.
# Are there any user-facing changes?
Yes, all opt-in. Defaults are unchanged, so an existing deployment behaves
exactly as before.
**New config value.** `--task-distribution` accepts `shuffle-affinity`
alongside `bias` and `round-robin`. The default is still `bias`.
**New metrics.** Six `shuffle_locality_*` counters appear on `/api/metrics`,
and only when the new policy is selected. Existing metrics are untouched.
**New log line.** Selecting `shuffle-affinity` together with
`--scheduler-policy pull-staged` logs a warning at startup, because the policy
cannot do anything useful in that mode.
**New public API** for embedded schedulers: the `cluster::affinity` module,
`ShuffleAffinityPolicy` with `new`, `stats` and `attach_metrics`,
`LocalityStats`, and `TaskDistributionPolicy::{shuffle_affinity,
shuffle_affinity_policy}`.
**One source-breaking change.** `TaskDistributionPolicy` is public and gains
a `ShuffleAffinity` variant, so code that matches on it without a wildcard arm
will no longer compile. Adding `_ => ...` fixes it. Nothing else breaks: the
new `SchedulerMetricsCollector::record_shuffle_locality` has a default body, so
existing implementations keep compiling untouched.
--
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]
