# 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