From e0ef44aa89564e234e7ed150f72fdf9e70671ae2 Mon Sep 17 00:00:00 2001 From: Jason Lowe Date: Thu, 14 Mar 2024 08:45:55 -0500 Subject: [PATCH] Turn on transition logging in HostAllocSuite (#10590) Signed-off-by: Jason Lowe --- .../scala/com/nvidia/spark/rapids/HostAllocSuite.scala | 7 +++++-- 1 file changed, 5 insertions(+), 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 index 13994c2a997..f0b541f6244 100644 --- 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)