diff --git a/cpp/include/cuvs/neighbors/cagra.hpp b/cpp/include/cuvs/neighbors/cagra.hpp index 43ae7a6235..2c3e5f0ef9 100644 --- a/cpp/include/cuvs/neighbors/cagra.hpp +++ b/cpp/include/cuvs/neighbors/cagra.hpp @@ -74,7 +74,10 @@ struct ace_params { * * Used when `use_disk` is true or when the graph does not fit in host and GPU * memory. This should be the fastest disk in the system and hold enough space - * for twice the dataset, final graph, and label mapping. + * for twice the dataset, final graph, and label mapping. The directory may + * already exist, but ACE's named artifacts must not already exist. Simultaneous + * builds must use different directories. On failure, ACE removes only artifacts + * it created and never deletes unrelated directory contents. */ std::string build_dir = "/tmp/ace_build"; /** diff --git a/cpp/include/cuvs/util/file_io.hpp b/cpp/include/cuvs/util/file_io.hpp index a7d67ec2c0..7c90d155b0 100644 --- a/cpp/include/cuvs/util/file_io.hpp +++ b/cpp/include/cuvs/util/file_io.hpp @@ -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 @@ -180,14 +180,17 @@ class file_descriptor { * @tparam T Data type for the numpy array * @param path File path to create * @param shape Shape of the numpy array (e.g., {rows, cols} for 2D) + * @param exclusive Fail when the file already exists instead of truncating it * @return Pair of (file_descriptor, header_size) */ template std::pair create_numpy_file(const std::string& path, - const std::vector& shape) + const std::vector& shape, + bool exclusive = false) { // Open file - file_descriptor fd(path, O_CREAT | O_RDWR | O_TRUNC, 0644); + const int flags = O_CREAT | O_RDWR | (exclusive ? O_EXCL : O_TRUNC); + file_descriptor fd(path, flags, 0644); // Build header const auto dtype = raft::numpy_serializer::get_numpy_dtype(); diff --git a/cpp/src/neighbors/detail/cagra/cagra_build.cuh b/cpp/src/neighbors/detail/cagra/cagra_build.cuh index c06f9b12e3..0a3c2b98cc 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_build.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_build.cuh @@ -36,8 +36,11 @@ #include +#include +#include #include #include +#include #include #include #include @@ -55,6 +58,94 @@ namespace cuvs::neighbors::cagra::detail { constexpr double to_mib(size_t bytes) { return static_cast(bytes) / (1 << 20); } constexpr double to_gib(size_t bytes) { return static_cast(bytes) / (1 << 30); } +class ace_disk_workspace { + public: + enum class artifact : size_t { + reordered_dataset, + augmented_dataset, + dataset_mapping, + cagra_graph, + }; + + explicit ace_disk_workspace(std::string build_dir) + : build_dir_(std::move(build_dir)), + artifacts_{build_dir_ / "reordered_dataset.npy", + build_dir_ / "augmented_dataset.npy", + build_dir_ / "dataset_mapping.npy", + build_dir_ / "cagra_graph.npy"} + { + } + + void initialize() + { + if (mkdir(build_dir_.c_str(), 0755) == 0) { + directory_created_by_this_build_ = true; + return; + } + + if (errno != EEXIST) { + RAFT_FAIL("Failed to create ACE build directory: %s (errno: %d, %s)", + build_dir_.c_str(), + errno, + strerror(errno)); + } + + std::error_code error; + const bool is_directory = std::filesystem::is_directory(build_dir_, error); + RAFT_EXPECTS(!error, + "Failed to inspect ACE build directory: %s (%s)", + build_dir_.c_str(), + error.message().c_str()); + RAFT_EXPECTS(is_directory, "ACE build path is not a directory: %s", build_dir_.c_str()); + } + + [[nodiscard]] std::string artifact_path(artifact which) const + { + return artifacts_[static_cast(which)].string(); + } + + void mark_artifact_created(artifact which) noexcept + { + artifacts_created_[static_cast(which)] = true; + } + + void commit() noexcept { committed_ = true; } + + void cleanup() noexcept + { + if (committed_) { return; } + + for (size_t i = artifacts_.size(); i > 0; --i) { + if (!artifacts_created_[i - 1]) { continue; } + + std::error_code error; + std::filesystem::remove(artifacts_[i - 1], error); + if (error) { + RAFT_LOG_WARN("ACE: Failed to remove build artifact %s: %s", + artifacts_[i - 1].c_str(), + error.message().c_str()); + } + } + + if (directory_created_by_this_build_) { + std::error_code error; + std::filesystem::remove(build_dir_, error); + if (error) { + RAFT_LOG_WARN("ACE: Failed to remove empty build directory %s: %s", + build_dir_.c_str(), + error.message().c_str()); + } + } + } + + private: + std::filesystem::path build_dir_; + std::array artifacts_; + std::array artifacts_created_{}; + bool directory_created_by_this_build_ = false; + bool committed_ = false; +}; + template void check_graph_degree(size_t& intermediate_degree, size_t& graph_degree, size_t dataset_size) { @@ -826,7 +917,7 @@ constexpr double vector_expansion_factor = 2.0; template bool ace_check_use_disk_mode(raft::resources const& res, bool use_disk, - std::string& build_dir, + const std::string& build_dir, size_t dataset_size, size_t dataset_dim, size_t n_partitions, @@ -924,19 +1015,7 @@ bool ace_check_use_disk_mode(raft::resources const& res, to_gib(mem.available_gpu_memory)); bool use_disk_mode = use_disk || host_memory_limited || gpu_memory_limited; - if (use_disk_mode) { - bool valid_build_dir = !build_dir.empty(); - valid_build_dir &= build_dir.length() <= 255; - valid_build_dir &= build_dir.find('\0') == std::string::npos; - valid_build_dir &= build_dir.find("//") == std::string::npos; - if (!valid_build_dir) { - RAFT_LOG_WARN("ACE: Invalid build_dir path, resetting to default: /tmp/ace_build"); - build_dir = "/tmp/ace_build"; - } - if (mkdir(build_dir.c_str(), 0755) != 0 && errno != EEXIST) { - RAFT_EXPECTS(false, "Failed to create ACE build directory: %s", build_dir.c_str()); - } - } + if (use_disk_mode) { RAFT_EXPECTS(!build_dir.empty(), "ACE build directory must not be empty"); } if (host_memory_limited && gpu_memory_limited) { RAFT_LOG_INFO( @@ -1192,8 +1271,7 @@ auto build_ace(raft::resources const& res, size_t intermediate_degree = params.intermediate_graph_degree; size_t graph_degree = params.graph_degree; - // Track whether to clean up build directory on failure - bool cleanup_on_failure = false; + ace_disk_workspace workspace(build_dir); try { check_graph_degree(intermediate_degree, graph_degree, dataset_size); @@ -1236,24 +1314,32 @@ auto build_ace(raft::resources const& res, size_t graph_header_size = 0; if (use_disk_mode) { - if (mkdir(build_dir.c_str(), 0755) != 0 && errno != EEXIST) { - RAFT_EXPECTS(false, "Failed to create ACE build directory: %s", build_dir.c_str()); - } - // Mark for cleanup if we fail after creating the directory - cleanup_on_failure = true; + workspace.initialize(); // Create numpy files with pre-allocated space std::tie(reordered_fd, reordered_header_size) = cuvs::util::create_numpy_file( - build_dir + "/reordered_dataset.npy", {dataset_size, dataset_dim}); + workspace.artifact_path(ace_disk_workspace::artifact::reordered_dataset), + {dataset_size, dataset_dim}, + true); + workspace.mark_artifact_created(ace_disk_workspace::artifact::reordered_dataset); std::tie(augmented_fd, augmented_header_size) = cuvs::util::create_numpy_file( - build_dir + "/augmented_dataset.npy", {dataset_size, dataset_dim}); + workspace.artifact_path(ace_disk_workspace::artifact::augmented_dataset), + {dataset_size, dataset_dim}, + true); + workspace.mark_artifact_created(ace_disk_workspace::artifact::augmented_dataset); - std::tie(mapping_fd, mapping_header_size) = - cuvs::util::create_numpy_file(build_dir + "/dataset_mapping.npy", {dataset_size}); + std::tie(mapping_fd, mapping_header_size) = cuvs::util::create_numpy_file( + workspace.artifact_path(ace_disk_workspace::artifact::dataset_mapping), + {dataset_size}, + true); + workspace.mark_artifact_created(ace_disk_workspace::artifact::dataset_mapping); std::tie(graph_fd, graph_header_size) = cuvs::util::create_numpy_file( - build_dir + "/cagra_graph.npy", {dataset_size, graph_degree}); + workspace.artifact_path(ace_disk_workspace::artifact::cagra_graph), + {dataset_size, graph_degree}, + true); + workspace.mark_artifact_created(ace_disk_workspace::artifact::cagra_graph); RAFT_LOG_DEBUG( "ACE: Wrote numpy headers (reordered: %zu, augmented: %zu, mapping: %zu, graph: %zu bytes)", @@ -1535,20 +1621,14 @@ auto build_ace(raft::resources const& res, std::chrono::duration_cast(total_end - total_start).count(); RAFT_LOG_INFO("ACE: Partitioned CAGRA build completed in %ld ms total", total_elapsed); + workspace.commit(); return std::move(idx); } catch (const std::exception& e) { - // Clean up build directory on failure if we created it RAFT_LOG_ERROR("ACE: Build failed with exception: %s", e.what()); - if (cleanup_on_failure && !build_dir.empty()) { - RAFT_LOG_INFO("ACE: Cleaning up build directory: %s", build_dir.c_str()); - try { - std::filesystem::remove_all(build_dir); - RAFT_LOG_INFO("ACE: Successfully removed build directory"); - } catch (const std::exception& cleanup_error) { - RAFT_LOG_WARN("ACE: Failed to clean up build directory: %s", cleanup_error.what()); - } - } - // Re-throw the original exception + workspace.cleanup(); + throw; + } catch (...) { + workspace.cleanup(); throw; } } diff --git a/cpp/tests/neighbors/ann_hnsw_ace.cuh b/cpp/tests/neighbors/ann_hnsw_ace.cuh index f71c577762..f65cffd57e 100644 --- a/cpp/tests/neighbors/ann_hnsw_ace.cuh +++ b/cpp/tests/neighbors/ann_hnsw_ace.cuh @@ -10,8 +10,16 @@ #include +#include +#include +#include #include +#include #include +#include +#include + +#include namespace cuvs::neighbors::hnsw { @@ -47,6 +55,119 @@ inline ::std::ostream& operator<<(::std::ostream& os, const AnnHnswAceInputs& p) return os; } +namespace test_detail { + +class ace_workspace_directory { + public: + ace_workspace_directory() + : path_{std::filesystem::temp_directory_path() / + ("cuvs_ace_workspace_" + std::to_string(getpid()) + "_" + + std::to_string(std::time(nullptr)) + "_" + + std::to_string(reinterpret_cast(this)) + "_" + + std::to_string(counter_++))} + { + std::filesystem::create_directories(path_); + } + + ~ace_workspace_directory() + { + std::error_code error; + std::filesystem::remove_all(path_, error); + } + + [[nodiscard]] const std::filesystem::path& path() const { return path_; } + + private: + std::filesystem::path path_; + static inline std::atomic counter_{0}; +}; + +template +void build_ace_with_workspace(raft::resources const& resources, + raft::host_matrix_view dataset, + const std::filesystem::path& workspace) +{ + cagra::index_params params; + params.intermediate_graph_degree = 32; + params.graph_degree = 16; + + auto ace_params = cagra::graph_build_params::ace_params(); + ace_params.npartitions = 2; + ace_params.ef_construction = 50; + ace_params.build_dir = workspace.string(); + ace_params.use_disk = true; + params.graph_build_params = ace_params; + + auto dataset_view = cuvs::neighbors::make_host_standard_dataset_view(dataset); + [[maybe_unused]] auto index = cagra::build(resources, params, dataset_view); +} + +template +raft::host_matrix make_workspace_test_dataset() +{ + auto dataset = raft::make_host_matrix(1000, 8); + std::fill_n(dataset.data_handle(), dataset.size(), DataT{}); + return dataset; +} + +} // namespace test_detail + +template +void test_ace_workspace_failure_preserves_caller_directory() +{ + raft::resources resources; + auto dataset = test_detail::make_workspace_test_dataset(); + test_detail::ace_workspace_directory workspace; + const auto sentinel = workspace.path() / "sentinel.txt"; + const auto existing_artifact = workspace.path() / "reordered_dataset.npy"; + + { + std::ofstream sentinel_file(sentinel); + ASSERT_TRUE(sentinel_file.is_open()); + sentinel_file << "keep me"; + } + ASSERT_TRUE(std::filesystem::create_directory(existing_artifact)); + + EXPECT_THROW(test_detail::build_ace_with_workspace( + resources, raft::make_const_mdspan(dataset.view()), workspace.path()), + raft::logic_error); + + EXPECT_TRUE(std::filesystem::is_directory(workspace.path())); + EXPECT_TRUE(std::filesystem::is_directory(existing_artifact)); + std::ifstream sentinel_file(sentinel); + std::string sentinel_contents; + std::getline(sentinel_file, sentinel_contents); + EXPECT_EQ(sentinel_contents, "keep me"); +} + +template +void test_ace_workspace_failure_does_not_truncate_existing_artifact() +{ + raft::resources resources; + auto dataset = test_detail::make_workspace_test_dataset(); + test_detail::ace_workspace_directory workspace; + const auto existing_graph = workspace.path() / "cagra_graph.npy"; + constexpr const char* expected_contents = "preexisting-graph-contents"; + + { + std::ofstream graph_file(existing_graph, std::ios::binary); + ASSERT_TRUE(graph_file.is_open()); + graph_file << expected_contents; + } + + EXPECT_THROW(test_detail::build_ace_with_workspace( + resources, raft::make_const_mdspan(dataset.view()), workspace.path()), + raft::logic_error); + + std::ifstream graph_file(existing_graph, std::ios::binary); + std::string graph_contents; + graph_file >> graph_contents; + EXPECT_EQ(graph_contents, expected_contents); + EXPECT_FALSE(std::filesystem::exists(workspace.path() / "reordered_dataset.npy")); + EXPECT_FALSE(std::filesystem::exists(workspace.path() / "augmented_dataset.npy")); + EXPECT_FALSE(std::filesystem::exists(workspace.path() / "dataset_mapping.npy")); +} + template class AnnHnswAceTest : public ::testing::TestWithParam { public: 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..a3dd6959d6 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 */ @@ -7,6 +7,16 @@ namespace cuvs::neighbors::hnsw { +TEST(CagraAceWorkspace, FailurePreservesCallerDirectory) +{ + test_ace_workspace_failure_preserves_caller_directory(); +} + +TEST(CagraAceWorkspace, FailureDoesNotTruncateExistingArtifact) +{ + test_ace_workspace_failure_does_not_truncate_existing_artifact(); +} + typedef AnnHnswAceTest AnnHnswAceTest_float; TEST_P(AnnHnswAceTest_float, AnnHnswAceBuild) { this->testHnswAceBuild(); } 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..37efc9ebb0 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 */ @@ -7,6 +7,16 @@ namespace cuvs::neighbors::hnsw { +TEST(CagraAceWorkspaceHalf, FailurePreservesCallerDirectory) +{ + test_ace_workspace_failure_preserves_caller_directory(); +} + +TEST(CagraAceWorkspaceHalf, FailureDoesNotTruncateExistingArtifact) +{ + test_ace_workspace_failure_does_not_truncate_existing_artifact(); +} + typedef AnnHnswAceTest AnnHnswAceTest_half; TEST_P(AnnHnswAceTest_half, AnnHnswAceBuild) { this->testHnswAceBuild(); } 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..a437a3f7cb 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 */ @@ -7,6 +7,16 @@ namespace cuvs::neighbors::hnsw { +TEST(CagraAceWorkspaceInt8, FailurePreservesCallerDirectory) +{ + test_ace_workspace_failure_preserves_caller_directory(); +} + +TEST(CagraAceWorkspaceInt8, FailureDoesNotTruncateExistingArtifact) +{ + test_ace_workspace_failure_does_not_truncate_existing_artifact(); +} + typedef AnnHnswAceTest AnnHnswAceTest_int8_t; TEST_P(AnnHnswAceTest_int8_t, AnnHnswAceBuild) { this->testHnswAceBuild(); } 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..2b512e4291 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 */ @@ -7,6 +7,16 @@ namespace cuvs::neighbors::hnsw { +TEST(CagraAceWorkspaceUint8, FailurePreservesCallerDirectory) +{ + test_ace_workspace_failure_preserves_caller_directory(); +} + +TEST(CagraAceWorkspaceUint8, FailureDoesNotTruncateExistingArtifact) +{ + test_ace_workspace_failure_does_not_truncate_existing_artifact(); +} + typedef AnnHnswAceTest AnnHnswAceTest_uint8_t; TEST_P(AnnHnswAceTest_uint8_t, AnnHnswAceBuild) { this->testHnswAceBuild(); } diff --git a/python/cuvs/cuvs/neighbors/cagra/cagra.pyx b/python/cuvs/cuvs/neighbors/cagra/cagra.pyx index dd481df259..d10895bcf8 100644 --- a/python/cuvs/cuvs/neighbors/cagra/cagra.pyx +++ b/python/cuvs/cuvs/neighbors/cagra/cagra.pyx @@ -91,7 +91,10 @@ cdef class AceParams: graph). Used when `use_disk` is true or when the graph does not fit in host and GPU memory. This should be the fastest disk in the system and hold enough space for twice the dataset, final graph, and label - mapping. + mapping. The directory may already exist, but ACE's named artifacts + must not already exist. Simultaneous builds must use different + directories. On failure, ACE removes only artifacts it created and + never deletes unrelated directory contents. use_disk : bool, default = False Whether to use disk-based storage for ACE build. When true, enables disk-based operations for memory-efficient graph construction. diff --git a/python/cuvs/cuvs/tests/test_cagra_ace.py b/python/cuvs/cuvs/tests/test_cagra_ace.py index a2323e3baa..5fd0eedb8b 100644 --- a/python/cuvs/cuvs/tests/test_cagra_ace.py +++ b/python/cuvs/cuvs/tests/test_cagra_ace.py @@ -4,6 +4,7 @@ import os import tempfile +from pathlib import Path import cupy as cp import numpy as np @@ -13,6 +14,7 @@ from sklearn.preprocessing import normalize from cuvs.common import make_device_padded_dataset +from cuvs.common.exceptions import CuvsException from cuvs.neighbors import cagra, hnsw from cuvs.tests.ann_utils import ( calc_recall, @@ -158,6 +160,57 @@ def run_cagra_ace_build_search_test( assert recall > 0.7 +def _build_ace_with_disk_workspace(dataset, build_dir): + ace_params = cagra.AceParams( + npartitions=2, + ef_construction=50, + build_dir=str(build_dir), + use_disk=True, + ) + build_params = cagra.IndexParams( + intermediate_graph_degree=32, + graph_degree=16, + build_algo="ace", + ace_params=ace_params, + ) + cagra.build(build_params, dataset) + + +def test_cagra_ace_workspace_failure_preserves_caller_directory(): + """ACE failures must not delete unrelated contents in an existing workspace.""" + dataset = np.zeros((1000, 8), dtype=np.float32) + + with tempfile.TemporaryDirectory() as temp_dir: + workspace = Path(temp_dir) + sentinel = workspace / "sentinel.txt" + sentinel.write_text("keep me") + preexisting_artifact = workspace / "reordered_dataset.npy" + preexisting_artifact.mkdir() + + with pytest.raises(CuvsException): + _build_ace_with_disk_workspace(dataset, workspace) + + assert workspace.is_dir() + assert sentinel.read_text() == "keep me" + assert preexisting_artifact.is_dir() + + +def test_cagra_ace_workspace_failure_does_not_truncate_existing_artifact(): + """ACE must fail safely when a named artifact is already present.""" + dataset = np.zeros((1000, 8), dtype=np.float32) + + with tempfile.TemporaryDirectory() as temp_dir: + workspace = Path(temp_dir) + existing_graph = workspace / "cagra_graph.npy" + expected_contents = b"preexisting graph contents" + existing_graph.write_bytes(expected_contents) + + with pytest.raises(CuvsException): + _build_ace_with_disk_workspace(dataset, workspace) + + assert existing_graph.read_bytes() == expected_contents + + @pytest.mark.parametrize("dtype", [np.float32, np.float16, np.int8, np.uint8]) @pytest.mark.parametrize("metric", ["sqeuclidean", "inner_product"]) @pytest.mark.parametrize("use_disk", [False, True])