Files
foxhunt/docs/plans/2026-03-08-ml-crate-split-phase2.md
jgrusewski 4709ca8bc2 feat(dqn): enable Branching DQN with 45 factored actions (5×3×3)
Restore 45-action factored space via Branching DQN (Tavakoli 2018),
outputting 11 Q-values (5+3+3) instead of 45. This was reduced to 5
exposure-only actions during debugging and was never intended as permanent.

- Enable use_branching: true by default in DQNConfig and DQNHyperparameters
- Add branching paths to select_action_with_confidence and select_action_inference
- Update agent.rs select_action_factored for branching-aware selection
- Expand CountBonus to per-branch tracking with bonuses_branched()
- Add order_type + urgency distribution tracking in monitoring
- Add DQN_ORDER_ACTIONS=3, DQN_URGENCY_ACTIONS=3, DQN_TOTAL_ACTIONS=45 to CUDA header
- Fix 7 pre-existing clippy doc_markdown errors in regime_conditional.rs
- Fix pre-existing cognitive_complexity in replay_buffer_type.rs (extract helpers)
- Fix flaky GPU test OOM under parallel execution (CPU fallback + test VRAM safety)
- Delete unused flash_attention submodules (block_sparse, causal_masking, etc.)
- Add GPU hot-path guard scripts and ensemble/hyperopt adapter improvements

Tests: ml-dqn 416/0, ml 905/0, clippy 0 errors on both crates

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-11 22:00:13 +01:00

18 KiB

ML Crate Split Phase 2 — Implementation Plan

For Claude: REQUIRED SUB-SKILL: Use superpowers:executing-plans to implement this plan task-by-task.

Goal: Reduce ml crate from 91K to ~12K LOC by extracting trainers, hyperopt adapters, model impls, and infrastructure into existing sub-crates.

Architecture: Extract in dependency order: foundational modules first (cuda_pipeline, features, flash_attention), then model impls, then trainers (which depend on foundations), then hyperopt/ensemble adapters (which depend on trainers). ml becomes an orchestration layer owning inference, model factory, and training pipeline.

Tech Stack: Rust workspace, Cargo.toml workspace deps, pub use re-exports for backward compat

Worktree: .worktrees/ml-split-p2 on branch feature/ml-crate-split-phase2

Build command: SQLX_OFFLINE=true cargo check --workspace (no DB required) Test command: SQLX_OFFLINE=true cargo test -p <crate> --lib Clippy: SQLX_OFFLINE=true cargo clippy --workspace -- -D warnings


Dependency Order & Parallelism

[Parallel Group A]          [Parallel Group B]       [Parallel Group C]
Task 1: cuda_pipeline→core  Task 3: model impls→     Task 4: benchmark→
Task 1b: flash_attn→core       sub-crates               ml-benchmark
Task 1c: features→ml-feat   Task 3b: dqn/→ml-dqn     Task 4b: deployment→
Task 1d: microstructure→     Task 3c: ppo/→ml-ppo        ml-deployment
    ml-features              Task 3d: model stubs→
                                ml-supervised
         │                        │
         ▼                        ▼
[Sequential: Task 2]     [Sequential: Task 5]
data_loaders→ml-data      trainers→model sub-crates
data_pipeline→ml-data      (depends on A + B)
training/→ml-data               │
                                ▼
                        [Sequential: Task 6]
                        hyperopt adapters→ml-hyperopt
                        ensemble adapters→ml-ensemble
                        (depends on Task 5)
                                │
                                ▼
                        [Final: Task 7]
                        Cleanup ml lib.rs,
                        update services, fix all imports

Task 1: Foundation Extractions → ml-core & ml-features

Parallel subtasks — no interdependencies.

Task 1a: cuda_pipeline/ → ml-core (4,084 LOC)

Files to move:

  • crates/ml/src/cuda_pipeline/mod.rs (1,122)
  • crates/ml/src/cuda_pipeline/gpu_experience_collector.rs (795)
  • crates/ml/src/cuda_pipeline/gpu_ppo_collector.rs (662)
  • crates/ml/src/cuda_pipeline/gpu_weights.rs (887)
  • crates/ml/src/cuda_pipeline/gpu_portfolio.rs (254)
  • crates/ml/src/cuda_pipeline/double_buffer.rs (161)
  • crates/ml/src/cuda_pipeline/multi_gpu.rs (203)

