Expose raft::memory_tracking_resources to C/Python/Java/Rust

c/include/cuvs/core/c_api.h

@@ -87,6 +87,33 @@ typedef uintptr_t cuvsResources_t;
  */
 CUVS_EXPORT cuvsError_t cuvsResourcesCreate(cuvsResources_t* res);
 
+/**
+ * @brief Create an opaque C handle for C++ type `raft::resources` whose memory
+ *        allocations are tracked and written as CSV samples from a background
+ *        thread.
+ *
+ * The returned handle wraps all reachable memory resources (host, pinned,
+ * managed, device, workspace, large_workspace) with allocation-tracking
+ * adaptors and replaces the global host and device memory resources for the
+ * lifetime of the handle. It is otherwise indistinguishable from a handle
+ * created by ::cuvsResourcesCreate and can be used wherever a
+ * ::cuvsResources_t is accepted. The CSV reporter is stopped and the global
+ * memory resources are restored when the handle is destroyed via
+ * ::cuvsResourcesDestroy.
+ *
+ * @param[out] res                 cuvsResources_t opaque C handle
+ * @param[in]  csv_path            Path to the output CSV file
+ *                                 (created/truncated). Must be a non-empty,
+ *                                 null-terminated UTF-8 string.
+ * @param[in]  sample_interval_ms  Minimum time in milliseconds between
+ *                                 successive CSV samples. Pass 10 to match the
+ *                                 C++ default.
+ * @return cuvsError_t
+ */
+CUVS_EXPORT cuvsError_t cuvsResourcesCreateWithMemoryTracking(cuvsResources_t* res,
+                                                              const char* csv_path,
+                                                              int64_t sample_interval_ms);
+
 /**
  * @brief Destroy and de-allocate opaque C handle for C++ type `raft::resources`
  *

c/src/core/c_api.cpp

@@ -12,6 +12,7 @@
 #include <raft/core/resource/resource_types.hpp>
 #include <raft/core/resources.hpp>
 #include <raft/util/cudart_utils.hpp>
+#include <raft/util/memory_tracking_resources.hpp>
 #include <rapids_logger/logger.hpp>
 #include <rmm/cuda_stream_view.hpp>
 #include <rmm/mr/cuda_memory_resource.hpp>
@@ -23,8 +24,11 @@
 
 #include "../core/exceptions.hpp"
 
+#include <chrono>
 #include <cstdint>
 #include <memory>
+#include <stdexcept>
+#include <string>
 #include <thread>
 
 extern "C" cuvsError_t cuvsResourcesCreate(cuvsResources_t* res)
@@ -35,6 +39,23 @@ extern "C" cuvsError_t cuvsResourcesCreate(cuvsResources_t* res)
   });
 }
 
+extern "C" cuvsError_t cuvsResourcesCreateWithMemoryTracking(cuvsResources_t* res,
+                                                             const char* csv_path,
+                                                             int64_t sample_interval_ms)
+{
+  return cuvs::core::translate_exceptions([=] {
+    if (csv_path == nullptr || csv_path[0] == '\0') {
+      throw std::invalid_argument("csv_path must be a non-empty string");
+    }
+    if (sample_interval_ms < 0) {
+      throw std::invalid_argument("sample_interval_ms must be >= 0");
+    }
+    auto res_ptr = new raft::memory_tracking_resources{
+      std::string{csv_path}, std::chrono::milliseconds{sample_interval_ms}};
+    *res = reinterpret_cast<uintptr_t>(res_ptr);
+  });
+}
+
 extern "C" cuvsError_t cuvsResourcesDestroy(cuvsResources_t res)
 {
   return cuvs::core::translate_exceptions([=] {

java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSResources.java

@@ -1,11 +1,12 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 package com.nvidia.cuvs;
 
 import com.nvidia.cuvs.spi.CuVSProvider;
 import java.nio.file.Path;
+import java.time.Duration;
 
 /**
  * Used for allocating resources for cuVS
@@ -78,4 +79,34 @@ static CuVSResources create() throws Throwable {
   static CuVSResources create(Path tempDirectory) throws Throwable {
     return CuVSProvider.provider().newCuVSResources(tempDirectory);
   }
+
+  /**
+   * Creates a new resources whose memory allocations are tracked and written as
+   * CSV samples from a background thread.
+   * <p>
+   * The returned handle wraps all reachable memory resources (host, pinned,
+   * managed, device, workspace, large_workspace) with allocation-tracking
+   * adaptors and replaces the global host and device memory resources for the
+   * lifetime of the handle. It is otherwise indistinguishable from a handle
+   * created by {@link #create(Path)} and can be used wherever a
+   * {@link CuVSResources} is accepted. The CSV reporter is stopped and the
+   * global memory resources are restored when the handle is closed.
+   *
+   * @param tempDirectory                   the temporary directory to use for
+   *                                        intermediate operations
+   * @param memoryTrackingCsvPath           path to the output CSV file
+   *                                        (created/truncated)
+   * @param memoryTrackingSampleInterval    minimum interval between successive
+   *                                        CSV samples
+   * @throws UnsupportedOperationException if the provider does not support cuvs
+   * @throws LibraryException if the native library cannot be loaded
+   */
+  static CuVSResources create(
+      Path tempDirectory,
+      Path memoryTrackingCsvPath,
+      Duration memoryTrackingSampleInterval) throws Throwable {
+    return CuVSProvider.provider()
+        .newCuVSResources(
+            tempDirectory, memoryTrackingCsvPath, memoryTrackingSampleInterval);
+  }
 }

