Skip to content

Commit

Permalink
Add classloader diagnostics to initShuffleManager error message (NVID…
Browse files Browse the repository at this point in the history
…IA#10871)

Add classloader diagnostics to initShuffleManager error message

---------

Signed-off-by: Zach Puller <zpuller@nvidia.com>
Co-authored-by: Jason Lowe <jlowe@nvidia.com>
Co-authored-by: Gera Shegalov <gshegalov@nvidia.com>
Co-authored-by: Alessandro Bellina <abellina@gmail.com>
  • Loading branch information
4 people authored and SurajAralihalli committed Jul 12, 2024
1 parent b0f81d8 commit 09ae905
Showing 1 changed file with 14 additions and 3 deletions.
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*
* Copyright (c) 2019-2023, NVIDIA CORPORATION.
* Copyright (c) 2019-2024, 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 @@ -104,11 +104,22 @@ object GpuShuffleEnv extends Logging {
// this forces the initialization when we know we are ready in the driver and executor.
//
def initShuffleManager(): Unit = {
SparkEnv.get.shuffleManager match {
val shuffleManager = SparkEnv.get.shuffleManager
shuffleManager match {
case rapidsShuffleManager: RapidsShuffleManagerLike =>
rapidsShuffleManager.initialize
case _ =>
throw new IllegalStateException(s"Cannot initialize the RAPIDS Shuffle Manager")
val rsmLoaderViaShuffleManager = shuffleManager.getClass.getSuperclass.getInterfaces
.collectFirst {
case c if c.getName == classOf[RapidsShuffleManagerLike].getName => c.getClassLoader
}
val rsmLoaderDirect = classOf[RapidsShuffleManagerLike].getClassLoader

throw new IllegalStateException(s"Cannot initialize the RAPIDS Shuffle Manager " +
s"${shuffleManager}! Expected: an instance of RapidsShuffleManagerLike loaded by " +
s"${rsmLoaderDirect}. Actual: ${shuffleManager} tagged with RapidsShuffleManagerLike " +
s"loaded by: ${rsmLoaderViaShuffleManager}"
)
}
}

Expand Down

0 comments on commit 09ae905

Please sign in to comment.