CAGRA-ACE: Make Disk Workspaces Transactional and Non-destructive

cpp/include/cuvs/neighbors/cagra.hpp

@@ -69,7 +69,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";
   /**

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 <typename T>
 std::pair<file_descriptor, size_t> create_numpy_file(const std::string& path,
-                                                     const std::vector<size_t>& shape)
+                                                     const std::vector<size_t>& 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<T>();

cpp/src/neighbors/detail/cagra/cagra_build.cuh

@@ -35,8 +35,11 @@
 
 #include <rmm/resource_ref.hpp>
 
+#include <array>
+#include <cerrno>
 #include <chrono>
 #include <cstdio>
+#include <cstring>
 #include <filesystem>
 #include <omp.h>
 #include <type_traits>
@@ -53,6 +56,94 @@ namespace cuvs::neighbors::cagra::detail {
 constexpr double to_mib(size_t bytes) { return static_cast<double>(bytes) / (1 << 20); }
 constexpr double to_gib(size_t bytes) { return static_cast<double>(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<size_t>(which)].string();
+  }
+
+  void mark_artifact_created(artifact which) noexcept
+  {
+    artifacts_created_[static_cast<size_t>(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<std::filesystem::path, 4> artifacts_;
+  std::array<bool, 4> artifacts_created_{};
+  bool directory_created_by_this_build_ = false;
+  bool committed_                       = false;
+};
+
 template <typename T, typename IdxT>
 void check_graph_degree(size_t& intermediate_degree, size_t& graph_degree, size_t dataset_size)
 {
@@ -824,7 +915,7 @@ constexpr double vector_expansion_factor = 2.0;
 template <typename T, typename IdxT>
 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,
@@ -922,19 +1013,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(
@@ -1181,8 +1260,7 @@ index<T, IdxT> 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<T, IdxT>(intermediate_degree, graph_degree, dataset_size);
@@ -1225,24 +1303,32 @@ index<T, IdxT> 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<T>(
-        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<T>(
-        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<IdxT>(build_dir + "/dataset_mapping.npy", {dataset_size});
+      std::tie(mapping_fd, mapping_header_size) = cuvs::util::create_numpy_file<IdxT>(
+        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<IdxT>(
-        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)",
@@ -1540,20 +1626,14 @@ index<T, IdxT> build_ace(raft::resources const& res,
       std::chrono::duration_cast<std::chrono::milliseconds>(total_end - total_start).count();
     RAFT_LOG_INFO("ACE: Partitioned CAGRA build completed in %ld ms total", total_elapsed);
 
+    workspace.commit();
     return 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;
   }
 }

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
@@ -10,8 +10,16 @@
 
 #include <rmm/mr/managed_memory_resource.hpp>
 
+#include <algorithm>
+#include <atomic>
+#include <cstdint>
 #include <cstdio>
+#include <ctime>
 #include <filesystem>
+#include <fstream>
+#include <string>
+
+#include <unistd.h>
 
 namespace cuvs::neighbors::hnsw {
 
@@ -47,6 +55,118 @@ 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<std::uintptr_t>(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<uint64_t> counter_{0};
+};
+
+template <typename DataT>
+void build_ace_with_workspace(raft::resources const& resources,
+                              raft::host_matrix_view<const DataT, int64_t, raft::row_major> 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;
+
+  [[maybe_unused]] auto index = cagra::build(resources, params, dataset);
+}
+
+template <typename DataT>
+raft::host_matrix<DataT, int64_t> make_workspace_test_dataset()
+{
+  auto dataset = raft::make_host_matrix<DataT, int64_t>(1000, 8);
+  std::fill_n(dataset.data_handle(), dataset.size(), DataT{});
+  return dataset;
+}
+
+}  // namespace test_detail
+
+template <typename DataT>
+void test_ace_workspace_failure_preserves_caller_directory()
+{
+  raft::resources resources;
+  auto dataset = test_detail::make_workspace_test_dataset<DataT>();
+  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 <typename DataT>
+void test_ace_workspace_failure_does_not_truncate_existing_artifact()
+{
+  raft::resources resources;
+  auto dataset = test_detail::make_workspace_test_dataset<DataT>();
+  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 <typename DistanceT, typename DataT, typename IdxT>
 class AnnHnswAceTest : public ::testing::TestWithParam<AnnHnswAceInputs> {
  public: