Northern-Thailand-Ping-River-Monitor/docs/RUST_REWRITE.md

The Rust rewrite

Status: done (2026-10-18). The service is ping-collector (rust/README.md); Python keeps model training, evaluation and the archive backfills. Phases 1-3 replaced the Python collector, the forecast + notifier and the website/API one at a time, each checked against the Python side on production data before the switch; phase 4 removed the Python service code. The migration tooling (scripts/collector_parity.py, scripts/api_parity.py, scripts/install_collector.sh with its cutover/rollback modes, the spikes/rust-forecast spike) is in git history before the cleanup commit.

The assessment below was written before phase 1 and is kept as the record of why and how.

Short version

• Feasible. The risky part, serving the scikit-learn forecast from Rust with no train/serve skew, works: the spike reproduces the pandas feature matrix and the HistGradientBoosting regression output bit for bit (classifier heads to the last ulp).
• Training stays Python. There is no Rust HGB that trains like scikit-learn, and the evaluation harness (scripts/evaluate_variants.py) is where model work happens. A "Rust rewrite" means a Rust service (collectors, API, forecast serving, ntfy) plus a Python offline job (monthly retrain and evaluation) that exports models for Rust.
• Performance is not a reason. The forecast takes 158 ms per station in Python and 0.25 ms in Rust, but it runs once an hour. Load tests showed the site is bandwidth-bound once caches are hit.
• Real reasons, if any: one process instead of leader-elected uvicorn workers (the root of per-worker state bugs like /scraping/status), a fraction of the memory (324 MB today), typed parsing and SQL, one deployable binary, and the same stack as proplan (axum + sqlx).
• Recommendation: no big-bang rewrite. If Rust is wanted, start with the collector + notifier daemon (phase 1 below), which stands alone and can be abandoned without loss. The cheap Python alternative is at the end.

What exists (Python, src/)

Area	LOC	Fate in Rust
Collectors: RID hourly, HII/ThaiWater, RID dams, Open-Meteo	2,500	port; typed serde payloads
Web API (20 routes), schemas, health, metrics	2,300	port with identical JSON (golden-file tests)
DB adapters: Postgres, SQLite, MySQL, InfluxDB, VictoriaMetrics	1,300	Postgres only via sqlx, existing 13 tables, no data migration
ntfy notifications; Matrix alerting (CLI path)	1,050	port ntfy; Matrix only if still used
ML serving: predict, features, rain/dam/HII fill	~1,300	port; parity-tested (spike)
ML training + evaluation	~1,700	stays Python
CLI (main.py)	650	small clap CLI (fill-gaps, test cycle)
static/dashboard.html (164 KB)		unchanged, served as a static file
tests	3,900	Rust unit tests + parity fixtures + API golden files

Target shape

One binary, one process (tokio):

• tasks: RID / HII / dams / rain collectors on their schedules; forecast precompute when new readings land; ntfy notifier. Shared state is an Arc<RwLock<..>>, so there's no leader election, no status hand-off file and no executor sizing.
• axum + tower-http: same routes and JSON, compression (gzip/brotli also covers the 257 KB river GeoJSON), static dashboard, CORS; utoipa if /docs stays.
• sqlx against the existing schema. Python and Rust can run against the same database during migration.
• models: train_all() also exports each bundle (models/flood_<station>.json.gz, 1.2-2.3 MB), scripts/retrain.sh promotes it. Rust reloads changed files on its next run. A bundle whose feature_names Rust cannot build falls back to the heuristic, as Python does.

The permanent cost: the Python/Rust model contract

Every feature added in features.py must be ported to Rust in the same change, or the new bundle is refused. That has to be enforced by CI: Python writes a parity fixture per station (as export_fixture.py does), and a Rust test asserts bit-identical features and regression output. The spike found two traps that silently flip tree branches (HGB thresholds often equal real data values): serde_json's default float parsing (1 ulp off; enable float_roundtrip) and pandas' Kahan-compensated rolling sums/means (replicated in rolling_kahan()). This slows model iteration, which is the most active part of the repo.

Benefits and costs

For	Against
Single process: no leader election, no per-worker state	Two languages forever; parity CI is mandatory
Memory ~324 MB -> est. 60-80 MB (mostly tree arrays; less with u32/u16 nodes)	Weeks of work with little user-visible change
Typed RID/HII parsing, compile-checked SQL	Re-introduces bugs into recently hardened code (tz, alerts, security)
One static binary; no venv on the server for the API	Server still needs Python + sklearn for the monthly retrain
Same stack as proplan	Feature work = Python + Rust change every time

