viirya commented on code in PR #58097:
URL: https://github.com/apache/spark/pull/58097#discussion_r3832936449
##########
core/src/main/scala/org/apache/spark/ContextCleaner.scala:
##########
@@ -262,7 +262,23 @@ private[spark] class ContextCleaner(
listeners.asScala.foreach(_.shuffleCleaned(shuffleId))
logDebug("Cleaned pipelined shuffle " + shuffleId)
} else {
- logDebug("Asked to cleanup non-existent shuffle (maybe it was already
removed)")
+ // The shuffle is in NEITHER tracker. This is either a genuinely
non-existent shuffle
+ // (already removed) OR an in-process pipelined shuffle whose manager
keeps NO output
+ // tracker
(PipelinedChannelShuffleManager.usesStreamingShuffleOutputTracker = false):
+ // such a shuffle registers with no tracker at all, so the two
branches above miss it,
+ // yet its process-wide rendezvous queues (ChannelShuffleRendezvous)
still need freeing.
+ // Call shuffleDriverComponents.removeShuffle unconditionally: the
RemoveShuffle it issues
+ // routes to SparkEnv.unregisterShuffleFromAllManagers, which reaches
that manager's
+ // unregisterShuffle (-> ChannelShuffleRendezvous.removeShuffle) and
frees the queues.
+ // Safe for a truly non-existent id:
BlockManagerMasterEndpoint.removeShuffle finds no
+ // blocks and unregisterShuffle for an unknown id is a no-op on every
manager. No
+ // shuffle-manager type check here on purpose -- the cleaner stays
transport-agnostic and
+ // just balances the registerShuffle the ShuffleDependency constructor
issues for every
+ // shuffle.
+ logDebug("Cleaning tracker-less shuffle " + shuffleId)
+ shuffleDriverComponents.removeShuffle(shuffleId, blocking)
Review Comment:
Follow-up: the manager-capability gate I described wasn't quite enough. With
the feature on, a session can still run regular shuffles (e.g. a plan the rules
leave regular, including the coalesce fallback above); an already-cleaned one
reaching this arm would still fire a duplicate `RemoveShuffle`. I've scoped the
arm to shuffles the manager actually holds:
`PipelinedShuffleManager.holdsShuffle(id)` (the channel manager tracks its
registered ids), gated as `!usesStreamingShuffleOutputTracker &&
holdsShuffle(id)`. A regular shuffle is never held, so it stays a no-op.
##########
core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala:
##########
@@ -2604,6 +2687,60 @@ private[spark] class DAGScheduler(
.getOrElse(new Properties())
addPySparkConfigsToProperties(stage, properties)
+ // For a pipelined PRODUCER stage, tell its tasks which reduce partitions
the job actually
+ // reads (the result stage's partitions). The in-process channel writer
drops records
+ // routed to partitions no consumer will drain -- otherwise a partial-read
job (LIMIT /
+ // executeTake reads a subset) fills the unread partitions' bounded queues
and deadlocks
+ // the writer. The result stage is created before submitStage, so its
partitions are known
+ // here. Only the map-side producer needs this; a regular full-read job's
result stage
+ // covers every partition, so the property (if written) lists them all and
drops nothing.
+ // The live set is per-SHUFFLE-EDGE, not per-job: it is the reduce
partitions the
+ // consumer of THIS shuffle reads. It equals the job's result partitions
ONLY for the
+ // producer whose shuffle the result stage reads DIRECTLY (result
partition i maps to
+ // that producer's reduce partition i). A middle pipelined exchange in a
chain (e.g. a
+ // subquery's hash below a single-partition agg) is consumed by another
map stage that
+ // reads ALL its partitions, so it must stay fully live -- setting the
result's subset
+ // there would make it drop partitions the downstream stage still needs,
deadlocking.
+ // So set the property only on the result-feeding producer.
+ stage match {
+ case sms: ShuffleMapStage if isPipelinedProducer(stage) =>
+ // Notify the pipelined manager that this producer stage is being
(re)submitted, before
+ // any of its map tasks start. A transport with per-run state keyed by
shuffleId resets it
+ // here -- the one point with no live task of the new run, so the
reset cannot race the
+ // run's own writers/readers. The in-process channel transport clears
a prior run's
+ // abandoned-partition marks; the RPC streaming manager keeps no such
state (no-op default).
+ // The scheduler stays transport-agnostic: it calls the
PipelinedShuffleManager trait, not
+ // a concrete transport.
+ SparkEnv.get.pipelinedShuffleManager
+ .onPipelinedProducerStageSubmit(sms.shuffleDep.shuffleId)
+ jobIdToActiveJob.get(jobId).map(_.finalStage).collect { case rs:
ResultStage => rs }
+ .foreach { rs =>
+ // The live set is the result stage's partition ids. Those are the
shuffle's reduce
+ // partition ids -- so partition i means reduce partition i --
ONLY when the result RDD
+ // reaches this shuffle through an IDENTITY-PRESERVING chain:
every hop a 1:1,
+ // same-index OneToOneDependency (e.g. the mapPartitions wrapper
executeTake/collect
+ // adds over the ShuffledRowRDD, which preserves partition count
and index). It is NOT
+ // enough for the shuffle to be merely reachable: a narrow
operator that REMAPS
+ // partitions (coalesce -> a custom NarrowDependency, union -> a
RangeDependency offset)
+ // makes result partition ids differ from reduce partition ids,
and treating them as
+ // the live set would drop partitions a downstream operator still
pulls -- hanging the
+ // reader. When the chain is not identity-preserving the property
is left unset (all
+ // partitions live) and that operator drains every reduce
partition.
+ def readsShuffleByIdentity(rdd: RDD[_]): Boolean =
rdd.dependencies match {
Review Comment:
Follow-up: this is now subsumed by the `getParents`-based `liveReduceSet`
(see the reply above) -- the index mapping is derived from the narrow chain
rather than assumed, so an offset spec maps correctly instead of degrading to
all-live.
--
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]
