0e3f7be856
Eliminated 7 downloads:
- dqn.rs dead neuron: GPU abs→le→sum (test-only, marked #[cold])
- ppo.rs compute_losses: GPU gather_rows kernel for per-action log-prob
(eliminated 3 full-batch downloads)
- ppo.rs update_gpu: GPU gather_rows + GpuTensor::symlog()
(sign(x)*ln(|x|+1) via 6 elementwise GPU kernels)
- Test assertions: annotated with // test-only readback
Marked #[cold] + annotated 8 checkpoint/API methods:
- CudaLinear::get_weights(), CudaVec::to_vec(), GpuTensor::to_host(),
GpuVarStore::{all_vars,flatten,export_to_host},
GpuLinear::{weight_to_vec,bias_to_vec}
New GPU infrastructure:
- ElementwiseKernels: gather_rows + gather_rows_u32 CUDA kernels
- GpuTensor::symlog() — fully GPU-native sign*log transform
Every remaining memcpy_dtoh is annotated: // gpu-exit: or // test-only readback
Verification: `rg "memcpy_dtoh" | grep -v "gpu-exit\|test.*readback"` = 0
1,116 tests pass across 5 sub-crates. Zero failures.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
851 lines
29 KiB
Rust
851 lines
29 KiB
Rust
//! Stream-aware ensemble with per-model CUDA device handles.
|
|
//!
|
|
//! Each model runs on its own CUDA stream for true GPU-level parallelism.
|
|
//! Falls back to standard rayon parallelism on CPU.
|
|
//!
|
|
//! # Design (Option C — per-model CudaDevice handles)
|
|
//!
|
|
//! On CUDA: each model adapter gets its own `CudaDevice::new(0)` handle so the
|
|
//! CUDA driver assigns a unique default stream. `predict()` runs all models via
|
|
//! rayon with per-device streams, syncs all, then aggregates on the primary device.
|
|
//!
|
|
//! On CPU: all adapters share the same device, rayon provides parallelism (same
|
|
//! behaviour as [`InferenceEnsemble`](super::inference_ensemble::InferenceEnsemble)).
|
|
|
|
use rayon::prelude::*;
|
|
|
|
use ml_core::device::MlDevice;
|
|
#[cfg(feature = "cuda")]
|
|
use ml_core::cuda_autograd::{ActivationKernels, GpuTensor};
|
|
|
|
use crate::cuda_streams::CudaStreamPool;
|
|
use crate::inference_adapter::{
|
|
EnsemblePrediction, FeatureVector, ModelInferenceAdapter, PredictionMeta, RawPrediction,
|
|
};
|
|
use crate::{MLError, MLResult};
|
|
|
|
/// Ensemble that runs models on separate CUDA streams for true
|
|
/// GPU-level parallelism. Falls back to rayon on CPU.
|
|
///
|
|
/// Uses [`CudaStreamPool`] for stream management and synchronisation.
|
|
#[allow(missing_debug_implementations)]
|
|
pub struct StreamAwareEnsemble {
|
|
/// Model adapters (each may have its own device handle)
|
|
adapters: Vec<Box<dyn ModelInferenceAdapter>>,
|
|
/// Weights per model (normalised, indexed by adapter position)
|
|
weights: Vec<f64>,
|
|
/// Stream pool for CUDA synchronisation (no-op on CPU)
|
|
stream_pool: CudaStreamPool,
|
|
}
|
|
|
|
impl StreamAwareEnsemble {
|
|
/// Create a new stream-aware ensemble.
|
|
///
|
|
/// `weights` are normalised so they sum to 1.0. If the provided weights
|
|
/// slice is shorter than `adapters`, missing entries default to equal
|
|
/// share. If `weights` is empty, all models get equal weight.
|
|
pub fn new(
|
|
adapters: Vec<Box<dyn ModelInferenceAdapter>>,
|
|
weights: Vec<f64>,
|
|
device: &MlDevice,
|
|
) -> MLResult<Self> {
|
|
let n = adapters.len();
|
|
let stream_pool = CudaStreamPool::new(device, n)?;
|
|
|
|
let raw_weights = if weights.is_empty() || n == 0 {
|
|
vec![1.0_f64; n]
|
|
} else {
|
|
let mut w = weights;
|
|
w.resize(n, 1.0);
|
|
w
|
|
};
|
|
|
|
let weight_sum: f64 = raw_weights.iter().sum();
|
|
let normalized = if weight_sum > 0.0 {
|
|
raw_weights.iter().map(|w| w / weight_sum).collect()
|
|
} else {
|
|
vec![1.0 / n.max(1) as f64; n]
|
|
};
|
|
|
|
Ok(Self {
|
|
adapters,
|
|
weights: normalized,
|
|
stream_pool,
|
|
})
|
|
}
|
|
|
|
/// Run ensemble prediction with stream-level parallelism.
|
|
///
|
|
/// 1. Runs all ready adapters in parallel via rayon (each on its own
|
|
/// CUDA stream if GPU, or rayon thread if CPU).
|
|
/// 2. Synchronises all CUDA streams (no-op on CPU).
|
|
/// 3. Aggregates predictions via confidence-weighted voting with
|
|
/// GPU-logit and CPU-scalar paths (same algorithm as
|
|
/// [`InferenceEnsemble`]).
|
|
pub fn predict(&self, features: &FeatureVector) -> MLResult<EnsemblePrediction> {
|
|
// Parallel inference via rayon — each adapter on its own stream if CUDA
|
|
let predictions: Vec<(usize, String, RawPrediction)> = self
|
|
.adapters
|
|
.par_iter()
|
|
.enumerate()
|
|
.filter(|(_, a)| a.is_ready())
|
|
.filter_map(|(i, adapter)| {
|
|
let name = adapter.model_name().to_string();
|
|
match adapter.predict_raw(features) {
|
|
Ok(pred) => {
|
|
if !pred.direction_scalar.is_finite() || !pred.confidence.is_finite() {
|
|
tracing::warn!(
|
|
model = %name,
|
|
direction = %pred.direction_scalar,
|
|
confidence = %pred.confidence,
|
|
"StreamAwareEnsemble: model returned NaN/Inf, skipping"
|
|
);
|
|
None
|
|
} else {
|
|
Some((i, name, pred))
|
|
}
|
|
}
|
|
Err(e) => {
|
|
tracing::warn!(
|
|
model = %name,
|
|
error = %e,
|
|
"StreamAwareEnsemble: model prediction failed, skipping"
|
|
);
|
|
None
|
|
}
|
|
}
|
|
})
|
|
.collect();
|
|
|
|
if predictions.is_empty() {
|
|
return Ok(EnsemblePrediction {
|
|
model_name: "STREAM_ENSEMBLE(empty)".to_owned(),
|
|
direction: 0.0,
|
|
confidence: 0.0,
|
|
metadata: PredictionMeta::default(),
|
|
});
|
|
}
|
|
|
|
// Synchronise all CUDA streams before aggregation (no-op on CPU)
|
|
self.stream_pool.sync_all()?;
|
|
|
|
// Partition into logit-bearing vs CPU-scalar predictions
|
|
let (logit_preds, cpu_preds): (Vec<_>, Vec<_>) = predictions
|
|
.into_iter()
|
|
.partition(|(_, _, p)| p.logits.is_some());
|
|
|
|
let mut model_names: Vec<String> = Vec::new();
|
|
let mut total_confidence_sum = 0.0_f64;
|
|
let mut total_count: usize = 0;
|
|
|
|
// --- Logits path: sigmoid + weighted-sum on host f32 vectors ---
|
|
let logit_result = if !logit_preds.is_empty() {
|
|
match self.aggregate_logits(&logit_preds) {
|
|
Ok((direction, count)) => {
|
|
for (_, name, pred) in &logit_preds {
|
|
model_names.push(name.clone());
|
|
total_confidence_sum += pred.confidence.clamp(0.0, 1.0);
|
|
}
|
|
total_count += logit_preds.len();
|
|
Some((direction, count))
|
|
}
|
|
Err(e) => {
|
|
tracing::warn!(
|
|
error = %e,
|
|
"Logits aggregation failed, falling back to CPU for {} models",
|
|
logit_preds.len()
|
|
);
|
|
let fallback = self.aggregate_cpu(&logit_preds);
|
|
for (_, name, pred) in &logit_preds {
|
|
model_names.push(name.clone());
|
|
total_confidence_sum += pred.confidence.clamp(0.0, 1.0);
|
|
}
|
|
total_count += logit_preds.len();
|
|
fallback
|
|
}
|
|
}
|
|
} else {
|
|
None
|
|
};
|
|
|
|
// --- CPU path: 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 logits and CPU directions by model-count-weighted average
|
|
let direction = match (logit_result, cpu_result) {
|
|
(Some((logit_dir, logit_n)), Some((cpu_dir, cpu_n))) => {
|
|
let total_n = (logit_n + cpu_n) as f64;
|
|
(logit_dir * logit_n as f64 + cpu_dir * cpu_n as f64) / total_n
|
|
}
|
|
(Some((dir, _)), None) | (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!("STREAM_ENSEMBLE({})", model_names.join("+"));
|
|
|
|
Ok(EnsemblePrediction {
|
|
model_name: ensemble_name,
|
|
direction,
|
|
confidence: avg_confidence,
|
|
metadata: PredictionMeta::default(),
|
|
})
|
|
}
|
|
|
|
/// Number of adapters in the ensemble.
|
|
pub fn adapter_count(&self) -> usize {
|
|
self.adapters.len()
|
|
}
|
|
|
|
/// Number of adapters whose `is_ready()` returns true.
|
|
pub fn ready_count(&self) -> usize {
|
|
self.adapters.iter().filter(|a| a.is_ready()).count()
|
|
}
|
|
|
|
/// Whether the underlying stream pool is on a CUDA device.
|
|
pub fn is_cuda(&self) -> bool {
|
|
self.stream_pool.is_cuda()
|
|
}
|
|
|
|
// -----------------------------------------------------------------------
|
|
// Private aggregation helpers
|
|
// -----------------------------------------------------------------------
|
|
|
|
/// Logits aggregation on GPU: concatenate all models' logits into a single
|
|
/// GPU tensor, apply sigmoid via `ActivationKernels` CUDA kernel, compute
|
|
/// per-model mean and weighted sum via a fused CUDA kernel, then read back
|
|
/// only the final scalar.
|
|
///
|
|
/// Falls back to a CPU path when CUDA is not available.
|
|
fn aggregate_logits(
|
|
&self,
|
|
preds: &[(usize, String, RawPrediction)],
|
|
) -> Result<(f64, usize), MLError> {
|
|
#[cfg(feature = "cuda")]
|
|
{
|
|
if let Some(result) = self.aggregate_logits_gpu(preds)? {
|
|
return Ok(result);
|
|
}
|
|
}
|
|
|
|
// CPU fallback (non-CUDA builds or when stream pool has no CUDA streams)
|
|
self.aggregate_logits_cpu(preds)
|
|
}
|
|
|
|
/// GPU path: sigmoid + weighted-mean-reduce entirely on device.
|
|
///
|
|
/// 1. Concatenate all models' logit vectors into one flat `GpuTensor`.
|
|
/// 2. Apply sigmoid via `ActivationKernels::sigmoid_fwd()` (one kernel launch).
|
|
/// 3. Launch a fused reduction kernel that computes per-model means and the
|
|
/// final confidence-weighted sum in a single pass.
|
|
/// 4. Read back only the single f32 result.
|
|
///
|
|
/// Returns `Ok(None)` if no CUDA stream is available (caller should fall back
|
|
/// to CPU).
|
|
#[cfg(feature = "cuda")]
|
|
#[allow(unsafe_code)] // CUDA kernel launches require unsafe FFI.
|
|
fn aggregate_logits_gpu(
|
|
&self,
|
|
preds: &[(usize, String, RawPrediction)],
|
|
) -> Result<Option<(f64, usize)>, MLError> {
|
|
use std::sync::Arc;
|
|
use cudarc::driver::{LaunchConfig, PushKernelArg};
|
|
|
|
// Need a CUDA stream from the pool
|
|
let stream = match self.stream_pool.get_stream(0) {
|
|
Some(s) => Arc::clone(s),
|
|
None => return Ok(None),
|
|
};
|
|
|
|
// Build per-model metadata: (weight*confidence, logit_offset, logit_len)
|
|
let mut all_logits: Vec<f32> = Vec::new();
|
|
let mut segment_offsets: Vec<i32> = Vec::new();
|
|
let mut segment_lengths: Vec<i32> = Vec::new();
|
|
let mut model_weights: Vec<f32> = Vec::new();
|
|
|
|
let mut weight_sum = 0.0_f32;
|
|
|
|
for (i, (idx, _, pred)) in preds.iter().enumerate() {
|
|
let logits = match pred.logits.as_ref() {
|
|
Some(l) if !l.is_empty() => l,
|
|
_ => continue,
|
|
};
|
|
let conf = pred.confidence.clamp(0.0, 1.0) as f32;
|
|
let w = self.weights.get(*idx).copied().unwrap_or(1.0) as f32;
|
|
let wc = w * conf;
|
|
|
|
segment_offsets.push(all_logits.len() as i32);
|
|
segment_lengths.push(logits.len() as i32);
|
|
model_weights.push(wc);
|
|
weight_sum += wc;
|
|
|
|
all_logits.extend_from_slice(logits);
|
|
let _ = i; // suppress unused
|
|
}
|
|
|
|
if all_logits.is_empty() || weight_sum.abs() < f32::EPSILON {
|
|
return Ok(Some((0.0, preds.len())));
|
|
}
|
|
|
|
let n_models = segment_offsets.len();
|
|
|
|
// Normalize weights
|
|
for w in &mut model_weights {
|
|
*w /= weight_sum;
|
|
}
|
|
|
|
// 1) Upload concatenated logits to GPU
|
|
let logit_tensor = GpuTensor::from_host(
|
|
&all_logits,
|
|
vec![all_logits.len()],
|
|
&stream,
|
|
)?;
|
|
|
|
// 2) Apply sigmoid on GPU via ActivationKernels
|
|
let act_kernels = ActivationKernels::new(&stream)?;
|
|
let (sigmoid_tensor, _saved) = act_kernels.sigmoid_fwd(&logit_tensor, &stream)?;
|
|
|
|
// 3) Upload segment metadata and weights to GPU, then launch a
|
|
// fused per-model-mean + weighted-sum reduction kernel.
|
|
let mut d_offsets = stream.alloc_zeros::<i32>(n_models).map_err(|e| {
|
|
MLError::ModelError(format!("alloc offsets: {e}"))
|
|
})?;
|
|
stream.memcpy_htod(&segment_offsets, &mut d_offsets).map_err(|e| {
|
|
MLError::ModelError(format!("htod offsets: {e}"))
|
|
})?;
|
|
|
|
let mut d_lengths = stream.alloc_zeros::<i32>(n_models).map_err(|e| {
|
|
MLError::ModelError(format!("alloc lengths: {e}"))
|
|
})?;
|
|
stream.memcpy_htod(&segment_lengths, &mut d_lengths).map_err(|e| {
|
|
MLError::ModelError(format!("htod lengths: {e}"))
|
|
})?;
|
|
|
|
let mut d_weights = stream.alloc_zeros::<f32>(n_models).map_err(|e| {
|
|
MLError::ModelError(format!("alloc weights: {e}"))
|
|
})?;
|
|
stream.memcpy_htod(&model_weights, &mut d_weights).map_err(|e| {
|
|
MLError::ModelError(format!("htod weights: {e}"))
|
|
})?;
|
|
|
|
// Output: single f32 (weighted sum of per-model sigmoid means)
|
|
let d_output = stream.alloc_zeros::<f32>(1).map_err(|e| {
|
|
MLError::ModelError(format!("alloc output: {e}"))
|
|
})?;
|
|
|
|
let context = stream.context();
|
|
let ptx = ml_core::cuda_compile::compile_ptx_for_device(
|
|
ENSEMBLE_REDUCE_CUDA_SRC,
|
|
context,
|
|
).map_err(|e| MLError::ModelError(format!("ensemble reduce compile: {e}")))?;
|
|
|
|
let module = context.load_module(ptx).map_err(|e| {
|
|
MLError::ModelError(format!("ensemble reduce module load: {e}"))
|
|
})?;
|
|
let kernel = module.load_function("weighted_sigmoid_mean_reduce").map_err(|e| {
|
|
MLError::ModelError(format!("ensemble reduce kernel load: {e}"))
|
|
})?;
|
|
|
|
let n_models_i32 = n_models as i32;
|
|
let cfg = LaunchConfig {
|
|
grid_dim: (1, 1, 1),
|
|
block_dim: (n_models.min(256) as u32, 1, 1),
|
|
shared_mem_bytes: 0,
|
|
};
|
|
|
|
// SAFETY: kernel arguments match the CUDA kernel signature
|
|
// (sigmoid_vals, offsets, lengths, weights, output, n_models).
|
|
// All buffers are GPU-allocated with correct sizes above.
|
|
unsafe {
|
|
stream
|
|
.launch_builder(&kernel)
|
|
.arg(sigmoid_tensor.data())
|
|
.arg(&d_offsets)
|
|
.arg(&d_lengths)
|
|
.arg(&d_weights)
|
|
.arg(&d_output)
|
|
.arg(&n_models_i32)
|
|
.launch(cfg)
|
|
.map_err(|e| MLError::ModelError(format!("ensemble reduce launch: {e}")))?;
|
|
}
|
|
|
|
// 4) Read back the single scalar result
|
|
let mut result_host = [0.0_f32];
|
|
stream.memcpy_dtoh(&d_output, &mut result_host).map_err(|e| { // gpu-exit: final prediction scalar exits system
|
|
MLError::ModelError(format!("ensemble reduce dtoh: {e}"))
|
|
})?;
|
|
|
|
// Sigmoid output is [0,1]; remap to [-1,1] direction space
|
|
let direction = (f64::from(result_host[0]) * 2.0) - 1.0;
|
|
|
|
Ok(Some((direction, preds.len())))
|
|
}
|
|
|
|
/// CPU fallback for logits aggregation (non-CUDA builds).
|
|
fn aggregate_logits_cpu(
|
|
&self,
|
|
preds: &[(usize, String, RawPrediction)],
|
|
) -> Result<(f64, usize), MLError> {
|
|
let weights_f32: Vec<f32> = preds
|
|
.iter()
|
|
.map(|(idx, _, p)| {
|
|
let conf = p.confidence.clamp(0.0, 1.0) as f32;
|
|
let w = self.weights.get(*idx).copied().unwrap_or(1.0) as f32;
|
|
w * conf
|
|
})
|
|
.collect();
|
|
|
|
let weight_sum: f32 = weights_f32.iter().sum();
|
|
if weight_sum.abs() < f32::EPSILON {
|
|
return Ok((0.0, preds.len()));
|
|
}
|
|
|
|
let mut weighted_sum = 0.0_f64;
|
|
|
|
for (i, (_, _, pred)) in preds.iter().enumerate() {
|
|
let logits = match pred.logits.as_ref() {
|
|
Some(l) if !l.is_empty() => l,
|
|
_ => continue,
|
|
};
|
|
|
|
// CPU sigmoid + mean (fallback only)
|
|
let sum_sigmoid: f32 = logits
|
|
.iter()
|
|
.map(|&x| 1.0_f32 / (1.0_f32 + (-x).exp()))
|
|
.sum();
|
|
let mean_sigmoid = sum_sigmoid / logits.len() as f32;
|
|
|
|
let normalized_w = weights_f32.get(i).copied().unwrap_or(0.0) / weight_sum;
|
|
weighted_sum += f64::from(mean_sigmoid) * f64::from(normalized_w);
|
|
}
|
|
|
|
let direction = (weighted_sum * 2.0) - 1.0;
|
|
Ok((direction, preds.len()))
|
|
}
|
|
|
|
/// CPU-side scalar weighted average.
|
|
fn aggregate_cpu(
|
|
&self,
|
|
preds: &[(usize, 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 (idx, _, pred) in preds {
|
|
let confidence = pred.confidence.clamp(0.0, 1.0);
|
|
let w = self.weights.get(*idx).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()))
|
|
}
|
|
}
|
|
|
|
// ── CUDA source for ensemble weighted-sigmoid-mean reduction ─────────────
|
|
//
|
|
// One thread per model. Each thread loops over its logit segment (already
|
|
// post-sigmoid), computes the mean, multiplies by the normalized model weight,
|
|
// and atomically adds to a single output scalar.
|
|
//
|
|
// This keeps the entire aggregation on GPU after sigmoid_fwd() — only one f32
|
|
// is read back to CPU.
|
|
#[cfg(feature = "cuda")]
|
|
const ENSEMBLE_REDUCE_CUDA_SRC: &str = r#"
|
|
extern "C" __global__
|
|
void weighted_sigmoid_mean_reduce(
|
|
const float* __restrict__ sigmoid_vals,
|
|
const int* __restrict__ offsets,
|
|
const int* __restrict__ lengths,
|
|
const float* __restrict__ weights,
|
|
float* __restrict__ output,
|
|
int n_models)
|
|
{
|
|
int model = blockIdx.x * blockDim.x + threadIdx.x;
|
|
if (model >= n_models) return;
|
|
|
|
int off = offsets[model];
|
|
int len = lengths[model];
|
|
if (len <= 0) return;
|
|
|
|
// Compute mean of post-sigmoid values for this model's segment
|
|
float sum = 0.0f;
|
|
for (int j = 0; j < len; ++j) {
|
|
sum += sigmoid_vals[off + j];
|
|
}
|
|
float mean_val = sum / (float)len;
|
|
|
|
// Weighted contribution
|
|
float contribution = mean_val * weights[model];
|
|
atomicAdd(output, contribution);
|
|
}
|
|
"#;
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
/// A simple test adapter with configurable direction, confidence, and readiness.
|
|
struct DummyAdapter {
|
|
name: String,
|
|
direction: f64,
|
|
confidence: f64,
|
|
ready: bool,
|
|
}
|
|
|
|
impl ModelInferenceAdapter for DummyAdapter {
|
|
fn model_name(&self) -> &str {
|
|
&self.name
|
|
}
|
|
|
|
fn predict(&self, _features: &FeatureVector) -> MLResult<EnsemblePrediction> {
|
|
Ok(EnsemblePrediction {
|
|
model_name: self.name.clone(),
|
|
direction: self.direction,
|
|
confidence: self.confidence,
|
|
metadata: PredictionMeta::default(),
|
|
})
|
|
}
|
|
|
|
fn is_ready(&self) -> bool {
|
|
self.ready
|
|
}
|
|
}
|
|
|
|
/// An adapter that always fails prediction.
|
|
struct FailingAdapter {
|
|
name: String,
|
|
}
|
|
|
|
impl ModelInferenceAdapter for FailingAdapter {
|
|
fn model_name(&self) -> &str {
|
|
&self.name
|
|
}
|
|
|
|
fn predict(&self, _features: &FeatureVector) -> MLResult<EnsemblePrediction> {
|
|
Err(MLError::InferenceError("intentional test failure".to_owned()))
|
|
}
|
|
|
|
fn is_ready(&self) -> bool {
|
|
true
|
|
}
|
|
}
|
|
|
|
fn make_features() -> FeatureVector {
|
|
FeatureVector {
|
|
values: vec![0.0; 51],
|
|
timestamp: 1_700_000_000_000_000,
|
|
}
|
|
}
|
|
|
|
fn test_device() -> MlDevice {
|
|
// Use CPU for tests — no CUDA dependency in unit tests
|
|
MlDevice::Cpu
|
|
}
|
|
|
|
#[test]
|
|
fn test_stream_ensemble_empty() {
|
|
let ensemble = StreamAwareEnsemble::new(vec![], vec![], &test_device())
|
|
.expect("empty ensemble should create");
|
|
assert_eq!(ensemble.adapter_count(), 0);
|
|
assert_eq!(ensemble.ready_count(), 0);
|
|
|
|
let pred = ensemble
|
|
.predict(&make_features())
|
|
.expect("empty predict should succeed");
|
|
assert_eq!(pred.direction, 0.0);
|
|
assert_eq!(pred.confidence, 0.0);
|
|
assert_eq!(pred.model_name, "STREAM_ENSEMBLE(empty)");
|
|
}
|
|
|
|
#[test]
|
|
fn test_stream_ensemble_single_model() {
|
|
let adapters: Vec<Box<dyn ModelInferenceAdapter>> = vec![Box::new(DummyAdapter {
|
|
name: "DQN".to_owned(),
|
|
direction: 0.7,
|
|
confidence: 0.9,
|
|
ready: true,
|
|
})];
|
|
|
|
let ensemble = StreamAwareEnsemble::new(adapters, vec![1.0], &test_device())
|
|
.expect("single model should create");
|
|
assert_eq!(ensemble.adapter_count(), 1);
|
|
assert_eq!(ensemble.ready_count(), 1);
|
|
|
|
let pred = ensemble
|
|
.predict(&make_features())
|
|
.expect("predict should succeed");
|
|
|
|
// Single model: direction comes straight through the CPU path
|
|
assert!(
|
|
(pred.direction - 0.7).abs() < 1e-9,
|
|
"expected direction 0.7, got {}",
|
|
pred.direction
|
|
);
|
|
assert!(
|
|
(pred.confidence - 0.9).abs() < 1e-9,
|
|
"expected confidence 0.9, got {}",
|
|
pred.confidence
|
|
);
|
|
assert!(pred.model_name.contains("DQN"));
|
|
}
|
|
|
|
#[test]
|
|
fn test_stream_ensemble_two_models_equal_weight() {
|
|
// Model A: bullish (dir=1.0, conf=0.8)
|
|
// Model B: bearish (dir=-1.0, conf=0.6)
|
|
// CPU path: weighted_direction = (1.0*0.5*0.8 + (-1.0)*0.5*0.6) / (0.5*0.8 + 0.5*0.6)
|
|
// = (0.4 - 0.3) / (0.4 + 0.3) = 0.1 / 0.7 ~ 0.1429
|
|
let adapters: Vec<Box<dyn ModelInferenceAdapter>> = vec![
|
|
Box::new(DummyAdapter {
|
|
name: "A".to_owned(),
|
|
direction: 1.0,
|
|
confidence: 0.8,
|
|
ready: true,
|
|
}),
|
|
Box::new(DummyAdapter {
|
|
name: "B".to_owned(),
|
|
direction: -1.0,
|
|
confidence: 0.6,
|
|
ready: true,
|
|
}),
|
|
];
|
|
|
|
let ensemble = StreamAwareEnsemble::new(adapters, vec![1.0, 1.0], &test_device())
|
|
.expect("two-model ensemble should create");
|
|
|
|
let pred = ensemble
|
|
.predict(&make_features())
|
|
.expect("predict should succeed");
|
|
|
|
// Net direction should be positive (bullish model has higher confidence)
|
|
assert!(
|
|
pred.direction > 0.0,
|
|
"direction should be positive, got {}",
|
|
pred.direction
|
|
);
|
|
assert!(
|
|
pred.direction < 0.5,
|
|
"direction should be < 0.5, got {}",
|
|
pred.direction
|
|
);
|
|
// Confidence = average = (0.8 + 0.6) / 2 = 0.7
|
|
assert!(
|
|
(pred.confidence - 0.7).abs() < 1e-9,
|
|
"expected confidence 0.7, got {}",
|
|
pred.confidence
|
|
);
|
|
assert!(pred.model_name.contains('A'));
|
|
assert!(pred.model_name.contains('B'));
|
|
}
|
|
|
|
#[test]
|
|
fn test_stream_ensemble_skips_unready_models() {
|
|
let adapters: Vec<Box<dyn ModelInferenceAdapter>> = vec![
|
|
Box::new(DummyAdapter {
|
|
name: "Ready".to_owned(),
|
|
direction: 1.0,
|
|
confidence: 0.9,
|
|
ready: true,
|
|
}),
|
|
Box::new(DummyAdapter {
|
|
name: "NotReady".to_owned(),
|
|
direction: -1.0,
|
|
confidence: 0.9,
|
|
ready: false,
|
|
}),
|
|
];
|
|
|
|
let ensemble = StreamAwareEnsemble::new(adapters, vec![1.0, 1.0], &test_device())
|
|
.expect("ensemble should create");
|
|
assert_eq!(ensemble.ready_count(), 1);
|
|
|
|
let pred = ensemble
|
|
.predict(&make_features())
|
|
.expect("predict should succeed");
|
|
|
|
assert!(pred.model_name.contains("Ready"));
|
|
assert!(!pred.model_name.contains("NotReady"));
|
|
assert!(pred.direction > 0.5);
|
|
}
|
|
|
|
#[test]
|
|
fn test_stream_ensemble_skips_nan_predictions() {
|
|
let adapters: Vec<Box<dyn ModelInferenceAdapter>> = vec![
|
|
Box::new(DummyAdapter {
|
|
name: "Valid".to_owned(),
|
|
direction: 0.8,
|
|
confidence: 0.7,
|
|
ready: true,
|
|
}),
|
|
Box::new(DummyAdapter {
|
|
name: "NaN_Model".to_owned(),
|
|
direction: f64::NAN,
|
|
confidence: 0.9,
|
|
ready: true,
|
|
}),
|
|
];
|
|
|
|
let ensemble = StreamAwareEnsemble::new(adapters, vec![1.0, 1.0], &test_device())
|
|
.expect("ensemble should create");
|
|
|
|
let pred = ensemble
|
|
.predict(&make_features())
|
|
.expect("predict should succeed");
|
|
|
|
assert!(pred.model_name.contains("Valid"));
|
|
assert!(!pred.model_name.contains("NaN_Model"));
|
|
assert!(
|
|
(pred.direction - 0.8).abs() < 1e-9,
|
|
"expected direction 0.8, got {}",
|
|
pred.direction
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_stream_ensemble_skips_failing_models() {
|
|
let adapters: Vec<Box<dyn ModelInferenceAdapter>> = vec![
|
|
Box::new(DummyAdapter {
|
|
name: "Good".to_owned(),
|
|
direction: 0.5,
|
|
confidence: 0.8,
|
|
ready: true,
|
|
}),
|
|
Box::new(FailingAdapter {
|
|
name: "Broken".to_owned(),
|
|
}),
|
|
];
|
|
|
|
let ensemble = StreamAwareEnsemble::new(adapters, vec![1.0, 1.0], &test_device())
|
|
.expect("ensemble should create");
|
|
|
|
let pred = ensemble
|
|
.predict(&make_features())
|
|
.expect("predict should succeed despite one failing model");
|
|
|
|
assert!(pred.model_name.contains("Good"));
|
|
assert!(!pred.model_name.contains("Broken"));
|
|
}
|
|
|
|
#[test]
|
|
fn test_stream_ensemble_weight_normalisation() {
|
|
// Weights [3.0, 1.0] should normalise to [0.75, 0.25]
|
|
let adapters: Vec<Box<dyn ModelInferenceAdapter>> = vec![
|
|
Box::new(DummyAdapter {
|
|
name: "Heavy".to_owned(),
|
|
direction: 1.0,
|
|
confidence: 0.8,
|
|
ready: true,
|
|
}),
|
|
Box::new(DummyAdapter {
|
|
name: "Light".to_owned(),
|
|
direction: -1.0,
|
|
confidence: 0.8,
|
|
ready: true,
|
|
}),
|
|
];
|
|
|
|
let ensemble = StreamAwareEnsemble::new(adapters, vec![3.0, 1.0], &test_device())
|
|
.expect("ensemble should create");
|
|
|
|
let pred = ensemble
|
|
.predict(&make_features())
|
|
.expect("predict should succeed");
|
|
|
|
// Heavy (0.75) bullish should dominate over Light (0.25) bearish
|
|
assert!(
|
|
pred.direction > 0.3,
|
|
"heavy bullish model should dominate, got {}",
|
|
pred.direction
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_stream_ensemble_all_models_fail_returns_neutral() {
|
|
let adapters: Vec<Box<dyn ModelInferenceAdapter>> = vec![
|
|
Box::new(FailingAdapter {
|
|
name: "Fail1".to_owned(),
|
|
}),
|
|
Box::new(FailingAdapter {
|
|
name: "Fail2".to_owned(),
|
|
}),
|
|
];
|
|
|
|
let ensemble = StreamAwareEnsemble::new(adapters, vec![1.0, 1.0], &test_device())
|
|
.expect("ensemble should create");
|
|
|
|
let pred = ensemble
|
|
.predict(&make_features())
|
|
.expect("predict should return neutral when all fail");
|
|
|
|
assert_eq!(pred.direction, 0.0);
|
|
assert_eq!(pred.confidence, 0.0);
|
|
assert_eq!(pred.model_name, "STREAM_ENSEMBLE(empty)");
|
|
}
|
|
|
|
#[test]
|
|
fn test_stream_ensemble_missing_weights_default() {
|
|
// Provide fewer weights than adapters — missing ones should default
|
|
let adapters: Vec<Box<dyn ModelInferenceAdapter>> = vec![
|
|
Box::new(DummyAdapter {
|
|
name: "A".to_owned(),
|
|
direction: 0.5,
|
|
confidence: 0.8,
|
|
ready: true,
|
|
}),
|
|
Box::new(DummyAdapter {
|
|
name: "B".to_owned(),
|
|
direction: 0.5,
|
|
confidence: 0.8,
|
|
ready: true,
|
|
}),
|
|
Box::new(DummyAdapter {
|
|
name: "C".to_owned(),
|
|
direction: 0.5,
|
|
confidence: 0.8,
|
|
ready: true,
|
|
}),
|
|
];
|
|
|
|
let ensemble = StreamAwareEnsemble::new(adapters, vec![1.0], &test_device())
|
|
.expect("missing weights should use defaults");
|
|
|
|
let pred = ensemble
|
|
.predict(&make_features())
|
|
.expect("predict should succeed");
|
|
|
|
// All models agree on 0.5, so direction should be ~0.5
|
|
assert!(
|
|
(pred.direction - 0.5).abs() < 1e-9,
|
|
"expected direction 0.5, got {}",
|
|
pred.direction
|
|
);
|
|
}
|
|
}
|