Turn on transition logging in HostAllocSuite (#10590)

Signed-off-by: Jason Lowe <[email protected]>
NVIDIA · Mar 14, 2024 · e0ef44a · e0ef44a
1 parent 9105fd7
commit e0ef44a
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/sql-plugin/src/test/scala/com/nvidia/spark/rapids/HostAllocSuite.scala b/sql-plugin/src/test/scala/com/nvidia/spark/rapids/HostAllocSuite.scala
@@ -36,6 +36,7 @@ import org.apache.spark.sql.rapids.execution.TrampolineUtil
 class HostAllocSuite extends AnyFunSuite with BeforeAndAfterEach with
     BeforeAndAfterAll with TimeLimits {
   private val sqlConf = new SQLConf()
+  sqlConf.setConfString("spark.rapids.memory.gpu.state.debug", "stderr")
   private val rc = new RapidsConf(sqlConf)
   private val timeoutMs = 10000
 
@@ -332,8 +333,10 @@ class HostAllocSuite extends AnyFunSuite with BeforeAndAfterEach with
   override def afterAll(): Unit = {
     RapidsBufferCatalog.close()
     PinnedMemoryPool.shutdown()
-    if (!rmmWasInitialized) {
-      Rmm.shutdown()
+    Rmm.shutdown()
+    if (rmmWasInitialized) {
+      // put RMM back for other tests to use
+      Rmm.initialize(RmmAllocationMode.CUDA_DEFAULT, null, 512 * 1024 * 1024)
     }
     // 1 GiB
     PinnedMemoryPool.initialize(1 * 1024 * 1024 * 1024)