BM-2107: env variable to set a redis memory cap for queueing more prove segments

.env.broker-template

@@ -3,7 +3,7 @@
 RUST_LOG=info
 
 REDIS_HOST=redis
-REDIS_IMG=redis:7.2.5-alpine3.19
+REDIS_IMG=valkey/valkey:9.0
 
 POSTGRES_HOST=postgres
 POSTGRES_IMG=postgres:16.3-bullseye

bento/Cargo.toml

@@ -22,6 +22,7 @@ aws-sdk-s3 = "1.34"                                                       # used
 bento-client = { path = "crates/bento-client" }
 bincode = "1.3.3"
 bonsai-sdk = { version = "1.4.1", features = ["non_blocking"] }
+byte-unit = "5.1"
 bytemuck = "1.16"
 clap = { version = "4.5", features = ["derive", "env"] }
 deadpool-redis = "0.15"

bento/crates/workflow/Cargo.toml

@@ -22,6 +22,7 @@ cuda = ["risc0-zkvm/cuda"]
 anyhow = { workspace = true }
 bincode = { workspace = true }
 bytemuck = { workspace = true }
+byte-unit = { workspace = true }
 clap = { workspace = true, features = ["env", "derive"] }
 deadpool-redis = { workspace = true }
 hex = { workspace = true }

bento/crates/workflow/src/bin/agent.rs

@@ -14,7 +14,6 @@ async fn main() -> Result<()> {
     tracing_subscriber::fmt().with_env_filter(EnvFilter::from_default_env()).init();
 
     let args = Args::parse();
-    let task_stream = args.task_stream.clone();
     let agent = Agent::new(args).await.context("[BENTO-AGENT-001] Failed to initialize Agent")?;
 
     sqlx::migrate!("../taskdb/migrations")

bento/crates/workflow/src/lib.rs

@@ -9,6 +9,7 @@
 
 use crate::redis::RedisPool;
 use anyhow::{Context, Result};
+use byte_unit::Byte;
 use clap::Parser;
 use risc0_zkvm::{ProverOpts, ProverServer, VerifierContext, get_prover_server};
 use sqlx::postgres::{PgPool, PgPoolOptions};
@@ -31,6 +32,11 @@ pub use workflow_common::{
     SNARK_TIMEOUT_DEFAULT, s3::S3Client,
 };
 
+/// Format a byte count into human-readable units (e.g. `15.5 GB`)
+pub(crate) fn format_bytes(bytes: u64) -> String {
+    format!("{:#}", Byte::from_u64(bytes))
+}
+
 /// Workflow agent
 ///
 /// Monitors taskdb for new tasks on the selected stream and processes the work.
@@ -72,6 +78,18 @@ pub struct Args {
     #[clap(env,long, default_value_t = 8 * 60 * 60)]
     pub redis_ttl: u64,
 
+    /// Optional redis memory limit (in bytes) for queueing segments; executor will wait if usage exceeds this
+    #[clap(env, long, env = "REDIS_MAX_QUEUE_MEMORY")]
+    pub redis_max_queue_memory: Option<Byte>,
+
+    /// Poll interval, in milliseconds, while waiting for redis memory usage to drop below the queueing limit
+    #[clap(env, long, default_value_t = 250)]
+    pub redis_queue_memory_poll_ms: u64,
+
+    /// Interval, in seconds, for redis memory warning checks
+    #[clap(env, long, default_value_t = 5)]
+    pub redis_memory_warning_interval: u64,
+
     /// Executor limit, in millions of cycles
     #[clap(env, short, long, default_value_t = 100_000)]
     pub exec_cycle_limit: u64,
@@ -339,6 +357,52 @@ impl Agent {
                     }
                 });
             }
+
+            if let Some(threshold) =
+                self.args.redis_max_queue_memory.filter(|value| value.as_u64() > 0)
+            {
+                let term_sig_copy = term_sig.clone();
+                let redis_pool = self.redis_pool.clone();
+                let poll_interval = self.args.redis_memory_warning_interval;
+                tokio::spawn(async move {
+                    let poll_duration = time::Duration::from_secs(poll_interval.max(1));
+                    while !term_sig_copy.load(Ordering::Relaxed) {
+                        match redis_pool.get().await {
+                            Ok(mut conn) => {
+                                match crate::redis::fetch_memory_info(&mut conn).await {
+                                    Ok(info) => {
+                                        if info.used_memory >= threshold.as_u64() {
+                                            tracing::warn!(
+                                                total_system_memory = ?info.total_system_memory,
+                                                "[BENTO-WF-134] Redis memory usage {} exceeded warning threshold {:#} (maxmemory: {})",
+                                                format_bytes(info.used_memory),
+                                                threshold,
+                                                format_bytes(info.maxmemory),
+                                            );
+                                        } else {
+                                            tracing::debug!(
+                                                "Redis memory usage {} is below warning threshold {:#} (maxmemory: {})",
+                                                format_bytes(info.used_memory),
+                                                threshold,
+                                                format_bytes(info.maxmemory),
+                                            );
+                                        }
+                                    }
+                                    Err(err) => {
+                                        tracing::warn!(
+                                            "[BENTO-WF-135] Failed to query redis memory for warning check: {err:?}"
+                                        );
+                                    }
+                                }
+                            }
+                            Err(err) => tracing::warn!(
+                                "[BENTO-WF-136] Failed to acquire redis connection for memory warning check: {err:?}"
+                            ),
+                        }
+                        time::sleep(poll_duration).await;
+                    }
+                });
+            }
         }
 
         while !term_sig.load(Ordering::Relaxed) {