java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java

@@ -8,6 +8,7 @@
 import java.lang.invoke.MethodHandle;
 import java.lang.invoke.MethodType;
 import java.nio.file.Path;
+import java.time.Duration;
 
 /**
  * A provider of low-level cuvs resources and builders.
@@ -35,6 +36,31 @@ default Path nativeLibraryPath() {
   /** Creates a new CuVSResources. */
   CuVSResources newCuVSResources(Path tempDirectory) throws Throwable;
 
+  /**
+   * Creates a new CuVSResources whose memory allocations are tracked and
+   * written as CSV samples from a background thread.
+   *
+   * <p>This method is declared as a {@code default} method so that adding it
+   * does not break binary compatibility with providers compiled against an
+   * earlier version of this interface; the default implementation throws
+   * {@link UnsupportedOperationException} and providers must override it to
+   * opt in.
+   *
+   * @param tempDirectory                the temporary directory to use for
+   *                                     intermediate operations
+   * @param memoryTrackingCsvPath        path to the output CSV file
+   *                                     (created/truncated)
+   * @param memoryTrackingSampleInterval minimum interval between successive
+   *                                     CSV samples
+   */
+  default CuVSResources newCuVSResources(
+      Path tempDirectory,
+      Path memoryTrackingCsvPath,
+      Duration memoryTrackingSampleInterval) throws Throwable {
+    throw new UnsupportedOperationException(
+        "Memory-tracking resources are not supported by this provider");
+  }
+
   /** Create a {@link CuVSMatrix.Builder} instance for a host memory matrix **/
   CuVSMatrix.Builder<CuVSHostMatrix> newHostMatrixBuilder(
       long size, long dimensions, CuVSMatrix.DataType dataType);

java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java

@@ -7,6 +7,7 @@
 import com.nvidia.cuvs.*;
 import java.lang.invoke.MethodHandle;
 import java.nio.file.Path;
