dongjoon-hyun commented on code in PR #58097: URL: https://github.com/apache/spark/pull/58097#discussion_r3823446866
########## core/src/main/scala/org/apache/spark/shuffle/local/pipelined/ChannelShuffleWriterReader.scala: ########## @@ -0,0 +1,304 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.shuffle.local.pipelined + +import java.util.Arrays + +import org.apache.spark.{SparkContext, SparkEnv, TaskContext} +import org.apache.spark.scheduler.MapStatus +import org.apache.spark.shuffle.{BaseShuffleHandle, ShuffleReader, ShuffleReadMetricsReporter, ShuffleWriteMetricsReporter, ShuffleWriter} + +/** + * Map-side of the in-process pipelined shuffle. Each input record is routed to the reduce + * partition its key hashes to and accumulated in a per-partition batch; a FULL batch (an + * `Array[AnyRef]` of `batchSize` pairs) is pushed onto that partition's shared queue in one + * queue operation (see [[ChannelShuffleRendezvous]]) -- the consumer stage, running + * concurrently, drains it batch by batch. No serialization, no disk, no network. + * + * Batching is what makes the transport viable for large unaggregated shuffles: the queue + * costs a lock acquisition per operation (~hundreds of ns under producer/consumer + * contention), so handing rows across one at a time costs that PER ROW -- measured at ~19x + * slower than a regular shuffle on a 20M-row repartition. Batching divides the lock traffic + * by `batchSize`, the same lesson as any object-batch transport. A batch + * array is handed off to the consumer and never touched again by the writer (a fresh array + * is allocated after each put), so ownership transfer is clean across threads. + */ +private[spark] class ChannelShuffleWriter[K, V]( + handle: BaseShuffleHandle[K, V, _], + mapId: Long, + batchSize: Int, + writeMetrics: ShuffleWriteMetricsReporter) + extends ShuffleWriter[K, V] { + + require(batchSize > 0, s"batchSize must be positive, got $batchSize") + + private val dep = handle.dependency + private val partitioner = dep.partitioner + private val numPartitions = partitioner.numPartitions + private val shuffleId = handle.shuffleId + + // The reduce partitions this job actually reads, from the producer stage's task property + // (set by the DAGScheduler from the result stage's partitions). A record routed to a + // partition NOT in this set has no consumer -- putting it would fill that partition's + // bounded queue and, because the writer interleaves all partitions on one thread, block + // the writer before it can feed even the read partitions or emit their end-of-stream, + // deadlocking the job. So such records are dropped. Absent property (None) means every + // partition is live (the normal full-read case: collect, count, a full-partition job) and + // nothing is dropped. + private val liveReducePartitions: Option[Set[Int]] = + Option(TaskContext.get()) + .flatMap(tc => + Option(tc.getLocalProperty(SparkContext.SPARK_PIPELINED_LIVE_REDUCE_PARTITIONS))) + .map(_.split(",").filter(_.nonEmpty).map(_.toInt).toSet) + + // Per-partition liveness, precomputed ONCE from the (static) live set: true iff a consumer + // reads this reduce partition at all. This is the hot-path gate -- checked per input record -- + // so it is a plain Array[Boolean] load, not a boxed Set lookup: on a large repartition the + // per-record path must not allocate (the transport's whole point is amortizing per-row cost). + // Absent property means every partition is live. The OTHER half of "worth writing" -- + // abandonment, which happens at runtime when a reader departs early (e.g. LIMIT) -- is dynamic + // and is checked where it matters (at hand-off, in putUnlessAbandoned), NOT per record: + // accumulating a few more rows into an in-memory batch for a since-abandoned partition is + // harmless because that batch is never put (putUnlessAbandoned drops it). + private val liveMask: Array[Boolean] = { + val mask = Array.fill(numPartitions)(true) + liveReducePartitions.foreach { live => + var p = 0 + while (p < numPartitions) { mask(p) = live.contains(p); p += 1 } + } + mask + } + + // Hand a batch to a partition's queue, but do NOT block forever if its reader departs: + // poll with a short timeout and bail out the moment the partition becomes abandoned. This + // is the cooperative unblock for the early-stop case -- abandon() also drains the queue to + // release a parked put, and this re-check ensures the writer then stops rather than + // re-filling. Returns false if the partition was abandoned before the batch was accepted. + // On a successful hand-off, records the batch's records and the time spent (including any + // backpressure wait) against the write metrics; a dropped/abandoned batch counts nothing, + // since those records are never shuffled out. `records` is the number of pairs in `batch` + // (a full batch is `batchSize`, a trimmed tail is shorter; the end-of-stream marker is 0). + private def putUnlessAbandoned(pid: Int, batch: AnyRef, records: Int): Boolean = { + val q = ChannelShuffleRendezvous.queue(shuffleId, pid) + val start = System.nanoTime() + while (!ChannelShuffleRendezvous.isAbandoned(shuffleId, pid)) { Review Comment: `putUnlessAbandoned` has no kill/interrupt check, so this is the writer-side twin of the reader hang fixed in the last round. The only exit from this loop is `isAbandoned`, and that mark is set exclusively by `ChannelShuffleReader`'s task-completion listener -- i.e. only if the reduce task for `pid` actually **started**. If it never starts, the loop never terminates. That is reachable on the opt-in path, because the producer is submitted before the consumer. In `submitStage` the recursion runs `submitStage(parent)` first and the consumer's `submitMissingTasks` afterwards, so the producer's map tasks can already be filling queues before the consumer TaskSet exists. Anything that aborts the group in that window -- `submitMissingTasks` throwing on the consumer (unserializable task), an early failure of another group member, a job cancel -- leaves the producer running with no reader for some or all partitions. The abort then goes through `cancelTasks` / `killAllTaskAttempts` with `shouldInterruptTaskThread(job)`, which reads `spark.job.interruptOnCancel` and defaults to `false`, so the thread is never interrupted and this `offer` stays parked. The consequence is the one you fixed on the reader: the task never dies and pins an executor slot for the life of the application. A partial abort is enough -- with 4 reduce partitions and only one reduce task launched before the abort, the writers for the other three park here. The symmetric fix is the one already applied to `takeItem`: check the kill flag each cycle, e.g. `Option(TaskContext.get()).foreach(_.killTaskIfInterrupted())` at the top of the loop body. ########## core/src/main/scala/org/apache/spark/shuffle/local/pipelined/ChannelShuffleRendezvous.scala: ########## @@ -0,0 +1,155 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.shuffle.local.pipelined + +import java.util.concurrent.{ConcurrentHashMap, LinkedBlockingQueue} + +/** + * Process-wide rendezvous between the map (writer) and reduce (reader) sides of an + * in-process pipelined shuffle. One bounded queue exists per + * `(shuffleId, reducePartitionId)`; every map task writing to a given reduce partition + * shares the queue with the single reduce task that drains it. Queue elements are BATCHES + * of records (`Array[AnyRef]` of pairs, see [[ChannelShuffleWriter]]) or the + * [[EndOfStream]] marker, so the queue's per-operation lock cost is paid per batch, not + * per row. + * + * This is correct only when producer and consumer tasks are co-resident in the same JVM, + * i.e. a single executor (local mode). The concurrent-stage scheduler co-schedules the two + * stages so both are running, but it does not by itself guarantee co-location; the + * [[PipelinedChannelShuffleManager]] is only intended for single-executor deployments, + * where co-location is automatic. Cross-executor pipelined shuffle is served by the RPC + * streaming shuffle instead. + * + * The queues are bounded, so a fast producer blocks on `put` when the consumer lags -- + * this is the backpressure that keeps the pipelined hand-off memory-bounded. + */ +private[spark] object ChannelShuffleRendezvous { + + /** + * Marker placed on a queue by each map task when it finishes writing to that reduce + * partition. A reader stops once it has seen one marker per map task. + */ + val EndOfStream: AnyRef = new AnyRef + + // Keyed by (shuffleId, reducePartitionId). Values are AnyRef because the queue carries + // both record batches (Array[AnyRef]) and the EndOfStream marker. + private val queues = + new ConcurrentHashMap[(Int, Int), LinkedBlockingQueue[AnyRef]]() + + // (shuffleId, reducePartitionId) keys whose reader has departed (its reduce task finished) + // and will drain no more. A writer stops feeding an abandoned partition and drops the rest. + // This covers the LIVE-partition early-stop case (e.g. a LIMIT reader that pulled enough + // and quit): without it the writer fills the partition's bounded queue and blocks forever. + private val abandoned = + java.util.concurrent.ConcurrentHashMap.newKeySet[(Int, Int)]() + + /** + * Per-queue capacity in BATCHES (not rows), the backpressure bound and the heap-residency + * knob (see spark.shuffle.pipelined.channel.queueCapacity). Set once by the channel manager at + * construction from that conf; defaults to 64 (with the default 1024-row batch, ~64K rows per + * reduce partition in flight) until a manager sets it. `@volatile` because the manager sets it + * on the driver while writer/reader threads read it. + */ + @volatile private var capacity = 64 + + /** Set the per-queue capacity in batches. Called by the channel manager from its conf. */ + private[pipelined] def setCapacity(batches: Int): Unit = { capacity = batches } + + /** The queue for one `(shuffleId, reducePartitionId)`, created on first access. */ + def queue(shuffleId: Int, reducePartitionId: Int): LinkedBlockingQueue[AnyRef] = + queues.computeIfAbsent( + (shuffleId, reducePartitionId), + _ => new LinkedBlockingQueue[AnyRef](capacity)) + + /** Whether this reduce partition's reader has departed (see [[abandon]]). */ + def isAbandoned(shuffleId: Int, reducePartitionId: Int): Boolean = + abandoned.contains((shuffleId, reducePartitionId)) + + /** + * Clear ALL abandoned marks for a shuffle. A pipelined producer is RE-RUN for the same + * shuffleId within one query (a RangePartitioner sampling job then the main job; executeTake's + * per-batch jobs), and abandonment is a PER-JOB fact -- a mark left by a previous run's reader + * must not make the re-run's fresh writers think a partition is dead. + * + * This is called ONCE by the DAGScheduler when it submits the producer stage, i.e. BEFORE any + * map task of the new run has started, so it can never race a concurrently running writer or + * reader of the SAME run. (An earlier design cleared marks from inside each writer's write(); + * that raced across the run's own map tasks -- a late-starting map task erased a departure a + * sibling's reader had legitimately recorded, re-hanging the writer. Clearing at the stage's + * submission, the one point with no live task of that run, removes the race by construction.) + * Queues are left intact -- only the marks are reset. + */ + def clearAbandonedForShuffle(shuffleId: Int): Unit = { Review Comment: Resetting the marks without the queues leaves a previous run's data readable by the next run of the same `shuffleId`. The scaladoc is explicit that "Queues are left intact", and `abandon` clears a queue only for a partition whose reader actually departed. A partition whose reduce task never started therefore keeps whatever the writers pushed into it. In classic mode a `Dataset` holds one `executedPlan`, so `ShuffleExchangeExec.shuffleDependency` -- and its `shuffleId` -- is reused across actions. Sequence: 1. An action runs and the group aborts partway (a consumer task fails, or the job is cancelled). Partitions whose readers never started keep their batches, and -- if their writers got that far -- their `EndOfStream` markers too. 2. The user re-runs the action. Same `shuffleId`, new producer stage; `onPipelinedProducerStageSubmit` clears the marks only. 3. The new writers append to the same queues. 4. The new reader drains the stale batches first and counts the stale `EndOfStream` markers toward `numMaps`, so it can stop before the new run's data arrives. That is a silent wrong result (duplicated and/or dropped rows) rather than a hang. Calling `removeShuffle(shuffleId)` here instead of `clearAbandonedForShuffle` would cover it, since the hook is documented to run before any map task of the run starts -- but only if a pipelined producer stage can never be resubmitted with partial progress, otherwise the clear would drop in-flight data. A sturdier option is to scope the rendezvous keys by a per-run epoch, `(shuffleId, epoch, reducePartitionId)`. That also closes the related race the mark reset still has: right after step 2 clears the marks, a straggler writer from the aborted run that has not yet noticed its kill sees `isAbandoned == false` again and can push into the new run's queue. -- 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]
