Fallback to ARENA if ASYNC configured and driver < 11.5.0 (#4947)

* Warn if using CUDA driver less than 11.5.0 with ASYNC Signed-off-by: Alessandro Bellina <[email protected]> * Update sql-plugin/src/main/scala/com/nvidia/spark/rapids/RapidsConf.scala Co-authored-by: Jason Lowe <[email protected]> Co-authored-by: Jason Lowe <[email protected]>
NVIDIA · Mar 14, 2022 · 611f54d · 611f54d
1 parent eda23d8
commit 611f54d
Showing 1 changed file with 16 additions and 7 deletions.
diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/RapidsConf.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/RapidsConf.scala
@@ -1492,14 +1492,23 @@ class RapidsConf(conf: Map[String, String]) extends Logging {
   lazy val isPooledMemEnabled: Boolean = get(POOLED_MEM)
 
   lazy val rmmPool: String = {
-    val pool = get(RMM_POOL)
-    if ("ASYNC".equalsIgnoreCase(pool) &&
-        (Cuda.getRuntimeVersion < 11020 || Cuda.getDriverVersion < 11020)) {
-      logWarning("CUDA runtime/driver does not support the ASYNC allocator, falling back to ARENA")
-      "ARENA"
-    } else {
-      pool
+    var pool = get(RMM_POOL)
+    if ("ASYNC".equalsIgnoreCase(pool)) {
+      val driverVersion = Cuda.getDriverVersion
+      val runtimeVersion = Cuda.getRuntimeVersion
+      var fallbackMessage: Option[String] = None
+      if (runtimeVersion < 11020 || driverVersion < 11020) {
+        fallbackMessage = Some("CUDA runtime/driver does not support the ASYNC allocator")
+      } else if (driverVersion < 11050) {
+        fallbackMessage = Some("CUDA drivers before 11.5 have known incompatibilities with " +
+          "the ASYNC allocator")
+      }
+      if (fallbackMessage.isDefined) {
+        logWarning(s"${fallbackMessage.get}, falling back to ARENA")
+        pool = "ARENA"
+      }
     }
+    pool
   }
 
   lazy val rmmAllocFraction: Double = get(RMM_ALLOC_FRACTION)