[WIP] Stack/pipelined shuffle pr8 rtm - #57692
Open
jerrypeng wants to merge 24 commits into
Open
Conversation
…utput tracker A pipelined shuffle streams records from a live producer stage to a co-scheduled consumer stage, so the consumer must look up its producer tasks' locations while both stages run. That routing lives in the StreamingShuffleOutputTracker, but createShuffleMapStage only registered the shuffle with the MapOutputTracker (which, for a transient pipelined shuffle, tracks an empty never-materialized output). As a result a consumer task querying for writer locations found the shuffle unregistered and failed. Register a PipelinedShuffleDependency's shuffle with the StreamingShuffleOutputTracker inside the same once-per-shuffleId guard as the MapOutputTracker registration. Inert for a regular shuffle (and when no streaming tracker is configured). Note: this is a latent gap in the co-scheduling layer independent of RTM; it belongs with the pipelined-scheduling change earlier in the stack and should be re-homed there (with a real-transport regression test) when the stack is next reorganized. Co-authored-by: Isaac
… output tracker on cleanup A pipelined shuffle is registered with the driver-only StreamingShuffleOutputTracker in DAGScheduler.createShuffleMapStage, alongside its MapOutputTracker registration, but nothing ever removed it. Over a long-running Real-Time Mode query -- which allocates a fresh shuffle id per micro-batch -- the tracker's shuffleInfos map grows without bound, and a re-registration of a reused id would hit the "registered twice" guard. Unregister it in ContextCleaner.doCleanupShuffle, right beside the existing mapOutputTrackerMaster.unregisterShuffle. The streaming tracker's state is driver-only (the worker tracker caches nothing and just RPCs the master), so a direct driver-side call is the correct counterpart -- unlike the MapOutputTracker, there is no per-executor cache to invalidate, so this does not belong in the RemoveShuffle RPC handler. The call is guarded by a new containsShuffle predicate so regular (non-pipelined) shuffles, which are never registered here, are skipped without emitting a spurious "not registered" warning on every shuffle cleanup. Also correct the StreamingShuffleManager.unregisterShuffle comment, which previously claimed the tracker was unregistered in BlockManagerStorageEndpoint's RemoveShuffle handler -- untrue in OSS (that path routes to this no-op manager method). Tests: a containsShuffle unit test, plus two ContextCleanerSuite integration tests (a pipelined shuffle is unregistered on cleanup; a regular shuffle is left untouched and quiet). Both integration tests were revert-checked -- each fails when its respective production change (the unregister call, or the containsShuffle guard) is removed. Co-authored-by: Isaac
…putTracker A pipelined shuffle produces no durable, addressable map output: its producer/consumer task-location routing lives entirely in the StreamingShuffleOutputTracker, its availability is tracked on the stage (pipelinedCompletedPartitions), and its reader locates writers through the streaming tracker. Yet createShuffleMapStage still registered an empty ShuffleStatus for it in the MapOutputTracker, purely to keep containsShuffle / unregisterShuffle bookkeeping uniform. That coupled the streaming tracker's lifecycle to MapOutputTracker state: registration and cleanup of the streaming entry were both nested inside MapOutputTracker.containsShuffle guards, so a future divergence between the two trackers could double-register (throwing "registered twice") or leak the streaming entry. Split the two trackers cleanly by dependency type, with no overlap: - createShuffleMapStage: a PipelinedShuffleDependency registers ONLY with the StreamingShuffleOutputTracker (self-guarded on that tracker's own containsShuffle); a regular shuffle registers ONLY with the MapOutputTracker, exactly as before. No empty MapOutputTracker shell is created for a pipelined shuffle. - ContextCleaner.doCleanupShuffle: two independent branches, each keyed on its own tracker's membership -- MapOutputTracker cleanup for a regular shuffle, a direct driver-side StreamingShuffleOutputTracker unregister for a pipelined one (no RemoveShuffle RPC: a streaming shuffle has no block-manager-served files). Neither branch depends on the other tracker's state. This is safe because every MapOutputTracker path a pipelined shuffle could reach is either already type-gated out (output registration, availability, FetchFailed invalidation which routes to a group-atomic abort, push-merge, indeterminate/checksum rollback, map-stage-job stats) or tolerant of the entry's absence (executor/host-loss strips iterate shuffleStatuses and simply skip it). Consumer task locality never reads the tracker for a shuffle dependency (getPreferredLocsInternal recurses only through narrow deps). Tests (ContextCleanerSuite): a pipelined shuffle registered ONLY in the streaming tracker is still cleaned up by doCleanupShuffle even though it is absent from the MapOutputTracker (revert-checked: fails without the streaming cleanup branch); a regular shuffle takes the MapOutputTracker branch and never touches the streaming tracker. Co-authored-by: Isaac
… (before any state mutation) with no streaming tracker, and abort the group on a base-path pipelined FetchFailed Three hardening fixes to "Fully separate pipelined shuffles from the MapOutputTracker", all found by review of that change: 1. createShuffleMapStage registered a pipelined dependency only inside a streamingShuffleOutputTracker.foreach, with the MapOutputTracker registration in the else-if branch. If the streaming tracker were absent (a non-streaming custom incremental manager), the shuffle was registered in NEITHER tracker -- a consumer would then fail with an opaque NullPointerException in the shuffle writer. Fail loud instead: a PipelinedShuffleDependency requires a StreamingShuffleOutputTracker. Crucially, resolve the tracker UP FRONT -- before createShuffleMapStage mutates stageIdToStage / shuffleIdToMapStage / updateJobIdStageIdMaps -- so the fail-loud throw leaves no partial scheduler state. (handleJobSubmitted's catch does no rollback, so a throw after those mutations would leak a half-created stage; a re-submission would then reuse the stale cached stage and fail with a misleading cross-job-reuse error instead of re-throwing the real misconfiguration.) 2. A pipelined-group FetchFailed is normally intercepted by the group-atomic abort branch, whose guard is keyed on the failed stage's active job carrying hasPipelinedDependency. If a pipelined-shuffle FetchFailed reaches the base (non-group-atomic) path instead, that path is invalid for a pipelined shuffle in two ways: (a) its MapOutputTracker invalidation would throw ShuffleStatusNotFoundException (a pipelined shuffle is no longer registered there); and (b) its resubmit branch would enqueue a lone-stage resubmit of the transient producer, which cannot be re-read and would deadlock the group. Guard the whole base-path tail on mapStage.isPipelined and abort the group instead -- matching the group-atomic branch's outcome. Tests (revert-checked): - PipelinedShuffleRoutingSuite: a job with a pipelined dependency and no streaming tracker fails loud with the IllegalStateException, and a re-submission re-throws the SAME error -- proving no partial stage leaked (revert-checked against the throw-after-mutation placement, which instead leaks and yields a cross-job-reuse error on re-submit). - DAGSchedulerSuite: a FetchFailed for a pipelined shuffle reaching the base path (group-atomic guard forced to miss) aborts the group -- the job fails, the throwing MapOutputTracker calls are never invoked, and the transient producer is not resubmitted (revert-checked against both the no-guard and the invalidation-only-guard variants). Co-authored-by: Isaac
…ng test Co-authored-by: Isaac
…ale availability comment - Reflow four 101-char comment lines in the FetchFailed group-abort branch and the unregisterOutputsOnFetchFailedExecutor scaladoc to <=100. - Correct ShuffleMapStage.numAvailableOutputs: a pipelined shuffle is not registered with the MapOutputTracker at all (it was leftover "registered but empty" wording that contradicted this file's own scaladoc). Comment-only; no behavior change. Co-authored-by: Isaac
cloud-fan's nits on the merged predecessor PR, applied here per his request to address them in the follow-up: - DAGScheduler: reword the push-based-shuffle-merge rejection comment ("as the" -> "on a" pipelined shuffle) and the slot-check-disabled warning comment ("Keeps" -> "Limits the warning to once"). - TaskSetManagerSuite: the pipelined-decommission test comment said the shuffle is "NEVER registered in the MapOutputTracker". With the MapOutputTracker now fully decoupled that is literally true, but state it precisely: registered only with the StreamingShuffleOutputTracker, never the MapOutputTracker. Comment-only; no behavior change. Co-authored-by: Isaac
…ype via a common interface Introduce a small ShuffleOutputTrackerMaster trait (registerShuffle / containsShuffle / unregisterShuffle) implemented by both MapOutputTrackerMaster and StreamingShuffleOutputTrackerMaster. DAGScheduler.createShuffleMapStage now resolves the owning tracker by dependency type through a private outputTrackerMaster(shuffleDep) helper (which isolates the fail-loud-if-no-streaming-tracker check) and registers through the common interface, replacing the awkward Option[StreamingShuffleOutputTrackerMaster] + isInstanceOf/asInstanceOf and the type-specific registration branches. Behavior is unchanged. Co-authored-by: Isaac
…anup Close two gaps in the pipelined-shuffle leak coverage: - ContextCleanerSuite: a GC-triggered end-to-end cleanup test. A pipelined shuffle's only cleanup channel is the ContextCleaner weak-reference path fired when its PipelinedShuffleDependency is garbage-collected; the existing tests call doCleanupShuffle directly. The new test builds a real dependency, asserts cleanup does NOT fire while it is strongly referenced, then dereferences it and asserts GC drives the streaming-tracker unregister. - DAGSchedulerSuite: a fan-in (two-producer) deferral leak test that drives a producer failure to full job teardown and asserts all scheduler state (including the two-parent deferral) is emptied. Co-authored-by: Isaac
… a pipelined shuffle on cleanup The ShuffleDependency constructor calls shuffleDriverComponents.registerShuffle for EVERY shuffle, including a pipelined one, but the pipelined cleanup branch in ContextCleaner.doCleanupShuffle only unregistered from the StreamingShuffleOutputTracker -- it skipped the matching removeShuffle. Benign for the default (no-op registerShuffle) impl, but a custom ShuffleDriverComponents tracking per-shuffle driver state would leak it across micro-batches. Call removeShuffle in the pipelined branch too, matching the regular branch's ordering. (Found by Isaac review.) Co-authored-by: Isaac
…dependency Add a `pipelined` flag to ShuffleExchangeExec. When set, prepareShuffleDependency constructs a PipelinedShuffleDependency instead of a regular ShuffleDependency, so the DAGScheduler co-schedules the exchange's producer and consumer stages and routes the shuffle to the streaming shuffle manager. The flag defaults to false, so every existing shuffle is unchanged. Modeling this as a node field (rather than out-of-band metadata) keeps the pipelined/non-pipelined distinction part of the exchange's identity, so it participates in exchange-reuse and canonicalization -- correct for a transient shuffle that must not be reused as a durable one. Adding the field shifts ShuffleExchangeExec's positional arity, so update the pattern matches that destructure it positionally (EnsureRequirements, AQEUtils, InsertSortForLimitAndOffset and several test suites) to account for the new field; EnsureRequirements' repartition rewrite now uses copy() so the flag is preserved. Co-authored-by: Isaac
…d shuffle Enable a stateful streaming query (starting with deduplication) to run in Real-Time Mode by pipelining its repartition shuffle: - IncrementalExecution adds a Real-Time Mode preparation rule that marks every streaming shuffle exchange pipelined (ShuffleExchangeExec.pipelined = true). Real-Time Mode is detected structurally by a RealTimeStreamScanExec leaf. The producer (source scan) and consumer (stateful operator) stages are then co-scheduled and stream records through a transient shuffle instead of the consumer waiting for the producer to fully materialize. Inert for a non-RTM batch, so the ordinary microbatch path is unchanged. (AQE is already disabled for RTM batches upstream, so the pipelined shuffle is only ever submitted as a result job, never an AQE map-stage job.) - RealTimeModeAllowlist admits the operators a streaming dropDuplicates plans into (StreamingDeduplicateExec, StateStoreRestoreExec, StateStoreSaveExec) plus ShuffleExchangeExec, so the query is no longer rejected before it runs. Co-authored-by: Isaac
…e over a pipelined shuffle Add an end-to-end test: a streaming dropDuplicates query in Real-Time Mode whose repartition shuffle is pipelined. It asserts correct dedup output across two batches, that every shuffle exchange in the executed plan is pipelined, and that the producer and consumer stages are co-scheduled (>= 2 stages running at once). The Real-Time Mode operator allowlist check is left enabled, so the test also covers admission of the dedup plan's operators. Co-authored-by: Isaac
… over a pipelined shuffle Add a fault-tolerance case to the pipelined-shuffle Real-Time Mode suite: a stateful dropDuplicates query whose repartition shuffle is pipelined, that hits a task failure mid-batch and then restarts from checkpoint. It asserts the query fails on the task error (Real-Time Mode does not retry tasks; a member failure aborts the whole co-scheduled group), and that on restart the deduplication state recovered from the last committed batch -- already-seen keys are not re-emitted. Switches the suite to the manual-clock base so batch boundaries are deterministic across the failure and restart, and uses CheckAnswerWithTimeout (which polls the sink in real time) with advanceRealTimeClock to drive Real-Time Mode batches. Co-authored-by: Isaac
… in Real-Time Mode Add a second fault-tolerance case: a dedup query over a pipelined shuffle whose commit log write is fault-injected (via FailureInjectionCheckpointFileManager) so a batch cannot commit and the query fails. On restart from the same checkpoint, the committed batch's deduplication state survives and the uncommitted batch is re-run, so its data is still emitted. The injected close() failure surfaces as the underlying IOException in this path (the COMMIT_LOG_WRITE_FAILURE error-class wrapping is not open source), so the test expects that. Co-authored-by: Isaac
MarkPipelinedShuffleForRealTimeMode marks every shuffle exchange in an RTM plan as pipelined, so a chain of several pipelined shuffles (e.g. two repartitions, or a repartition feeding a keyed stateful operator) is co-scheduled as one pipelined group. The one unsupported shape -- a pipelined shuffle with more than one consumer (fan-out, e.g. a reused exchange) -- is rejected up front by the DAGScheduler, so no ReusedExchangeExec reasoning is needed here. Co-authored-by: Isaac
`repartition(n)` uses RoundRobinPartitioning, and by default (SPARK-23207) Spark inserts a local sort before a round-robin shuffle so the row-to-partition assignment is deterministic across task retries. That sort is fully blocking: it drains its entire input before emitting a row. A Real-Time Mode task reads an unbounded stream, so the input never ends. The producer stage therefore never emits a row, never reaches the shuffle write, and every downstream stage waits forever -- the query returns no rows at all. This is not visible to the Real-Time Mode operator allowlist because the sort is not a SortExec node; it lives inside ShuffleExchangeExec's RDD construction. Determinism is not needed here: a pipelined task set is capped at a single attempt (TaskSetManager), and a pipelined group is aborted as a unit rather than recomputed, so the retry hazard the sort guards against cannot arise. Disable the config for Real-Time Mode unless the user has set it explicitly. Co-authored-by: Isaac
…ite and add coverage Move the pipelined-shuffle Real-Time Mode tests out of the standalone RealTimeModePipelinedShuffleSuite into StreamRealTimeModeSuite.scala, which already exists in both this repo and the Databricks runtime. The clock-driven tests land in StreamRealTimeModeWithManualClockSuite, whose base supplies advanceRealTimeClock; the reused-broadcast test needs no manual clock and goes in the parent suite. This keeps the tests in a file that both sides share. New coverage: - A chain of two pipelined shuffles (a round-robin repartition feeding a keyed dropDuplicates) forms a single co-scheduled pipelined group. This is the multi-shuffle-per-group shape, previously untested, and is the regression test for the blocking sort fixed in the previous commit. - Multi-key dropDuplicates over a pipelined shuffle. - An explicit repartition-by-key in Real-Time Mode. - Multiple broadcast joins against the same static table, which produces a ReusedExchangeExec (SC-209926). Co-authored-by: Isaac
…ensitive Turning off the deterministic local sort for Real-Time Mode (previous commit) has a second effect: it flips the same condition's other branch on, marking the map RDD of a round-robin shuffle order-sensitive. That marking exists only so a retry is handled correctly -- an order-sensitive RDD whose own input is UNORDERED becomes INDETERMINATE, which tells the scheduler the stage must be rolled back and recomputed. Any shuffle output is UNORDERED, so in a chain of two round-robin repartitions the second shuffle's producer reads an UNORDERED input and becomes INDETERMINATE. The DAGScheduler rejects an indeterminate pipelined producer outright, so such a chain failed with PIPELINED_SHUFFLE_UNSUPPORTED (before the previous commit it hung silently instead). A pipelined stage is never retried -- a pipelined task set gets a single attempt -- and is never rolled back and recomputed, so the marking protects nothing there. Exempt pipelined shuffles from it, which lets a chain of round-robin repartitions run. Non-pipelined shuffles are unaffected. Co-authored-by: Isaac
… Mode The previous guard only set the config when the user had not set it explicitly. Since it defaults to true, a user (or a spark-defaults deployment) that sets it explicitly to true kept the fully-blocking local sort before a round-robin repartition -- which never drains an unbounded stream, so the query hangs forever with no error. That is the exact hazard this change set removes. Override it unconditionally for Real-Time Mode. A query that never produces a row is worse than ignoring a config that only guards against a task retry that cannot happen here. Log a warning when overriding an explicit setting so the operator can see their config is not the one in force. Co-authored-by: Isaac
The marking rule flipped every ShuffleExchangeExec in any plan that contained a RealTimeStreamScanExec anywhere, without checking the exchange was on the streaming path. A plan can hold a static subtree alongside the streaming one: the static side of a broadcast stream-static join is planned in the same physical plan and may carry its own shuffle. That shuffle must materialize normally -- a static side runs to completion rather than streaming -- so marking it pipelined pulled it into the pipelined group and demanded concurrent slots for stages that must instead finish, failing admission with CONCURRENT_SCHEDULER_INSUFFICIENT_SLOT. Mark only shuffles whose subtree reaches the real-time scan. This mirrors the streaming-path detection the operator allowlist already uses, so the rule no longer marks a wider set of exchanges than the allowlist validates. Co-authored-by: Isaac
… what they claim Two tests could pass without exercising the behaviour they describe: - The co-scheduling tests counted every stage in the shared SparkContext, so an unrelated concurrent stage was enough to satisfy `maxConcurrentStages >= 2` even if this query's producer and consumer had run one after the other. Count only stages belonging to the query under test, identified by the query-id property StreamExecution tags its jobs with. - The reused-broadcast test asserted only the join answer, so it would pass whether or not exchange reuse occurred. Assert the plan actually contains a ReusedExchangeExec and no shuffle exchange, which is the shape the test is about. Co-authored-by: Isaac
Building a RangePartitioner runs a separate job that samples the input to compute range bounds, which cannot complete on an unbounded stream. Marking such a shuffle pipelined would pull that sampling job into a pipelined group, so leave it as a normal materializing shuffle instead. This does not make repartitionByRange work in Real-Time Mode -- the sampling job cannot terminate on an unbounded source either way -- it only stops the rule from labelling a shuffle the pipelined path cannot serve. Co-authored-by: Isaac
Allowlisting ShuffleExchangeExec by class name admits every partitioning, but only some can run in Real-Time Mode. Building a RangePartitioner runs a separate job that samples the input to compute range bounds, and that job cannot complete while the source keeps producing. Before this change set the allowlist rejected every shuffle, so a repartitionByRange failed fast; admitting shuffles turned that into a query that stalls indefinitely with no error. Reject a range-partitioned shuffle at admission so it fails fast with the usual allowlist error again. Co-authored-by: Isaac
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What changes were proposed in this pull request?
This PR activates pipelined (streaming) shuffles for Real-Time Mode (RTM) stateful streaming queries. It is the payload PR of the pipelined-shuffle stack: earlier PRs built the transient-shuffle machinery (group-atomic scheduling, MapOutputTracker decoupling) as dormant infrastructure with no production caller; this PR is the first to actually construct and use a
PipelinedShuffleDependency.It is stacked on top of #57646 ("Fully decouple pipelined shuffles from the MapOutputTracker"); please review/merge that first. This PR's own diff is SQL/streaming only.
Concretely:
ShuffleExchangeExec): a newpipelined: Booleanflag onShuffleExchangeExec; when set, it constructs aPipelinedShuffleDependencyinstead of a regularShuffleDependency. The dependency type is the entire opt-in -- routing to the streaming shuffle manager and pipelined-group co-scheduling in the DAGScheduler both follow from it. Plumbed throughEnsureRequirements,AQEUtils, andInsertSortForLimitAndOffset.IncrementalExecution): aMarkPipelinedShuffleForRealTimeModerule marks everyShuffleExchangeExecin a Real-Time Mode plan (detected structurally by aRealTimeStreamScanExecleaf) as pipelined. Inert for a non-RTM batch. A chain of several pipelined shuffles (e.g. two repartitions, or a repartition feeding a keyed stateful operator) is handled -- the whole all-pipelined job is co-scheduled as one pipelined group. The one unsupported shape (a pipelined shuffle read by more than one consumer -- fan-out) is rejected up front by the DAGScheduler; a reused shuffle requires referencing the same streaming source twice, which RTM already rejects (IDENTICAL_SOURCES_IN_UNION_NOT_SUPPORTED).RealTimeModeAllowlist): admitShuffleExchangeExecplus the stateful operators it enables --StateStoreRestoreExec,StateStoreSaveExec,StreamingDeduplicateExec-- as supported members of a pipelined group rather than materialization barriers.The net effect: a stateful streaming dedup query runs in Real-Time Mode with its input repartitioned through a transient, incrementally-read shuffle, instead of the consumer waiting for the producer to fully materialize.
Why are the changes needed?
Real-Time Mode aims for low end-to-end latency by co-scheduling a stateful query's producer and consumer stages so records stream through rather than materializing at a shuffle boundary. Until now RTM rejected shuffles and the stateful operators behind them (they were not in the allowlist), so a keyed stateful query could not run in RTM. This PR turns on the pipelined-shuffle path -- built and made safe by the earlier PRs in the stack -- for exactly those queries.
Does this PR introduce any user-facing change?
Yes, for Real-Time Mode streaming queries: a stateful query that repartitions by key (e.g. streaming deduplication) is now admitted and runs over a pipelined shuffle in RTM, where it was previously rejected by the RTM operator allowlist. Batch and non-RTM streaming execution are unchanged (the marking rule is inert without a
RealTimeStreamScanExec).How was this patch tested?
RealTimeModePipelinedShuffleSuite(new): an end-to-end RTM stateful-dedup query over a pipelined shuffle produces correct results; and two fault-tolerance scenarios -- recovery from a task failure via checkpoint restart, and recovery from a commit-log write failure.StreamRealTimeModeAllowlistSuite: updated for the new allowlist -- the generic operator-allowlist test dropsShuffleExchangeExecfrom the expected violations; the "repartition not allowed" test is removed (a single-shuffle repartition now runs in RTM, per SPARK-54237); the "stateful queries not allowed" test is narrowed toHashAggregateExec(still not allowlisted), verified against actual runtime behavior.EnsureRequirementsSuite,PlannerSuite,DataFrameWindowFunctionsSuite,DatasetSuite,StreamingJoinSuite: updated for theShuffleExchangeExec.pipelinedfield (defaulted false, so non-RTM plans are unchanged).Was this patch authored or co-authored using generative AI tooling?
Co-authored with Claude Code (Opus 4.8)