diff --git a/cpp/src/neighbors/detail/cagra/cagra_build.cuh b/cpp/src/neighbors/detail/cagra/cagra_build.cuh index a7c15b4161..4cdddc9fc3 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_build.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_build.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once @@ -959,6 +959,30 @@ bool ace_check_use_disk_mode(raft::resources const& res, return use_disk_mode; } +// Resolve the ACE partition count while preserving 0 as the auto-selection sentinel. +inline size_t ace_resolve_partition_count(size_t n_partitions) +{ + if (n_partitions == 0) { return 2; } + if (n_partitions == 1) { + RAFT_LOG_WARN( + "ACE: Requested 1 partition; adjusted to 2 before applying partitioning heuristics"); + return 2; + } + return n_partitions; +} + +// Validate the structural ACE partition-count invariants required by the labeler. +inline void ace_validate_partition_count(size_t n_partitions, + size_t dataset_size, + bool adjusted_for_memory = false) +{ + RAFT_EXPECTS(n_partitions <= dataset_size, + adjusted_for_memory + ? "ACE: configured memory limit is unsatisfiable because the requested partition " + "count cannot exceed dataset size" + : "ACE: number of partitions cannot exceed dataset size"); +} + // Validate and adjust partitions for disk mode memory requirements template void ace_validate_disk_mode_partitions(raft::resources const& res, @@ -1154,11 +1178,9 @@ index build_ace(raft::resources const& res, "ACE: Intermediate graph degree must be greater than 0"); RAFT_EXPECTS(params.graph_degree > 0, "ACE: Graph degree must be greater than 0"); - size_t n_partitions = npartitions; - if (n_partitions == 0) { - // Default: start with 2 partitions and increase if needed (minimum for ACE to make sense). - n_partitions = 2; - } + size_t n_partitions = ace_resolve_partition_count(npartitions); + + ace_validate_partition_count(n_partitions, dataset_size); size_t min_required_per_partition = 1000; if (n_partitions > dataset_size / min_required_per_partition) { @@ -1211,6 +1233,7 @@ index build_ace(raft::resources const& res, graph_degree, params.guarantee_connectivity, mem); + ace_validate_partition_count(n_partitions, dataset_size, true); } // Preallocate space for files for better performance and fail early if not enough space. diff --git a/cpp/tests/neighbors/ann_hnsw_ace.cuh b/cpp/tests/neighbors/ann_hnsw_ace.cuh index c75b3555f6..4637c4f6ad 100644 --- a/cpp/tests/neighbors/ann_hnsw_ace.cuh +++ b/cpp/tests/neighbors/ann_hnsw_ace.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once @@ -262,6 +262,31 @@ class AnnHnswAceTest : public ::testing::TestWithParam { std::filesystem::remove_all(temp_dir); } + void testHnswAceRejectsTooManyPartitions() + { + auto database_host = raft::make_host_matrix(ps.n_rows, ps.dim); + raft::copy(database_host.data_handle(), database_dev.data(), ps.n_rows * ps.dim, stream_); + raft::resource::sync_stream(handle_); + + hnsw::index_params hnsw_params; + hnsw_params.metric = ps.metric; + hnsw_params.hierarchy = hnsw::HnswHierarchy::GPU; + hnsw_params.M = 32; + hnsw_params.ef_construction = ps.ef_construction; + + auto ace_params = graph_build_params::ace_params(); + ace_params.npartitions = static_cast(ps.n_rows) + 1; + hnsw_params.graph_build_params = ace_params; + + try { + [[maybe_unused]] auto hnsw_index = + hnsw::build(handle_, hnsw_params, raft::make_const_mdspan(database_host.view())); + FAIL() << "ACE accepted more partitions than dataset rows"; + } catch (const std::exception& e) { + EXPECT_NE(std::string(e.what()).find("cannot exceed dataset size"), std::string::npos); + } + } + // Verify the in-memory CAGRA -> HNSW conversion spills to disk when the resulting HNSW // index would not fit in (an artificially constrained) host memory. This exercises // serialize_to_hnswlib_from_inmem and the batched serializer core, covering: @@ -472,7 +497,7 @@ inline std::vector generate_hnsw_ace_inputs() {5000}, // n_rows {64, 128}, // dim {10}, // k - {2, 4}, // npartitions + {0, 1, 2, 4}, // npartitions (auto, adjusted, and explicit) {100}, // ef_construction {false, true}, // use_disk (test both modes) {cuvs::distance::DistanceType::L2Expanded, @@ -502,6 +527,19 @@ inline std::vector generate_hnsw_ace_memory_fallback_inputs() }; } +inline std::vector generate_hnsw_ace_invalid_partition_inputs() +{ + return {{10, + 5000, + 64, + 10, + 0, // Set to n_rows + 1 by testHnswAceRejectsTooManyPartitions. + 100, + false, + cuvs::distance::DistanceType::L2Expanded, + 0.0}}; +} + // Inputs for testing the in-memory CAGRA -> HNSW disk-spill conversion path. inline std::vector generate_hnsw_inmem_spill_inputs() { @@ -524,6 +562,8 @@ inline std::vector generate_hnsw_inmem_spill_inputs() const std::vector hnsw_ace_inputs = generate_hnsw_ace_inputs(); const std::vector hnsw_ace_memory_fallback_inputs = generate_hnsw_ace_memory_fallback_inputs(); +const std::vector hnsw_ace_invalid_partition_inputs = + generate_hnsw_ace_invalid_partition_inputs(); const std::vector hnsw_inmem_spill_inputs = generate_hnsw_inmem_spill_inputs(); } // namespace cuvs::neighbors::hnsw diff --git a/cpp/tests/neighbors/ann_hnsw_ace/test_float_uint32_t.cu b/cpp/tests/neighbors/ann_hnsw_ace/test_float_uint32_t.cu index 4cde210d62..65563070fd 100644 --- a/cpp/tests/neighbors/ann_hnsw_ace/test_float_uint32_t.cu +++ b/cpp/tests/neighbors/ann_hnsw_ace/test_float_uint32_t.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -12,6 +12,16 @@ TEST_P(AnnHnswAceTest_float, AnnHnswAceBuild) { this->testHnswAceBuild(); } INSTANTIATE_TEST_CASE_P(AnnHnswAceTest, AnnHnswAceTest_float, ::testing::ValuesIn(hnsw_ace_inputs)); +typedef AnnHnswAceTest AnnHnswAceInvalidPartitionTest_float; +TEST_P(AnnHnswAceInvalidPartitionTest_float, RejectsTooManyPartitions) +{ + this->testHnswAceRejectsTooManyPartitions(); +} + +INSTANTIATE_TEST_CASE_P(AnnHnswAceInvalidPartitionTest, + AnnHnswAceInvalidPartitionTest_float, + ::testing::ValuesIn(hnsw_ace_invalid_partition_inputs)); + // Test for memory limit fallback to disk mode typedef AnnHnswAceTest AnnHnswAceMemoryFallbackTest_float; TEST_P(AnnHnswAceMemoryFallbackTest_float, AnnHnswAceMemoryLimitFallback) diff --git a/cpp/tests/neighbors/ann_hnsw_ace/test_half_uint32_t.cu b/cpp/tests/neighbors/ann_hnsw_ace/test_half_uint32_t.cu index d8664d4e14..744004a761 100644 --- a/cpp/tests/neighbors/ann_hnsw_ace/test_half_uint32_t.cu +++ b/cpp/tests/neighbors/ann_hnsw_ace/test_half_uint32_t.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -12,6 +12,16 @@ TEST_P(AnnHnswAceTest_half, AnnHnswAceBuild) { this->testHnswAceBuild(); } INSTANTIATE_TEST_CASE_P(AnnHnswAceTest, AnnHnswAceTest_half, ::testing::ValuesIn(hnsw_ace_inputs)); +typedef AnnHnswAceTest AnnHnswAceInvalidPartitionTest_half; +TEST_P(AnnHnswAceInvalidPartitionTest_half, RejectsTooManyPartitions) +{ + this->testHnswAceRejectsTooManyPartitions(); +} + +INSTANTIATE_TEST_CASE_P(AnnHnswAceInvalidPartitionTest, + AnnHnswAceInvalidPartitionTest_half, + ::testing::ValuesIn(hnsw_ace_invalid_partition_inputs)); + // Test for memory limit fallback to disk mode typedef AnnHnswAceTest AnnHnswAceMemoryFallbackTest_half; TEST_P(AnnHnswAceMemoryFallbackTest_half, AnnHnswAceMemoryLimitFallback) diff --git a/cpp/tests/neighbors/ann_hnsw_ace/test_int8_t_uint32_t.cu b/cpp/tests/neighbors/ann_hnsw_ace/test_int8_t_uint32_t.cu index 4c95192d8a..f967099679 100644 --- a/cpp/tests/neighbors/ann_hnsw_ace/test_int8_t_uint32_t.cu +++ b/cpp/tests/neighbors/ann_hnsw_ace/test_int8_t_uint32_t.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -14,6 +14,16 @@ INSTANTIATE_TEST_CASE_P(AnnHnswAceTest, AnnHnswAceTest_int8_t, ::testing::ValuesIn(hnsw_ace_inputs)); +typedef AnnHnswAceTest AnnHnswAceInvalidPartitionTest_int8_t; +TEST_P(AnnHnswAceInvalidPartitionTest_int8_t, RejectsTooManyPartitions) +{ + this->testHnswAceRejectsTooManyPartitions(); +} + +INSTANTIATE_TEST_CASE_P(AnnHnswAceInvalidPartitionTest, + AnnHnswAceInvalidPartitionTest_int8_t, + ::testing::ValuesIn(hnsw_ace_invalid_partition_inputs)); + // Test for memory limit fallback to disk mode typedef AnnHnswAceTest AnnHnswAceMemoryFallbackTest_int8_t; TEST_P(AnnHnswAceMemoryFallbackTest_int8_t, AnnHnswAceMemoryLimitFallback) diff --git a/cpp/tests/neighbors/ann_hnsw_ace/test_uint8_t_uint32_t.cu b/cpp/tests/neighbors/ann_hnsw_ace/test_uint8_t_uint32_t.cu index 3e4b91e759..80b75ca249 100644 --- a/cpp/tests/neighbors/ann_hnsw_ace/test_uint8_t_uint32_t.cu +++ b/cpp/tests/neighbors/ann_hnsw_ace/test_uint8_t_uint32_t.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -14,6 +14,16 @@ INSTANTIATE_TEST_CASE_P(AnnHnswAceTest, AnnHnswAceTest_uint8_t, ::testing::ValuesIn(hnsw_ace_inputs)); +typedef AnnHnswAceTest AnnHnswAceInvalidPartitionTest_uint8_t; +TEST_P(AnnHnswAceInvalidPartitionTest_uint8_t, RejectsTooManyPartitions) +{ + this->testHnswAceRejectsTooManyPartitions(); +} + +INSTANTIATE_TEST_CASE_P(AnnHnswAceInvalidPartitionTest, + AnnHnswAceInvalidPartitionTest_uint8_t, + ::testing::ValuesIn(hnsw_ace_invalid_partition_inputs)); + // Test for memory limit fallback to disk mode typedef AnnHnswAceTest AnnHnswAceMemoryFallbackTest_uint8_t; TEST_P(AnnHnswAceMemoryFallbackTest_uint8_t, AnnHnswAceMemoryLimitFallback) diff --git a/python/cuvs/cuvs/tests/test_cagra_ace.py b/python/cuvs/cuvs/tests/test_cagra_ace.py index c1633e3cad..439aec7f2b 100644 --- a/python/cuvs/cuvs/tests/test_cagra_ace.py +++ b/python/cuvs/cuvs/tests/test_cagra_ace.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # @@ -163,7 +163,7 @@ def test_cagra_ace_dtypes_and_metrics(dtype, metric, use_disk): ) -@pytest.mark.parametrize("npartitions", [2, 3, 8]) +@pytest.mark.parametrize("npartitions", [0, 1, 2, 3, 8]) def test_cagra_ace_partitions(npartitions): """Test ACE with different partition sizes (disk mode only).""" run_cagra_ace_build_search_test( @@ -172,6 +172,25 @@ def test_cagra_ace_partitions(npartitions): ) +@pytest.mark.parametrize( + ("npartitions", "message"), + [ + (33, "cannot exceed dataset size"), + ], +) +def test_cagra_ace_rejects_invalid_partition_count(npartitions, message): + """ACE rejects partition counts that would make partition labeling invalid.""" + dataset = np.zeros((32, 8), dtype=np.float32) + ace_params = cagra.AceParams(npartitions=npartitions, use_disk=False) + build_params = cagra.IndexParams( + build_algo="ace", + ace_params=ace_params, + ) + + with pytest.raises(Exception, match=message): + cagra.build(build_params, dataset) + + @pytest.mark.parametrize("ef_construction", [50, 100, 200]) def test_cagra_ace_ef_construction(ef_construction): """Test ACE with different ef_construction values (disk mode only)."""