From 022e6928bc6df7f022552c678502684a628acce9 Mon Sep 17 00:00:00 2001 From: Graham McIntire Date: Sun, 19 Apr 2026 17:31:12 -0500 Subject: [PATCH] perf(prop-grid-rs): claim N forecast hours concurrently MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replaces the serial claim→process loop with a JoinSet of N worker tasks, each running the same claim→fetch→decode→score→complete flow against grid_tasks. FOR UPDATE SKIP LOCKED already prevents contention between workers so no additional coordination is needed. talos5 is 4c/4t (i5-7260U); PROP_GRID_RS_PARALLELISM=4 turns the full f01..f18 chain from ~18 min serial into roughly ~5 min. Per-task peak memory after the streamed-bands refactor is ~250 Mi, so 4 concurrent is ~1 Gi — well under the 2 Gi container limit. Postgres pool sized to parallelism + 2 so NOTIFY and an opportunistic claim retry share the pool without blocking a worker transaction. --- k8s/deployment-grid-rs.yaml | 11 ++++- rust/prop_grid_rs/src/bin/worker.rs | 64 +++++++++++++++++++++++------ 2 files changed, 61 insertions(+), 14 deletions(-) diff --git a/k8s/deployment-grid-rs.yaml b/k8s/deployment-grid-rs.yaml index 9e5ff98c..3cd0af27 100644 --- a/k8s/deployment-grid-rs.yaml +++ b/k8s/deployment-grid-rs.yaml @@ -61,8 +61,17 @@ spec: value: "/data/scores" - name: RUST_LOG value: "info" - - name: PROP_GRID_RS_PG_CONNS + # talos5 is 4c/4t (i5-7260U). Claiming 4 forecast hours in + # parallel drops full-chain wall time from ~18 min serial to + # ~5 min. Each concurrent task peaks around 250 Mi after the + # streamed-bands refactor; 4×250 ≈ 1 Gi is comfortably under + # the 2 Gi memory limit. + - name: PROP_GRID_RS_PARALLELISM value: "4" + # Pool size = parallelism + 2 (NOTIFY + retry slack). Kept in + # sync with PROP_GRID_RS_PARALLELISM when either changes. + - name: PROP_GRID_RS_PG_CONNS + value: "6" envFrom: - secretRef: name: prop-secrets diff --git a/rust/prop_grid_rs/src/bin/worker.rs b/rust/prop_grid_rs/src/bin/worker.rs index 3328dd25..e8db59e6 100644 --- a/rust/prop_grid_rs/src/bin/worker.rs +++ b/rust/prop_grid_rs/src/bin/worker.rs @@ -2,11 +2,13 @@ //! //! Flow: //! 1. Connect to Postgres (DATABASE_URL env). -//! 2. Loop: claim one `grid_tasks` row (fh>0, status='queued'), -//! run the fetch → decode → score → write-scores-file chain, mark done + -//! `NOTIFY propagation_ready ''`. On error mark +//! 2. Spawn N concurrent worker tasks (`PROP_GRID_RS_PARALLELISM`, default +//! = num_cpus). Each worker independently claims one `grid_tasks` row +//! (fh>0, status='queued') via `FOR UPDATE SKIP LOCKED`, runs the +//! fetch → decode → score → write-scores-file chain, and marks done + +//! `NOTIFY propagation_ready ''`. On error marks //! failed; hourly cron reseeds. -//! 3. Graceful shutdown on SIGTERM: requeue current row and exit. +//! 3. Graceful shutdown on SIGTERM: workers exit after their current task. use std::path::PathBuf; use std::sync::Arc; @@ -14,6 +16,7 @@ use std::sync::atomic::{AtomicBool, Ordering}; use std::time::Duration; use prop_grid_rs::{db, fetcher::HrrrClient, pipeline}; +use tokio::task::JoinSet; use tracing::{error, info, warn}; use tracing_subscriber::EnvFilter; @@ -33,10 +36,18 @@ async fn main() -> Result<(), Box> { .unwrap_or_else(|_| PathBuf::from("/data/scores")); let hrrr_base = std::env::var("HRRR_BASE_URL") .unwrap_or_else(|_| prop_grid_rs::fetcher::HRRR_BASE_DEFAULT.to_string()); + let parallelism: usize = std::env::var("PROP_GRID_RS_PARALLELISM") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or_else(num_cpus::get) + .max(1); + // +2 slots over parallelism: each worker holds one connection during + // claim/complete transactions; the spare lets NOTIFY and an opportunistic + // retry share the pool without blocking. let max_conn: u32 = std::env::var("PROP_GRID_RS_PG_CONNS") .ok() .and_then(|v| v.parse().ok()) - .unwrap_or(4); + .unwrap_or((parallelism as u32) + 2); let pool = db::connect(&database_url, max_conn).await?; let client = Arc::new(HrrrClient::new(hrrr_base)?); @@ -44,8 +55,37 @@ async fn main() -> Result<(), Box> { let shutdown = Arc::new(AtomicBool::new(false)); spawn_signal_handler(shutdown.clone()); - info!(scores_dir = %scores_dir.display(), "prop-grid-rs started"); + info!( + scores_dir = %scores_dir.display(), + parallelism, + pg_conns = max_conn, + "prop-grid-rs started" + ); + let mut set: JoinSet<()> = JoinSet::new(); + for worker_id in 0..parallelism { + let pool = pool.clone(); + let client = client.clone(); + let scores_dir = scores_dir.clone(); + let shutdown = shutdown.clone(); + set.spawn(async move { + worker_loop(worker_id, pool, client, scores_dir, shutdown).await; + }); + } + + while set.join_next().await.is_some() {} + + info!("prop-grid-rs shutting down"); + Ok(()) +} + +async fn worker_loop( + worker_id: usize, + pool: sqlx::PgPool, + client: Arc, + scores_dir: PathBuf, + shutdown: Arc, +) { while !shutdown.load(Ordering::Relaxed) { match db::claim_next(&pool).await { Ok(Some(task)) => { @@ -62,6 +102,7 @@ async fn main() -> Result<(), Box> { { Ok(stats) => { info!( + worker_id, task_id = %task.id, run_time = %task.run_time, forecast_hour = task.forecast_hour, @@ -70,13 +111,13 @@ async fn main() -> Result<(), Box> { "chain step done" ); if let Err(e) = db::complete(&pool, &task).await { - error!(%task_id, error = %e, "failed to mark task done"); + error!(worker_id, %task_id, error = %e, "failed to mark task done"); } } Err(e) => { - warn!(%task_id, error = %e, "chain step failed"); + warn!(worker_id, %task_id, error = %e, "chain step failed"); if let Err(dbe) = db::fail(&pool, task_id, &e.to_string()).await { - error!(%task_id, error = %dbe, "failed to mark task failed"); + error!(worker_id, %task_id, error = %dbe, "failed to mark task failed"); } } } @@ -85,14 +126,11 @@ async fn main() -> Result<(), Box> { tokio::time::sleep(IDLE_SLEEP).await; } Err(e) => { - error!(error = %e, "claim_next failed; backing off"); + error!(worker_id, error = %e, "claim_next failed; backing off"); tokio::time::sleep(IDLE_SLEEP).await; } } } - - info!("prop-grid-rs shutting down"); - Ok(()) } #[cfg(unix)]