Steps:

  1. Create crates/ml-core/src/cuda_pipeline/ directory
  2. Move all 7 files, update mod.rs module paths
  3. Replace use crate::MLErroruse crate::error::MLError (already in ml-core)
  4. Replace use crate::dqn::mixed_precision::training_dtypeuse crate::mixed_precision::training_dtype (already in ml-core)
  5. Add pub mod cuda_pipeline; to crates/ml-core/src/lib.rs
  6. In crates/ml/src/lib.rs: replace pub mod cuda_pipeline; with pub use ml_core::cuda_pipeline;
  7. Run: SQLX_OFFLINE=true cargo check -p ml-core -p ml
  8. Commit: refactor(ml): move cuda_pipeline to ml-core

Task 1b: flash_attention/ → ml-core (1,225 LOC)

Files to move:

  • crates/ml/src/flash_attention/*.rs

Steps:

  1. Create crates/ml-core/src/flash_attention/
  2. Move files, fix crate:: → appropriate ml-core paths
  3. Add pub mod flash_attention; to ml-core lib.rs
  4. In ml lib.rs: replace pub mod flash_attention; with pub use ml_core::flash_attention;
  5. Run: SQLX_OFFLINE=true cargo check -p ml-core -p ml
  6. Commit: refactor(ml): move flash_attention to ml-core

Task 1c: features/ remainder → ml-features (4,348 LOC)

Files to move:

  • crates/ml/src/features/extraction.rs (1,285)
  • crates/ml/src/features/regime_adaptive.rs (673)
  • crates/ml/src/features/multi_timeframe.rs (630)
  • crates/ml/src/features/sample_weights.rs (362)
  • crates/ml/src/features/unified.rs (532)
  • crates/ml/src/features/regime_cusum.rs (372)
  • crates/ml/src/features/regime_transition.rs (333)
  • crates/ml/src/features/production_adapter.rs (123)
  • crates/ml/src/features/mod.rs (38)

Steps:

  1. Check what ml-features already has vs. what's in ml/src/features/
  2. Move non-duplicate files into crates/ml-features/src/
  3. Fix imports: use crate::use crate:: (within ml-features) or add deps
  4. Update ml-features Cargo.toml if new deps needed (e.g., ml-regime)
  5. In ml lib.rs: replace pub mod features; with pub use ml_features as features; (or re-export)
  6. Run: SQLX_OFFLINE=true cargo check -p ml-features -p ml
  7. Commit: refactor(ml): merge remaining features into ml-features

Task 1d: microstructure/ → ml-features (682 LOC)

Files to move:

  • crates/ml/src/microstructure/*.rs

Steps:

  1. Move to crates/ml-features/src/microstructure/
  2. Fix imports
  3. Re-export from ml: pub use ml_features::microstructure;
  4. Run: SQLX_OFFLINE=true cargo check -p ml-features -p ml
  5. Commit: refactor(ml): move microstructure to ml-features

Task 2: Data Modules → ml-data (7,041 LOC)

Depends on: Task 1c (features in ml-features)

Files to move:

  • crates/ml/src/data_loaders/*.rs (4,365 — 7 files)
  • crates/ml/src/data_pipeline/*.rs (1,500)
  • crates/ml/src/training/unified_data_loader.rs + other training/ files (1,176)

Steps:

  1. Add ml-features dependency to crates/ml-data/Cargo.toml
  2. Add dbn, zstd dependencies if not already present
  3. Move data_loaders/ to crates/ml-data/src/data_loaders/
  4. Move data_pipeline/ to crates/ml-data/src/data_pipeline/
  5. Move training/ to crates/ml-data/src/training/
  6. Fix imports: use crate::features::use ml_features::, use crate::types::use ml_core::types::
  7. Fix ensemble::MarketRegime import in dbn_sequence_loader (may need re-export from ml-core)
  8. In ml lib.rs: replace module declarations with re-exports
  9. Run: SQLX_OFFLINE=true cargo check -p ml-data -p ml
  10. Run: SQLX_OFFLINE=true cargo test -p ml-data --lib
  11. Commit: refactor(ml): move data_loaders, data_pipeline, training to ml-data

Task 3: Model Implementation Remnants → Model Sub-crates (~5,700 LOC)

Parallel subtasks — independent per model.

Task 3a: dqn/ → ml-dqn (959 LOC)

Files: crates/ml/src/dqn/*.rs (action_space, circuit_breaker, curiosity, logging, mixed_precision duplicate check, order_routing, portfolio_tracker, regime_conditional, reward, target_update)

Steps:

  1. Check for duplicates: compare ml/src/dqn/ contents with crates/ml-dqn/src/
  2. Move non-duplicate files to crates/ml-dqn/src/
  3. For duplicates: verify ml-dqn version is canonical, delete from ml/src/dqn/
  4. Re-export from ml: pub use ml_dqn as dqn; (already done partially)
  5. Run: SQLX_OFFLINE=true cargo check -p ml-dqn -p ml
  6. Commit: refactor(ml): merge dqn remnants into ml-dqn

Task 3b: ppo/ → ml-ppo (1,004 LOC)

Files: crates/ml/src/ppo/mod.rs, trainable_adapter.rs, stress_testing.rs

Steps:

  1. Move to crates/ml-ppo/src/
  2. Fix imports
  3. Re-export from ml
  4. Check + test
  5. Commit: refactor(ml): merge ppo remnants into ml-ppo

Task 3c: Supervised model dirs → ml-supervised (~4,900 LOC)

Files:

  • crates/ml/src/tft/ (1,796) — training.rs, trainable_adapter.rs
  • crates/ml/src/liquid/ (672)
  • crates/ml/src/tgnn/ (633)
  • crates/ml/src/kan/ (517)
  • crates/ml/src/xlstm/ (439)
  • crates/ml/src/diffusion/ (443)
  • crates/ml/src/mamba/ (399)
  • crates/ml/src/tlob/ (700)
  • crates/ml/src/transformers/ (821)

Steps:

  1. For each model dir: compare with ml-supervised/src/<model>/
  2. Move training logic and trainable adapters to ml-supervised
  3. Move transformers/ to ml-supervised/src/transformers/
  4. Re-export from ml
  5. Run: SQLX_OFFLINE=true cargo check -p ml-supervised -p ml
  6. Commit: refactor(ml): merge supervised model remnants into ml-supervised

Task 4: New Sub-crates (8,305 LOC)

Parallel — independent of other tasks.

Task 4a: benchmark/ → ml-benchmark (5,969 LOC)

Steps:

  1. Create crates/ml-benchmark/ with Cargo.toml
  2. Add to workspace members in root Cargo.toml
  3. Move crates/ml/src/benchmark/ + benchmarks.rs to new crate
  4. Add dependencies: ml-core, ml-dqn, ml-ppo, ml-supervised (model benchmarks reference models)
  5. Re-export from ml: pub use ml_benchmark as benchmark;
  6. Run: SQLX_OFFLINE=true cargo check -p ml-benchmark -p ml
  7. Commit: refactor(ml): extract benchmark into ml-benchmark

Task 4b: deployment/ → ml-deployment (2,336 LOC)

Steps:

  1. Create crates/ml-deployment/ with Cargo.toml
  2. Add to workspace members
  3. Move crates/ml/src/deployment/ to new crate
  4. Add dependencies: ml-core
  5. Re-export from ml
  6. Check + test
  7. Commit: refactor(ml): extract deployment into ml-deployment

Task 5: Trainers → Model Sub-crates (18,034 LOC)

Depends on: Tasks 1 (foundations extracted), Task 3 (model impls merged)

This is the largest and most complex extraction. The trainers have heavy cross-module dependencies. After Tasks 1+3, those dependencies are now in sub-crates, so trainers can follow.

Task 5a: trainers/dqn/ → ml-dqn (9,554 LOC)

Files:

  • trainers/dqn/trainer.rs (5,177)
  • trainers/dqn/config.rs (1,223)
  • trainers/dqn/data_loading.rs (920)
  • trainers/dqn/monitoring.rs (284)
  • trainers/dqn/early_stopping.rs (256)
  • trainers/dqn/features.rs (163)
  • trainers/dqn/financials.rs (178)
  • trainers/dqn/lr_scheduler.rs (228)
  • trainers/dqn/risk.rs (145)
  • trainers/dqn/statistics.rs (134)
  • trainers/dqn/mod.rs (45)

Steps:

  1. Add dependencies to crates/ml-dqn/Cargo.toml:
    • ml-features = { workspace = true } (feature extraction)
    • ml-labeling = { workspace = true } (triple barrier)
    • risk = { workspace = true } (DrawdownMonitor, position limiter)
    • Any missing deps (evaluation metrics — may need to move evaluation too)
  2. Create crates/ml-dqn/src/trainer/ directory
  3. Move all 11 files from crates/ml/src/trainers/dqn/crates/ml-dqn/src/trainer/
  4. Fix imports:
    • use crate::cuda_pipeline::use ml_core::cuda_pipeline::
    • use crate::dqn::use crate:: (now in same crate)
    • use crate::evaluation:: → determine home (may need to move evaluation first)
    • use crate::features::use ml_features::
    • use crate::labeling::use ml_labeling::
    • use crate::memory_optimization::use ml_core::memory_optimization::
  5. Add pub mod trainer; to ml-dqn lib.rs
  6. In ml trainers/mod.rs: pub use ml_dqn::trainer as dqn; (or similar re-export)
  7. Run: SQLX_OFFLINE=true cargo check -p ml-dqn -p ml
  8. Run: SQLX_OFFLINE=true cargo test -p ml-dqn --lib
  9. Commit: refactor(ml): move DQN trainer to ml-dqn

Note: The evaluation/ module (metrics) is used by DQN trainer. It should move to ml-dqn or ml-core depending on whether other trainers also use it. Check with: grep -r 'crate::evaluation' crates/ml/src/trainers/

Task 5b: trainers/ppo.rs → ml-ppo (1,832 LOC)

Steps:

  1. Add deps to ml-ppo Cargo.toml: ml-core (cuda_pipeline, batch_size_resolver)
  2. Move trainers/ppo.rscrates/ml-ppo/src/trainer.rs
  3. Fix imports: use crate::cuda_pipeline::PpoGpuDatause ml_core::cuda_pipeline::PpoGpuData
  4. Fix: use crate::dqn::mixed_precision::training_dtypeuse ml_core::mixed_precision::training_dtype
  5. Add pub mod trainer; to ml-ppo lib.rs
  6. Re-export from ml
  7. Check + test
  8. Commit: refactor(ml): move PPO trainer to ml-ppo

Task 5c: trainers/tft/ → ml-supervised (2,957 LOC)

Steps:

  1. Add deps to ml-supervised: ml-checkpoint, ml-core (memory_optimization)
  2. Move trainers/tft/crates/ml-supervised/src/tft/trainer/ (or tft_trainer/)
  3. Fix imports
  4. Re-export from ml
  5. Check + test
  6. Commit: refactor(ml): move TFT trainer to ml-supervised

Task 5d: Remaining trainers (3,691 LOC)

Files:

  • trainers/mamba2.rs (581) → ml-supervised
  • trainers/liquid.rs (530) → ml-supervised
  • trainers/tlob.rs (769) → ml-supervised
  • trainers/online_learning.rs (982) → ml-core
  • trainers/curriculum.rs (565) → ml-core (shared training infrastructure)
  • trainers/validation_metrics.rs (454) → ml-core (shared)
  • trainers/tft_parquet.rs (302) → ml-supervised

Steps:

  1. Move model-specific trainers to ml-supervised
  2. Move shared training infra to ml-core
  3. Fix imports in each
  4. Update ml trainers/mod.rs to re-export from sub-crates
  5. Check + test
  6. Commit: refactor(ml): move remaining trainers to sub-crates

Task 6: Adapters → ml-hyperopt & ml-ensemble (~16K LOC)

Depends on: Task 5 (trainers in sub-crates)

Task 6a: hyperopt/adapters/ → ml-hyperopt (12,752 LOC + 1,295 supporting)

Steps:

  1. Add dependencies to crates/ml-hyperopt/Cargo.toml:
    • ml-dqn = { workspace = true } (DQN adapter)
    • ml-ppo = { workspace = true } (PPO adapter)
    • ml-supervised = { workspace = true } (8 supervised adapters)
    • ml-features = { workspace = true } (feature extraction)
    • ml-data = { workspace = true } (data loading)
    • dbn, zstd (DBN file loading)
  2. Create crates/ml-hyperopt/src/adapters/ directory
  3. Move all 14 adapter files + mod.rs
  4. Move campaign.rs, shared_data.rs, tests_argmin.rs
  5. Fix imports: use crate::trainers::dqn::use ml_dqn::trainer::
  6. Fix: use crate::evaluation:: → determine new path
  7. Update ml hyperopt/mod.rs to re-export from ml-hyperopt
  8. Run: SQLX_OFFLINE=true cargo check -p ml-hyperopt -p ml
  9. Run: SQLX_OFFLINE=true cargo test -p ml-hyperopt --lib
  10. Commit: refactor(ml): move hyperopt adapters to ml-hyperopt

Task 6b: ensemble/adapters/ → ml-ensemble (3,427 LOC)

Steps:

  1. Add deps to ml-ensemble: ml-dqn, ml-ppo, ml-supervised, ml-features
  2. Move ensemble/adapters/ to crates/ml-ensemble/src/adapters/
  3. Move ensemble/model_adapter.rs to ml-ensemble
  4. Fix imports: model-specific use crate::tft::use ml_supervised::tft::
  5. Fix: use crate::dqn::mixed_precision::use ml_core::mixed_precision::
  6. Re-export from ml
  7. Check + test
  8. Commit: refactor(ml): move ensemble adapters to ml-ensemble

Task 7: Final Cleanup

Depends on: All previous tasks

Task 7a: Small module extraction

Move remaining small modules to existing sub-crates:

  • model_registry/ + model_registry.rs + registry/ → ml-core or ml-checkpoint
  • evaluation/ → ml-core (shared metrics used by multiple trainers)
  • data_validation/ stub → delete (ml-data-validation crate exists)
  • Other stubs (backtesting/, checkpoint/, labeling/, etc.) → delete re-export stubs

Task 7b: Update ml lib.rs

Rewrite crates/ml/src/lib.rs to be a clean orchestration facade:

  • Remove all extracted pub mod declarations
  • Replace with pub use re-exports from sub-crates
  • Keep local modules: integration/, inference.rs, model_factory.rs, training_pipeline.rs, etc.
  • Keep From impls that bridge local types to ml-core types

Task 7c: Update downstream consumers

Check and fix imports in:

  • services/trading_service/ (uses ml:: paths)
  • services/trading_agent_service/
  • services/backtesting_service/
  • services/ml_training_service/
  • bin/fxt/ (CLI binary)
  • Training binaries

For each: grep -r 'use ml::' <service>/src/ and fix any broken paths.

Task 7d: Workspace Cargo.toml

  • Add new crates to [workspace.members]: ml-benchmark, ml-deployment
  • Add workspace deps if needed
  • Verify feature flag propagation (cuda feature must chain through all sub-crates)

Task 7e: Final verification

# Full workspace check
SQLX_OFFLINE=true cargo check --workspace

# Full workspace clippy
SQLX_OFFLINE=true cargo clippy --workspace -- -D warnings

# Run all ml-* tests
for crate in crates/ml-*/; do
  name=$(basename "$crate")
  echo "Testing $name..."
  SQLX_OFFLINE=true cargo test -p "$name" --lib 2>&1 | tail -1