+import java.time.Duration;
 import java.util.logging.Level;
 
 /**
@@ -25,6 +26,14 @@ public CuVSResources newCuVSResources(Path tempDirectory) {
     throw new UnsupportedOperationException(reasons);
   }
 
+  @Override
+  public CuVSResources newCuVSResources(
+      Path tempDirectory,
+      Path memoryTrackingCsvPath,
+      Duration memoryTrackingSampleInterval) {
+    throw new UnsupportedOperationException(reasons);
+  }
+
   @Override
   public BruteForceIndex.Builder newBruteForceIndexBuilder(CuVSResources cuVSResources) {
     throw new UnsupportedOperationException(reasons);
@@ -47,8 +56,8 @@ public HnswIndex hnswIndexFromCagra(HnswIndexParams hnswParams, CagraIndex cagra
   }
 
   @Override
-  public HnswIndex hnswIndexBuild(CuVSResources resources, HnswIndexParams hnswParams, CuVSMatrix dataset)
-      throws Throwable {
+  public HnswIndex hnswIndexBuild(
+      CuVSResources resources, HnswIndexParams hnswParams, CuVSMatrix dataset) throws Throwable {
     throw new UnsupportedOperationException(reasons);
   }
 

java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSResourcesImpl.java

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 package com.nvidia.cuvs.internal;
@@ -13,7 +13,10 @@
 import com.nvidia.cuvs.internal.common.PinnedMemoryBuffer;
 import java.lang.foreign.Arena;
 import java.lang.foreign.MemorySegment;
+import java.lang.foreign.ValueLayout;
+import java.nio.charset.StandardCharsets;
 import java.nio.file.Path;
+import java.time.Duration;
 
 /**
  * Used for allocating resources for cuVS
@@ -46,6 +49,45 @@ public CuVSResourcesImpl(Path tempDirectory) {
     }
   }
 
+  /**
+   * Constructor that allocates a tracking resources handle. All memory
+   * allocations made through this handle are written as CSV samples to
+   * {@code memoryTrackingCsvPath} from a background thread, restoring the
+   * global memory resources on {@link #close()}.
+   *
+   * @param tempDirectory                the temporary directory to use for
+   *                                     intermediate operations
+   * @param memoryTrackingCsvPath        path to the output CSV file
+   *                                     (created/truncated)
+   * @param memoryTrackingSampleInterval minimum interval between successive
+   *                                     CSV samples
+   */
+  public CuVSResourcesImpl(
+      Path tempDirectory,
+      Path memoryTrackingCsvPath,
+      Duration memoryTrackingSampleInterval) {
+    this.tempDirectory = tempDirectory;
+    try (var localArena = Arena.ofConfined()) {
+      var resourcesMemorySegment = localArena.allocate(cuvsResources_t);
+      byte[] pathBytes =
+          memoryTrackingCsvPath.toString().getBytes(StandardCharsets.UTF_8);
+      var pathSegment = localArena.allocate(pathBytes.length + 1L);
+      MemorySegment.copy(
+          pathBytes, 0, pathSegment, ValueLayout.JAVA_BYTE, 0, pathBytes.length);
+      pathSegment.set(ValueLayout.JAVA_BYTE, pathBytes.length, (byte) 0);
+      long sampleIntervalMs = memoryTrackingSampleInterval.toMillis();
+      checkCuVSError(
+          cuvsResourcesCreateWithMemoryTracking(
+              resourcesMemorySegment, pathSegment, sampleIntervalMs),
+          "cuvsResourcesCreateWithMemoryTracking");
+      this.resourceHandle = resourcesMemorySegment.get(cuvsResources_t, 0);
+      var deviceIdPtr = localArena.allocate(C_INT);
+      checkCuVSError(cuvsDeviceIdGet(resourceHandle, deviceIdPtr), "cuvsDeviceIdGet");
+      this.deviceId = deviceIdPtr.get(C_INT, 0);
+      this.access = new ScopedAccessWithHostBuffer(resourceHandle, hostBuffer.address());
+    }
+  }
+
   @Override
   public ScopedAccess access() {
     return this.access;

java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java

@@ -26,6 +26,7 @@
 import java.lang.invoke.MethodType;
 import java.nio.file.Files;
 import java.nio.file.Path;
+import java.time.Duration;
 import java.util.Locale;
 import java.util.Objects;
 import java.util.jar.JarFile;
@@ -233,6 +234,24 @@ public CuVSResources newCuVSResources(Path tempDirectory) {
     return new CuVSResourcesImpl(tempDirectory);
   }
 
+  @Override
+  public CuVSResources newCuVSResources(
+      Path tempDirectory,
+      Path memoryTrackingCsvPath,
+      Duration memoryTrackingSampleInterval) {
+    Objects.requireNonNull(tempDirectory);
+    Objects.requireNonNull(memoryTrackingCsvPath);
+    Objects.requireNonNull(memoryTrackingSampleInterval);
+    if (Files.notExists(tempDirectory)) {
+      throw new IllegalArgumentException("does not exist:" + tempDirectory);
+    }
+    if (!Files.isDirectory(tempDirectory)) {
+      throw new IllegalArgumentException("not a directory:" + tempDirectory);
+    }
+    return new CuVSResourcesImpl(
+        tempDirectory, memoryTrackingCsvPath, memoryTrackingSampleInterval);
+  }
+
   @Override
   public BruteForceIndex.Builder newBruteForceIndexBuilder(CuVSResources cuVSResources) {
     return BruteForceIndexImpl.newBuilder(Objects.requireNonNull(cuVSResources));
@@ -255,8 +274,8 @@ public HnswIndex hnswIndexFromCagra(HnswIndexParams hnswParams, CagraIndex cagra
   }
 
   @Override
-  public HnswIndex hnswIndexBuild(CuVSResources resources, HnswIndexParams hnswParams, CuVSMatrix dataset)
-      throws Throwable {
+  public HnswIndex hnswIndexBuild(
+      CuVSResources resources, HnswIndexParams hnswParams, CuVSMatrix dataset) throws Throwable {
     return HnswIndexImpl.build(resources, hnswParams, dataset);
   }
 

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

@@ -0,0 +1,58 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package com.nvidia.cuvs;
+
+import static com.carrotsearch.randomizedtesting.RandomizedTest.assumeTrue;
+import static org.junit.Assert.assertTrue;
+
+import com.nvidia.cuvs.spi.CuVSProvider;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.time.Duration;
+import org.junit.Before;
+import org.junit.Test;
+
+public class MemoryTrackingResourcesIT extends CuVSTestCase {
+
+  @Before
+  public void setup() {
+    assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxAmd64());
+  }
+
+  @Test
+  public void writesNonEmptyCsv() throws Throwable {
+    Path csv = Files.createTempFile("cuvs-mtrack", ".csv");
+    try {
+      try (var resources =
+          CuVSResources.create(
+              CuVSProvider.tempDirectory(), csv, Duration.ofMillis(2))) {
+
+        // Allocate / release a couple of small device buffers so the
+        // background CSV reporter has something to report.
+        var b1 =
+            CuVSMatrix.deviceBuilder(resources, 64, 32, CuVSMatrix.DataType.FLOAT);
+        for (int i = 0; i < 64; ++i) {
+          b1.addVector(new float[32]);
+        }
+        try (var m1 = b1.build()) {
+          var b2 =
+              CuVSMatrix.deviceBuilder(resources, 32, 16, CuVSMatrix.DataType.FLOAT);
+          for (int i = 0; i < 32; ++i) {
+            b2.addVector(new float[16]);
+          }
+          try (var m2 = b2.build()) {
+            // Allow the background CSV reporter at least a few ticks
+            // before the matrices are released and the handle closed.
+            Thread.sleep(20);
+          }
+        }
+      }
+      // closing the resources flushes the CSV and restores globals
+      assertTrue("csv should be non-empty", Files.size(csv) > 0);
+    } finally {
+      Files.deleteIfExists(csv);
+    }
+  }
+}

python/cuvs/cuvs/common/c_api.pxd

@@ -1,5 +1,5 @@
 #
-# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION.
 # SPDX-License-Identifier: Apache-2.0
 #
 # cython: language_level=3
@@ -19,6 +19,10 @@ cdef extern from "cuvs/core/c_api.h":
         CUVS_SUCCESS
 
     cuvsError_t cuvsResourcesCreate(cuvsResources_t* res)
+    cuvsError_t cuvsResourcesCreateWithMemoryTracking(
+        cuvsResources_t* res,
+        const char* csv_path,
+        int64_t sample_interval_ms)
     cuvsError_t cuvsResourcesDestroy(cuvsResources_t res)
     cuvsError_t cuvsStreamSet(cuvsResources_t res, cudaStream_t stream)
     cuvsError_t cuvsStreamSync(cuvsResources_t res)