# 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`