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 1b2a91b58..33069cc0b 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" @@ -1158,6 +1164,7 @@ dependencies = [ "utoipa", "utoipa-scalar", "uuid", + "villar-pso", "which", "zstd", "zune-inflate", @@ -1583,6 +1590,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" @@ -1698,6 +1716,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" @@ -2146,7 +2185,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]] @@ -2155,6 +2215,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" @@ -3323,6 +3389,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" @@ -3383,6 +3458,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" @@ -3694,6 +3784,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" @@ -3733,7 +3832,7 @@ checksum = "a45fa2aa886c42762255da344f0a0d313e254066c46aad76f300c3d3da62d967" dependencies = [ "bitflags 2.11.1", "cfg-if", - "foreign-types", + "foreign-types 0.3.2", "libc", "openssl-macros", "openssl-sys", @@ -3946,6 +4045,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" @@ -6096,6 +6201,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" diff --git a/Cargo.toml b/Cargo.toml index 2d846e082..8a680d7f6 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -63,13 +63,24 @@ zune-inflate = { version = "0.2", default-features = false, features = [ "std", ] } -# on mac os silicon platform, install ort with the coreml feature +[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`. +# 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 +# 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 with the cuda 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 } [workspace] members = ["apache-avro-macros"] diff --git a/Dockerfile.gpu b/Dockerfile.gpu index 54d77e504..df2a21e78 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 --bin enrich_reprocess +RUN cargo build --release --features gpu --bin scheduler --bin migrate_fp_flux --bin migrate_snr --bin reprocess_crossmatch --bin enrich_reprocess FROM nvidia/cuda:12.8.1-cudnn-runtime-ubuntu24.04 AS app 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. diff --git a/config.yaml b/config.yaml index a14f2f7f1..5d35f5061 100644 --- a/config.yaml +++ b/config.yaml @@ -100,6 +100,10 @@ workers: n_workers: 3 enrichment: n_workers: 3 + # 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/config/prod/caltech/config.yaml b/config/prod/caltech/config.yaml index d1c4eae31..1403c9064 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 diff --git a/src/bin/scheduler.rs b/src/bin/scheduler.rs index c8dc86582..f04d1cb1d 100644 --- a/src/bin/scheduler.rs +++ b/src/bin/scheduler.rs @@ -27,6 +27,132 @@ use tokio::sync::oneshot; use tracing::{info, info_span, warn, Instrument}; use uuid::Uuid; +#[cfg(target_os = "linux")] +const ZTF_MIN_FREE_VRAM_MIB: u64 = 10 * 1024; + +#[cfg(target_os = "linux")] +fn validate_linux_gpu_runtime_preconditions() -> Result<(), &'static str> { + // fail fast if the runtime library path is not explicitly configured. + if std::env::var("ORT_DYLIB_PATH").map_or(true, |v| v.trim().is_empty()) { + return Err("GPU is enabled but ORT_DYLIB_PATH is not set. \ +Set ORT_DYLIB_PATH to a valid libonnxruntime.so path before starting scheduler."); + } + + Ok(()) +} + +#[cfg(target_os = "linux")] +fn validate_gpu_inference(device_ids: &[i32]) -> Result<(), Box> { + info!("Validating GPU inference: running one inference per configured CUDA device"); + + use boom::enrichment::models::{BtsBotModel, Model}; + for &device_id in device_ids { + info!(device_id, "Running BTSBotModel inference on device"); + // Standalone validation; no shared stream needed. + let mut model = BtsBotModel::new_on_device( + "data/models/btsbot-v1.0.1.onnx", + device_id, + std::ptr::null_mut(), + )?; + let metadata = ndarray::Array::from_shape_vec((1, 25), vec![0.5; 25])?; + let triplet = ndarray::Array::from_shape_vec((1, 63, 63, 3), vec![0.5; 63 * 63 * 3])?; + let _ = model.predict(&metadata, &triplet)?; + } + Ok(()) +} + +#[cfg(target_os = "linux")] +fn parse_nvidia_smi_memory_free_output( + output: &str, +) -> Result, Box> { + let mut values = Vec::new(); + for line in output.lines() { + let trimmed = line.trim(); + if trimmed.is_empty() { + continue; + } + let value = trimmed.parse::().map_err(|e| { + std::io::Error::other(format!( + "failed to parse nvidia-smi free memory value '{trimmed}': {e}" + )) + })?; + values.push(value); + } + + if values.is_empty() { + return Err(std::io::Error::other("nvidia-smi returned no GPU free-memory values").into()); + } + + Ok(values) +} + +#[cfg(target_os = "linux")] +fn query_nvidia_smi_free_memory_mib() -> Result, Box> { + let output = std::process::Command::new("nvidia-smi") + .args(["--query-gpu=memory.free", "--format=csv,noheader,nounits"]) + .output() + .map_err(|e| { + std::io::Error::other(format!( + "failed to execute nvidia-smi for GPU memory validation: {e}" + )) + })?; + + if !output.status.success() { + let stderr = String::from_utf8_lossy(&output.stderr); + return Err(std::io::Error::other(format!( + "nvidia-smi failed while validating free GPU memory: {}", + stderr.trim() + )) + .into()); + } + + let stdout = String::from_utf8_lossy(&output.stdout); + parse_nvidia_smi_memory_free_output(&stdout) +} + +#[cfg(target_os = "linux")] +fn validate_gpu_free_vram( + device_ids: &[i32], + min_free_vram_mib: u64, +) -> Result<(), Box> { + let free_by_gpu = query_nvidia_smi_free_memory_mib()?; + + for &device_id in device_ids { + if device_id < 0 { + return Err(std::io::Error::other(format!( + "invalid CUDA device id {device_id}; device ids must be >= 0" + )) + .into()); + } + + let index = device_id as usize; + let Some(&free_mib) = free_by_gpu.get(index) else { + return Err(std::io::Error::other(format!( + "configured CUDA device id {device_id} is out of range; nvidia-smi reported {} device(s)", + free_by_gpu.len() + )) + .into()); + }; + + if free_mib < min_free_vram_mib { + return Err(std::io::Error::other(format!( + "configured CUDA device {device_id} has only {free_mib} MiB free VRAM; ZTF enrichment requires at least {min_free_vram_mib} MiB ({:.1} GiB) free per device", + min_free_vram_mib as f64 / 1024.0 + )) + .into()); + } + + info!( + device_id, + free_vram_mib = free_mib, + min_required_mib = min_free_vram_mib, + "validated free GPU VRAM for ZTF enrichment" + ); + } + + Ok(()) +} + /// Sample one aux record at random and warn if it's missing crossmatches for /// any catalog declared under `crossmatch.` in the config. The live /// pipeline only crossmatches at first insert, so newly added catalogs never diff --git a/src/conf.rs b/src/conf.rs index 4de417044..5c316ad4c 100644 --- a/src/conf.rs +++ b/src/conf.rs @@ -595,6 +595,24 @@ pub struct WorkerConfig { pub n_workers: usize, } +fn default_enrichment_batch_size() -> usize { + 750 +} + +#[derive(Deserialize, Debug, Clone)] +pub struct EnrichmentWorkerConfig { + pub n_workers: usize, + /// 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, +} + fn default_filter_refresh_interval_minutes() -> u64 { 15 } @@ -687,7 +705,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/models/acai.rs b/src/enrichment/models/acai.rs index 8644fe556..0a1cdd894 100644 --- a/src/enrichment/models/acai.rs +++ b/src/enrichment/models/acai.rs @@ -39,10 +39,15 @@ impl Model for AcaiModel { } impl AcaiModel { - /// Load on a specific CUDA device. - pub fn new_on_device(path: &str, device_id: i32) -> 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..e66cdb307 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,32 @@ 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..1e389bebc 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,106 @@ 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( - "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, + )?, + 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")?, + 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 2598057cc..439d8b29a 100644 --- a/src/enrichment/ztf.rs +++ b/src/enrichment/ztf.rs @@ -20,6 +20,10 @@ use mongodb::options::{UpdateOneModel, WriteModel}; use serde::{Deserialize, Deserializer}; use std::sync::Arc; 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"))] +use villar_pso::gpu_metal::{GpuBatchData, SourceData}; #[serdavro] #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] @@ -390,10 +394,30 @@ 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, + /// Alerts per batch — also the fixed ONNX inference shape (see + /// [`EnrichmentWorkerConfig::batch_size`] in `conf.rs`). + batch_size: usize, +} + +#[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 +451,13 @@ impl EnrichmentWorker for ZtfEnrichmentWorker { None => Some(SharedModels::load(None)?), }; + let batch_size = config + .workers + .get(&Survey::Ztf) + .ok_or(EnrichmentWorkerError::WorkerConfigMissing(Survey::Ztf))? + .enrichment + .batch_size; + Ok(ZtfEnrichmentWorker { input_queue, output_queue, @@ -437,6 +468,7 @@ impl EnrichmentWorker for ZtfEnrichmentWorker { models, babamul, gpu_enabled: config.gpu.enabled, + batch_size, }) } @@ -495,13 +527,28 @@ 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))?; - let (properties, all_bands_properties, programid, _lightcurve) = + + // Compute numerical and boolean features from lightcurve and candidate analysis + #[cfg_attr(not(feature = "gpu"), allow(unused_variables))] + let (properties, all_bands_properties, programid, lightcurve) = self.get_alert_properties(&alert).await?; + #[cfg(feature = "gpu")] + if self + .models + .as_ref() + .and_then(|m| m.gpu_ctx.as_ref()) + .is_some() + { + villar_inputs.push((candid, lightcurve)); + } + work_items.push(AlertWork { candid, programid, @@ -554,6 +601,96 @@ impl EnrichmentWorker for ZtfEnrichmentWorker { let _ = self.client.bulk_write(updates).await?.modified_count; + // 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.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 = { + 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(); + + 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) => { + 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?; @@ -688,7 +825,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) @@ -749,6 +886,7 @@ impl ZtfEnrichmentWorker { } fn classify_gpu_batch( + &self, models: &SharedModels, work_items: &[AlertWork], ) -> Result>, EnrichmentWorkerError> { @@ -801,61 +939,66 @@ 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 (`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.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"); + 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.batch_size { + return Err(EnrichmentWorkerError::ConfigurationError(format!( + "model {} returned {} scores for {} padded inputs", + name, got, self.batch_size + ))); + } } - } - 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) diff --git a/src/utils/gpu.rs b/src/utils/gpu.rs index b2a18e296..f06367025 100644 --- a/src/utils/gpu.rs +++ b/src/utils/gpu.rs @@ -29,7 +29,11 @@ fn validate_gpu_inference(device_ids: &[i32]) -> Result<(), Box