Plan (each phase ships on its own)

Phase	Scope	Status
0	Spike: P.1 features + HGB inference	done, bit-identical
1	ping-collector: RID, HII, dams, Open-Meteo	in production since 2026-10-17 23:09 (Python: COLLECTION_ENABLED=false)
2	Forecast crate + ntfy	in production since 2026-10-18 15:08 (FORECAST_IN_COLLECTOR=true); live check on production data: 48/48 rows identical, 96/96 probabilities bit-identical
3	Website + read API in ping-collector, Matrix like ntfy	in production since 2026-10-18 16:40 (WEB_LISTEN, water-monitor.service disabled); on production data 53 of 53 compared routes identical to FastAPI (scripts/api_parity.py)
4	Cleanup: Python service code removed	done 2026-10-18

0. Spike: done. P.1 features + HGB inference, bit-identical.
1. Collector daemon (rust/collector, ~1.5k LOC incl. tests): RID with the fixed semantics (store late readings, retry while behind, re-fetch yesterday's holes), HII, dams, Open-Meteo, into the same tables. The Python API runs with COLLECTION_ENABLED=false and re-forecasts when the newest RID reading in the database advances. ntfy moved to phase 2: the outlook alert needs the forecasts, which are still Python. Cutover and rollback were one command each (scripts/install_collector.sh, since removed).
2. Forecast crate + ntfy (rust/forecast, run by ping-collector with FORECAST_IN_COLLECTOR=true): all stations, HII/rain/dam inputs, heuristic fallback, probabilities and stages, forecast_history, forecast_current (served by the API) and notification_state. Parity: fixtures from the Python production path (CI: P.1 model; locally all 15), a live comparison script, and 751 notifications replayed word for word. Porting found a Python serving bug, fixed first: forecast rain (rain_fc24) was always NaN at serving time (docs/FLOOD_FORECASTING.md section 5). The ntfy change: ThaiWater mirror readings keep level alerts and the feed check current while RID lags at night. The Python API then only serves.
3. Website + read API (rust/collector/src/web, axum): every FastAPI route with the same parameters, JSON fields and order, status codes, FastAPI-style errors ({"detail": ...}, pydantic's 422 list), cache headers, gzip, CORS, Umami API tracking and the stale-while-revalidate response caches (now with byte budgets). /docs is the same Swagger UI; its OpenAPI document is generated from the route table whose parameter declarations also do the validation. One process serves, collects, forecasts and notifies (WEB_LISTEN), so the status file, the leader election and the forecast_current hand-off between processes are gone in effect. Checked by scripts/api_parity.py, which runs ping-collector serve next to the FastAPI app on the same database and compares every route (query values vary per run so FastAPI's caches answer fresh); --cutover-api refuses to switch on a difference. Intended differences: /health counts ThaiWater telemetry of the same gauges while RID lags (as the notifier); /scraping/status adds the collector version and the last forecast run; /forecast serves an old stored forecast marked X-Data-Stale instead of recomputing (same process); /sensors/thaiwater is cached for a minute; database errors serve the last good response where FastAPI answered []. Matrix: the notifier's transitions go to the room too (per station, outlook, feed; not the per-topic summaries), with their own matrix: state keys seeded from ntfy's, replacing the run.py --alert-check cron job that repeated alerts every 15 minutes. Python keeps training, evaluation and the monthly retrain timer.
4. Cleanup (done): FastAPI/uvicorn, the Python collector, notifier, Matrix alerting, CLI (run.py), the multi-backend adapters (SQLite, MySQL, InfluxDB, VictoriaMetrics), demo databases, Docker/Grafana files and the migration scripts are gone, about 17k lines. The Python environment is the training stack (pandas, scikit-learn, SQLAlchemy/psycopg2): src/ml, the backfill clients/stores (src/hii_collector.py, src/rid_reservoir.py) and src/ml/predict.py as the reference implementation the Rust forecaster is tested against. PostgreSQL only; the Python database tests run against a throwaway schema (TEST_DATABASE_URL, as the Rust ones). One installer, scripts/install.sh.

Decided before phase 3 (2026-10-18): Matrix stays and behaves like ntfy; /docs stays; SQLite dev mode goes; the monthly retrain stays on the same machine (Python, water-monitor-retrain.timer).

Cheaper alternative (stay Python)

Most of the operational gain comes from taking collection out of the web workers: run the collector as its own systemd unit (one process) and the uvicorn workers as read-only API. That's about a day of work. It removes leader election and per-worker state, but doesn't bring the memory, typing or single-binary benefits.