done

# Run ml orchestration crate tests
SQLX_OFFLINE=true cargo test -p ml --lib

# Verify line count reduction
find crates/ml/src -name '*.rs' -type f | xargs wc -l | tail -1
# Expected: ~12,000 (down from 91,000)

Commit: refactor(ml): finalize phase 2 split — ml reduced to orchestration layer


Key Risks & Mitigations

  1. Circular dependencies: Sub-crates must never depend on ml. All deps flow: ml → sub-crates → ml-core. Verified: no sub-crate currently depends on ml.

  2. evaluation module: Used by DQN trainer AND hyperopt adapters. Must decide home (ml-core or ml-dqn) BEFORE moving trainers. Check: grep -r 'crate::evaluation' crates/ml/src/

  3. mixed_precision::training_dtype: Referenced from DQN, PPO, ensemble. Already in ml-core. Verify re-export path works.

  4. Feature flags: cuda feature must propagate. After split, ml's [features] cuda = [...] must include all new sub-crate cuda features.

  5. Re-export backward compat: Services use ml::trainers::dqn::DQNTrainer. After move, must re-export: pub use ml_dqn::trainer as dqn; in trainers/mod.rs.

  6. Test count regression: Track test counts before/after. Expected: same total, redistributed across sub-crates.

Pre-Split Baseline

# Record before starting
find crates/ml/src -name '*.rs' -type f | xargs wc -l | tail -1  # ~91,000
SQLX_OFFLINE=true cargo test -p ml --lib 2>&1 | grep 'test result'