diff --git a/c/include/cuvs/neighbors/cagra.h b/c/include/cuvs/neighbors/cagra.h index 57d063ef44..25555867b9 100644 --- a/c/include/cuvs/neighbors/cagra.h +++ b/c/include/cuvs/neighbors/cagra.h @@ -168,7 +168,10 @@ struct cuvsAceParams { * * 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. */ const char* build_dir; /** diff --git a/c/include/cuvs/neighbors/hnsw.h b/c/include/cuvs/neighbors/hnsw.h index 15eb1b0569..1e912fb59d 100644 --- a/c/include/cuvs/neighbors/hnsw.h +++ b/c/include/cuvs/neighbors/hnsw.h @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -66,7 +66,11 @@ struct cuvsHnswAceParams { size_t npartitions; /** * Directory to store ACE build artifacts (e.g., KNN graph, optimized graph). - * Used when `use_disk` is true or when the graph does not fit in memory. + * Used when `use_disk` is true or when the graph does not fit in memory. The + * directory may already exist, but ACE's named artifacts and `hnsw_index.bin` + * must not already exist. Simultaneous builds must use different directories. + * On failure, ACE removes only its uncommitted CAGRA artifacts; a completed + * CAGRA stage is retained if creating the HNSW index fails. */ const char* build_dir; /** diff --git a/cpp/include/cuvs/neighbors/cagra.hpp b/cpp/include/cuvs/neighbors/cagra.hpp index bfeafcfbe2..243a57f5bc 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/neighbors/hnsw.hpp b/cpp/include/cuvs/neighbors/hnsw.hpp index 39db5285f5..d5215637c2 100644 --- a/cpp/include/cuvs/neighbors/hnsw.hpp +++ b/cpp/include/cuvs/neighbors/hnsw.hpp @@ -79,6 +79,13 @@ struct index_params : cuvs::neighbors::index_params { * ace.use_disk = true; * ace.build_dir = "/tmp/hnsw_ace_build"; * @endcode + * + * When ACE writes to disk, `build_dir` may already exist, but ACE's named CAGRA + * artifacts and `hnsw_index.bin` must not already exist. Simultaneous builds + * must use different directories. The HNSW output is published only after it + * is fully serialized; however, the complete build is not transactional: if + * HNSW conversion fails after CAGRA succeeds, the completed CAGRA artifacts + * remain in the directory. */ std::variant graph_build_params; }; diff --git a/cpp/include/cuvs/util/file_io.hpp b/cpp/include/cuvs/util/file_io.hpp index a7d67ec2c0..38eec3f7e4 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,47 +180,58 @@ 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. + * If creation succeeds but pre-allocation or header writing + * fails, remove the newly created file. * @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); - - // Build header - const auto dtype = raft::numpy_serializer::get_numpy_dtype(); - const bool fortran_order = false; - const raft::numpy_serializer::header_t header = {dtype, fortran_order, shape}; - - std::stringstream ss; - raft::numpy_serializer::write_header(ss, header); - std::string header_str = ss.str(); - size_t header_size = header_str.size(); - - // Calculate data size from shape - size_t data_bytes = sizeof(T); - for (auto dim : shape) { - data_bytes *= dim; - } + const int flags = O_CREAT | O_RDWR | (exclusive ? O_EXCL : O_TRUNC); + file_descriptor fd(path, flags, 0644); + + try { + // Build header + const auto dtype = raft::numpy_serializer::get_numpy_dtype(); + const bool fortran_order = false; + const raft::numpy_serializer::header_t header = {dtype, fortran_order, shape}; + + std::stringstream ss; + raft::numpy_serializer::write_header(ss, header); + std::string header_str = ss.str(); + size_t header_size = header_str.size(); + + // Calculate data size from shape + size_t data_bytes = sizeof(T); + for (auto dim : shape) { + data_bytes *= dim; + } - // Pre-allocate file space - if (posix_fallocate(fd.get(), 0, header_size + data_bytes) != 0) { - RAFT_FAIL("Failed to pre-allocate space for file: %s", path.c_str()); - } + // Pre-allocate file space + if (posix_fallocate(fd.get(), 0, header_size + data_bytes) != 0) { + RAFT_FAIL("Failed to pre-allocate space for file: %s", path.c_str()); + } - // Seek to beginning and write header - if (lseek(fd.get(), 0, SEEK_SET) == -1) { - RAFT_FAIL("Failed to seek to beginning of file: %s", path.c_str()); - } + // Seek to beginning and write header + if (lseek(fd.get(), 0, SEEK_SET) == -1) { + RAFT_FAIL("Failed to seek to beginning of file: %s", path.c_str()); + } - ssize_t written = write(fd.get(), header_str.data(), header_str.size()); - if (written < 0 || static_cast(written) != header_str.size()) { - RAFT_FAIL("Failed to write numpy header to file: %s", path.c_str()); - } + ssize_t written = write(fd.get(), header_str.data(), header_str.size()); + if (written < 0 || static_cast(written) != header_str.size()) { + RAFT_FAIL("Failed to write numpy header to file: %s", path.c_str()); + } - return {std::move(fd), header_size}; + return {std::move(fd), header_size}; + } catch (...) { + fd.close(); + if (exclusive) { (void)::unlink(path.c_str()); } + throw; + } } /** diff --git a/cpp/src/neighbors/detail/cagra/cagra_build.cuh b/cpp/src/neighbors/detail/cagra/cagra_build.cuh index 8705926a41..715d4945a8 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) { @@ -836,7 +927,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, @@ -934,19 +1025,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( @@ -1208,8 +1287,7 @@ auto build_ace(raft::resources const& res, const index_params& params, DatasetVi 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); @@ -1252,24 +1330,32 @@ auto build_ace(raft::resources const& res, const index_params& params, DatasetVi 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)", @@ -1562,20 +1648,14 @@ auto build_ace(raft::resources const& res, const index_params& params, DatasetVi 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/src/neighbors/detail/hnsw.hpp b/cpp/src/neighbors/detail/hnsw.hpp index 8328fac5d5..f2f7542cec 100644 --- a/cpp/src/neighbors/detail/hnsw.hpp +++ b/cpp/src/neighbors/detail/hnsw.hpp @@ -26,7 +26,9 @@ #include +#include #include +#include #include #include #include @@ -41,6 +43,72 @@ namespace cuvs::neighbors::hnsw::detail { +class exclusive_hnsw_output_file { + public: + explicit exclusive_hnsw_output_file(std::filesystem::path output_path) + : output_path_{std::move(output_path)} + { + std::string temporary_path = output_path_.string() + ".tmp.XXXXXX"; + int fd = ::mkstemp(temporary_path.data()); + RAFT_EXPECTS(fd != -1, + "Cannot create temporary file for %s (errno: %d, %s)", + output_path_.c_str(), + errno, + strerror(errno)); + temporary_path_ = std::move(temporary_path); + + if (::close(fd) != 0) { + const int error = errno; + cleanup(); + RAFT_FAIL("Cannot close temporary file for %s (errno: %d, %s)", + output_path_.c_str(), + error, + strerror(error)); + } + + stream_.open(temporary_path_, std::ios::out | std::ios::binary | std::ios::trunc); + if (!stream_) { + cleanup(); + RAFT_FAIL("Cannot open temporary file for %s", output_path_.c_str()); + } + } + + exclusive_hnsw_output_file(const exclusive_hnsw_output_file&) = delete; + exclusive_hnsw_output_file& operator=(const exclusive_hnsw_output_file&) = delete; + + ~exclusive_hnsw_output_file() + { + stream_.close(); + cleanup(); + } + + std::ostream& stream() { return stream_; } + + void publish() + { + stream_.close(); + RAFT_EXPECTS(stream_, "Error writing output %s", output_path_.c_str()); + + if (::link(temporary_path_.c_str(), output_path_.c_str()) != 0) { + const int error = errno; + RAFT_FAIL("Cannot publish HNSW index %s (errno: %d, %s)", + output_path_.c_str(), + error, + strerror(error)); + } + } + + private: + void cleanup() noexcept + { + if (!temporary_path_.empty()) { (void)::unlink(temporary_path_.c_str()); } + } + + std::filesystem::path output_path_; + std::string temporary_path_; + std::ofstream stream_; +}; + template inline constexpr bool is_cagra_hnsw_export_index_v = std::is_same_v> || @@ -1298,15 +1366,10 @@ std::unique_ptr> from_cagra( index_directory.c_str()); std::string index_filename = (std::filesystem::path(index_directory) / "hnsw_index.bin").string(); + exclusive_hnsw_output_file output(index_filename); - std::ofstream of(index_filename, std::ios::out | std::ios::binary); - - RAFT_EXPECTS(of, "Cannot open file %s", index_filename.c_str()); - - serialize_to_hnswlib_from_disk(res, of, params, cagra_index); - - of.close(); - RAFT_EXPECTS(of, "Error writing output %s", index_filename.c_str()); + serialize_to_hnswlib_from_disk(res, output.stream(), params, cagra_index); + output.publish(); // Create an empty HNSW index that holds the file descriptor auto hnsw_index = @@ -1398,14 +1461,10 @@ std::unique_ptr> from_cagra( std::string index_filename = (std::filesystem::path(index_directory) / "hnsw_index.bin").string(); + exclusive_hnsw_output_file output(index_filename); - std::ofstream of(index_filename, std::ios::out | std::ios::binary); - RAFT_EXPECTS(of, "Cannot open file %s", index_filename.c_str()); - - serialize_to_hnswlib_from_inmem(res, of, params, cagra_index, dataset); - - of.close(); - RAFT_EXPECTS(of, "Error writing output %s", index_filename.c_str()); + serialize_to_hnswlib_from_inmem(res, output.stream(), params, cagra_index, dataset); + output.publish(); // Create an empty HNSW index that holds the file descriptor auto hnsw_index = diff --git a/cpp/tests/neighbors/ann_hnsw_ace.cuh b/cpp/tests/neighbors/ann_hnsw_ace.cuh index f71c577762..bb66d8a6c2 100644 --- a/cpp/tests/neighbors/ann_hnsw_ace.cuh +++ b/cpp/tests/neighbors/ann_hnsw_ace.cuh @@ -7,11 +7,21 @@ #include "ann_cagra.cuh" #include +#include #include +#include +#include +#include #include +#include #include +#include +#include +#include + +#include namespace cuvs::neighbors::hnsw { @@ -47,6 +57,170 @@ 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")); +} + +void test_exclusive_numpy_create_failure_removes_partial_file() +{ + test_detail::ace_workspace_directory workspace; + const auto path = workspace.path() / "partial.npy"; + const auto too_large = + static_cast(std::numeric_limits::max()) - static_cast(4096); + + EXPECT_THROW(cuvs::util::create_numpy_file(path.string(), {too_large}, true), + raft::logic_error); + + EXPECT_FALSE(std::filesystem::exists(path)); +} + +template +void test_hnsw_ace_build_does_not_truncate_existing_index() +{ + raft::resources resources; + auto dataset = test_detail::make_workspace_test_dataset(); + for (size_t i = 0; i < dataset.size(); ++i) { + dataset.data_handle()[i] = static_cast(i % 251); + } + test_detail::ace_workspace_directory workspace; + const auto index_path = workspace.path() / "hnsw_index.bin"; + constexpr const char* expected_contents = "preexisting-hnsw-index"; + + { + std::ofstream index_file(index_path, std::ios::binary); + ASSERT_TRUE(index_file.is_open()); + index_file << expected_contents; + } + + hnsw::index_params params; + params.M = 8; + params.ef_construction = 50; + auto ace_params = graph_build_params::ace_params(); + ace_params.npartitions = 2; + ace_params.ef_construction = 50; + ace_params.build_dir = workspace.path().string(); + ace_params.use_disk = true; + params.graph_build_params = ace_params; + + EXPECT_THROW(hnsw::build(resources, params, raft::make_const_mdspan(dataset.view())), + raft::logic_error); + + std::ifstream index_file(index_path, std::ios::binary); + std::string contents; + index_file >> contents; + EXPECT_EQ(contents, expected_contents); + EXPECT_TRUE(std::filesystem::exists(workspace.path() / "cagra_graph.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..832db29d91 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,26 @@ 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(); +} + +TEST(FileIo, ExclusiveNumpyCreateFailureRemovesPartialFile) +{ + test_exclusive_numpy_create_failure_removes_partial_file(); +} + +TEST(HnswAceWorkspace, ExistingIndexIsNotTruncated) +{ + test_hnsw_ace_build_does_not_truncate_existing_index(); +} + 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])