viirya commented on code in PR #58097: URL: https://github.com/apache/spark/pull/58097#discussion_r3825469940
########## 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: Fixed. `putUnlessAbandoned` now calls `killTaskIfInterrupted()` at the top of each loop cycle, so a writer whose reader never started wakes on the group abort interrupt flag instead of parking on `offer` forever -- symmetric to `takeItem`. Added a regression test: a producer map task that throws now fails the job promptly (under a deadline) rather than pinning the slot. ########## 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: Fixed with per-run epoch isolation. The rendezvous is now keyed `(shuffleId, epoch, reducePartitionId)`, where the epoch is the jobId: `DAGScheduler.handleJobSubmitted` stamps it into a pipelined job properties before any stage, and `submitMissingTasks` clones it per stage so the gang producer and consumer read the same value. A re-run of the same shuffleId is a different job, hence a different epoch and physically separate queues -- so it never drains a prior run leftover batches or counts its `EndOfStream` markers toward `numMaps`, and a straggler writer from the aborted run (still on the old epoch) cannot push into the new run queue. This removes the whole "clear the marks between runs" mechanism: `clearAbandonedForShuffle` and the `onPipelinedProducerStageSubmit` hook are both deleted, and `removeShuffle(shuffleId)` now drops every epoch of a shuffle. New regression test asserts the isolation. -- 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]
