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..44bf39249e 100644 --- a/c/src/core/c_api.cpp +++ b/c/src/core/c_api.cpp @@ -7,6 +7,7 @@ #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,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(res_ptr); + }); +} + extern "C" cuvsError_t cuvsResourcesDestroy(cuvsResources_t res) { return cuvs::core::translate_exceptions([=] { diff --git a/fern/pages/api_basics.md b/fern/pages/api_basics.md index f7889ac7bd..a194acfb4b 100644 --- a/fern/pages/api_basics.md +++ b/fern/pages/api_basics.md @@ -2,6 +2,7 @@ - [Memory management](#memory-management) - [Resource management](#resource-management) +- [Memory tracking](#memory-tracking) ## Memory management @@ -78,3 +79,85 @@ res = pylibraft.common.DeviceResources() ```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. + + + +- 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 + +```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++ + +```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 + +```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 + +```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 + +```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..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 @@ -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. + * + *

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 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..685817fa31 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,50 @@ 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()}. + * + *

Note: the ~8MB pinned host buffer backing this handle is allocated via a + * raw {@code cudaMallocHost} (see {@code PinnedMemoryBuffer.createPinnedBuffer()}) + * outside the tracking infrastructure, so it is not reflected in the CSV + * samples. + * + * @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 16a02c61b0..349ff8a0e8 100644 --- a/rust/cuvs/Cargo.toml +++ b/rust/cuvs/Cargo.toml @@ -18,6 +18,7 @@ ndarray = "0.15" [dev-dependencies] ndarray-rand = "0.14" +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..9537833de8 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,38 @@ 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: impl AsRef, + sample_interval: Option, + ) -> Result { + let c_path = + CString::new(csv_path.as_ref().as_os_str().as_encoded_bytes()).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 +96,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()); + } }