This is an automated email from the ASF dual-hosted git repository.
cloud-fan pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new c15aa124de83 [SPARK-58185][CORE] Define PipelinedShuffleDependency and
route shuffles to a ShuffleManager by dependency type
c15aa124de83 is described below
commit c15aa124de83584cdb652df637acf8d96506e6ed
Author: Boyang Jerry Peng <[email protected]>
AuthorDate: Wed Jul 22 20:40:42 2026 +0800
[SPARK-58185][CORE] Define PipelinedShuffleDependency and route shuffles to
a ShuffleManager by dependency type
### What changes were proposed in this pull request?
This PR routes each shuffle to a `ShuffleManager` **by its dependency
type**, so a
`PipelinedShuffleDependency` (added in the previous PR) is served by an
incremental shuffle
implementation while every other shuffle continues to use the regular one.
`SparkEnv` now holds two shuffle managers, keyed by **kind** rather than a
default and an override:
- `spark.shuffle.manager` — the **blocking** manager (a
`BlockingShuffleManager`), serving all
regular (materialized) shuffles and owning block-by-id resolution.
Defaults to `sort`, as before.
- `spark.shuffle.manager.incremental` — the **pipelined** manager (a
`PipelinedShuffleManager`),
serving pipelined (incrementally-readable) shuffle dependencies. Defaults
to the built-in
`streaming` manager, just as the blocking manager defaults to `sort`.
Routing is a single pure function,
`SparkEnv.shuffleManagerFor(dependency)`: a
`PipelinedShuffleDependency` goes to the pipelined manager, every other
`ShuffleDependency` to the
blocking manager. Both managers always exist, so there is no fallback path.
Because the decision is
a pure function of the (serialized, deterministic) dependency, the driver
at `registerShuffle` and
every executor at `getWriter` / `getReader` agree without any shared
routing state or config
re-read. All shuffle-I/O sites route through it
(`Dependency.shuffleHandle`, `ShuffleWriteProcessor`,
`ShuffledRDD`, `CoGroupedRDD`, `SubtractedRDD`, `ShuffledRowRDD`,
`ShuffleExchangeExec`).
To make the manager APIs harder to misuse, `ShuffleManager` is split by
kind:
```
ShuffleManager (declare a kind by extending one of the two
subtypes)
BlockingShuffleManager <- materialized output, served as
block-manager blocks via a
ShuffleBlockResolver (SortShuffleManager)
PipelinedShuffleManager <- incrementally-readable, served
out-of-band, no resolver
(StreamingShuffleManager)
```
`shuffleBlockResolver` moves off `ShuffleManager` onto
`BlockingShuffleManager`, so a manager's type
now says whether it serves block-manager-addressed blocks at all -- instead
of the old contract
where every manager had a `shuffleBlockResolver` that some implementations
threw from. Block-by-id
resolution (reads, push-based merge, decommission migration,
external-shuffle-service cleanup) is
served through a new `SparkEnv.shuffleBlockResolver:
Option[ShuffleBlockResolver]`, sourced from the
blocking manager; it is `None` only before the shuffle manager is
initialized. This is correct
because a pipelined shuffle is served out-of-band and produces no
block-manager blocks, so those
paths only ever resolve regular shuffles.
`initializeShuffleManager` validates each slot's kind and fails fast
otherwise: the
`spark.shuffle.manager` slot must be a `BlockingShuffleManager` and the
`spark.shuffle.manager.incremental` slot a `PipelinedShuffleManager`. A
manager that declares
neither kind is rejected from both slots. This is the "declare your kind"
enforcement, and it
applies to reflectively-loaded managers too, so `ShuffleManager` itself is
**not** sealed -- the
SPARK-45762 third-party extension point (a user-jar manager implementing
the interface) stays open.
Supporting changes:
- `SparkEnv.blockingShuffleManager` and `SparkEnv.pipelinedShuffleManager`
are the explicit
by-kind accessors (e.g. sort-shuffle detection in `ShuffleExchangeExec`
reads
`blockingShuffleManager`); the bare `SparkEnv.shuffleManager` accessor is
deprecated in favor of
`shuffleManagerFor` / `blockingShuffleManager` so the routing decision is
explicit at each call
site.
- `spark.shuffle.manager.incremental` accepts the same style of short
aliases as
`spark.shuffle.manager` -- `streaming` (the default) resolves to
`StreamingShuffleManager`, and
`sort` / `tungsten-sort` resolve to `SortShuffleManager` (though a
blocking manager is then
rejected from this slot by the kind validation).
- The `StreamingShuffleOutputTracker` is initialized when the pipelined
manager is a
`StreamingShuffleManager` (the default) or the blocking manager is a
`MultiShuffleManager`.
- `MultiShuffleManager` (the legacy cluster-level single-slot way to mix
streaming and regular
shuffle) is deprecated in favor of this per-dependency routing. The
routing code does not use it;
`SparkEnv` still recognizes it for tracker initialization (with a TODO to
drop that once
`MultiShuffleManager` is removed).
- `spark.shuffle.manager.incremental` declares
`ConfigBindingPolicy.NOT_APPLICABLE` (a shuffle
manager is a physical-execution choice and does not change how a view/UDF
body resolves).
This is a follow-up in the stack that begins with
`PipelinedShuffleDependency`; a later
`DAGScheduler` change adds the concurrent (pipelined-group) stage
scheduling that actually
constructs these dependencies.
### Why are the changes needed?
`PipelinedShuffleDependency` marks a shuffle whose consumer reads output
incrementally, but a marker
type alone does nothing until the shuffle layer can serve it with an
incremental implementation
while regular shuffles keep using the existing one. A cluster needs both to
coexist and to be chosen
per shuffle, not per cluster. This PR provides that: a dependency-typed
routing point plus two
coexisting managers keyed by kind, so batch (sort) and pipelined (e.g.
streaming real-time mode)
shuffles can run in the same application, each served by the right
implementation. Splitting
`ShuffleManager` by kind (`BlockingShuffleManager` /
`PipelinedShuffleManager`) encodes in the type
system the invariant the routing relies on -- that only a blocking manager
serves
block-manager-addressed blocks -- so callers can no longer accidentally
reach a resolver that does
not exist.
### Does this PR introduce _any_ user-facing change?
No user-visible behavior change for existing workloads: regular shuffles
are still served by
`spark.shuffle.manager` (default `sort`) exactly as before, query results
are unchanged, and the
block-by-id paths (reads, push-based merge, decommission, ESS cleanup) are
unchanged. Nothing
constructs a `PipelinedShuffleDependency` yet, so no shuffle is routed to
the pipelined manager in
this PR.
One internal behavior change worth calling out:
`spark.shuffle.manager.incremental` now defaults to
the built-in streaming manager (rather than being unset), so a
`StreamingShuffleOutputTracker` is
initialized in every `SparkEnv`. The new config, the
`BlockingShuffleManager` /
`PipelinedShuffleManager` traits, and the `SparkEnv` routing accessors are
`private[spark]`; the only
externally visible API change is that the `DeveloperApi`
`SparkEnv.shuffleManager` accessor is now
`deprecated` (still functional, pointing callers to the explicit accessors).
### How was this patch tested?
New unit tests in `PipelinedShuffleRoutingSuite` (16 tests) verify the
routing and the manager-kind
model:
- `shuffleManagerFor` routes a regular dependency to the blocking manager
and a
`PipelinedShuffleDependency` to the pipelined manager;
- the handle a dependency mints comes from the routed manager (so driver
and executor agree), and
`blockingShuffleManager` is the plain configured manager (no wrapper
installed in front of it);
- the id-only cleanup path notifies both managers (and is a no-op, not an
NPE, before the managers
are initialized), and both managers are stopped on `SparkContext` stop;
- `spark.shuffle.manager.incremental` resolves the same short aliases as
`spark.shuffle.manager`, a
bad incremental class name fails fast at startup, a blocking manager in
the incremental slot is
rejected, and a pipelined manager in the default slot is rejected;
- a manager's type declares its kind (`SortShuffleManager` is a
`BlockingShuffleManager`;
`StreamingShuffleManager` is a `PipelinedShuffleManager` and not a
`BlockingShuffleManager`), and
`SparkEnv.shuffleBlockResolver` is defined for the blocking manager.
`StreamingShuffleManagerSuite` verifies the streaming output tracker is
initialized by default (the
streaming incremental manager), for an explicit incremental
`StreamingShuffleManager`, and for a
`MultiShuffleManager` default. Existing coverage of the block-by-id and
cleanup paths was updated
and rerun: `BlockManagerSuite` (including the deferred-`ShuffleManager`
init and unsupported-resolver
paths), `ContextCleanerSuite` (the real `RemoveShuffle` RPC through
`BlockManagerMasterEndpoint`),
`SortShuffleSuite`, `SparkSubmitSuite` (the SPARK-45762 user-jar
`ShuffleManager` plugin, now
implementing `BlockingShuffleManager`), and `ShuffleDependencySuite` (which
also asserts a
`PipelinedShuffleDependency` never allows push-based shuffle merge).
```
build/sbt 'core/testOnly
org.apache.spark.shuffle.PipelinedShuffleRoutingSuite'
...
Tests: succeeded 16, failed 0
```
### Was this patch authored or co-authored using generative AI tooling?
Co-authored: Claude Code (Opus 4.8)
Closes #57286 from jerrypeng/stack/pipelined-shuffle-pr2-routing.
Authored-by: Boyang Jerry Peng <[email protected]>
Signed-off-by: Wenchen Fan <[email protected]>
---
.../main/scala/org/apache/spark/Dependency.scala | 73 ++++-
.../src/main/scala/org/apache/spark/SparkEnv.scala | 208 +++++++++++---
.../org/apache/spark/errors/SparkCoreErrors.scala | 6 +-
.../org/apache/spark/internal/config/package.scala | 15 +
.../scala/org/apache/spark/rdd/CoGroupedRDD.scala | 2 +-
.../scala/org/apache/spark/rdd/ShuffledRDD.scala | 2 +-
.../scala/org/apache/spark/rdd/SubtractedRDD.scala | 2 +-
.../org/apache/spark/shuffle/ShuffleManager.scala | 59 +++-
.../spark/shuffle/ShuffleWriteProcessor.scala | 71 ++---
.../spark/shuffle/sort/SortShuffleManager.scala | 3 +-
.../shuffle/streaming/MultiShuffleManager.scala | 12 +-
.../streaming/StreamingShuffleManager.scala | 8 +-
.../org/apache/spark/storage/BlockManager.scala | 44 +--
.../spark/storage/BlockManagerMasterEndpoint.scala | 14 +-
.../storage/BlockManagerStorageEndpoint.scala | 11 +-
.../test/scala/org/apache/spark/ShuffleSuite.scala | 2 +-
.../org/apache/spark/deploy/SparkSubmitSuite.scala | 2 +-
.../internal/plugin/PluginContainerSuite.scala | 2 +-
.../apache/spark/scheduler/AQEShuffledRDD.scala | 2 +-
.../shuffle/PipelinedShuffleRoutingSuite.scala | 317 +++++++++++++++++++++
.../spark/shuffle/ShuffleDependencySuite.scala | 86 ++++++
.../streaming/StreamingShuffleManagerSuite.scala | 34 ++-
.../streaming/StreamingShuffleReaderSuite.scala | 7 +-
.../shuffle/streaming/StreamingShuffleSuite.scala | 19 +-
.../streaming/StreamingShuffleWriterSuite.scala | 7 +-
.../storage/BlockManagerReplicationSuite.scala | 4 +-
.../apache/spark/storage/BlockManagerSuite.scala | 9 +-
.../spark/sql/execution/ShuffledRowRDD.scala | 8 +-
.../execution/exchange/ShuffleExchangeExec.scala | 5 +-
.../streaming/ReceivedBlockHandlerSuite.scala | 2 +-
30 files changed, 878 insertions(+), 158 deletions(-)
diff --git a/core/src/main/scala/org/apache/spark/Dependency.scala
b/core/src/main/scala/org/apache/spark/Dependency.scala
index da97aff5b344..ac2a662e2317 100644
--- a/core/src/main/scala/org/apache/spark/Dependency.scala
+++ b/core/src/main/scala/org/apache/spark/Dependency.scala
@@ -134,7 +134,10 @@ class ShuffleDependency[K: ClassTag, V: ClassTag, C:
ClassTag](
// is enabled
private[this] var _shuffleMergeAllowed = canShuffleMergeBeEnabled()
- val shuffleHandle: ShuffleHandle =
_rdd.context.env.shuffleManager.registerShuffle(
+ // Route to the manager for this dependency's type (the incremental manager
for a
+ // PipelinedShuffleDependency, the default manager otherwise). The handle is
minted once here on
+ // the driver, so executors that later read it are served by the same
manager.
+ val shuffleHandle: ShuffleHandle =
_rdd.context.env.shuffleManagerFor(this).registerShuffle(
shuffleId, this)
private[spark] def setShuffleMergeAllowed(shuffleMergeAllowed: Boolean):
Unit = {
@@ -258,6 +261,74 @@ class ShuffleDependency[K: ClassTag, V: ClassTag, C:
ClassTag](
}
+/**
+ * :: DeveloperApi ::
+ * A [[ShuffleDependency]] whose output can be read incrementally: a consumer
stage may begin
+ * reading the shuffle output while the producer stage is still running,
rather than waiting for the
+ * producer's full, materialized output.
+ *
+ * This is a subtype of [[ShuffleDependency]] -- and thus, like it, a
first-class dependency kind
+ * alongside [[NarrowDependency]] under [[Dependency]]. It is intended to be
the marker the
+ * `DAGScheduler` will use to decide that the producer and consumer
+ * stages connected by this edge may run concurrently (a "pipelined group"),
and that the shuffle
+ * layer should serve this shuffle with an incremental shuffle implementation.
A plain
+ * [[ShuffleDependency]] keeps the existing semantics: its output is fully
materialized before any
+ * consumer reads it.
+ *
+ * Two behaviors are active from construction, driven by this type: shuffle
registration routes to
+ * the incremental (pipelined) shuffle manager via
`SparkEnv.shuffleManagerFor` (a plain
+ * [[ShuffleDependency]] goes to the blocking manager), and push-based shuffle
merge is
+ * unconditionally disabled (`setShuffleMergeAllowed(false)`; see below)
because merge exposes
+ * output only after a post-completion finalize step and would register merge
results for a
+ * transient shuffle. Beyond those, the concurrent scheduling behavior is
added separately by the
+ * `DAGScheduler` components that match on this type; code that only matches
the parent
+ * `ShuffleDependency` still treats it as an ordinary (materialized) shuffle
for everything else.
+ *
+ * The name is *pipelined* rather than *streaming*: reading producer output as
it is produced is a
+ * general execution capability (software-pipelining of dependent stages), not
specific to
+ * streaming. Streaming / real-time mode is the first caller, but nothing here
is
+ * streaming-specific.
+ *
+ * Note: the parent's `checksumMismatchFullRetryEnabled` /
+ * `checksumMismatchQueryLevelRollbackEnabled` parameters are intentionally
not exposed here, so
+ * they stay at their `false` defaults for a pipelined shuffle. Their checksum
retry / query-level
+ * rollback recomputes and re-runs succeeding stages after a mismatch; in a
pipelined group any
+ * failure aborts the whole group and the caller reruns from scratch, so that
stage-level recompute
+ * never fires -- the mechanism is moot by construction (it is also
incompatible with a consumer
+ * that has already read the output incrementally). Leaving the params unset
keeps the idiom
+ * unreachable.
+ */
+@DeveloperApi
+class PipelinedShuffleDependency[K: ClassTag, V: ClassTag, C: ClassTag](
+ _rdd: RDD[_ <: Product2[K, V]],
+ partitioner: Partitioner,
+ serializer: Serializer = SparkEnv.get.serializer,
+ keyOrdering: Option[Ordering[K]] = None,
+ aggregator: Option[Aggregator[K, V, C]] = None,
+ mapSideCombine: Boolean = false,
+ shuffleWriterProcessor: ShuffleWriteProcessor = new ShuffleWriteProcessor,
+ rowBasedChecksums: Array[RowBasedChecksum] =
ShuffleDependency.EMPTY_ROW_BASED_CHECKSUMS)
+ extends ShuffleDependency[K, V, C](
+ _rdd,
+ partitioner,
+ serializer,
+ keyOrdering,
+ aggregator,
+ mapSideCombine,
+ shuffleWriterProcessor,
+ rowBasedChecksums) {
+
+ // Push-based shuffle merge is incompatible with a pipelined
(incrementally-readable) shuffle: it
+ // exposes output only after a post-completion "finalize" step, the opposite
of incremental reads,
+ // and would register merge results in the MapOutputTracker for a transient
shuffle that must not
+ // outlive its group. ShuffleDependency enables merge by default whenever
push-based shuffle is on
+ // cluster-wide; disable it here so a pipelined producer never acquires
merger locations. Without
+ // this, ShuffleWriteProcessor.write would reach the push path and
dereference the incremental
+ // manager's shuffleBlockResolver, which the streaming manager does not
support.
+ setShuffleMergeAllowed(false)
+}
+
+
/**
* :: DeveloperApi ::
* Represents a one-to-one dependency between partitions of the parent and
child RDDs.
diff --git a/core/src/main/scala/org/apache/spark/SparkEnv.scala
b/core/src/main/scala/org/apache/spark/SparkEnv.scala
index 2bc10adcbe3a..4cc2c45d7ee8 100644
--- a/core/src/main/scala/org/apache/spark/SparkEnv.scala
+++ b/core/src/main/scala/org/apache/spark/SparkEnv.scala
@@ -45,7 +45,8 @@ import org.apache.spark.scheduler.{LiveListenerBus,
OutputCommitCoordinator}
import
org.apache.spark.scheduler.OutputCommitCoordinator.OutputCommitCoordinatorEndpoint
import org.apache.spark.security.CryptoStreamUtils
import org.apache.spark.serializer.{JavaSerializer, Serializer,
SerializerManager}
-import org.apache.spark.shuffle.ShuffleManager
+import org.apache.spark.shuffle.{BlockingShuffleManager,
PipelinedShuffleManager}
+import org.apache.spark.shuffle.{ShuffleBlockResolver, ShuffleManager}
import org.apache.spark.shuffle.streaming.{MultiShuffleManager,
StreamingShuffleManager}
import org.apache.spark.storage._
import org.apache.spark.udf.worker.UDFWorkerSpecification
@@ -76,15 +77,105 @@ class SparkEnv (
val outputCommitCoordinator: OutputCommitCoordinator,
val conf: SparkConf) extends Logging {
- // We initialize the ShuffleManager later in SparkContext and Executor to
allow
- // user jars to define custom ShuffleManagers.
- @volatile private var _shuffleManager: ShuffleManager = _
+ // The two shuffle managers are peers keyed by kind, not a default and an
override: a shuffle is
+ // routed to one or the other by its dependency type via
`shuffleManagerFor`, so neither is ever
+ // installed "behind" the other.
+ //
+ // The blocking manager (spark.shuffle.manager) serves all regular,
materialized shuffles and owns
+ // block-by-id resolution. It is always a BlockingShuffleManager --
initializeShuffleManager
+ // rejects a non-blocking manager in this slot -- so block-resolution has a
single, well-typed
+ // source. We initialize it later in SparkContext and Executor to allow user
jars to define custom
+ // ShuffleManagers.
+ @volatile private var _blockingShuffleManager: BlockingShuffleManager = _
+
+ // The pipelined manager (spark.shuffle.manager.incremental) serves pipelined
+ // (incrementally-readable) shuffle dependencies. It defaults to the built-in
+ // StreamingShuffleManager -- just as the blocking manager defaults to sort
-- and is always a
+ // PipelinedShuffleManager (initializeShuffleManager rejects a blocking
manager in this slot), so
+ // its output is never expected to be reachable through the block-manager
resolver. Like the
+ // blocking manager, it is initialized later in SparkContext and Executor to
allow user jars.
+ @volatile private var _pipelinedShuffleManager: PipelinedShuffleManager = _
// Latch to signal when the ShuffleManager has been initialized.
// Used to allow callers to wait for initialization.
private val shuffleManagerInitLatch = new CountDownLatch(1)
- def shuffleManager: ShuffleManager = _shuffleManager
+ /**
+ * The `BlockingShuffleManager` (configured by spark.shuffle.manager), which
serves all regular,
+ * materialized shuffle dependencies and owns block-by-id resolution. Use
this when the intent is
+ * specifically the blocking manager -- e.g. inspecting its concrete type.
To serve a specific
+ * shuffle's reads/writes, use `shuffleManagerFor`, which routes by
dependency type; to resolve a
+ * block by id, use `shuffleBlockResolver`.
+ */
+ private[spark] def blockingShuffleManager: BlockingShuffleManager =
_blockingShuffleManager
+
+ /**
+ * The `PipelinedShuffleManager` (configured by
spark.shuffle.manager.incremental, defaulting to
+ * the built-in streaming manager) that serves pipelined shuffle
dependencies. A pipelined shuffle
+ * is read incrementally and served out-of-band, so this manager never
provides a
+ * `ShuffleBlockResolver` of its own.
+ */
+ private[spark] def pipelinedShuffleManager: PipelinedShuffleManager =
_pipelinedShuffleManager
+
+ /**
+ * The `ShuffleBlockResolver` used to resolve shuffle blocks by id (reads,
push-merge, and
+ * decommission migration), or `None` before the shuffle manager is
initialized. Block resolution
+ * is only ever needed for regular, materialized shuffles, so it always
comes from the blocking
+ * manager (the only manager that produces block-manager-addressed blocks);
a pipelined shuffle is
+ * served out-of-band and never resolved here. Callers that only need to
resolve a block should
+ * use this rather than reaching through a `ShuffleManager`.
+ */
+ private[spark] def shuffleBlockResolver: Option[ShuffleBlockResolver] =
+ Option(_blockingShuffleManager).map(_.shuffleBlockResolver)
+
+ /**
+ * Retained for binary compatibility; returns the blocking manager. Prefer
`shuffleManagerFor`
+ * (route a shuffle by its dependency) or `blockingShuffleManager` (the
blocking manager
+ * explicitly), so the routing decision is explicit at the call site.
+ */
+ @deprecated("use shuffleManagerFor(dependency) to route a shuffle by type,
or " +
+ "blockingShuffleManager for the blocking manager explicitly", "4.3.0")
+ def shuffleManager: ShuffleManager = _blockingShuffleManager
+
+ /**
+ * The `ShuffleManager` that serves the given shuffle, chosen by dependency
type: a
+ * [[PipelinedShuffleDependency]] is served by the pipelined manager
+ * (spark.shuffle.manager.incremental, defaulting to the built-in streaming
manager), every other
+ * [[ShuffleDependency]] by the blocking manager (spark.shuffle.manager).
This is the single
+ * routing point for shuffle I/O; the decision is a pure function of the
dependency, so the driver
+ * (at `registerShuffle`) and every executor (at `getWriter` / `getReader`)
agree without any
+ * shared routing state or config re-read -- the dependency (a serialized
field, deterministic on
+ * driver and executors) is always in hand at these sites, so no shuffleId
-> manager tracking is
+ * needed.
+ *
+ * Whether a job may use a pipelined dependency at all is a separate,
scheduler-level decision
+ * (see the fail-fast checks in `DAGScheduler`); this method only picks the
implementation.
+ */
+ private[spark] def shuffleManagerFor(dependency: ShuffleDependency[_, _,
_]): ShuffleManager =
+ dependency match {
+ case _: PipelinedShuffleDependency[_, _, _] =>
+ _pipelinedShuffleManager
+ case _ => _blockingShuffleManager
+ }
+
+ /**
+ * Unregisters the shuffle from every configured manager (default and, if
present, incremental).
+ * Used by the `RemoveShuffle` cleanup path, which holds only a shuffleId
and cannot know which
+ * manager owns it -- and must reach the owning manager on every node,
including one that never
+ * performed this shuffle's I/O. Notifying all managers is safe:
`unregisterShuffle` for an
+ * unknown id is a no-op. Returns true if any manager reports it removed
metadata; false if no
+ * manager is initialized yet (a RemoveShuffle can arrive before the
deferred init runs).
+ */
+ private[spark] def unregisterShuffleFromAllManagers(shuffleId: Int): Boolean
= {
+ // Both managers are null until initializeShuffleManager runs (deferred to
allow user jars), so
+ // guard them; a RemoveShuffle before init is a no-op, as it was before
routing.
+ val blockingResult =
+ _blockingShuffleManager != null &&
_blockingShuffleManager.unregisterShuffle(shuffleId)
+ // OR the pipelined result in without short-circuiting, so both are always
notified.
+ val pipelinedResult =
+ _pipelinedShuffleManager != null &&
_pipelinedShuffleManager.unregisterShuffle(shuffleId)
+ blockingResult || pipelinedResult
+ }
/**
* Wait for the ShuffleManager to be initialized within the specified
timeout.
@@ -102,7 +193,7 @@ class SparkEnv (
* @return true if the ShuffleManager is initialized, false otherwise
*/
private[spark] def isShuffleManagerInitialized: Boolean = {
- _shuffleManager != null
+ _blockingShuffleManager != null
}
// We initialize the MemoryManager later in SparkContext after DriverPlugin
is loaded
@@ -185,8 +276,11 @@ class SparkEnv (
udfDispatcherManager.foreach(_.close())
mapOutputTracker.stop()
_streamingShuffleOutputTracker.foreach(_.stop())
- if (shuffleManager != null) {
- shuffleManager.stop()
+ if (_blockingShuffleManager != null) {
+ _blockingShuffleManager.stop()
+ }
+ if (_pipelinedShuffleManager != null) {
+ _pipelinedShuffleManager.stop()
}
broadcastManager.stop()
blockManager.stop()
@@ -295,10 +389,42 @@ class SparkEnv (
}
private[spark] def initializeShuffleManager(): Unit = {
- Preconditions.checkState(null == _shuffleManager,
- "Shuffle manager already initialized to %s", _shuffleManager)
+ Preconditions.checkState(null == _blockingShuffleManager,
+ "Shuffle manager already initialized to %s", _blockingShuffleManager)
try {
- _shuffleManager = ShuffleManager.create(conf,
SparkContext.isDriver(executorId))
+ val isDriver = SparkContext.isDriver(executorId)
+ // The blocking manager (spark.shuffle.manager) serves all regular,
materialized shuffle
+ // dependencies and owns block-by-id resolution, so it must be a
BlockingShuffleManager. A
+ // manager that declares neither kind, or a pipelined-only one, is
rejected here rather than
+ // silently failing block resolution later.
+ _blockingShuffleManager = ShuffleManager.create(conf, isDriver) match {
+ case blocking: BlockingShuffleManager => blocking
+ case other =>
+ throw new IllegalArgumentException(
+ s"${config.SHUFFLE_MANAGER.key} must be a BlockingShuffleManager,
but " +
+ s"${other.getClass.getName} is not. A blocking manager serves
regular, " +
+ s"materialized shuffles and resolves blocks by id; configure a
pipelined manager " +
+ s"via ${config.SHUFFLE_MANAGER_INCREMENTAL.key} instead.")
+ }
+ // The pipelined manager (spark.shuffle.manager.incremental) serves
pipelined shuffle
+ // dependencies. It is a peer of the blocking manager -- see
`shuffleManagerFor` -- not a
+ // wrapper installed as the top-level manager, and is created the same
way: it defaults to the
+ // built-in streaming manager (just as the blocking manager defaults to
sort) and must be a
+ // PipelinedShuffleManager, since its output is read incrementally and
served out-of-band and
+ // so produces no block-manager blocks.
+ _pipelinedShuffleManager =
Utils.instantiateSerializerOrShuffleManager[ShuffleManager](
+ // Resolve short aliases ("sort", "tungsten-sort", "streaming") the
same way the blocking
+ // manager does, so spark.shuffle.manager.incremental accepts the same
values.
+
ShuffleManager.resolveShortName(conf.get(config.SHUFFLE_MANAGER_INCREMENTAL)),
+ conf, isDriver) match {
+ case pipelined: PipelinedShuffleManager => pipelined
+ case other =>
+ throw new IllegalArgumentException(
+ s"${config.SHUFFLE_MANAGER_INCREMENTAL.key} must be a
PipelinedShuffleManager, but " +
+ s"${other.getClass.getName} is not. Only a pipelined manager
(output read " +
+ s"incrementally and served out-of-band) belongs in this slot; a
blocking manager " +
+ s"belongs in ${config.SHUFFLE_MANAGER.key}.")
+ }
} finally {
// Signal that the ShuffleManager has been initialized
shuffleManagerInitLatch.countDown()
@@ -306,8 +432,9 @@ class SparkEnv (
initializeStreamingShuffleOutputTracker()
}
- // Holds the streaming shuffle output tracker, which is only present when
the configured
- // shuffle manager requires it (i.e., StreamingShuffleManager or
MultiShuffleManager).
+ // Holds the streaming shuffle output tracker, which is only present when
the configured shuffle
+ // managers require it (i.e., a StreamingShuffleManager as the pipelined
manager, or a
+ // MultiShuffleManager as the blocking manager).
@volatile private var _streamingShuffleOutputTracker:
Option[StreamingShuffleOutputTracker] =
None
@@ -323,28 +450,42 @@ class SparkEnv (
return
}
- val shuffleManagerName = ShuffleManager.getShuffleManagerClassName(conf)
- if (shuffleManagerName == classOf[StreamingShuffleManager].getName
- || shuffleManagerName == classOf[MultiShuffleManager].getName) {
- val tracker = if (SparkContext.isDriver(executorId)) {
- new StreamingShuffleOutputTrackerMaster(conf)
- } else {
- new StreamingShuffleOutputTrackerWorker(conf)
- }
+ // The tracker is needed when the pipelined manager
(spark.shuffle.manager.incremental) is a
+ // StreamingShuffleManager -- which is the default. Inspect the
already-instantiated manager
+ // rather than re-reading the config; this runs at the end of
initializeShuffleManager, so the
+ // manager is non-null here.
+ val incrementalIsStreaming =
+ _pipelinedShuffleManager.isInstanceOf[StreamingShuffleManager]
+ // It is also needed when a MultiShuffleManager is the blocking manager
(spark.shuffle.manager):
+ // it internally routes some shuffles to streaming. A bare
StreamingShuffleManager cannot be the
+ // blocking manager -- it is pipelined and rejected from that slot in
initializeShuffleManager.
+ // TODO: remove this MultiShuffleManager clause once MultiShuffleManager
is removed and
+ // streaming shuffles are served only through the incremental (pipelined)
manager slot.
+ val blockingIsMulti =
_blockingShuffleManager.isInstanceOf[MultiShuffleManager]
+ if (incrementalIsStreaming || blockingIsMulti) {
+ createStreamingShuffleOutputTracker()
+ }
+ }
- if (SparkContext.isDriver(executorId)) {
- tracker.trackerEndpoint = rpcEnv.setupEndpoint(
- StreamingShuffleOutputTracker.ENDPOINT_NAME,
- new StreamingShuffleOutputTrackerMasterEndpoint(
- rpcEnv,
- tracker.asInstanceOf[StreamingShuffleOutputTrackerMaster],
- conf))
- } else {
- tracker.trackerEndpoint = RpcUtils.makeDriverRef(
- StreamingShuffleOutputTracker.ENDPOINT_NAME, conf, rpcEnv)
- }
- _streamingShuffleOutputTracker = Some(tracker)
+ private def createStreamingShuffleOutputTracker(): Unit = {
+ val tracker = if (SparkContext.isDriver(executorId)) {
+ new StreamingShuffleOutputTrackerMaster(conf)
+ } else {
+ new StreamingShuffleOutputTrackerWorker(conf)
+ }
+
+ if (SparkContext.isDriver(executorId)) {
+ tracker.trackerEndpoint = rpcEnv.setupEndpoint(
+ StreamingShuffleOutputTracker.ENDPOINT_NAME,
+ new StreamingShuffleOutputTrackerMasterEndpoint(
+ rpcEnv,
+ tracker.asInstanceOf[StreamingShuffleOutputTrackerMaster],
+ conf))
+ } else {
+ tracker.trackerEndpoint = RpcUtils.makeDriverRef(
+ StreamingShuffleOutputTracker.ENDPOINT_NAME, conf, rpcEnv)
}
+ _streamingShuffleOutputTracker = Some(tracker)
}
private[spark] def initializeMemoryManager(
@@ -550,7 +691,6 @@ object SparkEnv extends Logging {
None
}, blockManagerInfo,
mapOutputTracker.asInstanceOf[MapOutputTrackerMaster],
- _shuffleManager = null,
isDriver)),
registerOrLookupEndpoint(
BlockManagerMaster.DRIVER_HEARTBEAT_ENDPOINT_NAME,
diff --git a/core/src/main/scala/org/apache/spark/errors/SparkCoreErrors.scala
b/core/src/main/scala/org/apache/spark/errors/SparkCoreErrors.scala
index 30660a177416..4813f3f94107 100644
--- a/core/src/main/scala/org/apache/spark/errors/SparkCoreErrors.scala
+++ b/core/src/main/scala/org/apache/spark/errors/SparkCoreErrors.scala
@@ -29,7 +29,7 @@ import org.apache.spark.internal.config.IO_COMPRESSION_CODEC
import org.apache.spark.io.CompressionCodec.FALLBACK_COMPRESSION_CODEC
import org.apache.spark.memory.SparkOutOfMemoryError
import
org.apache.spark.scheduler.{BarrierJobRunWithDynamicAllocationException,
BarrierJobSlotsNumberCheckFailed, BarrierJobUnsupportedRDDChainException}
-import org.apache.spark.shuffle.{FetchFailedException, ShuffleManager}
+import org.apache.spark.shuffle.{FetchFailedException, ShuffleBlockResolver}
import org.apache.spark.storage.{BlockId, BlockManagerId,
BlockNotFoundException, BlockSavedOnDecommissionedBlockManagerException,
RDDBlockId, UnrecognizedBlockId}
/**
@@ -335,13 +335,13 @@ private[spark] object SparkCoreErrors {
}
def unexpectedShuffleBlockWithUnsupportedResolverError(
- shuffleManager: ShuffleManager,
+ shuffleBlockResolver: ShuffleBlockResolver,
blockId: BlockId): Throwable = {
new SparkException(
errorClass = "_LEGACY_ERROR_TEMP_3035",
messageParameters = Map(
"blockId" -> s"$blockId",
- "shuffleBlockResolver" -> s"${shuffleManager.shuffleBlockResolver}"
+ "shuffleBlockResolver" -> s"$shuffleBlockResolver"
),
cause = null
)
diff --git a/core/src/main/scala/org/apache/spark/internal/config/package.scala
b/core/src/main/scala/org/apache/spark/internal/config/package.scala
index 685ef6e11329..c8bd6410be27 100644
--- a/core/src/main/scala/org/apache/spark/internal/config/package.scala
+++ b/core/src/main/scala/org/apache/spark/internal/config/package.scala
@@ -1774,6 +1774,21 @@ package object config {
.stringConf
.createWithDefault("sort")
+ private[spark] val SHUFFLE_MANAGER_INCREMENTAL =
+ ConfigBuilder("spark.shuffle.manager.incremental")
+ .doc("The ShuffleManager implementation used for pipelined
(incrementally-readable) " +
+ "shuffle dependencies, whose consumer stage may begin reading output
while the producer " +
+ "stage is still running. Regular shuffle dependencies continue to use
" +
+ "spark.shuffle.manager. The two managers coexist in the same JVM and
each shuffle is " +
+ "routed by its dependency type, so a cluster can run both regular and
pipelined (e.g. " +
+ "real-time mode) queries. Defaults to the built-in streaming shuffle
manager; accepts " +
+ "the short alias \"streaming\" (the default) or a fully-qualified " +
+ "PipelinedShuffleManager class name.")
+ .version("4.3.0")
+ .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE)
+ .stringConf
+ .createWithDefault("streaming")
+
private[spark] val SHUFFLE_REDUCE_LOCALITY_ENABLE =
ConfigBuilder("spark.shuffle.reduceLocality.enabled")
.doc("Whether to compute locality preferences for reduce tasks")
diff --git a/core/src/main/scala/org/apache/spark/rdd/CoGroupedRDD.scala
b/core/src/main/scala/org/apache/spark/rdd/CoGroupedRDD.scala
index 8bac6e736119..d9f117002c57 100644
--- a/core/src/main/scala/org/apache/spark/rdd/CoGroupedRDD.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/CoGroupedRDD.scala
@@ -143,7 +143,7 @@ class CoGroupedRDD[K: ClassTag](
case shuffleDependency: ShuffleDependency[_, _, _] =>
// Read map outputs of shuffle
val metrics = context.taskMetrics().createTempShuffleReadMetrics()
- val it = SparkEnv.get.shuffleManager
+ val it = SparkEnv.get.shuffleManagerFor(shuffleDependency)
.getReader(
shuffleDependency.shuffleHandle, split.index, split.index + 1,
context, metrics)
.read()
diff --git a/core/src/main/scala/org/apache/spark/rdd/ShuffledRDD.scala
b/core/src/main/scala/org/apache/spark/rdd/ShuffledRDD.scala
index 0930a5c9cfb9..0c48d4da0c28 100644
--- a/core/src/main/scala/org/apache/spark/rdd/ShuffledRDD.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/ShuffledRDD.scala
@@ -102,7 +102,7 @@ class ShuffledRDD[K: ClassTag, V: ClassTag, C: ClassTag](
override def compute(split: Partition, context: TaskContext): Iterator[(K,
C)] = {
val dep = dependencies.head.asInstanceOf[ShuffleDependency[K, V, C]]
val metrics = context.taskMetrics().createTempShuffleReadMetrics()
- SparkEnv.get.shuffleManager.getReader(
+ SparkEnv.get.shuffleManagerFor(dep).getReader(
dep.shuffleHandle, split.index, split.index + 1, context, metrics)
.read()
.asInstanceOf[Iterator[(K, C)]]
diff --git a/core/src/main/scala/org/apache/spark/rdd/SubtractedRDD.scala
b/core/src/main/scala/org/apache/spark/rdd/SubtractedRDD.scala
index 27dfdb4daa2c..f9562c3b1586 100644
--- a/core/src/main/scala/org/apache/spark/rdd/SubtractedRDD.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/SubtractedRDD.scala
@@ -108,7 +108,7 @@ private[spark] class SubtractedRDD[K: ClassTag, V:
ClassTag, W: ClassTag](
case shuffleDependency: ShuffleDependency[_, _, _] =>
val metrics = context.taskMetrics().createTempShuffleReadMetrics()
- val iter = SparkEnv.get.shuffleManager
+ val iter = SparkEnv.get.shuffleManagerFor(shuffleDependency)
.getReader(
shuffleDependency.shuffleHandle,
partition.index,
diff --git a/core/src/main/scala/org/apache/spark/shuffle/ShuffleManager.scala
b/core/src/main/scala/org/apache/spark/shuffle/ShuffleManager.scala
index 8e4636cfefbf..df47f7552688 100644
--- a/core/src/main/scala/org/apache/spark/shuffle/ShuffleManager.scala
+++ b/core/src/main/scala/org/apache/spark/shuffle/ShuffleManager.scala
@@ -28,12 +28,17 @@ import org.apache.spark.util.Utils
* and on each executor, based on the spark.shuffle.manager setting. The
driver registers shuffles
* with it, and executors (or tasks running locally in the driver) can ask to
read and write data.
*
+ * An implementation declares its kind by extending one of its two subtypes:
+ * [[BlockingShuffleManager]] (output is materialized as
block-manager-addressed blocks served
+ * through a [[ShuffleBlockResolver]]) or [[PipelinedShuffleManager]] (output
is read incrementally
+ * and served out-of-band).
+ *
* NOTE:
* 1. This will be instantiated by SparkEnv so its constructor can take a
SparkConf and
* boolean isDriver as parameters.
- * 2. This contains a method ShuffleBlockResolver which interacts with
External Shuffle Service
- * when it is enabled. Need to pay attention to that, if implementing a custom
ShuffleManager, to
- * make sure the custom ShuffleManager could co-exist with External Shuffle
Service.
+ * 2. A [[BlockingShuffleManager]] exposes a ShuffleBlockResolver which
interacts with the External
+ * Shuffle Service when it is enabled. Pay attention to this when implementing
a custom shuffle
+ * manager to make sure it can coexist with the External Shuffle Service.
*/
private[spark] trait ShuffleManager {
@@ -86,19 +91,46 @@ private[spark] trait ShuffleManager {
/**
* Remove a shuffle's metadata from the ShuffleManager.
+ *
+ * Implementations must treat an unknown or already-removed shuffleId as an
idempotent no-op --
+ * returning false rather than throwing or mutating unrelated state. The
id-only `RemoveShuffle`
+ * cleanup path cannot tell which manager owns a shuffle, so it broadcasts
to every configured
+ * manager (see `SparkEnv.unregisterShuffleFromAllManagers`), including ones
that never handled
+ * this shuffle.
+ *
* @return true if the metadata removed successfully, otherwise false.
*/
def unregisterShuffle(shuffleId: Int): Boolean
+ /** Shut down this ShuffleManager. */
+ def stop(): Unit
+}
+
+/**
+ * A [[ShuffleManager]] that materializes shuffle output as addressable blocks
served through the
+ * block manager (reads, push-based merge, and decommission migration all go
through its
+ * [[ShuffleBlockResolver]]). This is the traditional shuffle model: a
consumer stage reads the
+ * producer's output only after it is fully written.
+ * [[org.apache.spark.shuffle.sort.SortShuffleManager]] is the built-in
implementation. A manager's
+ * type declares its kind -- match on `BlockingShuffleManager` to reach the
resolver rather than
+ * assuming every `ShuffleManager` provides one.
+ */
+private[spark] trait BlockingShuffleManager extends ShuffleManager {
/**
* Return a resolver capable of retrieving shuffle block data based on block
coordinates.
*/
def shuffleBlockResolver: ShuffleBlockResolver
-
- /** Shut down this ShuffleManager. */
- def stop(): Unit
}
+/**
+ * A [[ShuffleManager]] whose output is read incrementally: a consumer stage
may begin reading while
+ * the producer is still running (see
[[org.apache.spark.PipelinedShuffleDependency]]). Such a
+ * manager serves its output out-of-band and does not produce
block-manager-addressed blocks, so it
+ * has no [[ShuffleBlockResolver]].
[[org.apache.spark.shuffle.streaming.StreamingShuffleManager]]
+ * is the built-in implementation.
+ */
+private[spark] trait PipelinedShuffleManager extends ShuffleManager
+
/**
* Utility companion object to create a ShuffleManager given a spark
configuration.
*/
@@ -108,12 +140,21 @@ private[spark] object ShuffleManager {
getShuffleManagerClassName(conf), conf, isDriver)
}
- def getShuffleManagerClassName(conf: SparkConf): String = {
+ def getShuffleManagerClassName(conf: SparkConf): String =
+ resolveShortName(conf.get(config.SHUFFLE_MANAGER))
+
+ /**
+ * Resolve a short shuffle-manager alias ("sort", "tungsten-sort",
"streaming") to its
+ * fully-qualified class name, passing any other value through unchanged.
Shared by the default
+ * manager (spark.shuffle.manager) and the incremental manager
(spark.shuffle.manager.incremental)
+ * so the same aliases are accepted for both. "streaming" is the default for
the incremental slot.
+ */
+ def resolveShortName(shuffleMgrName: String): String = {
val shortShuffleMgrNames = Map(
"sort" ->
classOf[org.apache.spark.shuffle.sort.SortShuffleManager].getName,
- "tungsten-sort" ->
classOf[org.apache.spark.shuffle.sort.SortShuffleManager].getName)
+ "tungsten-sort" ->
classOf[org.apache.spark.shuffle.sort.SortShuffleManager].getName,
+ "streaming" ->
classOf[org.apache.spark.shuffle.streaming.StreamingShuffleManager].getName)
- val shuffleMgrName = conf.get(config.SHUFFLE_MANAGER)
shortShuffleMgrNames.getOrElse(shuffleMgrName.toLowerCase(Locale.ROOT),
shuffleMgrName)
}
}
diff --git
a/core/src/main/scala/org/apache/spark/shuffle/ShuffleWriteProcessor.scala
b/core/src/main/scala/org/apache/spark/shuffle/ShuffleWriteProcessor.scala
index c397ac03ed92..14d2f3e9b11e 100644
--- a/core/src/main/scala/org/apache/spark/shuffle/ShuffleWriteProcessor.scala
+++ b/core/src/main/scala/org/apache/spark/shuffle/ShuffleWriteProcessor.scala
@@ -49,7 +49,7 @@ private[spark] class ShuffleWriteProcessor extends
Serializable with Logging {
context: TaskContext): MapStatus = {
var writer: ShuffleWriter[Any, Any] = null
try {
- val manager = SparkEnv.get.shuffleManager
+ val manager = SparkEnv.get.shuffleManagerFor(dep)
writer = manager.getWriter[Any, Any](
dep.shuffleHandle,
mapId,
@@ -72,37 +72,44 @@ private[spark] class ShuffleWriteProcessor extends
Serializable with Logging {
// block push requests. It delegates pushing the blocks to a different
thread-pool -
// ShuffleBlockPusher.BLOCK_PUSHER_POOL.
if (!dep.shuffleMergeFinalized) {
- manager.shuffleBlockResolver match {
- case resolver: IndexShuffleBlockResolver =>
- logInfo(log"Shuffle merge enabled with" +
- log" ${MDC(NUM_MERGER_LOCATIONS, dep.getMergerLocs.size)}
merger locations" +
- log" for stage ${MDC(STAGE_ID, context.stageId())}" +
- log" with shuffle ID ${MDC(SHUFFLE_ID, dep.shuffleId)}")
- logDebug(s"Starting pushing blocks for the task
${context.taskAttemptId()}")
- val dataFile = resolver.getDataFile(dep.shuffleId, mapId)
- val blockPusher = new ShuffleBlockPusher(SparkEnv.get.conf)
- // Register a post-status-update listener to defer push until
after the task
- // result has been sent to the driver. This ensures the driver
processes the
- // task result (and can mark stale partitions from speculative
duplicates)
- // before any push data reaches the merger, avoiding stale data
being merged
- // without detection.
- //
- // The listener callback runs on the Task thread and only does a
lightweight
- // submitTask; actual push I/O runs on BLOCK_PUSHER_POOL threads.
- context.addPostStatusUpdateListener(new PostStatusUpdateListener
{
- override def onStatusUpdateSent(context: TaskContext): Unit = {
- if (!context.isInterrupted() && !context.isFailed()) {
- logDebug(s"Task ${context.taskAttemptId()} status update
sent, " +
- s"proceeding with shuffle block push for shuffle
${dep.shuffleId}")
- blockPusher.initiateBlockPush(
- dataFile, writer.getPartitionLengths(), dep, mapIndex)
- } else {
- logInfo(s"Task ${context.taskAttemptId()} was " +
- s"${if (context.isInterrupted()) "killed" else
"failed"}, " +
- s"skipping shuffle block push for shuffle
${dep.shuffleId}")
- }
- }
- })
+ // Push-based merge only applies to a blocking shuffle that resolves
blocks via an
+ // IndexShuffleBlockResolver. A pipelined shuffle disables merge
(shuffleMergeFinalized is
+ // true above), so it never reaches here.
+ manager match {
+ case blocking: BlockingShuffleManager =>
+ blocking.shuffleBlockResolver match {
+ case resolver: IndexShuffleBlockResolver =>
+ logInfo(log"Shuffle merge enabled with" +
+ log" ${MDC(NUM_MERGER_LOCATIONS, dep.getMergerLocs.size)}
merger locations" +
+ log" for stage ${MDC(STAGE_ID, context.stageId())}" +
+ log" with shuffle ID ${MDC(SHUFFLE_ID, dep.shuffleId)}")
+ logDebug(s"Starting pushing blocks for the task
${context.taskAttemptId()}")
+ val dataFile = resolver.getDataFile(dep.shuffleId, mapId)
+ val blockPusher = new ShuffleBlockPusher(SparkEnv.get.conf)
+ // Register a post-status-update listener to defer push
until after the task
+ // result has been sent to the driver. This ensures the
driver processes the
+ // task result (and can mark stale partitions from
speculative duplicates)
+ // before any push data reaches the merger, avoiding stale
data being merged
+ // without detection.
+ //
+ // The listener callback runs on the Task thread and only
does a lightweight
+ // submitTask; actual push I/O runs on BLOCK_PUSHER_POOL
threads.
+ context.addPostStatusUpdateListener(new
PostStatusUpdateListener {
+ override def onStatusUpdateSent(context: TaskContext):
Unit = {
+ if (!context.isInterrupted() && !context.isFailed()) {
+ logDebug(s"Task ${context.taskAttemptId()} status
update sent, " +
+ s"proceeding with shuffle block push for shuffle
${dep.shuffleId}")
+ blockPusher.initiateBlockPush(
+ dataFile, writer.getPartitionLengths(), dep,
mapIndex)
+ } else {
+ logInfo(s"Task ${context.taskAttemptId()} was " +
+ s"${if (context.isInterrupted()) "killed" else
"failed"}, " +
+ s"skipping shuffle block push for shuffle
${dep.shuffleId}")
+ }
+ }
+ })
+ case _ =>
+ }
case _ =>
}
}
diff --git
a/core/src/main/scala/org/apache/spark/shuffle/sort/SortShuffleManager.scala
b/core/src/main/scala/org/apache/spark/shuffle/sort/SortShuffleManager.scala
index ca5f5bc28ae9..1644ab531d56 100644
--- a/core/src/main/scala/org/apache/spark/shuffle/sort/SortShuffleManager.scala
+++ b/core/src/main/scala/org/apache/spark/shuffle/sort/SortShuffleManager.scala
@@ -70,7 +70,8 @@ import org.apache.spark.util.collection.OpenHashSet
*
* For more details on these optimizations, see SPARK-7081.
*/
-private[spark] class SortShuffleManager(conf: SparkConf) extends
ShuffleManager with Logging {
+private[spark] class SortShuffleManager(conf: SparkConf)
+ extends BlockingShuffleManager with Logging {
import SortShuffleManager._
diff --git
a/core/src/main/scala/org/apache/spark/shuffle/streaming/MultiShuffleManager.scala
b/core/src/main/scala/org/apache/spark/shuffle/streaming/MultiShuffleManager.scala
index 9e63c9375955..3e2fa7ab4f68 100644
---
a/core/src/main/scala/org/apache/spark/shuffle/streaming/MultiShuffleManager.scala
+++
b/core/src/main/scala/org/apache/spark/shuffle/streaming/MultiShuffleManager.scala
@@ -22,7 +22,7 @@ import java.util.concurrent.ConcurrentHashMap
import org.apache.spark.{ShuffleDependency, SparkConf, SparkContext,
SparkException, TaskContext}
import org.apache.spark.internal.Logging
-import org.apache.spark.shuffle.{ShuffleBlockResolver, ShuffleHandle,
ShuffleManager, ShuffleReader, ShuffleReadMetricsReporter,
ShuffleWriteMetricsReporter, ShuffleWriter}
+import org.apache.spark.shuffle.{BlockingShuffleManager, ShuffleBlockResolver,
ShuffleHandle, ShuffleManager, ShuffleReader, ShuffleReadMetricsReporter,
ShuffleWriteMetricsReporter, ShuffleWriter}
import org.apache.spark.shuffle.sort.SortShuffleManager
import
org.apache.spark.shuffle.streaming.MultiShuffleManager.isStreamingShuffleEnabled
@@ -47,8 +47,14 @@ object MultiShuffleManager {
and normal queries that depends on sort shuffle to coexist in a cluster. Right
now, we only
allows configuration of shuffle manager at cluster level, so consider using
this shuffle
manager if you want to run batch and real time queries at the same time.
+
+Deprecated: this cluster-level single-slot approach is superseded by
per-dependency routing --
+configure spark.shuffle.manager (regular/blocking) and
spark.shuffle.manager.incremental
+(pipelined) so each shuffle is routed by its dependency type.
*/
-class MultiShuffleManager(conf: SparkConf) extends ShuffleManager with Logging
{
+@deprecated("Use per-dependency routing via spark.shuffle.manager and " +
+ "spark.shuffle.manager.incremental instead", "4.3.0")
+class MultiShuffleManager(conf: SparkConf) extends BlockingShuffleManager with
Logging {
// To make sure the type of shuffle manager used for a shuffle is the same
during its lifetime
private val shuffleIdToManager = new ConcurrentHashMap[Int, ShuffleManager]()
private var streamingShuffleManager: Option[StreamingShuffleManager] = None
@@ -129,6 +135,8 @@ class MultiShuffleManager(conf: SparkConf) extends
ShuffleManager with Logging {
}
}
+ // As a BlockingShuffleManager, MultiShuffleManager resolves blocks through
the inner
+ // SortShuffleManager it delegates regular shuffles to.
override def shuffleBlockResolver: ShuffleBlockResolver = {
shuffleIdToManager.synchronized {
if (sortShuffleManager.nonEmpty) {
diff --git
a/core/src/main/scala/org/apache/spark/shuffle/streaming/StreamingShuffleManager.scala
b/core/src/main/scala/org/apache/spark/shuffle/streaming/StreamingShuffleManager.scala
index 2bc9d83b6141..03ca257353c5 100644
---
a/core/src/main/scala/org/apache/spark/shuffle/streaming/StreamingShuffleManager.scala
+++
b/core/src/main/scala/org/apache/spark/shuffle/streaming/StreamingShuffleManager.scala
@@ -76,7 +76,7 @@ object StreamingShuffleManager extends Logging {
}
}
-private[spark] class StreamingShuffleManager extends ShuffleManager with
Logging {
+private[spark] class StreamingShuffleManager extends PipelinedShuffleManager
with Logging {
logInfo(log"Using StreamingShuffleManager")
@@ -118,11 +118,5 @@ private[spark] class StreamingShuffleManager extends
ShuffleManager with Logging
true
}
- override def shuffleBlockResolver: ShuffleBlockResolver = {
- // don't need to support this for the streaming shuffle implementation
- // since block manager is not used
- throw new UnsupportedOperationException()
- }
-
override def stop(): Unit = {}
}
diff --git a/core/src/main/scala/org/apache/spark/storage/BlockManager.scala
b/core/src/main/scala/org/apache/spark/storage/BlockManager.scala
index 21f932ba057a..38c391efad97 100644
--- a/core/src/main/scala/org/apache/spark/storage/BlockManager.scala
+++ b/core/src/main/scala/org/apache/spark/storage/BlockManager.scala
@@ -57,7 +57,7 @@ import org.apache.spark.network.util.TransportConf
import org.apache.spark.rpc.RpcEnv
import org.apache.spark.scheduler.ExecutorCacheTaskLocation
import org.apache.spark.serializer.{SerializerInstance, SerializerManager}
-import org.apache.spark.shuffle.{IndexShuffleBlockResolver,
MigratableResolver, ShuffleManager, ShuffleWriteMetricsReporter}
+import org.apache.spark.shuffle.{BlockingShuffleManager,
IndexShuffleBlockResolver, MigratableResolver, ShuffleBlockResolver,
ShuffleManager, ShuffleWriteMetricsReporter}
import
org.apache.spark.storage.BlockManagerMessages.{DecommissionBlockManager,
ReplicateBlock}
import org.apache.spark.storage.LogBlockType.LogBlockType
import org.apache.spark.storage.memory._
@@ -196,13 +196,25 @@ private[spark] class BlockManager(
// We initialize the ShuffleManager later in SparkContext and Executor, to
allow
// user jars to define custom ShuffleManagers, as such `_shuffleManager`
will be null here
- // (except for tests) and we ask for the instance from the SparkEnv.
- private lazy val shuffleManager = {
- Option(_shuffleManager).getOrElse {
- // Wait for ShuffleManager to be initialized before handling shuffle
operations.
- // Exception will be thrown if it is not initialized within the
configured timeout.
- waitForShuffleManagerInit()
- SparkEnv.get.shuffleManager
+ // (except for tests) and we ask for the resolver from the SparkEnv.
+ // The resolver used for block-by-id resolution (reads, push-merge,
decommission migration). Only
+ // a BlockingShuffleManager serves block-manager-addressed blocks; a
pipelined shuffle is served
+ // out-of-band and produces none, so these paths only ever resolve regular
shuffles. Prefer this
+ // over reaching through a ShuffleManager at the call site. Cached as a lazy
val so the
+ // init-wait and type match happen once, not on every (hot-path) block
access.
+ private lazy val shuffleBlockResolver: ShuffleBlockResolver = {
+ val resolver = Option(_shuffleManager) match {
+ case Some(blocking: BlockingShuffleManager) =>
Some(blocking.shuffleBlockResolver)
+ case Some(_) => None
+ case None =>
+ // Wait for ShuffleManager to be initialized before handling shuffle
operations.
+ // Exception will be thrown if it is not initialized within the
configured timeout.
+ waitForShuffleManagerInit()
+ SparkEnv.get.shuffleBlockResolver
+ }
+ resolver.getOrElse {
+ throw new UnsupportedOperationException(
+ "The configured shuffle manager does not provide a
ShuffleBlockResolver")
}
}
@@ -396,7 +408,7 @@ private[spark] class BlockManager(
// This is a lazy val so someone can migrating RDDs even if they don't have
a MigratableResolver
// for shuffles. Used in BlockManagerDecommissioner & block puts.
lazy val migratableResolver: MigratableResolver = {
- shuffleManager.shuffleBlockResolver.asInstanceOf[MigratableResolver]
+ shuffleBlockResolver.asInstanceOf[MigratableResolver]
}
// Timeout waiting for ShuffleManager initialization when receiving shuffle
migration requests
@@ -441,7 +453,7 @@ private[spark] class BlockManager(
assert(blockId.isInstanceOf[ShuffleBlockId],
s"Corruption diagnosis only supports shuffle block yet, but got
$blockId")
val shuffleBlock = blockId.asInstanceOf[ShuffleBlockId]
- val resolver =
shuffleManager.shuffleBlockResolver.asInstanceOf[IndexShuffleBlockResolver]
+ val resolver = shuffleBlockResolver.asInstanceOf[IndexShuffleBlockResolver]
val checksumFile =
resolver.getChecksumFile(shuffleBlock.shuffleId, shuffleBlock.mapId,
algorithm)
val reduceId = shuffleBlock.reduceId
@@ -871,7 +883,7 @@ private[spark] class BlockManager(
override def getHostLocalShuffleData(
blockId: BlockId,
dirs: Array[String]): ManagedBuffer = {
- shuffleManager.shuffleBlockResolver.getBlockData(blockId, Some(dirs))
+ shuffleBlockResolver.getBlockData(blockId, Some(dirs))
}
/**
@@ -882,7 +894,7 @@ private[spark] class BlockManager(
if (blockId.isShuffle) {
logDebug(s"Getting local shuffle block ${blockId}")
try {
- shuffleManager.shuffleBlockResolver.getBlockData(blockId)
+ shuffleBlockResolver.getBlockData(blockId)
} catch {
case e: IOException =>
if
(conf.get(config.STORAGE_DECOMMISSION_FALLBACK_STORAGE_PATH).isDefined) {
@@ -937,8 +949,8 @@ private[spark] class BlockManager(
return migratableResolver.putShuffleBlockAsStream(blockId,
serializerManager)
} catch {
case _: ClassCastException =>
- throw
SparkCoreErrors.unexpectedShuffleBlockWithUnsupportedResolverError(shuffleManager,
- blockId)
+ throw
SparkCoreErrors.unexpectedShuffleBlockWithUnsupportedResolverError(
+ shuffleBlockResolver, blockId)
}
}
logDebug(s"Putting regular block ${blockId}")
@@ -989,7 +1001,7 @@ private[spark] class BlockManager(
def getLocalMergedBlockData(
blockId: ShuffleMergedBlockId,
dirs: Array[String]): Seq[ManagedBuffer] = {
- shuffleManager.shuffleBlockResolver.getMergedBlockData(blockId, Some(dirs))
+ shuffleBlockResolver.getMergedBlockData(blockId, Some(dirs))
}
/**
@@ -998,7 +1010,7 @@ private[spark] class BlockManager(
def getLocalMergedBlockMeta(
blockId: ShuffleMergedBlockId,
dirs: Array[String]): MergedBlockMeta = {
- shuffleManager.shuffleBlockResolver.getMergedBlockMeta(blockId, Some(dirs))
+ shuffleBlockResolver.getMergedBlockMeta(blockId, Some(dirs))
}
/**
diff --git
a/core/src/main/scala/org/apache/spark/storage/BlockManagerMasterEndpoint.scala
b/core/src/main/scala/org/apache/spark/storage/BlockManagerMasterEndpoint.scala
index 58e1e899b11e..9c59d5ea49fb 100644
---
a/core/src/main/scala/org/apache/spark/storage/BlockManagerMasterEndpoint.scala
+++
b/core/src/main/scala/org/apache/spark/storage/BlockManagerMasterEndpoint.scala
@@ -38,7 +38,6 @@ import
org.apache.spark.network.shuffle.{ExternalBlockStoreClient, RemoteBlockPu
import org.apache.spark.rpc.{IsolatedThreadSafeRpcEndpoint, RpcCallContext,
RpcEndpointRef, RpcEnv}
import org.apache.spark.scheduler._
import org.apache.spark.scheduler.cluster.{CoarseGrainedClusterMessages,
CoarseGrainedSchedulerBackend}
-import org.apache.spark.shuffle.ShuffleManager
import org.apache.spark.storage.BlockManagerMessages._
import org.apache.spark.util.{RpcUtils, ThreadUtils, Utils}
import org.apache.spark.util.ArrayImplicits._
@@ -56,14 +55,9 @@ class BlockManagerMasterEndpoint(
externalBlockStoreClient: Option[ExternalBlockStoreClient],
blockManagerInfo: mutable.Map[BlockManagerId, BlockManagerInfo],
mapOutputTracker: MapOutputTrackerMaster,
- private val _shuffleManager: ShuffleManager,
isDriver: Boolean)
extends IsolatedThreadSafeRpcEndpoint with Logging {
- // We initialize the ShuffleManager later in SparkContext and Executor, to
allow
- // user jars to define custom ShuffleManagers, as such `_shuffleManager`
will be null here
- // (except for tests) and we ask for the instance from the SparkEnv.
- private lazy val shuffleManager =
Option(_shuffleManager).getOrElse(SparkEnv.get.shuffleManager)
// Mapping from executor id to the block manager's local disk directories.
private val executorIdToLocalDirs =
@@ -441,8 +435,12 @@ class BlockManagerMasterEndpoint(
mapStatuses.filter(_ != null).foreach { mapStatus =>
// Check if the executor has been deallocated
if
(!blockManagerIdByExecutor.contains(mapStatus.location.executorId)) {
- val blocksToDel =
-
shuffleManager.shuffleBlockResolver.getBlocksForShuffle(shuffleId,
mapStatus.mapId)
+ // Only a BlockingShuffleManager serves block-manager blocks; a
pipelined shuffle is
+ // served out-of-band and is not in the MapOutputTracker this
loop iterates, so this
+ // resolves only regular shuffles (None only before the manager
is initialized).
+ val blocksToDel = SparkEnv.get.shuffleBlockResolver
+ .map(_.getBlocksForShuffle(shuffleId, mapStatus.mapId))
+ .getOrElse(Seq.empty)
if (blocksToDel.nonEmpty) {
val blocks =
blocksToDeleteByShuffleService.getOrElseUpdate(mapStatus.location,
new mutable.HashSet[BlockId])
diff --git
a/core/src/main/scala/org/apache/spark/storage/BlockManagerStorageEndpoint.scala
b/core/src/main/scala/org/apache/spark/storage/BlockManagerStorageEndpoint.scala
index 54329c5b1e51..c4ee264ccf79 100644
---
a/core/src/main/scala/org/apache/spark/storage/BlockManagerStorageEndpoint.scala
+++
b/core/src/main/scala/org/apache/spark/storage/BlockManagerStorageEndpoint.scala
@@ -60,9 +60,14 @@ class BlockManagerStorageEndpoint(
if (mapOutputTracker != null) {
mapOutputTracker.unregisterShuffle(shuffleId)
}
- val shuffleManager = SparkEnv.get.shuffleManager
- if (shuffleManager != null) {
- shuffleManager.unregisterShuffle(shuffleId)
+ val env = SparkEnv.get
+ // The shuffle manager init is deferred (gated by a latch) so user
jars can load first; a
+ // RemoveShuffle can arrive before it runs. As before this routing,
skip in that case.
+ if (env != null && env.isShuffleManagerInitialized) {
+ // Only a shuffleId is available here and it does not identify the
owning manager, so
+ // notify every configured manager (default and, if present,
incremental). This reaches
+ // the owning manager on any node, including one that never
performed this shuffle's I/O.
+ env.unregisterShuffleFromAllManagers(shuffleId)
} else {
logDebug(log"Ignore remove shuffle ${MDC(SHUFFLE_ID, shuffleId)}")
true
diff --git a/core/src/test/scala/org/apache/spark/ShuffleSuite.scala
b/core/src/test/scala/org/apache/spark/ShuffleSuite.scala
index 18e909fce51c..ac64bcf2b715 100644
--- a/core/src/test/scala/org/apache/spark/ShuffleSuite.scala
+++ b/core/src/test/scala/org/apache/spark/ShuffleSuite.scala
@@ -357,7 +357,7 @@ abstract class ShuffleSuite extends SparkFunSuite with
Matchers with LocalRootDi
test("multiple simultaneous attempts for one task (SPARK-8029)") {
sc = new SparkContext("local", "test", conf)
val mapTrackerMaster =
sc.env.mapOutputTracker.asInstanceOf[MapOutputTrackerMaster]
- val manager = sc.env.shuffleManager
+ val manager = sc.env.blockingShuffleManager
val taskMemoryManager = new TaskMemoryManager(sc.env.memoryManager, 0L)
val metricsSystem = sc.env.metricsSystem
diff --git a/core/src/test/scala/org/apache/spark/deploy/SparkSubmitSuite.scala
b/core/src/test/scala/org/apache/spark/deploy/SparkSubmitSuite.scala
index 761e34571cf4..2627633126cd 100644
--- a/core/src/test/scala/org/apache/spark/deploy/SparkSubmitSuite.scala
+++ b/core/src/test/scala/org/apache/spark/deploy/SparkSubmitSuite.scala
@@ -1569,7 +1569,7 @@ class SparkSubmitSuite
"",
null,
Seq.empty,
- Seq("org.apache.spark.shuffle.ShuffleManager"),
+ Seq("org.apache.spark.shuffle.BlockingShuffleManager"),
shuffleManagerBody)
val jarUrl = TestUtils.createJar(
diff --git
a/core/src/test/scala/org/apache/spark/internal/plugin/PluginContainerSuite.scala
b/core/src/test/scala/org/apache/spark/internal/plugin/PluginContainerSuite.scala
index b1c5683cb225..d87e5c024542 100644
---
a/core/src/test/scala/org/apache/spark/internal/plugin/PluginContainerSuite.scala
+++
b/core/src/test/scala/org/apache/spark/internal/plugin/PluginContainerSuite.scala
@@ -318,7 +318,7 @@ class PluginContainerSuite extends SparkFunSuite with
LocalSparkContext {
// Ensures the shuffle manager specified in configuration was
// overridden by the Spark plugin.
-
assert(sc.env.shuffleManager.isInstanceOf[SetShuffleManagerPlugin.MyShuffleManager])
+
assert(sc.env.blockingShuffleManager.isInstanceOf[SetShuffleManagerPlugin.MyShuffleManager])
}
}
diff --git
a/core/src/test/scala/org/apache/spark/scheduler/AQEShuffledRDD.scala
b/core/src/test/scala/org/apache/spark/scheduler/AQEShuffledRDD.scala
index 84f9ef0d557e..3a5c0ff1d2c6 100644
--- a/core/src/test/scala/org/apache/spark/scheduler/AQEShuffledRDD.scala
+++ b/core/src/test/scala/org/apache/spark/scheduler/AQEShuffledRDD.scala
@@ -105,7 +105,7 @@ class AQEShuffledRDD[K, V, C](
override def compute(p: Partition, context: TaskContext): Iterator[(K, C)] =
{
val part = p.asInstanceOf[AQEShuffledRDDPartition]
val metrics = context.taskMetrics().createTempShuffleReadMetrics()
- SparkEnv.get.shuffleManager.getReader(
+ SparkEnv.get.shuffleManagerFor(dependency).getReader(
dependency.shuffleHandle, part.startIndexInParent,
part.endIndexInParent, context, metrics)
.read()
.asInstanceOf[Iterator[(K, C)]]
diff --git
a/core/src/test/scala/org/apache/spark/shuffle/PipelinedShuffleRoutingSuite.scala
b/core/src/test/scala/org/apache/spark/shuffle/PipelinedShuffleRoutingSuite.scala
new file mode 100644
index 000000000000..da74803a10eb
--- /dev/null
+++
b/core/src/test/scala/org/apache/spark/shuffle/PipelinedShuffleRoutingSuite.scala
@@ -0,0 +1,317 @@
+/*
+ * 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
+
+import scala.collection.mutable
+
+import org.mockito.Mockito.mock
+
+import org.apache.spark._
+import org.apache.spark.internal.config.{SHUFFLE_MANAGER,
SHUFFLE_MANAGER_INCREMENTAL}
+import org.apache.spark.rdd.RDD
+import org.apache.spark.shuffle.sort.SortShuffleManager
+import org.apache.spark.shuffle.streaming.StreamingShuffleManager
+
+/**
+ * A recording ShuffleManager test double: it records the shuffle ids it was
asked to register so a
+ * test can assert that [[SparkEnv.shuffleManagerFor]] routed a dependency to
the correct manager.
+ * Every instance registers itself in a process-wide registry keyed by class,
because the two
+ * managers live in private SparkEnv fields and a test otherwise has no
reference to them.
+ * Constructor takes (SparkConf, Boolean) to match how SparkEnv instantiates
managers. This base
+ * declares no shuffle kind; each concrete subclass mixes in
[[BlockingShuffleManager]] or
+ * [[PipelinedShuffleManager]] so it is accepted in the corresponding SparkEnv
slot.
+ */
+private abstract class RecordingShuffleManager(conf: SparkConf, isDriver:
Boolean)
+ extends ShuffleManager {
+ RecordingShuffleManager.register(this)
+ val registered = mutable.ArrayBuffer[Int]()
+ val unregistered = mutable.ArrayBuffer[Int]()
+ @volatile var stopped = false
+
+ override def registerShuffle[K, V, C](
+ shuffleId: Int, dependency: ShuffleDependency[K, V, C]): ShuffleHandle =
{
+ registered += shuffleId
+ new RecordingShuffleManager.RecordingHandle(shuffleId)
+ }
+ override def getWriter[K, V](
+ handle: ShuffleHandle, mapId: Long, context: TaskContext,
+ metrics: ShuffleWriteMetricsReporter): ShuffleWriter[K, V] = null
+ override def getReader[K, C](
+ handle: ShuffleHandle, startMapIndex: Int, endMapIndex: Int,
startPartition: Int,
+ endPartition: Int, context: TaskContext, metrics:
ShuffleReadMetricsReporter)
+ : ShuffleReader[K, C] = null
+ override def unregisterShuffle(shuffleId: Int): Boolean = {
+ unregistered += shuffleId
+ true
+ }
+ override def stop(): Unit = stopped = true
+}
+
+private object RecordingShuffleManager {
+ class RecordingHandle(shuffleId: Int) extends ShuffleHandle(shuffleId)
+
+ // Registry so tests can reach the instances SparkEnv built. Cleared per
test.
+ private val instances = mutable.ArrayBuffer[RecordingShuffleManager]()
+ def register(m: RecordingShuffleManager): Unit = synchronized { instances +=
m }
+ def clear(): Unit = synchronized { instances.clear() }
+ def of[T <: RecordingShuffleManager](cls: Class[T]): T = synchronized {
+ instances.find(cls.isInstance).getOrElse(
+ throw new NoSuchElementException(s"no ${cls.getSimpleName} was
instantiated")).asInstanceOf[T]
+ }
+}
+
+// Distinct subclasses so SparkEnv instantiates two different classes
(blocking vs pipelined) and
+// tests can tell which one a dependency was routed to. Each declares its kind
so it is accepted in
+// the matching SparkEnv slot: the blocking manager (spark.shuffle.manager)
must be a
+// BlockingShuffleManager, the pipelined manager
(spark.shuffle.manager.incremental) a
+// PipelinedShuffleManager. The routing tests never exercise the resolver, so
it returns a mock.
+private class DefaultRecordingManager(conf: SparkConf, isDriver: Boolean)
+ extends RecordingShuffleManager(conf, isDriver) with BlockingShuffleManager {
+ override def shuffleBlockResolver: ShuffleBlockResolver =
mock(classOf[ShuffleBlockResolver])
+}
+private class IncrementalRecordingManager(conf: SparkConf, isDriver: Boolean)
+ extends RecordingShuffleManager(conf, isDriver) with PipelinedShuffleManager
+
+/**
+ * Tests routing of shuffles to the default vs. incremental [[ShuffleManager]]
by dependency type,
+ * done natively in [[SparkEnv.shuffleManagerFor]] with no wrapping "router"
manager. A
+ * [[PipelinedShuffleDependency]] routes to the manager named by
spark.shuffle.manager.incremental;
+ * every other dependency routes to spark.shuffle.manager.
`SparkEnv.blockingShuffleManager` remains
+ * the plain configured manager (nothing is installed "in front" of it). The
id-only cleanup path
+ * (`unregisterShuffleFromAllManagers`) notifies both managers, since a
shuffleId alone does not
+ * identify the owner.
+ */
+class PipelinedShuffleRoutingSuite extends SparkFunSuite with
LocalSparkContext {
+
+ override def beforeEach(): Unit = {
+ super.beforeEach()
+ RecordingShuffleManager.clear()
+ }
+
+ private def newConf(): SparkConf = new SparkConf(loadDefaults = false)
+ .set(SHUFFLE_MANAGER, classOf[DefaultRecordingManager].getName)
+ .set(SHUFFLE_MANAGER_INCREMENTAL,
classOf[IncrementalRecordingManager].getName)
+
+ /** Start a SparkContext with the given conf and return the live SparkEnv. */
+ private def startEnv(conf: SparkConf = newConf()): SparkEnv = {
+ sc = new SparkContext("local", "test", conf)
+ SparkEnv.get
+ }
+
+ private def defaultMgr =
RecordingShuffleManager.of(classOf[DefaultRecordingManager])
+ private def incrementalMgr =
RecordingShuffleManager.of(classOf[IncrementalRecordingManager])
+
+ private def pipelinedDep(sc: SparkContext): PipelinedShuffleDependency[Int,
Int, Int] = {
+ val rdd: RDD[(Int, Int)] = sc.parallelize(1 to 4, 2).map(x => (x, x))
+ new PipelinedShuffleDependency[Int, Int, Int](rdd, new HashPartitioner(2))
+ }
+ private def regularDep(sc: SparkContext): ShuffleDependency[Int, Int, Int] =
{
+ val rdd: RDD[(Int, Int)] = sc.parallelize(1 to 4, 2).map(x => (x, x))
+ new ShuffleDependency[Int, Int, Int](rdd, new HashPartitioner(2))
+ }
+
+ /** Walk an exception's cause chain to find the first cause of the given
type, if any. */
+ private def findCause[T <: Throwable : reflect.ClassTag](t: Throwable):
Option[T] = {
+ val cls = implicitly[reflect.ClassTag[T]].runtimeClass
+ Iterator.iterate(t)(_.getCause).takeWhile(_ != null)
+ .find(cls.isInstance).map(_.asInstanceOf[T])
+ }
+
+ test("blockingShuffleManager is the plain configured manager (no wrapper is
installed)") {
+ val env = startEnv()
+ // The blocking manager is exactly the configured one -- not a
router/composite in front of it
+ // -- so callers that inspect its type (e.g. ShuffleExchangeExec) see the
real manager.
+ assert(env.blockingShuffleManager.isInstanceOf[DefaultRecordingManager])
+ }
+
+ test("shuffleManagerFor routes a regular dependency to the default manager")
{
+ val env = startEnv()
+ val dep = regularDep(sc)
+ assert(env.shuffleManagerFor(dep).isInstanceOf[DefaultRecordingManager])
+ // registerShuffle (in the dependency's constructor) went to the default
manager, not the
+ // incremental one.
+ assert(defaultMgr.registered.contains(dep.shuffleId))
+ assert(!incrementalMgr.registered.contains(dep.shuffleId))
+ }
+
+ test("shuffleManagerFor routes a pipelined dependency to the incremental
manager") {
+ val env = startEnv()
+ val dep = pipelinedDep(sc)
+
assert(env.shuffleManagerFor(dep).isInstanceOf[IncrementalRecordingManager])
+ // registerShuffle went to the incremental manager, not the default
+ assert(incrementalMgr.registered.contains(dep.shuffleId))
+ assert(!defaultMgr.registered.contains(dep.shuffleId))
+ }
+
+ test("unregisterShuffleFromAllManagers notifies both managers (owner unknown
from id alone)") {
+ val env = startEnv()
+ // The RemoveShuffle path holds only a shuffleId and cannot tell which
manager owns it, and must
+ // reach the owner even on a node that never did this shuffle's I/O -- so
it notifies both. A
+ // fresh id that was never registered exercises exactly that "don't know
the owner" case.
+ assert(env.unregisterShuffleFromAllManagers(4242))
+ assert(defaultMgr.unregistered.contains(4242), "the default manager must
be unregistered")
+ assert(incrementalMgr.unregistered.contains(4242),
+ "the incremental manager must be unregistered")
+ }
+
+ test("unregisterShuffleFromAllManagers notifies the default when no
incremental is configured") {
+ val without = new SparkConf(loadDefaults = false)
+ .set(SHUFFLE_MANAGER, classOf[DefaultRecordingManager].getName)
+ val env = startEnv(without)
+ assert(env.unregisterShuffleFromAllManagers(4242))
+ assert(defaultMgr.unregistered.contains(4242))
+ }
+
+ test("unregisterShuffleFromAllManagers is a no-op (not an NPE) before
managers are initialized") {
+ val env = startEnv()
+ // SPARK-45762 defers shuffle-manager init behind a latch (so user jars
load first), leaving
+ // both managers null until initializeShuffleManager runs. A RemoveShuffle
RPC can arrive in
+ // that window; it must be skipped, not dereference a null manager.
Reflectively reproduce the
+ // pre-init state (both slots null), since a live SparkContext finishes
init during startup.
+ val blockingField =
classOf[SparkEnv].getDeclaredField("_blockingShuffleManager")
+ val pipelinedField =
classOf[SparkEnv].getDeclaredField("_pipelinedShuffleManager")
+ blockingField.setAccessible(true)
+ pipelinedField.setAccessible(true)
+ val origBlocking = blockingField.get(env)
+ val origPipelined = pipelinedField.get(env)
+ blockingField.set(env, null)
+ pipelinedField.set(env, null)
+ try {
+ assert(!env.isShuffleManagerInitialized, "precondition: the manager must
look uninitialized")
+ // With both managers null there is nothing to notify: the call must
return false without
+ // throwing. Before the null guards this threw NullPointerException.
+ assert(!env.unregisterShuffleFromAllManagers(4242))
+ } finally {
+ blockingField.set(env, origBlocking)
+ pipelinedField.set(env, origPipelined)
+ }
+ }
+
+ test("the handle is minted by the routed manager (unwrapped) so driver and
executor agree") {
+ startEnv()
+ // The dependency's shuffleHandle is produced by the incremental manager's
registerShuffle and
+ // is that manager's own handle type -- there is no wrapping handle to
route on. Executors that
+ // later read dep.shuffleHandle route by the same dependency type, so they
hit the same manager.
+ val dep = pipelinedDep(sc)
+
assert(dep.shuffleHandle.isInstanceOf[RecordingShuffleManager.RecordingHandle])
+ assert(dep.shuffleHandle.shuffleId === dep.shuffleId)
+ }
+
+ test("both managers are stopped when the SparkContext stops") {
+ startEnv()
+ // Touch a pipelined dep so the incremental manager instance exists in the
registry.
+ pipelinedDep(sc)
+ val default = defaultMgr
+ val incremental = incrementalMgr
+ // Stopping the context tears down its SparkEnv, which must stop BOTH
managers (default and
+ // incremental). Clear `sc` afterward so LocalSparkContext's afterEach
does not double-stop.
+ sc.stop()
+ sc = null
+ assert(default.stopped, "the default manager must be stopped")
+ assert(incremental.stopped, "the incremental manager must be stopped")
+ }
+
+ test("SparkEnv initializes the streaming shuffle tracker when the
incremental manager is " +
+ "StreamingShuffleManager") {
+ val conf = new SparkConf(loadDefaults = false)
+ .set(SHUFFLE_MANAGER, classOf[DefaultRecordingManager].getName)
+ .set(SHUFFLE_MANAGER_INCREMENTAL,
classOf[StreamingShuffleManager].getName)
+ sc = new SparkContext("local", "test", conf)
+ assert(SparkEnv.get.streamingShuffleOutputTracker.isDefined)
+ }
+
+ test("SparkEnv does not initialize the tracker when the incremental manager
is not streaming") {
+ sc = new SparkContext("local", "test", newConf())
+ assert(SparkEnv.get.streamingShuffleOutputTracker.isEmpty)
+ }
+
+ test("spark.shuffle.manager.incremental resolves the same short aliases as
the default manager") {
+ // "sort" is a short alias the incremental slot must resolve (to
SortShuffleManager) rather than
+ // treat as a class name. SortShuffleManager is blocking, so it is
rejected from the pipelined
+ // slot -- proving the alias resolved (a ClassNotFoundException would mean
it did not). See the
+ // streaming-tracker test for a valid pipelined manager in this slot.
+ val conf = new SparkConf(loadDefaults = false)
+ .set(SHUFFLE_MANAGER, classOf[DefaultRecordingManager].getName)
+ .set(SHUFFLE_MANAGER_INCREMENTAL, "sort")
+ val e = intercept[Exception] {
+ sc = new SparkContext("local", "test", conf)
+ }
+ assert(findCause[IllegalArgumentException](e).exists(
+ _.getMessage.contains("must be a PipelinedShuffleManager")),
+ s"expected a PipelinedShuffleManager rejection, got: $e")
+ }
+
+ test("a blocking manager in the incremental slot is rejected at startup") {
+ val conf = new SparkConf(loadDefaults = false)
+ .set(SHUFFLE_MANAGER, classOf[DefaultRecordingManager].getName)
+ // DefaultRecordingManager is a BlockingShuffleManager; the pipelined
slot must reject it, so
+ // a shuffle routed there can never write blocks unreachable through the
block resolver.
+ .set(SHUFFLE_MANAGER_INCREMENTAL,
classOf[DefaultRecordingManager].getName)
+ val e = intercept[Exception] {
+ sc = new SparkContext("local", "test", conf)
+ }
+ assert(findCause[IllegalArgumentException](e).exists(
+ _.getMessage.contains("must be a PipelinedShuffleManager")),
+ s"expected a PipelinedShuffleManager rejection, got: $e")
+ }
+
+ test("a bad incremental manager class name surfaces a clear error at
startup") {
+ val badConf = new SparkConf(loadDefaults = false)
+ .set(SHUFFLE_MANAGER, "sort")
+ .set(SHUFFLE_MANAGER_INCREMENTAL, "org.apache.spark.NotAShuffleManager")
+ // The incremental manager is instantiated eagerly during SparkEnv init,
so a bad class name
+ // fails context startup rather than silently.
+ intercept[Exception] {
+ sc = new SparkContext("local", "test", badConf)
+ }
+ }
+
+ test("a manager's type declares its shuffle kind (blocking vs pipelined)") {
+ // SortShuffleManager materializes blocks -- it is a
BlockingShuffleManager with a resolver.
+ val sortMgr = new SortShuffleManager(new SparkConf(loadDefaults = false))
+ assert(sortMgr.isInstanceOf[BlockingShuffleManager])
+ // StreamingShuffleManager serves output out-of-band -- it is a
PipelinedShuffleManager with no
+ // resolver (the shuffleBlockResolver method was removed from the
ShuffleManager interface).
+ val streaming = new StreamingShuffleManager
+ assert(streaming.isInstanceOf[PipelinedShuffleManager])
+ assert(!streaming.isInstanceOf[BlockingShuffleManager])
+ }
+
+ test("SparkEnv.shuffleBlockResolver is defined for the blocking manager") {
+ // The blocking manager is SortShuffleManager -> a resolver is available.
It is the only manager
+ // that produces block-manager blocks, so block resolution always comes
from it.
+ val blockingConf = new SparkConf(loadDefaults =
false).set(SHUFFLE_MANAGER, "sort")
+ sc = new SparkContext("local", "test", blockingConf)
+ assert(SparkEnv.get.shuffleBlockResolver.isDefined,
+ "the blocking manager must expose a shuffleBlockResolver")
+ }
+
+ test("a pipelined manager in the default slot is rejected at startup") {
+ // StreamingShuffleManager is pipelined; it cannot be the blocking
manager, because block
+ // resolution (reads, push-merge, decommission migration) has no source
otherwise. Rejecting it
+ // at startup means shuffleBlockResolver always has exactly one well-typed
source.
+ val pipelinedConf = new SparkConf(loadDefaults = false)
+ .set(SHUFFLE_MANAGER, classOf[StreamingShuffleManager].getName)
+ val e = intercept[Exception] {
+ sc = new SparkContext("local", "test", pipelinedConf)
+ }
+ assert(findCause[IllegalArgumentException](e).exists(
+ _.getMessage.contains("must be a BlockingShuffleManager")),
+ s"expected a BlockingShuffleManager rejection, got: $e")
+ }
+}
diff --git
a/core/src/test/scala/org/apache/spark/shuffle/ShuffleDependencySuite.scala
b/core/src/test/scala/org/apache/spark/shuffle/ShuffleDependencySuite.scala
index 4d5f599fb12a..570dd96e3737 100644
--- a/core/src/test/scala/org/apache/spark/shuffle/ShuffleDependencySuite.scala
+++ b/core/src/test/scala/org/apache/spark/shuffle/ShuffleDependencySuite.scala
@@ -17,6 +17,9 @@
package org.apache.spark.shuffle
import org.apache.spark._
+import org.apache.spark.internal.config
+import org.apache.spark.internal.config.Tests
+import org.apache.spark.rdd.RDD
case class KeyClass()
@@ -64,4 +67,87 @@ class ShuffleDependencySuite extends SparkFunSuite with
LocalSparkContext {
assert(dep.combinerClassName == None)
}
+ test("PipelinedShuffleDependency is a ShuffleDependency and preserves its
fields") {
+ sc = new SparkContext("local", "test", conf.clone())
+ val rdd: RDD[(KeyClass, ValueClass)] =
+ sc.parallelize(1 to 5, 4).map(_ => (KeyClass(), ValueClass()))
+ val partitioner = new HashPartitioner(2)
+ val dep = new PipelinedShuffleDependency[KeyClass, ValueClass,
ValueClass](rdd, partitioner)
+
+ // It is a first-class dependency kind, but IS-A ShuffleDependency: code
that matches
+ // ShuffleDependency continues to see it as an ordinary shuffle, so it
changes no existing
+ // behavior. The concurrent-scheduling behavior is keyed on the subtype
elsewhere.
+ assert(dep.isInstanceOf[ShuffleDependency[_, _, _]])
+ assert(dep.partitioner === partitioner)
+ assert(dep.keyClassName == classOf[KeyClass].getName)
+ assert(dep.valueClassName == classOf[ValueClass].getName)
+ assert(dep.rdd === rdd)
+ // Construction goes through the normal ShuffleDependency path: the
shuffle is registered with
+ // the ShuffleManager (a handle is produced) and a second instance gets a
distinct shuffleId.
+ assert(dep.shuffleHandle != null)
+ val dep2 = new PipelinedShuffleDependency[KeyClass, ValueClass,
ValueClass](rdd, partitioner)
+ assert(dep2.shuffleId != dep.shuffleId)
+
+ // The checksum retry / query-level rollback params are intentionally not
exposed by the
+ // subclass (see PipelinedShuffleDependency scaladoc): their stage-level
recompute is moot for a
+ // pipelined group, so they must stay at their false defaults.
+ assert(!dep.checksumMismatchFullRetryEnabled)
+ assert(!dep.checksumMismatchQueryLevelRollbackEnabled)
+ }
+
+ test("PipelinedShuffleDependency forwards non-default constructor args to
ShuffleDependency") {
+ sc = new SparkContext("local", "test", conf.clone())
+ val rdd: RDD[(KeyClass, ValueClass)] =
+ sc.parallelize(1 to 5, 4).map(_ => (KeyClass(), ValueClass()))
+ val partitioner = new HashPartitioner(2)
+ val aggregator = new Aggregator[KeyClass, ValueClass, ValueClass](
+ v => v, (c, _) => c, (c1, _) => c1)
+ val dep = new PipelinedShuffleDependency[KeyClass, ValueClass, ValueClass](
+ rdd, partitioner, aggregator = Some(aggregator), mapSideCombine = true)
+
+ // The subclass forwards its constructor args positionally to
ShuffleDependency; assert the
+ // non-default ones land where expected rather than on the wrong parameter.
+ assert(dep.aggregator.contains(aggregator))
+ assert(dep.mapSideCombine)
+ }
+
+ test("an ordinary ShuffleDependency is not a PipelinedShuffleDependency") {
+ sc = new SparkContext("local", "test", conf.clone())
+ val rdd = sc.parallelize(1 to 5, 4).map(_ => (KeyClass(),
ValueClass())).groupByKey()
+ val dep = rdd.dependencies.head.asInstanceOf[ShuffleDependency[_, _, _]]
+ assert(!dep.isInstanceOf[PipelinedShuffleDependency[_, _, _]])
+ }
+
+ test("PipelinedShuffleDependency never allows push-based shuffle merge") {
+ // Push-based shuffle merge is incompatible with a pipelined
(incrementally-readable) shuffle
+ // (spec S9): it exposes output only after a post-completion finalize and
would register merge
+ // results for a transient shuffle. Even with push-based shuffle enabled
cluster-wide -- where
+ // an ordinary ShuffleDependency does allow merge -- a
PipelinedShuffleDependency must not, so a
+ // pipelined producer never acquires merger locations and its completion
never routes through
+ // the merge-finalize path in the DAGScheduler.
+ val pushConf = new SparkConf(loadDefaults = false)
+ .set(config.PUSH_BASED_SHUFFLE_ENABLED, true)
+ .set(Tests.IS_TESTING, true)
+ .set(config.IO_ENCRYPTION_ENABLED, false)
+ // Push-based shuffle requires a serializer that supports relocation of
serialized objects.
+ .set(config.SERIALIZER, "org.apache.spark.serializer.KryoSerializer")
+ sc = new SparkContext("local", "test", pushConf)
+ val rdd: RDD[(KeyClass, ValueClass)] =
+ sc.parallelize(1 to 5, 4).map(_ => (KeyClass(), ValueClass()))
+ val partitioner = new HashPartitioner(2)
+
+ // Control: an ordinary ShuffleDependency DOES allow merge under this
config, so the test is
+ // meaningful (the pipelined result is not vacuously false because
push-shuffle is off).
+ val regularDep = new ShuffleDependency[KeyClass, ValueClass,
ValueClass](rdd, partitioner)
+ assert(regularDep.shuffleMergeAllowed,
+ "precondition: push-based shuffle must be enabled so a regular dep
allows merge")
+
+ val pipelinedDep =
+ new PipelinedShuffleDependency[KeyClass, ValueClass, ValueClass](rdd,
partitioner)
+ assert(!pipelinedDep.shuffleMergeAllowed,
+ "a pipelined shuffle must never allow push-based shuffle merge")
+ assert(!pipelinedDep.shuffleMergeEnabled,
+ "a pipelined shuffle must never have merge enabled")
+ }
+
}
diff --git
a/core/src/test/scala/org/apache/spark/shuffle/streaming/StreamingShuffleManagerSuite.scala
b/core/src/test/scala/org/apache/spark/shuffle/streaming/StreamingShuffleManagerSuite.scala
index 181d779e8bb5..4ee83135e59a 100644
---
a/core/src/test/scala/org/apache/spark/shuffle/streaming/StreamingShuffleManagerSuite.scala
+++
b/core/src/test/scala/org/apache/spark/shuffle/streaming/StreamingShuffleManagerSuite.scala
@@ -24,9 +24,8 @@ import org.scalatestplus.mockito.MockitoSugar
import org.apache.spark._
import org.apache.spark.LocalSparkContext.withSpark
-import org.apache.spark.internal.config.SHUFFLE_MANAGER
+import org.apache.spark.internal.config.{SHUFFLE_MANAGER,
SHUFFLE_MANAGER_INCREMENTAL}
import org.apache.spark.network.shuffle.streaming.{DataMessage,
TerminationAckMessage, TerminationControlMessage}
-import org.apache.spark.shuffle.sort.SortShuffleManager
import org.apache.spark.shuffle.streaming.StreamingShuffleManager.{getQueryId,
getWriterId, QUERY_ID_PROPERTY_KEY}
class StreamingShuffleManagerSuite
@@ -97,10 +96,17 @@ class StreamingShuffleManagerSuite
// ---- SparkEnv tracker initialization gating ----
- private def assertTrackerInitialized(shuffleManager: Option[String],
expectPresent: Boolean):
- Unit = {
+ private def assertTrackerInitialized(
+ defaultManager: Option[String] = None,
+ incrementalManager: Option[String] = None,
+ expectPresent: Boolean): Unit = {
val conf = new SparkConf()
- shuffleManager.foreach(conf.set(SHUFFLE_MANAGER, _))
+ // The default slot must be a BlockingShuffleManager; the incremental slot
holds a
+ // PipelinedShuffleManager and defaults to the streaming manager. Each is
set only when the test
+ // provides it, so leaving the incremental slot unset exercises the
streaming default. The
+ // tracker is initialized when either slot requires it.
+ defaultManager.foreach(conf.set(SHUFFLE_MANAGER, _))
+ incrementalManager.foreach(conf.set(SHUFFLE_MANAGER_INCREMENTAL, _))
withSpark(new SparkContext("local", "StreamingShuffleManagerSuite", conf))
{ _ =>
val tracker = SparkEnv.get.streamingShuffleOutputTracker
assert(tracker.isDefined == expectPresent)
@@ -111,15 +117,21 @@ class StreamingShuffleManagerSuite
}
}
- test("SparkEnv initializes the streaming shuffle tracker for
StreamingShuffleManager") {
- assertTrackerInitialized(Some(classOf[StreamingShuffleManager].getName),
expectPresent = true)
+ test("SparkEnv initializes the streaming shuffle tracker by default
(streaming incremental " +
+ "manager)") {
+ // The incremental slot defaults to the streaming manager, so the tracker
is present with no
+ // explicit configuration.
+ assertTrackerInitialized(expectPresent = true)
}
- test("SparkEnv does not initialize the tracker for a non-streaming (sort)
manager") {
- assertTrackerInitialized(Some(classOf[SortShuffleManager].getName),
expectPresent = false)
+ test("SparkEnv initializes the streaming shuffle tracker for an explicit
incremental " +
+ "StreamingShuffleManager") {
+ assertTrackerInitialized(
+ incrementalManager = Some(classOf[StreamingShuffleManager].getName),
expectPresent = true)
}
- test("SparkEnv does not initialize the tracker for the default manager") {
- assertTrackerInitialized(None, expectPresent = false)
+ test("SparkEnv initializes the streaming shuffle tracker for a
MultiShuffleManager default") {
+ assertTrackerInitialized(
+ defaultManager = Some(classOf[MultiShuffleManager].getName),
expectPresent = true)
}
}
diff --git
a/core/src/test/scala/org/apache/spark/shuffle/streaming/StreamingShuffleReaderSuite.scala
b/core/src/test/scala/org/apache/spark/shuffle/streaming/StreamingShuffleReaderSuite.scala
index f035ffcdb055..b3b09156fe0c 100644
---
a/core/src/test/scala/org/apache/spark/shuffle/streaming/StreamingShuffleReaderSuite.scala
+++
b/core/src/test/scala/org/apache/spark/shuffle/streaming/StreamingShuffleReaderSuite.scala
@@ -26,7 +26,7 @@ import org.scalatestplus.mockito.MockitoSugar
import org.apache.spark._
import org.apache.spark.LocalSparkContext.withSpark
-import org.apache.spark.internal.config.SHUFFLE_MANAGER
+import org.apache.spark.internal.config.SHUFFLE_MANAGER_INCREMENTAL
import org.apache.spark.memory.{TaskMemoryManager, TestMemoryManager}
import org.apache.spark.metrics.MetricsSystem
import org.apache.spark.network.shuffle.streaming.{DataMessage,
StreamingShuffleMessage, TerminationControlMessage}
@@ -43,7 +43,10 @@ class StreamingShuffleReaderSuite
with MockitoSugar {
private def newConf(): SparkConf =
- new SparkConf().set(SHUFFLE_MANAGER,
classOf[StreamingShuffleManager].getName)
+ // StreamingShuffleManager is pipelined, so it belongs in the incremental
slot (the default
+ // spark.shuffle.manager must be a BlockingShuffleManager). This is what
initializes the
+ // streaming output tracker that the reader constructor asserts.
+ new SparkConf().set(SHUFFLE_MANAGER_INCREMENTAL,
classOf[StreamingShuffleManager].getName)
private def createTaskContext(conf: SparkConf, partitionId: Int):
TaskContextImpl = {
val properties = new Properties()
diff --git
a/core/src/test/scala/org/apache/spark/shuffle/streaming/StreamingShuffleSuite.scala
b/core/src/test/scala/org/apache/spark/shuffle/streaming/StreamingShuffleSuite.scala
index fe88478416d6..af8e074fb9f1 100644
---
a/core/src/test/scala/org/apache/spark/shuffle/streaming/StreamingShuffleSuite.scala
+++
b/core/src/test/scala/org/apache/spark/shuffle/streaming/StreamingShuffleSuite.scala
@@ -39,7 +39,7 @@ import org.scalatestplus.mockito.MockitoSugar
import org.apache.spark._
import org.apache.spark.LocalSparkContext.withSpark
import org.apache.spark.internal.Logging
-import org.apache.spark.internal.config.{SHUFFLE_MANAGER,
STREAMING_SHUFFLE_CHECKSUM_ENABLED, STREAMING_SHUFFLE_NETWORK_BUFFER_SIZE,
STREAMING_SHUFFLE_READER_MAX_MEMORY, STREAMING_SHUFFLE_WRITER_MAX_MEMORY}
+import org.apache.spark.internal.config.{SHUFFLE_MANAGER_INCREMENTAL,
STREAMING_SHUFFLE_CHECKSUM_ENABLED, STREAMING_SHUFFLE_NETWORK_BUFFER_SIZE,
STREAMING_SHUFFLE_READER_MAX_MEMORY, STREAMING_SHUFFLE_WRITER_MAX_MEMORY}
import org.apache.spark.memory.{TaskMemoryManager, TestMemoryManager}
import org.apache.spark.metrics.MetricsSystem
import org.apache.spark.network.client.{RpcResponseCallback, TransportClient,
TransportClientFactory}
@@ -85,7 +85,10 @@ class StreamingShuffleSuite
protected override def beforeEach(): Unit = {
super.beforeEach()
sparkConf = new SparkConf()
- .set(SHUFFLE_MANAGER, classOf[StreamingShuffleManager].getName)
+ // StreamingShuffleManager is pipelined, so it occupies the incremental
slot (the default
+ // spark.shuffle.manager must be a BlockingShuffleManager). This still
initializes the
+ // streaming output tracker these tests rely on.
+ .set(SHUFFLE_MANAGER_INCREMENTAL,
classOf[StreamingShuffleManager].getName)
.set(STREAMING_SHUFFLE_CHECKSUM_ENABLED, true)
}
@@ -127,7 +130,7 @@ class StreamingShuffleSuite
extends StreamingShuffleReader[K, C](handle, context, clientHandler,
errorNotifier)
private class ShuffleGroup[T: ClassTag](sc: SparkContext, mappers: Int,
reducers: Int) {
- assert(SparkEnv.get.shuffleManager.isInstanceOf[StreamingShuffleManager])
+
assert(SparkEnv.get.pipelinedShuffleManager.isInstanceOf[StreamingShuffleManager])
assert(SparkEnv.get.streamingShuffleOutputTracker.isDefined)
SparkEnv.get.streamingShuffleOutputTracker.get
.asInstanceOf[StreamingShuffleOutputTrackerMaster]
@@ -803,7 +806,7 @@ class StreamingShuffleSuite
}
withSpark(new SparkContext("local", "StreamingShuffleSuite", sparkConf)) {
sc =>
- assert(SparkEnv.get.shuffleManager.isInstanceOf[StreamingShuffleManager])
+
assert(SparkEnv.get.pipelinedShuffleManager.isInstanceOf[StreamingShuffleManager])
assert(SparkEnv.get.streamingShuffleOutputTracker.isDefined)
SparkEnv.get.streamingShuffleOutputTracker.get
.asInstanceOf[StreamingShuffleOutputTrackerMaster]
@@ -919,7 +922,7 @@ class StreamingShuffleSuite
test("credit control message send failure") {
withSpark(new SparkContext("local", "StreamingShuffleSuite", sparkConf)) {
sc =>
- assert(SparkEnv.get.shuffleManager.isInstanceOf[StreamingShuffleManager])
+
assert(SparkEnv.get.pipelinedShuffleManager.isInstanceOf[StreamingShuffleManager])
assert(SparkEnv.get.streamingShuffleOutputTracker.isDefined)
SparkEnv.get.streamingShuffleOutputTracker.get
.asInstanceOf[StreamingShuffleOutputTrackerMaster]
@@ -975,7 +978,7 @@ class StreamingShuffleSuite
test("term ack message message send failure") {
withSpark(new SparkContext("local", "StreamingShuffleSuite", sparkConf)) {
sc =>
- assert(SparkEnv.get.shuffleManager.isInstanceOf[StreamingShuffleManager])
+
assert(SparkEnv.get.pipelinedShuffleManager.isInstanceOf[StreamingShuffleManager])
assert(SparkEnv.get.streamingShuffleOutputTracker.isDefined)
SparkEnv.get.streamingShuffleOutputTracker.get
.asInstanceOf[StreamingShuffleOutputTrackerMaster]
@@ -1446,7 +1449,7 @@ class StreamingShuffleSuite
}
withSpark(new SparkContext("local", "StreamingShuffleSuite", sparkConf)) {
sc =>
- assert(SparkEnv.get.shuffleManager.isInstanceOf[StreamingShuffleManager])
+
assert(SparkEnv.get.pipelinedShuffleManager.isInstanceOf[StreamingShuffleManager])
assert(SparkEnv.get.streamingShuffleOutputTracker.isDefined)
SparkEnv.get.streamingShuffleOutputTracker.get
.asInstanceOf[StreamingShuffleOutputTrackerMaster]
@@ -1488,7 +1491,7 @@ class StreamingShuffleSuite
test("writer resources are cleaned up when server handler " +
"receive fails via errorNotifier") {
withSpark(new SparkContext("local", "StreamingShuffleSuite", sparkConf)) {
sc =>
- assert(SparkEnv.get.shuffleManager.isInstanceOf[StreamingShuffleManager])
+
assert(SparkEnv.get.pipelinedShuffleManager.isInstanceOf[StreamingShuffleManager])
assert(SparkEnv.get.streamingShuffleOutputTracker.isDefined)
SparkEnv.get.streamingShuffleOutputTracker.get
.asInstanceOf[StreamingShuffleOutputTrackerMaster]
diff --git
a/core/src/test/scala/org/apache/spark/shuffle/streaming/StreamingShuffleWriterSuite.scala
b/core/src/test/scala/org/apache/spark/shuffle/streaming/StreamingShuffleWriterSuite.scala
index 42ad68a5668a..602cdd0bfc48 100644
---
a/core/src/test/scala/org/apache/spark/shuffle/streaming/StreamingShuffleWriterSuite.scala
+++
b/core/src/test/scala/org/apache/spark/shuffle/streaming/StreamingShuffleWriterSuite.scala
@@ -30,7 +30,7 @@ import org.scalatestplus.mockito.MockitoSugar
import org.apache.spark._
import org.apache.spark.LocalSparkContext.withSpark
-import org.apache.spark.internal.config.{SHUFFLE_MANAGER,
STREAMING_SHUFFLE_CHECKSUM_ENABLED, STREAMING_SHUFFLE_NETWORK_BUFFER_SIZE}
+import org.apache.spark.internal.config.{SHUFFLE_MANAGER_INCREMENTAL,
STREAMING_SHUFFLE_CHECKSUM_ENABLED, STREAMING_SHUFFLE_NETWORK_BUFFER_SIZE}
import org.apache.spark.memory.{TaskMemoryManager, TestMemoryManager}
import org.apache.spark.metrics.MetricsSystem
import org.apache.spark.network.client.TransportClient
@@ -49,7 +49,10 @@ class StreamingShuffleWriterSuite
with MockitoSugar {
private def newConf(): SparkConf =
- new SparkConf().set(SHUFFLE_MANAGER,
classOf[StreamingShuffleManager].getName)
+ // StreamingShuffleManager is pipelined, so it belongs in the incremental
slot (the default
+ // spark.shuffle.manager must be a BlockingShuffleManager). This is what
initializes the
+ // streaming output tracker that the writer constructor asserts.
+ new SparkConf().set(SHUFFLE_MANAGER_INCREMENTAL,
classOf[StreamingShuffleManager].getName)
private def createTaskContext(conf: SparkConf, partitionId: Int):
TaskContextImpl = {
val properties = new Properties()
diff --git
a/core/src/test/scala/org/apache/spark/storage/BlockManagerReplicationSuite.scala
b/core/src/test/scala/org/apache/spark/storage/BlockManagerReplicationSuite.scala
index f5fca56e5ef7..07eb7b3da65c 100644
---
a/core/src/test/scala/org/apache/spark/storage/BlockManagerReplicationSuite.scala
+++
b/core/src/test/scala/org/apache/spark/storage/BlockManagerReplicationSuite.scala
@@ -106,7 +106,7 @@ trait BlockManagerReplicationBehavior extends SparkFunSuite
val blockManagerInfo = new mutable.HashMap[BlockManagerId,
BlockManagerInfo]()
master = new BlockManagerMaster(rpcEnv.setupEndpoint("blockmanager",
new BlockManagerMasterEndpoint(rpcEnv, true, conf,
- new LiveListenerBus(conf), None, blockManagerInfo, mapOutputTracker,
sc.env.shuffleManager,
+ new LiveListenerBus(conf), None, blockManagerInfo, mapOutputTracker,
isDriver = true)),
rpcEnv.setupEndpoint("blockmanagerHeartbeat",
new BlockManagerMasterHeartbeatEndpoint(rpcEnv, true,
blockManagerInfo)), conf, true)
@@ -310,7 +310,7 @@ trait BlockManagerReplicationBehavior extends SparkFunSuite
master = new BlockManagerMaster(rpcEnv.setupEndpoint("blockmanager-2",
new BlockManagerMasterEndpoint(rpcEnv, true, newConf,
new LiveListenerBus(newConf), shuffleClient, blockManagerInfo,
mapOutputTracker,
- sc.env.shuffleManager, isDriver = true)),
+ isDriver = true)),
rpcEnv.setupEndpoint("blockmanagerHeartbeat-2",
new BlockManagerMasterHeartbeatEndpoint(rpcEnv, true,
blockManagerInfo)), newConf, true)
diff --git
a/core/src/test/scala/org/apache/spark/storage/BlockManagerSuite.scala
b/core/src/test/scala/org/apache/spark/storage/BlockManagerSuite.scala
index aadf52d81853..7e09c17a5f76 100644
--- a/core/src/test/scala/org/apache/spark/storage/BlockManagerSuite.scala
+++ b/core/src/test/scala/org/apache/spark/storage/BlockManagerSuite.scala
@@ -62,7 +62,7 @@ import org.apache.spark.scheduler.{LiveListenerBus,
MapStatus, MergeStatus, Spar
import org.apache.spark.scheduler.cluster.{CoarseGrainedClusterMessages,
CoarseGrainedSchedulerBackend}
import org.apache.spark.security.{CryptoStreamUtils, EncryptionFunSuite}
import org.apache.spark.serializer.{DeserializationStream, JavaSerializer,
KryoDeserializationStream, KryoSerializer, KryoSerializerInstance,
SerializerInstance, SerializerManager}
-import org.apache.spark.shuffle.{MigratableResolver, ShuffleBlockInfo,
ShuffleBlockResolver, ShuffleManager}
+import org.apache.spark.shuffle.{BlockingShuffleManager, MigratableResolver,
ShuffleBlockInfo, ShuffleBlockResolver, ShuffleManager}
import org.apache.spark.shuffle.sort.SortShuffleManager
import org.apache.spark.storage.BlockManagerMessages._
import org.apache.spark.storage.LogBlockType.LogBlockType
@@ -199,7 +199,7 @@ class BlockManagerSuite extends SparkFunSuite with Matchers
with PrivateMethodTe
liveListenerBus = spy[LiveListenerBus](new LiveListenerBus(conf))
master = spy[BlockManagerMaster](new
BlockManagerMaster(rpcEnv.setupEndpoint("blockmanager",
new BlockManagerMasterEndpoint(rpcEnv, true, conf,
- liveListenerBus, None, blockManagerInfo, mapOutputTracker,
shuffleManager,
+ liveListenerBus, None, blockManagerInfo, mapOutputTracker,
isDriver = true)),
rpcEnv.setupEndpoint("blockmanagerHeartbeat",
new BlockManagerMasterHeartbeatEndpoint(rpcEnv, true,
blockManagerInfo)), conf, true))
@@ -2056,7 +2056,7 @@ class BlockManagerSuite extends SparkFunSuite with
Matchers with PrivateMethodTe
}
test("we reject putting blocks when we have the wrong shuffle resolver") {
- val badShuffleManager = mock(classOf[ShuffleManager])
+ val badShuffleManager = mock(classOf[BlockingShuffleManager])
val badShuffleResolver = mock(classOf[ShuffleBlockResolver])
when(badShuffleManager.shuffleBlockResolver).thenReturn(badShuffleResolver)
val shuffleBlockId = ShuffleDataBlockId(0, 0, 0)
@@ -2091,7 +2091,8 @@ class BlockManagerSuite extends SparkFunSuite with
Matchers with PrivateMethodTe
// Retry initialization should succeed once ShuffleManager is initialized
when(mockEnv.isShuffleManagerInitialized).thenReturn(true)
when(mockEnv.waitForShuffleManagerInit(mc.anyLong())).thenReturn(true)
- when(mockEnv.shuffleManager).thenReturn(sortShuffleMgr)
+ when(mockEnv.blockingShuffleManager).thenReturn(sortShuffleMgr)
+
when(mockEnv.shuffleBlockResolver).thenReturn(Some(sortShuffleMgr.shuffleBlockResolver))
assert(bm.migratableResolver != null)
}
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/ShuffledRowRDD.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/ShuffledRowRDD.scala
index 367732dbb205..3265d66853e5 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/ShuffledRowRDD.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/ShuffledRowRDD.scala
@@ -192,7 +192,7 @@ class ShuffledRowRDD(
val sqlMetricsReporter = new SQLShuffleReadMetricsReporter(tempMetrics,
metrics)
val reader = split.asInstanceOf[ShuffledRowRDDPartition].spec match {
case CoalescedPartitionSpec(startReducerIndex, endReducerIndex, _) =>
- SparkEnv.get.shuffleManager.getReader(
+ SparkEnv.get.shuffleManagerFor(dependency).getReader(
dependency.shuffleHandle,
startReducerIndex,
endReducerIndex,
@@ -200,7 +200,7 @@ class ShuffledRowRDD(
sqlMetricsReporter)
case PartialReducerPartitionSpec(reducerIndex, startMapIndex,
endMapIndex, _) =>
- SparkEnv.get.shuffleManager.getReader(
+ SparkEnv.get.shuffleManagerFor(dependency).getReader(
dependency.shuffleHandle,
startMapIndex,
endMapIndex,
@@ -210,7 +210,7 @@ class ShuffledRowRDD(
sqlMetricsReporter)
case PartialMapperPartitionSpec(mapIndex, startReducerIndex,
endReducerIndex) =>
- SparkEnv.get.shuffleManager.getReader(
+ SparkEnv.get.shuffleManagerFor(dependency).getReader(
dependency.shuffleHandle,
mapIndex,
mapIndex + 1,
@@ -220,7 +220,7 @@ class ShuffledRowRDD(
sqlMetricsReporter)
case CoalescedMapperPartitionSpec(startMapIndex, endMapIndex,
numReducers) =>
- SparkEnv.get.shuffleManager.getReader(
+ SparkEnv.get.shuffleManagerFor(dependency).getReader(
dependency.shuffleHandle,
startMapIndex,
endMapIndex,
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala
index 114f221c52f6..dd829e697df6 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala
@@ -300,7 +300,10 @@ object ShuffleExchangeExec {
// corner-cases where a partitioner constructed with `numPartitions`
partitions may output
// fewer partitions (like RangePartitioner, for example).
val conf = SparkEnv.get.conf
- val shuffleManager = SparkEnv.get.shuffleManager
+ // This decision concerns the regular (materialized) shuffle path only. A
pipelined shuffle is
+ // served by a separate pipelined manager (see SparkEnv.shuffleManagerFor)
and does not go
+ // through here, so inspect the blocking manager's type directly.
+ val shuffleManager = SparkEnv.get.blockingShuffleManager
val sortBasedShuffleOn = shuffleManager.isInstanceOf[SortShuffleManager]
val bypassMergeThreshold =
conf.get(config.SHUFFLE_SORT_BYPASS_MERGE_THRESHOLD)
val numParts = partitioner.numPartitions
diff --git
a/streaming/src/test/scala/org/apache/spark/streaming/ReceivedBlockHandlerSuite.scala
b/streaming/src/test/scala/org/apache/spark/streaming/ReceivedBlockHandlerSuite.scala
index 1bf74e6e9a36..7b2d159390ef 100644
---
a/streaming/src/test/scala/org/apache/spark/streaming/ReceivedBlockHandlerSuite.scala
+++
b/streaming/src/test/scala/org/apache/spark/streaming/ReceivedBlockHandlerSuite.scala
@@ -93,7 +93,7 @@ abstract class
BaseReceivedBlockHandlerSuite(enableEncryption: Boolean)
val blockManagerInfo = new mutable.HashMap[BlockManagerId,
BlockManagerInfo]()
blockManagerMaster = new
BlockManagerMaster(rpcEnv.setupEndpoint("blockmanager",
new BlockManagerMasterEndpoint(rpcEnv, true, conf,
- new LiveListenerBus(conf), None, blockManagerInfo, mapOutputTracker,
shuffleManager,
+ new LiveListenerBus(conf), None, blockManagerInfo, mapOutputTracker,
isDriver = true)),
rpcEnv.setupEndpoint("blockmanagerHeartbeat",
new BlockManagerMasterHeartbeatEndpoint(rpcEnv, true,
blockManagerInfo)), conf, true)
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
