From 74e9ea5bbeb03a90dc913a412cc7dbcab4de4c25 Mon Sep 17 00:00:00 2001 From: frenbox Date: Sun, 29 Mar 2026 20:04:04 -0500 Subject: [PATCH 01/20] adding villar light curve fitting using batch GPU --- .gitignore | 1 + Cargo.lock | 31 ++++++++ Cargo.toml | 10 ++- src/enrichment/ztf.rs | 159 ++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 200 insertions(+), 1 deletion(-) diff --git a/.gitignore b/.gitignore index b566bf661..e98363c7a 100644 --- a/.gitignore +++ b/.gitignore @@ -19,3 +19,4 @@ logs *.jsonl simbad_cache *.parquet +config.docker.yaml diff --git a/Cargo.lock b/Cargo.lock index bd69a7d2b..5a5a4781c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1157,6 +1157,7 @@ dependencies = [ "utoipa", "utoipa-scalar", "uuid", + "villar-pso", "which", "zstd", "zune-inflate", @@ -1697,6 +1698,27 @@ dependencies = [ "hybrid-array", ] +[[package]] +name = "csv" +version = "1.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52cd9d68cf7efc6ddfaaee42e7288d3a99d613d4b50f76ce9827ae0c6e14f938" +dependencies = [ + "csv-core", + "itoa", + "ryu", + "serde_core", +] + +[[package]] +name = "csv-core" +version = "0.1.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "704a3c26996a80471189265814dbc2c257598b96b8a7feae2d31ace646bb9782" +dependencies = [ + "memchr", +] + [[package]] name = "ctr" version = "0.9.2" @@ -6109,6 +6131,15 @@ checksum = "29790946404f91d9c5d06f9874efddea1dc06c5efe94541a7d6863108e3a5e4b" dependencies = [ "same-file", "winapi-util", +name = "villar-pso" +version = "0.1.0" +source = "git+https://github.com/frenbox/villar-pso#c73ed71b1110ca86ca0eb6ab9bb01f55c0e997aa" +dependencies = [ + "cc", + "csv", + "rand 0.9.4", + "rayon", + "serde", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index a01faf4f6..7900e4dde 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -62,13 +62,21 @@ zune-inflate = { version = "0.2", default-features = false, features = [ "std", ] } +[features] +default = [] +# The `gpu` feature only controls villar-pso. ONNX GPU is already enabled +# unconditionally on Linux via ort's `cuda` + `load-dynamic` features below, +# and is toggled purely at runtime via `config.gpu.enabled`. +gpu = ["dep:villar-pso"] + # on mac os silicon platform, install ort with the coreml feature [target.'cfg(target_os = "macos")'.dependencies] ort = { version = "=2.0.0-rc.12", features = ["coreml"] } -# otherwise, install ort with the cuda feature +# otherwise, install ort without cuda by default (enabled via gpu feature) [target.'cfg(target_os = "linux")'.dependencies] ort = { version = "=2.0.0-rc.12", features = ["cuda", "load-dynamic"] } +villar-pso = { git = "https://github.com/frenbox/villar-pso", features = ["cuda"], optional = true } [workspace] members = ["apache-avro-macros"] diff --git a/src/enrichment/ztf.rs b/src/enrichment/ztf.rs index 80478ad68..fae756244 100644 --- a/src/enrichment/ztf.rs +++ b/src/enrichment/ztf.rs @@ -19,7 +19,17 @@ use mongodb::bson::{doc, Document}; use mongodb::options::{UpdateOneModel, WriteModel}; use serde::{Deserialize, Deserializer}; use std::sync::Arc; +#[cfg(feature = "gpu")] +use std::sync::atomic::{AtomicI32, Ordering}; use tracing::{instrument, trace, warn}; +#[cfg(feature = "gpu")] +use tracing::info; +#[cfg(feature = "gpu")] +use villar_pso::gpu::{GpuBatchData, GpuContext, SourceData}; + +/// Atomic counter for round-robin GPU device assignment across worker threads. +#[cfg(feature = "gpu")] +static GPU_DEVICE_COUNTER: AtomicI32 = AtomicI32::new(0); #[serdavro] #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] @@ -394,6 +404,26 @@ pub struct ZtfEnrichmentWorker { models: Option>, babamul: Option, gpu_enabled: bool, + /// Per-worker villar-pso GPU context. `Some` only when `config.gpu.enabled` + /// is true and the binary was built with the `gpu` feature. Workers are + /// round-robin-assigned a device from `config.gpu.device_ids` at init time. + #[cfg(feature = "gpu")] + gpu_ctx: Option, +} + +#[cfg(feature = "gpu")] +fn to_villar_photometry(p: &PhotometryMag) -> Option { + let band = match p.band { + Band::G => villar_pso::Band::G, + Band::R => villar_pso::Band::R, + _ => return None, + }; + Some(villar_pso::PhotometryMag { + time: p.time, + mag: p.mag, + mag_err: p.mag_err, + band, + }) } #[async_trait::async_trait] @@ -427,6 +457,31 @@ impl EnrichmentWorker for ZtfEnrichmentWorker { None => Some(SharedModels::load(None)?), }; + // Assign a GPU device to this worker via round-robin over `config.gpu.device_ids`. + // Only initialize a villar-pso GpuContext when GPU usage is actually enabled at + // runtime; otherwise leave it `None` so this worker never touches CUDA. + #[cfg(feature = "gpu")] + let gpu_ctx: Option = if config.gpu.enabled { + let device_ids = &config.gpu.device_ids; + if device_ids.is_empty() { + return Err(EnrichmentWorkerError::ConfigurationError( + "config.gpu.enabled is true but config.gpu.device_ids is empty".to_string(), + )); + } + let idx = (GPU_DEVICE_COUNTER.fetch_add(1, Ordering::Relaxed) as usize) + % device_ids.len(); + let device_id = device_ids[idx]; + info!(device_id, "initializing villar-pso GPU context"); + Some(GpuContext::new(device_id).map_err(|e| { + EnrichmentWorkerError::ConfigurationError(format!( + "villar-pso GPU init failed for device {}: {}", + device_id, e + )) + })?) + } else { + None + }; + Ok(ZtfEnrichmentWorker { input_queue, output_queue, @@ -437,6 +492,8 @@ impl EnrichmentWorker for ZtfEnrichmentWorker { models, babamul, gpu_enabled: config.gpu.enabled, + #[cfg(feature = "gpu")] + gpu_ctx, }) } @@ -491,13 +548,26 @@ impl EnrichmentWorker for ZtfEnrichmentWorker { let mut enriched_alerts: Vec = Vec::new(); let mut work_items: Vec = Vec::with_capacity(alerts.len()); + #[cfg(feature = "gpu")] + let mut villar_inputs: Vec<(i64, Vec)> = Vec::new(); for alert in alerts { let candid = alert.candid; let cutouts = candid_to_cutouts .remove(&candid) .ok_or_else(|| EnrichmentWorkerError::MissingCutouts(candid))?; + + // Compute numerical and boolean features from lightcurve and candidate analysis + #[cfg(feature = "gpu")] + let (properties, all_bands_properties, programid, lightcurve) = + self.get_alert_properties(&alert).await?; + #[cfg(not(feature = "gpu"))] let (properties, all_bands_properties, programid, _lightcurve) = self.get_alert_properties(&alert).await?; + #[cfg(feature = "gpu")] + if self.gpu_ctx.is_some() { + villar_inputs.push((candid, lightcurve)); + } + work_items.push(AlertWork { candid, programid, @@ -550,6 +620,95 @@ impl EnrichmentWorker for ZtfEnrichmentWorker { let _ = self.client.bulk_write(updates).await?.modified_count; + // GPU batch Villar light curve fitting — only when a villar-pso GPU + // context was actually initialized for this worker (i.e. config.gpu.enabled + // is true). Otherwise `villar_inputs` is empty and we skip the entire block. + #[cfg(feature = "gpu")] + if let Some(gpu_ctx) = self.gpu_ctx.as_ref() { + // Document whose keys match a successful fit's keys but with all values NaN. + // Written whenever a fit can't be produced (bad photometry or GPU failure). + let nan_set_doc = { + let mut d = doc! { "villar_fit.reduced_chi2": f64::NAN }; + for filt in villar_pso::FILTERS { + for pname in villar_pso::PARAM_NAMES { + d.insert(format!("villar_fit.{}_{}", pname, filt), f64::NAN); + } + } + d + }; + + let alert_collection = &self.alert_collection; + let build_update = |candid: i64, set_doc: Document| { + WriteModel::UpdateOne( + UpdateOneModel::builder() + .namespace(alert_collection.namespace()) + .filter(doc! { "_id": candid }) + .update(doc! { "$set": set_doc }) + .build(), + ) + }; + + // Preprocess each lightcurve; split into fittable sources and + // NaN-update-only candids that failed preprocessing. + let mut villar_updates: Vec = Vec::new(); + let mut fittable: Vec<(i64, SourceData)> = Vec::new(); + for (candid, lc) in &villar_inputs { + let villar_lc: Vec = + lc.iter().filter_map(to_villar_photometry).collect(); + match villar_pso::preprocess_from_photometry(&villar_lc) { + Ok(preproc) => fittable.push(( + *candid, + SourceData { + name: candid.to_string(), + data: preproc, + }, + )), + Err(e) => { + trace!(candid, "skipping Villar fit: {}", e); + villar_updates.push(build_update(*candid, nan_set_doc.clone())); + } + } + } + + // Run GPU batch fit on the fittable sources. + if !fittable.is_empty() { + let (candids, sources): (Vec, Vec) = + fittable.into_iter().unzip(); + let source_refs: Vec<&SourceData> = sources.iter().collect(); + let pso_config = villar_pso::PsoConfig::default(); + + match GpuBatchData::new(&source_refs).and_then(|batch| { + gpu_ctx.batch_pso_multi_seed(&batch, &source_refs, &pso_config) + }) { + Ok(results) => { + for (result, candid) in results.iter().zip(candids) { + let mut set_doc = doc! { + "villar_fit.reduced_chi2": result.reduced_chi2, + }; + for (key, val) in &result.params_unnorm.to_named_map() { + set_doc.insert(format!("villar_fit.{}", key), *val); + } + villar_updates.push(build_update(candid, set_doc)); + } + } + Err(e) => { + warn!("GPU Villar batch fitting failed: {}", e); + villar_updates.extend( + candids + .into_iter() + .map(|c| build_update(c, nan_set_doc.clone())), + ); + } + } + } + + if !villar_updates.is_empty() { + if let Err(e) = self.client.bulk_write(villar_updates).await { + warn!("failed to write Villar fit results: {}", e); + } + } + } + // Send to Babamul for batch processing if let Some(babamul) = self.babamul.as_ref() { babamul.process_ztf_alerts(enriched_alerts).await?; From 80b00a83f9be4984f08a6c3b4586f6e05899b755 Mon Sep 17 00:00:00 2001 From: Sushant Sharma Chaudhary Date: Sun, 3 May 2026 21:27:08 -0500 Subject: [PATCH 02/20] adding apple's metal architecture support for villar pso --- Cargo.lock | 90 +++++++++++++++++++++++++++++++++++++++++-- Cargo.toml | 4 +- src/enrichment/ztf.rs | 15 ++++++-- 3 files changed, 102 insertions(+), 7 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 5a5a4781c..2146495b3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1048,6 +1048,12 @@ dependencies = [ "wyz", ] +[[package]] +name = "block" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d8c1fef690941d3e7788d328517591fecc684c084084702d6ff1641e993699a" + [[package]] name = "block-buffer" version = "0.10.4" @@ -1583,6 +1589,17 @@ version = "0.8.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" +[[package]] +name = "core-graphics-types" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "45390e6114f68f718cc7a830514a96f903cccd70d02a8f6d9f643ac4ba45afaf" +dependencies = [ + "bitflags 1.3.2", + "core-foundation 0.9.4", + "libc", +] + [[package]] name = "cpufeatures" version = "0.2.17" @@ -2167,7 +2184,28 @@ version = "0.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f6f339eb8adc052cd2ca78910fda869aefa38d22d5cb648e6485e4d3fc06f3b1" dependencies = [ - "foreign-types-shared", + "foreign-types-shared 0.1.1", +] + +[[package]] +name = "foreign-types" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d737d9aa519fb7b749cbc3b962edcf310a8dd1f4b67c91c4f83975dbdd17d965" +dependencies = [ + "foreign-types-macros", + "foreign-types-shared 0.3.1", +] + +[[package]] +name = "foreign-types-macros" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1a5c6c585bc94aaf2c7b51dd4c2ba22680844aba4c687be581871a6f518c5742" +dependencies = [ + "proc-macro2", + "quote", + "syn", ] [[package]] @@ -2176,6 +2214,12 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "00b0228411908ca8685dba7fc2cdd70ec9990a6e753e89b6ac91a84c40fbaf4b" +[[package]] +name = "foreign-types-shared" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aa9a19cbb55df58761df49b23516a86d432839add4af60fc256da840f66ed35b" + [[package]] name = "form_urlencoded" version = "1.2.2" @@ -3344,6 +3388,15 @@ dependencies = [ "syn", ] +[[package]] +name = "malloc_buf" +version = "0.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "62bb907fe88d54d8d9ce32a3cceab4218ed2f6b7d35617cafe9adf84e43919cb" +dependencies = [ + "libc", +] + [[package]] name = "mapproj" version = "0.4.0" @@ -3404,6 +3457,21 @@ dependencies = [ "libc", ] +[[package]] +name = "metal" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7ecfd3296f8c56b7c1f6fbac3c71cefa9d78ce009850c45000015f206dc7fa21" +dependencies = [ + "bitflags 2.11.1", + "block", + "core-graphics-types", + "foreign-types 0.5.0", + "log", + "objc", + "paste", +] + [[package]] name = "mime" version = "0.3.17" @@ -3715,6 +3783,15 @@ dependencies = [ "syn", ] +[[package]] +name = "objc" +version = "0.2.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "915b1b472bc21c53464d6c8461c9d3af805ba1ef837e1cac254428f4a77177b1" +dependencies = [ + "malloc_buf", +] + [[package]] name = "object" version = "0.37.3" @@ -3754,7 +3831,7 @@ checksum = "a45fa2aa886c42762255da344f0a0d313e254066c46aad76f300c3d3da62d967" dependencies = [ "bitflags 2.11.1", "cfg-if", - "foreign-types", + "foreign-types 0.3.2", "libc", "openssl-macros", "openssl-sys", @@ -3967,6 +4044,12 @@ dependencies = [ "windows-link", ] +[[package]] +name = "paste" +version = "1.0.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" + [[package]] name = "pathdiff" version = "0.2.3" @@ -6133,10 +6216,11 @@ dependencies = [ "winapi-util", name = "villar-pso" version = "0.1.0" -source = "git+https://github.com/frenbox/villar-pso#c73ed71b1110ca86ca0eb6ab9bb01f55c0e997aa" +source = "git+https://github.com/frenbox/villar-pso#daa2ff747af6eced53979c628183f11cf5b71fd1" dependencies = [ "cc", "csv", + "metal", "rand 0.9.4", "rayon", "serde", diff --git a/Cargo.toml b/Cargo.toml index 7900e4dde..0d97c90da 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -69,9 +69,11 @@ default = [] # and is toggled purely at runtime via `config.gpu.enabled`. gpu = ["dep:villar-pso"] -# on mac os silicon platform, install ort with the coreml feature +# on mac os silicon platform, install ort with the coreml feature, and +# villar-pso with the metal feature for Apple Silicon GPU light curve fitting [target.'cfg(target_os = "macos")'.dependencies] ort = { version = "=2.0.0-rc.12", features = ["coreml"] } +villar-pso = { git = "https://github.com/frenbox/villar-pso", features = ["metal"], optional = true } # otherwise, install ort without cuda by default (enabled via gpu feature) [target.'cfg(target_os = "linux")'.dependencies] diff --git a/src/enrichment/ztf.rs b/src/enrichment/ztf.rs index fae756244..2cb7a530a 100644 --- a/src/enrichment/ztf.rs +++ b/src/enrichment/ztf.rs @@ -24,8 +24,10 @@ use std::sync::atomic::{AtomicI32, Ordering}; use tracing::{instrument, trace, warn}; #[cfg(feature = "gpu")] use tracing::info; -#[cfg(feature = "gpu")] +#[cfg(all(feature = "gpu", target_os = "linux"))] use villar_pso::gpu::{GpuBatchData, GpuContext, SourceData}; +#[cfg(all(feature = "gpu", target_os = "macos"))] +use villar_pso::gpu_metal::{GpuBatchData, GpuContext, SourceData}; /// Atomic counter for round-robin GPU device assignment across worker threads. #[cfg(feature = "gpu")] @@ -459,7 +461,9 @@ impl EnrichmentWorker for ZtfEnrichmentWorker { // Assign a GPU device to this worker via round-robin over `config.gpu.device_ids`. // Only initialize a villar-pso GpuContext when GPU usage is actually enabled at - // runtime; otherwise leave it `None` so this worker never touches CUDA. + // runtime; otherwise leave it `None` so this worker never touches the GPU. + // Backend is selected by target_os: CUDA on Linux, Metal on macOS. The Metal + // backend ignores `device_id` since Apple Silicon exposes a single GPU. #[cfg(feature = "gpu")] let gpu_ctx: Option = if config.gpu.enabled { let device_ids = &config.gpu.device_ids; @@ -677,7 +681,12 @@ impl EnrichmentWorker for ZtfEnrichmentWorker { let source_refs: Vec<&SourceData> = sources.iter().collect(); let pso_config = villar_pso::PsoConfig::default(); - match GpuBatchData::new(&source_refs).and_then(|batch| { + #[cfg(target_os = "linux")] + let batch_result = GpuBatchData::new(&source_refs); + #[cfg(target_os = "macos")] + let batch_result = GpuBatchData::new(gpu_ctx, &source_refs); + + match batch_result.and_then(|batch| { gpu_ctx.batch_pso_multi_seed(&batch, &source_refs, &pso_config) }) { Ok(results) => { From 0a3a2b51531fc9994bab77875d9c0452691fd1b5 Mon Sep 17 00:00:00 2001 From: frenbox Date: Sat, 2 May 2026 17:08:38 -0500 Subject: [PATCH 03/20] adding apple metal implementation of villar fit --- Cargo.toml | 3 ++- src/enrichment/ztf.rs | 21 +++++++++++++++++++++ 2 files changed, 23 insertions(+), 1 deletion(-) diff --git a/Cargo.toml b/Cargo.toml index 0d97c90da..489fdbfe2 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -67,6 +67,7 @@ default = [] # The `gpu` feature only controls villar-pso. ONNX GPU is already enabled # unconditionally on Linux via ort's `cuda` + `load-dynamic` features below, # and is toggled purely at runtime via `config.gpu.enabled`. +# Backend selection for villar-pso is per-OS: cuda on Linux, metal on macOS. gpu = ["dep:villar-pso"] # on mac os silicon platform, install ort with the coreml feature, and @@ -75,7 +76,7 @@ gpu = ["dep:villar-pso"] ort = { version = "=2.0.0-rc.12", features = ["coreml"] } villar-pso = { git = "https://github.com/frenbox/villar-pso", features = ["metal"], optional = true } -# otherwise, install ort without cuda by default (enabled via gpu feature) +# Linux: ort with cuda + load-dynamic, villar-pso with the CUDA backend. [target.'cfg(target_os = "linux")'.dependencies] ort = { version = "=2.0.0-rc.12", features = ["cuda", "load-dynamic"] } villar-pso = { git = "https://github.com/frenbox/villar-pso", features = ["cuda"], optional = true } diff --git a/src/enrichment/ztf.rs b/src/enrichment/ztf.rs index 2cb7a530a..980d18f50 100644 --- a/src/enrichment/ztf.rs +++ b/src/enrichment/ztf.rs @@ -413,6 +413,27 @@ pub struct ZtfEnrichmentWorker { gpu_ctx: Option, } +/// Build a `GpuBatchData` for the active villar-pso backend. +/// +/// The CUDA backend's `GpuBatchData::new` only takes the sources, while the +/// Metal backend additionally requires the context. This shim absorbs the +/// signature difference so the call site can stay backend-agnostic. +#[cfg(all(feature = "gpu", target_os = "linux"))] +fn make_villar_batch( + _ctx: &GpuContext, + sources: &[&SourceData], +) -> Result { + GpuBatchData::new(sources) +} + +#[cfg(all(feature = "gpu", target_os = "macos"))] +fn make_villar_batch( + ctx: &GpuContext, + sources: &[&SourceData], +) -> Result { + GpuBatchData::new(ctx, sources) +} + #[cfg(feature = "gpu")] fn to_villar_photometry(p: &PhotometryMag) -> Option { let band = match p.band { From d0399b0d1071cd2a158da764cb948127e67954b5 Mon Sep 17 00:00:00 2001 From: frenbox Date: Thu, 7 May 2026 08:40:25 -0500 Subject: [PATCH 04/20] created a cudastream and added ML models and Villar PSO onto it --- src/bin/scheduler.rs | 7 +- src/enrichment/models/acai.rs | 11 +- src/enrichment/models/base.rs | 32 ++++- src/enrichment/models/btsbot.rs | 11 +- src/enrichment/models/mod.rs | 133 ++++++++++++++++---- src/enrichment/ztf.rs | 61 +++------ src/gpu/mod.rs | 3 - src/gpu/pool.rs | 216 -------------------------------- src/lib.rs | 1 - 9 files changed, 178 insertions(+), 297 deletions(-) delete mode 100644 src/gpu/mod.rs delete mode 100644 src/gpu/pool.rs diff --git a/src/bin/scheduler.rs b/src/bin/scheduler.rs index a4cf65e67..80fc5ccb9 100644 --- a/src/bin/scheduler.rs +++ b/src/bin/scheduler.rs @@ -46,7 +46,12 @@ fn validate_gpu_inference(device_ids: &[i32]) -> Result<(), Box Result { + /// Load on a specific CUDA device, optionally sharing a compute stream. + /// `cuda_stream` is a `cudaStream_t` (or null) — see [`load_model_on_device`]. + pub fn new_on_device( + path: &str, + device_id: i32, + cuda_stream: *mut std::ffi::c_void, + ) -> Result { Ok(Self { - model: load_model_on_device(path, Some(device_id))?, + model: load_model_on_device(path, Some(device_id), cuda_stream)?, }) } diff --git a/src/enrichment/models/base.rs b/src/enrichment/models/base.rs index 7945ea50e..c95a2de9d 100644 --- a/src/enrichment/models/base.rs +++ b/src/enrichment/models/base.rs @@ -32,7 +32,7 @@ pub enum ModelError { /// GPU usage is controlled by the `BOOM_GPU__ENABLED` environment variable (default: `"true"`). /// When `device_id` is `Some(id)`, that CUDA device is used; otherwise device 0. pub fn load_model(path: &str) -> Result { - load_model_on_device(path, None) + load_model_on_device(path, None, std::ptr::null_mut()) } fn env_truthy(value: &str) -> bool { @@ -42,7 +42,20 @@ fn env_truthy(value: &str) -> bool { ) } -pub fn load_model_on_device(path: &str, device_id: Option) -> Result { +/// Load an ONNX model on a specific device. On Linux+CUDA, `cuda_stream` (a +/// `cudaStream_t` cast to `*mut c_void`) lets the session share its compute +/// stream with other CUDA work — pass `std::ptr::null_mut()` to let ORT +/// allocate its own stream. The stream argument is ignored on macOS. +/// +/// # Safety +/// When non-null, `cuda_stream` must be a valid `cudaStream_t` belonging to +/// `device_id`'s device, and must outlive the returned [`Session`]. +pub fn load_model_on_device( + path: &str, + device_id: Option, + #[cfg_attr(not(target_os = "linux"), allow(unused_variables))] + cuda_stream: *mut std::ffi::c_void, +) -> Result { let mut builder = Session::builder()?; let use_gpu = env::var("BOOM_GPU__ENABLED") @@ -67,11 +80,20 @@ pub fn load_model_on_device(path: &str, device_id: Option) -> Result Result { + /// Load on a specific CUDA device, optionally sharing a compute stream. + /// `cuda_stream` is a `cudaStream_t` (or null) — see [`load_model_on_device`]. + pub fn new_on_device( + path: &str, + device_id: i32, + cuda_stream: *mut std::ffi::c_void, + ) -> Result { Ok(Self { - model: load_model_on_device(path, Some(device_id))?, + model: load_model_on_device(path, Some(device_id), cuda_stream)?, }) } diff --git a/src/enrichment/models/mod.rs b/src/enrichment/models/mod.rs index 806221426..4b39970c0 100644 --- a/src/enrichment/models/mod.rs +++ b/src/enrichment/models/mod.rs @@ -6,6 +6,11 @@ pub use acai::AcaiModel; pub use base::{load_model, load_model_on_device, Model, ModelError}; pub use btsbot::BtsBotModel; +#[cfg(all(feature = "gpu", target_os = "linux"))] +use villar_pso::gpu::{GpuContext, Stream}; +#[cfg(all(feature = "gpu", target_os = "macos"))] +use villar_pso::gpu_metal::GpuContext; + use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; use tracing::info; @@ -17,6 +22,14 @@ const SESSIONS_PER_DEVICE: usize = 1; /// Each model is wrapped in a `Mutex` because `Session::run` requires `&mut self`. /// Concurrent workers will serialize on the mutex, but model weights are loaded /// only once in memory (and on GPU VRAM if using CUDA). +/// +/// On Linux+GPU, all sessions share a single CUDA stream with the villar-pso +/// `GpuContext` so PSO and ONNX inference run on the same stream and avoid +/// the implicit cross-stream barriers of the legacy default stream. +/// +/// **Field ordering matters**: Rust drops struct fields in declaration order, +/// so the stream must be declared AFTER everything that uses it (the ORT +/// sessions and the `GpuContext`) to ensure `cudaStreamDestroy` fires last. pub struct SharedModels { pub acai_h: Mutex, pub acai_n: Mutex, @@ -24,6 +37,14 @@ pub struct SharedModels { pub acai_o: Mutex, pub acai_b: Mutex, pub btsbot: Mutex, + /// Villar-PSO GPU context bound to this device. `None` when running + /// without the `gpu` feature. + #[cfg(feature = "gpu")] + pub gpu_ctx: Option, + /// CUDA stream shared with the ORT sessions above. Must be dropped last + /// — see struct-level docstring. + #[cfg(all(feature = "gpu", target_os = "linux"))] + _stream: Option, } impl std::fmt::Debug for SharedModels { @@ -35,44 +56,110 @@ impl std::fmt::Debug for SharedModels { impl SharedModels { /// Load all ONNX models, optionally on a specific CUDA device. /// Returns an `Arc` for sharing across threads. + /// + /// On Linux+`gpu`, this creates a CUDA stream for `device_id` and binds + /// every ORT session to it via `with_compute_stream`. The same stream is + /// also handed to the villar-pso `GpuContext`, so all GPU work for this + /// device runs on one stream. pub fn load(device_id: Option) -> Result, ModelError> { info!(?device_id, "loading shared ONNX models"); - let models = match device_id { - Some(id) => Self { - acai_h: Mutex::new(AcaiModel::new_on_device( + + // Create the per-device CUDA stream first (Linux+gpu only). The raw + // pointer is shared between ORT sessions and villar-pso. + #[cfg(all(feature = "gpu", target_os = "linux"))] + let stream: Option = match device_id { + Some(id) => Some(Stream::new_on_device(id).map_err(|e| { + ModelError::Ort(ort::Error::new(format!( + "failed to create CUDA stream on device {}: {}", + id, e + ))) + })?), + None => None, + }; + #[cfg(all(feature = "gpu", target_os = "linux"))] + let stream_ptr: *mut std::ffi::c_void = stream + .as_ref() + .map(|s| s.as_ptr()) + .unwrap_or(std::ptr::null_mut()); + #[cfg(not(all(feature = "gpu", target_os = "linux")))] + let stream_ptr: *mut std::ffi::c_void = std::ptr::null_mut(); + + let (acai_h, acai_n, acai_v, acai_o, acai_b, btsbot) = match device_id { + Some(id) => ( + AcaiModel::new_on_device( "data/models/acai_h.d1_dnn_20201130.onnx", id, - )?), - acai_n: Mutex::new(AcaiModel::new_on_device( + stream_ptr, + )?, + AcaiModel::new_on_device( "data/models/acai_n.d1_dnn_20201130.onnx", id, - )?), - acai_v: Mutex::new(AcaiModel::new_on_device( + stream_ptr, + )?, + AcaiModel::new_on_device( "data/models/acai_v.d1_dnn_20201130.onnx", id, - )?), - acai_o: Mutex::new(AcaiModel::new_on_device( + stream_ptr, + )?, + AcaiModel::new_on_device( "data/models/acai_o.d1_dnn_20201130.onnx", id, - )?), - acai_b: Mutex::new(AcaiModel::new_on_device( + stream_ptr, + )?, + AcaiModel::new_on_device( "data/models/acai_b.d1_dnn_20201130.onnx", id, - )?), - btsbot: Mutex::new(BtsBotModel::new_on_device( + stream_ptr, + )?, + BtsBotModel::new_on_device( "data/models/btsbot-v1.0.1.onnx", id, - )?), - }, - None => Self { - acai_h: Mutex::new(AcaiModel::new("data/models/acai_h.d1_dnn_20201130.onnx")?), - acai_n: Mutex::new(AcaiModel::new("data/models/acai_n.d1_dnn_20201130.onnx")?), - acai_v: Mutex::new(AcaiModel::new("data/models/acai_v.d1_dnn_20201130.onnx")?), - acai_o: Mutex::new(AcaiModel::new("data/models/acai_o.d1_dnn_20201130.onnx")?), - acai_b: Mutex::new(AcaiModel::new("data/models/acai_b.d1_dnn_20201130.onnx")?), - btsbot: Mutex::new(BtsBotModel::new("data/models/btsbot-v1.0.1.onnx")?), - }, + stream_ptr, + )?, + ), + None => ( + AcaiModel::new("data/models/acai_h.d1_dnn_20201130.onnx")?, + AcaiModel::new("data/models/acai_n.d1_dnn_20201130.onnx")?, + AcaiModel::new("data/models/acai_v.d1_dnn_20201130.onnx")?, + AcaiModel::new("data/models/acai_o.d1_dnn_20201130.onnx")?, + AcaiModel::new("data/models/acai_b.d1_dnn_20201130.onnx")?, + BtsBotModel::new("data/models/btsbot-v1.0.1.onnx")?, + ), }; + + // Build the villar-pso GpuContext on the same device + stream. + #[cfg(feature = "gpu")] + let gpu_ctx: Option = match device_id { + #[cfg(target_os = "linux")] + Some(id) => Some(GpuContext::new(id, stream_ptr).map_err(|e| { + ModelError::Ort(ort::Error::new(format!( + "villar-pso GPU init failed for device {}: {}", + id, e + ))) + })?), + #[cfg(target_os = "macos")] + Some(id) => Some(GpuContext::new(id).map_err(|e| { + ModelError::Ort(ort::Error::new(format!( + "villar-pso GPU init failed for device {}: {}", + id, e + ))) + })?), + None => None, + }; + + let models = Self { + acai_h: Mutex::new(acai_h), + acai_n: Mutex::new(acai_n), + acai_v: Mutex::new(acai_v), + acai_o: Mutex::new(acai_o), + acai_b: Mutex::new(acai_b), + btsbot: Mutex::new(btsbot), + #[cfg(feature = "gpu")] + gpu_ctx, + #[cfg(all(feature = "gpu", target_os = "linux"))] + _stream: stream, + }; + info!("all ONNX models loaded successfully"); Ok(Arc::new(models)) } diff --git a/src/enrichment/ztf.rs b/src/enrichment/ztf.rs index 980d18f50..a9acfb46c 100644 --- a/src/enrichment/ztf.rs +++ b/src/enrichment/ztf.rs @@ -19,19 +19,13 @@ use mongodb::bson::{doc, Document}; use mongodb::options::{UpdateOneModel, WriteModel}; use serde::{Deserialize, Deserializer}; use std::sync::Arc; -#[cfg(feature = "gpu")] -use std::sync::atomic::{AtomicI32, Ordering}; use tracing::{instrument, trace, warn}; #[cfg(feature = "gpu")] use tracing::info; #[cfg(all(feature = "gpu", target_os = "linux"))] -use villar_pso::gpu::{GpuBatchData, GpuContext, SourceData}; +use villar_pso::gpu::{GpuBatchData, SourceData}; #[cfg(all(feature = "gpu", target_os = "macos"))] -use villar_pso::gpu_metal::{GpuBatchData, GpuContext, SourceData}; - -/// Atomic counter for round-robin GPU device assignment across worker threads. -#[cfg(feature = "gpu")] -static GPU_DEVICE_COUNTER: AtomicI32 = AtomicI32::new(0); +use villar_pso::gpu_metal::{GpuBatchData, SourceData}; #[serdavro] #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] @@ -402,36 +396,12 @@ pub struct ZtfEnrichmentWorker { alert_collection: mongodb::Collection, alert_cutout_storage: CutoutStorage, alert_pipeline: Vec, - /// Shared ONNX models (loaded once, shared across all enrichment workers via Arc). + /// Shared ONNX models (loaded once, shared across all enrichment workers + /// via Arc). On Linux+`gpu` this also owns the per-device CUDA stream and + /// villar-pso `GpuContext` so that PSO and ONNX inference share a stream. models: Option>, babamul: Option, gpu_enabled: bool, - /// Per-worker villar-pso GPU context. `Some` only when `config.gpu.enabled` - /// is true and the binary was built with the `gpu` feature. Workers are - /// round-robin-assigned a device from `config.gpu.device_ids` at init time. - #[cfg(feature = "gpu")] - gpu_ctx: Option, -} - -/// Build a `GpuBatchData` for the active villar-pso backend. -/// -/// The CUDA backend's `GpuBatchData::new` only takes the sources, while the -/// Metal backend additionally requires the context. This shim absorbs the -/// signature difference so the call site can stay backend-agnostic. -#[cfg(all(feature = "gpu", target_os = "linux"))] -fn make_villar_batch( - _ctx: &GpuContext, - sources: &[&SourceData], -) -> Result { - GpuBatchData::new(sources) -} - -#[cfg(all(feature = "gpu", target_os = "macos"))] -fn make_villar_batch( - ctx: &GpuContext, - sources: &[&SourceData], -) -> Result { - GpuBatchData::new(ctx, sources) } #[cfg(feature = "gpu")] @@ -517,8 +487,6 @@ impl EnrichmentWorker for ZtfEnrichmentWorker { models, babamul, gpu_enabled: config.gpu.enabled, - #[cfg(feature = "gpu")] - gpu_ctx, }) } @@ -589,7 +557,12 @@ impl EnrichmentWorker for ZtfEnrichmentWorker { let (properties, all_bands_properties, programid, _lightcurve) = self.get_alert_properties(&alert).await?; #[cfg(feature = "gpu")] - if self.gpu_ctx.is_some() { + if self + .models + .as_ref() + .and_then(|m| m.gpu_ctx.as_ref()) + .is_some() + { villar_inputs.push((candid, lightcurve)); } @@ -645,11 +618,15 @@ impl EnrichmentWorker for ZtfEnrichmentWorker { let _ = self.client.bulk_write(updates).await?.modified_count; - // GPU batch Villar light curve fitting — only when a villar-pso GPU - // context was actually initialized for this worker (i.e. config.gpu.enabled - // is true). Otherwise `villar_inputs` is empty and we skip the entire block. + // GPU batch Villar light curve fitting — only when SharedModels was + // loaded with a GPU device (i.e. config.gpu.enabled is true). + // Otherwise `villar_inputs` is empty and we skip the entire block. #[cfg(feature = "gpu")] - if let Some(gpu_ctx) = self.gpu_ctx.as_ref() { + if let Some(gpu_ctx) = self + .models + .as_ref() + .and_then(|m| m.gpu_ctx.as_ref()) + { // Document whose keys match a successful fit's keys but with all values NaN. // Written whenever a fit can't be produced (bad photometry or GPU failure). let nan_set_doc = { diff --git a/src/gpu/mod.rs b/src/gpu/mod.rs deleted file mode 100644 index 7d9607523..000000000 --- a/src/gpu/mod.rs +++ /dev/null @@ -1,3 +0,0 @@ -mod pool; - -pub use pool::{DeviceContext, DeviceGuard, GpuPool}; diff --git a/src/gpu/pool.rs b/src/gpu/pool.rs deleted file mode 100644 index 2e0c073f7..000000000 --- a/src/gpu/pool.rs +++ /dev/null @@ -1,216 +0,0 @@ -//! GPU resource pool for distributing work across multiple CUDA devices. -//! -//! # Architecture -//! -//! `GpuPool` holds one GPU context per configured CUDA device. Workers -//! call `acquire()` to get a reference to the least-contended device, -//! submit their batch, and release. This maximizes GPU utilization across -//! multi-GPU nodes without dedicated GPU worker threads. -//! -//! # Usage -//! -//! ```ignore -//! let pool = GpuPool::new(&[0, 1, 2, 3])?; -//! -//! // From any enrichment worker thread: -//! let device = pool.acquire(); -//! let results = device.context.batch_pso_multi_bazin(...)?; -//! drop(device); // releases the device back to the pool -//! ``` -//! -//! # Contention model -//! -//! Each device has a `Mutex` that serializes access. When all devices are -//! busy, `acquire()` blocks on the device with the fewest queued waiters. -//! For N GPUs and M workers, contention is reduced by a factor of N vs -//! a single shared mutex. - -use std::sync::atomic::{AtomicUsize, Ordering}; -use std::sync::{Arc, Mutex, MutexGuard}; -use tracing::info; - -/// A single GPU device with its context and contention tracking. -struct GpuDevice { - context: Mutex, - /// Approximate number of threads waiting or holding this device. - /// Used by `acquire()` to pick the least-contended device. - active_count: AtomicUsize, - device_id: i32, -} - -/// The actual GPU context held behind the mutex. -/// Currently a placeholder — will hold `lightcurve_fitting::gpu::GpuContext` -/// once the crate is integrated as a dependency. -pub struct DeviceContext { - pub device_id: i32, - // Future: pub lc_gpu: lightcurve_fitting::gpu::GpuContext, -} - -/// Guard returned by `GpuPool::acquire()`. Holds the mutex lock and -/// decrements the active count on drop. -pub struct DeviceGuard<'a> { - pub context: MutexGuard<'a, DeviceContext>, - active_count: &'a AtomicUsize, -} - -impl<'a> Drop for DeviceGuard<'a> { - fn drop(&mut self) { - self.active_count.fetch_sub(1, Ordering::Relaxed); - } -} - -/// Pool of GPU devices for distributing work across multiple CUDA devices. -pub struct GpuPool { - devices: Vec, -} - -impl GpuPool { - /// Create a new pool with one context per device ID. - pub fn new(device_ids: &[i32]) -> Result, String> { - if device_ids.is_empty() { - return Err("GpuPool requires at least one device ID".to_string()); - } - - let mut devices = Vec::with_capacity(device_ids.len()); - for &id in device_ids { - info!(device_id = id, "initializing GPU device context"); - // Future: let lc_gpu = lightcurve_fitting::gpu::GpuContext::new(id)?; - devices.push(GpuDevice { - context: Mutex::new(DeviceContext { - device_id: id, - // Future: lc_gpu, - }), - active_count: AtomicUsize::new(0), - device_id: id, - }); - } - - info!( - n_devices = devices.len(), - "GPU pool initialized with {} device(s)", - devices.len() - ); - Ok(Arc::new(Self { devices })) - } - - /// Acquire the least-contended GPU device. Blocks until available. - /// - /// The returned `DeviceGuard` holds the device lock. Drop it to release - /// the device back to the pool. - pub fn acquire(&self) -> DeviceGuard<'_> { - // Find the device with the lowest active count - let best = self - .devices - .iter() - .min_by_key(|d| d.active_count.load(Ordering::Relaxed)) - .expect("GpuPool is non-empty by construction"); - - // Increment before locking so other threads see the contention - best.active_count.fetch_add(1, Ordering::Relaxed); - let guard = best.context.lock().unwrap(); - - DeviceGuard { - context: guard, - active_count: &best.active_count, - } - } - - /// Number of devices in the pool. - pub fn len(&self) -> usize { - self.devices.len() - } - - /// Returns true if the pool has no devices (should never happen after construction). - pub fn is_empty(&self) -> bool { - self.devices.is_empty() - } - - /// Get the device IDs in the pool. - pub fn device_ids(&self) -> Vec { - self.devices.iter().map(|d| d.device_id).collect() - } -} - -#[cfg(test)] -mod tests { - use super::*; - use std::sync::Arc; - - #[test] - fn test_pool_creation() { - let pool = GpuPool::new(&[0]).unwrap(); - assert_eq!(pool.len(), 1); - assert!(!pool.is_empty()); - assert_eq!(pool.device_ids(), vec![0]); - } - - #[test] - fn test_pool_multi_device() { - let pool = GpuPool::new(&[0, 1, 2, 3]).unwrap(); - assert_eq!(pool.len(), 4); - assert_eq!(pool.device_ids(), vec![0, 1, 2, 3]); - } - - #[test] - fn test_pool_empty_fails() { - assert!(GpuPool::new(&[]).is_err()); - } - - #[test] - fn test_acquire_returns_device() { - let pool = GpuPool::new(&[5]).unwrap(); - let guard = pool.acquire(); - assert_eq!(guard.context.device_id, 5); - } - - #[test] - fn test_acquire_least_contended() { - let pool = GpuPool::new(&[0, 1]).unwrap(); - - // Hold device — the next acquire should pick the other one - let guard1 = pool.acquire(); - let id1 = guard1.context.device_id; - - let guard2 = pool.acquire(); - let id2 = guard2.context.device_id; - - // With 2 devices, they should be different - assert_ne!(id1, id2); - } - - #[test] - fn test_acquire_from_multiple_threads() { - let pool = Arc::new(GpuPool::new(&[0, 1, 2, 3]).unwrap()); - let mut handles = Vec::new(); - - for _ in 0..8 { - let pool = Arc::clone(&pool); - handles.push(std::thread::spawn(move || { - let guard = pool.acquire(); - let id = guard.context.device_id; - // Simulate some work - std::thread::sleep(std::time::Duration::from_millis(10)); - drop(guard); - id - })); - } - - let ids: Vec = handles.into_iter().map(|h| h.join().unwrap()).collect(); - // All IDs should be valid device IDs - for id in &ids { - assert!([0, 1, 2, 3].contains(id)); - } - } - - #[test] - fn test_active_count_decrements_on_drop() { - let pool = GpuPool::new(&[0]).unwrap(); - - { - let _guard = pool.acquire(); - assert_eq!(pool.devices[0].active_count.load(Ordering::Relaxed), 1); - } - // Guard dropped - assert_eq!(pool.devices[0].active_count.load(Ordering::Relaxed), 0); - } -} diff --git a/src/lib.rs b/src/lib.rs index 143ecb62f..0f687bdc7 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -4,7 +4,6 @@ pub mod api; pub mod conf; pub mod enrichment; pub mod filter; -pub mod gpu; pub mod kafka; pub mod scheduler; pub mod utils; From 6e43ca04fc328086035efce94f8560f3e8793e61 Mon Sep 17 00:00:00 2001 From: frenbox Date: Fri, 8 May 2026 09:52:18 -0500 Subject: [PATCH 05/20] updating cargo for latest villar pso --- Cargo.lock | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Cargo.lock b/Cargo.lock index 2146495b3..a95c399ce 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6216,7 +6216,7 @@ dependencies = [ "winapi-util", name = "villar-pso" version = "0.1.0" -source = "git+https://github.com/frenbox/villar-pso#daa2ff747af6eced53979c628183f11cf5b71fd1" +source = "git+https://github.com/frenbox/villar-pso#f7187b98b6cb28adaa3b8a29eb4ff3c2718a02a0" dependencies = [ "cc", "csv", From 1e3fe45ca0f3eae629428cd18119a8330e71beae Mon Sep 17 00:00:00 2001 From: frenbox Date: Fri, 8 May 2026 09:55:27 -0500 Subject: [PATCH 06/20] minor readme changes --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 40e86693e..9a0e51d5d 100644 --- a/README.md +++ b/README.md @@ -17,7 +17,7 @@ BOOM is an alert broker. What sets it apart from other alert brokers is that it 1. The `Kafka` consumer(s), reading alerts from astronomical surveys' `Kafka` topics to transfer them to `Redis`/`Valkey` in-memory queues. 2. The Alert Ingestion workers, reading alerts from the `Redis`/`Valkey` queues, responsible of formatting them to BSON documents, and enriching them with crossmatches from archival astronomical catalogs and other surveys before writing the formatted alert packets to a `MongoDB` database. -3. The enrichment workers, running alerts through a series of enrichment classifiers, and writing the results back to the `MongoDB` database. +3. The enrichment workers, running alerts through a series of enrichment classifiers (ML inference) and per-alert light-curve fitting (Villar fits, GPU-accelerated when enabled), and writing the results back to the `MongoDB` database. 4. The Filter workers, running user-defined filters on the alerts, and sending the results to Kafka topics for other services to consume. Workers are managed by a Scheduler that can spawn or kill workers of each type. From b29e3043e72ccc290968fab20f135d672e10c85d Mon Sep 17 00:00:00 2001 From: frenbox Date: Fri, 8 May 2026 10:14:56 -0500 Subject: [PATCH 07/20] cargo fmt --- src/enrichment/models/mod.rs | 6 +----- src/enrichment/ztf.rs | 15 +++++---------- 2 files changed, 6 insertions(+), 15 deletions(-) diff --git a/src/enrichment/models/mod.rs b/src/enrichment/models/mod.rs index 4b39970c0..1e389bebc 100644 --- a/src/enrichment/models/mod.rs +++ b/src/enrichment/models/mod.rs @@ -111,11 +111,7 @@ impl SharedModels { id, stream_ptr, )?, - BtsBotModel::new_on_device( - "data/models/btsbot-v1.0.1.onnx", - id, - stream_ptr, - )?, + BtsBotModel::new_on_device("data/models/btsbot-v1.0.1.onnx", id, stream_ptr)?, ), None => ( AcaiModel::new("data/models/acai_h.d1_dnn_20201130.onnx")?, diff --git a/src/enrichment/ztf.rs b/src/enrichment/ztf.rs index a9acfb46c..2555071ac 100644 --- a/src/enrichment/ztf.rs +++ b/src/enrichment/ztf.rs @@ -19,9 +19,9 @@ use mongodb::bson::{doc, Document}; use mongodb::options::{UpdateOneModel, WriteModel}; use serde::{Deserialize, Deserializer}; use std::sync::Arc; -use tracing::{instrument, trace, warn}; #[cfg(feature = "gpu")] use tracing::info; +use tracing::{instrument, trace, warn}; #[cfg(all(feature = "gpu", target_os = "linux"))] use villar_pso::gpu::{GpuBatchData, SourceData}; #[cfg(all(feature = "gpu", target_os = "macos"))] @@ -463,8 +463,8 @@ impl EnrichmentWorker for ZtfEnrichmentWorker { "config.gpu.enabled is true but config.gpu.device_ids is empty".to_string(), )); } - let idx = (GPU_DEVICE_COUNTER.fetch_add(1, Ordering::Relaxed) as usize) - % device_ids.len(); + let idx = + (GPU_DEVICE_COUNTER.fetch_add(1, Ordering::Relaxed) as usize) % device_ids.len(); let device_id = device_ids[idx]; info!(device_id, "initializing villar-pso GPU context"); Some(GpuContext::new(device_id).map_err(|e| { @@ -622,11 +622,7 @@ impl EnrichmentWorker for ZtfEnrichmentWorker { // loaded with a GPU device (i.e. config.gpu.enabled is true). // Otherwise `villar_inputs` is empty and we skip the entire block. #[cfg(feature = "gpu")] - if let Some(gpu_ctx) = self - .models - .as_ref() - .and_then(|m| m.gpu_ctx.as_ref()) - { + if let Some(gpu_ctx) = self.models.as_ref().and_then(|m| m.gpu_ctx.as_ref()) { // Document whose keys match a successful fit's keys but with all values NaN. // Written whenever a fit can't be produced (bad photometry or GPU failure). let nan_set_doc = { @@ -674,8 +670,7 @@ impl EnrichmentWorker for ZtfEnrichmentWorker { // Run GPU batch fit on the fittable sources. if !fittable.is_empty() { - let (candids, sources): (Vec, Vec) = - fittable.into_iter().unzip(); + let (candids, sources): (Vec, Vec) = fittable.into_iter().unzip(); let source_refs: Vec<&SourceData> = sources.iter().collect(); let pso_config = villar_pso::PsoConfig::default(); From b8eebc714cee715e7306b78525848b20b825184b Mon Sep 17 00:00:00 2001 From: frenbox Date: Sun, 10 May 2026 08:30:58 -0500 Subject: [PATCH 08/20] removing GPU context (dead code) --- src/enrichment/ztf.rs | 32 -------------------------------- 1 file changed, 32 deletions(-) diff --git a/src/enrichment/ztf.rs b/src/enrichment/ztf.rs index 2555071ac..fddd8d303 100644 --- a/src/enrichment/ztf.rs +++ b/src/enrichment/ztf.rs @@ -19,8 +19,6 @@ use mongodb::bson::{doc, Document}; use mongodb::options::{UpdateOneModel, WriteModel}; use serde::{Deserialize, Deserializer}; use std::sync::Arc; -#[cfg(feature = "gpu")] -use tracing::info; use tracing::{instrument, trace, warn}; #[cfg(all(feature = "gpu", target_os = "linux"))] use villar_pso::gpu::{GpuBatchData, SourceData}; @@ -450,33 +448,6 @@ impl EnrichmentWorker for ZtfEnrichmentWorker { None => Some(SharedModels::load(None)?), }; - // Assign a GPU device to this worker via round-robin over `config.gpu.device_ids`. - // Only initialize a villar-pso GpuContext when GPU usage is actually enabled at - // runtime; otherwise leave it `None` so this worker never touches the GPU. - // Backend is selected by target_os: CUDA on Linux, Metal on macOS. The Metal - // backend ignores `device_id` since Apple Silicon exposes a single GPU. - #[cfg(feature = "gpu")] - let gpu_ctx: Option = if config.gpu.enabled { - let device_ids = &config.gpu.device_ids; - if device_ids.is_empty() { - return Err(EnrichmentWorkerError::ConfigurationError( - "config.gpu.enabled is true but config.gpu.device_ids is empty".to_string(), - )); - } - let idx = - (GPU_DEVICE_COUNTER.fetch_add(1, Ordering::Relaxed) as usize) % device_ids.len(); - let device_id = device_ids[idx]; - info!(device_id, "initializing villar-pso GPU context"); - Some(GpuContext::new(device_id).map_err(|e| { - EnrichmentWorkerError::ConfigurationError(format!( - "villar-pso GPU init failed for device {}: {}", - device_id, e - )) - })?) - } else { - None - }; - Ok(ZtfEnrichmentWorker { input_queue, output_queue, @@ -674,9 +645,6 @@ impl EnrichmentWorker for ZtfEnrichmentWorker { let source_refs: Vec<&SourceData> = sources.iter().collect(); let pso_config = villar_pso::PsoConfig::default(); - #[cfg(target_os = "linux")] - let batch_result = GpuBatchData::new(&source_refs); - #[cfg(target_os = "macos")] let batch_result = GpuBatchData::new(gpu_ctx, &source_refs); match batch_result.and_then(|batch| { From dff8966db8661d4b6231277e9d735c618886c31f Mon Sep 17 00:00:00 2001 From: frenbox Date: Wed, 20 May 2026 07:14:50 -0500 Subject: [PATCH 09/20] fixed batch size to 750 --- src/enrichment/base.rs | 12 +++ src/enrichment/models/base.rs | 14 +++- src/enrichment/ztf.rs | 143 +++++++++++++++++++++------------- 3 files changed, 113 insertions(+), 56 deletions(-) diff --git a/src/enrichment/base.rs b/src/enrichment/base.rs index 5e92f6eaa..a52ca9db7 100644 --- a/src/enrichment/base.rs +++ b/src/enrichment/base.rs @@ -223,6 +223,7 @@ pub async fn run_enrichment_worker( } ACTIVE.add(1, &active_attrs); +<<<<<<< HEAD let candids: Vec = retry_transient( "valkey_rpop", DEFAULT_MAX_RETRIES, @@ -243,6 +244,17 @@ pub async fn run_enrichment_worker( ACTIVE.add(-1, &active_attrs); BATCH_PROCESSED.add(1, &input_error_attrs); })?; +======= + let candids: Vec = con + // Capped at 750 to match GPU_INFER_SHAPE; shape 1000 OOMs the GPU. + // .rpop::<&str, Vec>(&input_queue, NonZero::new(1000)) + .rpop::<&str, Vec>(&input_queue, NonZero::new(750)) + .await + .inspect_err(|_| { + ACTIVE.add(-1, &active_attrs); + BATCH_PROCESSED.add(1, &input_error_attrs); + })?; +>>>>>>> 31bdd88 (fixed batch size to 750) if candids.is_empty() { ACTIVE.add(-1, &active_attrs); diff --git a/src/enrichment/models/base.rs b/src/enrichment/models/base.rs index c95a2de9d..e66cdb307 100644 --- a/src/enrichment/models/base.rs +++ b/src/enrichment/models/base.rs @@ -82,7 +82,19 @@ pub fn load_model_on_device( #[cfg(target_os = "linux")] let cuda_ep = { - let mut ep = ort::ep::CUDAExecutionProvider::default().with_device_id(dev); + // `with_conv_max_workspace(false)` caps the cuDNN conv + // algorithm-search workspace at 32 MB, shrinking the per-shape + // scratch buffers the dynamic batch dimension would otherwise grab. + // + // NOTE: `arena_extend_strategy = SameAsRequested` was tried here + // and removed. With the dynamic-batch models its exact-sized + // extensions wreck BFC arena reuse, so GPU memory climbs + // monotonically every batch and OOMs. The default `kNextPowerOfTwo` + // allocates reusable power-of-two blocks and the arena stabilizes. + // .with_arena_extend_strategy(ort::ep::ArenaExtendStrategy::SameAsRequested) + let mut ep = ort::ep::CUDAExecutionProvider::default() + .with_device_id(dev) + .with_conv_max_workspace(false); if !cuda_stream.is_null() { // Safety: caller guarantees the stream is valid for `dev` // and outlives the session (see fn-level safety comment). diff --git a/src/enrichment/ztf.rs b/src/enrichment/ztf.rs index fddd8d303..3c6b97850 100644 --- a/src/enrichment/ztf.rs +++ b/src/enrichment/ztf.rs @@ -19,7 +19,7 @@ use mongodb::bson::{doc, Document}; use mongodb::options::{UpdateOneModel, WriteModel}; use serde::{Deserialize, Deserializer}; use std::sync::Arc; -use tracing::{instrument, trace, warn}; +use tracing::{info, instrument, trace, warn}; #[cfg(all(feature = "gpu", target_os = "linux"))] use villar_pso::gpu::{GpuBatchData, SourceData}; #[cfg(all(feature = "gpu", target_os = "macos"))] @@ -548,12 +548,18 @@ impl EnrichmentWorker for ZtfEnrichmentWorker { } // Run ML classification using shared models + let classify_start = std::time::Instant::now(); let classifications_list: Vec> = if let Some(ref models) = self.models { self.classify(&models, &work_items)? } else { vec![None; work_items.len()] }; + info!( + batch_size = work_items.len(), + classify_ms = classify_start.elapsed().as_millis() as u64, + "ML classification complete" + ); for (item, classifications) in work_items.into_iter().zip(classifications_list) { let update_alert_document = if let Some(ref cls) = classifications { @@ -645,11 +651,19 @@ impl EnrichmentWorker for ZtfEnrichmentWorker { let source_refs: Vec<&SourceData> = sources.iter().collect(); let pso_config = villar_pso::PsoConfig::default(); + let villar_start = std::time::Instant::now(); let batch_result = GpuBatchData::new(gpu_ctx, &source_refs); - match batch_result.and_then(|batch| { + let fit_result = batch_result.and_then(|batch| { gpu_ctx.batch_pso_multi_seed(&batch, &source_refs, &pso_config) - }) { + }); + info!( + fittable = source_refs.len(), + villar_ms = villar_start.elapsed().as_millis() as u64, + "Villar PSO batch complete" + ); + + match fit_result { Ok(results) => { for (result, candid) in results.iter().zip(candids) { let mut set_doc = doc! { @@ -873,6 +887,18 @@ impl ZtfEnrichmentWorker { Ok(results) } + /// Fixed ONNX batch dimension for GPU inference. ORT builds — and never + /// releases — a separate memory plan per distinct input shape, so a + /// varying batch size makes GPU memory spike. Every inference is run at + /// exactly this many rows: partial batches are zero-padded and the padding + /// outputs discarded, so ORT only ever sees one shape and the arena is + /// stable. An RPOP batch larger than this is split into several fixed-size + /// chunks; only the final chunk carries any zero-padding. + // 1000 was tried and OOMs (~15.7 GB footprint exceeds the 16 GB card); + // 750 is stable at ~10.3 GB. + // const GPU_INFER_SHAPE: usize = 1000; + const GPU_INFER_SHAPE: usize = 750; + fn classify_gpu_batch( models: &SharedModels, work_items: &[AlertWork], @@ -926,61 +952,68 @@ impl ZtfEnrichmentWorker { return Ok(results); } - let mut triplet = ndarray::Array::zeros((selected_indices.len(), 63, 63, 3)); - let mut metadata = ndarray::Array::zeros((selected_indices.len(), 25)); - let mut btsbot_metadata = ndarray::Array::zeros((selected_indices.len(), 25)); - - for (row, idx) in selected_indices.iter().enumerate() { - let tpos = *triplet_pos.get(idx).expect("triplet position missing"); - let apos = *acai_pos.get(idx).expect("acai position missing"); - let bpos = *bts_pos.get(idx).expect("bts position missing"); - - triplet - .slice_mut(ndarray::s![row, .., .., ..]) - .assign(&triplet_all.slice(ndarray::s![tpos, .., .., ..])); - metadata.row_mut(row).assign(&acai_metadata_all.row(apos)); - btsbot_metadata - .row_mut(row) - .assign(&bts_metadata_all.row(bpos)); - } + // Run inference in fixed-size chunks so ORT always sees the same + // input shape (see `GPU_INFER_SHAPE`). The final chunk is zero-padded + // up to the fixed size; padding rows produce scores that are ignored. + for chunk in selected_indices.chunks(Self::GPU_INFER_SHAPE) { + let mut triplet = ndarray::Array::zeros((Self::GPU_INFER_SHAPE, 63, 63, 3)); + let mut metadata = ndarray::Array::zeros((Self::GPU_INFER_SHAPE, 25)); + let mut btsbot_metadata = ndarray::Array::zeros((Self::GPU_INFER_SHAPE, 25)); + + for (row, idx) in chunk.iter().enumerate() { + let tpos = *triplet_pos.get(idx).expect("triplet position missing"); + let apos = *acai_pos.get(idx).expect("acai position missing"); + let bpos = *bts_pos.get(idx).expect("bts position missing"); + + triplet + .slice_mut(ndarray::s![row, .., .., ..]) + .assign(&triplet_all.slice(ndarray::s![tpos, .., .., ..])); + metadata.row_mut(row).assign(&acai_metadata_all.row(apos)); + btsbot_metadata + .row_mut(row) + .assign(&bts_metadata_all.row(bpos)); + } - let acai_h_scores = models.acai_h.lock().unwrap().predict(&metadata, &triplet)?; - let acai_n_scores = models.acai_n.lock().unwrap().predict(&metadata, &triplet)?; - let acai_v_scores = models.acai_v.lock().unwrap().predict(&metadata, &triplet)?; - let acai_o_scores = models.acai_o.lock().unwrap().predict(&metadata, &triplet)?; - let acai_b_scores = models.acai_b.lock().unwrap().predict(&metadata, &triplet)?; - let btsbot_scores = models - .btsbot - .lock() - .unwrap() - .predict(&btsbot_metadata, &triplet)?; - - let expected = selected_indices.len(); - for (name, got) in [ - ("acai_h", acai_h_scores.len()), - ("acai_n", acai_n_scores.len()), - ("acai_v", acai_v_scores.len()), - ("acai_o", acai_o_scores.len()), - ("acai_b", acai_b_scores.len()), - ("btsbot", btsbot_scores.len()), - ] { - if got != expected { - return Err(EnrichmentWorkerError::ConfigurationError(format!( - "model {} returned {} scores for {} inputs", - name, got, expected - ))); + let acai_h_scores = models.acai_h.lock().unwrap().predict(&metadata, &triplet)?; + let acai_n_scores = models.acai_n.lock().unwrap().predict(&metadata, &triplet)?; + let acai_v_scores = models.acai_v.lock().unwrap().predict(&metadata, &triplet)?; + let acai_o_scores = models.acai_o.lock().unwrap().predict(&metadata, &triplet)?; + let acai_b_scores = models.acai_b.lock().unwrap().predict(&metadata, &triplet)?; + let btsbot_scores = models + .btsbot + .lock() + .unwrap() + .predict(&btsbot_metadata, &triplet)?; + + for (name, got) in [ + ("acai_h", acai_h_scores.len()), + ("acai_n", acai_n_scores.len()), + ("acai_v", acai_v_scores.len()), + ("acai_o", acai_o_scores.len()), + ("acai_b", acai_b_scores.len()), + ("btsbot", btsbot_scores.len()), + ] { + if got != Self::GPU_INFER_SHAPE { + return Err(EnrichmentWorkerError::ConfigurationError(format!( + "model {} returned {} scores for {} padded inputs", + name, + got, + Self::GPU_INFER_SHAPE + ))); + } } - } - for (batch_idx, &item_idx) in selected_indices.iter().enumerate() { - results[item_idx] = Some(ZtfAlertClassifications { - acai_h: acai_h_scores[batch_idx], - acai_n: acai_n_scores[batch_idx], - acai_v: acai_v_scores[batch_idx], - acai_o: acai_o_scores[batch_idx], - acai_b: acai_b_scores[batch_idx], - btsbot: btsbot_scores[batch_idx], - }); + // Map only the real rows back; padding rows (chunk.len()..) are dropped. + for (batch_idx, &item_idx) in chunk.iter().enumerate() { + results[item_idx] = Some(ZtfAlertClassifications { + acai_h: acai_h_scores[batch_idx], + acai_n: acai_n_scores[batch_idx], + acai_v: acai_v_scores[batch_idx], + acai_o: acai_o_scores[batch_idx], + acai_b: acai_b_scores[batch_idx], + btsbot: btsbot_scores[batch_idx], + }); + } } Ok(results) From 6be106017b5ae1f46da4f8d319374d8de6dc2f2f Mon Sep 17 00:00:00 2001 From: frenbox Date: Wed, 20 May 2026 08:55:22 -0500 Subject: [PATCH 10/20] removed extra debugging info logs --- src/enrichment/ztf.rs | 20 +++----------------- 1 file changed, 3 insertions(+), 17 deletions(-) diff --git a/src/enrichment/ztf.rs b/src/enrichment/ztf.rs index 3c6b97850..8221a968f 100644 --- a/src/enrichment/ztf.rs +++ b/src/enrichment/ztf.rs @@ -19,7 +19,7 @@ use mongodb::bson::{doc, Document}; use mongodb::options::{UpdateOneModel, WriteModel}; use serde::{Deserialize, Deserializer}; use std::sync::Arc; -use tracing::{info, instrument, trace, warn}; +use tracing::{instrument, trace, warn}; #[cfg(all(feature = "gpu", target_os = "linux"))] use villar_pso::gpu::{GpuBatchData, SourceData}; #[cfg(all(feature = "gpu", target_os = "macos"))] @@ -548,18 +548,12 @@ impl EnrichmentWorker for ZtfEnrichmentWorker { } // Run ML classification using shared models - let classify_start = std::time::Instant::now(); let classifications_list: Vec> = if let Some(ref models) = self.models { self.classify(&models, &work_items)? } else { vec![None; work_items.len()] }; - info!( - batch_size = work_items.len(), - classify_ms = classify_start.elapsed().as_millis() as u64, - "ML classification complete" - ); for (item, classifications) in work_items.into_iter().zip(classifications_list) { let update_alert_document = if let Some(ref cls) = classifications { @@ -651,19 +645,11 @@ impl EnrichmentWorker for ZtfEnrichmentWorker { let source_refs: Vec<&SourceData> = sources.iter().collect(); let pso_config = villar_pso::PsoConfig::default(); - let villar_start = std::time::Instant::now(); let batch_result = GpuBatchData::new(gpu_ctx, &source_refs); - let fit_result = batch_result.and_then(|batch| { + match batch_result.and_then(|batch| { gpu_ctx.batch_pso_multi_seed(&batch, &source_refs, &pso_config) - }); - info!( - fittable = source_refs.len(), - villar_ms = villar_start.elapsed().as_millis() as u64, - "Villar PSO batch complete" - ); - - match fit_result { + }) { Ok(results) => { for (result, candid) in results.iter().zip(candids) { let mut set_doc = doc! { From 774d3101d8ee50862a98a147e9bb8b82a7a65752 Mon Sep 17 00:00:00 2001 From: frenbox Date: Wed, 20 May 2026 14:16:21 -0500 Subject: [PATCH 11/20] made changes to config to include batch size --- config.yaml | 5 +++++ src/conf.rs | 29 ++++++++++++++++++++++++++++- src/enrichment/base.rs | 7 ++++--- src/enrichment/ztf.rs | 42 ++++++++++++++++++++---------------------- 4 files changed, 57 insertions(+), 26 deletions(-) diff --git a/config.yaml b/config.yaml index d4c970197..0240cc30b 100644 --- a/config.yaml +++ b/config.yaml @@ -100,6 +100,11 @@ workers: n_workers: 3 enrichment: n_workers: 3 + # Max alerts pulled from the input queue per worker iteration (RPOP cap). + batch_size: 1000 + # Fixed ONNX batch size for GPU inference. Partial batches are zero-padded + # so ORT only ever sees one shape and the BFC arena stays stable. + gpu_infer_shape: 750 filter: n_workers: 3 refresh_interval_minutes: 15 diff --git a/src/conf.rs b/src/conf.rs index 4de417044..f825c63ac 100644 --- a/src/conf.rs +++ b/src/conf.rs @@ -595,6 +595,33 @@ pub struct WorkerConfig { pub n_workers: usize, } +fn default_enrichment_batch_size() -> usize { + 750 +} + +fn default_gpu_infer_shape() -> usize { + 750 +} + +#[derive(Deserialize, Debug, Clone)] +pub struct EnrichmentWorkerConfig { + pub n_workers: usize, + /// Max alerts pulled from the input queue per worker iteration (the + /// RPOP cap). Decoupled from `gpu_infer_shape`: pops larger than the + /// inference shape are split into multiple inference chunks, with the + /// final chunk zero-padded. + #[serde(default = "default_enrichment_batch_size")] + pub batch_size: usize, + /// Fixed ONNX batch dimension for GPU inference. ORT builds — and never + /// releases — a separate memory plan per distinct input shape, so a + /// varying batch size makes GPU memory spike. Every inference runs at + /// exactly this many rows; partial batches are zero-padded so ORT only + /// ever sees one shape and the arena stays stable. 750 is the proven + /// stable shape on a 16 GB card (~10.3 GB footprint); 1000 OOMs (~15.7 GB). + #[serde(default = "default_gpu_infer_shape")] + pub gpu_infer_shape: usize, +} + fn default_filter_refresh_interval_minutes() -> u64 { 15 } @@ -687,7 +714,7 @@ pub struct SurveyWorkerConfig { #[serde(deserialize_with = "deserialize_command_interval")] pub command_interval: usize, // in milliseconds pub alert: WorkerConfig, - pub enrichment: WorkerConfig, + pub enrichment: EnrichmentWorkerConfig, pub filter: FilterWorkerConfig, } diff --git a/src/enrichment/base.rs b/src/enrichment/base.rs index a52ca9db7..cf4b2dd7d 100644 --- a/src/enrichment/base.rs +++ b/src/enrichment/base.rs @@ -246,9 +246,10 @@ pub async fn run_enrichment_worker( })?; ======= let candids: Vec = con - // Capped at 750 to match GPU_INFER_SHAPE; shape 1000 OOMs the GPU. - // .rpop::<&str, Vec>(&input_queue, NonZero::new(1000)) - .rpop::<&str, Vec>(&input_queue, NonZero::new(750)) + .rpop::<&str, Vec>( + &input_queue, + NonZero::new(worker_config.enrichment.batch_size), + ) .await .inspect_err(|_| { ACTIVE.add(-1, &active_attrs); diff --git a/src/enrichment/ztf.rs b/src/enrichment/ztf.rs index 8221a968f..5b4a09491 100644 --- a/src/enrichment/ztf.rs +++ b/src/enrichment/ztf.rs @@ -400,6 +400,9 @@ pub struct ZtfEnrichmentWorker { models: Option>, babamul: Option, gpu_enabled: bool, + /// Fixed ONNX batch dimension for GPU inference (see + /// [`EnrichmentWorkerConfig::gpu_infer_shape`] in `conf.rs`). + gpu_infer_shape: usize, } #[cfg(feature = "gpu")] @@ -448,6 +451,13 @@ impl EnrichmentWorker for ZtfEnrichmentWorker { None => Some(SharedModels::load(None)?), }; + let gpu_infer_shape = config + .workers + .get(&Survey::Ztf) + .ok_or(EnrichmentWorkerError::WorkerConfigMissing(Survey::Ztf))? + .enrichment + .gpu_infer_shape; + Ok(ZtfEnrichmentWorker { input_queue, output_queue, @@ -458,6 +468,7 @@ impl EnrichmentWorker for ZtfEnrichmentWorker { models, babamul, gpu_enabled: config.gpu.enabled, + gpu_infer_shape, }) } @@ -813,7 +824,7 @@ impl ZtfEnrichmentWorker { work_items: &[AlertWork], ) -> Result>, EnrichmentWorkerError> { if self.gpu_enabled { - return Self::classify_gpu_batch(models, work_items); + return self.classify_gpu_batch(models, work_items); } Self::classify_per_item(models, work_items) @@ -873,19 +884,8 @@ impl ZtfEnrichmentWorker { Ok(results) } - /// Fixed ONNX batch dimension for GPU inference. ORT builds — and never - /// releases — a separate memory plan per distinct input shape, so a - /// varying batch size makes GPU memory spike. Every inference is run at - /// exactly this many rows: partial batches are zero-padded and the padding - /// outputs discarded, so ORT only ever sees one shape and the arena is - /// stable. An RPOP batch larger than this is split into several fixed-size - /// chunks; only the final chunk carries any zero-padding. - // 1000 was tried and OOMs (~15.7 GB footprint exceeds the 16 GB card); - // 750 is stable at ~10.3 GB. - // const GPU_INFER_SHAPE: usize = 1000; - const GPU_INFER_SHAPE: usize = 750; - fn classify_gpu_batch( + &self, models: &SharedModels, work_items: &[AlertWork], ) -> Result>, EnrichmentWorkerError> { @@ -939,12 +939,12 @@ impl ZtfEnrichmentWorker { } // Run inference in fixed-size chunks so ORT always sees the same - // input shape (see `GPU_INFER_SHAPE`). The final chunk is zero-padded + // input shape (`self.gpu_infer_shape`). The final chunk is zero-padded // up to the fixed size; padding rows produce scores that are ignored. - for chunk in selected_indices.chunks(Self::GPU_INFER_SHAPE) { - let mut triplet = ndarray::Array::zeros((Self::GPU_INFER_SHAPE, 63, 63, 3)); - let mut metadata = ndarray::Array::zeros((Self::GPU_INFER_SHAPE, 25)); - let mut btsbot_metadata = ndarray::Array::zeros((Self::GPU_INFER_SHAPE, 25)); + for chunk in selected_indices.chunks(self.gpu_infer_shape) { + let mut triplet = ndarray::Array::zeros((self.gpu_infer_shape, 63, 63, 3)); + let mut metadata = ndarray::Array::zeros((self.gpu_infer_shape, 25)); + let mut btsbot_metadata = ndarray::Array::zeros((self.gpu_infer_shape, 25)); for (row, idx) in chunk.iter().enumerate() { let tpos = *triplet_pos.get(idx).expect("triplet position missing"); @@ -979,12 +979,10 @@ impl ZtfEnrichmentWorker { ("acai_b", acai_b_scores.len()), ("btsbot", btsbot_scores.len()), ] { - if got != Self::GPU_INFER_SHAPE { + if got != self.gpu_infer_shape { return Err(EnrichmentWorkerError::ConfigurationError(format!( "model {} returned {} scores for {} padded inputs", - name, - got, - Self::GPU_INFER_SHAPE + name, got, self.gpu_infer_shape ))); } } From bfe02cd7b4205803b07e0256abadc2e815979d9f Mon Sep 17 00:00:00 2001 From: frenbox Date: Wed, 20 May 2026 14:24:41 -0500 Subject: [PATCH 12/20] fix the cargo.toml --- Cargo.lock | 23 +++++++++++++---------- 1 file changed, 13 insertions(+), 10 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index a95c399ce..0bca2d528 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6200,6 +6200,19 @@ version = "0.9.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" +[[package]] +name = "villar-pso" +version = "0.1.0" +source = "git+https://github.com/frenbox/villar-pso#f7187b98b6cb28adaa3b8a29eb4ff3c2718a02a0" +dependencies = [ + "cc", + "csv", + "metal", + "rand 0.9.4", + "rayon", + "serde", +] + [[package]] name = "vsimd" version = "0.8.0" @@ -6214,16 +6227,6 @@ checksum = "29790946404f91d9c5d06f9874efddea1dc06c5efe94541a7d6863108e3a5e4b" dependencies = [ "same-file", "winapi-util", -name = "villar-pso" -version = "0.1.0" -source = "git+https://github.com/frenbox/villar-pso#f7187b98b6cb28adaa3b8a29eb4ff3c2718a02a0" -dependencies = [ - "cc", - "csv", - "metal", - "rand 0.9.4", - "rayon", - "serde", ] [[package]] From 32716e2bac8c771e1db4c3e9ce9e8b67f543a39e Mon Sep 17 00:00:00 2001 From: frenbox Date: Wed, 20 May 2026 14:32:59 -0500 Subject: [PATCH 13/20] fix the dockerfile gpu --- Dockerfile.gpu | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Dockerfile.gpu b/Dockerfile.gpu index 71823ff73..75acc2fb3 100644 --- a/Dockerfile.gpu +++ b/Dockerfile.gpu @@ -37,7 +37,7 @@ RUN mkdir -p /app/src && \ cargo build --release && rm -rf /app/src COPY ./src /app/src -RUN cargo build --release --bin scheduler --bin migrate_fp_flux --bin migrate_snr --bin reprocess_crossmatch +RUN cargo build --release --features gpu --bin scheduler --bin migrate_fp_flux --bin migrate_snr --bin reprocess_crossmatch FROM nvidia/cuda:12.8.1-cudnn-runtime-ubuntu24.04 AS app From 5777de9005b0815744188a9d6932ddd769717ea1 Mon Sep 17 00:00:00 2001 From: frenbox Date: Thu, 21 May 2026 06:50:31 -0500 Subject: [PATCH 14/20] benchmark fixes --- config.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/config.yaml b/config.yaml index 0240cc30b..c21991145 100644 --- a/config.yaml +++ b/config.yaml @@ -21,7 +21,7 @@ babamul: webapp_url: # Set via BOOM_BABAMUL__WEBAPP_URL environment variable retention_days: 3 gpu: - enabled: false # Set to true to load ONNX models on GPU (CUDA) instead of CPU. + enabled: true # Set to true to load ONNX models on GPU (CUDA) instead of CPU. # Models are loaded once at startup and shared across all enrichment workers. # When false, models are loaded on CPU (BOOM_GPU__ENABLED env var is still respected). device_ids: [0] # CUDA device(s) to use. E.g. [0,1,2,3,4,5,6,7] for 8 GPUs. From 177a9d12f40885421206ae3721edce1efd85a4e4 Mon Sep 17 00:00:00 2001 From: frenbox Date: Fri, 22 May 2026 17:54:22 -0500 Subject: [PATCH 15/20] setting one variable for batch size --- config.yaml | 9 ++++----- src/conf.rs | 23 +++++++---------------- src/enrichment/ztf.rs | 26 +++++++++++++------------- 3 files changed, 24 insertions(+), 34 deletions(-) diff --git a/config.yaml b/config.yaml index c21991145..2fcbff53b 100644 --- a/config.yaml +++ b/config.yaml @@ -100,11 +100,10 @@ workers: n_workers: 3 enrichment: n_workers: 3 - # Max alerts pulled from the input queue per worker iteration (RPOP cap). - batch_size: 1000 - # Fixed ONNX batch size for GPU inference. Partial batches are zero-padded - # so ORT only ever sees one shape and the BFC arena stays stable. - gpu_infer_shape: 750 + # Alerts per enrichment batch: the queue RPOP cap AND the fixed ONNX + # inference shape. Partial batches are zero-padded so ORT sees one shape + # and the GPU memory arena stays stable. 750 is stable for both VRAM => 12GB. + batch_size: 750 filter: n_workers: 3 refresh_interval_minutes: 15 diff --git a/src/conf.rs b/src/conf.rs index f825c63ac..5c316ad4c 100644 --- a/src/conf.rs +++ b/src/conf.rs @@ -599,27 +599,18 @@ fn default_enrichment_batch_size() -> usize { 750 } -fn default_gpu_infer_shape() -> usize { - 750 -} - #[derive(Deserialize, Debug, Clone)] pub struct EnrichmentWorkerConfig { pub n_workers: usize, - /// Max alerts pulled from the input queue per worker iteration (the - /// RPOP cap). Decoupled from `gpu_infer_shape`: pops larger than the - /// inference shape are split into multiple inference chunks, with the - /// final chunk zero-padded. + /// Alerts processed per enrichment batch. Serves two roles at once: the + /// queue RPOP cap (max alerts pulled per worker iteration) and the fixed + /// ONNX batch dimension. Every GPU inference runs at exactly this many + /// rows — partial batches are zero-padded — so ORT builds a single memory + /// plan and the BFC arena stays stable instead of growing per distinct + /// input shape. 750 is the proven stable shape on a 16 GB card + /// (~10.3 GB footprint); 1000 OOMs (~15.7 GB). #[serde(default = "default_enrichment_batch_size")] pub batch_size: usize, - /// Fixed ONNX batch dimension for GPU inference. ORT builds — and never - /// releases — a separate memory plan per distinct input shape, so a - /// varying batch size makes GPU memory spike. Every inference runs at - /// exactly this many rows; partial batches are zero-padded so ORT only - /// ever sees one shape and the arena stays stable. 750 is the proven - /// stable shape on a 16 GB card (~10.3 GB footprint); 1000 OOMs (~15.7 GB). - #[serde(default = "default_gpu_infer_shape")] - pub gpu_infer_shape: usize, } fn default_filter_refresh_interval_minutes() -> u64 { diff --git a/src/enrichment/ztf.rs b/src/enrichment/ztf.rs index 5b4a09491..96d495572 100644 --- a/src/enrichment/ztf.rs +++ b/src/enrichment/ztf.rs @@ -400,9 +400,9 @@ pub struct ZtfEnrichmentWorker { models: Option>, babamul: Option, gpu_enabled: bool, - /// Fixed ONNX batch dimension for GPU inference (see - /// [`EnrichmentWorkerConfig::gpu_infer_shape`] in `conf.rs`). - gpu_infer_shape: usize, + /// Alerts per batch — also the fixed ONNX inference shape (see + /// [`EnrichmentWorkerConfig::batch_size`] in `conf.rs`). + batch_size: usize, } #[cfg(feature = "gpu")] @@ -451,12 +451,12 @@ impl EnrichmentWorker for ZtfEnrichmentWorker { None => Some(SharedModels::load(None)?), }; - let gpu_infer_shape = config + let batch_size = config .workers .get(&Survey::Ztf) .ok_or(EnrichmentWorkerError::WorkerConfigMissing(Survey::Ztf))? .enrichment - .gpu_infer_shape; + .batch_size; Ok(ZtfEnrichmentWorker { input_queue, @@ -468,7 +468,7 @@ impl EnrichmentWorker for ZtfEnrichmentWorker { models, babamul, gpu_enabled: config.gpu.enabled, - gpu_infer_shape, + batch_size, }) } @@ -939,12 +939,12 @@ impl ZtfEnrichmentWorker { } // Run inference in fixed-size chunks so ORT always sees the same - // input shape (`self.gpu_infer_shape`). The final chunk is zero-padded + // input shape (`self.batch_size`). The final chunk is zero-padded // up to the fixed size; padding rows produce scores that are ignored. - for chunk in selected_indices.chunks(self.gpu_infer_shape) { - let mut triplet = ndarray::Array::zeros((self.gpu_infer_shape, 63, 63, 3)); - let mut metadata = ndarray::Array::zeros((self.gpu_infer_shape, 25)); - let mut btsbot_metadata = ndarray::Array::zeros((self.gpu_infer_shape, 25)); + for chunk in selected_indices.chunks(self.batch_size) { + let mut triplet = ndarray::Array::zeros((self.batch_size, 63, 63, 3)); + let mut metadata = ndarray::Array::zeros((self.batch_size, 25)); + let mut btsbot_metadata = ndarray::Array::zeros((self.batch_size, 25)); for (row, idx) in chunk.iter().enumerate() { let tpos = *triplet_pos.get(idx).expect("triplet position missing"); @@ -979,10 +979,10 @@ impl ZtfEnrichmentWorker { ("acai_b", acai_b_scores.len()), ("btsbot", btsbot_scores.len()), ] { - if got != self.gpu_infer_shape { + if got != self.batch_size { return Err(EnrichmentWorkerError::ConfigurationError(format!( "model {} returned {} scores for {} padded inputs", - name, got, self.gpu_infer_shape + name, got, self.batch_size ))); } } From cd6922285adb9ec27ada920c19e3cbca814d19a6 Mon Sep 17 00:00:00 2001 From: frenbox Date: Sat, 23 May 2026 05:00:33 -0500 Subject: [PATCH 16/20] putting the gpu config to false as default --- config.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/config.yaml b/config.yaml index 2fcbff53b..7cfdc822e 100644 --- a/config.yaml +++ b/config.yaml @@ -21,7 +21,7 @@ babamul: webapp_url: # Set via BOOM_BABAMUL__WEBAPP_URL environment variable retention_days: 3 gpu: - enabled: true # Set to true to load ONNX models on GPU (CUDA) instead of CPU. + enabled: false # Set to true to load ONNX models on GPU (CUDA) instead of CPU. # Models are loaded once at startup and shared across all enrichment workers. # When false, models are loaded on CPU (BOOM_GPU__ENABLED env var is still respected). device_ids: [0] # CUDA device(s) to use. E.g. [0,1,2,3,4,5,6,7] for 8 GPUs. From 7fab445da8c919870e13b92a33a3085fec09cec1 Mon Sep 17 00:00:00 2001 From: frenbox Date: Sun, 28 Jun 2026 18:20:15 -0500 Subject: [PATCH 17/20] rebase fixes --- src/enrichment/base.rs | 13 ------------- 1 file changed, 13 deletions(-) diff --git a/src/enrichment/base.rs b/src/enrichment/base.rs index cf4b2dd7d..5e92f6eaa 100644 --- a/src/enrichment/base.rs +++ b/src/enrichment/base.rs @@ -223,7 +223,6 @@ pub async fn run_enrichment_worker( } ACTIVE.add(1, &active_attrs); -<<<<<<< HEAD let candids: Vec = retry_transient( "valkey_rpop", DEFAULT_MAX_RETRIES, @@ -244,18 +243,6 @@ pub async fn run_enrichment_worker( ACTIVE.add(-1, &active_attrs); BATCH_PROCESSED.add(1, &input_error_attrs); })?; -======= - let candids: Vec = con - .rpop::<&str, Vec>( - &input_queue, - NonZero::new(worker_config.enrichment.batch_size), - ) - .await - .inspect_err(|_| { - ACTIVE.add(-1, &active_attrs); - BATCH_PROCESSED.add(1, &input_error_attrs); - })?; ->>>>>>> 31bdd88 (fixed batch size to 750) if candids.is_empty() { ACTIVE.add(-1, &active_attrs); From 2eb1076e5bd227faff42d0a7bc950303ff0a0338 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" <41898282+github-actions[bot]@users.noreply.github.com> Date: Sun, 28 Jun 2026 23:29:25 +0000 Subject: [PATCH 18/20] Update generated deployment configs --- config/prod/caltech/config.yaml | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/config/prod/caltech/config.yaml b/config/prod/caltech/config.yaml index 4d24b703c..3d8b002b0 100644 --- a/config/prod/caltech/config.yaml +++ b/config/prod/caltech/config.yaml @@ -101,6 +101,10 @@ workers: n_workers: 8 enrichment: n_workers: 12 + # Alerts per enrichment batch: the queue RPOP cap AND the fixed ONNX + # inference shape. Partial batches are zero-padded so ORT sees one shape + # and the GPU memory arena stays stable. 750 is stable for both VRAM => 12GB. + batch_size: 750 filter: n_workers: 4 refresh_interval_minutes: 15 From c961a5791e959aa1a4064aee8ec023f514daddc5 Mon Sep 17 00:00:00 2001 From: antoine-le-calloch Date: Tue, 21 Jul 2026 11:02:54 +0200 Subject: [PATCH 19/20] Fix merge. --- src/utils/gpu.rs | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/src/utils/gpu.rs b/src/utils/gpu.rs index b2a18e296..665f6e1e2 100644 --- a/src/utils/gpu.rs +++ b/src/utils/gpu.rs @@ -29,7 +29,12 @@ fn validate_gpu_inference(device_ids: &[i32]) -> Result<(), Box Date: Tue, 21 Jul 2026 11:04:07 +0200 Subject: [PATCH 20/20] Clean up. --- src/enrichment/ztf.rs | 5 +---- src/utils/gpu.rs | 1 - 2 files changed, 1 insertion(+), 5 deletions(-) diff --git a/src/enrichment/ztf.rs b/src/enrichment/ztf.rs index 20783f0d9..439d8b29a 100644 --- a/src/enrichment/ztf.rs +++ b/src/enrichment/ztf.rs @@ -536,12 +536,9 @@ impl EnrichmentWorker for ZtfEnrichmentWorker { .ok_or_else(|| EnrichmentWorkerError::MissingCutouts(candid))?; // Compute numerical and boolean features from lightcurve and candidate analysis - #[cfg(feature = "gpu")] + #[cfg_attr(not(feature = "gpu"), allow(unused_variables))] let (properties, all_bands_properties, programid, lightcurve) = self.get_alert_properties(&alert).await?; - #[cfg(not(feature = "gpu"))] - let (properties, all_bands_properties, programid, _lightcurve) = - self.get_alert_properties(&alert).await?; #[cfg(feature = "gpu")] if self .models diff --git a/src/utils/gpu.rs b/src/utils/gpu.rs index 665f6e1e2..f06367025 100644 --- a/src/utils/gpu.rs +++ b/src/utils/gpu.rs @@ -29,7 +29,6 @@ fn validate_gpu_inference(device_ids: &[i32]) -> Result<(), Box