Villar-PSO light curve fitting on GPU in ZTF enrichment workers

.gitignore

@@ -19,3 +19,4 @@ logs
 *.jsonl
 simbad_cache
 *.parquet
+config.docker.yaml

Cargo.toml

@@ -62,13 +62,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"]

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
 

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.

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

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

src/bin/scheduler.rs

@@ -46,7 +46,12 @@ fn validate_gpu_inference(device_ids: &[i32]) -> Result<(), Box<dyn std::error::
     use boom::enrichment::models::{BtsBotModel, Model};
     for &device_id in device_ids {
         info!(device_id, "Running BTSBotModel inference on device");
-        let mut model = BtsBotModel::new_on_device("data/models/btsbot-v1.0.1.onnx", device_id)?;
+        // 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)?;

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,
 }
 

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<Self, ModelError> {
+    /// 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<Self, ModelError> {
         Ok(Self {
-            model: load_model_on_device(path, Some(device_id))?,
+            model: load_model_on_device(path, Some(device_id), cuda_stream)?,
         })
     }