From 65f7963ca37e6fd3474fa9ebd45e5e65ad895565 Mon Sep 17 00:00:00 2001 From: Graham McIntire Date: Sun, 19 Apr 2026 18:10:18 -0500 Subject: [PATCH] feat(prop-grid-rs): analysis pipeline + kind-aware worker dispatch run_analysis_step chains the Rust-side f00 work end-to-end: - fetch surface + pressure GRIB2 + native duct in parallel (tokio::try_join!) with per-leg error-lift to PipelineError - merge pressure into surface, fold native duct metrics per cell - NEXRAD composite reflectivity overlay via the IEM n0q PNG path - commercial-link degradation via the Postgres per-link baseline - score all 23 bands in parallel (rayon) then write band files via parallel spawn_blocking (try_join_all) - write the ProfilesFile as MessagePack alongside the score files worker.rs main loop now claims kind='analysis' first, falls back to kind='forecast'. Analysis tasks dispatch to run_analysis_step; forecast tasks still go through run_chain_step unchanged. ChainOutcome enum tags log events (kind=analysis|forecast) and carries analysis-specific counters (profile_cells_written, duct_cells, nexrad_cells_with_echo, commercial_cells_boosted) so Phase 3 cutover observability is readable from kubectl logs alone. 118 Rust tests green. Elixir seeder flip + f00 code deletion lands next. --- rust/prop_grid_rs/src/bin/worker.rs | 91 ++++++-- rust/prop_grid_rs/src/pipeline.rs | 330 ++++++++++++++++++++++++++++ 2 files changed, 399 insertions(+), 22 deletions(-) diff --git a/rust/prop_grid_rs/src/bin/worker.rs b/rust/prop_grid_rs/src/bin/worker.rs index a0372bae..edfa40f5 100644 --- a/rust/prop_grid_rs/src/bin/worker.rs +++ b/rust/prop_grid_rs/src/bin/worker.rs @@ -23,6 +23,11 @@ use tracing_subscriber::EnvFilter; const IDLE_SLEEP: Duration = Duration::from_secs(5); +enum ChainOutcome { + Forecast(pipeline::ChainStepStats), + Analysis(pipeline::AnalysisStepStats), +} + #[tokio::main] async fn main() -> Result<(), Box> { tracing_subscriber::fmt() @@ -99,34 +104,76 @@ async fn worker_loop( shutdown: Arc, ) { while !shutdown.load(Ordering::Relaxed) { - match db::claim_next(&pool).await { + // Analysis tasks (kind='analysis', f00) get priority — a /map + // user values the current hour more than +1..+18h. Forecast is + // only attempted when there is no analysis row ready. + let claimed = match db::claim_next_analysis(&pool).await { + Ok(Some(task)) => Ok(Some(task)), + Ok(None) => db::claim_next(&pool).await, + Err(e) => Err(e), + }; + + match claimed { Ok(Some(task)) => { let task_id = task.id; let _guard = metrics::InFlightGuard::new(); let started = Instant::now(); - match pipeline::run_chain_step( - &client, - &scores_dir, - &pipeline::ChainStepInput { - run_time: task.run_time, - forecast_hour: task.forecast_hour as u8, - }, - ) - .await - { - Ok(stats) => { + let result = match task.kind { + db::TaskKind::Forecast => { + pipeline::run_chain_step( + &client, + &scores_dir, + &pipeline::ChainStepInput { + run_time: task.run_time, + forecast_hour: task.forecast_hour as u8, + }, + ) + .await + .map(ChainOutcome::Forecast) + } + db::TaskKind::Analysis => { + pipeline::run_analysis_step( + &client, + &pool, + &scores_dir, + &pipeline::AnalysisStepInput { run_time: task.run_time }, + ) + .await + .map(ChainOutcome::Analysis) + } + }; + + match result { + Ok(outcome) => { let elapsed = started.elapsed(); metrics::record_chain_step(elapsed, true); - info!( - worker_id, - task_id = %task.id, - run_time = %task.run_time, - forecast_hour = task.forecast_hour, - files = stats.score_files_written, - points = stats.point_count, - duration_s = elapsed.as_secs_f64(), - "chain step done" - ); + match &outcome { + ChainOutcome::Forecast(stats) => info!( + worker_id, + task_id = %task.id, + kind = "forecast", + run_time = %task.run_time, + forecast_hour = task.forecast_hour, + files = stats.score_files_written, + points = stats.point_count, + duration_s = elapsed.as_secs_f64(), + "chain step done" + ), + ChainOutcome::Analysis(stats) => info!( + worker_id, + task_id = %task.id, + kind = "analysis", + run_time = %task.run_time, + files = stats.score_files_written, + points = stats.point_count, + profile_cells = stats.profile_cells_written, + duct_cells = stats.duct_cells, + nexrad_cells = stats.nexrad_cells_with_echo, + commercial_cells = stats.commercial_cells_boosted, + duration_s = elapsed.as_secs_f64(), + "analysis step done" + ), + } if let Err(e) = db::complete(&pool, &task).await { error!(worker_id, %task_id, error = %e, "failed to mark task done"); } diff --git a/rust/prop_grid_rs/src/pipeline.rs b/rust/prop_grid_rs/src/pipeline.rs index d7345bd2..50a7aa90 100644 --- a/rust/prop_grid_rs/src/pipeline.rs +++ b/rust/prop_grid_rs/src/pipeline.rs @@ -12,9 +12,13 @@ use std::sync::Arc; use chrono::{DateTime, Datelike, Timelike, Utc}; use crate::band_config; +use crate::commercial::{self, LinkLookupEntry}; use crate::decoder::{self, CellValues, PointGrid}; use crate::fetcher::{self, HrrrClient, Product}; use crate::grid::wgrib2_grid_spec; +use crate::native_duct; +use crate::nexrad::{self, NexradObservation}; +use crate::profiles_file::{self, CellEntry}; use crate::scores_file::{self, ScorePoint}; use crate::scorer::{self, Conditions}; use crate::sounding_params::{self, Level}; @@ -27,6 +31,14 @@ pub enum PipelineError { Decode(#[from] decoder::DecodeError), #[error("write: {0}")] Write(#[from] scores_file::WriteError), + #[error("profile write: {0}")] + ProfileWrite(#[from] profiles_file::WriteError), + #[error("native duct: {0}")] + NativeDuct(#[from] native_duct::NativeDuctError), + #[error("nexrad: {0}")] + Nexrad(#[from] nexrad::NexradError), + #[error("commercial: {0}")] + Commercial(#[from] commercial::CommercialError), #[error("forecast hour 0 is reserved for Elixir")] F00Reserved, #[error("surface grib was empty")] @@ -326,3 +338,321 @@ mod tests { assert!(matches!(err, PipelineError::F00Reserved)); } } + +// ===================================================================== +// Analysis pipeline (f00): fetch surface+pressure + native duct, merge +// with NEXRAD + commercial-link degradation, score all bands, write the +// ProfilesFile (MessagePack) and the band score files. +// ===================================================================== + +#[derive(Debug, Clone)] +pub struct AnalysisStepInput { + pub run_time: DateTime, +} + +impl AnalysisStepInput { + pub fn valid_time(&self) -> DateTime { + // f00 analysis valid_time == run_time. + self.run_time + } +} + +#[derive(Debug, Clone)] +pub struct AnalysisStepStats { + pub score_files_written: u32, + pub point_count: u32, + pub band_count: u32, + pub profile_cells_written: u32, + pub duct_cells: u32, + pub nexrad_cells_with_echo: u32, + pub commercial_cells_boosted: u32, +} + +pub async fn run_analysis_step( + client: &HrrrClient, + pool: &sqlx::PgPool, + scores_dir: &Path, + step: &AnalysisStepInput, +) -> Result { + let valid_time = step.valid_time(); + let date = step.run_time.date_naive(); + let hour = step.run_time.hour() as u8; + let grid_spec = wgrib2_grid_spec(); + + // Parallel fetch + decode of surface + pressure + native duct. All + // three are independent HRRR products; overlap the HTTP legs and + // the subsequent wgrib2 decodes so the long pole is a single + // decode, not three serial decodes. + let sfc_wanted = fetcher::surface_messages_owned(); + let prs_wanted = fetcher::pressure_messages_grid(); + let sfc_pattern = match_pattern(&sfc_wanted); + let prs_pattern = match_pattern(&prs_wanted); + + // All three futures return different error types; lift into + // PipelineError inside each branch so try_join! stays happy. + let sfc_fut = async { + client + .fetch_product_blob(date, hour, Product::Surface, 0, &sfc_wanted) + .await + .map_err(PipelineError::Fetch) + }; + let prs_fut = async { + client + .fetch_product_blob(date, hour, Product::Pressure, 0, &prs_wanted) + .await + .map_err(PipelineError::Fetch) + }; + let duct_fut = async { + native_duct::fetch_native_duct_grid(client, date, hour, 0, grid_spec) + .await + .map_err(PipelineError::NativeDuct) + }; + + let (sfc_blob, prs_blob, duct_map) = tokio::try_join!(sfc_fut, prs_fut, duct_fut)?; + + let duct_count = duct_map.len() as u32; + + let sfc_grid = tokio::task::spawn_blocking(move || { + decoder::extract_grid(&sfc_blob, &sfc_pattern, grid_spec) + }) + .await + .expect("blocking join")?; + + let prs_grid = tokio::task::spawn_blocking(move || { + decoder::extract_grid(&prs_blob, &prs_pattern, grid_spec) + }) + .await + .expect("blocking join")?; + + // Merge pressure into surface, then fold duct metrics into every + // cell that has both. The duct values live alongside the raw + // CellValues strings so build_grid_cache_rows can pick them up + // under their atom keys after the Elixir reader normalises. + let mut merged = merge_grids(sfc_grid, prs_grid); + merged = native_duct::merge_duct_grid(merged, &duct_map); + + // NEXRAD composite reflectivity overlay — only valid for f00 + // because forecast hours can't see the future radar image. + let nexrad_http = reqwest::Client::builder() + .user_agent("prop-grid-rs/0.1") + .build() + .expect("reqwest client"); + let nexrad_points: Vec<(f64, f64)> = merged + .keys() + .map(|&k| decoder::key_to_latlon(k)) + .collect(); + let nexrad_obs: Vec = + match nexrad::fetch_frame(&nexrad_http, valid_time, &nexrad_points).await { + Ok(obs) => obs, + Err(e) => { + tracing::warn!(error = %e, "NEXRAD fetch failed — continuing without radar overlay"); + Vec::new() + } + }; + let nexrad_cells_with_echo = apply_nexrad(&mut merged, &nexrad_obs); + + // Commercial-link degradation — precompute once, then fold into + // every in-range cell. Only ~7 links cluster around DFW so most + // cells stay untouched. + let commercial_lookup: Vec = commercial::build_link_lookup(pool, valid_time) + .await + .map_err(PipelineError::Commercial)?; + let commercial_cells_boosted = + apply_commercial(&mut merged, &commercial_lookup); + + let point_count = merged.len() as u32; + + // Build the ProfilesFile cell list (the raw enriched grid) and + // the Conditions vector (for scoring) in one pass. Profile keeps + // atmospheric + duct + NEXRAD + commercial fields under + // string-keyed msgpack; Rust doesn't need to mutate those further. + let mut profile_entries: Vec = Vec::with_capacity(merged.len()); + let mut prepared: Vec<(f64, f64, Conditions, scorer::BandInvariants)> = Vec::with_capacity(merged.len()); + for (key, cell) in merged.iter() { + let (lat, lon) = decoder::key_to_latlon(*key); + if let Some(conditions) = cell_to_conditions(cell, lat, lon, &valid_time) { + let invariants = scorer::precompute_band_invariants(&conditions); + prepared.push((lat, lon, conditions, invariants)); + profile_entries.push(cell_to_profile_entry(lat, lon, cell)); + } + } + + // Write the profile file on the blocking pool — msgpack encode + + // gzip + atomic rename. One file per run, ~2 MB compressed, so the + // write runs well inside the scoring loop's wall time. + let scores_dir_owned = scores_dir.to_path_buf(); + let scores_dir_for_profiles = scores_dir_owned.clone(); + let profile_future = { + let profiles = profile_entries; + tokio::task::spawn_blocking(move || { + profiles_file::write_atomic(&scores_dir_for_profiles, valid_time, &profiles) + .map(|_| 0u32) + }) + }; + + // Bands and score writes run in parallel, same shape as + // run_chain_step. Rayon across bands, try_join_all across writes. + let bands = band_config::all_bands(); + let prepared_arc = Arc::new(prepared); + let scored: Vec<(u32, Vec)> = { + use rayon::prelude::*; + bands + .par_iter() + .map(|band| { + let mut scores: Vec = Vec::with_capacity(prepared_arc.len()); + for (lat, lon, conditions, invariants) in prepared_arc.iter() { + let r = scorer::composite_score_with(conditions, band, Some(*invariants)); + scores.push(ScorePoint { lat: *lat, lon: *lon, score: r.score }); + } + (band.freq_mhz, scores) + }) + .collect() + }; + + let write_futs = scored.into_iter().map(|(band_mhz, scores)| { + let dir = scores_dir_owned.clone(); + tokio::task::spawn_blocking(move || { + scores_file::write_atomic(&dir, band_mhz, valid_time, &scores) + }) + }); + let results = futures::future::try_join_all(write_futs).await.expect("blocking join"); + for r in results { + r?; + } + profile_future + .await + .expect("blocking join") + .map_err(PipelineError::ProfileWrite)?; + + let files_written = bands.len() as u32; + Ok(AnalysisStepStats { + score_files_written: files_written, + point_count, + band_count: bands.len() as u32, + profile_cells_written: prepared_arc.len() as u32, + duct_cells: duct_count, + nexrad_cells_with_echo, + commercial_cells_boosted, + }) +} + +fn apply_nexrad(grid: &mut PointGrid, obs: &[NexradObservation]) -> u32 { + let mut with_echo = 0u32; + for o in obs { + let key = ( + (o.lat * 1000.0).round() as i32, + (o.lon * 1000.0).round() as i32, + ); + // merged's keys are whatever extract_grid produced; use direct + // lat/lon match by scanning. This is O(N*M) worst-case but N is + // ~92k and obs mirrors the same point list, so in practice each + // observation matches exactly one cell by int-key. + if let Some(cell) = grid.get_mut(&key) { + if o.max_reflectivity_dbz > 0.0 { + with_echo += 1; + } + cell.insert( + "nexrad_max_reflectivity_dbz".into(), + o.max_reflectivity_dbz as f32, + ); + } + } + with_echo +} + +fn apply_commercial(grid: &mut PointGrid, lookup: &[LinkLookupEntry]) -> u32 { + let mut boosted = 0u32; + for (key, cell) in grid.iter_mut() { + let (lat, lon) = decoder::key_to_latlon(*key); + if let Some(d) = commercial::degradation_at((lat, lon), lookup) { + cell.insert("commercial_degradation_db".into(), d.degradation_db as f32); + cell.insert("commercial_baseline_dbm".into(), d.baseline_dbm as f32); + cell.insert("commercial_current_dbm".into(), d.current_dbm as f32); + cell.insert("commercial_n_links".into(), d.n_links as f32); + boosted += 1; + } + } + boosted +} + +/// Turn one enriched cell into a ProfilesFile entry. The msgpack +/// `profile` map keeps string keys; the Elixir reader atomizes the +/// whitelisted subset (surface_*, hpbl_m, pwat_mm, duct metrics, …). +fn cell_to_profile_entry(lat: f64, lon: f64, cell: &CellValues) -> CellEntry { + use profiles_file::value_map; + use rmpv::Value as V; + + let mut kv: Vec<(&'static str, V)> = Vec::with_capacity(16); + + if let Some(&v) = cell.get("TMP:2 m above ground") { + kv.push(("surface_temp_c", V::F64((v as f64) - 273.15))); + } + if let Some(&v) = cell.get("DPT:2 m above ground") { + kv.push(("surface_dewpoint_c", V::F64((v as f64) - 273.15))); + } + if let Some(&v) = cell.get("PRES:surface") { + kv.push(("surface_pressure_mb", V::F64((v as f64) / 100.0))); + } + if let Some(&v) = cell.get("HPBL:surface") { + kv.push(("hpbl_m", V::F64(v as f64))); + } + if let Some(&v) = cell.get("PWAT:entire atmosphere (considered as a single layer)") { + kv.push(("pwat_mm", V::F64(v as f64))); + } + if let Some(&v) = cell.get("native_min_gradient") { + kv.push(("native_min_gradient", V::F64(v as f64))); + } + if let Some(&v) = cell.get("best_duct_freq_ghz") { + kv.push(("best_duct_freq_ghz", V::F64(v as f64))); + } + if let Some(&v) = cell.get("max_duct_thickness_m") { + kv.push(("max_duct_thickness_m", V::F64(v as f64))); + } + if let Some(&v) = cell.get("duct_count") { + kv.push(("duct_count", V::Integer((v as i64).into()))); + } + if let Some(&v) = cell.get("nexrad_max_reflectivity_dbz") { + kv.push(("nexrad_max_reflectivity_dbz", V::F64(v as f64))); + } + if let Some(&d) = cell.get("commercial_degradation_db") { + let links = cell.get("commercial_n_links").copied().unwrap_or(0.0); + let baseline = cell.get("commercial_baseline_dbm").copied().unwrap_or(0.0); + let current = cell.get("commercial_current_dbm").copied().unwrap_or(0.0); + kv.push(( + "commercial_link_degradation", + value_map([ + ("degradation_db", V::F64(d as f64)), + ("baseline_dbm", V::F64(baseline as f64)), + ("current_dbm", V::F64(current as f64)), + ("n_links", V::Integer((links as i64).into())), + ]), + )); + } + + // Pressure-level profile list for SoundingParams.derive on the + // Elixir side. Mirrors the fields derive expects. + let levels: Vec = fetcher::GRID_PRESSURE_LEVELS + .iter() + .filter_map(|&p| { + let pres_mb = p as f64; + let t = cell.get(&format!("TMP:{p} mb")).copied()?; + let d = cell.get(&format!("DPT:{p} mb")).copied(); + let h = cell.get(&format!("HGT:{p} mb")).copied()?; + let mut pairs: Vec<(V, V)> = Vec::with_capacity(4); + pairs.push((V::String("pres_mb".into()), V::F64(pres_mb))); + pairs.push((V::String("hght_m".into()), V::F64(h as f64))); + pairs.push((V::String("tmpc".into()), V::F64((t as f64) - 273.15))); + if let Some(dv) = d { + pairs.push((V::String("dwpc".into()), V::F64((dv as f64) - 273.15))); + } + Some(V::Map(pairs)) + }) + .collect(); + if !levels.is_empty() { + kv.push(("profile", V::Array(levels))); + } + + let profile = value_map(kv); + + CellEntry { lat, lon, profile } +}