From 48cd2de8c84e926c358b0f769760a0098c227b79 Mon Sep 17 00:00:00 2001 From: achirkin Date: Tue, 12 May 2026 12:43:31 +0200 Subject: [PATCH 1/3] Expose raft::memory_tracking_resources to C/Python/Java/Rust --- c/include/cuvs/core/c_api.h | 27 +++++ c/src/core/c_api.cpp | 18 +++ docs/source/api_basics.rst | 105 ++++++++++++++++++ .../java/com/nvidia/cuvs/CuVSResources.java | 33 +++++- .../com/nvidia/cuvs/spi/CuVSProvider.java | 17 +++ .../nvidia/cuvs/spi/UnsupportedProvider.java | 13 ++- .../cuvs/internal/CuVSResourcesImpl.java | 44 +++++++- .../com/nvidia/cuvs/spi/JDKProvider.java | 23 +++- .../cuvs/MemoryTrackingResourcesIT.java | 58 ++++++++++ python/cuvs/cuvs/common/c_api.pxd | 6 +- python/cuvs/cuvs/common/resources.pyx | 35 +++++- .../cuvs/cuvs/tests/test_memory_tracking.py | 40 +++++++ rust/cuvs-sys/src/bindings.rs | 9 ++ rust/cuvs/Cargo.toml | 1 + rust/cuvs/src/resources.rs | 54 ++++++++- 15 files changed, 472 insertions(+), 11 deletions(-) create mode 100644 java/cuvs-java/src/test/java/com/nvidia/cuvs/MemoryTrackingResourcesIT.java create mode 100644 python/cuvs/cuvs/tests/test_memory_tracking.py diff --git a/c/include/cuvs/core/c_api.h b/c/include/cuvs/core/c_api.h index 00d4729481..39a0148abe 100644 --- a/c/include/cuvs/core/c_api.h +++ b/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` * diff --git a/c/src/core/c_api.cpp b/c/src/core/c_api.cpp index f4e3664482..026c6a3553 100644 --- a/c/src/core/c_api.cpp +++ b/c/src/core/c_api.cpp @@ -12,6 +12,7 @@ #include #include #include +#include #include #include #include @@ -23,8 +24,11 @@ #include "../core/exceptions.hpp" +#include #include #include +#include +#include #include extern "C" cuvsError_t cuvsResourcesCreate(cuvsResources_t* res) @@ -35,6 +39,20 @@ 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"); + } + auto res_ptr = new raft::memory_tracking_resources{ + std::string{csv_path}, std::chrono::milliseconds{sample_interval_ms}}; + *res = reinterpret_cast(res_ptr); + }); +} + extern "C" cuvsError_t cuvsResourcesDestroy(cuvsResources_t res) { return cuvs::core::translate_exceptions([=] { diff --git a/docs/source/api_basics.rst b/docs/source/api_basics.rst index 5ffb1da630..f99b10a23c 100644 --- a/docs/source/api_basics.rst +++ b/docs/source/api_basics.rst @@ -3,6 +3,7 @@ cuVS API Basics - `Memory management`_ - `Resource management`_ +- `Memory tracking`_ Memory management ----------------- @@ -88,3 +89,107 @@ Rust .. code-block:: rust let res = cuvs::Resources::new()?; + + +Memory tracking +--------------- + +A resources handle whose memory allocations are tracked and written as CSV +samples from a background thread can be created in any of the supported +languages. The 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 regular resources handle +and can be passed to every cuVS API that accepts one. The CSV reporter is +stopped and the global memory resources are restored when the handle is +destroyed. + +.. note:: + + - The handle replaces the **global** host and device memory resources while + it is alive. Do not create multiple tracking handles concurrently and make + sure the handle outlives every consumer (matrices, indexes, search results, + ...) that allocates memory through cuVS. + - The CSV file is flushed eagerly: the header is flushed on construction and + every sample row is flushed as soon as it is written, so the file can be + tailed while the handle is alive. Destroying the handle stops the + background sampler and writes one final row. + - The sample interval is a *minimum* time between samples. The background + thread blocks until an allocation/deallocation occurs, then sleeps for at + least ``sample_interval`` before writing the next row; quiescent periods do + not produce extra rows. + +C +^ + +.. code-block:: c + + #include + #include + + cuvsResources_t res; + // 10 ms sampling matches the C++ default. + cuvsResourcesCreateWithMemoryTracking(&res, "/tmp/allocations.csv", 10); + + // ... do some processing ... + + cuvsResourcesDestroy(res); + +C++ +^^^ + +.. code-block:: c++ + + #include + + // Sample interval defaults to std::chrono::milliseconds{10}. + raft::memory_tracking_resources res{"/tmp/allocations.csv"}; + + // ... do some processing ... + // `res` is implicitly convertible to raft::resources& and can be passed + // to any cuVS / raft API that accepts a resources handle. + +Python +^^^^^^ + +.. code-block:: python + + from cuvs.common import Resources + + res = Resources( + memory_tracking_csv_path="/tmp/allocations.csv", + memory_tracking_sample_interval_ms=10, + ) + + # ... do some processing ... + + del res # flushes the CSV and restores the global memory resources + +Java +^^^^ + +.. code-block:: java + + import com.nvidia.cuvs.CuVSResources; + import com.nvidia.cuvs.spi.CuVSProvider; + import java.nio.file.Path; + import java.time.Duration; + + try (var res = CuVSResources.create( + CuVSProvider.tempDirectory(), + Path.of("/tmp/allocations.csv"), + Duration.ofMillis(10))) { + // ... do some processing ... + } + +Rust +^^^^ + +.. code-block:: rust + + use std::time::Duration; + + let res = cuvs::Resources::with_memory_tracking( + "/tmp/allocations.csv", + Some(Duration::from_millis(10)), + )?; diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSResources.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSResources.java index b105580328..51c7074a98 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSResources.java +++ b/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. + *

+ * 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); + } } diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java index c39578755c..ed89786308 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java +++ b/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,22 @@ 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. + * + * @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 + */ + CuVSResources newCuVSResources( + Path tempDirectory, + Path memoryTrackingCsvPath, + Duration memoryTrackingSampleInterval) throws Throwable; + /** Create a {@link CuVSMatrix.Builder} instance for a host memory matrix **/ CuVSMatrix.Builder newHostMatrixBuilder( long size, long dimensions, CuVSMatrix.DataType dataType); diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java index 7cbeee4e75..3dd3b2f1c2 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java +++ b/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); } diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSResourcesImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSResourcesImpl.java index efdf7283ac..fb14d07eea 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSResourcesImpl.java +++ b/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; diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java index 1d3199f26f..1f6d21af3e 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java +++ b/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); } diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/MemoryTrackingResourcesIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/MemoryTrackingResourcesIT.java new file mode 100644 index 0000000000..0e70033002 --- /dev/null +++ b/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); + } + } +} diff --git a/python/cuvs/cuvs/common/c_api.pxd b/python/cuvs/cuvs/common/c_api.pxd index 6ccfe47159..7cd74d980a 100644 --- a/python/cuvs/cuvs/common/c_api.pxd +++ b/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) diff --git a/python/cuvs/cuvs/common/resources.pyx b/python/cuvs/cuvs/common/resources.pyx index cf6f3284a6..18233a56e0 100644 --- a/python/cuvs/cuvs/common/resources.pyx +++ b/python/cuvs/cuvs/common/resources.pyx @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. # SPDX-License-Identifier: Apache-2.0 # # cython: language_level=3 @@ -7,10 +7,12 @@ import functools from cuda.bindings.cyruntime cimport cudaStream_t +from libc.stdint cimport int64_t from cuvs.common.c_api cimport ( cuvsResources_t, cuvsResourcesCreate, + cuvsResourcesCreateWithMemoryTracking, cuvsResourcesDestroy, cuvsStreamSet, cuvsStreamSync, @@ -29,6 +31,18 @@ cdef class Resources: Parameters ---------- stream : Optional stream to use for ordering CUDA instructions + memory_tracking_csv_path : Optional path-like + If provided, the handle wraps all reachable memory resources + (host, pinned, managed, device, workspace, large_workspace) + with allocation-tracking adaptors and logs CSV samples to the + given file from a background thread. The CSV file is created + or truncated. The global host and device memory resources are + replaced for the lifetime of the handle and restored when the + handle is destroyed. + memory_tracking_sample_interval_ms : int, default ``10`` + Minimum interval between successive CSV samples, in + milliseconds. Ignored when ``memory_tracking_csv_path`` is + ``None``. Examples -------- @@ -50,10 +64,25 @@ cdef class Resources: >>> >>> cupy_stream = cupy.cuda.Stream() >>> handle = Resources(stream=cupy_stream.ptr) + + Tracking memory allocations to a CSV file: + + >>> from cuvs.common import Resources + >>> handle = Resources(memory_tracking_csv_path="/tmp/allocations.csv", + ... memory_tracking_sample_interval_ms=10) # doctest: +SKIP """ - def __cinit__(self, stream=None): - check_cuvs(cuvsResourcesCreate(&self.c_obj)) + def __cinit__(self, stream=None, memory_tracking_csv_path=None, + memory_tracking_sample_interval_ms=10): + cdef bytes csv_bytes + if memory_tracking_csv_path: + csv_bytes = str(memory_tracking_csv_path).encode("utf-8") + check_cuvs(cuvsResourcesCreateWithMemoryTracking( + &self.c_obj, + csv_bytes, + memory_tracking_sample_interval_ms)) + else: + check_cuvs(cuvsResourcesCreate(&self.c_obj)) if stream: check_cuvs(cuvsStreamSet(self.c_obj, stream)) diff --git a/python/cuvs/cuvs/tests/test_memory_tracking.py b/python/cuvs/cuvs/tests/test_memory_tracking.py new file mode 100644 index 0000000000..e9321cc278 --- /dev/null +++ b/python/cuvs/cuvs/tests/test_memory_tracking.py @@ -0,0 +1,40 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 +# + +import time + +import cupy as cp + +from cuvs.common import Resources + + +def test_memory_tracking_writes_csv(tmp_path): + """Allocate a couple of small device buffers under a tracking + Resources handle and confirm that the CSV reporter wrote at least + one row before the handle was destroyed. + """ + csv = tmp_path / "alloc.csv" + + res = Resources( + memory_tracking_csv_path=str(csv), + memory_tracking_sample_interval_ms=2, + ) + try: + a = cp.zeros((1024,), dtype=cp.float32) + b = cp.zeros((2048,), dtype=cp.float32) + res.sync() + # Give the background reporter enough time to emit at least one + # sample (interval is 2 ms above). + time.sleep(0.05) + del a, b + finally: + # Destroying the handle flushes the CSV and restores the + # global host/device memory resources. + del res + + assert csv.exists(), f"expected csv file at {csv}" + assert csv.stat().st_size > 0, "tracking csv should be non-empty" + + lines = csv.read_text().splitlines() + assert lines, "expected at least one line (header) in the csv" diff --git a/rust/cuvs-sys/src/bindings.rs b/rust/cuvs-sys/src/bindings.rs index 0498b77f3a..171af6f422 100644 --- a/rust/cuvs-sys/src/bindings.rs +++ b/rust/cuvs-sys/src/bindings.rs @@ -186,6 +186,15 @@ unsafe extern "C" { #[doc = " @brief Create an Initialized opaque C handle for C++ type `raft::resources`\n\n @param[in] res cuvsResources_t opaque C handle\n @return cuvsError_t"] pub fn cuvsResourcesCreate(res: *mut cuvsResources_t) -> cuvsError_t; } +unsafe extern "C" { + #[must_use] + #[doc = " @brief Create an opaque C handle for C++ type `raft::resources` whose memory\n allocations are tracked and written as CSV samples from a background\n thread.\n\n The returned handle wraps all reachable memory resources (host, pinned,\n managed, device, workspace, large_workspace) with allocation-tracking\n adaptors and replaces the global host and device memory resources for the\n lifetime of the handle. It is otherwise indistinguishable from a handle\n created by ::cuvsResourcesCreate and can be used wherever a\n ::cuvsResources_t is accepted. The CSV reporter is stopped and the global\n memory resources are restored when the handle is destroyed via\n ::cuvsResourcesDestroy.\n\n @param[out] res cuvsResources_t opaque C handle\n @param[in] csv_path Path to the output CSV file\n (created/truncated). Must be a non-empty,\n null-terminated UTF-8 string.\n @param[in] sample_interval_ms Minimum time in milliseconds between\n successive CSV samples. Pass 10 to match the\n C++ default.\n @return cuvsError_t"] + pub fn cuvsResourcesCreateWithMemoryTracking( + res: *mut cuvsResources_t, + csv_path: *const ::std::os::raw::c_char, + sample_interval_ms: i64, + ) -> cuvsError_t; +} unsafe extern "C" { #[must_use] #[doc = " @brief Destroy and de-allocate opaque C handle for C++ type `raft::resources`\n\n @param[in] res cuvsResources_t opaque C handle\n @return cuvsError_t"] diff --git a/rust/cuvs/Cargo.toml b/rust/cuvs/Cargo.toml index 2a8a4a3ffa..bd55817438 100644 --- a/rust/cuvs/Cargo.toml +++ b/rust/cuvs/Cargo.toml @@ -19,6 +19,7 @@ ndarray = "0.15" [dev-dependencies] ndarray-rand = "0.14" mark-flaky-tests = "1" +tempfile = "3" [package.metadata.docs.rs] features = ["doc-only"] diff --git a/rust/cuvs/src/resources.rs b/rust/cuvs/src/resources.rs index 70f128abb7..72e23d5cdd 100644 --- a/rust/cuvs/src/resources.rs +++ b/rust/cuvs/src/resources.rs @@ -3,8 +3,11 @@ * SPDX-License-Identifier: Apache-2.0 */ -use crate::error::{Result, check_cuvs}; +use crate::error::{Error, Result, check_cuvs}; +use std::ffi::CString; use std::io::{Write, stderr}; +use std::path::Path; +use std::time::Duration; /// Resources are objects that are shared between function calls, /// and includes things like CUDA streams, cuBLAS handles and other @@ -22,6 +25,40 @@ impl Resources { Ok(Resources(res)) } + /// Returns a new `Resources` object whose memory allocations are tracked + /// and written as CSV samples to `csv_path` from a background thread. + /// + /// The 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. The CSV reporter is stopped and the global + /// memory resources are restored when the handle is dropped. + /// + /// `sample_interval` controls the minimum time between successive CSV + /// samples; when `None`, the C++ default of 10 ms is used. + pub fn with_memory_tracking>( + csv_path: P, + sample_interval: Option, + ) -> Result { + let path_str = csv_path.as_ref().to_str().ok_or_else(|| { + Error::InvalidArgument(format!("csv_path is not valid UTF-8: {:?}", csv_path.as_ref())) + })?; + let c_path = CString::new(path_str).map_err(|e| { + Error::InvalidArgument(format!("csv_path contains an interior NUL byte: {}", e)) + })?; + let sample_interval_ms = + sample_interval.unwrap_or(Duration::from_millis(10)).as_millis() as i64; + let mut res: ffi::cuvsResources_t = 0; + unsafe { + check_cuvs(ffi::cuvsResourcesCreateWithMemoryTracking( + &mut res, + c_path.as_ptr(), + sample_interval_ms, + ))?; + } + Ok(Resources(res)) + } + /// Sets the current cuda stream pub fn set_cuda_stream(&self, stream: ffi::cudaStream_t) -> Result<()> { unsafe { check_cuvs(ffi::cuvsStreamSet(self.0, stream)) } @@ -61,4 +98,19 @@ mod tests { fn test_resources_create() { let _ = Resources::new(); } + + #[test] + fn test_resources_with_memory_tracking() { + let dir = tempfile::tempdir().unwrap(); + let csv = dir.path().join("alloc.csv"); + { + let _r = Resources::with_memory_tracking(&csv, Some(Duration::from_millis(2))) + .expect("with_memory_tracking should succeed"); + // closing _r at end of scope flushes the CSV reporter and + // restores the global host/device memory resources. + } + let meta = std::fs::metadata(&csv).expect("csv file should exist after drop"); + // at minimum, the header row should have been written before drop + assert!(meta.len() > 0, "tracking csv should be non-empty (got {} bytes)", meta.len()); + } } From ffb392192b2d9d391cb7915241709799092a26d7 Mon Sep 17 00:00:00 2001 From: achirkin Date: Tue, 12 May 2026 13:30:03 +0200 Subject: [PATCH 2/3] Address coderabbitai comments --- c/src/core/c_api.cpp | 3 +++ .../main/java/com/nvidia/cuvs/spi/CuVSProvider.java | 13 +++++++++++-- 2 files changed, 14 insertions(+), 2 deletions(-) diff --git a/c/src/core/c_api.cpp b/c/src/core/c_api.cpp index 026c6a3553..40d35a4af2 100644 --- a/c/src/core/c_api.cpp +++ b/c/src/core/c_api.cpp @@ -47,6 +47,9 @@ extern "C" cuvsError_t cuvsResourcesCreateWithMemoryTracking(cuvsResources_t* re 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(res_ptr); diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java index ed89786308..0d52e06b45 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java @@ -40,6 +40,12 @@ default Path nativeLibraryPath() { * Creates a new CuVSResources whose memory allocations are tracked and * written as CSV samples from a background thread. * + *

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 @@ -47,10 +53,13 @@ default Path nativeLibraryPath() { * @param memoryTrackingSampleInterval minimum interval between successive * CSV samples */ - CuVSResources newCuVSResources( + default CuVSResources newCuVSResources( Path tempDirectory, Path memoryTrackingCsvPath, - Duration memoryTrackingSampleInterval) throws Throwable; + 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 newHostMatrixBuilder( From 364e61ac16c88e657046588381c64e82a2a12e7c Mon Sep 17 00:00:00 2001 From: "Artem M. Chirkin" <9253178+achirkin@users.noreply.github.com> Date: Wed, 13 May 2026 07:00:37 +0200 Subject: [PATCH 3/3] Revert api_basics.rst --- docs/source/api_basics.rst | 105 ------------------------------------- 1 file changed, 105 deletions(-) diff --git a/docs/source/api_basics.rst b/docs/source/api_basics.rst index f99b10a23c..5ffb1da630 100644 --- a/docs/source/api_basics.rst +++ b/docs/source/api_basics.rst @@ -3,7 +3,6 @@ cuVS API Basics - `Memory management`_ - `Resource management`_ -- `Memory tracking`_ Memory management ----------------- @@ -89,107 +88,3 @@ Rust .. code-block:: rust let res = cuvs::Resources::new()?; - - -Memory tracking ---------------- - -A resources handle whose memory allocations are tracked and written as CSV -samples from a background thread can be created in any of the supported -languages. The 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 regular resources handle -and can be passed to every cuVS API that accepts one. The CSV reporter is -stopped and the global memory resources are restored when the handle is -destroyed. - -.. note:: - - - The handle replaces the **global** host and device memory resources while - it is alive. Do not create multiple tracking handles concurrently and make - sure the handle outlives every consumer (matrices, indexes, search results, - ...) that allocates memory through cuVS. - - The CSV file is flushed eagerly: the header is flushed on construction and - every sample row is flushed as soon as it is written, so the file can be - tailed while the handle is alive. Destroying the handle stops the - background sampler and writes one final row. - - The sample interval is a *minimum* time between samples. The background - thread blocks until an allocation/deallocation occurs, then sleeps for at - least ``sample_interval`` before writing the next row; quiescent periods do - not produce extra rows. - -C -^ - -.. code-block:: c - - #include - #include - - cuvsResources_t res; - // 10 ms sampling matches the C++ default. - cuvsResourcesCreateWithMemoryTracking(&res, "/tmp/allocations.csv", 10); - - // ... do some processing ... - - cuvsResourcesDestroy(res); - -C++ -^^^ - -.. code-block:: c++ - - #include - - // Sample interval defaults to std::chrono::milliseconds{10}. - raft::memory_tracking_resources res{"/tmp/allocations.csv"}; - - // ... do some processing ... - // `res` is implicitly convertible to raft::resources& and can be passed - // to any cuVS / raft API that accepts a resources handle. - -Python -^^^^^^ - -.. code-block:: python - - from cuvs.common import Resources - - res = Resources( - memory_tracking_csv_path="/tmp/allocations.csv", - memory_tracking_sample_interval_ms=10, - ) - - # ... do some processing ... - - del res # flushes the CSV and restores the global memory resources - -Java -^^^^ - -.. code-block:: java - - import com.nvidia.cuvs.CuVSResources; - import com.nvidia.cuvs.spi.CuVSProvider; - import java.nio.file.Path; - import java.time.Duration; - - try (var res = CuVSResources.create( - CuVSProvider.tempDirectory(), - Path.of("/tmp/allocations.csv"), - Duration.ofMillis(10))) { - // ... do some processing ... - } - -Rust -^^^^ - -.. code-block:: rust - - use std::time::Duration; - - let res = cuvs::Resources::with_memory_tracking( - "/tmp/allocations.csv", - Some(Duration::from_millis(10)), - )?;