Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Handle ShuffleManager api calls when plugin is not fully initialized #1540

Merged
merged 3 commits into from
Jan 19, 2021
Merged

Handle ShuffleManager api calls when plugin is not fully initialized #1540

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
216f3e3
Handle ShuffleManager api calls when plugin is not fully initialized
abellina Jan 14, 2021
3c0f39b
Address review comments
abellina Jan 16, 2021
87c44a6
Copyright - add 2021
abellina Jan 19, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 14 additions & 5 deletions sql-plugin/src/main/scala/org/apache/spark/sql/rapids/RapidsShuffleInternalManager.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*
* Copyright (c) 2019-2020, NVIDIA CORPORATION.
* Copyright (c) 2019-2021, NVIDIA CORPORATION.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
Expand Down Expand Up @@ -215,7 +215,6 @@ abstract class RapidsShuffleInternalManagerBase(conf: SparkConf, isDriver: Boole
// so they all need to be lazy values that are executed when things are first called
jlowe marked this conversation as resolved.
Show resolved Hide resolved

// NOTE: this can be null in the driver side.
private[this] lazy val catalog = GpuShuffleEnv.getCatalog
private lazy val env = SparkEnv.get
private lazy val blockManager = env.blockManager
private lazy val shouldFallThroughOnEverything = {
Expand All @@ -232,10 +231,17 @@ abstract class RapidsShuffleInternalManagerBase(conf: SparkConf, isDriver: Boole

private lazy val localBlockManagerId = blockManager.blockManagerId

// Code that expects the shuffle catalog to be initialized gets it this way,
// with error checking in case we are in a bad state.
private def getCatalogOrThrow: ShuffleBufferCatalog =
Option(GpuShuffleEnv.getCatalog).getOrElse(
throw new IllegalStateException("The ShuffleBufferCatalog is not initialized but the " +
"RapidsShuffleManager is configured"))

private lazy val resolver = if (shouldFallThroughOnEverything) {
wrapped.shuffleBlockResolver
} else {
new GpuShuffleBlockResolver(wrapped.shuffleBlockResolver, catalog)
new GpuShuffleBlockResolver(wrapped.shuffleBlockResolver, getCatalogOrThrow)
}

private[this] lazy val transport: Option[RapidsShuffleTransport] = {
Expand All @@ -248,6 +254,7 @@ abstract class RapidsShuffleInternalManagerBase(conf: SparkConf, isDriver: Boole

private[this] lazy val server: Option[RapidsShuffleServer] = {
if (rapidsConf.shuffleTransportEnabled && !isDriver) {
val catalog = getCatalogOrThrow
val requestHandler = new RapidsShuffleRequestHandler() {
override def acquireShuffleBuffer(tableId: Int): RapidsBuffer = {
val shuffleBufferId = catalog.getShuffleBufferId(tableId)
Expand Down Expand Up @@ -295,7 +302,7 @@ abstract class RapidsShuffleInternalManagerBase(conf: SparkConf, isDriver: Boole
gpu.asInstanceOf[GpuShuffleHandle[K, V]],
mapId,
metricsReporter,
catalog,
getCatalogOrThrow,
RapidsBufferCatalog.getDeviceStorage,
server,
gpu.dependency.metrics)
Expand Down Expand Up @@ -335,7 +342,7 @@ abstract class RapidsShuffleInternalManagerBase(conf: SparkConf, isDriver: Boole
context,
metrics,
transport,
catalog,
getCatalogOrThrow,
gpu.dependency.sparkTypes)
case other => {
val shuffleHandle = RapidsShuffleInternalManagerBase.unwrapHandle(other)
Expand All @@ -345,6 +352,7 @@ abstract class RapidsShuffleInternalManagerBase(conf: SparkConf, isDriver: Boole
}

def registerGpuShuffle(shuffleId: Int): Unit = {
val catalog = GpuShuffleEnv.getCatalog
if (catalog != null) {
// Note that in local mode this can be called multiple times.
logInfo(s"Registering shuffle $shuffleId")
Expand All @@ -353,6 +361,7 @@ abstract class RapidsShuffleInternalManagerBase(conf: SparkConf, isDriver: Boole
}

def unregisterGpuShuffle(shuffleId: Int): Unit = {
val catalog = GpuShuffleEnv.getCatalog
if (catalog != null) {
logInfo(s"Unregistering shuffle $shuffleId")
catalog.unregisterShuffle(shuffleId)
Expand Down