[WIP] Stack/pipelined shuffle pr8 rtm - #57692
Open
jerrypeng wants to merge 16 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
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)