foxhunt/docs/plans/2026-03-07-ml-crate-split-design.md

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):

// 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