feat: improve storage budget with Docker-style size formats

- Changed BB_STORAGE_BUDGET_GB env var to BB_STORAGE_BUDGET
- Changed --storage_budget_gb CLI flag to --storage_budget
- Now accepts Docker-style size formats: '500m', '2g', '1024k'
- Added proper error handling with clear messages for invalid formats
- Created reusable parse_size_string() function
- Updated docker_isolate to accept TMPFS_SIZE parameter

Author: johnathan79717

barretenberg/cpp/src/barretenberg/api/api.hpp

@@ -23,7 +23,7 @@ class API {
         bool write_vk{ false };    // should we addditionally write the verification key when writing the proof
         bool include_gates_per_opcode{ false }; // should we include gates_per_opcode in the gates command output
         bool slow_low_memory{ false };          // use file backed memory for polynomials
-        size_t storage_budget_gb{ 0 };          // storage budget in GB for file backed memory (0 = no limit)
+        std::string storage_budget;             // storage budget for file backed memory (e.g. "500m", "2g")
         bool update_inputs{ false };            // update inputs when check fails
 
         friend std::ostream& operator<<(std::ostream& os, const Flags& flags)
@@ -41,7 +41,7 @@ class API {
                << "  write_vk " << flags.write_vk << "\n"
                << "  include_gates_per_opcode " << flags.include_gates_per_opcode << "\n"
                << "  slow_low_memory " << flags.slow_low_memory << "\n"
-               << "  storage_budget_gb " << flags.storage_budget_gb << "\n"
+               << "  storage_budget " << flags.storage_budget << "\n"
                << "]" << std::endl;
             return os;
         }

barretenberg/cpp/src/barretenberg/bb/cli.cpp

@@ -279,10 +279,10 @@ int parse_and_run_cli_command(int argc, char* argv[])
     };
 
     const auto add_storage_budget_option = [&](CLI::App* subcommand) {
-        return subcommand->add_option("--storage_budget_gb",
-                                      flags.storage_budget_gb,
-                                      "Storage budget in GB for FileBackedMemory. When exceeded, falls back to RAM "
-                                      "(requires --slow_low_memory).");
+        return subcommand->add_option("--storage_budget",
+                                      flags.storage_budget,
+                                      "Storage budget for FileBackedMemory (e.g. '500m', '2g'). When exceeded, falls "
+                                      "back to RAM (requires --slow_low_memory).");
     };
 
     const auto add_update_inputs_flag = [&](CLI::App* subcommand) {
@@ -566,8 +566,8 @@ int parse_and_run_cli_command(int argc, char* argv[])
     debug_logging = flags.debug;
     verbose_logging = debug_logging || flags.verbose;
     slow_low_memory = flags.slow_low_memory;
-    if (flags.storage_budget_gb > 0) {
-        storage_budget = flags.storage_budget_gb * 1024ULL * 1024ULL * 1024ULL;
+    if (!flags.storage_budget.empty()) {
+        storage_budget = parse_size_string(flags.storage_budget);
     }
 #ifndef __wasm__
     if (print_op_counts || !op_counts_out.empty()) {

barretenberg/cpp/src/barretenberg/polynomials/backing_memory.cpp

@@ -1,30 +1,73 @@
 #include "barretenberg/polynomials/backing_memory.hpp"
+#include "barretenberg/common/throw_or_abort.hpp"
 #include <atomic>
+#include <cctype>
 #include <cstdlib>
 #include <limits>
+#include <string>
 
 // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
 bool slow_low_memory =
     std::getenv("BB_SLOW_LOW_MEMORY") == nullptr ? false : std::string(std::getenv("BB_SLOW_LOW_MEMORY")) == "1";
 
-// Parse storage budget from environment variable (in GB)
-static size_t parse_storage_budget()
+// Parse storage size string (e.g., "500m", "2g", "1024k")
+size_t parse_size_string(const std::string& size_str)
 {
-    const char* env_val = std::getenv("BB_STORAGE_BUDGET_GB");
-    if (env_val == nullptr) {
-        return std::numeric_limits<size_t>::max(); // No limit by default
+    if (size_str.empty()) {
+        return std::numeric_limits<size_t>::max();
     }
 
     try {
-        size_t gb = std::stoull(env_val);
-        return gb * 1024ULL * 1024ULL * 1024ULL; // Convert GB to bytes
-    } catch (...) {
-        return std::numeric_limits<size_t>::max(); // Invalid value, no limit
+        std::string str = size_str;
+
+        // Convert to lowercase for case-insensitive comparison
+        char suffix = static_cast<char>(std::tolower(static_cast<unsigned char>(str.back())));
+        size_t multiplier = 1;
+
+        // Check for unit suffix
+        if (suffix == 'k') {
+            multiplier = 1024ULL;
+            str.pop_back();
+        } else if (suffix == 'm') {
+            multiplier = 1024ULL * 1024ULL;
+            str.pop_back();
+        } else if (suffix == 'g') {
+            multiplier = 1024ULL * 1024ULL * 1024ULL;
+            str.pop_back();
+        } else if (std::isdigit(static_cast<unsigned char>(suffix)) == 0) {
+            // Invalid suffix
+            throw_or_abort("Invalid storage size format: '" + size_str + "'. Use format like '500m', '2g', or '1024k'");
+        }
+
+        // Check if remaining string is a valid number
+        if (str.empty()) {
+            throw_or_abort("Invalid storage size format: '" + size_str + "'. No numeric value provided");
+        }
+
+        size_t value = std::stoull(str);
+        return value * multiplier;
+    } catch (const std::invalid_argument&) {
+        throw_or_abort("Invalid storage size format: '" + size_str + "'. Not a valid number");
+    } catch (const std::out_of_range&) {
+        throw_or_abort("Invalid storage size format: '" + size_str + "'. Value out of range");
     }
 }
 
+namespace {
+// Parse storage budget from environment variable (supports k/m/g suffixes like Docker)
+size_t parse_storage_budget()
+{
+    const char* env_val = std::getenv("BB_STORAGE_BUDGET");
+    if (env_val == nullptr) {
+        return std::numeric_limits<size_t>::max(); // No limit by default
+    }
+
+    return parse_size_string(std::string(env_val));
+}
+} // namespace
+
 // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
 size_t storage_budget = parse_storage_budget();
 
 // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
-std::atomic<size_t> current_storage_usage{ 0 };
+std::atomic<size_t> current_storage_usage{ 0 };

barretenberg/cpp/src/barretenberg/polynomials/backing_memory.hpp

@@ -27,6 +27,9 @@ extern size_t storage_budget;
 // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
 extern std::atomic<size_t> current_storage_usage;
 
+// Parse storage size string (e.g., "500m", "2g", "1024k")
+size_t parse_size_string(const std::string& size_str);
+
 template <typename T> class AlignedMemory;
 
 #ifndef __wasm__

ci3/docker_isolate

@@ -6,6 +6,7 @@ NO_CD=1 source $(git rev-parse --show-toplevel)/ci3/source
 cmd=$1
 export CPUS=${CPUS:-2}
 export MEM=${MEM:-$((CPUS * 4))g}
+export TMPFS_SIZE=${TMPFS_SIZE:-1g}
 
 function cleanup {
   if [ -n "${cid:-}" ]; then
@@ -54,14 +55,15 @@ cid=$(docker run -d \
   --pid=host \
   --user $(id -u):$(id -g) \
   -v$HOME:$HOME \
-  --mount type=tmpfs,target=/tmp,tmpfs-size=1g \
+  --mount type=tmpfs,target=/tmp,tmpfs-size=$TMPFS_SIZE \
   --workdir $PWD \
   -e HOME \
   -e VERBOSE \
   -e GIT_CONFIG_GLOBAL=$root/build-images/src/home/.gitconfig \
   -e FORCE_COLOR=true \
   -e CPUS \
   -e MEM \
+  -e TMPFS_SIZE \
   "${arg_env_vars[@]}" \
   aztecprotocol/build:3.0 \
   /bin/bash -c "$cmd")