Temp fix: enable/disable RMM pooling to work around cuvsTieredIndexBuild/cuvsBruteForceSearch

Author: ldematte

java/cuvs-java/src/test/java/com/nvidia/cuvs/BruteForceAndSearchIT.java

@@ -6,6 +6,7 @@
 
 import static com.carrotsearch.randomizedtesting.RandomizedTest.assumeTrue;
 
+import com.nvidia.cuvs.spi.CuVSProvider;
 import java.io.*;
 import java.nio.file.Files;
 import java.nio.file.Path;
@@ -14,6 +15,7 @@
 import java.util.Map;
 import java.util.UUID;
 import java.util.function.LongToIntFunction;
+import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 
@@ -22,6 +24,12 @@ public class BruteForceAndSearchIT extends CuVSTestCase {
   @Before
   public void setup() {
     assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxAmd64());
+    CuVSProvider.provider().enableRMMPooledMemory(10, 60);
+  }
+
+  @After
+  public void cleanup() {
+    CuVSProvider.provider().resetRMMPooledMemory();
   }
 
   // Sample data and query

java/cuvs-java/src/test/java/com/nvidia/cuvs/BruteForceRandomizedIT.java

@@ -7,9 +7,11 @@
 import static com.carrotsearch.randomizedtesting.RandomizedTest.assumeTrue;
 
 import com.carrotsearch.randomizedtesting.RandomizedRunner;
+import com.nvidia.cuvs.spi.CuVSProvider;
 import java.lang.invoke.MethodHandles;
 import java.util.BitSet;
 import java.util.List;
+import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -26,6 +28,12 @@ public void setup() {
     assumeTrue(isLinuxAmd64());
     initializeRandom();
     log.trace("Random context initialized for test.");
+    CuVSProvider.provider().enableRMMPooledMemory(10, 60);
+  }
+
+  @After
+  public void cleanup() {
+    CuVSProvider.provider().resetRMMPooledMemory();
   }
 
   @Test

java/cuvs-java/src/test/java/com/nvidia/cuvs/TieredIndexIT.java

@@ -10,12 +10,14 @@
 import static org.junit.Assume.assumeTrue;
 
 import com.carrotsearch.randomizedtesting.RandomizedRunner;
+import com.nvidia.cuvs.spi.CuVSProvider;
 import java.lang.invoke.MethodHandles;
 import java.util.Arrays;
 import java.util.BitSet;
 import java.util.List;
 import java.util.Map;
 import java.util.stream.Collectors;
+import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -31,9 +33,15 @@ public class TieredIndexIT extends CuVSTestCase {
   public void setup() {
     assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxAmd64());
     initializeRandom();
+    CuVSProvider.provider().enableRMMPooledMemory(10, 60);
     log.debug("Random context initialized for test");
   }
 
+  @After
+  public void cleanup() {
+    CuVSProvider.provider().resetRMMPooledMemory();
+  }
+
   @Test
   public void testBasicOperations() throws Throwable {
     float[][] initialDataset = {