docs: add operational maturity system design (3 pillars)
7-gate conviction system, autonomous feedback loop with kill switch, and Rust-native model registry — backed by PostgreSQL + QuestDB. Resolve merge conflicts in hyperopt/adapters/mod.rs and enhanced_ml.rs. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
153
docs/infra/scaleway-gpu-training.md
Normal file
153
docs/infra/scaleway-gpu-training.md
Normal file
@@ -0,0 +1,153 @@
|
||||
# Scaleway GPU Training Infrastructure
|
||||
|
||||
## Instance Types
|
||||
|
||||
| Instance | GPU | VRAM | vCPUs | RAM | Cost/hr |
|
||||
|----------|-----|------|-------|-----|---------|
|
||||
| GPU-3070-S | RTX 3070 | 8GB | 8 | 32GB | ~EUR 0.90 |
|
||||
| L4-1-24G | L4 | 24GB | 8 | 48GB | ~EUR 1.20 |
|
||||
| L4-2-48G | 2x L4 | 48GB | 16 | 96GB | ~EUR 2.40 |
|
||||
|
||||
**Recommendation:** GPU-3070-S for single-model training runs. L4-1-24G for hyperopt parallel trials and large batch training.
|
||||
|
||||
## Setup
|
||||
|
||||
### 1. Instance provisioning
|
||||
|
||||
```bash
|
||||
# Create GPU instance via Scaleway CLI
|
||||
scw instance server create \
|
||||
type=GPU-3070-S \
|
||||
image=ubuntu_jammy \
|
||||
name=foxhunt-training \
|
||||
root-volume=l:100G
|
||||
```
|
||||
|
||||
### 2. Initial setup (run once after provisioning)
|
||||
|
||||
```bash
|
||||
# Install Rust toolchain
|
||||
curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y
|
||||
source ~/.cargo/env
|
||||
|
||||
# Install CUDA toolkit (Ubuntu 22.04)
|
||||
wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.1-1_all.deb
|
||||
sudo dpkg -i cuda-keyring_1.1-1_all.deb
|
||||
sudo apt-get update
|
||||
sudo apt-get -y install cuda-toolkit-12-4
|
||||
|
||||
# Clone and build
|
||||
git clone ssh://gitea@git.fxhnt.ai:2222/foxhunt/foxhunt.git /opt/foxhunt
|
||||
cd /opt/foxhunt
|
||||
SQLX_OFFLINE=true cargo build --release -p ml --features cuda
|
||||
```
|
||||
|
||||
### 3. Tailscale access
|
||||
|
||||
```bash
|
||||
# Install Tailscale for secure access
|
||||
curl -fsSL https://tailscale.com/install.sh | sh
|
||||
sudo tailscale up --ssh
|
||||
# Add to Tailscale ACL: tag:gpu-training
|
||||
```
|
||||
|
||||
### 4. DNS (optional)
|
||||
|
||||
Add a Tailscale MagicDNS entry or Scaleway DNS record:
|
||||
|
||||
```
|
||||
gpu.fxhnt.ai -> <tailscale-ip>
|
||||
```
|
||||
|
||||
The `train_launcher.sh` script defaults to `FOXHUNT_GPU_HOST=gpu.fxhnt.ai`.
|
||||
|
||||
## Training workflow
|
||||
|
||||
### Local development then cloud production
|
||||
|
||||
```bash
|
||||
# 1. Develop and test locally (RTX 3050 Ti, 4GB)
|
||||
./scripts/train_launcher.sh --model dqn --epochs 5
|
||||
|
||||
# 2. Push code to Gitea
|
||||
git push origin feat/model-improvements
|
||||
|
||||
# 3. Train on cloud GPU
|
||||
./scripts/train_launcher.sh --model dqn --epochs 100 --cloud
|
||||
|
||||
# 4. Sync trained model back
|
||||
rsync -avz training@gpu.fxhnt.ai:/opt/foxhunt/checkpoints/ ./checkpoints/
|
||||
rsync -avz training@gpu.fxhnt.ai:/opt/foxhunt/ml/trained_models/ ./ml/trained_models/
|
||||
```
|
||||
|
||||
### Model artifact sync
|
||||
|
||||
```bash
|
||||
# Upload training data to cloud
|
||||
rsync -avz ./data/databento/ training@gpu.fxhnt.ai:/opt/foxhunt/data/databento/
|
||||
rsync -avz ./test_data/ training@gpu.fxhnt.ai:/opt/foxhunt/test_data/
|
||||
|
||||
# Download trained models
|
||||
rsync -avz training@gpu.fxhnt.ai:/opt/foxhunt/checkpoints/ ./checkpoints/
|
||||
rsync -avz training@gpu.fxhnt.ai:/opt/foxhunt/ml/trained_models/ ./ml/trained_models/
|
||||
```
|
||||
|
||||
## Batch sizes by model and GPU
|
||||
|
||||
These values are used by `scripts/train_launcher.sh` auto-detection.
|
||||
See `scripts/measure_vram.sh` for VRAM profiling methodology.
|
||||
|
||||
| Model | RTX 3050 Ti (4GB) | RTX 3070 (8GB) | L4 (24GB) |
|
||||
|-------|-------------------|----------------|-----------|
|
||||
| DQN | 128 | 256 | 512 |
|
||||
| PPO | 230 (max) | 512 | 1024 |
|
||||
| TFT | 32 | 64 | 256 |
|
||||
| Mamba2 | 64 | 128 | 512 |
|
||||
| CfC | 128 | 256 | 512 |
|
||||
|
||||
Notes:
|
||||
- PPO on 4GB is capped at 230 batches (empirically verified, see MEMORY.md).
|
||||
- TFT and Mamba2 are memory-intensive due to attention/state-space layers.
|
||||
- All values assume `--features cuda` and `--release` builds.
|
||||
|
||||
## Cargo example targets
|
||||
|
||||
The launcher maps model names to these `ml` crate examples:
|
||||
|
||||
| Model | Cargo example | Source |
|
||||
|-------|--------------|--------|
|
||||
| dqn | `train_dqn` | `ml/examples/train_dqn.rs` |
|
||||
| ppo | `train_ppo_parquet` | `ml/examples/train_ppo_parquet.rs` |
|
||||
| tft | `train_tft_parquet` | `ml/examples/train_tft_parquet.rs` |
|
||||
| mamba2 | `train_mamba2_parquet` | `ml/examples/train_mamba2_parquet.rs` |
|
||||
| cfc | `train_liquid_dbn` | `ml/examples/train_liquid_dbn.rs` |
|
||||
|
||||
## Cost estimates
|
||||
|
||||
| Training run | Instance | Duration | Est. cost |
|
||||
|-------------|----------|----------|-----------|
|
||||
| DQN 100 epochs | GPU-3070-S | ~2 hours | ~EUR 1.80 |
|
||||
| PPO 100 epochs | GPU-3070-S | ~3 hours | ~EUR 2.70 |
|
||||
| TFT 100 epochs | L4-1-24G | ~4 hours | ~EUR 4.80 |
|
||||
| Full hyperopt (50 trials) | L4-1-24G | ~12 hours | ~EUR 14.40 |
|
||||
| All 5 models training | L4-1-24G | ~8 hours | ~EUR 9.60 |
|
||||
|
||||
## Environment variables
|
||||
|
||||
| Variable | Default | Description |
|
||||
|----------|---------|-------------|
|
||||
| `FOXHUNT_GPU_HOST` | `gpu.fxhnt.ai` | Cloud GPU hostname |
|
||||
| `FOXHUNT_CLOUD_DIR` | `/opt/foxhunt` | Remote project directory |
|
||||
| `SQLX_OFFLINE` | `true` | Required (no PostgreSQL on GPU instances) |
|
||||
|
||||
## Shutdown procedure
|
||||
|
||||
GPU instances are billed per hour. Always stop when not training:
|
||||
|
||||
```bash
|
||||
# From local machine
|
||||
scw instance server stop foxhunt-training
|
||||
|
||||
# Or from the instance itself
|
||||
sudo shutdown -h now
|
||||
```
|
||||
380
docs/plans/2026-02-23-operational-maturity-design.md
Normal file
380
docs/plans/2026-02-23-operational-maturity-design.md
Normal file
@@ -0,0 +1,380 @@
|
||||
# Foxhunt Operational Maturity System — Design Document
|
||||
|
||||
**Date:** 2026-02-23
|
||||
**Status:** Approved
|
||||
**Philosophy:** "Be right, not always" — high-conviction, quality-over-quantity trading
|
||||
|
||||
## Overview
|
||||
|
||||
Three integrated pillars that close the feedback loop from prediction to execution to learning:
|
||||
|
||||
1. **7-Gate Conviction System** — enforce ensemble quality before trading
|
||||
2. **Autonomous Feedback Loop** — self-adjusting weights, thresholds, retraining triggers
|
||||
3. **Rust-Native Model Registry** — experiment tracking and lifecycle management
|
||||
|
||||
Data layer: PostgreSQL (audit, registry) + QuestDB (time-series analytics).
|
||||
|
||||
## Architecture
|
||||
|
||||
```
|
||||
┌─────────────────────────────────────────────────────────────────┐
|
||||
│ TRADING PIPELINE (critical path) │
|
||||
│ │
|
||||
│ Market Data → Features(51) → Ensemble(10 models) → 7 Gates │
|
||||
│ ↓ │
|
||||
│ Gate Pass? ──No──→ HOLD
|
||||
│ ↓ Yes │
|
||||
│ Conviction Sizing │
|
||||
│ ↓ │
|
||||
│ Risk Validation │
|
||||
│ ↓ │
|
||||
│ Order Execution │
|
||||
│ ↓ │
|
||||
│ PostgreSQL (audit) │
|
||||
└─────────────────────────────────────────────────────────────────┘
|
||||
│ (writes metrics) │ (trade close events)
|
||||
▼ ▼
|
||||
┌───────────────────────────────────────────────────────────────┐
|
||||
│ ANALYTICS LAYER (non-critical path) │
|
||||
│ │
|
||||
│ Ring Buffer (10K) ──→ QuestDB │
|
||||
│ - model_predictions │
|
||||
│ - trade_attribution │
|
||||
│ - gate_performance │
|
||||
│ - system_health │
|
||||
│ │ │
|
||||
│ ▼ │
|
||||
│ AUTONOMOUS FEEDBACK LOOP (periodic, every 24h) │
|
||||
│ 1. Query rolling metrics from QuestDB │
|
||||
│ 2. Calculate new weights (EMA of Sharpe, bounded) │
|
||||
│ 3. Optimize gate thresholds (win rate by bucket) │
|
||||
│ 4. Check retraining triggers (drift, accuracy, Sharpe) │
|
||||
│ 5. Apply changes (max ±0.03/cycle, 24h cooldown) │
|
||||
│ 6. Snapshot config to PostgreSQL (rollback point) │
|
||||
│ │
|
||||
│ KILL SWITCH: 7-day Sharpe < -1.0 → freeze + revert + alert │
|
||||
└───────────────────────────────────────────────────────────────┘
|
||||
│ (model lifecycle events)
|
||||
▼
|
||||
┌───────────────────────────────────────────────────────────────┐
|
||||
│ MODEL REGISTRY (PostgreSQL) │
|
||||
│ │
|
||||
│ ml_model_versions (extended) │
|
||||
│ model_stages (new) │
|
||||
│ training_runs (extended from existing) │
|
||||
│ │
|
||||
│ CheckpointManager → register_run() → promote() → deploy() │
|
||||
└───────────────────────────────────────────────────────────────┘
|
||||
```
|
||||
|
||||
## Pillar 1: 7-Gate Conviction System
|
||||
|
||||
Gates are evaluated in order. Any gate failure results in HOLD (no trade).
|
||||
|
||||
| # | Gate | Metric | Default Threshold | Auto-Adjusted |
|
||||
|---|------|--------|-------------------|---------------|
|
||||
| 1 | Model Health | `healthy_models / total_models` | >= 0.70 | No (safety) |
|
||||
| 2 | Time-of-Day | `current_session in allowed_sessions` | Regular hours only | No (policy) |
|
||||
| 3 | Confidence | `ensemble_confidence` | >= 0.60 | Yes |
|
||||
| 4 | Agreement | `disagreement_rate` | <= 0.40 | Yes |
|
||||
| 5 | Quorum | `agreeing_models / active_models` | >= 0.60 | Yes |
|
||||
| 6 | Regime | Volatility-adjusted multiplier on gates 3-5 | 0.80x in high-vol | Yes |
|
||||
| 7 | Conviction Sizing | `confidence × (1 - disagreement) × quorum × health` | 0→1 scalar | Yes |
|
||||
|
||||
### Gate Details
|
||||
|
||||
**Gate 1 — Model Health:** Checks that enough models produced valid, recent predictions. A model is "unhealthy" if its last prediction errored, timed out (>100ms), or is stale (>5 minutes old). If <70% of models are healthy, system holds rather than trading on incomplete ensemble.
|
||||
|
||||
**Gate 2 — Time-of-Day:** Configurable per-session thresholds. Default: only trade during regular market hours (09:30-16:00 ET). Pre-market and after-hours can be enabled with tighter thresholds on gates 3-5.
|
||||
|
||||
**Gate 3 — Confidence:** Weighted consensus strength from ensemble. Already calculated in `EnsembleCoordinator` but not enforced. Default minimum: 0.60.
|
||||
|
||||
**Gate 4 — Agreement:** Percentage of models contradicting the mean signal direction. Already calculated as `disagreement_rate`. Default maximum: 0.40 (at most 40% of models can disagree).
|
||||
|
||||
**Gate 5 — Quorum:** Count of models agreeing on direction (buy/sell/hold) divided by active models. Different from disagreement (which measures magnitude). Default minimum: 0.60 (6/10 models must agree on direction).
|
||||
|
||||
**Gate 6 — Regime Filter:** Uses existing regime features (dimensions 48-50 in feature vector). In high-volatility regimes, multiplies thresholds for gates 3-5 by a tightening factor (default 0.80x, meaning confidence must be 0.75 instead of 0.60).
|
||||
|
||||
**Gate 7 — Conviction Sizing:** Composite scalar that multiplies the base position size. Formula: `confidence × (1 - disagreement) × quorum_ratio × health_ratio`. Range: 0.0 to 1.0. A trade passing all gates with marginal scores gets smaller position size than one passing with strong scores.
|
||||
|
||||
### Configuration
|
||||
|
||||
```rust
|
||||
pub struct ConvictionGateConfig {
|
||||
pub model_health_threshold: f64, // 0.70
|
||||
pub allowed_sessions: Vec<TradingSession>,
|
||||
pub min_confidence: f64, // 0.60
|
||||
pub max_disagreement: f64, // 0.40
|
||||
pub min_quorum: f64, // 0.60
|
||||
pub regime_tightening_factor: f64, // 0.80
|
||||
pub high_vol_threshold: f64, // Regime volatility cutoff
|
||||
pub conviction_scaling_enabled: bool, // true
|
||||
}
|
||||
```
|
||||
|
||||
### Location
|
||||
|
||||
`ml/src/ensemble/conviction_gates.rs` — new file. Wired into `EnsembleCoordinator::generate_decision()` in `coordinator.rs`.
|
||||
|
||||
## Pillar 2: Autonomous Feedback Loop
|
||||
|
||||
### Attribution Calculator
|
||||
|
||||
Triggered on every trade close/exit. Decomposes realized P&L into per-model contributions.
|
||||
|
||||
**Attribution formula:**
|
||||
- `signal_alignment = sign(model_signal) == sign(realized_direction) ? 1.0 : -1.0`
|
||||
- `model_contribution = model_weight × signal_alignment × abs(realized_pnl)`
|
||||
- If model voted "buy" and trade was profitable → positive attribution
|
||||
- If model voted "sell" but ensemble went "buy" and lost → positive attribution (model was right)
|
||||
|
||||
**Output:** Written to QuestDB `trade_attribution` table.
|
||||
|
||||
### Rolling Performance Engine
|
||||
|
||||
QuestDB queries running on configurable windows (1-day, 7-day, 30-day, 90-day).
|
||||
|
||||
Per-model metrics:
|
||||
- Rolling Sharpe ratio: `avg(pnl) / stddev(pnl) * sqrt(252)`
|
||||
- Win rate: `count(correct) / count(total)` per directional call
|
||||
- Attribution P&L: Cumulative contribution
|
||||
- Signal accuracy: Predicted direction vs actual price movement
|
||||
- Prediction stability: Variance of signals over time
|
||||
|
||||
Per-gate metrics:
|
||||
- Win rate at each confidence bucket (0.5-0.6, 0.6-0.7, 0.7-0.8, 0.8+)
|
||||
- Average P&L by disagreement level
|
||||
- Performance by time-of-day session
|
||||
|
||||
### Weight Optimizer
|
||||
|
||||
Runs periodically (configurable: default every 24 hours, or after every N trades).
|
||||
|
||||
**Algorithm:** Exponentially-weighted moving average of rolling Sharpe ratios.
|
||||
- `raw_weight[i] = ema(sharpe_30d[i], alpha=0.1)`
|
||||
- `new_weight[i] = normalize(clamp(raw_weight[i], MIN_WEIGHT, MAX_WEIGHT))`
|
||||
- Sum normalized to 1.0
|
||||
- Applied via existing `EnsembleCoordinator::update_model_weights()`
|
||||
|
||||
### Gate Threshold Optimizer
|
||||
|
||||
Same periodic cycle as weight optimizer.
|
||||
|
||||
**Algorithm:** Analyze win rate by confidence bucket.
|
||||
- If trades at confidence 0.55-0.65 have <45% win rate → raise min_confidence by 0.03
|
||||
- If regime=high_vol trades underperform by >20% → increase regime tightening factor by 0.03
|
||||
- Bounded: no gate can be tighter than 0.90 or looser than 0.30
|
||||
|
||||
### Retraining Triggers
|
||||
|
||||
Conditions for automatic retraining (emits `RetrainRequest` to ml_training_service):
|
||||
- Model's 30-day Sharpe drops below -0.5 (consistently losing)
|
||||
- Model's prediction accuracy falls below 45% (worse than random)
|
||||
- Model's signal correlation with best-performing model exceeds 0.9 (redundant)
|
||||
|
||||
### Safety Rails
|
||||
|
||||
| Rail | Value |
|
||||
|------|-------|
|
||||
| Max weight change per cycle | ±0.03 |
|
||||
| Max threshold change per cycle | ±0.03 |
|
||||
| Min observations before first adjustment | 100 trades |
|
||||
| Cooldown between adjustments | 24 hours |
|
||||
| New model deployment grace period | 7 days |
|
||||
| Kill switch trigger | 7-day Sharpe < -1.0 |
|
||||
| Kill switch action | Freeze all adjustments, revert to last-known-good config, alert via webhook |
|
||||
| Kill switch recovery | Requires human acknowledgment to re-enable autonomy |
|
||||
| Weight bounds | [0.05, 0.40] |
|
||||
| Gate threshold bounds | [0.30, 0.90] |
|
||||
|
||||
### Location
|
||||
|
||||
- `ml/src/ensemble/weight_optimizer.rs` — weight adjustment logic
|
||||
- `ml/src/ensemble/gate_optimizer.rs` — threshold adjustment logic
|
||||
- `services/trading_service/src/attribution.rs` — per-trade P&L attribution
|
||||
- `services/trading_service/src/feedback_loop.rs` — orchestrates periodic optimization cycle
|
||||
|
||||
## Pillar 3: Rust-Native Model Registry
|
||||
|
||||
### Schema Changes
|
||||
|
||||
**Extend `ml_model_versions`:**
|
||||
```sql
|
||||
ALTER TABLE ml_model_versions ADD COLUMN experiment_id UUID;
|
||||
ALTER TABLE ml_model_versions ADD COLUMN git_commit VARCHAR;
|
||||
ALTER TABLE ml_model_versions ADD COLUMN data_hash VARCHAR;
|
||||
ALTER TABLE ml_model_versions ADD COLUMN run_status VARCHAR DEFAULT 'completed';
|
||||
ALTER TABLE ml_model_versions ADD COLUMN started_at TIMESTAMP;
|
||||
ALTER TABLE ml_model_versions ADD COLUMN finished_at TIMESTAMP;
|
||||
```
|
||||
|
||||
**New table — `model_stages`:**
|
||||
```sql
|
||||
CREATE TABLE model_stages (
|
||||
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
||||
model_version_id UUID REFERENCES ml_model_versions(id),
|
||||
model_type VARCHAR NOT NULL,
|
||||
stage VARCHAR NOT NULL, -- 'candidate' | 'staging' | 'production' | 'archived'
|
||||
promoted_at TIMESTAMP DEFAULT NOW(),
|
||||
promoted_by VARCHAR NOT NULL, -- 'ab_test' | 'manual' | 'auto_optimizer' | 'retraining'
|
||||
performance_snapshot JSONB, -- metrics at promotion time
|
||||
reverted_at TIMESTAMP, -- NULL unless reverted
|
||||
revert_reason VARCHAR
|
||||
);
|
||||
CREATE INDEX idx_model_stages_type_stage ON model_stages(model_type, stage);
|
||||
CREATE INDEX idx_model_stages_version ON model_stages(model_version_id);
|
||||
```
|
||||
|
||||
### Rust API
|
||||
|
||||
```rust
|
||||
pub trait ModelRegistry {
|
||||
async fn log_run(&self, run: &TrainingRun) -> Result<RunId>;
|
||||
async fn log_metrics(&self, run_id: RunId, metrics: &HashMap<String, f64>, step: u32) -> Result<()>;
|
||||
async fn promote(&self, model_version_id: Uuid, to_stage: ModelStage, promoted_by: &str) -> Result<()>;
|
||||
async fn get_production_model(&self, model_type: ModelType) -> Result<Option<ModelVersion>>;
|
||||
async fn revert(&self, model_version_id: Uuid, reason: &str) -> Result<()>;
|
||||
async fn compare_runs(&self, run_ids: &[Uuid]) -> Result<RunComparison>;
|
||||
}
|
||||
```
|
||||
|
||||
### Integration Points
|
||||
|
||||
- `CheckpointManager::register_checkpoint()` → calls `log_run()` with hyperparameters, metrics, git commit
|
||||
- `DeploymentPipeline::deploy()` → calls `promote(staging → production)` with performance snapshot
|
||||
- Feedback loop retraining trigger → creates new run with `promoted_by: "auto_optimizer"`
|
||||
- Kill switch revert → calls `revert()` with reason, rolls back to previous production model
|
||||
|
||||
### Location
|
||||
|
||||
- `ml/src/registry/mod.rs` — `ModelRegistry` trait + `PostgresModelRegistry` impl
|
||||
|
||||
## Data Layer: QuestDB Integration
|
||||
|
||||
### Deployment
|
||||
|
||||
- Docker container on Scaleway (alongside Gitea and PostgreSQL)
|
||||
- ILP ingestion on port 9009 (high-throughput writes)
|
||||
- PostgreSQL wire protocol on port 8812 (queries via sqlx)
|
||||
- Web console on port 9000 (optional, for debugging)
|
||||
|
||||
### QuestDB Tables
|
||||
|
||||
```sql
|
||||
-- Every ensemble prediction (written per-tick when model runs)
|
||||
CREATE TABLE model_predictions (
|
||||
timestamp TIMESTAMP,
|
||||
model_id SYMBOL,
|
||||
signal DOUBLE,
|
||||
confidence DOUBLE,
|
||||
direction SHORT, -- -1 (sell), 0 (hold), 1 (buy)
|
||||
latency_us LONG,
|
||||
was_healthy BOOLEAN
|
||||
) TIMESTAMP(timestamp) PARTITION BY DAY;
|
||||
|
||||
-- Per-trade P&L attribution (written on trade close)
|
||||
CREATE TABLE trade_attribution (
|
||||
timestamp TIMESTAMP,
|
||||
trade_id SYMBOL,
|
||||
model_id SYMBOL,
|
||||
model_weight DOUBLE,
|
||||
model_signal DOUBLE,
|
||||
signal_alignment DOUBLE, -- 1.0 if model agreed with outcome, -1.0 if not
|
||||
pnl_contribution DOUBLE,
|
||||
realized_pnl DOUBLE,
|
||||
symbol SYMBOL
|
||||
) TIMESTAMP(timestamp) PARTITION BY DAY;
|
||||
|
||||
-- Gate decision log (written per ensemble decision)
|
||||
CREATE TABLE gate_performance (
|
||||
timestamp TIMESTAMP,
|
||||
gate_name SYMBOL, -- 'confidence', 'agreement', 'quorum', etc.
|
||||
gate_value DOUBLE,
|
||||
gate_threshold DOUBLE,
|
||||
gate_passed BOOLEAN,
|
||||
subsequent_pnl DOUBLE -- filled post-trade for analysis
|
||||
) TIMESTAMP(timestamp) PARTITION BY DAY;
|
||||
|
||||
-- System health (written every 30s)
|
||||
CREATE TABLE system_health (
|
||||
timestamp TIMESTAMP,
|
||||
questdb_connected BOOLEAN,
|
||||
buffer_size LONG,
|
||||
healthy_models SHORT,
|
||||
total_models SHORT,
|
||||
ensemble_sharpe_7d DOUBLE,
|
||||
kill_switch_active BOOLEAN
|
||||
) TIMESTAMP(timestamp) PARTITION BY DAY;
|
||||
```
|
||||
|
||||
### Failure Handling
|
||||
|
||||
QuestDB is non-critical path. If unavailable:
|
||||
- Trading continues normally (gates use live data, not QuestDB)
|
||||
- Metrics buffer in 10K-entry ring buffer
|
||||
- Autonomous feedback loop pauses (freezes current config)
|
||||
- Alert emitted after 60s disconnection
|
||||
|
||||
### Health Monitoring
|
||||
|
||||
- Periodic ping every 30s via `SELECT 1`
|
||||
- Prometheus metrics: `questdb_connected`, `questdb_buffer_size`, `questdb_buffer_age_seconds`
|
||||
- Alert thresholds: buffer_size > 5000 (warning), buffer_age > 300s (warning), disconnected > 60s (alert)
|
||||
|
||||
### Rust Client
|
||||
|
||||
```rust
|
||||
pub struct QuestDBClient {
|
||||
pg_pool: PgPool, // For queries (port 8812)
|
||||
ilp_sender: IlpSender, // For writes (port 9009)
|
||||
buffer: RingBuffer<MetricEntry>, // 10K capacity
|
||||
health: Arc<AtomicBool>, // Connected status
|
||||
}
|
||||
```
|
||||
|
||||
Using `sqlx::PgPool` for queries (QuestDB speaks Postgres wire protocol) and ILP over TCP for high-throughput writes.
|
||||
|
||||
### Retention
|
||||
|
||||
- Detailed data: 90 days
|
||||
- Daily aggregates: indefinite (downsampled via SAMPLE BY 1d)
|
||||
|
||||
## New Files
|
||||
|
||||
| File | Purpose | Est. Lines |
|
||||
|------|---------|-----------|
|
||||
| `ml/src/ensemble/conviction_gates.rs` | 7-gate system, ConvictionGateConfig | ~300 |
|
||||
| `ml/src/ensemble/weight_optimizer.rs` | Autonomous weight adjustment | ~250 |
|
||||
| `ml/src/ensemble/gate_optimizer.rs` | Autonomous threshold adjustment | ~200 |
|
||||
| `ml/src/registry/mod.rs` | ModelRegistry trait + PostgreSQL impl | ~300 |
|
||||
| `common/src/questdb.rs` | QuestDB client, ring buffer, health monitor | ~400 |
|
||||
| `services/trading_service/src/attribution.rs` | Per-trade P&L attribution calculator | ~200 |
|
||||
| `services/trading_service/src/feedback_loop.rs` | Periodic optimization orchestrator | ~350 |
|
||||
|
||||
## Modified Files
|
||||
|
||||
| File | Change |
|
||||
|------|--------|
|
||||
| `ml/src/ensemble/coordinator.rs` | Wire conviction gates into generate_decision() |
|
||||
| `ml/src/ensemble/mod.rs` | Add conviction_gates, weight_optimizer, gate_optimizer modules |
|
||||
| `services/trading_service/src/main.rs` | Init QuestDB client, feedback loop, gate config |
|
||||
| `services/ml_training_service/src/checkpoint_manager.rs` | Log to model registry |
|
||||
| `services/ml_training_service/src/deployment_pipeline.rs` | Promote via registry stages |
|
||||
| Database migration | Extend ml_model_versions, add model_stages table |
|
||||
|
||||
## Infrastructure
|
||||
|
||||
| Component | Where | Docker Image | Ports |
|
||||
|-----------|-------|-------------|-------|
|
||||
| QuestDB | Scaleway DEV1-S | `questdb/questdb:latest` | 9009 (ILP), 8812 (PG), 9000 (console) |
|
||||
|
||||
## Verification Plan
|
||||
|
||||
1. Unit tests for each gate in conviction_gates.rs
|
||||
2. Unit tests for weight optimizer (bounded, normalized, cooldown enforced)
|
||||
3. Unit tests for gate optimizer (threshold bounds, step limits)
|
||||
4. Integration test: full feedback loop cycle (mock QuestDB or embedded)
|
||||
5. Integration test: kill switch triggers freeze + revert
|
||||
6. Integration test: QuestDB failure → graceful degradation
|
||||
7. `SQLX_OFFLINE=true cargo check --workspace` — 0 errors
|
||||
8. Clippy clean on all modified crates
|
||||
199
docs/plans/2026-02-23-production-hardening-phase2-design.md
Normal file
199
docs/plans/2026-02-23-production-hardening-phase2-design.md
Normal file
@@ -0,0 +1,199 @@
|
||||
# Production Hardening Phase 2 — $100K Live Trading Readiness
|
||||
|
||||
> **For Claude:** REQUIRED SUB-SKILL: Use superpowers:executing-plans to implement this plan task-by-task.
|
||||
|
||||
**Goal:** Eliminate all remaining safety gaps, wire real data into placeholder paths, verify correctness with integration tests, and establish training infrastructure — making the system ready for $100K live trading.
|
||||
|
||||
**Architecture:** Risk-prioritized 4-layer approach: Safety Net (crash prevention) → Correctness (accurate calculations) → Verification (test coverage) → Training Infrastructure (GPU strategy + data pipeline). The liquid-cfc-v2 ensemble extension is integrated as part of Layer 2 correctness work.
|
||||
|
||||
**Tech Stack:** Rust (37+ crate workspace), Candle v0.9.1 (ML), Databento (market data), Scaleway (cloud GPU), safetensors (model checkpoints)
|
||||
|
||||
---
|
||||
|
||||
## Audit Context
|
||||
|
||||
### What Was Done (Phase 1)
|
||||
- 53-task production hardening merged to main (+2,952/-1,023 lines across 52 files)
|
||||
- 80+ TODO/FIXME items resolved, all clippy deny rules enforced
|
||||
- 4,234+ tests passing across 30 crate targets, 0 warnings
|
||||
|
||||
### What the Audit Found
|
||||
Four parallel agents audited the codebase and identified:
|
||||
- **92 remaining TODOs** — 4 CRITICAL, 3 HIGH, 8 MEDIUM
|
||||
- **7 `std::process::exit()` calls** in temporal_guard.rs (panic-equivalent)
|
||||
- **28 critical files with 0 tests**, 49 vacuous `assert!(true)` tests, 107 ignored tests
|
||||
- **Model loading has no integrity checks**, ensemble has no per-model circuit breakers
|
||||
- **GPU OOM = hard crash** (no runtime VRAM monitoring)
|
||||
|
||||
---
|
||||
|
||||
## Layer 1 — Safety Net (Crash Prevention)
|
||||
|
||||
**Rationale:** These issues can crash the system or cause unrecoverable failures during live trading. Fix first.
|
||||
|
||||
### 1.1 Replace `std::process::exit()` in temporal_guard.rs
|
||||
- **File:** `ml/src/validation/temporal_guard.rs`
|
||||
- **Problem:** 7 calls to `std::process::exit(1)` — kills the entire process without cleanup
|
||||
- **Fix:** Replace with `Result`-based error propagation using a new `TemporalGuardError` enum
|
||||
- **Risk:** HIGH — process exit during trading = lost positions, no graceful shutdown
|
||||
|
||||
### 1.2 Wire real position data into trading agent
|
||||
- **File:** `services/trading_agent_service/src/service.rs:804-809`
|
||||
- **Problem:** Position data hardcoded to `current_weight: 0.0, current_quantity: 0.0`
|
||||
- **Fix:** Query position manager for real portfolio weights and quantities
|
||||
- **Risk:** CRITICAL — ensemble allocations are meaningless without real position data
|
||||
|
||||
### 1.3 Wire real VaR calculation
|
||||
- **Files:** `services/trading_service/src/services/risk.rs:315,415`
|
||||
- **Problem:** VaR uses placeholder formula `confidence_level * 1_000_000.0`
|
||||
- **Fix:** Wire to the real VaR calculator in `risk/src/var/` which already has parametric, historical, and Monte Carlo methods
|
||||
- **Risk:** CRITICAL — risk limits are not enforced if VaR is fake
|
||||
|
||||
### 1.4 Kill switch Redis monitoring
|
||||
- **File:** `risk/src/safety/kill_switch.rs:378`
|
||||
- **Problem:** Redis monitoring task comment says it should be spawned but never is
|
||||
- **Fix:** Spawn the monitoring task in `start()`, or document that Redis monitoring is deferred and the local kill switch is sufficient for Phase 1
|
||||
- **Risk:** HIGH — distributed kill switch won't propagate across services
|
||||
|
||||
### 1.5 GPU OOM detection and CPU fallback
|
||||
- **Files:** `ml/src/inference/inference.rs`, `ml/src/inference/inference_engine.rs`
|
||||
- **Problem:** GPU out-of-memory = hard crash, no runtime VRAM monitoring
|
||||
- **Fix:** Add VRAM usage check before GPU inference, fall back to CPU with warning log if VRAM > 80% threshold
|
||||
- **Risk:** HIGH — GPU OOM during live trading = system crash
|
||||
|
||||
### 1.6 Per-model circuit breakers in ensemble
|
||||
- **File:** `ml/src/inference/inference_ensemble.rs`
|
||||
- **Problem:** If one model returns garbage (NaN, extreme values), it contaminates the ensemble vote
|
||||
- **Fix:** Add per-model circuit breaker that trips on NaN, repeated identical outputs, or extreme value divergence; ensemble continues with remaining healthy models
|
||||
- **Risk:** HIGH — one bad model can cause the entire ensemble to generate bad trades
|
||||
|
||||
---
|
||||
|
||||
## Layer 2 — Correctness (Accurate Calculations)
|
||||
|
||||
### 2.1 Portfolio correlation matrix for Markowitz allocation
|
||||
- **File:** `services/trading_agent_service/src/allocation.rs`
|
||||
- **Problem:** Uses diagonal covariance (ignores correlations between assets)
|
||||
- **Fix:** Implement rolling correlation matrix from historical returns, use in Markowitz optimization
|
||||
- **Risk:** MEDIUM — suboptimal allocation but not dangerous (diagonal is conservative)
|
||||
|
||||
### 2.2 Feature extraction NaN/Inf guards
|
||||
- **File:** `services/trading_service/src/services/enhanced_ml.rs`
|
||||
- **Problem:** Feature extraction can produce NaN/Inf from division by zero (e.g., zero volume, zero range), which propagates through model inference
|
||||
- **Fix:** Add NaN/Inf check after feature extraction, before normalization. Replace with 0.0 and log warning.
|
||||
- **Risk:** HIGH — NaN in model input = NaN in output = unpredictable trades
|
||||
|
||||
### 2.3 Model file integrity validation
|
||||
- **Problem:** Model files loaded from disk with no integrity verification — corrupted file = silent bad predictions
|
||||
- **Fix:** SHA-256 checksum stored alongside safetensors files, verified on load. Schema validation ensures tensor shapes match expected architecture.
|
||||
- **Risk:** MEDIUM — unlikely but catastrophic if it happens
|
||||
|
||||
### 2.4 Model versioning and rollback
|
||||
- **Problem:** No way to roll back to a previous model version if a new one performs poorly
|
||||
- **Fix:** Model registry with version tracking, A/B comparison metrics, and one-command rollback
|
||||
- **Risk:** MEDIUM — operational risk during model updates
|
||||
|
||||
### 2.5 Liquid CfC v2 ensemble integration
|
||||
- **Branch:** `worktree-liquid-cfc-v2` (10 commits, +15,014/-2,666 lines)
|
||||
- **Status:** LiquidInferenceAdapter and LiquidTrainableAdapter already implemented
|
||||
- **Work needed:**
|
||||
- Merge liquid-cfc-v2 branch to main
|
||||
- Register CfC adapter in EnsembleCoordinator (alongside DQN/PPO/TFT/Mamba2)
|
||||
- Add CfC to hyperopt adapter registry
|
||||
- Verify CfC inference latency is within ensemble budget (<10ms)
|
||||
- Add CfC-specific circuit breaker configuration
|
||||
- **Risk:** LOW — additive change, ensemble already handles N models
|
||||
|
||||
### 2.6 Hyperopt Phase B unblocking
|
||||
- **File:** `ml/src/hyperopt/optimizer.rs`
|
||||
- **Problem:** Hyperopt works for individual models but multi-model orchestration (Phase B) is blocked
|
||||
- **Fix:** Wire ensemble-level hyperopt that optimizes model weights and per-model hyperparameters jointly
|
||||
- **Risk:** MEDIUM — training without hyperopt = suboptimal model configurations
|
||||
|
||||
---
|
||||
|
||||
## Layer 3 — Verification (Test Coverage)
|
||||
|
||||
### 3.1 Execution path integration tests (~53 tests)
|
||||
- **Target files (0-test critical files):**
|
||||
- `services/trading_service/src/services/risk.rs` (VaR, risk limits)
|
||||
- `services/trading_service/src/services/enhanced_ml.rs` (feature extraction, inference)
|
||||
- `services/trading_agent_service/src/service.rs` (allocation, order generation)
|
||||
- `services/trading_agent_service/src/allocation.rs` (Markowitz optimization)
|
||||
- `services/trading_service/src/core/execution_engine.rs` (order execution)
|
||||
- **Approach:** Property-based tests for numerical code, scenario tests for trading logic
|
||||
|
||||
### 3.2 Replace vacuous `assert!(true)` tests
|
||||
- **Count:** 49 tests that just `assert!(true)` or test trivial construction
|
||||
- **Fix:** Replace each with meaningful assertions testing actual behavior
|
||||
- **Priority:** Focus on tests in critical paths first (risk, ML, execution)
|
||||
|
||||
### 3.3 Integration test: ML → Order → Fill
|
||||
- **Scope:** End-to-end test from feature extraction through model inference, ensemble voting, order generation, and simulated fill
|
||||
- **Purpose:** Verify the complete trading pipeline produces valid orders from market data
|
||||
|
||||
### 3.4 Integration test: Risk cascade → Kill switch
|
||||
- **Scope:** Test that risk limit violations properly cascade through circuit breakers to kill switch activation
|
||||
- **Purpose:** Verify safety mechanisms actually trigger under stress conditions
|
||||
|
||||
---
|
||||
|
||||
## Layer 4 — Training Infrastructure
|
||||
|
||||
### 4.1 GPU training strategy
|
||||
- **Local (RTX 3050 Ti, 4GB VRAM):**
|
||||
- Dev/debug training, small batch sizes (max 230 for PPO)
|
||||
- Rapid iteration on model architecture changes
|
||||
- Feature extraction and data preprocessing
|
||||
- **Cloud (Scaleway GPU instances):**
|
||||
- Production training runs with full datasets
|
||||
- Hyperparameter optimization (parallel trials)
|
||||
- Large batch training for all 5 model types
|
||||
- **Deliverables:**
|
||||
- Training launcher script that detects GPU and selects local/cloud path
|
||||
- Scaleway instance provisioning configuration
|
||||
- Model artifact sync between local and cloud storage
|
||||
|
||||
### 4.2 Data pipeline — Databento OHLCV + MBP-10
|
||||
- **Data source:** Databento API for historical market data
|
||||
- **Schemas needed:**
|
||||
- **OHLCV (ohlcv-1m, ohlcv-1h):** Primary training data for all models. ~50MB/symbol/year at 1-min bars
|
||||
- **MBP-10 (market-by-price, 10 levels):** Microstructure features, order book depth. ~2-5GB/symbol/year compressed
|
||||
- **Data quantity research:**
|
||||
- Minimum: 2 years OHLCV for regime diversity (bull, bear, sideways, high-vol)
|
||||
- Recommended: 5 years OHLCV + 1 year MBP-10 for microstructure features
|
||||
- Symbols: Start with 5-10 liquid instruments (ES, NQ, CL, GC, EUR/USD equivalent)
|
||||
- **Storage:** Databento DBN format, already supported by `dbn_sequence_loader.rs` and `dbn_data_source.rs`
|
||||
- **Pipeline:**
|
||||
- Databento API client for historical data download
|
||||
- DBN file management (metadata caching already implemented)
|
||||
- Train/validation/test temporal splits (temporal_guard.rs enforces no leakage)
|
||||
|
||||
### 4.3 Training pipeline safety
|
||||
- **Checkpointing:** Already implemented (safetensors), verify all 5 model types checkpoint correctly
|
||||
- **NaN detection:** Add gradient NaN checks during training, auto-halt with last good checkpoint
|
||||
- **LR scheduling:** Implement cosine annealing with warmup for production training runs
|
||||
- **Metric logging:** Training metrics to structured logs for monitoring dashboards
|
||||
- **Reproducibility:** Seed management for all random operations (data shuffling, model init, exploration)
|
||||
|
||||
---
|
||||
|
||||
## Scope Exclusions
|
||||
|
||||
- **Broker integration:** Separate workstream, not blocked by this work
|
||||
- **FIX protocol production config:** Phase 2 of broker_gateway_service
|
||||
- **Web dashboard enhancements:** Already functional, not in scope
|
||||
- **TLI replacement:** Already completed (deleted in Phase 1)
|
||||
|
||||
## Success Criteria
|
||||
|
||||
1. Zero `std::process::exit()` calls in the codebase
|
||||
2. All placeholder data replaced with real calculations (VaR, positions, features)
|
||||
3. GPU OOM handled gracefully with CPU fallback
|
||||
4. Ensemble survives individual model failures (circuit breakers)
|
||||
5. 5 model types (DQN, PPO, TFT, Mamba2, CfC) all in ensemble
|
||||
6. >90% code coverage on critical trading paths
|
||||
7. End-to-end integration tests pass (ML→Order→Fill, Risk→KillSwitch)
|
||||
8. Training runs complete successfully on both local GPU and Scaleway cloud
|
||||
9. Databento data pipeline downloads, processes, and feeds into training
|
||||
10. Model checksums verified on every load
|
||||
1144
docs/plans/2026-02-23-production-hardening-phase2-implementation.md
Normal file
1144
docs/plans/2026-02-23-production-hardening-phase2-implementation.md
Normal file
File diff suppressed because it is too large
Load Diff
@@ -64,7 +64,7 @@ pub use signer::{CheckpointSigner, SignatureInfo};
|
||||
#[cfg(feature = "s3-storage")]
|
||||
pub use storage::S3CheckpointStorage;
|
||||
pub use storage::{CheckpointStorage, FileSystemStorage, MemoryStorage, StorageStats};
|
||||
pub use validation::ValidationManager;
|
||||
pub use validation::{verify_checksum, write_checksum, ValidationManager};
|
||||
pub use versioning::VersionManager;
|
||||
|
||||
/// Checkpoint format options
|
||||
|
||||
@@ -3,6 +3,7 @@
|
||||
//! Provides checksum validation and corruption detection for checkpoints.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::path::Path;
|
||||
|
||||
use sha2::{Digest, Sha256};
|
||||
use tracing::{debug, error, warn};
|
||||
@@ -10,6 +11,43 @@ use tracing::{debug, error, warn};
|
||||
use super::{CheckpointMetadata, ModelType};
|
||||
use crate::MLError;
|
||||
|
||||
/// Write SHA-256 checksum sidecar file alongside a safetensors checkpoint.
|
||||
/// Creates `{path}.sha256` containing the hex digest.
|
||||
pub fn write_checksum(safetensors_path: &Path) -> Result<(), MLError> {
|
||||
let bytes = std::fs::read(safetensors_path).map_err(|e| {
|
||||
MLError::CheckpointError(format!("Failed to read file for checksum: {}", e))
|
||||
})?;
|
||||
let hash = Sha256::digest(&bytes);
|
||||
let hex = format!("{:x}", hash);
|
||||
let checksum_path = safetensors_path.with_extension("sha256");
|
||||
std::fs::write(&checksum_path, hex.as_bytes()).map_err(|e| {
|
||||
MLError::CheckpointError(format!("Failed to write checksum: {}", e))
|
||||
})?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Verify SHA-256 checksum of a safetensors file against its `.sha256` sidecar.
|
||||
/// Returns `Ok(true)` if valid, `Ok(false)` if mismatch.
|
||||
/// Returns `Ok(true)` if no sidecar exists (backwards compatible).
|
||||
pub fn verify_checksum(safetensors_path: &Path) -> Result<bool, MLError> {
|
||||
let checksum_path = safetensors_path.with_extension("sha256");
|
||||
if !checksum_path.exists() {
|
||||
warn!(
|
||||
"No checksum file for {}, skipping integrity check",
|
||||
safetensors_path.display()
|
||||
);
|
||||
return Ok(true);
|
||||
}
|
||||
let expected = std::fs::read_to_string(&checksum_path).map_err(|e| {
|
||||
MLError::CheckpointError(format!("Failed to read checksum: {}", e))
|
||||
})?;
|
||||
let bytes = std::fs::read(safetensors_path).map_err(|e| {
|
||||
MLError::CheckpointError(format!("Failed to read file for verification: {}", e))
|
||||
})?;
|
||||
let actual = format!("{:x}", Sha256::digest(&bytes));
|
||||
Ok(actual.trim() == expected.trim())
|
||||
}
|
||||
|
||||
/// Validation manager for checkpoint integrity
|
||||
#[derive(Debug)]
|
||||
pub struct ValidationManager {
|
||||
@@ -522,4 +560,44 @@ mod tests {
|
||||
assert!(summary.contains("1 errors"));
|
||||
assert!(summary.contains("1 warnings"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sidecar_checksum_roundtrip() -> Result<(), Box<dyn std::error::Error>> {
|
||||
let dir = tempfile::TempDir::new()?;
|
||||
let path = dir.path().join("test_model.safetensors");
|
||||
std::fs::write(&path, b"fake model data")?;
|
||||
|
||||
write_checksum(&path)?;
|
||||
|
||||
let checksum_path = path.with_extension("sha256");
|
||||
assert!(checksum_path.exists());
|
||||
|
||||
assert!(verify_checksum(&path)?);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sidecar_checksum_detects_corruption() -> Result<(), Box<dyn std::error::Error>> {
|
||||
let dir = tempfile::TempDir::new()?;
|
||||
let path = dir.path().join("test_model.safetensors");
|
||||
std::fs::write(&path, b"original data")?;
|
||||
|
||||
write_checksum(&path)?;
|
||||
|
||||
// Corrupt the file
|
||||
std::fs::write(&path, b"corrupted data")?;
|
||||
assert!(!verify_checksum(&path)?);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_verify_without_sidecar_returns_true() -> Result<(), Box<dyn std::error::Error>> {
|
||||
let dir = tempfile::TempDir::new()?;
|
||||
let path = dir.path().join("test_model.safetensors");
|
||||
std::fs::write(&path, b"no checksum file")?;
|
||||
|
||||
// No .sha256 file exists -- should return true (backwards compatible)
|
||||
assert!(verify_checksum(&path)?);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -76,15 +76,27 @@ impl InferenceEnsemble {
|
||||
let model_name = adapter.model_name().to_string();
|
||||
match adapter.predict(features) {
|
||||
Ok(pred) => {
|
||||
// Circuit breaker: skip NaN/Inf predictions
|
||||
if !pred.direction.is_finite() || !pred.confidence.is_finite() {
|
||||
tracing::warn!(
|
||||
model = %model_name,
|
||||
direction = %pred.direction,
|
||||
confidence = %pred.confidence,
|
||||
"Model returned NaN/Inf prediction, skipping (circuit breaker)"
|
||||
);
|
||||
continue;
|
||||
}
|
||||
// Clamp confidence to valid range
|
||||
let confidence = pred.confidence.clamp(0.0, 1.0);
|
||||
let w = self
|
||||
.weights
|
||||
.get(&model_name)
|
||||
.copied()
|
||||
.unwrap_or(1.0);
|
||||
let wc = w * pred.confidence;
|
||||
let wc = w * confidence;
|
||||
weighted_direction_sum += pred.direction * wc;
|
||||
weight_confidence_sum += wc;
|
||||
confidence_sum += pred.confidence;
|
||||
confidence_sum += confidence;
|
||||
successful_count += 1;
|
||||
model_names.push(model_name);
|
||||
}
|
||||
@@ -294,4 +306,92 @@ mod tests {
|
||||
pred.direction
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_ensemble_filters_nan_predictions() {
|
||||
// NaN model should be skipped; the valid model's prediction stands alone.
|
||||
let adapters: Vec<Box<dyn ModelInferenceAdapter>> = vec![
|
||||
Box::new(DummyAdapter {
|
||||
name: "Valid".to_string(),
|
||||
direction: 0.8,
|
||||
confidence: 0.7,
|
||||
ready: true,
|
||||
}),
|
||||
Box::new(DummyAdapter {
|
||||
name: "NaN_Model".to_string(),
|
||||
direction: f64::NAN,
|
||||
confidence: 0.9,
|
||||
ready: true,
|
||||
}),
|
||||
];
|
||||
|
||||
let ensemble = InferenceEnsemble::new(adapters);
|
||||
let features = make_features();
|
||||
let pred = ensemble.predict(&features).expect("predict should succeed");
|
||||
|
||||
// Only the Valid model should contribute
|
||||
assert!(
|
||||
pred.model_name.contains("Valid"),
|
||||
"model_name should contain 'Valid', got {}",
|
||||
pred.model_name
|
||||
);
|
||||
assert!(
|
||||
!pred.model_name.contains("NaN_Model"),
|
||||
"model_name should NOT contain 'NaN_Model', got {}",
|
||||
pred.model_name
|
||||
);
|
||||
// direction should come entirely from the Valid model
|
||||
assert!(
|
||||
(pred.direction - 0.8).abs() < 1e-9,
|
||||
"direction should be 0.8, got {}",
|
||||
pred.direction
|
||||
);
|
||||
// confidence should be from the single valid model
|
||||
assert!(
|
||||
(pred.confidence - 0.7).abs() < 1e-9,
|
||||
"confidence should be 0.7, got {}",
|
||||
pred.confidence
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_ensemble_filters_extreme_confidence() {
|
||||
// Model with confidence 5.0 should be clamped to 1.0
|
||||
let adapters: Vec<Box<dyn ModelInferenceAdapter>> = vec![
|
||||
Box::new(DummyAdapter {
|
||||
name: "Normal".to_string(),
|
||||
direction: 1.0,
|
||||
confidence: 0.6,
|
||||
ready: true,
|
||||
}),
|
||||
Box::new(DummyAdapter {
|
||||
name: "Extreme".to_string(),
|
||||
direction: 1.0,
|
||||
confidence: 5.0,
|
||||
ready: true,
|
||||
}),
|
||||
];
|
||||
|
||||
let ensemble = InferenceEnsemble::new(adapters);
|
||||
let features = make_features();
|
||||
let pred = ensemble.predict(&features).expect("predict should succeed");
|
||||
|
||||
// Both models should contribute (extreme confidence is clamped, not skipped)
|
||||
assert!(
|
||||
pred.model_name.contains("Normal"),
|
||||
"model_name should contain 'Normal', got {}",
|
||||
pred.model_name
|
||||
);
|
||||
assert!(
|
||||
pred.model_name.contains("Extreme"),
|
||||
"model_name should contain 'Extreme', got {}",
|
||||
pred.model_name
|
||||
);
|
||||
// Confidence should be clamped: avg of 0.6 and 1.0 = 0.8
|
||||
assert!(
|
||||
(pred.confidence - 0.8).abs() < 1e-9,
|
||||
"confidence should be 0.8 (avg of 0.6 and clamped 1.0), got {}",
|
||||
pred.confidence
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -105,3 +105,70 @@ pub fn scale_grads(grads: &mut GradStore, vars: &[Var], scale: f64) -> Result<()
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Check if any gradient in the GradStore contains NaN or Inf.
|
||||
/// Returns `Err` with the variable index if any gradient is non-finite.
|
||||
pub fn check_gradients_finite(grads: &GradStore, vars: &[Var]) -> Result<(), MLError> {
|
||||
for (idx, var) in vars.iter().enumerate() {
|
||||
if let Some(grad) = grads.get(var) {
|
||||
let flat = grad.flatten_all().map_err(|e| {
|
||||
MLError::TrainingError(format!("Failed to flatten gradient {}: {}", idx, e))
|
||||
})?;
|
||||
let values = flat.to_vec1::<f32>().map_err(|e| {
|
||||
MLError::TrainingError(format!("Failed to read gradient {}: {}", idx, e))
|
||||
})?;
|
||||
for val in &values {
|
||||
if !val.is_finite() {
|
||||
return Err(MLError::TrainingError(format!(
|
||||
"NaN/Inf gradient detected in parameter {} (shape: {:?}). \
|
||||
Halting training to prevent model corruption.",
|
||||
idx,
|
||||
grad.shape()
|
||||
)));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use candle_core::Device;
|
||||
|
||||
#[test]
|
||||
fn test_finite_gradients_pass() {
|
||||
let device = Device::Cpu;
|
||||
let var = Var::from_tensor(
|
||||
&candle_core::Tensor::new(&[1.0f32, 2.0, 3.0], &device)
|
||||
.expect("failed to create tensor"),
|
||||
)
|
||||
.expect("failed to create var");
|
||||
let loss = var.mul(&var).expect("mul failed").sum_all().expect("sum failed");
|
||||
let grads = loss.backward().expect("backward failed");
|
||||
let result = check_gradients_finite(&grads, &[var]);
|
||||
assert!(result.is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_nan_gradients_detected() {
|
||||
let device = Device::Cpu;
|
||||
let var = Var::from_tensor(
|
||||
&candle_core::Tensor::new(&[0.0f32], &device).expect("failed to create tensor"),
|
||||
)
|
||||
.expect("failed to create var");
|
||||
let zero = candle_core::Tensor::new(&[0.0f32], &device).expect("failed to create zero");
|
||||
let nan_result = var.div(&zero).expect("div failed");
|
||||
let loss = nan_result.sum_all().expect("sum failed");
|
||||
let grads = loss.backward().expect("backward failed");
|
||||
let result = check_gradients_finite(&grads, &[var]);
|
||||
assert!(result.is_err());
|
||||
let err_msg = format!("{}", result.expect_err("expected error"));
|
||||
assert!(
|
||||
err_msg.contains("NaN") || err_msg.contains("Inf"),
|
||||
"Expected NaN/Inf mention in: {}",
|
||||
err_msg
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,9 +1,11 @@
|
||||
//! Gradient utilities for Candle framework
|
||||
//!
|
||||
//! Provides gradient clipping and other gradient-related operations
|
||||
//! that are missing from candle_nn::optim
|
||||
//! Provides gradient clipping, NaN/Inf detection, and other gradient-related
|
||||
//! operations that are missing from candle_nn::optim
|
||||
|
||||
use candle_core::{Error, backprop::GradStore, Var};
|
||||
use candle_core::{backprop::GradStore, Error, Var};
|
||||
|
||||
use crate::MLError;
|
||||
|
||||
/// Clip gradients by global L2 norm (similar to PyTorch's clip_grad_norm_)
|
||||
///
|
||||
@@ -32,7 +34,11 @@ use candle_core::{Error, backprop::GradStore, Var};
|
||||
/// println!("Gradient norm: {} -> {}", actual_norm, clipped_norm);
|
||||
/// # Ok::<(), candle_core::Error>(())
|
||||
/// ```
|
||||
pub fn clip_grad_norm(vars: &[Var], grads: &mut GradStore, max_norm: f64) -> Result<(f64, f64), Error> {
|
||||
pub fn clip_grad_norm(
|
||||
vars: &[Var],
|
||||
grads: &mut GradStore,
|
||||
max_norm: f64,
|
||||
) -> Result<(f64, f64), Error> {
|
||||
let mut total_norm_sq = 0.0f64;
|
||||
|
||||
// First pass: Calculate the total L2 norm of all gradients
|
||||
@@ -61,3 +67,103 @@ pub fn clip_grad_norm(vars: &[Var], grads: &mut GradStore, max_norm: f64) -> Res
|
||||
Ok((total_norm, total_norm))
|
||||
}
|
||||
}
|
||||
|
||||
/// Check if any gradient in the GradStore contains NaN or Inf values.
|
||||
///
|
||||
/// Returns `Ok(())` if all gradients are finite, or an error naming
|
||||
/// the first parameter with non-finite values. Uses `sum_all` as an
|
||||
/// efficient check: if any element is NaN/Inf, the sum will be non-finite.
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `vars` - Slice of Var containing model parameters
|
||||
/// * `grads` - Reference to GradStore from loss.backward()
|
||||
///
|
||||
/// # Returns
|
||||
/// `Ok(())` if all gradients are finite, `Err(MLError::TrainingError)` otherwise
|
||||
pub fn check_gradients_finite(vars: &[Var], grads: &GradStore) -> Result<(), MLError> {
|
||||
for (idx, var) in vars.iter().enumerate() {
|
||||
if let Some(grad) = grads.get(var) {
|
||||
// Sum all elements — if any element is NaN, the sum will be NaN
|
||||
let sum = grad
|
||||
.sum_all()
|
||||
.and_then(|t| t.to_scalar::<f32>())
|
||||
.map_err(|e| {
|
||||
MLError::TrainingError(format!(
|
||||
"Failed to check gradient for param {}: {}",
|
||||
idx, e
|
||||
))
|
||||
})?;
|
||||
|
||||
if !sum.is_finite() {
|
||||
return Err(MLError::TrainingError(format!(
|
||||
"NaN/Inf gradient detected in parameter index {}. \
|
||||
Training is numerically unstable — halting to prevent model corruption. \
|
||||
Consider reducing learning rate or checking input data for anomalies.",
|
||||
idx
|
||||
)));
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use candle_core::{Device, Tensor};
|
||||
|
||||
#[test]
|
||||
fn test_finite_gradients_pass() {
|
||||
let device = Device::Cpu;
|
||||
let var = Var::from_tensor(
|
||||
&Tensor::new(&[1.0f32, 2.0, 3.0], &device).expect("failed to create tensor"),
|
||||
)
|
||||
.expect("failed to create var");
|
||||
|
||||
let loss = var
|
||||
.mul(&var)
|
||||
.expect("mul failed")
|
||||
.sum_all()
|
||||
.expect("sum failed");
|
||||
let grads = loss.backward().expect("backward failed");
|
||||
|
||||
let result = check_gradients_finite(&[var], &grads);
|
||||
assert!(result.is_ok(), "Finite gradients should pass check");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_nan_gradient_detected() {
|
||||
let device = Device::Cpu;
|
||||
let var = Var::from_tensor(
|
||||
&Tensor::new(&[0.0f32], &device).expect("failed to create tensor"),
|
||||
)
|
||||
.expect("failed to create var");
|
||||
|
||||
// 0/0 produces NaN
|
||||
let zero = Tensor::new(&[0.0f32], &device).expect("failed to create zero");
|
||||
let nan_result = var.div(&zero).expect("div failed");
|
||||
let loss = nan_result.sum_all().expect("sum failed");
|
||||
let grads = loss.backward().expect("backward failed");
|
||||
|
||||
let result = check_gradients_finite(&[var], &grads);
|
||||
assert!(result.is_err(), "NaN/Inf gradients should be detected");
|
||||
let err_msg = format!("{}", result.expect_err("expected error"));
|
||||
assert!(
|
||||
err_msg.contains("NaN/Inf gradient detected"),
|
||||
"Error should mention NaN/Inf: {}",
|
||||
err_msg
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_empty_vars_passes() {
|
||||
let device = Device::Cpu;
|
||||
// Create an empty GradStore by computing backward on a constant
|
||||
let loss = Tensor::new(1.0f32, &device).expect("failed to create tensor");
|
||||
let grads = loss.backward().expect("backward failed");
|
||||
|
||||
let vars: Vec<Var> = vec![];
|
||||
let result = check_gradients_finite(&vars, &grads);
|
||||
assert!(result.is_ok(), "Empty vars should pass");
|
||||
}
|
||||
}
|
||||
|
||||
515
ml/src/hyperopt/adapters/ensemble.rs
Normal file
515
ml/src/hyperopt/adapters/ensemble.rs
Normal file
@@ -0,0 +1,515 @@
|
||||
//! Ensemble-level hyperopt adapter
|
||||
//!
|
||||
//! Adds N ensemble weight parameters (one per model) to a combined parameter space,
|
||||
//! enabling joint optimization of model weights alongside per-model hyperparameters.
|
||||
//!
|
||||
//! ## Design
|
||||
//!
|
||||
//! The [`ParameterSpace`] trait uses static methods (`continuous_bounds()`, `param_names()`),
|
||||
//! which means the dimensionality must be known at compile time. To support variable-size
|
||||
//! ensemble configurations at runtime, this module provides:
|
||||
//!
|
||||
//! - [`EnsembleSpaceConfig`]: Runtime configuration specifying model names and per-model
|
||||
//! parameter dimensions. Stored in a thread-local so that static trait methods can
|
||||
//! access it.
|
||||
//! - [`EnsembleParameterSpace`]: The parameter space struct implementing [`ParameterSpace`].
|
||||
//! Holds a flat continuous vector (per-model params concatenated + N weight values).
|
||||
//!
|
||||
//! ## Weight Normalization
|
||||
//!
|
||||
//! Raw weight values in `[0, 1]` are normalized via softmax-style division so they
|
||||
//! sum to 1.0. If all raw weights are zero, equal weights are assigned.
|
||||
//!
|
||||
//! ## Usage
|
||||
//!
|
||||
//! ```rust,no_run
|
||||
//! use ml::hyperopt::adapters::ensemble::{EnsembleSpaceConfig, EnsembleParameterSpace};
|
||||
//! use ml::hyperopt::ParameterSpace;
|
||||
//!
|
||||
//! // Configure for a 3-model ensemble (DQN=11D, PPO=5D, TFT=6D)
|
||||
//! let config = EnsembleSpaceConfig::new(
|
||||
//! vec!["dqn".into(), "ppo".into(), "tft".into()],
|
||||
//! vec![11, 5, 6],
|
||||
//! // Per-model bounds: 11 DQN bounds + 5 PPO bounds + 6 TFT bounds
|
||||
//! vec![
|
||||
//! // ... 22 bounds total from individual model ParameterSpaces
|
||||
//! ],
|
||||
//! // Per-model param names
|
||||
//! vec![
|
||||
//! // ... 22 names total
|
||||
//! ],
|
||||
//! );
|
||||
//!
|
||||
//! // Install config so ParameterSpace static methods can read it
|
||||
//! config.install();
|
||||
//!
|
||||
//! // Now EnsembleParameterSpace implements ParameterSpace
|
||||
//! let bounds = EnsembleParameterSpace::continuous_bounds();
|
||||
//! // bounds.len() == 22 (model params) + 3 (weights) = 25
|
||||
//! ```
|
||||
|
||||
use std::cell::RefCell;
|
||||
|
||||
use crate::hyperopt::traits::ParameterSpace;
|
||||
use crate::MLError;
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Runtime configuration
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Runtime configuration for the ensemble parameter space.
|
||||
///
|
||||
/// Must be [`install()`](EnsembleSpaceConfig::install)ed before calling
|
||||
/// `EnsembleParameterSpace` trait methods.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct EnsembleSpaceConfig {
|
||||
/// Human-readable model names (e.g., "dqn", "ppo", "tft").
|
||||
pub model_names: Vec<String>,
|
||||
/// Number of continuous parameters per model.
|
||||
pub model_param_dims: Vec<usize>,
|
||||
/// Concatenated per-model bounds `[(min, max), ...]`.
|
||||
/// Length must equal `model_param_dims.iter().sum()`.
|
||||
pub model_bounds: Vec<(f64, f64)>,
|
||||
/// Concatenated per-model parameter names.
|
||||
/// Length must equal `model_param_dims.iter().sum()`.
|
||||
pub model_param_names: Vec<String>,
|
||||
}
|
||||
|
||||
thread_local! {
|
||||
static ENSEMBLE_CONFIG: RefCell<Option<EnsembleSpaceConfig>> = const { RefCell::new(None) };
|
||||
}
|
||||
|
||||
impl EnsembleSpaceConfig {
|
||||
/// Create a new ensemble space configuration.
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `model_names` - Names of models in the ensemble.
|
||||
/// * `model_param_dims` - Number of hyperparameters per model.
|
||||
/// * `model_bounds` - Concatenated bounds for all model parameters.
|
||||
/// * `model_param_names` - Concatenated parameter names for all models.
|
||||
///
|
||||
/// # Panics
|
||||
///
|
||||
/// Panics if lengths are inconsistent.
|
||||
pub fn new(
|
||||
model_names: Vec<String>,
|
||||
model_param_dims: Vec<usize>,
|
||||
model_bounds: Vec<(f64, f64)>,
|
||||
model_param_names: Vec<String>,
|
||||
) -> Self {
|
||||
let total_model_params: usize = model_param_dims.iter().sum();
|
||||
assert_eq!(
|
||||
model_names.len(),
|
||||
model_param_dims.len(),
|
||||
"model_names and model_param_dims must have the same length"
|
||||
);
|
||||
assert_eq!(
|
||||
model_bounds.len(),
|
||||
total_model_params,
|
||||
"model_bounds length must equal sum of model_param_dims"
|
||||
);
|
||||
assert_eq!(
|
||||
model_param_names.len(),
|
||||
total_model_params,
|
||||
"model_param_names length must equal sum of model_param_dims"
|
||||
);
|
||||
Self {
|
||||
model_names,
|
||||
model_param_dims,
|
||||
model_bounds,
|
||||
model_param_names,
|
||||
}
|
||||
}
|
||||
|
||||
/// Install this configuration in the thread-local slot.
|
||||
///
|
||||
/// Must be called before using `EnsembleParameterSpace` trait methods.
|
||||
pub fn install(&self) {
|
||||
ENSEMBLE_CONFIG.with(|cell| {
|
||||
*cell.borrow_mut() = Some(self.clone());
|
||||
});
|
||||
}
|
||||
|
||||
/// Remove the installed configuration.
|
||||
pub fn uninstall() {
|
||||
ENSEMBLE_CONFIG.with(|cell| {
|
||||
*cell.borrow_mut() = None;
|
||||
});
|
||||
}
|
||||
|
||||
/// Total dimension = sum(model_param_dims) + num_models (weight params).
|
||||
pub fn total_dim(&self) -> usize {
|
||||
let model_params: usize = self.model_param_dims.iter().sum();
|
||||
model_params + self.model_names.len()
|
||||
}
|
||||
|
||||
/// Number of models in the ensemble.
|
||||
pub fn num_models(&self) -> usize {
|
||||
self.model_names.len()
|
||||
}
|
||||
}
|
||||
|
||||
/// Read the installed config, returning an error if none is installed.
|
||||
fn with_config<T>(f: impl FnOnce(&EnsembleSpaceConfig) -> T) -> Result<T, MLError> {
|
||||
ENSEMBLE_CONFIG.with(|cell| {
|
||||
let borrow = cell.borrow();
|
||||
match borrow.as_ref() {
|
||||
Some(cfg) => Ok(f(cfg)),
|
||||
None => Err(MLError::ConfigError {
|
||||
reason: "EnsembleSpaceConfig not installed. Call config.install() first."
|
||||
.to_string(),
|
||||
}),
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Parameter space
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Combined parameter space for ensemble optimization.
|
||||
///
|
||||
/// Holds a flat vector of continuous values:
|
||||
/// `[model_0_params..., model_1_params..., ..., weight_0, weight_1, ...]`
|
||||
///
|
||||
/// The last N values (where N = number of models) are raw ensemble weights
|
||||
/// in `[0, 1]`. Use [`extract_weights()`](Self::extract_weights) to get
|
||||
/// normalized weights that sum to 1.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct EnsembleParameterSpace {
|
||||
/// Flat continuous parameter vector.
|
||||
pub values: Vec<f64>,
|
||||
}
|
||||
|
||||
impl EnsembleParameterSpace {
|
||||
/// Extract normalized ensemble weights from the parameter vector.
|
||||
///
|
||||
/// Weights are the last N values, normalized to sum to 1.
|
||||
/// If all raw weights are zero (or negative after clamping), returns equal weights.
|
||||
pub fn extract_weights(&self, num_models: usize) -> Vec<f64> {
|
||||
let weight_start = self.values.len().saturating_sub(num_models);
|
||||
let raw_weights: Vec<f64> = (0..num_models)
|
||||
.filter_map(|i| self.values.get(weight_start + i).copied())
|
||||
.map(|w| w.max(0.0))
|
||||
.collect();
|
||||
let sum: f64 = raw_weights.iter().sum();
|
||||
if sum > f64::EPSILON {
|
||||
raw_weights.iter().map(|w| w / sum).collect()
|
||||
} else {
|
||||
// Fallback: equal weights
|
||||
let equal = 1.0 / num_models.max(1) as f64;
|
||||
vec![equal; num_models]
|
||||
}
|
||||
}
|
||||
|
||||
/// Extract per-model parameter slices from the combined vector.
|
||||
///
|
||||
/// Returns a vector of slices, one per model. The slices are in the same
|
||||
/// order as the model names in the config.
|
||||
pub fn extract_model_params(&self, model_param_dims: &[usize]) -> Vec<Vec<f64>> {
|
||||
let mut result = Vec::with_capacity(model_param_dims.len());
|
||||
let mut offset = 0;
|
||||
for &dim in model_param_dims {
|
||||
let end = (offset + dim).min(self.values.len());
|
||||
let start = offset.min(end);
|
||||
result.push(self.values.get(start..end).unwrap_or_default().to_vec());
|
||||
offset += dim;
|
||||
}
|
||||
result
|
||||
}
|
||||
}
|
||||
|
||||
impl ParameterSpace for EnsembleParameterSpace {
|
||||
fn continuous_bounds() -> Vec<(f64, f64)> {
|
||||
// Read from thread-local config; if missing, return empty
|
||||
// (caller should have installed config first)
|
||||
with_config(|cfg| {
|
||||
let mut bounds = cfg.model_bounds.clone();
|
||||
// Append weight bounds: [0, 1] for each model
|
||||
for _ in 0..cfg.num_models() {
|
||||
bounds.push((0.0, 1.0));
|
||||
}
|
||||
bounds
|
||||
})
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
fn from_continuous(x: &[f64]) -> Result<Self, MLError> {
|
||||
let expected_dim = with_config(|cfg| cfg.total_dim())?;
|
||||
if x.len() != expected_dim {
|
||||
return Err(MLError::ConfigError {
|
||||
reason: format!(
|
||||
"EnsembleParameterSpace: expected {} params, got {}",
|
||||
expected_dim,
|
||||
x.len()
|
||||
),
|
||||
});
|
||||
}
|
||||
Ok(Self {
|
||||
values: x.to_vec(),
|
||||
})
|
||||
}
|
||||
|
||||
fn to_continuous(&self) -> Vec<f64> {
|
||||
self.values.clone()
|
||||
}
|
||||
|
||||
fn param_names() -> Vec<&'static str> {
|
||||
// Build names from config. Since the trait requires &'static str,
|
||||
// we leak the strings. This is acceptable because hyperopt configs
|
||||
// are created once per optimization run (not in a hot loop).
|
||||
with_config(|cfg| {
|
||||
let mut names: Vec<&'static str> = cfg
|
||||
.model_param_names
|
||||
.iter()
|
||||
.map(|s| -> &'static str { Box::leak(s.clone().into_boxed_str()) })
|
||||
.collect();
|
||||
for model_name in &cfg.model_names {
|
||||
let weight_name = format!("weight_{}", model_name);
|
||||
names.push(Box::leak(weight_name.into_boxed_str()));
|
||||
}
|
||||
names
|
||||
})
|
||||
.unwrap_or_default()
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Tests
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
/// Helper: create a simple 2-model config for testing.
|
||||
fn test_config() -> EnsembleSpaceConfig {
|
||||
EnsembleSpaceConfig::new(
|
||||
vec!["model_a".into(), "model_b".into()],
|
||||
vec![2, 3], // model_a has 2 params, model_b has 3 params
|
||||
vec![
|
||||
(0.0, 1.0),
|
||||
(0.0, 10.0),
|
||||
(-1.0, 1.0),
|
||||
(-1.0, 1.0),
|
||||
(0.0, 100.0),
|
||||
],
|
||||
vec![
|
||||
"a_lr".into(),
|
||||
"a_batch".into(),
|
||||
"b_x".into(),
|
||||
"b_y".into(),
|
||||
"b_z".into(),
|
||||
],
|
||||
)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_ensemble_parameter_space_dimensions() {
|
||||
let config = test_config();
|
||||
config.install();
|
||||
|
||||
// 2 + 3 model params + 2 weights = 7 total
|
||||
assert_eq!(config.total_dim(), 7);
|
||||
|
||||
let bounds = EnsembleParameterSpace::continuous_bounds();
|
||||
assert_eq!(bounds.len(), 7);
|
||||
|
||||
// First 5 bounds are model params
|
||||
assert_eq!(bounds.first().copied(), Some((0.0, 1.0))); // a_lr
|
||||
assert_eq!(bounds.get(4).copied(), Some((0.0, 100.0))); // b_z
|
||||
|
||||
// Last 2 bounds are weight params [0, 1]
|
||||
assert_eq!(bounds.get(5).copied(), Some((0.0, 1.0))); // weight_model_a
|
||||
assert_eq!(bounds.get(6).copied(), Some((0.0, 1.0))); // weight_model_b
|
||||
|
||||
let names = EnsembleParameterSpace::param_names();
|
||||
assert_eq!(names.len(), 7);
|
||||
assert_eq!(names.first().copied(), Some("a_lr"));
|
||||
assert_eq!(names.get(4).copied(), Some("b_z"));
|
||||
assert_eq!(names.get(5).copied(), Some("weight_model_a"));
|
||||
assert_eq!(names.get(6).copied(), Some("weight_model_b"));
|
||||
|
||||
EnsembleSpaceConfig::uninstall();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_weights_normalizes() {
|
||||
let config = test_config();
|
||||
config.install();
|
||||
|
||||
// 7 params: [a_lr, a_batch, b_x, b_y, b_z, weight_a, weight_b]
|
||||
let params = EnsembleParameterSpace::from_continuous(&[
|
||||
0.5, 5.0, 0.0, 0.0, 50.0, // model params
|
||||
0.3, 0.7, // raw weights
|
||||
])
|
||||
.unwrap_or_else(|e| panic!("from_continuous failed: {}", e));
|
||||
|
||||
let weights = params.extract_weights(2);
|
||||
assert_eq!(weights.len(), 2);
|
||||
|
||||
// 0.3 / (0.3 + 0.7) = 0.3
|
||||
assert!((weights.first().copied().unwrap_or(0.0) - 0.3).abs() < 1e-10);
|
||||
// 0.7 / (0.3 + 0.7) = 0.7
|
||||
assert!((weights.get(1).copied().unwrap_or(0.0) - 0.7).abs() < 1e-10);
|
||||
|
||||
// Verify weights sum to 1
|
||||
let sum: f64 = weights.iter().sum();
|
||||
assert!(
|
||||
(sum - 1.0).abs() < 1e-10,
|
||||
"Weights should sum to 1.0, got {}",
|
||||
sum
|
||||
);
|
||||
|
||||
EnsembleSpaceConfig::uninstall();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_weights_handles_zeros() {
|
||||
let config = test_config();
|
||||
config.install();
|
||||
|
||||
let params = EnsembleParameterSpace::from_continuous(&[
|
||||
0.5, 5.0, 0.0, 0.0, 50.0, // model params
|
||||
0.0, 0.0, // all-zero weights
|
||||
])
|
||||
.unwrap_or_else(|e| panic!("from_continuous failed: {}", e));
|
||||
|
||||
let weights = params.extract_weights(2);
|
||||
assert_eq!(weights.len(), 2);
|
||||
|
||||
// Should fall back to equal weights: 0.5, 0.5
|
||||
assert!(
|
||||
(weights.first().copied().unwrap_or(0.0) - 0.5).abs() < 1e-10,
|
||||
"Expected 0.5 for zero-weight fallback, got {:?}",
|
||||
weights.first()
|
||||
);
|
||||
assert!(
|
||||
(weights.get(1).copied().unwrap_or(0.0) - 0.5).abs() < 1e-10,
|
||||
"Expected 0.5 for zero-weight fallback, got {:?}",
|
||||
weights.get(1)
|
||||
);
|
||||
|
||||
EnsembleSpaceConfig::uninstall();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_roundtrip_continuous() {
|
||||
let config = test_config();
|
||||
config.install();
|
||||
|
||||
let original = vec![0.5, 5.0, 0.0, -0.5, 50.0, 0.3, 0.7];
|
||||
let params = EnsembleParameterSpace::from_continuous(&original)
|
||||
.unwrap_or_else(|e| panic!("from_continuous failed: {}", e));
|
||||
let recovered = params.to_continuous();
|
||||
|
||||
assert_eq!(original.len(), recovered.len());
|
||||
for (a, b) in original.iter().zip(recovered.iter()) {
|
||||
assert!(
|
||||
(a - b).abs() < 1e-10,
|
||||
"Round-trip mismatch: {} vs {}",
|
||||
a,
|
||||
b
|
||||
);
|
||||
}
|
||||
|
||||
EnsembleSpaceConfig::uninstall();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_wrong_dimension_returns_error() {
|
||||
let config = test_config();
|
||||
config.install();
|
||||
|
||||
let result = EnsembleParameterSpace::from_continuous(&[1.0, 2.0]);
|
||||
assert!(result.is_err(), "Should reject wrong dimension");
|
||||
|
||||
EnsembleSpaceConfig::uninstall();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_model_params() {
|
||||
let config = test_config();
|
||||
config.install();
|
||||
|
||||
let params = EnsembleParameterSpace::from_continuous(&[
|
||||
0.5, 5.0, // model_a params (dim=2)
|
||||
0.1, -0.3, 75.0, // model_b params (dim=3)
|
||||
0.4, 0.6, // weights
|
||||
])
|
||||
.unwrap_or_else(|e| panic!("from_continuous failed: {}", e));
|
||||
|
||||
let model_params = params.extract_model_params(&config.model_param_dims);
|
||||
assert_eq!(model_params.len(), 2);
|
||||
|
||||
// model_a: [0.5, 5.0]
|
||||
assert_eq!(model_params.first().map(|v| v.len()), Some(2));
|
||||
assert!(
|
||||
(model_params
|
||||
.first()
|
||||
.and_then(|v| v.first().copied())
|
||||
.unwrap_or(0.0)
|
||||
- 0.5)
|
||||
.abs()
|
||||
< 1e-10
|
||||
);
|
||||
|
||||
// model_b: [0.1, -0.3, 75.0]
|
||||
assert_eq!(model_params.get(1).map(|v| v.len()), Some(3));
|
||||
assert!(
|
||||
(model_params
|
||||
.get(1)
|
||||
.and_then(|v| v.get(2).copied())
|
||||
.unwrap_or(0.0)
|
||||
- 75.0)
|
||||
.abs()
|
||||
< 1e-10
|
||||
);
|
||||
|
||||
EnsembleSpaceConfig::uninstall();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_no_config_installed_returns_empty_or_error() {
|
||||
// Ensure no config is installed
|
||||
EnsembleSpaceConfig::uninstall();
|
||||
|
||||
// continuous_bounds returns empty when no config
|
||||
let bounds = EnsembleParameterSpace::continuous_bounds();
|
||||
assert!(bounds.is_empty());
|
||||
|
||||
// from_continuous returns error when no config
|
||||
let result = EnsembleParameterSpace::from_continuous(&[1.0]);
|
||||
assert!(result.is_err());
|
||||
|
||||
// param_names returns empty when no config
|
||||
let names = EnsembleParameterSpace::param_names();
|
||||
assert!(names.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_single_model_ensemble() {
|
||||
let config = EnsembleSpaceConfig::new(
|
||||
vec!["only_model".into()],
|
||||
vec![1],
|
||||
vec![(0.0, 1.0)],
|
||||
vec!["lr".into()],
|
||||
);
|
||||
config.install();
|
||||
|
||||
assert_eq!(config.total_dim(), 2); // 1 param + 1 weight
|
||||
|
||||
let params = EnsembleParameterSpace::from_continuous(&[0.5, 0.8])
|
||||
.unwrap_or_else(|e| panic!("from_continuous failed: {}", e));
|
||||
let weights = params.extract_weights(1);
|
||||
assert_eq!(weights.len(), 1);
|
||||
// Single model: weight normalizes to 1.0
|
||||
assert!(
|
||||
(weights.first().copied().unwrap_or(0.0) - 1.0).abs() < 1e-10,
|
||||
"Single model weight should be 1.0"
|
||||
);
|
||||
|
||||
EnsembleSpaceConfig::uninstall();
|
||||
}
|
||||
}
|
||||
@@ -52,6 +52,7 @@
|
||||
pub mod async_data_loader;
|
||||
pub mod continuous_ppo;
|
||||
pub mod dqn;
|
||||
pub mod ensemble;
|
||||
pub mod kan;
|
||||
pub mod liquid;
|
||||
pub mod mamba2;
|
||||
@@ -68,6 +69,7 @@ pub use kan::{KANMetrics, KANParams};
|
||||
pub use liquid::LiquidParams;
|
||||
pub use continuous_ppo::{ContinuousPPOMetrics, ContinuousPPOParams, ContinuousPPOTrainer};
|
||||
pub use dqn::{DQNMetrics, DQNParams, DQNTrainer};
|
||||
pub use ensemble::{EnsembleParameterSpace, EnsembleSpaceConfig};
|
||||
pub use mamba2::{Mamba2Metrics, Mamba2Params, Mamba2Trainer};
|
||||
pub use ppo::{PPOMetrics, PPOParams, PPOTrainer};
|
||||
pub use tft::{TFTMetrics, TFTParams, TFTTrainer as TFTHyperoptTrainer};
|
||||
|
||||
@@ -544,17 +544,37 @@ impl RealMLInferenceEngine {
|
||||
let device = match self.config.device_preference.as_str() {
|
||||
"cuda" | "gpu" => match Device::new_cuda(0) {
|
||||
Ok(cuda_device) => {
|
||||
info!("✅ Using CUDA device for model: {}", model_id);
|
||||
cuda_device
|
||||
},
|
||||
match crate::memory_optimization::auto_batch_size::detect_gpu_memory() {
|
||||
Ok((_, free_mb, _)) if free_mb > 500.0 => {
|
||||
info!(
|
||||
"Using CUDA device for model: {} (free VRAM: {:.0}MB)",
|
||||
model_id, free_mb
|
||||
);
|
||||
cuda_device
|
||||
}
|
||||
Ok((_, free_mb, _)) => {
|
||||
warn!(
|
||||
"GPU VRAM too low ({:.0}MB free), falling back to CPU for model: {}",
|
||||
free_mb, model_id
|
||||
);
|
||||
Device::Cpu
|
||||
}
|
||||
Err(e) => {
|
||||
warn!(
|
||||
"Cannot detect GPU memory ({}), falling back to CPU for model: {}",
|
||||
e, model_id
|
||||
);
|
||||
Device::Cpu
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
return Err(MLSafetyError::from(RealInferenceError::GpuRequired {
|
||||
reason: format!(
|
||||
"GPU acceleration required for production model {}: {}",
|
||||
model_id, e
|
||||
),
|
||||
}));
|
||||
},
|
||||
warn!(
|
||||
"CUDA not available ({}), falling back to CPU for model: {}",
|
||||
e, model_id
|
||||
);
|
||||
Device::Cpu
|
||||
}
|
||||
},
|
||||
_ => {
|
||||
info!("Using CPU device for model: {}", model_id);
|
||||
|
||||
@@ -697,6 +697,310 @@ pub struct RegistryStatistics {
|
||||
pub earliest_training_date: Option<DateTime<Utc>>,
|
||||
}
|
||||
|
||||
/// A rollback event recorded in the audit log
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct RollbackEvent {
|
||||
/// Model name that was rolled back
|
||||
pub model_name: String,
|
||||
/// Version that was active before the rollback
|
||||
pub from_version: String,
|
||||
/// Version that became active after the rollback
|
||||
pub to_version: String,
|
||||
/// Timestamp of the rollback
|
||||
pub timestamp: DateTime<Utc>,
|
||||
/// Optional reason for the rollback
|
||||
pub reason: Option<String>,
|
||||
}
|
||||
|
||||
/// Version entry stored per model in the in-memory registry
|
||||
#[derive(Debug, Clone)]
|
||||
struct VersionEntry {
|
||||
/// All registered versions (ordered by registration time)
|
||||
versions: Vec<ModelVersionMetadata>,
|
||||
/// Index into `versions` for the currently active version
|
||||
active_index: usize,
|
||||
}
|
||||
|
||||
/// In-memory model registry with version tracking and rollback support.
|
||||
///
|
||||
/// This registry does not require a database. It stores all model versions
|
||||
/// in memory, tracks which version is "active" for each model name, and
|
||||
/// records rollback events in an audit log.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct InMemoryModelRegistry {
|
||||
/// Per-model version entries keyed by model name (e.g. "dqn", "ppo")
|
||||
models: Arc<RwLock<HashMap<String, VersionEntry>>>,
|
||||
/// Audit log of rollback events
|
||||
rollback_log: Arc<RwLock<Vec<RollbackEvent>>>,
|
||||
}
|
||||
|
||||
impl InMemoryModelRegistry {
|
||||
/// Create a new empty in-memory model registry.
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
models: Arc::new(RwLock::new(HashMap::new())),
|
||||
rollback_log: Arc::new(RwLock::new(Vec::new())),
|
||||
}
|
||||
}
|
||||
|
||||
/// Register a new model version.
|
||||
///
|
||||
/// The first version registered for a given model name automatically becomes
|
||||
/// the active version. Subsequent registrations are stored but do not change
|
||||
/// the active version (use [`rollback_to_version`] or [`promote_version`] for that).
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `model_name` - Logical model name (e.g. "dqn", "ppo")
|
||||
/// * `metadata` - Full version metadata
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns `MLError::ModelError` if a version with the same version string
|
||||
/// is already registered for this model name.
|
||||
pub async fn register_version(
|
||||
&self,
|
||||
model_name: &str,
|
||||
metadata: ModelVersionMetadata,
|
||||
) -> MLResult<()> {
|
||||
let mut models = self.models.write().await;
|
||||
let entry = models
|
||||
.entry(model_name.to_string())
|
||||
.or_insert_with(|| VersionEntry {
|
||||
versions: Vec::new(),
|
||||
active_index: 0,
|
||||
});
|
||||
|
||||
// Check for duplicate version strings
|
||||
let version_str = metadata.version.clone();
|
||||
for existing in &entry.versions {
|
||||
if existing.version == version_str {
|
||||
return Err(MLError::ModelError(format!(
|
||||
"Version {} already registered for model {}",
|
||||
version_str, model_name
|
||||
)));
|
||||
}
|
||||
}
|
||||
|
||||
entry.versions.push(metadata);
|
||||
|
||||
// First version auto-becomes active (active_index is already 0)
|
||||
// Subsequent versions do not change active_index
|
||||
|
||||
tracing::info!(
|
||||
model_name = model_name,
|
||||
version = version_str.as_str(),
|
||||
total_versions = entry.versions.len(),
|
||||
"Registered model version"
|
||||
);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Roll back a model to a previously registered version.
|
||||
///
|
||||
/// This sets the specified version as the active version and records
|
||||
/// a rollback event in the audit log.
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `model_name` - Logical model name
|
||||
/// * `target_version` - Semantic version string to roll back to
|
||||
/// * `reason` - Optional reason for the rollback
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns `MLError::ModelNotFound` if the model name or target version
|
||||
/// is not found in the registry.
|
||||
pub async fn rollback_to_version(
|
||||
&self,
|
||||
model_name: &str,
|
||||
target_version: &str,
|
||||
reason: Option<String>,
|
||||
) -> MLResult<()> {
|
||||
let mut models = self.models.write().await;
|
||||
let entry = models.get_mut(model_name).ok_or_else(|| {
|
||||
MLError::ModelNotFound(format!("Model {} not found in registry", model_name))
|
||||
})?;
|
||||
|
||||
// Find the target version index
|
||||
let target_index = entry
|
||||
.versions
|
||||
.iter()
|
||||
.position(|v| v.version == target_version)
|
||||
.ok_or_else(|| {
|
||||
MLError::ModelNotFound(format!(
|
||||
"Version {} not found for model {}",
|
||||
target_version, model_name
|
||||
))
|
||||
})?;
|
||||
|
||||
let from_version = entry
|
||||
.versions
|
||||
.get(entry.active_index)
|
||||
.map(|v| v.version.clone())
|
||||
.unwrap_or_default();
|
||||
|
||||
if entry.active_index == target_index {
|
||||
tracing::warn!(
|
||||
model_name = model_name,
|
||||
version = target_version,
|
||||
"Rollback requested to already-active version (no-op)"
|
||||
);
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
entry.active_index = target_index;
|
||||
|
||||
tracing::info!(
|
||||
model_name = model_name,
|
||||
from_version = from_version.as_str(),
|
||||
to_version = target_version,
|
||||
reason = reason.as_deref().unwrap_or("none"),
|
||||
"Rolled back model version"
|
||||
);
|
||||
|
||||
// Record rollback event
|
||||
let event = RollbackEvent {
|
||||
model_name: model_name.to_string(),
|
||||
from_version,
|
||||
to_version: target_version.to_string(),
|
||||
timestamp: Utc::now(),
|
||||
reason,
|
||||
};
|
||||
// Drop models lock before acquiring rollback_log lock to avoid deadlock
|
||||
drop(models);
|
||||
self.rollback_log.write().await.push(event);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Promote a version to be the active version (same as rollback but with
|
||||
/// clearer semantics for forward version changes).
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns `MLError::ModelNotFound` if the model or version is not found.
|
||||
pub async fn promote_version(
|
||||
&self,
|
||||
model_name: &str,
|
||||
target_version: &str,
|
||||
) -> MLResult<()> {
|
||||
self.rollback_to_version(model_name, target_version, Some("promoted".to_string()))
|
||||
.await
|
||||
}
|
||||
|
||||
/// List all registered versions for a model, ordered by registration time.
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `model_name` - Logical model name
|
||||
///
|
||||
/// # Returns
|
||||
///
|
||||
/// A vector of `(version_string, is_active)` tuples.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns `MLError::ModelNotFound` if the model name is not found.
|
||||
pub async fn list_versions(
|
||||
&self,
|
||||
model_name: &str,
|
||||
) -> MLResult<Vec<(String, bool)>> {
|
||||
let models = self.models.read().await;
|
||||
let entry = models.get(model_name).ok_or_else(|| {
|
||||
MLError::ModelNotFound(format!("Model {} not found in registry", model_name))
|
||||
})?;
|
||||
|
||||
let result = entry
|
||||
.versions
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, v)| (v.version.clone(), i == entry.active_index))
|
||||
.collect();
|
||||
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
/// Get the active version metadata for a model.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns `MLError::ModelNotFound` if the model is not found or has no versions.
|
||||
pub async fn get_active_version(
|
||||
&self,
|
||||
model_name: &str,
|
||||
) -> MLResult<ModelVersionMetadata> {
|
||||
let models = self.models.read().await;
|
||||
let entry = models.get(model_name).ok_or_else(|| {
|
||||
MLError::ModelNotFound(format!("Model {} not found in registry", model_name))
|
||||
})?;
|
||||
|
||||
entry
|
||||
.versions
|
||||
.get(entry.active_index)
|
||||
.cloned()
|
||||
.ok_or_else(|| {
|
||||
MLError::ModelNotFound(format!("No versions registered for model {}", model_name))
|
||||
})
|
||||
}
|
||||
|
||||
/// Get a specific version's metadata for a model.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns `MLError::ModelNotFound` if the model or version is not found.
|
||||
pub async fn get_version(
|
||||
&self,
|
||||
model_name: &str,
|
||||
version: &str,
|
||||
) -> MLResult<ModelVersionMetadata> {
|
||||
let models = self.models.read().await;
|
||||
let entry = models.get(model_name).ok_or_else(|| {
|
||||
MLError::ModelNotFound(format!("Model {} not found in registry", model_name))
|
||||
})?;
|
||||
|
||||
entry
|
||||
.versions
|
||||
.iter()
|
||||
.find(|v| v.version == version)
|
||||
.cloned()
|
||||
.ok_or_else(|| {
|
||||
MLError::ModelNotFound(format!(
|
||||
"Version {} not found for model {}",
|
||||
version, model_name
|
||||
))
|
||||
})
|
||||
}
|
||||
|
||||
/// Get all model names in the registry.
|
||||
pub async fn list_models(&self) -> Vec<String> {
|
||||
self.models.read().await.keys().cloned().collect()
|
||||
}
|
||||
|
||||
/// Get the rollback audit log.
|
||||
pub async fn get_rollback_log(&self) -> Vec<RollbackEvent> {
|
||||
self.rollback_log.read().await.clone()
|
||||
}
|
||||
|
||||
/// Get rollback events for a specific model.
|
||||
pub async fn get_rollback_log_for_model(&self, model_name: &str) -> Vec<RollbackEvent> {
|
||||
self.rollback_log
|
||||
.read()
|
||||
.await
|
||||
.iter()
|
||||
.filter(|e| e.model_name == model_name)
|
||||
.cloned()
|
||||
.collect()
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for InMemoryModelRegistry {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
@@ -747,4 +1051,300 @@ mod tests {
|
||||
assert_eq!(retrieved.model_id, "dqn-test-v1.0.0");
|
||||
assert_eq!(retrieved.version, "1.0.0");
|
||||
}
|
||||
|
||||
// ---- In-memory registry tests (no database required) ----
|
||||
|
||||
fn make_version(model_id: &str, version: &str) -> ModelVersionMetadata {
|
||||
ModelVersionMetadata::new(
|
||||
model_id.to_string(),
|
||||
ModelType::DQN,
|
||||
version.to_string(),
|
||||
"test_data".to_string(),
|
||||
format!("s3://models/{}/{}/", model_id, version),
|
||||
)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_inmemory_register_and_get_active() {
|
||||
let registry = InMemoryModelRegistry::new();
|
||||
|
||||
let v1 = make_version("dqn-v1", "1.0.0");
|
||||
registry.register_version("dqn", v1).await.unwrap();
|
||||
|
||||
let active = registry.get_active_version("dqn").await.unwrap();
|
||||
assert_eq!(active.version, "1.0.0");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_inmemory_first_version_is_active() {
|
||||
let registry = InMemoryModelRegistry::new();
|
||||
|
||||
let v1 = make_version("dqn-v1", "1.0.0");
|
||||
let v2 = make_version("dqn-v2", "2.0.0");
|
||||
|
||||
registry.register_version("dqn", v1).await.unwrap();
|
||||
registry.register_version("dqn", v2).await.unwrap();
|
||||
|
||||
// First registered version stays active
|
||||
let active = registry.get_active_version("dqn").await.unwrap();
|
||||
assert_eq!(active.version, "1.0.0");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_inmemory_duplicate_version_rejected() {
|
||||
let registry = InMemoryModelRegistry::new();
|
||||
|
||||
let v1 = make_version("dqn-v1", "1.0.0");
|
||||
let v1_dup = make_version("dqn-v1-dup", "1.0.0");
|
||||
|
||||
registry.register_version("dqn", v1).await.unwrap();
|
||||
let result = registry.register_version("dqn", v1_dup).await;
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_inmemory_list_versions() {
|
||||
let registry = InMemoryModelRegistry::new();
|
||||
|
||||
let v1 = make_version("dqn-v1", "1.0.0");
|
||||
let v2 = make_version("dqn-v2", "2.0.0");
|
||||
let v3 = make_version("dqn-v3", "3.0.0");
|
||||
|
||||
registry.register_version("dqn", v1).await.unwrap();
|
||||
registry.register_version("dqn", v2).await.unwrap();
|
||||
registry.register_version("dqn", v3).await.unwrap();
|
||||
|
||||
let versions = registry.list_versions("dqn").await.unwrap();
|
||||
assert_eq!(versions.len(), 3);
|
||||
assert_eq!(versions.first().map(|v| v.0.as_str()), Some("1.0.0"));
|
||||
assert_eq!(versions.first().map(|v| v.1), Some(true)); // active
|
||||
assert_eq!(versions.get(1).map(|v| v.0.as_str()), Some("2.0.0"));
|
||||
assert_eq!(versions.get(1).map(|v| v.1), Some(false)); // not active
|
||||
assert_eq!(versions.get(2).map(|v| v.0.as_str()), Some("3.0.0"));
|
||||
assert_eq!(versions.get(2).map(|v| v.1), Some(false)); // not active
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_inmemory_list_versions_unknown_model() {
|
||||
let registry = InMemoryModelRegistry::new();
|
||||
let result = registry.list_versions("nonexistent").await;
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_inmemory_rollback_to_version() {
|
||||
let registry = InMemoryModelRegistry::new();
|
||||
|
||||
let v1 = make_version("dqn-v1", "1.0.0");
|
||||
let v2 = make_version("dqn-v2", "2.0.0");
|
||||
let v3 = make_version("dqn-v3", "3.0.0");
|
||||
|
||||
registry.register_version("dqn", v1).await.unwrap();
|
||||
registry.register_version("dqn", v2).await.unwrap();
|
||||
registry.register_version("dqn", v3).await.unwrap();
|
||||
|
||||
// Promote to v3 first
|
||||
registry.promote_version("dqn", "3.0.0").await.unwrap();
|
||||
let active = registry.get_active_version("dqn").await.unwrap();
|
||||
assert_eq!(active.version, "3.0.0");
|
||||
|
||||
// Roll back to v1
|
||||
registry
|
||||
.rollback_to_version("dqn", "1.0.0", Some("regression in v3".to_string()))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let active = registry.get_active_version("dqn").await.unwrap();
|
||||
assert_eq!(active.version, "1.0.0");
|
||||
|
||||
// Check rollback log
|
||||
let log = registry.get_rollback_log().await;
|
||||
// Two events: promote to v3 and rollback to v1
|
||||
assert_eq!(log.len(), 2);
|
||||
|
||||
let last = log.get(1);
|
||||
assert!(last.is_some());
|
||||
if let Some(event) = last {
|
||||
assert_eq!(event.model_name, "dqn");
|
||||
assert_eq!(event.from_version, "3.0.0");
|
||||
assert_eq!(event.to_version, "1.0.0");
|
||||
assert_eq!(event.reason.as_deref(), Some("regression in v3"));
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_inmemory_rollback_unknown_model() {
|
||||
let registry = InMemoryModelRegistry::new();
|
||||
let result = registry
|
||||
.rollback_to_version("nonexistent", "1.0.0", None)
|
||||
.await;
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_inmemory_rollback_unknown_version() {
|
||||
let registry = InMemoryModelRegistry::new();
|
||||
|
||||
let v1 = make_version("dqn-v1", "1.0.0");
|
||||
registry.register_version("dqn", v1).await.unwrap();
|
||||
|
||||
let result = registry
|
||||
.rollback_to_version("dqn", "99.0.0", None)
|
||||
.await;
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_inmemory_rollback_to_same_version_is_noop() {
|
||||
let registry = InMemoryModelRegistry::new();
|
||||
|
||||
let v1 = make_version("dqn-v1", "1.0.0");
|
||||
registry.register_version("dqn", v1).await.unwrap();
|
||||
|
||||
// Rolling back to already-active version succeeds silently
|
||||
registry
|
||||
.rollback_to_version("dqn", "1.0.0", None)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// No rollback event recorded for no-op
|
||||
let log = registry.get_rollback_log().await;
|
||||
assert!(log.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_inmemory_get_version() {
|
||||
let registry = InMemoryModelRegistry::new();
|
||||
|
||||
let v1 = make_version("dqn-v1", "1.0.0");
|
||||
let v2 = make_version("dqn-v2", "2.0.0");
|
||||
|
||||
registry.register_version("dqn", v1).await.unwrap();
|
||||
registry.register_version("dqn", v2).await.unwrap();
|
||||
|
||||
let retrieved = registry.get_version("dqn", "2.0.0").await.unwrap();
|
||||
assert_eq!(retrieved.version, "2.0.0");
|
||||
assert_eq!(retrieved.model_id, "dqn-v2");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_inmemory_get_version_not_found() {
|
||||
let registry = InMemoryModelRegistry::new();
|
||||
|
||||
let v1 = make_version("dqn-v1", "1.0.0");
|
||||
registry.register_version("dqn", v1).await.unwrap();
|
||||
|
||||
let result = registry.get_version("dqn", "99.0.0").await;
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_inmemory_list_models() {
|
||||
let registry = InMemoryModelRegistry::new();
|
||||
|
||||
let dqn = make_version("dqn-v1", "1.0.0");
|
||||
let ppo = make_version("ppo-v1", "1.0.0");
|
||||
|
||||
registry.register_version("dqn", dqn).await.unwrap();
|
||||
registry.register_version("ppo", ppo).await.unwrap();
|
||||
|
||||
let mut models = registry.list_models().await;
|
||||
models.sort();
|
||||
assert_eq!(models, vec!["dqn", "ppo"]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_inmemory_rollback_log_per_model() {
|
||||
let registry = InMemoryModelRegistry::new();
|
||||
|
||||
let dqn_v1 = make_version("dqn-v1", "1.0.0");
|
||||
let dqn_v2 = make_version("dqn-v2", "2.0.0");
|
||||
let ppo_v1 = make_version("ppo-v1", "1.0.0");
|
||||
let ppo_v2 = make_version("ppo-v2", "2.0.0");
|
||||
|
||||
registry.register_version("dqn", dqn_v1).await.unwrap();
|
||||
registry.register_version("dqn", dqn_v2).await.unwrap();
|
||||
registry.register_version("ppo", ppo_v1).await.unwrap();
|
||||
registry.register_version("ppo", ppo_v2).await.unwrap();
|
||||
|
||||
// Roll back both
|
||||
registry
|
||||
.rollback_to_version("dqn", "2.0.0", None)
|
||||
.await
|
||||
.unwrap();
|
||||
registry
|
||||
.rollback_to_version("ppo", "2.0.0", None)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Filter by model
|
||||
let dqn_log = registry.get_rollback_log_for_model("dqn").await;
|
||||
assert_eq!(dqn_log.len(), 1);
|
||||
assert_eq!(dqn_log.first().map(|e| e.model_name.as_str()), Some("dqn"));
|
||||
|
||||
let ppo_log = registry.get_rollback_log_for_model("ppo").await;
|
||||
assert_eq!(ppo_log.len(), 1);
|
||||
assert_eq!(ppo_log.first().map(|e| e.model_name.as_str()), Some("ppo"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_inmemory_default_trait() {
|
||||
let registry = InMemoryModelRegistry::default();
|
||||
let models = registry.list_models().await;
|
||||
assert!(models.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_inmemory_multiple_rollbacks() {
|
||||
let registry = InMemoryModelRegistry::new();
|
||||
|
||||
let v1 = make_version("dqn-v1", "1.0.0");
|
||||
let v2 = make_version("dqn-v2", "2.0.0");
|
||||
let v3 = make_version("dqn-v3", "3.0.0");
|
||||
|
||||
registry.register_version("dqn", v1).await.unwrap();
|
||||
registry.register_version("dqn", v2).await.unwrap();
|
||||
registry.register_version("dqn", v3).await.unwrap();
|
||||
|
||||
// v1 -> v3 -> v2 -> v1 -> v3
|
||||
registry.promote_version("dqn", "3.0.0").await.unwrap();
|
||||
registry
|
||||
.rollback_to_version("dqn", "2.0.0", None)
|
||||
.await
|
||||
.unwrap();
|
||||
registry
|
||||
.rollback_to_version("dqn", "1.0.0", None)
|
||||
.await
|
||||
.unwrap();
|
||||
registry.promote_version("dqn", "3.0.0").await.unwrap();
|
||||
|
||||
let active = registry.get_active_version("dqn").await.unwrap();
|
||||
assert_eq!(active.version, "3.0.0");
|
||||
|
||||
let log = registry.get_rollback_log().await;
|
||||
assert_eq!(log.len(), 4);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_inmemory_version_list_reflects_active_after_rollback() {
|
||||
let registry = InMemoryModelRegistry::new();
|
||||
|
||||
let v1 = make_version("dqn-v1", "1.0.0");
|
||||
let v2 = make_version("dqn-v2", "2.0.0");
|
||||
|
||||
registry.register_version("dqn", v1).await.unwrap();
|
||||
registry.register_version("dqn", v2).await.unwrap();
|
||||
|
||||
// Initially v1 is active
|
||||
let versions = registry.list_versions("dqn").await.unwrap();
|
||||
assert_eq!(versions.first().map(|v| v.1), Some(true));
|
||||
assert_eq!(versions.get(1).map(|v| v.1), Some(false));
|
||||
|
||||
// Promote v2
|
||||
registry.promote_version("dqn", "2.0.0").await.unwrap();
|
||||
|
||||
let versions = registry.list_versions("dqn").await.unwrap();
|
||||
assert_eq!(versions.first().map(|v| v.1), Some(false));
|
||||
assert_eq!(versions.get(1).map(|v| v.1), Some(true));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -20,7 +20,7 @@ use serde::{Deserialize, Serialize};
|
||||
use std::path::PathBuf;
|
||||
use tracing::{debug, info, warn};
|
||||
|
||||
use crate::gradient_accumulation::{accumulate_grads, scale_grads};
|
||||
use crate::gradient_accumulation::{accumulate_grads, check_gradients_finite, scale_grads};
|
||||
use crate::tensor_ops::TensorOps;
|
||||
|
||||
use super::gae::GAEConfig;
|
||||
@@ -978,6 +978,8 @@ impl PPO {
|
||||
&actor_vars,
|
||||
1.0 / accumulation_steps as f64,
|
||||
)?;
|
||||
check_gradients_finite(policy_grads, &actor_vars)
|
||||
.map_err(|e| MLError::TrainingError(format!("Policy gradient NaN: {}", e)))?;
|
||||
|
||||
let policy_grad_norm =
|
||||
Self::compute_gradient_norm(&actor_vars, policy_grads)?;
|
||||
@@ -1008,6 +1010,8 @@ impl PPO {
|
||||
&critic_vars,
|
||||
1.0 / accumulation_steps as f64,
|
||||
)?;
|
||||
check_gradients_finite(value_grads, &critic_vars)
|
||||
.map_err(|e| MLError::TrainingError(format!("Value gradient NaN: {}", e)))?;
|
||||
|
||||
let value_grad_norm =
|
||||
Self::compute_gradient_norm(&critic_vars, value_grads)?;
|
||||
@@ -1049,6 +1053,8 @@ impl PPO {
|
||||
|
||||
if let Some(ref mut policy_grads) = policy_grad_accumulator {
|
||||
scale_grads(policy_grads, &actor_vars, 1.0 / accum_step as f64)?;
|
||||
check_gradients_finite(policy_grads, &actor_vars)
|
||||
.map_err(|e| MLError::TrainingError(format!("Policy gradient NaN: {}", e)))?;
|
||||
|
||||
let policy_grad_norm =
|
||||
Self::compute_gradient_norm(&actor_vars, policy_grads)?;
|
||||
@@ -1077,6 +1083,8 @@ impl PPO {
|
||||
|
||||
if let Some(ref mut value_grads) = value_grad_accumulator {
|
||||
scale_grads(value_grads, &critic_vars, 1.0 / accum_step as f64)?;
|
||||
check_gradients_finite(value_grads, &critic_vars)
|
||||
.map_err(|e| MLError::TrainingError(format!("Value gradient NaN: {}", e)))?;
|
||||
|
||||
let value_grad_norm =
|
||||
Self::compute_gradient_norm(&critic_vars, value_grads)?;
|
||||
@@ -1333,6 +1341,8 @@ impl PPO {
|
||||
let policy_grads = policy_loss.backward().map_err(|e| {
|
||||
MLError::TrainingError(format!("Policy backward failed: {}", e))
|
||||
})?;
|
||||
check_gradients_finite(&policy_grads, &actor_vars)
|
||||
.map_err(|e| MLError::TrainingError(format!("Policy gradient NaN: {}", e)))?;
|
||||
|
||||
let policy_grad_norm =
|
||||
Self::compute_gradient_norm(&actor_vars, &policy_grads)?;
|
||||
@@ -1357,6 +1367,8 @@ impl PPO {
|
||||
let value_grads = scaled_value_loss.backward().map_err(|e| {
|
||||
MLError::TrainingError(format!("Value backward failed: {}", e))
|
||||
})?;
|
||||
check_gradients_finite(&value_grads, &critic_vars)
|
||||
.map_err(|e| MLError::TrainingError(format!("Value gradient NaN: {}", e)))?;
|
||||
|
||||
let value_grad_norm =
|
||||
Self::compute_gradient_norm(&critic_vars, &value_grads)?;
|
||||
|
||||
@@ -2875,6 +2875,9 @@ impl DQNTrainer {
|
||||
)
|
||||
.map_err(|e| anyhow::anyhow!("Gradient scaling failed: {}", e))?;
|
||||
|
||||
crate::gradient_accumulation::check_gradients_finite(grads, &vars)
|
||||
.map_err(|e| anyhow::anyhow!("Training halted: {}", e))?;
|
||||
|
||||
agent
|
||||
.apply_accumulated_gradients(grads)
|
||||
.map_err(|e| anyhow::anyhow!("Apply accumulated gradients failed: {}", e))?;
|
||||
|
||||
@@ -278,35 +278,16 @@ mod tests {
|
||||
}
|
||||
|
||||
/// Helper: create a simple ascending-price time series.
|
||||
fn make_test_data(n: usize) -> TimeSeriesData {
|
||||
fn make_test_data(n: usize) -> Result<TimeSeriesData, Box<dyn std::error::Error>> {
|
||||
let prices: Vec<f64> = (0..n).map(|i| 100.0 + i as f64).collect();
|
||||
TimeSeriesData::new(make_timestamps(n), make_features(n, 3), prices)
|
||||
.unwrap_or_else(|e| {
|
||||
// We cannot panic due to clippy deny, but this is test code.
|
||||
// Use a fallback that will never actually be reached.
|
||||
eprintln!("Test data creation failed: {e}");
|
||||
// Return minimal valid data
|
||||
TimeSeriesData::new(
|
||||
make_timestamps(2),
|
||||
make_features(2, 3),
|
||||
vec![100.0, 101.0],
|
||||
)
|
||||
.unwrap_or_else(|_| std::process::exit(1))
|
||||
})
|
||||
Ok(TimeSeriesData::new(make_timestamps(n), make_features(n, 3), prices)?)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_training_slice_returns_correct_range() {
|
||||
let data = make_test_data(10);
|
||||
let guard =
|
||||
TemporalGuard::new(&data, 5).unwrap_or_else(|e| {
|
||||
eprintln!("Guard creation failed: {e}");
|
||||
std::process::exit(1)
|
||||
});
|
||||
let train = guard.training_slice().unwrap_or_else(|e| {
|
||||
eprintln!("training_slice failed: {e}");
|
||||
std::process::exit(1)
|
||||
});
|
||||
fn test_training_slice_returns_correct_range() -> Result<(), Box<dyn std::error::Error>> {
|
||||
let data = make_test_data(10)?;
|
||||
let guard = TemporalGuard::new(&data, 5)?;
|
||||
let train = guard.training_slice()?;
|
||||
|
||||
assert_eq!(train.len(), 5);
|
||||
// First price should be 100.0, last should be 104.0
|
||||
@@ -314,16 +295,13 @@ mod tests {
|
||||
let last = train.prices.last().copied().unwrap_or(0.0);
|
||||
assert!((first - 100.0).abs() < 1e-12);
|
||||
assert!((last - 104.0).abs() < 1e-12);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_test_slice_rejects_before_cutoff() {
|
||||
let data = make_test_data(10);
|
||||
let guard =
|
||||
TemporalGuard::new(&data, 5).unwrap_or_else(|e| {
|
||||
eprintln!("Guard creation failed: {e}");
|
||||
std::process::exit(1)
|
||||
});
|
||||
fn test_test_slice_rejects_before_cutoff() -> Result<(), Box<dyn std::error::Error>> {
|
||||
let data = make_test_data(10)?;
|
||||
let guard = TemporalGuard::new(&data, 5)?;
|
||||
|
||||
// start=3 is before cutoff=5 → should fail
|
||||
let result = guard.test_slice(3, 8);
|
||||
@@ -333,16 +311,13 @@ mod tests {
|
||||
err_msg.contains("leak"),
|
||||
"Expected 'leak' in error, got: {err_msg}",
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_slice_rejects_cross_boundary() {
|
||||
let data = make_test_data(10);
|
||||
let guard =
|
||||
TemporalGuard::new(&data, 5).unwrap_or_else(|e| {
|
||||
eprintln!("Guard creation failed: {e}");
|
||||
std::process::exit(1)
|
||||
});
|
||||
fn test_slice_rejects_cross_boundary() -> Result<(), Box<dyn std::error::Error>> {
|
||||
let data = make_test_data(10)?;
|
||||
let guard = TemporalGuard::new(&data, 5)?;
|
||||
|
||||
// [3, 7) crosses cutoff=5
|
||||
let result = guard.slice(3, 7);
|
||||
@@ -360,31 +335,25 @@ mod tests {
|
||||
// [5, 8) is entirely in test → should succeed
|
||||
let test_ok = guard.slice(5, 8);
|
||||
assert!(test_ok.is_ok());
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_audit_detects_no_leakage_in_sorted_data() {
|
||||
let data = make_test_data(10);
|
||||
let guard =
|
||||
TemporalGuard::new(&data, 5).unwrap_or_else(|e| {
|
||||
eprintln!("Guard creation failed: {e}");
|
||||
std::process::exit(1)
|
||||
});
|
||||
fn test_audit_detects_no_leakage_in_sorted_data() -> Result<(), Box<dyn std::error::Error>> {
|
||||
let data = make_test_data(10)?;
|
||||
let guard = TemporalGuard::new(&data, 5)?;
|
||||
|
||||
let report = guard.audit_leakage();
|
||||
assert!(!report.has_future_timestamps);
|
||||
assert_eq!(report.training_bars, 5);
|
||||
assert!(report.cutoff_timestamp.is_some());
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_normalization_stats_from_training_only() {
|
||||
let data = make_test_data(10);
|
||||
let guard =
|
||||
TemporalGuard::new(&data, 5).unwrap_or_else(|e| {
|
||||
eprintln!("Guard creation failed: {e}");
|
||||
std::process::exit(1)
|
||||
});
|
||||
fn test_normalization_stats_from_training_only() -> Result<(), Box<dyn std::error::Error>> {
|
||||
let data = make_test_data(10)?;
|
||||
let guard = TemporalGuard::new(&data, 5)?;
|
||||
|
||||
let stats = guard.compute_normalization_stats();
|
||||
assert_eq!(stats.sample_count, 5);
|
||||
@@ -399,5 +368,6 @@ mod tests {
|
||||
(got_mean_0 - expected_mean_0).abs() < 1e-10,
|
||||
"Expected mean_0={expected_mean_0}, got {got_mean_0}",
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -375,8 +375,24 @@ impl AtomicKillSwitch {
|
||||
"Kill switch monitoring started"
|
||||
);
|
||||
|
||||
// TODO: spawn a tokio background task that periodically checks Redis
|
||||
// connectivity and logs health status (interval from config.health_check_interval).
|
||||
// ARCHITECTURAL DECISION: Redis health monitoring is deferred to multi-service deployment.
|
||||
//
|
||||
// Rationale:
|
||||
// - The local AtomicBool kill switch provides immediate, zero-latency process-level
|
||||
// protection that works even when Redis is unreachable. This is the correct
|
||||
// fail-safe for a single-process HFT system.
|
||||
// - Background Redis monitoring adds complexity and a tokio task that is not
|
||||
// useful until multiple services need coordinated kill switch state.
|
||||
//
|
||||
// When multi-service coordination is needed, implement:
|
||||
// 1. Periodic Redis PING (interval from config.health_check_interval)
|
||||
// 2. On disconnect: log warning, increment failure_count, continue operating
|
||||
// with local AtomicBool state (fail-safe: local state is authoritative)
|
||||
// 3. On reconnect: re-sync local state to Redis, publish current status
|
||||
// 4. Expose health via is_healthy() (already implemented above)
|
||||
//
|
||||
// See also: unix_socket_kill_switch.rs for an alternative IPC-based kill
|
||||
// switch that avoids the Redis dependency entirely for same-host deployments.
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
226
scripts/train_launcher.sh
Executable file
226
scripts/train_launcher.sh
Executable file
@@ -0,0 +1,226 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
# GPU Training Launcher -- detects GPU and selects local or cloud path
|
||||
#
|
||||
# Usage:
|
||||
# ./scripts/train_launcher.sh --model dqn [--data path/to/data] [--cloud] [--epochs N]
|
||||
# ./scripts/train_launcher.sh --model ppo --data data/databento/ES.FUT/ohlcv-1m/ --epochs 50
|
||||
# ./scripts/train_launcher.sh --model tft --cloud
|
||||
#
|
||||
# Models: dqn, ppo, tft, mamba2, cfc
|
||||
# Flags:
|
||||
# --cloud Force cloud training even if local GPU is available
|
||||
# --epochs Number of training epochs (default: 20)
|
||||
# --data Path to training data (parquet file or directory)
|
||||
# --lr Learning rate (default: 0.0001)
|
||||
# --output Output directory for trained models (default: ml/trained_models)
|
||||
|
||||
MODEL=""
|
||||
DATA_PATH=""
|
||||
CLOUD=false
|
||||
EPOCHS=20
|
||||
LEARNING_RATE=0.0001
|
||||
OUTPUT_DIR="ml/trained_models"
|
||||
EXTRA_ARGS=()
|
||||
|
||||
while [[ $# -gt 0 ]]; do
|
||||
case "$1" in
|
||||
--model) MODEL="$2"; shift 2 ;;
|
||||
--data) DATA_PATH="$2"; shift 2 ;;
|
||||
--cloud) CLOUD=true; shift ;;
|
||||
--epochs) EPOCHS="$2"; shift 2 ;;
|
||||
--lr) LEARNING_RATE="$2"; shift 2 ;;
|
||||
--output) OUTPUT_DIR="$2"; shift 2 ;;
|
||||
--help|-h) usage; exit 0 ;;
|
||||
*) EXTRA_ARGS+=("$1"); shift ;;
|
||||
esac
|
||||
done
|
||||
|
||||
usage() {
|
||||
cat <<'USAGE'
|
||||
GPU Training Launcher -- detects GPU and routes to local or cloud
|
||||
|
||||
Usage:
|
||||
./scripts/train_launcher.sh --model <model> [OPTIONS]
|
||||
|
||||
Models: dqn, ppo, tft, mamba2, cfc
|
||||
|
||||
Options:
|
||||
--model <name> Model to train (required)
|
||||
--data <path> Training data parquet file or directory
|
||||
--epochs <N> Number of training epochs (default: 20)
|
||||
--lr <rate> Learning rate (default: 0.0001)
|
||||
--output <dir> Output directory (default: ml/trained_models)
|
||||
--cloud Force cloud training
|
||||
--help Show this help
|
||||
|
||||
Examples:
|
||||
./scripts/train_launcher.sh --model dqn --epochs 100
|
||||
./scripts/train_launcher.sh --model ppo --data test_data/ES_FUT_180d.parquet --epochs 50
|
||||
./scripts/train_launcher.sh --model tft --cloud
|
||||
USAGE
|
||||
}
|
||||
|
||||
if [[ -z "$MODEL" ]]; then
|
||||
echo "Error: --model is required (dqn, ppo, tft, mamba2, cfc)"
|
||||
echo "Run with --help for usage information"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
VALID_MODELS="dqn ppo tft mamba2 cfc"
|
||||
if ! echo "$VALID_MODELS" | grep -qw "$MODEL"; then
|
||||
echo "Error: Invalid model '$MODEL'. Valid: $VALID_MODELS"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# Map model names to cargo example targets
|
||||
# These are the actual ml/examples/train_*.rs binaries
|
||||
declare -A EXAMPLE_MAP
|
||||
EXAMPLE_MAP[dqn]="train_dqn"
|
||||
EXAMPLE_MAP[ppo]="train_ppo_parquet"
|
||||
EXAMPLE_MAP[tft]="train_tft_parquet"
|
||||
EXAMPLE_MAP[mamba2]="train_mamba2_parquet"
|
||||
EXAMPLE_MAP[cfc]="train_liquid_dbn"
|
||||
|
||||
EXAMPLE_NAME="${EXAMPLE_MAP[$MODEL]}"
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# GPU Detection
|
||||
# ---------------------------------------------------------------------------
|
||||
GPU_DETECTED=false
|
||||
VRAM_MB=0
|
||||
GPU_NAME="none"
|
||||
|
||||
if command -v nvidia-smi &>/dev/null; then
|
||||
VRAM_MB=$(nvidia-smi --query-gpu=memory.total --format=csv,noheader,nounits 2>/dev/null | head -1 || echo "0")
|
||||
if [[ "$VRAM_MB" -gt 0 ]]; then
|
||||
GPU_DETECTED=true
|
||||
GPU_NAME=$(nvidia-smi --query-gpu=name --format=csv,noheader 2>/dev/null | head -1 || echo "unknown")
|
||||
echo "Local GPU detected: $GPU_NAME (${VRAM_MB}MB VRAM)"
|
||||
fi
|
||||
else
|
||||
echo "No NVIDIA GPU detected (nvidia-smi not found)"
|
||||
fi
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Batch size selection based on VRAM and model
|
||||
# ---------------------------------------------------------------------------
|
||||
# These values are empirically tested; PPO is constrained on 4GB (max 230).
|
||||
# See scripts/measure_vram.sh for VRAM profiling methodology.
|
||||
declare -A BS_TIER_LOW # < 4 GB
|
||||
declare -A BS_TIER_MED # 4-8 GB (e.g. RTX 3050 Ti 4GB)
|
||||
declare -A BS_TIER_HIGH # 8-16 GB (e.g. RTX 3070 8GB)
|
||||
declare -A BS_TIER_MAX # 16+ GB (e.g. L4 24GB)
|
||||
|
||||
BS_TIER_LOW[dqn]=32; BS_TIER_MED[dqn]=128; BS_TIER_HIGH[dqn]=256; BS_TIER_MAX[dqn]=512
|
||||
BS_TIER_LOW[ppo]=32; BS_TIER_MED[ppo]=230; BS_TIER_HIGH[ppo]=512; BS_TIER_MAX[ppo]=1024
|
||||
BS_TIER_LOW[tft]=8; BS_TIER_MED[tft]=32; BS_TIER_HIGH[tft]=64; BS_TIER_MAX[tft]=256
|
||||
BS_TIER_LOW[mamba2]=8; BS_TIER_MED[mamba2]=64; BS_TIER_HIGH[mamba2]=128; BS_TIER_MAX[mamba2]=512
|
||||
BS_TIER_LOW[cfc]=32; BS_TIER_MED[cfc]=128; BS_TIER_HIGH[cfc]=256; BS_TIER_MAX[cfc]=512
|
||||
|
||||
BATCH_SIZE=64
|
||||
if [[ "$GPU_DETECTED" == true ]]; then
|
||||
if [[ "$VRAM_MB" -lt 4096 ]]; then
|
||||
BATCH_SIZE="${BS_TIER_LOW[$MODEL]}"
|
||||
elif [[ "$VRAM_MB" -lt 8192 ]]; then
|
||||
BATCH_SIZE="${BS_TIER_MED[$MODEL]}"
|
||||
elif [[ "$VRAM_MB" -lt 16384 ]]; then
|
||||
BATCH_SIZE="${BS_TIER_HIGH[$MODEL]}"
|
||||
else
|
||||
BATCH_SIZE="${BS_TIER_MAX[$MODEL]}"
|
||||
fi
|
||||
fi
|
||||
|
||||
echo ""
|
||||
echo "Configuration:"
|
||||
echo " Model: $MODEL (example: $EXAMPLE_NAME)"
|
||||
echo " Epochs: $EPOCHS"
|
||||
echo " Batch size: $BATCH_SIZE"
|
||||
echo " Learning rate: $LEARNING_RATE"
|
||||
echo " Output: $OUTPUT_DIR"
|
||||
if [[ -n "$DATA_PATH" ]]; then
|
||||
echo " Data: $DATA_PATH"
|
||||
fi
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Build the cargo command arguments
|
||||
# ---------------------------------------------------------------------------
|
||||
build_cargo_args() {
|
||||
local args=()
|
||||
args+=(--epochs "$EPOCHS")
|
||||
args+=(--batch-size "$BATCH_SIZE")
|
||||
args+=(--learning-rate "$LEARNING_RATE")
|
||||
args+=(--output-dir "$OUTPUT_DIR")
|
||||
|
||||
if [[ -n "$DATA_PATH" ]]; then
|
||||
args+=(--parquet-file "$DATA_PATH")
|
||||
fi
|
||||
|
||||
# Pass through any extra arguments
|
||||
if [[ ${#EXTRA_ARGS[@]} -gt 0 ]]; then
|
||||
args+=("${EXTRA_ARGS[@]}")
|
||||
fi
|
||||
|
||||
echo "${args[@]}"
|
||||
}
|
||||
|
||||
CARGO_EXTRA=$(build_cargo_args)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Route to local or cloud
|
||||
# ---------------------------------------------------------------------------
|
||||
CLOUD_HOST="${FOXHUNT_GPU_HOST:-gpu.fxhnt.ai}"
|
||||
CLOUD_DIR="${FOXHUNT_CLOUD_DIR:-/opt/foxhunt}"
|
||||
|
||||
need_cloud() {
|
||||
# Cloud is needed when: forced, no GPU, or insufficient VRAM (< 4GB)
|
||||
[[ "$CLOUD" == true ]] || [[ "$GPU_DETECTED" == false ]] || [[ "$VRAM_MB" -lt 4096 ]]
|
||||
}
|
||||
|
||||
if need_cloud; then
|
||||
echo ""
|
||||
echo "==> Routing to cloud training at $CLOUD_HOST"
|
||||
echo ""
|
||||
|
||||
# Build the remote command
|
||||
REMOTE_CMD="cd $CLOUD_DIR && git pull --ff-only && SQLX_OFFLINE=true cargo run --release -p ml --example $EXAMPLE_NAME --features cuda -- $CARGO_EXTRA"
|
||||
|
||||
echo "To run on cloud GPU:"
|
||||
echo ""
|
||||
echo " ssh training@$CLOUD_HOST '$REMOTE_CMD'"
|
||||
echo ""
|
||||
echo "To sync model artifacts back:"
|
||||
echo ""
|
||||
echo " rsync -avz training@$CLOUD_HOST:$CLOUD_DIR/checkpoints/ ./checkpoints/"
|
||||
echo " rsync -avz training@$CLOUD_HOST:$CLOUD_DIR/ml/trained_models/ ./ml/trained_models/"
|
||||
echo ""
|
||||
|
||||
# If we have SSH access, offer to run it directly
|
||||
if [[ "$CLOUD" == true ]]; then
|
||||
read -rp "Execute on cloud now? [y/N] " CONFIRM
|
||||
if [[ "$CONFIRM" =~ ^[Yy]$ ]]; then
|
||||
echo "Connecting to $CLOUD_HOST..."
|
||||
# shellcheck disable=SC2029
|
||||
ssh "training@$CLOUD_HOST" "$REMOTE_CMD"
|
||||
echo ""
|
||||
echo "Syncing model artifacts..."
|
||||
rsync -avz "training@$CLOUD_HOST:$CLOUD_DIR/$OUTPUT_DIR/" "./$OUTPUT_DIR/"
|
||||
rsync -avz "training@$CLOUD_HOST:$CLOUD_DIR/checkpoints/" "./checkpoints/"
|
||||
echo "Done."
|
||||
else
|
||||
echo "Skipped. Run the SSH command above manually."
|
||||
fi
|
||||
fi
|
||||
else
|
||||
echo ""
|
||||
echo "==> Running local training on $GPU_NAME"
|
||||
echo ""
|
||||
|
||||
mkdir -p "$OUTPUT_DIR"
|
||||
|
||||
export SQLX_OFFLINE=true
|
||||
|
||||
# shellcheck disable=SC2086
|
||||
exec cargo run --release -p ml --example "$EXAMPLE_NAME" --features cuda -- $CARGO_EXTRA
|
||||
fi
|
||||
@@ -75,6 +75,51 @@ impl PortfolioAllocator {
|
||||
}
|
||||
}
|
||||
|
||||
/// Allocate capital across assets using a correlation matrix
|
||||
///
|
||||
/// Like [`allocate`](Self::allocate), but accepts an N x N correlation matrix
|
||||
/// to build a full covariance matrix for mean-variance optimization.
|
||||
/// Only meaningful when the allocation method is `MeanVariance` or `MLOptimized`;
|
||||
/// other methods ignore the correlation matrix.
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `assets` - Asset information (returns, volatility, ML scores)
|
||||
/// * `total_capital` - Total capital to allocate
|
||||
/// * `correlations` - N x N correlation matrix (must be symmetric, 1.0 on diagonal)
|
||||
///
|
||||
/// # Errors
|
||||
/// Returns an error if the correlation matrix dimensions do not match the asset count.
|
||||
pub fn allocate_with_correlations(
|
||||
&self,
|
||||
assets: &[AssetInfo],
|
||||
total_capital: Decimal,
|
||||
correlations: &DMatrix<f64>,
|
||||
) -> Result<HashMap<String, Decimal>> {
|
||||
if assets.is_empty() {
|
||||
return Ok(HashMap::new());
|
||||
}
|
||||
|
||||
match &self.method {
|
||||
AllocationMethod::MeanVariance { lambda } => {
|
||||
self.mean_variance_with_corr(assets, total_capital, *lambda, Some(correlations))
|
||||
}
|
||||
AllocationMethod::MLOptimized => {
|
||||
// Use ML scores as expected returns, then apply correlated mean-variance
|
||||
let ml_assets: Vec<AssetInfo> = assets
|
||||
.iter()
|
||||
.map(|a| {
|
||||
let mut asset = a.clone();
|
||||
asset.expected_return = a.ml_score;
|
||||
asset
|
||||
})
|
||||
.collect();
|
||||
self.mean_variance_with_corr(&ml_assets, total_capital, 1.0, Some(correlations))
|
||||
}
|
||||
// Other methods don't use correlations — delegate to standard allocate
|
||||
_ => self.allocate(assets, total_capital),
|
||||
}
|
||||
}
|
||||
|
||||
/// Strategy 1: Equal Weight (Baseline)
|
||||
///
|
||||
/// Allocates capital equally across all assets (1/N portfolio).
|
||||
@@ -126,35 +171,69 @@ impl PortfolioAllocator {
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `lambda` - Risk aversion parameter (higher = more conservative)
|
||||
/// * `correlations` - Optional N x N correlation matrix. When `None`, assumes
|
||||
/// independent assets (diagonal covariance). When provided, builds full
|
||||
/// covariance: `Sigma[i][j] = corr[i][j] * vol_i * vol_j`.
|
||||
fn mean_variance(
|
||||
&self,
|
||||
assets: &[AssetInfo],
|
||||
total_capital: Decimal,
|
||||
lambda: f64,
|
||||
) -> Result<HashMap<String, Decimal>> {
|
||||
self.mean_variance_with_corr(assets, total_capital, lambda, None)
|
||||
}
|
||||
|
||||
/// Mean-Variance optimization with optional correlation matrix.
|
||||
///
|
||||
/// When `correlations` is `Some`, builds the full covariance matrix from the
|
||||
/// correlation matrix and per-asset volatilities. Falls back to diagonal
|
||||
/// covariance if the correlation matrix is ill-conditioned.
|
||||
fn mean_variance_with_corr(
|
||||
&self,
|
||||
assets: &[AssetInfo],
|
||||
total_capital: Decimal,
|
||||
lambda: f64,
|
||||
correlations: Option<&DMatrix<f64>>,
|
||||
) -> Result<HashMap<String, Decimal>> {
|
||||
let n = assets.len();
|
||||
|
||||
// Expected returns vector
|
||||
let mu = DVector::from_vec(assets.iter().map(|a| a.expected_return).collect());
|
||||
|
||||
// Covariance matrix — currently diagonal (independent assets).
|
||||
//
|
||||
// A diagonal covariance matrix assumes zero correlation between all asset pairs,
|
||||
// which is a simplification. For full Markowitz optimization with correlation
|
||||
// support, the following steps are needed:
|
||||
//
|
||||
// 1. Accept a `correlations: Option<&DMatrix<f64>>` parameter (N x N correlation matrix)
|
||||
// 2. When provided, build full covariance: Sigma[i][j] = corr[i][j] * vol_i * vol_j
|
||||
// 3. Ensure the correlation matrix is symmetric positive-definite (Cholesky check)
|
||||
// 4. Fall back to diagonal if the matrix is ill-conditioned (det < epsilon)
|
||||
//
|
||||
// nalgebra's DMatrix is already available in this crate, so the implementation
|
||||
// is straightforward once historical return data is available to estimate
|
||||
// pairwise correlations (e.g., via a rolling Pearson correlation window).
|
||||
let mut sigma = DMatrix::zeros(n, n);
|
||||
for (i, asset) in assets.iter().enumerate() {
|
||||
sigma[(i, i)] = asset.volatility.powi(2);
|
||||
}
|
||||
// Build covariance matrix
|
||||
let mut sigma = if let Some(corr) = correlations {
|
||||
// Validate dimensions
|
||||
if corr.nrows() != n || corr.ncols() != n {
|
||||
anyhow::bail!(
|
||||
"Correlation matrix dimensions ({}, {}) do not match asset count {}",
|
||||
corr.nrows(),
|
||||
corr.ncols(),
|
||||
n
|
||||
);
|
||||
}
|
||||
// Build full covariance: Sigma[i][j] = corr[i][j] * vol_i * vol_j
|
||||
let mut cov = DMatrix::zeros(n, n);
|
||||
for i in 0..n {
|
||||
let vol_i = assets.get(i).map(|a| a.volatility).unwrap_or(0.0);
|
||||
for j in 0..n {
|
||||
let vol_j = assets.get(j).map(|a| a.volatility).unwrap_or(0.0);
|
||||
let corr_ij = corr.get((i, j)).copied().unwrap_or(0.0);
|
||||
if let Some(cell) = cov.get_mut((i, j)) {
|
||||
*cell = corr_ij * vol_i * vol_j;
|
||||
}
|
||||
}
|
||||
}
|
||||
cov
|
||||
} else {
|
||||
// Diagonal covariance (independent assets)
|
||||
let mut cov = DMatrix::zeros(n, n);
|
||||
for (i, asset) in assets.iter().enumerate() {
|
||||
if let Some(cell) = cov.get_mut((i, i)) {
|
||||
*cell = asset.volatility.powi(2);
|
||||
}
|
||||
}
|
||||
cov
|
||||
};
|
||||
|
||||
// Add small regularization to diagonal for numerical stability
|
||||
for i in 0..n {
|
||||
@@ -642,4 +721,146 @@ mod tests {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Identity correlation matrix (diagonal = 1.0) should produce the same result
|
||||
/// as the default diagonal covariance path (no correlations).
|
||||
#[test]
|
||||
fn test_mean_variance_identity_correlation_matches_diagonal() {
|
||||
let allocator = PortfolioAllocator::new(AllocationMethod::MeanVariance { lambda: 2.0 });
|
||||
let assets = create_test_assets();
|
||||
let total_capital = Decimal::from(100_000);
|
||||
let n = assets.len();
|
||||
|
||||
// Identity correlation matrix
|
||||
let identity = DMatrix::identity(n, n);
|
||||
|
||||
let alloc_diagonal = allocator.allocate(&assets, total_capital).unwrap();
|
||||
let alloc_identity = allocator
|
||||
.allocate_with_correlations(&assets, total_capital, &identity)
|
||||
.unwrap();
|
||||
|
||||
// Both should produce identical allocations
|
||||
for asset in &assets {
|
||||
let diag_val = alloc_diagonal.get(&asset.symbol).unwrap();
|
||||
let ident_val = alloc_identity.get(&asset.symbol).unwrap();
|
||||
let diff = (*diag_val - *ident_val).abs();
|
||||
assert!(
|
||||
diff < Decimal::from_f64_retain(0.01).unwrap(),
|
||||
"Symbol {} differs: diagonal={}, identity={}",
|
||||
asset.symbol,
|
||||
diag_val,
|
||||
ident_val,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// When two assets are highly correlated, the optimizer should allocate
|
||||
/// differently compared to the uncorrelated (diagonal) case.
|
||||
#[test]
|
||||
fn test_correlated_allocation_differs_from_diagonal() {
|
||||
let allocator = PortfolioAllocator::new(AllocationMethod::MeanVariance { lambda: 2.0 });
|
||||
let assets = create_test_assets(); // ES, NQ, ZN
|
||||
let total_capital = Decimal::from(100_000);
|
||||
let n = assets.len();
|
||||
|
||||
// High correlation between ES and NQ (both equity futures), low with ZN (bonds)
|
||||
let corr_data = vec![
|
||||
1.0, 0.90, 0.10, // ES row
|
||||
0.90, 1.0, 0.10, // NQ row
|
||||
0.10, 0.10, 1.0, // ZN row
|
||||
];
|
||||
let corr = DMatrix::from_row_slice(n, n, &corr_data);
|
||||
|
||||
let alloc_diagonal = allocator.allocate(&assets, total_capital).unwrap();
|
||||
let alloc_correlated = allocator
|
||||
.allocate_with_correlations(&assets, total_capital, &corr)
|
||||
.unwrap();
|
||||
|
||||
// Correlated allocation should differ from diagonal
|
||||
let mut any_differs = false;
|
||||
for asset in &assets {
|
||||
let diag_val = alloc_diagonal.get(&asset.symbol).unwrap();
|
||||
let corr_val = alloc_correlated.get(&asset.symbol).unwrap();
|
||||
if (*diag_val - *corr_val).abs() > Decimal::from_f64_retain(1.0).unwrap() {
|
||||
any_differs = true;
|
||||
}
|
||||
}
|
||||
assert!(
|
||||
any_differs,
|
||||
"Correlated allocation should differ from diagonal allocation"
|
||||
);
|
||||
|
||||
// With high ES-NQ correlation, ZN (diversifier) should get relatively more weight
|
||||
// compared to the diagonal case
|
||||
let zn_diag = alloc_diagonal.get("ZN.FUT").unwrap();
|
||||
let zn_corr = alloc_correlated.get("ZN.FUT").unwrap();
|
||||
assert!(
|
||||
zn_corr > zn_diag,
|
||||
"ZN (uncorrelated diversifier) should get more weight with correlations: corr={}, diag={}",
|
||||
zn_corr,
|
||||
zn_diag,
|
||||
);
|
||||
}
|
||||
|
||||
/// Correlation matrix with wrong dimensions should return an error.
|
||||
#[test]
|
||||
fn test_invalid_correlation_matrix_dimensions() {
|
||||
let allocator = PortfolioAllocator::new(AllocationMethod::MeanVariance { lambda: 2.0 });
|
||||
let assets = create_test_assets(); // 3 assets
|
||||
let total_capital = Decimal::from(100_000);
|
||||
|
||||
// 2x2 matrix for 3 assets — wrong dimensions
|
||||
let bad_corr = DMatrix::identity(2, 2);
|
||||
let result = allocator.allocate_with_correlations(&assets, total_capital, &bad_corr);
|
||||
assert!(result.is_err(), "Should fail with mismatched dimensions");
|
||||
let err_msg = format!("{}", result.unwrap_err());
|
||||
assert!(
|
||||
err_msg.contains("do not match"),
|
||||
"Error should mention dimension mismatch: {}",
|
||||
err_msg
|
||||
);
|
||||
|
||||
// 4x4 matrix for 3 assets — also wrong
|
||||
let bad_corr_large = DMatrix::identity(4, 4);
|
||||
let result = allocator.allocate_with_correlations(&assets, total_capital, &bad_corr_large);
|
||||
assert!(result.is_err(), "Should fail with oversized dimensions");
|
||||
}
|
||||
|
||||
/// Non-square correlation matrix should also fail.
|
||||
#[test]
|
||||
fn test_non_square_correlation_matrix() {
|
||||
let allocator = PortfolioAllocator::new(AllocationMethod::MeanVariance { lambda: 2.0 });
|
||||
let assets = create_test_assets();
|
||||
let total_capital = Decimal::from(100_000);
|
||||
|
||||
// 3x2 matrix — not square
|
||||
let bad_corr = DMatrix::zeros(3, 2);
|
||||
let result = allocator.allocate_with_correlations(&assets, total_capital, &bad_corr);
|
||||
assert!(result.is_err(), "Should fail with non-square matrix");
|
||||
}
|
||||
|
||||
/// Allocate with correlations on non-MeanVariance methods should delegate
|
||||
/// to standard allocate (correlations ignored).
|
||||
#[test]
|
||||
fn test_correlations_ignored_for_equal_weight() {
|
||||
let allocator = PortfolioAllocator::new(AllocationMethod::EqualWeight);
|
||||
let assets = create_test_assets();
|
||||
let total_capital = Decimal::from(100_000);
|
||||
let n = assets.len();
|
||||
|
||||
let corr = DMatrix::identity(n, n);
|
||||
let alloc_std = allocator.allocate(&assets, total_capital).unwrap();
|
||||
let alloc_corr = allocator
|
||||
.allocate_with_correlations(&assets, total_capital, &corr)
|
||||
.unwrap();
|
||||
|
||||
for asset in &assets {
|
||||
assert_eq!(
|
||||
alloc_std.get(&asset.symbol),
|
||||
alloc_corr.get(&asset.symbol),
|
||||
"EqualWeight should ignore correlations for {}",
|
||||
asset.symbol
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -269,6 +269,37 @@ impl TradingAgentServiceImpl {
|
||||
}
|
||||
}
|
||||
|
||||
/// Fetch current net positions from the `agent_orders` table.
|
||||
///
|
||||
/// Returns a map of symbol -> net quantity (buys positive, sells negative)
|
||||
/// derived from non-cancelled orders. Returns an empty map on DB error
|
||||
/// so callers degrade gracefully to zero-position assumptions.
|
||||
async fn fetch_current_positions(&self) -> HashMap<String, f64> {
|
||||
let rows: Result<Vec<(String, f64)>, _> = sqlx::query_as(
|
||||
r#"
|
||||
SELECT symbol, COALESCE(SUM(
|
||||
CASE WHEN side = 'Buy' THEN CAST(quantity AS DOUBLE PRECISION)
|
||||
WHEN side = 'Sell' THEN -CAST(quantity AS DOUBLE PRECISION)
|
||||
ELSE 0.0
|
||||
END
|
||||
), 0.0) as net_quantity
|
||||
FROM agent_orders
|
||||
WHERE status != 'CANCELLED'
|
||||
GROUP BY symbol
|
||||
"#,
|
||||
)
|
||||
.fetch_all(&self.db_pool)
|
||||
.await;
|
||||
|
||||
match rows {
|
||||
Ok(data) => data.into_iter().collect(),
|
||||
Err(e) => {
|
||||
warn!("Failed to fetch current positions (defaulting to empty): {}", e);
|
||||
HashMap::new()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Convert internal Instrument to proto
|
||||
fn convert_instrument(&self, inst: &crate::universe::Instrument) -> Instrument {
|
||||
Instrument {
|
||||
@@ -771,14 +802,24 @@ impl trading_agent_service_server::TradingAgentService for TradingAgentServiceIm
|
||||
Status::internal(format!("Allocation failed: {}", e))
|
||||
})?;
|
||||
|
||||
// 4. Convert to proto — compute target_quantity from last price
|
||||
// 4. Convert to proto -- compute target_quantity from last price,
|
||||
// and populate current position data from agent_orders.
|
||||
//
|
||||
// BLOCKER: current_weight and current_quantity require live position data from
|
||||
// a portfolio/position-tracking service (or a positions table). This service does
|
||||
// not currently have access to live position state. When a PositionService gRPC
|
||||
// client is added, these should be fetched per-symbol.
|
||||
// rebalance_delta = target_quantity - current_quantity (set to target_quantity
|
||||
// until current positions are available).
|
||||
// Current positions are derived from the agent_orders table (net buy - sell
|
||||
// quantities per symbol). Current weight is the symbol's share of total
|
||||
// portfolio exposure valued at last close prices. Falls back to zeros when
|
||||
// position or price data is unavailable.
|
||||
let current_positions = self.fetch_current_positions().await;
|
||||
|
||||
// Compute total portfolio value from current positions * last prices
|
||||
let total_current_value: f64 = current_positions
|
||||
.iter()
|
||||
.map(|(sym, qty)| {
|
||||
let price = symbol_last_prices.get(sym).copied().unwrap_or(0.0);
|
||||
qty.abs() * price
|
||||
})
|
||||
.sum();
|
||||
|
||||
let proto_allocations: Vec<AssetAllocation> = allocations
|
||||
.iter()
|
||||
.map(|(symbol, capital)| {
|
||||
@@ -796,17 +837,28 @@ impl trading_agent_service_server::TradingAgentService for TradingAgentServiceIm
|
||||
.map(|price| capital_f64 / price)
|
||||
.unwrap_or(0.0); // 0.0 if no price data available
|
||||
|
||||
// Current position from agent_orders (net quantity)
|
||||
let current_quantity = current_positions
|
||||
.get(symbol)
|
||||
.copied()
|
||||
.unwrap_or(0.0);
|
||||
|
||||
// Current weight: position value / total portfolio value
|
||||
let current_weight = if total_current_value > 0.0 {
|
||||
let price = symbol_last_prices.get(symbol).copied().unwrap_or(0.0);
|
||||
(current_quantity.abs() * price) / total_current_value
|
||||
} else {
|
||||
0.0
|
||||
};
|
||||
|
||||
AssetAllocation {
|
||||
symbol: symbol.clone(),
|
||||
target_weight: weight,
|
||||
target_capital: capital_f64,
|
||||
target_quantity,
|
||||
// TODO(positions): Fetch from live position service / positions table.
|
||||
// Requires PositionService gRPC client or position-tracking DB query.
|
||||
current_weight: 0.0,
|
||||
current_quantity: 0.0,
|
||||
// Once current_quantity is available: target_quantity - current_quantity
|
||||
rebalance_delta: target_quantity,
|
||||
current_weight,
|
||||
current_quantity,
|
||||
rebalance_delta: target_quantity - current_quantity,
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
@@ -1057,25 +1109,8 @@ impl trading_agent_service_server::TradingAgentService for TradingAgentServiceIm
|
||||
.collect();
|
||||
|
||||
// 3. Load current positions from agent_orders to derive current weights
|
||||
let current_rows: Vec<(String, f64)> = sqlx::query_as(
|
||||
r#"
|
||||
SELECT symbol, COALESCE(SUM(
|
||||
CASE WHEN side = 'Buy' THEN CAST(quantity AS DOUBLE PRECISION)
|
||||
WHEN side = 'Sell' THEN -CAST(quantity AS DOUBLE PRECISION)
|
||||
ELSE 0.0
|
||||
END
|
||||
), 0.0) as net_quantity
|
||||
FROM agent_orders
|
||||
WHERE status != 'CANCELLED'
|
||||
GROUP BY symbol
|
||||
"#,
|
||||
)
|
||||
.fetch_all(&self.db_pool)
|
||||
.await
|
||||
.map_err(|e| {
|
||||
warn!("Failed to query current positions (non-fatal): {}", e);
|
||||
Status::internal(format!("Failed to query positions: {e}"))
|
||||
})?;
|
||||
let current_positions = self.fetch_current_positions().await;
|
||||
let current_rows: Vec<(String, f64)> = current_positions.into_iter().collect();
|
||||
|
||||
let total_current: f64 = current_rows.iter().map(|(_, q)| q.abs()).sum();
|
||||
let current_weights: HashMap<String, f64> = if total_current > 0.0 {
|
||||
|
||||
@@ -114,9 +114,15 @@ impl FeaturePreprocessor {
|
||||
|
||||
/// Normalize a feature value using z-score normalization
|
||||
pub fn normalize(&self, feature_name: &str, value: f64) -> f64 {
|
||||
if !value.is_finite() {
|
||||
warn!(feature = %feature_name, value = %value, "NaN/Inf feature detected, replacing with 0.0");
|
||||
return 0.0;
|
||||
}
|
||||
|
||||
if let Some(stats) = self.stats.get(feature_name) {
|
||||
if stats.std_dev > 0.0 {
|
||||
(value - stats.mean) / stats.std_dev
|
||||
let normalized = (value - stats.mean) / stats.std_dev;
|
||||
normalized.clamp(-10.0, 10.0)
|
||||
} else {
|
||||
value
|
||||
}
|
||||
@@ -253,7 +259,7 @@ impl EnhancedMLServiceImpl {
|
||||
"LIQUID"
|
||||
} else {
|
||||
return Err(Status::invalid_argument(format!(
|
||||
"Unknown model type in model_id: {}",
|
||||
"Unknown model type in model_id: {}. Supported: DQN, PPO, TFT, MAMBA2, CFC/Liquid",
|
||||
model_id
|
||||
)));
|
||||
};
|
||||
@@ -1985,3 +1991,283 @@ impl MLModel for LiquidModel {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod feature_preprocessor_tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_normalize_nan_returns_zero() {
|
||||
let preprocessor = FeaturePreprocessor::new();
|
||||
let result = preprocessor.normalize("price_momentum", f64::NAN);
|
||||
assert!(result.is_finite());
|
||||
assert_eq!(result, 0.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_normalize_inf_returns_zero() {
|
||||
let preprocessor = FeaturePreprocessor::new();
|
||||
let result = preprocessor.normalize("price_momentum", f64::INFINITY);
|
||||
assert!(result.is_finite());
|
||||
assert_eq!(result, 0.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_normalize_neg_inf_returns_zero() {
|
||||
let preprocessor = FeaturePreprocessor::new();
|
||||
let result = preprocessor.normalize("price_momentum", f64::NEG_INFINITY);
|
||||
assert!(result.is_finite());
|
||||
assert_eq!(result, 0.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_normalize_clamps_extreme_values() {
|
||||
let preprocessor = FeaturePreprocessor::new();
|
||||
// Price momentum has mean=0.0, std_dev=0.1, so value=100.0 would be z=1000
|
||||
let result = preprocessor.normalize("price_momentum", 100.0);
|
||||
assert!(result <= 10.0);
|
||||
assert!(result >= -10.0);
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[allow(clippy::unwrap_used)]
|
||||
mod enhanced_ml_tests {
|
||||
use super::*;
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// 1. EnsembleConfig::default() values
|
||||
// -----------------------------------------------------------------------
|
||||
#[test]
|
||||
fn test_ensemble_config_defaults() {
|
||||
let config = EnsembleConfig::default();
|
||||
assert_eq!(config.min_models, 2, "min_models should default to 2");
|
||||
assert!(
|
||||
(config.confidence_threshold - 0.7).abs() < 1e-10,
|
||||
"confidence_threshold should default to 0.7"
|
||||
);
|
||||
assert!(
|
||||
config.use_weighted_voting,
|
||||
"use_weighted_voting should default to true"
|
||||
);
|
||||
assert_eq!(
|
||||
config.fallback_timeout_ms, 50,
|
||||
"fallback_timeout_ms should default to 50"
|
||||
);
|
||||
assert!(
|
||||
(config.consensus_threshold - 0.6).abs() < 1e-10,
|
||||
"consensus_threshold should default to 0.6"
|
||||
);
|
||||
}
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// 2. FeaturePreprocessor::classify_feature_type
|
||||
// -----------------------------------------------------------------------
|
||||
#[test]
|
||||
fn test_classify_price_features() {
|
||||
let pp = FeaturePreprocessor::new();
|
||||
assert_eq!(pp.classify_feature_type("price_momentum") as i32, FeatureType::Price as i32);
|
||||
assert_eq!(pp.classify_feature_type("close_price") as i32, FeatureType::Price as i32);
|
||||
assert_eq!(pp.classify_feature_type("momentum_5m") as i32, FeatureType::Price as i32);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_classify_volume_features() {
|
||||
let pp = FeaturePreprocessor::new();
|
||||
assert_eq!(pp.classify_feature_type("volume") as i32, FeatureType::Volume as i32);
|
||||
assert_eq!(pp.classify_feature_type("volume_ratio") as i32, FeatureType::Volume as i32);
|
||||
assert_eq!(pp.classify_feature_type("orderbook_depth") as i32, FeatureType::Volume as i32);
|
||||
assert_eq!(pp.classify_feature_type("depth_imbalance") as i32, FeatureType::Volume as i32);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_classify_technical_features() {
|
||||
let pp = FeaturePreprocessor::new();
|
||||
assert_eq!(pp.classify_feature_type("volatility") as i32, FeatureType::Technical as i32);
|
||||
assert_eq!(pp.classify_feature_type("rsi_14") as i32, FeatureType::Technical as i32);
|
||||
assert_eq!(pp.classify_feature_type("ma_20") as i32, FeatureType::Technical as i32);
|
||||
assert_eq!(pp.classify_feature_type("spread_bps") as i32, FeatureType::Technical as i32);
|
||||
assert_eq!(pp.classify_feature_type("liquidity_score") as i32, FeatureType::Technical as i32);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_classify_sentiment_features() {
|
||||
let pp = FeaturePreprocessor::new();
|
||||
assert_eq!(pp.classify_feature_type("sentiment_score") as i32, FeatureType::Sentiment as i32);
|
||||
assert_eq!(pp.classify_feature_type("news_impact") as i32, FeatureType::Sentiment as i32);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_classify_unknown_defaults_to_technical() {
|
||||
let pp = FeaturePreprocessor::new();
|
||||
assert_eq!(pp.classify_feature_type("foo_bar_baz") as i32, FeatureType::Technical as i32);
|
||||
assert_eq!(pp.classify_feature_type("") as i32, FeatureType::Technical as i32);
|
||||
}
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// 3. FeaturePreprocessor normalization (z-score, tanh fallback)
|
||||
// -----------------------------------------------------------------------
|
||||
#[test]
|
||||
fn test_normalize_known_feature_zscore() {
|
||||
let pp = FeaturePreprocessor::new();
|
||||
// price_momentum: mean=0.0, std_dev=0.1
|
||||
// z-score for value=0.05: (0.05 - 0.0) / 0.1 = 0.5
|
||||
let result = pp.normalize("price_momentum", 0.05);
|
||||
assert!((result - 0.5).abs() < 1e-10, "Expected 0.5; got {}", result);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_normalize_volume_zscore() {
|
||||
let pp = FeaturePreprocessor::new();
|
||||
// volume: mean=1_000_000, std_dev=500_000
|
||||
// z-score for value=1_500_000: (1_500_000 - 1_000_000) / 500_000 = 1.0
|
||||
let result = pp.normalize("volume", 1_500_000.0);
|
||||
assert!((result - 1.0).abs() < 1e-10, "Expected 1.0; got {}", result);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_normalize_unknown_feature_uses_tanh() {
|
||||
let pp = FeaturePreprocessor::new();
|
||||
// Unknown feature uses value.tanh()
|
||||
let value: f64 = 0.5;
|
||||
let expected = value.tanh();
|
||||
let result = pp.normalize("unknown_feature_xyz", value);
|
||||
assert!((result - expected).abs() < 1e-10, "Expected tanh({})={}; got {}", value, expected, result);
|
||||
}
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// 4. FeaturePreprocessor default stats
|
||||
// -----------------------------------------------------------------------
|
||||
#[test]
|
||||
fn test_preprocessor_default_has_three_features() {
|
||||
let pp = FeaturePreprocessor::default();
|
||||
assert_eq!(pp.stats.len(), 3, "Default preprocessor should have 3 feature stats");
|
||||
assert!(pp.stats.contains_key("price_momentum"));
|
||||
assert!(pp.stats.contains_key("volume"));
|
||||
assert!(pp.stats.contains_key("volatility"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_preprocessor_new_equals_default() {
|
||||
let pp_new = FeaturePreprocessor::new();
|
||||
let pp_default = FeaturePreprocessor::default();
|
||||
assert_eq!(pp_new.stats.len(), pp_default.stats.len());
|
||||
for (key, new_stat) in &pp_new.stats {
|
||||
let default_stat = pp_default.stats.get(key).unwrap();
|
||||
assert!((new_stat.mean - default_stat.mean).abs() < 1e-10);
|
||||
assert!((new_stat.std_dev - default_stat.std_dev).abs() < 1e-10);
|
||||
}
|
||||
}
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// 5. RuntimeModelInfo creation
|
||||
// -----------------------------------------------------------------------
|
||||
#[test]
|
||||
fn test_runtime_model_info_creation() {
|
||||
let info = RuntimeModelInfo {
|
||||
model_id: "test-dqn-v1".to_string(),
|
||||
version: "1.0.0".to_string(),
|
||||
load_time: SystemTime::now(),
|
||||
last_inference: None,
|
||||
inference_count: 0,
|
||||
error_count: 0,
|
||||
avg_latency_us: 0.0,
|
||||
confidence_threshold: 0.7,
|
||||
weight_in_ensemble: 1.0,
|
||||
fallback_priority: 0,
|
||||
model_type: ModelType::DQN,
|
||||
supported_symbols: vec!["EURUSD".to_string(), "GBPUSD".to_string()],
|
||||
supported_horizons: vec![1, 5, 15],
|
||||
feature_count: 16,
|
||||
model_instance: None,
|
||||
};
|
||||
|
||||
assert_eq!(info.model_id, "test-dqn-v1");
|
||||
assert_eq!(info.model_type, ModelType::DQN);
|
||||
assert_eq!(info.feature_count, 16);
|
||||
assert_eq!(info.supported_symbols.len(), 2);
|
||||
assert_eq!(info.supported_horizons.len(), 3);
|
||||
assert!(info.last_inference.is_none());
|
||||
assert!(info.model_instance.is_none());
|
||||
assert_eq!(info.inference_count, 0);
|
||||
assert_eq!(info.error_count, 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_runtime_model_info_various_model_types() {
|
||||
let model_types = vec![
|
||||
(ModelType::DQN, "DQN"),
|
||||
(ModelType::PPO, "PPO"),
|
||||
(ModelType::TFT, "TFT"),
|
||||
(ModelType::Mamba, "Mamba"),
|
||||
(ModelType::LNN, "LNN"),
|
||||
];
|
||||
|
||||
for (model_type, name) in model_types {
|
||||
let info = RuntimeModelInfo {
|
||||
model_id: format!("test-{}", name),
|
||||
version: "1.0.0".to_string(),
|
||||
load_time: SystemTime::now(),
|
||||
last_inference: None,
|
||||
inference_count: 0,
|
||||
error_count: 0,
|
||||
avg_latency_us: 0.0,
|
||||
confidence_threshold: 0.7,
|
||||
weight_in_ensemble: 0.25,
|
||||
fallback_priority: 1,
|
||||
model_type,
|
||||
supported_symbols: vec![],
|
||||
supported_horizons: vec![],
|
||||
feature_count: 16,
|
||||
model_instance: None,
|
||||
};
|
||||
assert_eq!(info.model_type, model_type, "Model type mismatch for {}", name);
|
||||
assert_eq!(info.weight_in_ensemble, 0.25);
|
||||
}
|
||||
}
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// 6. ModelPerformanceMetrics defaults
|
||||
// -----------------------------------------------------------------------
|
||||
#[test]
|
||||
fn test_model_performance_metrics_default() {
|
||||
let metrics = ModelPerformanceMetrics::default();
|
||||
assert_eq!(metrics.total_predictions, 0);
|
||||
assert_eq!(metrics.successful_predictions, 0);
|
||||
assert_eq!(metrics.failed_predictions, 0);
|
||||
assert!((metrics.avg_latency_us - 0.0).abs() < 1e-10);
|
||||
assert!((metrics.p95_latency_us - 0.0).abs() < 1e-10);
|
||||
assert!((metrics.accuracy_percentage - 0.0).abs() < 1e-10);
|
||||
assert!(metrics.last_health_check.is_none());
|
||||
}
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// 7. FeatureNormStats bounds validation
|
||||
// -----------------------------------------------------------------------
|
||||
#[test]
|
||||
fn test_feature_norm_stats_volatility_defaults() {
|
||||
let pp = FeaturePreprocessor::default();
|
||||
let vol_stats = pp.stats.get("volatility").unwrap();
|
||||
assert!((vol_stats.mean - 0.02).abs() < 1e-10, "Volatility mean should be 0.02");
|
||||
assert!((vol_stats.std_dev - 0.01).abs() < 1e-10, "Volatility std_dev should be 0.01");
|
||||
assert!((vol_stats.min - 0.0).abs() < 1e-10, "Volatility min should be 0.0");
|
||||
assert!((vol_stats.max - 0.5).abs() < 1e-10, "Volatility max should be 0.5");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_normalize_zero_std_dev_returns_raw() {
|
||||
// Create a feature with zero std_dev -- should return value as-is
|
||||
let mut pp = FeaturePreprocessor::new();
|
||||
pp.stats.insert(
|
||||
"zero_std".to_string(),
|
||||
FeatureNormStats {
|
||||
mean: 5.0,
|
||||
std_dev: 0.0,
|
||||
min: 0.0,
|
||||
max: 10.0,
|
||||
},
|
||||
);
|
||||
let result = pp.normalize("zero_std", 7.0);
|
||||
assert!((result - 7.0).abs() < 1e-10, "Zero std_dev should return raw value; got {}", result);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -308,23 +308,42 @@ impl RiskService for RiskServiceImpl {
|
||||
let lookback_days = req.lookback_days;
|
||||
let method = VaRMethod::try_from(req.method).unwrap_or(VaRMethod::VarMethodHistorical);
|
||||
|
||||
// Fetch live positions FIRST so we can compute real portfolio notional for VaR.
|
||||
let positions = self.fetch_positions().await;
|
||||
let position_map: std::collections::HashMap<String, f64> = positions
|
||||
.iter()
|
||||
.map(|p| (p.symbol.clone(), p.quantity.abs()))
|
||||
.collect();
|
||||
|
||||
// Portfolio notional = sum of |quantity * avg_price| across all open positions.
|
||||
// Falls back to 100_000.0 when no positions exist (conservative default for
|
||||
// empty-portfolio queries so VaR still returns a meaningful estimate).
|
||||
let portfolio_notional: f64 = positions
|
||||
.iter()
|
||||
.map(|p| (p.quantity * p.average_price).abs())
|
||||
.sum();
|
||||
let portfolio_notional = if portfolio_notional > 0.0 {
|
||||
portfolio_notional
|
||||
} else {
|
||||
100_000.0
|
||||
};
|
||||
|
||||
// Delegate to the real RiskEngine for marginal VaR.
|
||||
// calculate_comprehensive_var requires full historical price data that is not available
|
||||
// at the gRPC boundary, so we use calculate_marginal_var as the portfolio-level estimate
|
||||
// with a representative notional value.
|
||||
// TODO: Feed real position data from position_manager when available.
|
||||
// with the real portfolio notional derived from open positions above.
|
||||
let risk_engine = self.state.risk_engine.read().await;
|
||||
let portfolio_var = match risk_engine
|
||||
.calculate_marginal_var(
|
||||
"portfolio",
|
||||
"PORTFOLIO",
|
||||
confidence_level * 1_000_000.0, // notional proxy scaled by confidence
|
||||
portfolio_notional,
|
||||
1.0,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(var) => {
|
||||
info!("VaR calculated via RiskEngine: {:.4}", var);
|
||||
info!("VaR calculated via RiskEngine: notional={:.2}, var={:.4}", portfolio_notional, var);
|
||||
var
|
||||
},
|
||||
Err(e) => {
|
||||
@@ -332,23 +351,11 @@ impl RiskService for RiskServiceImpl {
|
||||
"RiskEngine VaR calculation failed ({}), falling back to parametric estimate",
|
||||
e
|
||||
);
|
||||
// Parametric fallback: confidence_level * 2% daily volatility assumption
|
||||
confidence_level * 0.02
|
||||
// Parametric fallback: 2% daily volatility assumption on real notional
|
||||
portfolio_notional * 0.02
|
||||
},
|
||||
};
|
||||
|
||||
// Fetch live positions to get real position sizes per symbol.
|
||||
// Drop the risk_engine read lock first to avoid holding it across the await.
|
||||
drop(risk_engine);
|
||||
let positions = self.fetch_positions().await;
|
||||
let position_map: std::collections::HashMap<String, f64> = positions
|
||||
.iter()
|
||||
.map(|p| (p.symbol.clone(), p.quantity.abs()))
|
||||
.collect();
|
||||
|
||||
// Re-acquire the risk engine lock for per-symbol VaR calculations.
|
||||
let risk_engine = self.state.risk_engine.read().await;
|
||||
|
||||
// Build per-symbol marginal VaRs using the same engine.
|
||||
// Each symbol's contribution is calculated individually; if a symbol fails we skip it.
|
||||
let num_symbols = req.symbols.len();
|
||||
@@ -406,23 +413,40 @@ impl RiskService for RiskServiceImpl {
|
||||
&self,
|
||||
_request: Request<GetRiskMetricsRequest>,
|
||||
) -> Result<Response<GetRiskMetricsResponse>, Status> {
|
||||
// Fetch live position data FIRST so we can compute real portfolio notional for VaR.
|
||||
let positions = self.fetch_positions().await;
|
||||
|
||||
// Portfolio notional = sum of |quantity * avg_price| across all open positions.
|
||||
// Falls back to 100_000.0 when no positions exist (conservative default so VaR
|
||||
// still returns a meaningful estimate for empty portfolios).
|
||||
let portfolio_notional: f64 = positions
|
||||
.iter()
|
||||
.map(|p| (p.quantity * p.average_price).abs())
|
||||
.sum();
|
||||
let portfolio_notional = if portfolio_notional > 0.0 {
|
||||
portfolio_notional
|
||||
} else {
|
||||
100_000.0
|
||||
};
|
||||
|
||||
// Use the real RiskEngine's configured confidence level and max VaR limit
|
||||
// to produce the 1d VaR estimate. Longer horizons scale by sqrt(T).
|
||||
let risk_engine = self.state.risk_engine.read().await;
|
||||
let confidence = risk_engine.var_confidence();
|
||||
|
||||
// Compute a representative 1-day portfolio VaR via marginal VaR.
|
||||
// TODO: Replace with calculate_comprehensive_var once position_manager
|
||||
// provides real PositionInfo and historical price data.
|
||||
// Compute a representative 1-day portfolio VaR via marginal VaR using the
|
||||
// real portfolio notional derived from open positions above.
|
||||
// TODO: Replace with calculate_comprehensive_var once historical price data
|
||||
// is available at the gRPC boundary for full parametric/Monte Carlo VaR.
|
||||
let portfolio_var_1d = match risk_engine
|
||||
.calculate_marginal_var("portfolio", "PORTFOLIO", confidence * 1_000_000.0, 1.0)
|
||||
.calculate_marginal_var("portfolio", "PORTFOLIO", portfolio_notional, 1.0)
|
||||
.await
|
||||
{
|
||||
Ok(var) => var,
|
||||
Err(e) => {
|
||||
warn!("RiskEngine marginal VaR failed for get_risk_metrics: {}", e);
|
||||
// Parametric fallback using configured confidence level
|
||||
confidence * 0.02
|
||||
// Parametric fallback: 2% daily volatility assumption on real notional
|
||||
portfolio_notional * 0.02
|
||||
},
|
||||
};
|
||||
// Drop the read lock before fetching from repositories
|
||||
@@ -440,9 +464,6 @@ impl RiskService for RiskServiceImpl {
|
||||
.await
|
||||
.map_err(|e| Status::internal(format!("Failed to get max drawdown config: {}", e)))?
|
||||
.unwrap_or(0.10);
|
||||
|
||||
// Fetch live position data and execution history for metric calculations
|
||||
let positions = self.fetch_positions().await;
|
||||
let executions = self.fetch_executions().await;
|
||||
|
||||
// Current drawdown from open positions' unrealized PnL
|
||||
@@ -970,3 +991,335 @@ impl RiskService for RiskServiceImpl {
|
||||
)))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[allow(clippy::unwrap_used, clippy::indexing_slicing)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::repositories::{ExecutionEvent, TradingPosition};
|
||||
use common::OrderSide;
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// Helper: create a TradingPosition
|
||||
// -----------------------------------------------------------------------
|
||||
fn position(symbol: &str, qty: f64, avg_price: f64, market_value: f64, pnl: f64) -> TradingPosition {
|
||||
TradingPosition {
|
||||
account_id: "test-account".to_string(),
|
||||
symbol: symbol.to_string(),
|
||||
quantity: qty,
|
||||
average_price: avg_price,
|
||||
market_value,
|
||||
unrealized_pnl: pnl,
|
||||
timestamp: 1_700_000_000,
|
||||
}
|
||||
}
|
||||
|
||||
fn execution(ts: i64, price: f64) -> ExecutionEvent {
|
||||
ExecutionEvent {
|
||||
id: format!("exec-{}", ts),
|
||||
order_id: "order-1".to_string(),
|
||||
account_id: "test-account".to_string(),
|
||||
symbol: "EURUSD".to_string(),
|
||||
side: OrderSide::Buy,
|
||||
quantity: 100.0,
|
||||
price,
|
||||
timestamp: ts,
|
||||
}
|
||||
}
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// 1. Parametric VaR fallback formula
|
||||
// -----------------------------------------------------------------------
|
||||
#[test]
|
||||
fn test_parametric_var_fallback_formula() {
|
||||
// When RiskEngine fails, VaR = portfolio_notional * 0.02
|
||||
let portfolio_notional = 500_000.0_f64;
|
||||
let fallback_var = portfolio_notional * 0.02;
|
||||
assert!((fallback_var - 10_000.0).abs() < 1e-6,
|
||||
"VaR fallback for 500k notional should be 10,000; got {}", fallback_var);
|
||||
|
||||
// Zero notional
|
||||
let zero_var = 0.0_f64 * 0.02;
|
||||
assert!((zero_var - 0.0).abs() < 1e-10);
|
||||
|
||||
// Very large notional
|
||||
let large = 1_000_000_000.0_f64 * 0.02;
|
||||
assert!((large - 20_000_000.0).abs() < 1e-6);
|
||||
}
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// 2. equal_contribution_pct calculation
|
||||
// -----------------------------------------------------------------------
|
||||
#[test]
|
||||
fn test_equal_contribution_pct_for_n_symbols() {
|
||||
// 4 symbols: each contributes 25%
|
||||
let num_symbols = 4_usize;
|
||||
let pct = if num_symbols > 0 { 100.0 / num_symbols as f64 } else { 0.0 };
|
||||
assert!((pct - 25.0).abs() < 1e-10);
|
||||
|
||||
// 1 symbol: 100%
|
||||
let pct1 = 100.0 / 1.0_f64;
|
||||
assert!((pct1 - 100.0).abs() < 1e-10);
|
||||
|
||||
// 10 symbols: 10%
|
||||
let pct10 = 100.0 / 10.0_f64;
|
||||
assert!((pct10 - 10.0).abs() < 1e-10);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_equal_contribution_pct_empty_symbols() {
|
||||
let num_symbols = 0_usize;
|
||||
let pct = if num_symbols > 0 { 100.0 / num_symbols as f64 } else { 0.0 };
|
||||
assert!((pct - 0.0).abs() < 1e-10);
|
||||
}
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// 3. compute_current_drawdown
|
||||
// -----------------------------------------------------------------------
|
||||
#[test]
|
||||
fn test_drawdown_empty_positions() {
|
||||
let dd = RiskServiceImpl::compute_current_drawdown(&[]);
|
||||
assert!((dd - 0.0).abs() < 1e-10);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_drawdown_positive_pnl_is_zero() {
|
||||
let positions = vec![
|
||||
position("EURUSD", 1000.0, 1.10, 11000.0, 500.0),
|
||||
];
|
||||
let dd = RiskServiceImpl::compute_current_drawdown(&positions);
|
||||
assert!((dd - 0.0).abs() < 1e-10, "Positive PnL should yield zero drawdown");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_drawdown_negative_pnl() {
|
||||
// market_value = 10000, unrealized_pnl = -2000
|
||||
// drawdown = 2000 / 10000 = 0.20
|
||||
let positions = vec![
|
||||
position("EURUSD", 1000.0, 1.10, 10000.0, -2000.0),
|
||||
];
|
||||
let dd = RiskServiceImpl::compute_current_drawdown(&positions);
|
||||
assert!((dd - 0.20).abs() < 1e-10, "Drawdown should be 0.20; got {}", dd);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_drawdown_multiple_positions_mixed_pnl() {
|
||||
// pos1: market_value=10000, pnl=-3000
|
||||
// pos2: market_value=5000, pnl=+1000
|
||||
// total_market_value = 15000, total_pnl = -2000
|
||||
// drawdown = 2000 / 15000 = 0.1333...
|
||||
let positions = vec![
|
||||
position("EURUSD", 1000.0, 10.0, 10000.0, -3000.0),
|
||||
position("GBPUSD", 500.0, 10.0, 5000.0, 1000.0),
|
||||
];
|
||||
let dd = RiskServiceImpl::compute_current_drawdown(&positions);
|
||||
let expected = 2000.0 / 15000.0;
|
||||
assert!((dd - expected).abs() < 1e-10, "Expected drawdown {:.6}; got {:.6}", expected, dd);
|
||||
}
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// 4. compute_returns_from_executions
|
||||
// -----------------------------------------------------------------------
|
||||
#[test]
|
||||
fn test_returns_empty_executions() {
|
||||
let returns = RiskServiceImpl::compute_returns_from_executions(&[]);
|
||||
assert!(returns.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_returns_single_execution() {
|
||||
let execs = vec![execution(1, 100.0)];
|
||||
let returns = RiskServiceImpl::compute_returns_from_executions(&execs);
|
||||
assert!(returns.is_empty(), "Single execution should yield no returns");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_returns_two_executions() {
|
||||
// price goes from 100 to 110 => return = 110/100 - 1 = 0.10
|
||||
let execs = vec![execution(1, 100.0), execution(2, 110.0)];
|
||||
let returns = RiskServiceImpl::compute_returns_from_executions(&execs);
|
||||
assert_eq!(returns.len(), 1);
|
||||
assert!((returns[0] - 0.10).abs() < 1e-10);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_returns_negative_price_filtered() {
|
||||
// Negative and zero prices should be filtered out
|
||||
let execs = vec![
|
||||
execution(1, 100.0),
|
||||
execution(2, 0.0), // zero price filtered
|
||||
execution(3, 120.0),
|
||||
];
|
||||
let returns = RiskServiceImpl::compute_returns_from_executions(&execs);
|
||||
// After filtering zero price: prices = [(1, 100), (3, 120)]
|
||||
// return = 120/100 - 1 = 0.20
|
||||
assert_eq!(returns.len(), 1);
|
||||
assert!((returns[0] - 0.20).abs() < 1e-10);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_returns_sorted_by_timestamp() {
|
||||
// Executions provided out of order should be sorted
|
||||
let execs = vec![
|
||||
execution(3, 120.0),
|
||||
execution(1, 100.0),
|
||||
execution(2, 110.0),
|
||||
];
|
||||
let returns = RiskServiceImpl::compute_returns_from_executions(&execs);
|
||||
assert_eq!(returns.len(), 2);
|
||||
// 100 -> 110: +10%
|
||||
assert!((returns[0] - 0.10).abs() < 1e-10);
|
||||
// 110 -> 120: +9.09%
|
||||
assert!((returns[1] - (120.0 / 110.0 - 1.0)).abs() < 1e-10);
|
||||
}
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// 5. compute_volatility
|
||||
// -----------------------------------------------------------------------
|
||||
#[test]
|
||||
fn test_volatility_empty() {
|
||||
assert!((RiskServiceImpl::compute_volatility(&[]) - 0.0).abs() < 1e-10);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_volatility_single_return() {
|
||||
assert!((RiskServiceImpl::compute_volatility(&[0.01]) - 0.0).abs() < 1e-10);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_volatility_constant_returns_is_zero() {
|
||||
// All returns the same => std dev = 0 => vol = 0
|
||||
let returns = vec![0.01, 0.01, 0.01, 0.01, 0.01];
|
||||
let vol = RiskServiceImpl::compute_volatility(&returns);
|
||||
assert!((vol - 0.0).abs() < 1e-10, "Constant returns should have zero vol");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_volatility_known_series() {
|
||||
// Returns: [0.01, -0.01, 0.02, -0.02]
|
||||
// mean = 0.0
|
||||
// variance = (0.0001 + 0.0001 + 0.0004 + 0.0004) / 3 = 0.001 / 3
|
||||
// daily_vol = sqrt(0.001/3)
|
||||
// annualized = daily_vol * sqrt(252)
|
||||
let returns = vec![0.01, -0.01, 0.02, -0.02];
|
||||
let vol = RiskServiceImpl::compute_volatility(&returns);
|
||||
let expected_variance: f64 = 0.001 / 3.0;
|
||||
let expected = expected_variance.sqrt() * 252.0_f64.sqrt();
|
||||
assert!((vol - expected).abs() < 1e-10, "Expected vol {:.6}; got {:.6}", expected, vol);
|
||||
}
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// 6. compute_sharpe_ratio
|
||||
// -----------------------------------------------------------------------
|
||||
#[test]
|
||||
fn test_sharpe_insufficient_data() {
|
||||
// Fewer than MIN_RETURN_OBSERVATIONS (5) returns
|
||||
let returns = vec![0.01, 0.02, 0.01, -0.01];
|
||||
let sharpe = RiskServiceImpl::compute_sharpe_ratio(&returns);
|
||||
assert!((sharpe - 0.0).abs() < 1e-10, "Should return 0.0 for insufficient data");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sharpe_zero_volatility() {
|
||||
let returns = vec![0.01; 10];
|
||||
let sharpe = RiskServiceImpl::compute_sharpe_ratio(&returns);
|
||||
assert!((sharpe - 0.0).abs() < 1e-10, "Zero volatility should return 0.0 Sharpe");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sharpe_positive_returns() {
|
||||
// Use enough data points with positive mean return
|
||||
let returns = vec![0.01, 0.02, 0.015, 0.005, 0.012, 0.008];
|
||||
let sharpe = RiskServiceImpl::compute_sharpe_ratio(&returns);
|
||||
// Mean return is clearly positive and above risk-free => Sharpe should be positive
|
||||
assert!(sharpe > 0.0, "Sharpe should be positive for consistently positive returns; got {}", sharpe);
|
||||
}
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// 7. compute_sortino_ratio
|
||||
// -----------------------------------------------------------------------
|
||||
#[test]
|
||||
fn test_sortino_insufficient_data() {
|
||||
let returns = vec![0.01, -0.01, 0.02];
|
||||
let sortino = RiskServiceImpl::compute_sortino_ratio(&returns);
|
||||
assert!((sortino - 0.0).abs() < 1e-10);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sortino_no_downside() {
|
||||
// All returns well above risk-free => no downside deviation => 0
|
||||
let daily_rf = RISK_FREE_RATE_ANNUAL / 252.0;
|
||||
let high_return = daily_rf + 0.01; // well above risk-free
|
||||
let returns = vec![high_return; 10];
|
||||
let sortino = RiskServiceImpl::compute_sortino_ratio(&returns);
|
||||
assert!((sortino - 0.0).abs() < 1e-10, "No downside should yield zero Sortino");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sortino_with_downside() {
|
||||
// Mix of positive and negative returns
|
||||
let returns = vec![-0.02, 0.03, -0.01, 0.02, -0.015, 0.01, 0.005];
|
||||
let sortino = RiskServiceImpl::compute_sortino_ratio(&returns);
|
||||
// Just verify it produces a finite number (not NaN/Inf)
|
||||
assert!(sortino.is_finite(), "Sortino should be finite; got {}", sortino);
|
||||
}
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// 8. VaR scaling: square-root-of-time rule
|
||||
// -----------------------------------------------------------------------
|
||||
#[test]
|
||||
fn test_var_scaling_sqrt_time() {
|
||||
// The service uses: var_5d = var_1d * sqrt(5), var_30d = var_1d * sqrt(30)
|
||||
let var_1d = 10_000.0_f64;
|
||||
let var_5d = var_1d * 5_f64.sqrt();
|
||||
let var_30d = var_1d * 30_f64.sqrt();
|
||||
|
||||
assert!((var_5d - 22_360.679).abs() < 1.0,
|
||||
"5d VaR should be ~22360.68; got {:.3}", var_5d);
|
||||
assert!((var_30d - 54_772.256).abs() < 1.0,
|
||||
"30d VaR should be ~54772.26; got {:.3}", var_30d);
|
||||
// Verify ordering: 1d < 5d < 30d
|
||||
assert!(var_1d < var_5d);
|
||||
assert!(var_5d < var_30d);
|
||||
}
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// 9. Concentration risk level classification
|
||||
// -----------------------------------------------------------------------
|
||||
#[test]
|
||||
fn test_concentration_risk_level_thresholds() {
|
||||
// The service classifies: >50% Critical, >30% High, >15% Medium, else Low
|
||||
let classify = |concentration: f64| -> RiskLevel {
|
||||
if concentration > 50.0 {
|
||||
RiskLevel::Critical
|
||||
} else if concentration > 30.0 {
|
||||
RiskLevel::High
|
||||
} else if concentration > 15.0 {
|
||||
RiskLevel::Medium
|
||||
} else {
|
||||
RiskLevel::Low
|
||||
}
|
||||
};
|
||||
|
||||
assert_eq!(classify(60.0) as i32, RiskLevel::Critical as i32);
|
||||
assert_eq!(classify(50.1) as i32, RiskLevel::Critical as i32);
|
||||
assert_eq!(classify(50.0) as i32, RiskLevel::High as i32);
|
||||
assert_eq!(classify(35.0) as i32, RiskLevel::High as i32);
|
||||
assert_eq!(classify(30.0) as i32, RiskLevel::Medium as i32);
|
||||
assert_eq!(classify(20.0) as i32, RiskLevel::Medium as i32);
|
||||
assert_eq!(classify(15.0) as i32, RiskLevel::Low as i32);
|
||||
assert_eq!(classify(5.0) as i32, RiskLevel::Low as i32);
|
||||
assert_eq!(classify(0.0) as i32, RiskLevel::Low as i32);
|
||||
}
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// 10. Constants are sane
|
||||
// -----------------------------------------------------------------------
|
||||
#[test]
|
||||
fn test_risk_constants() {
|
||||
assert!((RISK_FREE_RATE_ANNUAL - 0.05).abs() < 1e-10,
|
||||
"Risk-free rate should be 5%");
|
||||
assert_eq!(MIN_RETURN_OBSERVATIONS, 5,
|
||||
"Min return observations should be 5");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -142,6 +142,14 @@ path = "integration/checkpoint_roundtrip.rs"
|
||||
name = "feature_pipeline"
|
||||
path = "integration/feature_pipeline.rs"
|
||||
|
||||
[[test]]
|
||||
name = "ml_order_pipeline_test"
|
||||
path = "integration/ml_order_pipeline_test.rs"
|
||||
|
||||
[[test]]
|
||||
name = "risk_killswitch_test"
|
||||
path = "integration/risk_killswitch_test.rs"
|
||||
|
||||
[target.'cfg(target_os = "linux")'.dependencies]
|
||||
# Linux-specific performance monitoring
|
||||
perf-event = { version = "0.4", optional = true }
|
||||
|
||||
359
tests/integration/ml_order_pipeline_test.rs
Normal file
359
tests/integration/ml_order_pipeline_test.rs
Normal file
@@ -0,0 +1,359 @@
|
||||
//! Integration test: ML inference -> ensemble vote -> order generation
|
||||
//!
|
||||
//! Verifies the complete pipeline from feature vector through model
|
||||
//! inference, ensemble aggregation, and order signal generation.
|
||||
//! This is a critical path test for the ML -> Order execution pipeline.
|
||||
|
||||
use ml::ensemble::inference_adapter::{
|
||||
EnsemblePrediction, FeatureVector, ModelInferenceAdapter, PredictionMeta,
|
||||
};
|
||||
use ml::ensemble::inference_ensemble::InferenceEnsemble;
|
||||
use ml::MLResult;
|
||||
|
||||
/// Test adapter that simulates a bullish DQN model
|
||||
struct MockDQNAdapter;
|
||||
|
||||
impl ModelInferenceAdapter for MockDQNAdapter {
|
||||
fn model_name(&self) -> &str {
|
||||
"DQN-v1"
|
||||
}
|
||||
|
||||
fn predict(&self, _features: &FeatureVector) -> MLResult<EnsemblePrediction> {
|
||||
Ok(EnsemblePrediction {
|
||||
model_name: "DQN-v1".to_string(),
|
||||
direction: 0.8,
|
||||
confidence: 0.85,
|
||||
metadata: PredictionMeta::default(),
|
||||
})
|
||||
}
|
||||
|
||||
fn is_ready(&self) -> bool {
|
||||
true
|
||||
}
|
||||
}
|
||||
|
||||
/// Test adapter that simulates a bearish PPO model
|
||||
struct MockPPOAdapter;
|
||||
|
||||
impl ModelInferenceAdapter for MockPPOAdapter {
|
||||
fn model_name(&self) -> &str {
|
||||
"PPO-v1"
|
||||
}
|
||||
|
||||
fn predict(&self, _features: &FeatureVector) -> MLResult<EnsemblePrediction> {
|
||||
Ok(EnsemblePrediction {
|
||||
model_name: "PPO-v1".to_string(),
|
||||
direction: -0.3,
|
||||
confidence: 0.6,
|
||||
metadata: PredictionMeta::default(),
|
||||
})
|
||||
}
|
||||
|
||||
fn is_ready(&self) -> bool {
|
||||
true
|
||||
}
|
||||
}
|
||||
|
||||
/// Test adapter that returns an error (simulates model failure)
|
||||
struct FailingAdapter;
|
||||
|
||||
impl ModelInferenceAdapter for FailingAdapter {
|
||||
fn model_name(&self) -> &str {
|
||||
"FailingModel"
|
||||
}
|
||||
|
||||
fn predict(&self, _features: &FeatureVector) -> MLResult<EnsemblePrediction> {
|
||||
Err(ml::MLError::InferenceError(
|
||||
"Simulated model failure".to_string(),
|
||||
))
|
||||
}
|
||||
|
||||
fn is_ready(&self) -> bool {
|
||||
true
|
||||
}
|
||||
}
|
||||
|
||||
/// Test adapter that returns NaN direction and confidence
|
||||
struct NaNAdapter;
|
||||
|
||||
impl ModelInferenceAdapter for NaNAdapter {
|
||||
fn model_name(&self) -> &str {
|
||||
"NaN-model"
|
||||
}
|
||||
|
||||
fn predict(&self, _features: &FeatureVector) -> MLResult<EnsemblePrediction> {
|
||||
Ok(EnsemblePrediction {
|
||||
model_name: "NaN-model".to_string(),
|
||||
direction: f64::NAN,
|
||||
confidence: f64::NAN,
|
||||
metadata: PredictionMeta::default(),
|
||||
})
|
||||
}
|
||||
|
||||
fn is_ready(&self) -> bool {
|
||||
true
|
||||
}
|
||||
}
|
||||
|
||||
/// Test adapter that is never ready (simulates an unloaded model)
|
||||
struct NotReadyAdapter;
|
||||
|
||||
impl ModelInferenceAdapter for NotReadyAdapter {
|
||||
fn model_name(&self) -> &str {
|
||||
"NotReady"
|
||||
}
|
||||
|
||||
fn predict(&self, _features: &FeatureVector) -> MLResult<EnsemblePrediction> {
|
||||
Ok(EnsemblePrediction {
|
||||
model_name: "NotReady".to_string(),
|
||||
direction: 1.0,
|
||||
confidence: 1.0,
|
||||
metadata: PredictionMeta::default(),
|
||||
})
|
||||
}
|
||||
|
||||
fn is_ready(&self) -> bool {
|
||||
false
|
||||
}
|
||||
}
|
||||
|
||||
fn make_feature_vector() -> FeatureVector {
|
||||
FeatureVector {
|
||||
values: vec![0.1; 51],
|
||||
timestamp: 1_700_000_000_000_000,
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Test 1: Full ML -> Order pipeline
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[test]
|
||||
fn test_ml_to_order_pipeline() {
|
||||
// 1. Create a canonical 51-dim feature vector
|
||||
let features = make_feature_vector();
|
||||
assert_eq!(features.values.len(), 51, "Feature vector must be 51-dim");
|
||||
|
||||
// 2. Create ensemble with mock adapters (bullish DQN + bearish PPO)
|
||||
let adapters: Vec<Box<dyn ModelInferenceAdapter>> = vec![
|
||||
Box::new(MockDQNAdapter),
|
||||
Box::new(MockPPOAdapter),
|
||||
];
|
||||
let ensemble = InferenceEnsemble::new(adapters);
|
||||
|
||||
// 3. Verify both models are ready
|
||||
assert_eq!(ensemble.ready_count(), 2, "Both mock models should be ready");
|
||||
|
||||
// 4. Run ensemble prediction
|
||||
let prediction = ensemble.predict(&features);
|
||||
assert!(prediction.is_ok(), "Ensemble prediction should succeed");
|
||||
let pred = prediction.unwrap_or_else(|e| panic!("Prediction failed: {e}"));
|
||||
|
||||
// 5. Verify prediction properties are well-formed
|
||||
assert!(pred.direction.is_finite(), "Direction must be finite");
|
||||
assert!(pred.confidence.is_finite(), "Confidence must be finite");
|
||||
assert!(
|
||||
pred.confidence >= 0.0 && pred.confidence <= 1.0,
|
||||
"Confidence {} should be in [0.0, 1.0]",
|
||||
pred.confidence
|
||||
);
|
||||
assert!(
|
||||
pred.direction >= -1.0 && pred.direction <= 1.0,
|
||||
"Direction {} should be in [-1.0, 1.0]",
|
||||
pred.direction
|
||||
);
|
||||
|
||||
// 6. Generate order signal from prediction
|
||||
// The DQN model (dir=0.8, conf=0.85) dominates the PPO model (dir=-0.3, conf=0.6)
|
||||
// because higher confidence gives it more weight in the ensemble.
|
||||
// Expected net direction: positive (bullish).
|
||||
let order_side = if pred.direction > 0.0 { "Buy" } else { "Sell" };
|
||||
let order_size = (pred.confidence * 100.0).round();
|
||||
|
||||
assert_eq!(
|
||||
order_side, "Buy",
|
||||
"Net bullish ensemble (DQN dominates) should generate Buy, got direction={}",
|
||||
pred.direction
|
||||
);
|
||||
assert!(
|
||||
order_size > 0.0,
|
||||
"Order size should be positive, got {}",
|
||||
order_size
|
||||
);
|
||||
assert!(
|
||||
order_size <= 100.0,
|
||||
"Order size should be <= 100, got {}",
|
||||
order_size
|
||||
);
|
||||
|
||||
// 7. Verify model name reflects aggregation
|
||||
assert!(
|
||||
pred.model_name.contains("ENSEMBLE"),
|
||||
"Aggregated prediction model_name should contain 'ENSEMBLE', got '{}'",
|
||||
pred.model_name
|
||||
);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Test 2: Ensemble handles all-NaN models gracefully
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[test]
|
||||
fn test_ensemble_handles_all_models_returning_nan() {
|
||||
let adapters: Vec<Box<dyn ModelInferenceAdapter>> = vec![Box::new(NaNAdapter)];
|
||||
let ensemble = InferenceEnsemble::new(adapters);
|
||||
let features = make_feature_vector();
|
||||
|
||||
// The NaN circuit breaker should filter out the NaN model,
|
||||
// leaving zero successful predictions -> error.
|
||||
let result = ensemble.predict(&features);
|
||||
assert!(
|
||||
result.is_err(),
|
||||
"All-NaN ensemble should return error, but got: {:?}",
|
||||
result
|
||||
);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Test 3: Ensemble handles a mix of good + failing models
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[test]
|
||||
fn test_ensemble_survives_partial_model_failure() {
|
||||
let adapters: Vec<Box<dyn ModelInferenceAdapter>> = vec![
|
||||
Box::new(MockDQNAdapter),
|
||||
Box::new(FailingAdapter),
|
||||
];
|
||||
let ensemble = InferenceEnsemble::new(adapters);
|
||||
let features = make_feature_vector();
|
||||
|
||||
// The failing model is skipped; DQN alone should produce a valid prediction.
|
||||
let result = ensemble.predict(&features);
|
||||
assert!(
|
||||
result.is_ok(),
|
||||
"Ensemble with one good model should succeed, got: {:?}",
|
||||
result
|
||||
);
|
||||
|
||||
let pred = result.unwrap_or_else(|e| panic!("Prediction failed: {e}"));
|
||||
assert!(
|
||||
pred.direction.is_finite(),
|
||||
"Direction must be finite after partial failure"
|
||||
);
|
||||
assert!(
|
||||
pred.confidence.is_finite(),
|
||||
"Confidence must be finite after partial failure"
|
||||
);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Test 4: Ensemble with no ready models
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[test]
|
||||
fn test_ensemble_no_ready_models_returns_error() {
|
||||
let adapters: Vec<Box<dyn ModelInferenceAdapter>> = vec![Box::new(NotReadyAdapter)];
|
||||
let ensemble = InferenceEnsemble::new(adapters);
|
||||
let features = make_feature_vector();
|
||||
|
||||
let result = ensemble.predict(&features);
|
||||
assert!(
|
||||
result.is_err(),
|
||||
"Ensemble with no ready models should return error"
|
||||
);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Test 5: Ensemble with custom weights changes outcome
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[test]
|
||||
fn test_ensemble_custom_weights_affect_direction() {
|
||||
// Two opposing models with equal confidence
|
||||
struct BullAdapter;
|
||||
impl ModelInferenceAdapter for BullAdapter {
|
||||
fn model_name(&self) -> &str { "Bull" }
|
||||
fn predict(&self, _: &FeatureVector) -> MLResult<EnsemblePrediction> {
|
||||
Ok(EnsemblePrediction {
|
||||
model_name: "Bull".to_string(),
|
||||
direction: 1.0,
|
||||
confidence: 0.8,
|
||||
metadata: PredictionMeta::default(),
|
||||
})
|
||||
}
|
||||
fn is_ready(&self) -> bool { true }
|
||||
}
|
||||
|
||||
struct BearAdapter;
|
||||
impl ModelInferenceAdapter for BearAdapter {
|
||||
fn model_name(&self) -> &str { "Bear" }
|
||||
fn predict(&self, _: &FeatureVector) -> MLResult<EnsemblePrediction> {
|
||||
Ok(EnsemblePrediction {
|
||||
model_name: "Bear".to_string(),
|
||||
direction: -1.0,
|
||||
confidence: 0.8,
|
||||
metadata: PredictionMeta::default(),
|
||||
})
|
||||
}
|
||||
fn is_ready(&self) -> bool { true }
|
||||
}
|
||||
|
||||
let adapters: Vec<Box<dyn ModelInferenceAdapter>> = vec![
|
||||
Box::new(BullAdapter),
|
||||
Box::new(BearAdapter),
|
||||
];
|
||||
|
||||
// Without custom weights, equal confidence => direction ~ 0.0
|
||||
let ensemble_equal = InferenceEnsemble::new(adapters);
|
||||
let features = make_feature_vector();
|
||||
let pred_equal = ensemble_equal
|
||||
.predict(&features)
|
||||
.unwrap_or_else(|e| panic!("Equal-weight prediction failed: {e}"));
|
||||
assert!(
|
||||
pred_equal.direction.abs() < 0.01,
|
||||
"Equal weight/confidence opposing models should cancel out, got {}",
|
||||
pred_equal.direction
|
||||
);
|
||||
|
||||
// With Bull weighted 3x heavier, direction should be strongly positive
|
||||
let adapters2: Vec<Box<dyn ModelInferenceAdapter>> = vec![
|
||||
Box::new(BullAdapter),
|
||||
Box::new(BearAdapter),
|
||||
];
|
||||
let mut ensemble_weighted = InferenceEnsemble::new(adapters2);
|
||||
ensemble_weighted.set_weight("Bull", 3.0);
|
||||
ensemble_weighted.set_weight("Bear", 1.0);
|
||||
|
||||
let pred_weighted = ensemble_weighted
|
||||
.predict(&features)
|
||||
.unwrap_or_else(|e| panic!("Weighted prediction failed: {e}"));
|
||||
assert!(
|
||||
pred_weighted.direction > 0.3,
|
||||
"Bull-weighted ensemble should have positive direction, got {}",
|
||||
pred_weighted.direction
|
||||
);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Test 6: Order sizing from confidence
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[test]
|
||||
fn test_order_sizing_from_confidence() {
|
||||
// Verify that different confidence levels produce proportional order sizes
|
||||
for (conf, expected_min, expected_max) in [
|
||||
(0.0, 0.0, 0.0),
|
||||
(0.5, 49.0, 51.0),
|
||||
(1.0, 99.0, 101.0),
|
||||
] {
|
||||
let size = (conf * 100.0_f64).round();
|
||||
assert!(
|
||||
size >= expected_min && size <= expected_max,
|
||||
"Confidence {} -> size {}, expected [{}, {}]",
|
||||
conf,
|
||||
size,
|
||||
expected_min,
|
||||
expected_max
|
||||
);
|
||||
}
|
||||
}
|
||||
339
tests/integration/risk_killswitch_test.rs
Normal file
339
tests/integration/risk_killswitch_test.rs
Normal file
@@ -0,0 +1,339 @@
|
||||
//! Integration test: Risk limit violation -> Kill switch activation
|
||||
//!
|
||||
//! Verifies that risk safety mechanisms properly trigger under stress,
|
||||
//! correctly scope kill switch activations, and block trading when active.
|
||||
//! This is a critical path test for the Risk -> Kill Switch pipeline.
|
||||
|
||||
use risk::safety::kill_switch::{AtomicKillSwitch, TradingGate};
|
||||
use risk::safety::KillSwitchConfig;
|
||||
use risk::risk_types::KillSwitchScope;
|
||||
|
||||
fn create_test_kill_switch() -> AtomicKillSwitch {
|
||||
let config = KillSwitchConfig::default();
|
||||
AtomicKillSwitch::new_test(config)
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Test 1: Kill switch activation and trading blocking
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_kill_switch_activation_and_blocking() {
|
||||
let kill_switch = create_test_kill_switch();
|
||||
|
||||
// Initially not triggered
|
||||
assert!(
|
||||
!kill_switch.is_triggered(),
|
||||
"Kill switch should start inactive"
|
||||
);
|
||||
assert!(
|
||||
kill_switch.is_trading_allowed(&KillSwitchScope::Global),
|
||||
"Trading should be allowed initially"
|
||||
);
|
||||
|
||||
// Trigger the kill switch via global activation
|
||||
let result = kill_switch
|
||||
.activate_global(
|
||||
"Max drawdown exceeded: -5.2%".to_string(),
|
||||
"risk_monitor".to_string(),
|
||||
)
|
||||
.await;
|
||||
assert!(result.is_ok(), "Triggering kill switch should succeed");
|
||||
|
||||
// Now it should be triggered
|
||||
assert!(
|
||||
kill_switch.is_triggered(),
|
||||
"Kill switch should be active after trigger"
|
||||
);
|
||||
|
||||
// Trading should be blocked
|
||||
assert!(
|
||||
!kill_switch.is_trading_allowed(&KillSwitchScope::Global),
|
||||
"Trading must be blocked after global kill switch trigger"
|
||||
);
|
||||
|
||||
// Also blocked for any scoped query (global takes precedence)
|
||||
assert!(
|
||||
!kill_switch.is_trading_allowed(&KillSwitchScope::Symbol("ES.FUT".to_string())),
|
||||
"Symbol-scoped trading must also be blocked by global kill switch"
|
||||
);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Test 2: Scoped kill switch (symbol-level)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_scoped_kill_switch_symbol() {
|
||||
let kill_switch = create_test_kill_switch();
|
||||
|
||||
// Trigger for a specific symbol scope only
|
||||
let result = kill_switch
|
||||
.engage(
|
||||
KillSwitchScope::Symbol("ES.FUT".to_string()),
|
||||
"Symbol-level risk limit breached".to_string(),
|
||||
"symbol_monitor".to_string(),
|
||||
false, // no cascade
|
||||
)
|
||||
.await;
|
||||
assert!(result.is_ok(), "Scoped trigger should succeed");
|
||||
|
||||
// Global should NOT be triggered
|
||||
assert!(
|
||||
!kill_switch.is_triggered(),
|
||||
"Global kill switch should NOT be triggered by symbol-scoped engagement"
|
||||
);
|
||||
assert!(
|
||||
kill_switch.is_trading_allowed(&KillSwitchScope::Global),
|
||||
"Global trading should still be allowed"
|
||||
);
|
||||
|
||||
// The specific symbol should be blocked
|
||||
assert!(
|
||||
!kill_switch.is_trading_allowed(&KillSwitchScope::Symbol("ES.FUT".to_string())),
|
||||
"ES.FUT trading should be blocked"
|
||||
);
|
||||
|
||||
// Other symbols should NOT be blocked
|
||||
assert!(
|
||||
kill_switch.is_trading_allowed(&KillSwitchScope::Symbol("NQ.FUT".to_string())),
|
||||
"NQ.FUT trading should still be allowed"
|
||||
);
|
||||
|
||||
// Health metrics should be finite
|
||||
let (error_rate, _failures) = kill_switch.get_health_metrics();
|
||||
assert!(
|
||||
error_rate.is_finite(),
|
||||
"Error rate should be finite, got {}",
|
||||
error_rate
|
||||
);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Test 3: Kill switch reset restores trading
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_kill_switch_reset_restores_trading() {
|
||||
let kill_switch = create_test_kill_switch();
|
||||
|
||||
// Trigger globally
|
||||
kill_switch.trigger();
|
||||
assert!(kill_switch.is_triggered());
|
||||
assert!(!kill_switch.is_trading_allowed(&KillSwitchScope::Global));
|
||||
|
||||
// Reset
|
||||
let result = kill_switch.reset(Some(KillSwitchScope::Global)).await;
|
||||
assert!(result.is_ok(), "Reset should succeed");
|
||||
|
||||
// Trading should be restored
|
||||
assert!(!kill_switch.is_triggered(), "Kill switch should be cleared");
|
||||
assert!(
|
||||
kill_switch.is_trading_allowed(&KillSwitchScope::Global),
|
||||
"Trading should be allowed after reset"
|
||||
);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Test 4: Multiple scoped activations and selective reset
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_multiple_scoped_activations_and_selective_reset() {
|
||||
let kill_switch = create_test_kill_switch();
|
||||
|
||||
// Engage two different scopes
|
||||
kill_switch
|
||||
.engage(
|
||||
KillSwitchScope::Symbol("AAPL".to_string()),
|
||||
"AAPL halt".to_string(),
|
||||
"user".to_string(),
|
||||
false,
|
||||
)
|
||||
.await
|
||||
.expect("AAPL engage should succeed");
|
||||
|
||||
kill_switch
|
||||
.engage(
|
||||
KillSwitchScope::Account("ACC-001".to_string()),
|
||||
"Account risk limit".to_string(),
|
||||
"user".to_string(),
|
||||
false,
|
||||
)
|
||||
.await
|
||||
.expect("Account engage should succeed");
|
||||
|
||||
// Both should be blocked
|
||||
assert!(
|
||||
!kill_switch.is_trading_allowed(&KillSwitchScope::Symbol("AAPL".to_string())),
|
||||
"AAPL should be blocked"
|
||||
);
|
||||
assert!(
|
||||
!kill_switch.is_trading_allowed(&KillSwitchScope::Account("ACC-001".to_string())),
|
||||
"ACC-001 should be blocked"
|
||||
);
|
||||
|
||||
// Global still allowed
|
||||
assert!(
|
||||
kill_switch.is_trading_allowed(&KillSwitchScope::Global),
|
||||
"Global should still be allowed"
|
||||
);
|
||||
|
||||
// Reset only the symbol scope
|
||||
kill_switch
|
||||
.reset(Some(KillSwitchScope::Symbol("AAPL".to_string())))
|
||||
.await
|
||||
.expect("AAPL reset should succeed");
|
||||
|
||||
// AAPL should be restored, account still blocked
|
||||
assert!(
|
||||
kill_switch.is_trading_allowed(&KillSwitchScope::Symbol("AAPL".to_string())),
|
||||
"AAPL should be allowed after reset"
|
||||
);
|
||||
assert!(
|
||||
!kill_switch.is_trading_allowed(&KillSwitchScope::Account("ACC-001".to_string())),
|
||||
"ACC-001 should remain blocked"
|
||||
);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Test 5: Cascade behavior (portfolio -> strategies)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_cascade_portfolio_to_strategies() {
|
||||
let kill_switch = create_test_kill_switch();
|
||||
|
||||
// Trigger a portfolio-level kill switch with cascade=true
|
||||
kill_switch
|
||||
.engage(
|
||||
KillSwitchScope::Portfolio("portfolio1".to_string()),
|
||||
"Portfolio drawdown exceeded".to_string(),
|
||||
"risk_engine".to_string(),
|
||||
true, // cascade
|
||||
)
|
||||
.await
|
||||
.expect("Portfolio engage should succeed");
|
||||
|
||||
// The portfolio itself should be blocked
|
||||
assert!(
|
||||
!kill_switch.is_trading_allowed(&KillSwitchScope::Portfolio("portfolio1".to_string())),
|
||||
"Portfolio should be blocked"
|
||||
);
|
||||
|
||||
// The system is active
|
||||
let is_active = kill_switch.is_active().await.expect("is_active should work");
|
||||
assert!(is_active, "Kill switch should report as active");
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Test 6: Metrics tracking
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_kill_switch_metrics_tracking() {
|
||||
let kill_switch = create_test_kill_switch();
|
||||
|
||||
// Initial metrics should be zero
|
||||
let (checks_before, commands_before) = kill_switch.get_metrics();
|
||||
assert_eq!(checks_before, 0, "Initial health checks should be 0");
|
||||
assert_eq!(commands_before, 0, "Initial commands should be 0");
|
||||
|
||||
// Perform operations that increment counters
|
||||
kill_switch
|
||||
.engage(
|
||||
KillSwitchScope::Global,
|
||||
"Test".to_string(),
|
||||
"user".to_string(),
|
||||
false,
|
||||
)
|
||||
.await
|
||||
.expect("engage should succeed");
|
||||
|
||||
kill_switch
|
||||
.reset(Some(KillSwitchScope::Global))
|
||||
.await
|
||||
.expect("reset should succeed");
|
||||
|
||||
// Commands should have incremented (engage + reset = 2)
|
||||
let (_checks_after, commands_after) = kill_switch.get_metrics();
|
||||
assert_eq!(
|
||||
commands_after, 2,
|
||||
"Two commands (engage + reset) should be tracked, got {}",
|
||||
commands_after
|
||||
);
|
||||
|
||||
// Health metrics: no failures (no Redis in test mode)
|
||||
let (error_rate, failures) = kill_switch.get_health_metrics();
|
||||
assert_eq!(error_rate, 0.0, "Error rate should be 0.0 with no Redis");
|
||||
assert_eq!(failures, 0, "Failures should be 0 with no Redis");
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Test 7: Trading gate integration
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_trading_gate_lifecycle() {
|
||||
let gate = TradingGate::new(true);
|
||||
assert!(gate.is_open(), "Gate should start open");
|
||||
|
||||
// Close gate (simulating risk event)
|
||||
gate.close();
|
||||
assert!(!gate.is_open(), "Gate should be closed");
|
||||
|
||||
// Reopen gate (simulating risk clearance)
|
||||
gate.open();
|
||||
assert!(gate.is_open(), "Gate should be reopened");
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Test 8: Kill switch health check (no Redis = healthy)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_kill_switch_health_in_test_mode() {
|
||||
let kill_switch = create_test_kill_switch();
|
||||
|
||||
let healthy = kill_switch
|
||||
.is_healthy()
|
||||
.await
|
||||
.expect("is_healthy should succeed");
|
||||
assert!(
|
||||
healthy,
|
||||
"Kill switch without Redis should report healthy (test mode)"
|
||||
);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Test 9: Deactivate restores scoped trading
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_deactivate_restores_scoped_trading() {
|
||||
let kill_switch = create_test_kill_switch();
|
||||
|
||||
let scope = KillSwitchScope::Strategy("momentum_v2".to_string());
|
||||
|
||||
// Activate
|
||||
kill_switch
|
||||
.activate(scope.clone(), "Test halt".to_string(), "user".to_string(), false)
|
||||
.await
|
||||
.expect("activate should succeed");
|
||||
|
||||
assert!(
|
||||
!kill_switch.is_trading_allowed(&scope),
|
||||
"Strategy should be blocked after activation"
|
||||
);
|
||||
|
||||
// Deactivate
|
||||
kill_switch
|
||||
.deactivate(scope.clone(), "user".to_string())
|
||||
.await
|
||||
.expect("deactivate should succeed");
|
||||
|
||||
assert!(
|
||||
kill_switch.is_trading_allowed(&scope),
|
||||
"Strategy should be allowed after deactivation"
|
||||
);
|
||||
}
|
||||
Reference in New Issue
Block a user