viirya commented on code in PR #58097: URL: https://github.com/apache/spark/pull/58097#discussion_r3816668638
########## core/src/main/scala/org/apache/spark/shuffle/local/pipelined/ChannelShuffleWriterReader.scala: ########## @@ -0,0 +1,272 @@ +/* + * 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) + + // A partition is worth writing only while a consumer still wants its data: it must be in + // the job's live set (Half 1: no-reader partitions), AND its reader must not have departed + // (Half 2: a live reader that stopped early, e.g. LIMIT). `wants` is re-checked as data + // flows because abandonment happens at runtime when the reader task completes. + private def wants(pid: Int): Boolean = + liveReducePartitions.forall(_.contains(pid)) && + !ChannelShuffleRendezvous.isAbandoned(shuffleId, pid) + + // 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)) { + if (q.offer(batch, 100, java.util.concurrent.TimeUnit.MILLISECONDS)) { + // A successful offer can race abandon(): abandon does `add(mark)` then `q.clear()`, so if + // it ran between the isAbandoned check above and this offer, our batch lands AFTER the + // clear and would be stranded in the queue (no reader will ever drain it). Re-check and + // clear it ourselves so nothing is left behind. The reader has departed, so discarding is + // correct; and it keeps the queue empty for removeShuffle rather than pinning a batch. + if (ChannelShuffleRendezvous.isAbandoned(shuffleId, pid)) { + q.clear() + return false + } + if (records > 0) { + writeMetrics.incRecordsWritten(records.toLong) + writeMetrics.incWriteTime(System.nanoTime() - start) + } + return true + } + } + false + } + + override def write(records: Iterator[Product2[K, V]]): Unit = { + // This is a fresh producer attempt. Clear any abandoned marks left on our partitions by + // an EARLIER job that reused this shuffleId (RangePartitioner sampling job -> main job; + // executeTake batches) so we do not mistake a prior reader's departure for our own + // partitions being dead. Only abandonment happening during THIS attempt then counts. + var pc = 0 + while (pc < numPartitions) { + ChannelShuffleRendezvous.clearAbandoned(shuffleId, pc) Review Comment: Fixed, exactly your diagnosis -- a late map task erasing a departure a sibling reader recorded. The per-`write()` reset is gone; the reset now happens once, in the DAGScheduler when it submits the producer stage (before any map task of that run starts), via a new `PipelinedShuffleManager.onPipelinedProducerStageSubmit` hook that the channel manager implements. That point has no live task of the run, so it cannot race the run own writers or readers, while still resetting a prior run marks (the sample-job/main-job and executeTake re-runs). Verified `PipelinedLimitHangSuite` completes in both AQE modes. ########## core/src/main/scala/org/apache/spark/shuffle/local/pipelined/ChannelShuffleWriterReader.scala: ########## @@ -0,0 +1,272 @@ +/* + * 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) + + // A partition is worth writing only while a consumer still wants its data: it must be in + // the job's live set (Half 1: no-reader partitions), AND its reader must not have departed + // (Half 2: a live reader that stopped early, e.g. LIMIT). `wants` is re-checked as data + // flows because abandonment happens at runtime when the reader task completes. + private def wants(pid: Int): Boolean = Review Comment: Fixed. The static live set is now a precomputed `Array[Boolean]` seeded once from the property, so the per-record path is a single array load with no boxing or `Tuple2`. Abandonment (the dynamic half) is not re-read per record -- it is checked at hand-off in `putUnlessAbandoned`, which is correct because accumulating a few more rows into an in-memory batch for a since-abandoned partition is harmless (the batch is never put). -- 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]
