Automatic Partition Count Derivation for ACE (rapidsai#1603)

This PR adds automatic partition count derivation for ACE (Augmented Core Extraction) graph builds based on available system memory.

Previously, users had to manually calculate and specify the number of partitions based on their dataset size and available memory. This was error-prone and required understanding the internal memory requirements of the ACE algorithm. With auto-derivation, users get optimal partitioning out of the box while still having the option to override if needed.

**Changes**
- When `npartitions` is set to `0` (new default), ACE automatically calculates the optimal number of partitions based on available host and GPU memory.
- When `npartitions` is set to a positive value, the specified count is used but may be automatically increased if it would exceed memory limits.
- Added `max_host_memory_gb` and `max_gpu_memory_gb` parameters to allow users to constrain memory usage (useful for shared systems or testing).

This builds on top of PR rapidsai#1597, which should be merged first.

Authors:
  - Julian Miller (https://github.com/julianmi)
  - Tamas Bela Feher (https://github.com/tfeher)

Approvers:
  - Kyle Edwards (https://github.com/KyleFromNVIDIA)
  - Tamas Bela Feher (https://github.com/tfeher)
  - Ben Frederickson (https://github.com/benfred)

URL: rapidsai#1603

Author: julianmi

c/include/cuvs/neighbors/cagra.h

@@ -137,13 +137,20 @@ struct cuvsAceParams {
   /**
    * Number of partitions for ACE (Augmented Core Extraction) partitioned build.
    *
+   * When set to 0 (default), the number of partitions is automatically derived
+   * based on available host and GPU memory to maximize partition size while
+   * ensuring the build fits in memory.
+   *
    * Small values might improve recall but potentially degrade performance and
    * increase memory usage. Partitions should not be too small to prevent issues
-   * in KNN graph construction. 100k - 5M vectors per partition is recommended
-   * depending on the available host and GPU memory. The partition size is on
-   * average 2 * (n_rows / npartitions) * dim * sizeof(T). 2 is because of the
-   * core and augmented vectors. Please account for imbalance in the partition
-   * sizes (up to 3x in our tests).
+   * in KNN graph construction. The partition size is on average 2 * (n_rows /
+   * npartitions) * dim * sizeof(T). 2 is because of the core and augmented
+   * vectors. Please account for imbalance in the partition sizes (up to 3x in
+   * our tests).
+   *
+   * If the specified number of partitions results in partitions that exceed
+   * available memory, the value will be automatically increased to fit memory
+   * constraints and a warning will be issued.
    */
   size_t npartitions;
   /**
@@ -167,6 +174,22 @@ struct cuvsAceParams {
    * When true, enables disk-based operations for memory-efficient graph construction.
    */
   bool use_disk;
+  /**
+   * Maximum host memory to use for ACE build in GiB.
+   *
+   * When set to 0 (default), uses available host memory.
+   * When set to a positive value, limits host memory usage to the specified amount.
+   * Useful for testing or when running alongside other memory-intensive processes.
+   */
+  double max_host_memory_gb;
+  /**
+   * Maximum GPU memory to use for ACE build in GiB.
+   *
+   * When set to 0 (default), uses available GPU memory.
+   * When set to a positive value, limits GPU memory usage to the specified amount.
+   * Useful for testing or when running alongside other memory-intensive processes.
+   */
+  double max_gpu_memory_gb;
 };
 
 typedef struct cuvsAceParams* cuvsAceParams_t;

c/include/cuvs/neighbors/hnsw.h

@@ -46,8 +46,20 @@ enum cuvsHnswHierarchy {
 struct cuvsHnswAceParams {
   /**
    * Number of partitions for ACE partitioned build.
+   *
+   * When set to 0 (default), the number of partitions is automatically derived
+   * based on available host and GPU memory to maximize partition size while
+   * ensuring the build fits in memory.
+   *
    * Small values might improve recall but potentially degrade performance and
-   * increase memory usage. 100k - 5M vectors per partition is recommended.
+   * increase memory usage. The partition size is on average 2 * (n_rows /
+   * npartitions) * dim * sizeof(T). 2 is because of the core and augmented
+   * vectors. Please account for imbalance in the partition sizes (up to 3x in
+   * our tests).
+   *
+   * If the specified number of partitions results in partitions that exceed
+   * available memory, the value will be automatically increased to fit memory
+   * constraints and a warning will be issued.
    */
   size_t npartitions;
   /**
@@ -60,6 +72,18 @@ struct cuvsHnswAceParams {
    * When true, enables disk-based operations for memory-efficient graph construction.
    */
   bool use_disk;
+  /**
+   * Maximum host memory to use for ACE build in GiB.
+   * When set to 0 (default), uses available host memory.
+   * Useful for testing or when running alongside other memory-intensive processes.
+   */
+  double max_host_memory_gb;
+  /**
+   * Maximum GPU memory to use for ACE build in GiB.
+   * When set to 0 (default), uses available GPU memory.
+   * Useful for testing or when running alongside other memory-intensive processes.
+   */
+  double max_gpu_memory_gb;
 };
 
 typedef struct cuvsHnswAceParams* cuvsHnswAceParams_t;

c/src/neighbors/cagra.cpp

@@ -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
  */
 
@@ -88,10 +88,12 @@ static void _set_graph_build_params(
       cuvs::neighbors::cagra::graph_build_params::ace_params ace_p;
       if (params.graph_build_params) {
         auto ace_params_c             = static_cast<cuvsAceParams*>(params.graph_build_params);
-        ace_p.npartitions     = ace_params_c->npartitions;
-        ace_p.ef_construction = ace_params_c->ef_construction;
-        ace_p.build_dir       = std::string(ace_params_c->build_dir);
-        ace_p.use_disk        = ace_params_c->use_disk;
+        ace_p.npartitions         = ace_params_c->npartitions;
+        ace_p.ef_construction     = ace_params_c->ef_construction;
+        ace_p.build_dir           = std::string(ace_params_c->build_dir);
+        ace_p.use_disk            = ace_params_c->use_disk;
+        ace_p.max_host_memory_gb  = ace_params_c->max_host_memory_gb;
+        ace_p.max_gpu_memory_gb   = ace_params_c->max_gpu_memory_gb;
       }
       out_params = ace_p;
       break;
@@ -778,10 +780,12 @@ extern "C" cuvsError_t cuvsAceParamsCreate(cuvsAceParams_t* params)
     // Allocate and copy the build directory string
     const char* build_dir = strdup(ps.build_dir.c_str());
 
-    *params = new cuvsAceParams{.npartitions     = ps.npartitions,
-                                .ef_construction = ps.ef_construction,
-                                .build_dir       = build_dir,
-                                .use_disk        = ps.use_disk};
+    *params = new cuvsAceParams{.npartitions         = ps.npartitions,
+                                .ef_construction     = ps.ef_construction,
+                                .build_dir           = build_dir,
+                                .use_disk            = ps.use_disk,
+                                .max_host_memory_gb  = ps.max_host_memory_gb,
+                                .max_gpu_memory_gb   = ps.max_gpu_memory_gb};
   });
 }
 

c/src/neighbors/hnsw.cpp

@@ -43,6 +43,8 @@ void _build(cuvsResources_t res,
   ace_params.npartitions          = params->ace_params->npartitions;
   ace_params.build_dir            = params->ace_params->build_dir ? params->ace_params->build_dir : "/tmp/hnsw_ace_build";
   ace_params.use_disk             = params->ace_params->use_disk;
+  ace_params.max_host_memory_gb   = params->ace_params->max_host_memory_gb;
+  ace_params.max_gpu_memory_gb    = params->ace_params->max_gpu_memory_gb;
   cpp_params.graph_build_params   = ace_params;
 
   using dataset_mdspan_type = raft::host_matrix_view<T const, int64_t, raft::row_major>;
@@ -151,9 +153,11 @@ void* _deserialize(cuvsResources_t res,
 extern "C" cuvsError_t cuvsHnswAceParamsCreate(cuvsHnswAceParams_t* params)
 {
   return cuvs::core::translate_exceptions([=] {
-    *params = new cuvsHnswAceParams{.npartitions     = 1,
-                                    .build_dir       = "/tmp/hnsw_ace_build",
-                                    .use_disk        = false};
+    *params = new cuvsHnswAceParams{.npartitions         = 0,
+                                    .build_dir           = "/tmp/hnsw_ace_build",
+                                    .use_disk            = false,
+                                    .max_host_memory_gb  = 0,
+                                    .max_gpu_memory_gb   = 0};
   });
 }
 

cpp/include/cuvs/neighbors/graph_build_types.hpp

@@ -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
  */
 
@@ -99,15 +99,21 @@ struct ace_params {
   /**
    * Number of partitions for ACE (Augmented Core Extraction) partitioned build.
    *
+   * When set to 0 (default), the number of partitions is automatically derived
+   * based on available host and GPU memory to maximize partition size while
+   * ensuring the build fits in memory.
+   *
    * Small values might improve recall but potentially degrade performance and
    * increase memory usage. Partitions should not be too small to prevent issues
-   * in KNN graph construction. 100k - 5M vectors per partition is recommended
-   * depending on the available host and GPU memory. The partition size is on
-   * average 2 * (n_rows / npartitions) * dim * sizeof(T). 2 is because of the
-   * core and augmented vectors. Please account for imbalance in the partition
-   * sizes (up to 3x in our tests).
+   * in KNN graph construction. The partition size is on average 2 * (n_rows / npartitions) * dim *
+   * sizeof(T). 2 is because of the core and augmented vectors. Please account for imbalance in the
+   * partition sizes (up to 3x in our tests).
+   *
+   * If the specified number of partitions results in partitions that exceed
+   * available memory, the value will be automatically increased to fit memory
+   * constraints and a warning will be issued.
    */
-  size_t npartitions = 1;
+  size_t npartitions = 0;
   /**
    * The index quality for the ACE build.
    *
@@ -129,6 +135,22 @@ struct ace_params {
    * When true, enables disk-based operations for memory-efficient graph construction.
    */
   bool use_disk = false;
+  /**
+   * Maximum host memory to use for ACE build in GiB.
+   *
+   * When set to 0 (default), uses available host memory.
+   * When set to a positive value, limits host memory usage to the specified amount.
+   * Useful for testing or when running alongside other memory-intensive processes.
+   */
+  double max_host_memory_gb = 0;
+  /**
+   * Maximum GPU memory to use for ACE build in GiB.
+   *
+   * When set to 0 (default), uses available GPU memory.
+   * When set to a positive value, limits GPU memory usage to the specified amount.
+   * Useful for testing or when running alongside other memory-intensive processes.
+   */
+  double max_gpu_memory_gb = 0;
 
   ace_params() = default;
 };