092b1d2a24
Split ml-rl into ml-dqn and ml-ppo for 3-way parallel compilation and better sccache hit rate. Extract TradingAction to ml-core to enable full DQN/PPO independence. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
245 lines
10 KiB
Markdown
245 lines
10 KiB
Markdown
# ML Crate Split — Design Document
|
|
|
|
## Problem
|
|
|
|
The `ml` crate is 260K lines / 530 files. It compiles as a single `rustc` invocation taking **55.1s** — 69% of the total workspace build. Only 4 of 16 CPU cores are utilized during this phase. Incremental builds after touching any file in `ml/src/` recompile all 260K lines.
|
|
|
|
## Goal
|
|
|
|
Split `ml` into sub-crates that compile in parallel, reducing:
|
|
- Clean build time from ~55s to ~35-40s
|
|
- Incremental rebuild after editing a model from ~55s to ~10-15s
|
|
|
|
## Architecture
|
|
|
|
```
|
|
ml-core (shared types, traits, utilities, features, data)
|
|
├─→ ml-dqn (DQN + cuda_pipeline) ── parallel ──┐
|
|
├─→ ml-ppo (PPO) ── parallel ──┤
|
|
├─→ ml-supervised (TFT, Mamba, Liquid, etc.) ── parallel ──┤
|
|
│ ↓
|
|
└──────────────────────────→ ml-infra (hyperopt, ensemble, trainers)
|
|
↓
|
|
ml (thin facade, re-exports)
|
|
```
|
|
|
|
`ml-dqn`, `ml-ppo`, and `ml-supervised` compile in **3-way parallel** — they have zero dependencies on each other. This also maximizes sccache hit rate in CI: editing DQN never invalidates PPO's cache entry.
|
|
|
|
## Crate Assignment (Every Module Accounted For)
|
|
|
|
### ml-core (~50K lines)
|
|
|
|
Shared infrastructure that 2+ other sub-crates depend on.
|
|
|
|
**From lib.rs (types defined inline):**
|
|
- `Trade`, `MarketDataSnapshot`, `FeatureVector`, `IntegerTensor`, `UpdateSummary`
|
|
- `HealthStatus`, `MarketRegime`, `CommonTypeError`, `CommonError`
|
|
- `MLError`, `MLResult`, `UnifiedMLResult`
|
|
- `ModelPrediction`, `InferenceResult`, `ModelMetadata`, `Features`, `MLModel` (traits)
|
|
- `PRECISION_FACTOR`, `MAX_INFERENCE_LATENCY_US`
|
|
- `get_training_device()`, `get_training_device_at()`
|
|
- `HFTPerformanceProfile`, `ParallelExecutor`, `LatencyOptimizer`, `ModelRegistry` helpers
|
|
- All `pub use` re-exports (`ErrorCategory`, `Adam`, `ModelType`)
|
|
|
|
**Modules (moved as-is):**
|
|
- `error.rs` (66 lines)
|
|
- `types/` (49 lines) — `OHLCVBar` etc.
|
|
- `common/` (1,435 lines) — `FactoredAction`, `ExposureLevel`, `CircuitBreaker`
|
|
- `config/` (818 lines)
|
|
- `cuda_compat.rs` (511 lines)
|
|
- `tensor_ops.rs` (143 lines)
|
|
- `traits.rs` (199 lines)
|
|
- `model.rs` (4 lines)
|
|
- `optimizers/` (261 lines) — `Adam`
|
|
- `gradient_accumulation.rs` (274 lines)
|
|
- `gradient_utils.rs` (169 lines)
|
|
- `gpu/` (836 lines) — capabilities, `DeviceConfig`, `memory_profile`
|
|
- `safety/` (6,467 lines) — `MLSafetyManager`, tensor safety
|
|
- `security/` (1,407 lines)
|
|
- `memory_optimization/` (4,576 lines) — quantization (used by TFT)
|
|
- `checkpoint/` (7,867 lines) — `Checkpointable` trait, storage
|
|
- `training/` (1,565 lines) — `UnifiedTrainable`, `CheckpointMetadata`, `TrainingMetrics`
|
|
- `training.rs` (85 lines)
|
|
- `features/` (20,256 lines) — feature extraction, normalization, OFI, MBP-10
|
|
- `feature_cache.rs` (349 lines)
|
|
- `preprocessing.rs` (754 lines)
|
|
- `data_loader.rs` (713 lines)
|
|
- `data_loaders/` (4,365 lines) — `DbnSequenceLoader`, calibration
|
|
- `data_pipeline/` (1,512 lines)
|
|
- `data_validation/` (3,529 lines)
|
|
- `evaluation/` (1,011 lines) — `EvaluationEngine`, `PerformanceMetrics`
|
|
- `metrics/` (420 lines)
|
|
- `performance.rs` (481 lines)
|
|
- `observability/` (1,193 lines)
|
|
|
|
**Extracted from dqn/ → ml-core shared modules:**
|
|
- `dqn/mixed_precision.rs` (722 lines) → `ml-core/src/mixed_precision.rs`
|
|
- `dqn/xavier_init.rs` (200 lines) → `ml-core/src/xavier_init.rs`
|
|
- `dqn/portfolio_tracker.rs` (931 lines) → `ml-core/src/portfolio_tracker.rs`
|
|
- `dqn/action_space.rs` (245 lines) → `ml-core/src/action_space.rs`
|
|
- `dqn/order_router.rs` (175 lines) → `ml-core/src/order_router.rs`
|
|
|
|
**Extracted from dqn/agent.rs → ml-core:**
|
|
- `TradingAction` enum (Buy/Sell/Hold) — used by PPO tests, must be in ml-core for DQN/PPO independence
|
|
|
|
**NOT extracting:**
|
|
- `dqn/circuit_breaker.rs` — already a re-export of `common/circuit_breaker.rs` (which stays in ml-core)
|
|
- `liquid/FixedPoint` — only used by mamba (both in ml-supervised), no extraction needed
|
|
|
|
### ml-dqn (~35K lines)
|
|
|
|
DQN reinforcement learning. Depends on: `ml-core`.
|
|
|
|
- `dqn/` (33,116 lines minus ~2,500 extracted = ~30,600 lines)
|
|
- `cuda_pipeline/` (4,431 lines) — GPU data pipeline for DQN trainers
|
|
|
|
### ml-ppo (~14K lines)
|
|
|
|
PPO reinforcement learning. Depends on: `ml-core`. **Zero dependency on ml-dqn.**
|
|
|
|
- `ppo/` (14,182 lines)
|
|
|
|
After extraction of shared items to ml-core, PPO's only remaining DQN import was `TradingAction` (test-only) — now in ml-core.
|
|
|
|
### Import changes (ml-dqn and ml-ppo):
|
|
- `crate::dqn::mixed_precision::training_dtype` → `ml_core::mixed_precision::training_dtype`
|
|
- `crate::dqn::xavier_init::linear_xavier` → `ml_core::xavier_init::linear_xavier`
|
|
- `crate::dqn::portfolio_tracker` → `ml_core::portfolio_tracker`
|
|
- `crate::dqn::action_space` → `ml_core::action_space`
|
|
- `crate::dqn::order_router` → `ml_core::order_router`
|
|
- `crate::dqn::circuit_breaker` → `ml_core::common::circuit_breaker`
|
|
- All `crate::MLError` → `ml_core::MLError`
|
|
- All `crate::training::unified_trainer` → `ml_core::training::unified_trainer`
|
|
|
|
### ml-supervised (~36K lines)
|
|
|
|
Supervised/generative models. Depends on: `ml-core`. **Zero dependency on ml-rl.**
|
|
|
|
- `tft/` (11,440 lines)
|
|
- `mamba/` (6,465 lines)
|
|
- `liquid/` (5,228 lines) — includes `FixedPoint`
|
|
- `tgnn/` (4,410 lines)
|
|
- `tlob/` (2,189 lines)
|
|
- `kan/` (1,397 lines)
|
|
- `xlstm/` (1,404 lines)
|
|
- `diffusion/` (1,317 lines)
|
|
- `transformers/` (821 lines) — shared transformer utilities
|
|
- `flash_attention/` (1,225 lines) — shared attention implementation
|
|
|
|
**Import changes:**
|
|
- Same pattern as ml-rl: `crate::X` → `ml_core::X` for shared items
|
|
- `crate::liquid::FixedPoint` stays as `crate::liquid::FixedPoint` (both in same crate)
|
|
- `crate::dqn::mixed_precision::training_dtype` → `ml_core::mixed_precision::training_dtype`
|
|
|
|
### ml-infra (~67K lines)
|
|
|
|
Training infrastructure. Depends on: `ml-core`, `ml-dqn`, `ml-ppo`, `ml-supervised`.
|
|
|
|
- `hyperopt/` (19,934 lines) — Bayesian optimization, adapters for all 10 models
|
|
- `ensemble/` (13,744 lines) — inference adapters, confidence, voting
|
|
- `trainers/` (17,772 lines) — `DQNTrainer`, `PPOTrainer`, unified training loop
|
|
- `benchmark/` (5,969 lines) — model benchmarks
|
|
- `deployment/` (8,190 lines) — A/B testing, hot-swap, versioning
|
|
- `training_pipeline.rs` (847 lines)
|
|
- `walk_forward.rs` (485 lines)
|
|
|
|
**Import changes:**
|
|
- `crate::dqn::DQN` → `ml_dqn::dqn::DQN`
|
|
- `crate::ppo::PPO` → `ml_ppo::ppo::PPO`
|
|
- `crate::tft::TemporalFusionTransformer` → `ml_supervised::tft::TemporalFusionTransformer`
|
|
- (and so on for all model types)
|
|
- Shared items: `crate::X` → `ml_core::X`
|
|
|
|
### ml (thin facade, ~30K lines)
|
|
|
|
Re-exports everything from sub-crates. Contains modules that cross sub-crate boundaries or are only used externally.
|
|
|
|
**Modules (moved here):**
|
|
- `inference.rs` (1,677 lines) — depends on TFT (supervised) + safety (core)
|
|
- `inference_validator.rs` (528 lines)
|
|
- `model_factory.rs` (339 lines) — creates any model type
|
|
- `model_loader_integration.rs` (725 lines)
|
|
- `model_registry/` (558 lines) + `model_registry.rs` (1,350 lines)
|
|
- `registry/` (527 lines)
|
|
- `bridge.rs` (356 lines)
|
|
- `batch_processing.rs` (710 lines)
|
|
- `benchmarks.rs` (733 lines)
|
|
- `examples.rs` (892 lines)
|
|
- `portfolio_transformer.rs` (814 lines)
|
|
- `validation/` (4,337 lines) — depends on DQN + PPO types
|
|
- `risk/` (5,023 lines) — ML risk (Kelly, position sizing)
|
|
- `regime/` (4,943 lines) + `regime_detection/` (1,196 lines)
|
|
- `microstructure/` (2,598 lines)
|
|
- `labeling/` (3,340 lines)
|
|
- `universe/` (1,731 lines)
|
|
- `asset_selection/` (1,184 lines)
|
|
- `backtesting/` (1,120 lines)
|
|
- `paper_trading/` (389 lines)
|
|
- `stress_testing/` (1,254 lines)
|
|
- `integration/` (3,678 lines)
|
|
- `tests/` (1,844 lines)
|
|
- `gpu_benchmarks/` (21 lines)
|
|
- `cuda_common/` (0 lines — empty)
|
|
|
|
**Re-exports (preserves external API):**
|
|
```rust
|
|
// ml/src/lib.rs
|
|
pub use ml_core::*; // All types, errors, traits
|
|
pub use ml_dqn::dqn; // DQN module
|
|
pub use ml_dqn::cuda_pipeline; // GPU pipeline
|
|
pub use ml_ppo::ppo; // PPO module
|
|
pub use ml_supervised::tft; // All supervised models
|
|
pub use ml_supervised::mamba;
|
|
// ... etc for all model modules
|
|
pub use ml_infra::hyperopt;
|
|
pub use ml_infra::ensemble;
|
|
pub use ml_infra::trainers;
|
|
// ... etc
|
|
```
|
|
|
|
## Feature Flags
|
|
|
|
| Feature | Crate | Reason |
|
|
|---------|-------|--------|
|
|
| `cuda` | `ml-core` (propagated) | candle-core/cuda, candle-nn/cuda |
|
|
| `minimal-inference` | `ml` facade | controls which sub-crates are included |
|
|
| `financial` | `ml` facade | feature gate for services |
|
|
| `s3-storage` | `ml-core` | AWS SDK deps |
|
|
| `high-precision` | `ml-core` | rust_decimal feature |
|
|
| `mimalloc-allocator` | `ml-core` | allocator |
|
|
| `simd` | `ml-core` | SIMD features |
|
|
| `gc` | `ml-core` | GC features |
|
|
| `nccl` | `ml-dqn` | multi-GPU for RL training |
|
|
|
|
## External Consumer Migration
|
|
|
|
All external consumers (`trading_service`, `ml_training_service`, etc.) continue to depend on `ml` (the facade). The facade re-exports everything, so **zero import changes** in downstream crates. The workspace `Cargo.toml` entry stays: `ml = { path = "crates/ml" }`.
|
|
|
|
## Projected Compile Times
|
|
|
|
| Phase | Time | Parallelism |
|
|
|-------|------|-------------|
|
|
| ml-core | ~15s | serial (first) |
|
|
| ml-dqn + ml-ppo + ml-supervised | ~12s | **3-way parallel** |
|
|
| ml-infra | ~18s | serial (after models) |
|
|
| ml facade | ~3s | serial (last) |
|
|
| **Total clean** | **~48s** | vs 55s current |
|
|
| **Edit in dqn (incremental)** | **~12s** | vs 55s current |
|
|
| **Edit in ppo (incremental)** | **~6s** | vs 55s current |
|
|
| **Edit in tft (incremental)** | **~10s** | vs 55s current |
|
|
| **Edit in features (incremental)** | **~10s** | vs 55s current |
|
|
| **CI sccache (PPO edit)** | **~6s** | only ml-ppo rebuilt |
|
|
|
|
## Risks
|
|
|
|
1. **Circular dependencies**: `validation/` imports both DQN and PPO types → placed in facade
|
|
2. **Test placement**: Tests that cross model boundaries stay in the facade crate
|
|
3. **Feature flag propagation**: `cuda` must be forwarded from facade through sub-crates
|
|
4. **Re-export completeness**: Must verify every `use ml::X` in the workspace still resolves
|
|
|
|
## Non-Goals
|
|
|
|
- No stubs — every function is wired to its real implementation
|
|
- No new functionality — pure mechanical restructuring
|
|
- No API changes for external consumers
|