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:
95
docs/plans/2026-02-24-ml-production-hardening-design.md
Normal file
95
docs/plans/2026-02-24-ml-production-hardening-design.md
Normal 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
|
||||
Reference in New Issue
Block a user