fix: set core opts from SHARD* env vars (#889)

succinctlabs · Jun 7, 2024 · 0377fcb · 0377fcb
2 parents 7eaa793 + 722b805
commit 0377fcb
Showing 1 changed file with 13 additions and 2 deletions.
diff --git a/core/src/utils/options.rs b/core/src/utils/options.rs
@@ -1,3 +1,8 @@
+use std::env;
+
+const DEFAULT_SHARD_SIZE: usize = 1 << 22;
+const DEFAULT_SHARD_BATCH_SIZE: usize = 16;
+
 #[derive(Debug, Clone, Copy)]
 pub struct SP1CoreOpts {
     pub shard_size: usize,
@@ -9,8 +14,14 @@ pub struct SP1CoreOpts {
 impl Default for SP1CoreOpts {
     fn default() -> Self {
         Self {
-            shard_size: 1 << 22,
-            shard_batch_size: 16,
+            shard_size: env::var("SHARD_SIZE").map_or_else(
+                |_| DEFAULT_SHARD_SIZE,
+                |s| s.parse::<usize>().unwrap_or(DEFAULT_SHARD_SIZE),
+            ),
+            shard_batch_size: env::var("SHARD_BATCH_SIZE").map_or_else(
+                |_| DEFAULT_SHARD_BATCH_SIZE,
+                |s| s.parse::<usize>().unwrap_or(DEFAULT_SHARD_BATCH_SIZE),
+            ),
             shard_chunking_multiplier: 1,
             reconstruct_commitments: true,
         }