Automatic Partition Count Derivation for ACE

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;
 };