graham/prop
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

feat(prop-grid-rs): analysis pipeline + kind-aware worker dispatch

Browse source 
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.
This commit is contained in:
Graham McIntire 2026-04-19 18:10:18 -05:00
parent d41ced5850
commit 65f7963ca3
No known key found for this signature in database
GPG key ID: F4ABF488E6029E59
2 changed files with 399 additions and 22 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

91
rust/prop_grid_rs/src/bin/worker.rs
View file

@ -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<dyn std::error::Error>> {
tracing_subscriber::fmt()
@ -99,34 +104,76 @@ async fn worker_loop(
shutdown: Arc<AtomicBool>,
) {
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");
}

330
rust/prop_grid_rs/src/pipeline.rs
View file

@ -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<Utc>,
}
impl AnalysisStepInput {
pub fn valid_time(&self) -> DateTime<Utc> {
// 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<AnalysisStepStats, PipelineError> {
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<NexradObservation> =
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<LinkLookupEntry> = 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<CellEntry> = 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<ScorePoint>)> = {
use rayon::prelude::*;
bands
.par_iter()
.map(|band| {
let mut scores: Vec<ScorePoint> = 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<V> = 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 }
}