fix(ml): OOM hardening + battle-test KAN/xLSTM/Diffusion models

Replace 8 unbounded Vec accumulation patterns with bounded VecDeque
across ensemble, PPO, DQN, Mamba2, and data pipeline code to prevent
OOM on RTX 3050 Ti (4GB VRAM) during live trading and extended training.

Key OOM fixes:
- Ensemble price/volatility history: Vec → VecDeque with O(1) eviction
- Data pipeline: MAX_FEATURES=500K cap (~512MB) prevents unbounded loading
- DQN replay buffer: full-array shuffle → HashSet random sampling (8MB → 256B)
- PPO loss histories: bounded VecDeque (cap 1K), eliminated batch.clone()
- Mamba2 scan: pre-allocated Vecs, explicit drop() after Tensor::cat
- Mamba2 training history: capped at 100, Tensor::randn replaces Vec→Tensor
- Mamba2 SSM reset: 2 unwrap() violations replaced with proper error handling

Battle-testing (19 new integration tests):
- KAN: 5 tests (forward, 50-epoch training 89.9% loss reduction, checkpoint)
- xLSTM: 7 tests (2D+3D forward, 30-epoch training 82% reduction, checkpoint)
- Diffusion: 7 tests (2D+3D forward, 20-epoch pipeline, checkpoint, validation)

Bonus: fix pre-existing cache test failure (match .dbn.zst files, graceful skip)

All 2390 lib tests pass, 0 new clippy errors.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-02-24 09:55:39 +01:00
parent c3b5e124f0
commit 57bae2cb68
14 changed files with 872 additions and 148 deletions

View File

@@ -0,0 +1,95 @@
# ML Production Hardening Design
**Date**: 2026-02-24
**Branch**: `fix/ml-production-hardening`
**Scope**: OOM hardening (27 issues) + battle-testing KAN/xLSTM/Diffusion (3 models)
## Problem
1. **OOM risks**: 8 critical unbounded buffer patterns across ensemble, PPO, Mamba2, DQN, TFT, and data pipeline code. On RTX 3050 Ti (4GB VRAM), these can cause silent crashes during live trading or extended training.
2. **Untested models**: KAN, xLSTM, and Diffusion have zero integration tests and are not wired into the ensemble coordinator. They compile but have never been validated end-to-end.
## Strategy
Single worktree. Parallel coding on independent files. Sequential model validation (one at a time, 4GB VRAM constraint).
## Phase 1: OOM Hardening (parallel agents on independent files)
### Pattern
Replace all unbounded `Vec::push()` accumulation with bounded `VecDeque`:
```rust
use std::collections::VecDeque;
// Bounded push — drop oldest when at capacity
if buf.len() >= cap {
buf.pop_front();
}
buf.push_back(item);
```
### Fixes
| # | File | Issue | Fix | Severity |
|---|------|-------|-----|----------|
| 1 | `ensemble/adaptive_ml_integration.rs` | Unbounded price_history/volatility_history | VecDeque cap 10K | Critical |
| 2 | `ppo/trajectories.rs` | Unbounded trajectory collection | Add max_trajectories param, chunk | Critical |
| 3 | `mamba/scan_algorithms.rs` | Tensor .clone() in inner loop | Drop seq_results after cat() | Critical |
| 4 | `data_pipeline/manager.rs` | All features loaded into one Vec | Stream via chunks | Critical |
| 5 | `tft/hft_optimizations.rs` | AttentionCache no eviction | LRU cap 256 entries | Critical |
| 6 | `dqn/replay_buffer.rs` | Full index collect for sampling | Random index sampling | Critical |
| 7 | `mamba/mod.rs` | Unbounded training_history | VecDeque cap 100 | Critical |
| 8 | `trainers/ppo.rs` | Unbounded value_loss_history | VecDeque cap 1K | Critical |
| 9 | `dqn/ensemble_uncertainty.rs` | Unbounded metrics history | VecDeque cap 10K | High |
| 10 | `trainers/ppo.rs:816` | batch.clone() in training loop | Use &mut or move | High |
| 11 | `mamba/mod.rs:296-323` | Vec->Tensor for SSM reset | Tensor::randn() direct | High |
| 12 | `tgnn/message_passing.rs` | Attention recomputation | Cache forward pass scores | High |
| 13 | `dqn/hindsight_replay.rs` | All HER experiences materialized | Generate on-the-fly | High |
### Agent Assignment (parallel, no file conflicts)
- **Agent A**: Ensemble + Data Pipeline (#1, #4)
- **Agent B**: DQN fixes (#6, #9, #13)
- **Agent C**: PPO fixes (#2, #8, #10)
- **Agent D**: Mamba2 + TFT + TGGN (#3, #5, #7, #11, #12)
## Phase 2: Model Battle-Testing (sequential)
For each of KAN, xLSTM, Diffusion:
1. **Integration test** (`ml/tests/<model>_integration.rs`):
- Construct with default config
- Forward pass on synthetic [batch=4, seq=32, features=10]
- Train 5 epochs, assert loss decreases
- Checkpoint save/load roundtrip
- Wire into ensemble coordinator, verify prediction
2. **Ensemble wiring** in `ensemble/coordinator.rs`:
- Add model to InferenceAdapterBridge
- Register in EnsembleCoordinator initialization
### Order: KAN → xLSTM → Diffusion → smoke test all 10
## Phase 3: Validation
```bash
SQLX_OFFLINE=true cargo test -p ml --lib # All 2204+ tests pass
SQLX_OFFLINE=true cargo test -p ml --test kan_integration
SQLX_OFFLINE=true cargo test -p ml --test xlstm_integration
SQLX_OFFLINE=true cargo test -p ml --test diffusion_integration
```
## Files Touched
- 13 existing files (OOM fixes)
- 3 new test files (integration tests)
- 1-2 existing files (ensemble wiring)
## Success Criteria
- All 8 critical OOM patterns replaced with bounded alternatives
- KAN, xLSTM, Diffusion each have passing integration tests
- All 10 models wired into ensemble coordinator
- `cargo test -p ml --lib` passes (2204+ tests)
- Zero new clippy warnings