[Improvement] Allow Configurable DRAM in DiskANN SSD Wrapper (rapidsai#1598)

Notes: During the final param parsing to DiskANN index building API, we convert the dram budgets to GB

Authors:
  - Tarang Jain (https://github.com/tarang-jain)

Approvers:
  - MithunR (https://github.com/mythrocks)

URL: rapidsai#1598

Author: tarang-jain

cpp/bench/ann/src/diskann/diskann_benchmark.cpp

@@ -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
  */
 
@@ -36,9 +36,15 @@ void parse_build_param(const nlohmann::json& conf,
 {
   param.R       = conf.at("R");
   param.L_build = conf.at("L_build");
-  if (conf.contains("alpha")) { param.num_threads = conf.at("alpha"); }
+  if (conf.contains("alpha")) { param.alpha = conf.at("alpha"); }
   if (conf.contains("num_threads")) { param.num_threads = conf.at("num_threads"); }
   if (conf.contains("QD")) { param.QD = conf.at("QD"); }
+  if (conf.contains("build_dram_budget_megabytes")) {
+    param.build_dram_budget_megabytes = conf.at("build_dram_budget_megabytes");
+  }
+  if (conf.contains("search_dram_budget_megabytes")) {
+    param.search_dram_budget_megabytes = conf.at("search_dram_budget_megabytes");
+  }
   param.dataset_base_file = cuvs::bench::configuration::singleton().get_dataset_conf().base_file;
   for (const auto& index : cuvs::bench::configuration::singleton().get_indices()) {
     // The nlohmann::json operator== comparator compares each key and value by content. Reference:

cpp/bench/ann/src/diskann/diskann_wrapper.h

@@ -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
  */
 #pragma once
@@ -17,6 +17,8 @@
 #include <utils.h>
 
 #include <chrono>
+#include <cmath>
+#include <cstdio>
 #include <memory>
 #include <vector>
 
@@ -172,12 +174,14 @@ class diskann_ssd : public algo<T> {
   struct build_param {
     uint32_t R;
     uint32_t L_build;
-    uint32_t build_pq_bytes       = 0;
-    float alpha                   = 1.2;
-    int num_threads               = omp_get_max_threads();
-    uint32_t QD                   = 192;
-    std::string dataset_base_file = "";
-    std::string index_file        = "";
+    uint32_t build_pq_bytes               = 0;
+    float alpha                           = 1.2;
+    int num_threads                       = omp_get_max_threads();
+    uint32_t QD                           = 192;
+    std::string dataset_base_file         = "";
+    std::string index_file                = "";
+    uint32_t build_dram_budget_megabytes  = std::numeric_limits<uint32_t>::max();
+    uint32_t search_dram_budget_megabytes = std::numeric_limits<uint32_t>::max();
   };
   using search_param_base = typename algo<T>::search_param;
 
@@ -232,12 +236,17 @@ template <typename T>
 diskann_ssd<T>::diskann_ssd(Metric metric, int dim, const build_param& param) : algo<T>(metric, dim)
 {
   // Currently set the indexing RAM budget and the search RAM budget to max value to avoid sharding
-  uint32_t build_dram_budget  = std::numeric_limits<uint32_t>::max();
-  uint32_t search_dram_budget = std::numeric_limits<uint32_t>::max();
+  float build_dram_budget  = static_cast<float>(param.build_dram_budget_megabytes) / 1024.0f;
+  float search_dram_budget = static_cast<float>(param.search_dram_budget_megabytes) / 1024.0f;
+  char search_buf[16];
+  char build_buf[16];
+  std::snprintf(search_buf, sizeof(search_buf), "%.2f", search_dram_budget);
+  std::snprintf(build_buf, sizeof(build_buf), "%.2f", build_dram_budget);
+  const std::string search_dram_budget_gb(search_buf);
+  const std::string build_dram_budget_gb(build_buf);
   index_build_params_str =
     std::string(std::to_string(param.R)) + " " + std::string(std::to_string(param.L_build)) + " " +
-    std::string(std::to_string(search_dram_budget)) + " " +
-    std::string(std::to_string(build_dram_budget)) + " " +
+    search_dram_budget_gb + " " + build_dram_budget_gb + " " +
     std::string(std::to_string(param.num_threads)) + " " + std::string(std::to_string(false)) +
     " " + std::string(std::to_string(false)) + " " + std::string(std::to_string(0)) + " " +
     std::string(std::to_string(param.QD));