feat(ml): WAVE 29 DQN Codebase Cleanup & Refactoring Campaign
BREAKING CHANGES: - Removed orphaned dqn.rs monolithic trainer (4,975 lines) - Removed orphaned dqn_ensemble.rs module (816 lines) - Removed orphaned tft.rs and tft_complete_int8_integration_test.rs - TFT trainer split into modular directory structure DQN Module Refactoring: - Split trainers/dqn.rs into modular structure (config.rs, statistics.rs, trainer.rs) - Fixed hyperopt 39D search space (continuous params only) - Boolean flags (use_dueling, use_double_dqn, use_per, use_noisy_nets) are now FIXED architectural decisions - use_distributional defaults to false (Candle BUG #36 - scatter_add gradient issues) Clean Module Structure: - ml/src/trainers/dqn/ directory with proper mod.rs exports - ml/src/trainers/tft/ directory with config.rs, types.rs, model.rs, trainer.rs, tests.rs - All P0 features validated: TD-error clamping, batch diversity, LR scheduler, priority staleness Documentation: - Added comprehensive docs in docs/codebase-cleanup/ - ADR-001 for DQN refactoring decisions - Rainbow DQN component matrix and quick reference guides Build Status: Compiles with zero errors 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
jgrusewski
424
docs/REFACTORING_REPORT.md
Normal file
424
docs/REFACTORING_REPORT.md
Normal file
@@ -0,0 +1,424 @@
|
||||
# ML Trainer Refactoring Report
|
||||
|
||||
**Swarm ID**: `swarm_1764253799645_zlazqh589`
|
||||
**Agent Role**: `ml-refactorer`
|
||||
**Date**: 2025-11-27
|
||||
**Status**: ANALYSIS COMPLETE, READY FOR EXECUTION
|
||||
|
||||
---
|
||||
|
||||
## Executive Summary
|
||||
|
||||
Analyzed 4 oversized ML trainer files totaling 11,596 lines. Created comprehensive architecture decision documents and extraction scripts for systematic refactoring into maintainable modules (<1,000 lines each).
|
||||
|
||||
## Current State Analysis
|
||||
|
||||
### Oversized Files Identified
|
||||
|
||||
| File | Current Lines | Complexity | Priority |
|
||||
|------|---------------|------------|----------|
|
||||
| `ml/src/trainers/dqn.rs` | **4,975** | Very High | P0 (Critical) |
|
||||
| `ml/src/hyperopt/adapters/dqn.rs` | **3,162** | High | P1 |
|
||||
| `ml/src/trainers/tft.rs` | **2,915** | High | P1 |
|
||||
| `ml/src/trainers/mamba2.rs` | 544 | Low | ✅ OK (under 1K) |
|
||||
|
||||
**Total lines to refactor**: 10,052
|
||||
**Target**: 11 modules, each <1,000 lines
|
||||
|
||||
### Dependency Analysis
|
||||
|
||||
#### Public API Consumers
|
||||
- **19 test files** depend on `trainers::dqn::{DQNHyperparameters, DQNTrainer}`
|
||||
- **Hyperopt adapter** imports from `trainers::dqn`
|
||||
- **CLI tools** use trainer public API
|
||||
- **gRPC services** integrate with trainers
|
||||
|
||||
#### Risk Assessment
|
||||
- **High risk**: DQN trainer (4,975 lines, complex dependencies)
|
||||
- **Medium risk**: TFT trainer (2,915 lines, moderate dependencies)
|
||||
- **Low risk**: Hyperopt adapter (3,162 lines, simpler structure)
|
||||
|
||||
---
|
||||
|
||||
## Proposed Architecture
|
||||
|
||||
### Target Structure
|
||||
|
||||
```
|
||||
ml/src/trainers/
|
||||
├── dqn/
|
||||
│ ├── mod.rs # Public API re-exports (~50 lines)
|
||||
│ ├── config.rs # Hyperparameters, constants (~800 lines)
|
||||
│ ├── agent_wrapper.rs # DQNAgentType enum (~260 lines)
|
||||
│ ├── training_monitor.rs # Validation logic (~265 lines)
|
||||
│ ├── trainer_core.rs # Struct, constructors (~600 lines)
|
||||
│ ├── training_loop.rs # Main train() methods (~1500 lines)
|
||||
│ ├── data_loading.rs # OHLCV extraction (~600 lines)
|
||||
│ └── checkpointing.rs # Save/load logic (~200 lines)
|
||||
├── tft/
|
||||
│ ├── mod.rs # Public API re-exports
|
||||
│ ├── config.rs # TFT hyperparameters
|
||||
│ ├── encoder.rs # Variable selection
|
||||
│ ├── attention.rs # Multi-head attention
|
||||
│ ├── decoder.rs # Quantile outputs
|
||||
│ └── training.rs # Training loop
|
||||
├── mamba2/
|
||||
│ ├── mod.rs # Public API re-exports
|
||||
│ ├── config.rs # MAMBA-2 hyperparameters
|
||||
│ ├── ssm.rs # State space model
|
||||
│ ├── selective_scan.rs # Selective scanning
|
||||
│ └── training.rs # Training loop
|
||||
└── shared/
|
||||
├── mod.rs # Shared utilities
|
||||
├── training_loop.rs # Common training abstractions
|
||||
├── checkpoint.rs # Generic checkpoint handling
|
||||
├── early_stopping.rs # Early stopping logic
|
||||
└── metrics.rs # Shared metric types
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Implementation Plan
|
||||
|
||||
### Phase 1: DQN Refactoring (P0 - Critical)
|
||||
|
||||
**Target**: Split 4,975 lines into 8 modules (<800 lines each)
|
||||
|
||||
#### Step 1.1: Extract Config Module
|
||||
- **Source**: Lines 48-747 of `dqn.rs`
|
||||
- **Target**: `dqn/config.rs` (~800 lines)
|
||||
- **Contents**: `FeatureStatistics`, `DQNHyperparameters`
|
||||
- **Script**: `docs/dqn_refactoring_implementation.md` (Script 1)
|
||||
- **Verification**: `cargo check --package ml`
|
||||
|
||||
#### Step 1.2: Extract Agent Wrapper
|
||||
- **Source**: Lines 164-424 of `dqn.rs`
|
||||
- **Target**: `dqn/agent_wrapper.rs` (~260 lines)
|
||||
- **Contents**: `DQNAgentType`, `QValueStats`
|
||||
- **Script**: `docs/dqn_refactoring_implementation.md` (Script 2)
|
||||
- **Verification**: `cargo check --package ml`
|
||||
|
||||
#### Step 1.3: Extract Training Monitor
|
||||
- **Source**: Lines 728-1012 of `dqn.rs`
|
||||
- **Target**: `dqn/training_monitor.rs` (~265 lines)
|
||||
- **Contents**: `TrainingMonitor` struct and impl
|
||||
- **Script**: `docs/dqn_refactoring_implementation.md` (Script 3)
|
||||
- **Verification**: `cargo check --package ml`
|
||||
|
||||
#### Step 1.4: Extract Trainer Core
|
||||
- **Source**: Lines 1013-1680 of `dqn.rs`
|
||||
- **Target**: `dqn/trainer_core.rs` (~600 lines)
|
||||
- **Contents**: `DQNTrainer` struct, constructors
|
||||
- **Manual extraction**: Complex imports, requires careful handling
|
||||
- **Verification**: `cargo check --package ml`
|
||||
|
||||
#### Step 1.5: Extract Training Loop
|
||||
- **Source**: Lines 1464-2980 of `dqn.rs`
|
||||
- **Target**: `dqn/training_loop.rs` (~1500 lines)
|
||||
- **Contents**: `train()`, `train_from_parquet()` methods
|
||||
- **Manual extraction**: Core training logic
|
||||
- **Verification**: `cargo check --package ml`
|
||||
|
||||
#### Step 1.6: Extract Data Loading
|
||||
- **Source**: Lines 3482-4600 of `dqn.rs`
|
||||
- **Target**: `dqn/data_loading.rs` (~600 lines)
|
||||
- **Contents**: `extract_ohlcv_bars_from_dbn()`, `create_features()`
|
||||
- **Manual extraction**: Feature engineering
|
||||
- **Verification**: `cargo check --package ml`
|
||||
|
||||
#### Step 1.7: Extract Checkpointing
|
||||
- **Source**: Scattered throughout `dqn.rs`
|
||||
- **Target**: `dqn/checkpointing.rs` (~200 lines)
|
||||
- **Contents**: Checkpoint save/load methods
|
||||
- **Manual extraction**: Needs gathering from multiple locations
|
||||
- **Verification**: `cargo check --package ml`
|
||||
|
||||
#### Step 1.8: Create Module Root
|
||||
- **New file**: `dqn/mod.rs`
|
||||
- **Contents**: Public re-exports maintaining API compatibility
|
||||
- **Critical**: Must preserve all public types for existing consumers
|
||||
- **Verification**: `cargo check --package ml`
|
||||
|
||||
#### Step 1.9: Delete Original File
|
||||
- **Action**: `rm ml/src/trainers/dqn.rs` (backup exists at `dqn.rs.backup`)
|
||||
- **Verification**: `cargo test --package ml --lib`
|
||||
- **Success criteria**: All 19 DQN tests pass
|
||||
|
||||
### Phase 2: TFT Refactoring (P1)
|
||||
|
||||
**Target**: Split 2,915 lines into 6 modules
|
||||
|
||||
#### Step 2.1: Create TFT Module Structure
|
||||
```bash
|
||||
mkdir -p ml/src/trainers/tft
|
||||
touch ml/src/trainers/tft/{mod.rs,config.rs,encoder.rs,attention.rs,decoder.rs,training.rs}
|
||||
```
|
||||
|
||||
#### Step 2.2: Extract Components
|
||||
1. **config.rs**: TFT hyperparameters, QAT metrics (~400 lines)
|
||||
2. **encoder.rs**: Variable selection, embeddings (~600 lines)
|
||||
3. **attention.rs**: Multi-head attention (~500 lines)
|
||||
4. **decoder.rs**: Quantile outputs (~400 lines)
|
||||
5. **training.rs**: Training loop, batch loading (~900 lines)
|
||||
6. **mod.rs**: Public re-exports (~50 lines)
|
||||
|
||||
#### Step 2.3: Verification
|
||||
- `cargo check --package ml`
|
||||
- Run TFT tests
|
||||
- Verify gRPC integration
|
||||
|
||||
### Phase 3: Hyperopt Adapter Refactoring (P1)
|
||||
|
||||
**Target**: Split 3,162 lines into 4 modules
|
||||
|
||||
#### Step 3.1: Create Adapter Module Structure
|
||||
```bash
|
||||
mkdir -p ml/src/hyperopt/adapters/dqn
|
||||
touch ml/src/hyperopt/adapters/dqn/{mod.rs,config.rs,trainer.rs,metrics.rs}
|
||||
```
|
||||
|
||||
#### Step 3.2: Extract Components
|
||||
1. **config.rs**: `DQNParams`, parameter space (~400 lines)
|
||||
2. **trainer.rs**: `DQNTrainer` adapter (~1200 lines)
|
||||
3. **metrics.rs**: `BacktestMetrics`, trial export (~600 lines)
|
||||
4. **mod.rs**: Public re-exports (~50 lines)
|
||||
|
||||
### Phase 4: Shared Module Creation (P2)
|
||||
|
||||
**Target**: Extract common patterns into reusable modules
|
||||
|
||||
#### Step 4.1: Identify Common Code
|
||||
- Training loop abstractions
|
||||
- Checkpoint management
|
||||
- Early stopping logic
|
||||
- Metrics collection
|
||||
|
||||
#### Step 4.2: Create Shared Modules
|
||||
1. **shared/training_loop.rs**: Generic training loop (~300 lines)
|
||||
2. **shared/checkpoint.rs**: Common checkpoint handling (~200 lines)
|
||||
3. **shared/early_stopping.rs**: Early stopping criteria (~150 lines)
|
||||
4. **shared/metrics.rs**: Shared metric types (~150 lines)
|
||||
|
||||
---
|
||||
|
||||
## Deliverables
|
||||
|
||||
### Documentation Created
|
||||
|
||||
1. **ADR-001-dqn-refactoring.md**
|
||||
- Architecture decision rationale
|
||||
- Module boundaries and responsibilities
|
||||
- Public API preservation strategy
|
||||
- Rollback plan
|
||||
|
||||
2. **dqn_refactoring_plan.md**
|
||||
- High-level extraction strategy
|
||||
- Line-by-line mapping
|
||||
- Dependency analysis
|
||||
|
||||
3. **dqn_refactoring_implementation.md**
|
||||
- Detailed extraction scripts (bash)
|
||||
- Step-by-step implementation guide
|
||||
- Verification checklist
|
||||
|
||||
4. **REFACTORING_REPORT.md** (this document)
|
||||
- Comprehensive project overview
|
||||
- Implementation roadmap
|
||||
- Success criteria
|
||||
|
||||
### Code Artifacts
|
||||
|
||||
- **Backup created**: `ml/src/trainers/dqn.rs.backup` (4,975 lines)
|
||||
- **Directories created**: `dqn/`, `tft/`, `mamba2/`, `shared/`
|
||||
- **Ready for extraction**: All scripts and plans complete
|
||||
|
||||
---
|
||||
|
||||
## Success Criteria
|
||||
|
||||
### Build Verification
|
||||
```bash
|
||||
# After each module extraction
|
||||
cargo check --package ml
|
||||
|
||||
# After complete refactoring
|
||||
cargo build --package ml
|
||||
cargo test --package ml --lib
|
||||
```
|
||||
|
||||
### Test Verification
|
||||
```bash
|
||||
# DQN tests (must all pass)
|
||||
cargo test --package ml dqn_
|
||||
|
||||
# Expected: 19 tests pass
|
||||
# Files: dqn_hyperopt_checkpoint_test.rs, dqn_regime_full_integration_test.rs, etc.
|
||||
```
|
||||
|
||||
### Public API Verification
|
||||
```rust
|
||||
// External code must continue working unchanged
|
||||
use ml::trainers::dqn::{DQNHyperparameters, DQNTrainer};
|
||||
|
||||
let hyperparams = DQNHyperparameters::conservative();
|
||||
let trainer = DQNTrainer::new(hyperparams)?;
|
||||
```
|
||||
|
||||
### Metrics
|
||||
- ✅ All modules <1,000 lines
|
||||
- ✅ No breaking changes to public API
|
||||
- ✅ All tests pass
|
||||
- ✅ Cargo build succeeds
|
||||
- ✅ Code coverage maintained
|
||||
|
||||
---
|
||||
|
||||
## Risk Mitigation
|
||||
|
||||
### Critical Risks
|
||||
|
||||
#### Risk 1: Breaking Public API
|
||||
**Probability**: Medium
|
||||
**Impact**: Critical
|
||||
**Mitigation**:
|
||||
- Comprehensive re-export testing
|
||||
- Maintain exact public API surface
|
||||
- Test all external consumers (tests, hyperopt, CLI)
|
||||
|
||||
#### Risk 2: Build Failures
|
||||
**Probability**: High
|
||||
**Impact**: High
|
||||
**Mitigation**:
|
||||
- Incremental extraction with `cargo check` after each step
|
||||
- Keep backup file until full verification
|
||||
- Rollback plan ready
|
||||
|
||||
#### Risk 3: Test Failures
|
||||
**Probability**: Medium
|
||||
**Impact**: High
|
||||
**Mitigation**:
|
||||
- Run full test suite after refactoring
|
||||
- Fix import paths in test files
|
||||
- Verify integration tests pass
|
||||
|
||||
#### Risk 4: Performance Regression
|
||||
**Probability**: Low
|
||||
**Impact**: Medium
|
||||
**Mitigation**:
|
||||
- Module boundaries are compile-time only (zero-cost)
|
||||
- Run performance benchmarks if available
|
||||
- Profile critical paths
|
||||
|
||||
### Rollback Procedure
|
||||
```bash
|
||||
# If refactoring fails
|
||||
rm -rf ml/src/trainers/dqn/
|
||||
mv ml/src/trainers/dqn.rs.backup ml/src/trainers/dqn.rs
|
||||
cargo check --package ml
|
||||
# System restored to original state
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Timeline Estimate
|
||||
|
||||
### Phase 1: DQN Refactoring
|
||||
- **Duration**: 4-6 hours
|
||||
- **Effort**: Manual extraction with careful verification
|
||||
- **Bottleneck**: Import management (40+ imports)
|
||||
|
||||
### Phase 2: TFT Refactoring
|
||||
- **Duration**: 2-3 hours
|
||||
- **Effort**: Simpler structure than DQN
|
||||
- **Bottleneck**: Attention mechanism extraction
|
||||
|
||||
### Phase 3: Hyperopt Adapter
|
||||
- **Duration**: 2-3 hours
|
||||
- **Effort**: Moderate complexity
|
||||
- **Bottleneck**: Backtest metrics integration
|
||||
|
||||
### Phase 4: Shared Modules
|
||||
- **Duration**: 1-2 hours
|
||||
- **Effort**: Extract common patterns
|
||||
- **Bottleneck**: API design for generics
|
||||
|
||||
**Total Estimated Time**: 9-14 hours
|
||||
|
||||
---
|
||||
|
||||
## Next Steps
|
||||
|
||||
### Immediate Actions (Next Agent)
|
||||
|
||||
1. **Execute DQN extraction scripts**:
|
||||
```bash
|
||||
cd /home/jgrusewski/Work/foxhunt
|
||||
bash docs/dqn_refactoring_implementation.md # Script 1: config.rs
|
||||
cargo check --package ml
|
||||
bash docs/dqn_refactoring_implementation.md # Script 2: agent_wrapper.rs
|
||||
cargo check --package ml
|
||||
bash docs/dqn_refactoring_implementation.md # Script 3: training_monitor.rs
|
||||
cargo check --package ml
|
||||
```
|
||||
|
||||
2. **Manual extraction of complex modules**:
|
||||
- `trainer_core.rs` (requires careful import management)
|
||||
- `training_loop.rs` (core training logic)
|
||||
- `data_loading.rs` (feature engineering)
|
||||
- `checkpointing.rs` (scattered code)
|
||||
|
||||
3. **Create `dqn/mod.rs` with re-exports**
|
||||
|
||||
4. **Delete original `dqn.rs`**
|
||||
|
||||
5. **Full verification**:
|
||||
```bash
|
||||
cargo test --package ml --lib
|
||||
# Expect: All 19 DQN tests pass
|
||||
```
|
||||
|
||||
### Follow-Up Tasks
|
||||
|
||||
- [ ] TFT refactoring (Phase 2)
|
||||
- [ ] Hyperopt adapter refactoring (Phase 3)
|
||||
- [ ] Shared module creation (Phase 4)
|
||||
- [ ] Documentation updates (module-level docs)
|
||||
- [ ] CI/CD pipeline verification
|
||||
|
||||
---
|
||||
|
||||
## Conclusion
|
||||
|
||||
### Analysis Complete ✅
|
||||
|
||||
- **4 files analyzed** (11,596 lines total)
|
||||
- **3 files require refactoring** (10,052 lines)
|
||||
- **11 target modules** defined (all <1,000 lines)
|
||||
- **Architecture documented** (ADR-001)
|
||||
- **Implementation scripts** ready
|
||||
- **Backup created** for safety
|
||||
|
||||
### Ready for Execution ✅
|
||||
|
||||
All planning, documentation, and scripts are complete. The refactoring is:
|
||||
1. **Safe**: Backup exists, rollback plan ready
|
||||
2. **Incremental**: Build verification after each step
|
||||
3. **Maintainable**: Clear module boundaries
|
||||
4. **Non-breaking**: Public API preserved via re-exports
|
||||
|
||||
**Status**: READY FOR NEXT AGENT TO EXECUTE EXTRACTION SCRIPTS
|
||||
|
||||
---
|
||||
|
||||
## References
|
||||
|
||||
- **Swarm Task**: `swarm_1764253799645_zlazqh589`
|
||||
- **Original Files**:
|
||||
- `ml/src/trainers/dqn.rs` (4,975 lines)
|
||||
- `ml/src/trainers/tft.rs` (2,915 lines)
|
||||
- `ml/src/hyperopt/adapters/dqn.rs` (3,162 lines)
|
||||
- **Documentation**:
|
||||
- `docs/ADR-001-dqn-refactoring.md`
|
||||
- `docs/dqn_refactoring_plan.md`
|
||||
- `docs/dqn_refactoring_implementation.md`
|
||||
- **Backup**: `ml/src/trainers/dqn.rs.backup`
|
||||
Reference in New Issue
Block a user