Skip to content

Commit

Permalink
Add classloader diagnostics to initShuffleManager error message (#10871)
Browse files Browse the repository at this point in the history
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 May 29, 2024
1 parent f0b13ed commit dfcde72
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 dfcde72

Please sign in to comment.