From 58dc95cf54838b4307dedcf8f43cf2291cad82cd Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Mon, 2 Mar 2026 17:46:21 +0100 Subject: [PATCH] =?UTF-8?q?feat(ml):=20InferenceEnsemble=20GPU-aggregated?= =?UTF-8?q?=20prediction=20=E2=80=94=20N=20syncs=20=E2=86=92=201?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Use predict_raw() to collect raw GPU tensors from adapters. Stack, sigmoid, weighted-sum on GPU before single extraction. Falls back to CPU path for adapters without tensor output. Co-Authored-By: Claude Opus 4.6 --- crates/ml/src/ensemble/inference_ensemble.rs | 228 ++++++++++++++++--- 1 file changed, 195 insertions(+), 33 deletions(-) diff --git a/crates/ml/src/ensemble/inference_ensemble.rs b/crates/ml/src/ensemble/inference_ensemble.rs index bc72f9ba1..b6c5d0fa5 100644 --- a/crates/ml/src/ensemble/inference_ensemble.rs +++ b/crates/ml/src/ensemble/inference_ensemble.rs @@ -8,8 +8,10 @@ use std::collections::HashMap; use rayon::prelude::*; +use candle_core::Tensor; + use crate::ensemble::inference_adapter::{ - EnsemblePrediction, FeatureVector, ModelInferenceAdapter, PredictionMeta, + EnsemblePrediction, FeatureVector, ModelInferenceAdapter, PredictionMeta, RawPrediction, }; use crate::{MLError, MLResult}; @@ -57,6 +59,13 @@ impl InferenceEnsemble { /// Run inference across all ready adapters and aggregate via /// confidence-weighted voting. /// + /// Uses [`predict_raw()`](ModelInferenceAdapter::predict_raw) to collect + /// raw GPU tensors where available. Models that return a tensor are + /// aggregated on-device (stack → sigmoid → weighted-sum → single + /// extraction). Models without a tensor fall back to the original + /// scalar-weighted-average path. The two paths are merged by + /// model-count-weighted average. + /// /// Returns [`MLError::InferenceError`] if no models are ready. pub fn predict(&self, features: &FeatureVector) -> MLResult { let ready_adapters: Vec<&Box> = @@ -68,17 +77,17 @@ impl InferenceEnsemble { )); } - // Run inference in parallel across all ready adapters - let results: Vec<(String, EnsemblePrediction)> = ready_adapters + // Run predict_raw() in parallel across all ready adapters + let raw_predictions: Vec<(String, RawPrediction)> = ready_adapters .par_iter() .filter_map(|adapter| { let model_name = adapter.model_name().to_string(); - match adapter.predict(features) { + match adapter.predict_raw(features) { Ok(pred) => { - if !pred.direction.is_finite() || !pred.confidence.is_finite() { + if !pred.direction_scalar.is_finite() || !pred.confidence.is_finite() { tracing::warn!( model = %model_name, - direction = %pred.direction, + direction = %pred.direction_scalar, confidence = %pred.confidence, "Model returned NaN/Inf prediction, skipping (circuit breaker)" ); @@ -99,41 +108,82 @@ impl InferenceEnsemble { }) .collect(); - // Aggregate results sequentially (fast arithmetic) - let mut weighted_direction_sum = 0.0_f64; - let mut weight_confidence_sum = 0.0_f64; - let mut confidence_sum = 0.0_f64; - let mut model_names: Vec = Vec::new(); - - for (model_name, pred) in &results { - let confidence = pred.confidence.clamp(0.0, 1.0); - let w = self - .weights - .get(model_name) - .copied() - .unwrap_or(1.0); - let wc = w * confidence; - weighted_direction_sum += pred.direction * wc; - weight_confidence_sum += wc; - confidence_sum += confidence; - model_names.push(model_name.clone()); - } - - let successful_count = results.len(); - - if successful_count == 0 { + if raw_predictions.is_empty() { return Err(MLError::InferenceError( "All ready models failed during inference".to_owned(), )); } - let direction = if weight_confidence_sum.abs() < f64::EPSILON { - 0.0 + // Partition into GPU-tensor vs CPU-scalar predictions + let (gpu_preds, cpu_preds): (Vec<_>, Vec<_>) = raw_predictions + .into_iter() + .partition(|(_, p)| p.tensor.is_some()); + + // Collect model names from both paths for the ensemble label + let mut model_names: Vec = Vec::new(); + let mut total_confidence_sum = 0.0_f64; + let mut total_count: usize = 0; + + // --- GPU path: stack tensors, sigmoid, weighted-sum, single extraction --- + let gpu_result = if !gpu_preds.is_empty() { + match self.aggregate_gpu(&gpu_preds) { + Ok((direction, count)) => { + for (name, pred) in &gpu_preds { + model_names.push(name.clone()); + total_confidence_sum += pred.confidence.clamp(0.0, 1.0); + } + total_count += gpu_preds.len(); + Some((direction, count)) + } + Err(e) => { + tracing::warn!( + error = %e, + "GPU aggregation failed, falling back to CPU for {} models", + gpu_preds.len() + ); + // Fall back: treat GPU preds as CPU scalars + let fallback = self.aggregate_cpu(&gpu_preds); + for (name, pred) in &gpu_preds { + model_names.push(name.clone()); + total_confidence_sum += pred.confidence.clamp(0.0, 1.0); + } + total_count += gpu_preds.len(); + fallback + } + } } else { - weighted_direction_sum / weight_confidence_sum + None }; - let avg_confidence = confidence_sum / successful_count as f64; + // --- CPU path: existing scalar weighted average --- + let cpu_result = if !cpu_preds.is_empty() { + let result = self.aggregate_cpu(&cpu_preds); + for (name, pred) in &cpu_preds { + model_names.push(name.clone()); + total_confidence_sum += pred.confidence.clamp(0.0, 1.0); + } + total_count += cpu_preds.len(); + result + } else { + None + }; + + // Merge GPU and CPU directions by model-count-weighted average + let direction = match (gpu_result, cpu_result) { + (Some((gpu_dir, gpu_n)), Some((cpu_dir, cpu_n))) => { + let total_n = (gpu_n + cpu_n) as f64; + (gpu_dir * gpu_n as f64 + cpu_dir * cpu_n as f64) / total_n + } + (Some((dir, _)), None) => dir, + (None, Some((dir, _))) => dir, + (None, None) => 0.0, + }; + + let avg_confidence = if total_count > 0 { + total_confidence_sum / total_count as f64 + } else { + 0.0 + }; let ensemble_name = format!("ENSEMBLE({})", model_names.join("+")); @@ -144,6 +194,118 @@ impl InferenceEnsemble { metadata: PredictionMeta::default(), }) } + + /// Aggregate predictions on GPU: stack tensors, apply sigmoid, + /// compute confidence-weighted sum, extract single scalar. + fn aggregate_gpu( + &self, + preds: &[(String, RawPrediction)], + ) -> Result<(f64, usize), MLError> { + let tensors: Vec = preds + .iter() + .filter_map(|(_, p)| p.tensor.as_ref().cloned()) + .collect(); + + if tensors.is_empty() { + return Err(MLError::InferenceError( + "No GPU tensors available for aggregation".to_owned(), + )); + } + + // Stack all model outputs into a single tensor [N] + let stacked = Tensor::stack(&tensors, 0).map_err(|e| { + MLError::TensorOperationError(format!("Failed to stack GPU tensors: {e}")) + })?; + + // Apply sigmoid to convert logits → probabilities on device + let sigmoided = candle_nn::ops::sigmoid(&stacked).map_err(|e| { + MLError::TensorOperationError(format!("Sigmoid failed: {e}")) + })?; + + // Build confidence weights on the same device + let weights_f32: Vec = preds + .iter() + .map(|(name, p)| { + let conf = p.confidence.clamp(0.0, 1.0) as f32; + let w = self.weights.get(name).copied().unwrap_or(1.0) as f32; + w * conf + }) + .collect(); + + let device = stacked.device(); + let n = weights_f32.len(); + + let weight_t = Tensor::from_vec(weights_f32, n, device).map_err(|e| { + MLError::TensorCreationError { + operation: "weight_tensor".to_owned(), + reason: format!("{e}"), + } + })?; + + let weight_sum = weight_t.sum_all().map_err(|e| { + MLError::TensorOperationError(format!("Weight sum failed: {e}")) + })?; + + // Guard against zero-weight sum + let weight_sum_val = weight_sum.to_scalar::().map_err(|e| { + MLError::TensorOperationError(format!("Weight sum extraction failed: {e}")) + })?; + + if weight_sum_val.abs() < f32::EPSILON { + return Ok((0.0, preds.len())); + } + + let normalized = weight_t.broadcast_div(&weight_sum).map_err(|e| { + MLError::TensorOperationError(format!("Weight normalization failed: {e}")) + })?; + + // Weighted sum: dot product of sigmoided values and normalized weights + let weighted = sigmoided.mul(&normalized).map_err(|e| { + MLError::TensorOperationError(format!("Weighted multiply failed: {e}")) + })?; + + let result = weighted.sum_all().map_err(|e| { + MLError::TensorOperationError(format!("Sum extraction failed: {e}")) + })?; + + let direction_f32 = result.to_scalar::().map_err(|e| { + MLError::TensorOperationError(format!("Scalar extraction failed: {e}")) + })?; + + // Sigmoid output is [0,1]; remap to [-1,1] direction space + let direction = (f64::from(direction_f32) * 2.0) - 1.0; + + Ok((direction, preds.len())) + } + + /// CPU-side scalar weighted average (original algorithm). + fn aggregate_cpu( + &self, + preds: &[(String, RawPrediction)], + ) -> Option<(f64, usize)> { + if preds.is_empty() { + return None; + } + + let mut weighted_direction_sum = 0.0_f64; + let mut weight_confidence_sum = 0.0_f64; + + for (name, pred) in preds { + let confidence = pred.confidence.clamp(0.0, 1.0); + let w = self.weights.get(name).copied().unwrap_or(1.0); + let wc = w * confidence; + weighted_direction_sum += pred.direction_scalar * wc; + weight_confidence_sum += wc; + } + + let direction = if weight_confidence_sum.abs() < f64::EPSILON { + 0.0 + } else { + weighted_direction_sum / weight_confidence_sum + }; + + Some((direction, preds.len())) + } } #[cfg(test)]