dongjoon-hyun commented on code in PR #58097: URL: https://github.com/apache/spark/pull/58097#discussion_r3815166740
########## 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: This clears the abandoned mark for every partition of the shuffle, but the marks are process-wide and shared by **all concurrently running map tasks of the same stage** -- so a late-starting map task erases a departure recorded while an earlier sibling was already writing. LIMIT over a pipelined shuffle (the `PipelinedLimitHangSuite` shape): 1. `m0..m3` and the reduce task for partition 0 are gang-admitted together. 2. `m0` enters `write()` and pushes a full batch. 3. The reduce task satisfies `limit(10)`, completes, and its completion listener calls `abandon(shuffleId, 0)`. 4. `m3` is still deserializing; when it finally enters `write()`, this loop removes the mark for partition 0. 5. `wants(0)` is now permanently true for `m3`: it fills partition 0's 64-batch queue and then spins forever in `putUnlessAbandoned`'s `while (!isAbandoned) q.offer(..., 100ms)` with no drainer. The map stage never finishes, and since the consumer's completion is deferred until its pipelined producers finish, the job hangs -- the exact failure the abandon mechanism exists to prevent, now timing-dependent rather than fixed. The window is between task launch and entering `write()` (task deserialization, broadcast fetch), which is far longer than the time it takes a reader to satisfy a small LIMIT from one already-enqueued batch. The reset likely needs to be scoped to a job / stage-attempt generation (e.g. tag marks with the producing stage attempt id and ignore stale ones) rather than "every writer clears every partition". ########## 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: `wants(pid)` runs once per input record and both halves allocate: `abandoned.contains((shuffleId, pid))` builds a `Tuple2` and boxes `pid`, and `liveReducePartitions.forall(_.contains(pid))` boxes `pid` again for the `Set[Int]` lookup. On the 20M-row repartition this transport is benchmarked against that is roughly 40M short-lived allocations plus two hash lookups per row, on the hot loop of a design whose entire justification is amortizing per-row cost via batching. Neither fact needs to be re-read per record. An `Array[Boolean]` of per-partition liveness, seeded from `liveReducePartitions` and refreshed for a partition only when that partition's batch is flushed, reduces the per-record path to a single array load. ########## core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala: ########## @@ -2604,6 +2687,36 @@ 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) => Review Comment: This assumes the `ResultStage`'s partition ids **are** the shuffle's reduce partition ids, but that only holds when the result RDD is the `ShuffledRowRDD` itself. Any narrow operator in between that changes the partition count breaks it. `spark.range(1000).groupBy($"id" % 10).count().coalesce(2).collect()` (feature on, local mode): 1. `getShuffleDependenciesAndResourceProfiles(rs.rdd)` walks the `CoalescedRDD`'s narrow dep down to the pipelined dep, so `resultReadsThisShuffle` is true. 2. The property is set to `"0,1"` -- the `CoalescedRDD`'s partition ids, not reduce partition ids. 3. `ChannelShuffleWriter.wants()` then drops every record with `pid >= 2` and emits no `EndOfStream` for those partitions. 4. But `CoalescedRDD.compute` still pulls the parent iterator for reduce partitions `0..199`, so `ChannelShuffleReader` for `2..199` blocks forever in `q.take()` and the job hangs. The same class of failure applies to a partial-read `take` over a `UnionExec` whose shuffle branch is not the first child, where result ids are offset from reduce ids. Neither `EnablePipelinedShuffle` nor `AQEEnablePipelinedShuffle` rejects these shapes, so the check here probably needs to establish that the result stage's RDD *is* this shuffle's `ShuffledRowRDD` (identity), not merely that the dependency is reachable from it. ########## sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/PipelinedShuffleSqlSuite.scala: ########## @@ -0,0 +1,310 @@ +/* + * 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.sql.execution.exchange + +import org.apache.spark.SparkFunSuite +import org.apache.spark.sql.SparkSession +import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + +/** + * End-to-end SQL coverage of the pipelined channel path: a batch query whose hash + * exchange is rewritten to a pipelined shuffle (EnablePipelinedShuffle) and served by the + * in-process channel manager (PipelinedChannelShuffleManager), run through the + * concurrent-stage scheduler on a single executor. Self-manages its SparkSession because the + * shuffle manager and AQE-off gate are start-up configs. + */ +class PipelinedShuffleSqlSuite extends SparkFunSuite with AdaptiveSparkPlanHelper { + + private def withPipelinedSession(body: SparkSession => Unit): Unit = { + val spark = SparkSession.builder() + // High task-concurrency cap so the pipelined group's whole-group slot demand (the sum + // of every concurrent stage's partitions) is admitted. This is a correctness harness, + // not a perf one: on a smaller physical machine these logical slots oversubscribe the + // cores, which is fine for verifying results but meaningless for timing. + .master("local[16]") + .appName("pipelined-shuffle-sql") + .config("spark.shuffle.manager.incremental", + "org.apache.spark.shuffle.local.pipelined.PipelinedChannelShuffleManager") + .config("spark.sql.adaptive.enabled", "false") // rule only sees exchanges with AQE off + .config("spark.sql.pipelinedShuffle.enabled", "true") + .config("spark.speculation", "false") + .config("spark.sql.shuffle.partitions", "4") + .getOrCreate() Review Comment: Building the session with `SparkSession.builder()...getOrCreate()` inside a plain `SparkFunSuite`, without first clearing any pre-existing session, is fragile. Same pattern in `AQEPipelinedShuffleSuite` and `PipelinedLimitHangSuite`. sql/core suites share a JVM. If an earlier suite in the same fork leaves an active or default `SparkSession` behind, `getOrCreate` returns **that** session and silently ignores every `.config(...)` here -- including `spark.shuffle.manager.incremental` and `spark.sql.pipelinedShuffle.enabled`. The exchanges are then never flipped and `assert(pipelinedExchanges.nonEmpty)` fails with a message that points nowhere near the real cause. Worse, the `finally` block's `spark.stop()` would tear down a session this suite does not own. Extending `LocalSparkSession` (it does `clearActiveSession()` / `clearDefaultSession()` in `beforeEach` and stops in `afterEach`), or at minimum stopping any live `SparkContext` and clearing both session slots before the builder, removes the ordering dependency. -- 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]
