# TFT Trainer Split Analysis

**Agent**: tft-splitter
**Date**: 2025-11-27
**Target File**: `ml/src/trainers/tft.rs` (2,915 lines)
**Reference**: DQN split (4 files: config.rs, statistics.rs, trainer.rs, mod.rs)

---

## Executive Summary

The TFT trainer module can be split into **6 files** following a similar pattern to DQN, but with additional files for QAT metrics and the TFTModel trait abstraction.

**Proposed Structure:**
```
ml/src/trainers/tft/
├── mod.rs              (~50 lines)   - Module exports and re-exports
├── config.rs           (~350 lines)  - Configuration structs
├── types.rs            (~250 lines)  - QAT metrics and helper types
├── model.rs            (~100 lines)  - TFTModel trait and implementations
├── trainer.rs          (~2,000 lines) - Main TFTTrainer implementation
└── tests.rs            (~165 lines)  - Test module
```

---

## Current File Structure Analysis

### Line Distribution

| Section | Lines | Description |
|---------|-------|-------------|
| Header/Imports | 1-33 | Module documentation and use statements |
| QAT Metrics Types | 34-100 | QATMetrics, ScaleStatistics, ZeroPointStatistics, etc. |
| TFTModel Trait | 101-202 | Trait definition and implementations |
| TFTTrainer Struct | 203-277 | Main trainer struct |
| TrainingState | 278-332 | Internal state tracking |
| Progress Types | 333-385 | TrainingProgress, ResourceUsage |
| Config | 386-538 | TFTTrainerConfig struct and impls |
| Trainer Impl | 539-2,498 | Main implementation block with 30+ methods |
| ValidationMetrics | 2,499-2,507 | Helper struct |
| TrainingMetrics | 2,508-2,548 | Public metrics result |
| Tests | 2,550-2,915 | Test module (~365 lines) |

### Major Components Identified

1. **Configuration Structs** (Lines 386-538)
   - `TFTTrainerConfig` (main config)
   - `impl Default for TFTTrainerConfig`
   - `impl TFTTrainerConfig` (conversion methods)

2. **QAT Metrics Types** (Lines 34-100)
   - `QATMetrics`
   - `ScaleStatistics`
   - `ZeroPointStatistics`
   - `ObserverRangeStatistics`
   - `LayerQuantizationMetrics`

3. **TFTModel Trait** (Lines 101-202)
   - `pub trait TFTModel`
   - `impl TFTModel for TemporalFusionTransformer`
   - Commented out QAT implementation

4. **Progress/State Types** (Lines 278-385)
   - `TrainingState` (private)
   - `TrainingProgress` (public)
   - `ResourceUsage` (public)

5. **Statistics Types** (Lines 2,499-2,548)
   - `ValidationMetrics` (private)
   - `TrainingMetrics` (public)

6. **Main Trainer** (Lines 203-2,498)
   - `TFTTrainer` struct
   - `impl Debug for TFTTrainer`
   - `impl TFTTrainer` (30+ methods)

---

## Proposed Split Plan

### File 1: `mod.rs` (~50 lines)

**Purpose**: Module organization and re-exports

**Contents:**
```rust
//! TFT (Temporal Fusion Transformer) Trainer Module
//!
//! Production-grade TFT training pipeline with:
//! - GPU acceleration (4GB VRAM optimized)
//! - QAT (Quantization-Aware Training) support
//! - Real-time progress streaming
//! - Checkpoint persistence to MinIO/S3
//!
//! ## Module Structure
//!
//! - `config` - Training configuration and hyperparameters
//! - `types` - QAT metrics and helper types
//! - `model` - TFTModel trait abstraction
//! - `trainer` - Main TFTTrainer implementation
//! - `tests` - Integration tests

mod config;
mod model;
mod trainer;
mod types;

#[cfg(test)]
mod tests;

// Re-export public API
pub use config::TFTTrainerConfig;
pub use model::TFTModel;
pub use trainer::TFTTrainer;
pub use types::{
    LayerQuantizationMetrics, ObserverRangeStatistics, QATMetrics, ResourceUsage,
    ScaleStatistics, TrainingMetrics, TrainingProgress, ZeroPointStatistics,
};
```

**Lines**: ~50

---

### File 2: `config.rs` (~350 lines)

**Purpose**: Training configuration and hyperparameters

**Contents:**
- Lines 386-538: `TFTTrainerConfig` struct
- All conversion methods (`to_model_config()`, `to_training_config()`)
- Default implementation
- Builder methods (if any)

**Key Structs:**
- `pub struct TFTTrainerConfig`
- `impl Default for TFTTrainerConfig`
- `impl TFTTrainerConfig` (conversion helpers)

**Imports Needed:**
```rust
use serde::{Deserialize, Serialize};
use crate::checkpoint::CheckpointConfig;
use crate::tft::{TFTConfig, training::TFTTrainingConfig};
```

**Lines**: ~350

---

### File 3: `types.rs` (~250 lines)

**Purpose**: QAT metrics, progress reporting, and helper types

**Contents:**
- Lines 34-100: QAT metrics structs
  - `QATMetrics`
  - `ScaleStatistics`
  - `ZeroPointStatistics`
  - `ObserverRangeStatistics`
  - `LayerQuantizationMetrics`
- Lines 333-385: Progress types
  - `TrainingProgress`
  - `ResourceUsage`
- Lines 2,508-2,548: Result types
  - `TrainingMetrics`
- Internal types:
  - `TrainingState` (lines 278-332)
  - `ValidationMetrics` (lines 2,499-2,507)

**Key Structs:**
- All `pub` metric/progress types (exported)
- Internal state types (not exported, only pub(crate))

**Imports Needed:**
```rust
use std::collections::HashMap;
use serde::{Deserialize, Serialize};
```

**Lines**: ~250

---

### File 4: `model.rs` (~100 lines)

**Purpose**: TFTModel trait abstraction for polymorphic FP32/QAT support

**Contents:**
- Lines 101-202: TFTModel trait
  - `pub trait TFTModel: Send + Sync`
  - `impl TFTModel for TemporalFusionTransformer`
  - Commented QAT implementation (for future use)

**Key Components:**
- Trait definition with forward pass, device, config access
- FP32 implementation
- QAT placeholder (currently disabled)

**Imports Needed:**
```rust
use std::sync::Arc;
use candle_core::{Device, Tensor};
use candle_nn::VarMap;
use crate::tft::{TFTConfig, TemporalFusionTransformer};
use crate::MLError;
```

**Lines**: ~100

---

### File 5: `trainer.rs` (~2,000 lines)

**Purpose**: Main TFTTrainer implementation

**Contents:**
- Lines 203-277: `TFTTrainer` struct definition
- Lines 259-277: `impl Debug for TFTTrainer`
- Lines 539-2,498: `impl TFTTrainer` (all methods)
  - Constructor: `new()`
  - Public API: `train()`, `set_progress_callback()`
  - Internal methods: 30+ methods for training, validation, checkpointing, etc.

**Major Methods:**
1. `new()` - Constructor with auto-batch sizing
2. `train()` - Main training loop
3. `train_epoch()` - Single epoch training
4. `validate_epoch()` - Validation pass
5. `batch_to_tensors()` - Data conversion
6. `compute_quantile_loss()` - Loss computation
7. `compute_rmse()` - RMSE metric
8. `check_early_stopping()` - Early stopping logic
9. `save_checkpoint()` - Checkpoint persistence
10. `send_progress_update()` - Progress streaming
11. QAT methods:
    - `run_qat_calibration()`
    - `qat_to_quantized_checkpoint()`
    - `export_qat_metrics()`
    - `apply_qat_lr_schedule()`
12. Helper methods:
    - `initialize_optimizer()`
    - `sync_cuda_device()`
    - `get_resource_usage()`
    - Various getters

**Imports Needed:**
```rust
use std::collections::HashMap;
use std::sync::Arc;
use std::time::{Duration, Instant, SystemTime};

use candle_core::{Device, IndexOp, Tensor};
use candle_nn::VarMap;
use ndarray::Dimension;
use tokio::sync::mpsc;
use tracing::{debug, error, info, instrument, warn};

use crate::checkpoint::{CheckpointConfig, CheckpointManager, CheckpointMetadata, CheckpointStorage};
use crate::memory_optimization::{AutoBatchSizer, BatchSizeConfig, ModelPrecision, OptimizerType};
use crate::tft::training::{TFTBatch, TFTDataLoader};
use crate::{MLError, MLResult};

use super::config::TFTTrainerConfig;
use super::model::TFTModel;
use super::types::{QATMetrics, ResourceUsage, TrainingMetrics, TrainingProgress, TrainingState, ValidationMetrics};
```

**Lines**: ~2,000

---

### File 6: `tests.rs` (~165 lines)

**Purpose**: Integration tests (separated from main code)

**Contents:**
- Lines 2,550-2,915: Test module
  - `test_tft_trainer_creation()`
  - `test_training_config_conversion()`
  - `test_checkpoint_save_load()`
  - OOM retry tests (multiple tests)

**Structure:**
```rust
use super::*;
use crate::checkpoint::FileSystemStorage;
use std::path::PathBuf;
use tempfile::TempDir;

#[tokio::test]
async fn test_tft_trainer_creation() { ... }

#[tokio::test]
async fn test_training_config_conversion() { ... }

// ... more tests
```

**Lines**: ~165

---

## Dependencies Between Modules

### Import Graph

```
mod.rs
  ├─> config.rs (TFTTrainerConfig)
  ├─> types.rs (metrics, progress, state)
  ├─> model.rs (TFTModel trait)
  ├─> trainer.rs (TFTTrainer)
  └─> tests.rs (uses all above)

trainer.rs
  ├─> config.rs (TFTTrainerConfig)
  ├─> types.rs (TrainingState, ValidationMetrics, etc.)
  └─> model.rs (Box<dyn TFTModel>)

model.rs
  ├─> crate::tft (TemporalFusionTransformer, TFTConfig)
  └─> independent (no internal deps)

config.rs
  ├─> crate::tft (TFTConfig, TFTTrainingConfig)
  └─> independent (no internal deps)

types.rs
  └─> independent (no internal deps)
```

### Internal Type Visibility

**Private Types** (only used in trainer.rs):
- `TrainingState` → Keep in types.rs as `pub(crate)`
- `ValidationMetrics` → Keep in types.rs as `pub(crate)`

**Public Types** (exported from mod.rs):
- All QAT metrics types
- `TrainingProgress`
- `ResourceUsage`
- `TrainingMetrics`
- `TFTTrainerConfig`
- `TFTModel` trait
- `TFTTrainer`

---

## Migration Strategy

### Step 1: Create Directory Structure
```bash
mkdir -p ml/src/trainers/tft
```

### Step 2: Create Files in Order (preserves compilation)

1. **Create `types.rs`** (no dependencies on other split files)
   - Extract QAT metrics (lines 34-100)
   - Extract progress types (lines 333-385)
   - Extract TrainingState (lines 278-332)
   - Extract ValidationMetrics (lines 2,499-2,507)
   - Extract TrainingMetrics (lines 2,508-2,548)

2. **Create `model.rs`** (only depends on external crates)
   - Extract TFTModel trait (lines 101-202)

3. **Create `config.rs`** (only depends on external crates)
   - Extract TFTTrainerConfig (lines 386-538)

4. **Create `trainer.rs`** (depends on types, model, config)
   - Extract TFTTrainer struct (lines 203-277)
   - Extract impl blocks (lines 259-277, 539-2,498)
   - Update imports to use `super::{types, model, config}`

5. **Create `tests.rs`** (depends on all above)
   - Extract test module (lines 2,550-2,915)
   - Remove `#[cfg(test)]` attribute (entire file is tests)

6. **Create `mod.rs`** (coordinates all modules)
   - Module declarations
   - Re-exports

7. **Update `ml/src/trainers/mod.rs`**
   - Change from `pub mod tft;` to module with sub-modules
   - Or keep as `pub mod tft;` and let tft/mod.rs handle re-exports

### Step 3: Update References

**Files that import from tft.rs:**
- Likely in `ml/src/lib.rs` or other trainer files
- Should continue to work due to re-exports in `mod.rs`

**Example Import Changes:**
```rust
// Before (may still work)
use crate::trainers::tft::TFTTrainer;

// After (explicit, if needed)
use crate::trainers::tft::TFTTrainer;  // Same! Re-exported from mod.rs
```

### Step 4: Verification

1. Run `cargo check` after each file creation
2. Run full test suite: `cargo test --package ml trainers::tft`
3. Run clippy: `cargo clippy --package ml`
4. Format: `cargo fmt --package ml`

---

## Risk Assessment

### Low Risk
- **types.rs**: Pure data structures, no complex logic
- **config.rs**: Configuration only, well-defined boundaries
- **model.rs**: Trait abstraction, clean separation

### Medium Risk
- **trainer.rs**: Large impl block, many internal methods
  - **Mitigation**: Keep all trainer methods together, don't split impl block
  - **Mitigation**: Use `pub(crate)` for internal helpers in types.rs

### Potential Issues

1. **Circular Dependencies**
   - **Risk**: If types.rs needs TFTTrainer, and trainer.rs needs types
   - **Mitigation**: Use `pub(crate)` and keep types as pure data structures

2. **Import Path Changes**
   - **Risk**: External code may need updates
   - **Mitigation**: Re-export everything from mod.rs for backward compatibility

3. **Private Method Visibility**
   - **Risk**: Internal helpers in TFTTrainer may need visibility adjustments
   - **Mitigation**: Keep all trainer impl in one file, use `pub(crate)` sparingly

---

## Comparison with DQN Split

| Aspect | DQN | TFT (Proposed) |
|--------|-----|----------------|
| **Files** | 4 | 6 |
| **config.rs** | 566 lines | ~350 lines |
| **statistics.rs** | 134 lines | ~250 lines (types.rs) |
| **trainer.rs** | 4,007 lines | ~2,000 lines |
| **mod.rs** | 30 lines | ~50 lines |
| **Extra files** | - | model.rs (~100), tests.rs (~165) |

**Key Differences:**
1. TFT has TFTModel trait abstraction → separate `model.rs`
2. TFT has extensive QAT metrics → dedicated section in `types.rs`
3. TFT tests are larger → separate `tests.rs` file
4. DQN has simpler statistics → smaller statistics.rs

**Similarities:**
- Both split config into dedicated file
- Both extract statistics/types
- Both keep main trainer implementation in trainer.rs
- Both use mod.rs for re-exports

---

## Next Steps

### For Implementation Agent:

1. **Create `ml/src/trainers/tft/` directory**
2. **Extract files in order** (types → model → config → trainer → tests → mod)
3. **Verify compilation** after each file
4. **Update imports** in trainer.rs to use `super::{types, model, config}`
5. **Run test suite** to ensure no breakage
6. **Update `ml/src/trainers/mod.rs`** if needed

### Validation Checklist:
- [ ] All files compile independently
- [ ] `cargo test --package ml trainers::tft` passes
- [ ] No new clippy warnings
- [ ] Public API unchanged (backward compatible)
- [ ] File sizes reasonable (<2,500 lines each)
- [ ] No circular dependencies

---

## Conclusion

The TFT trainer can be cleanly split into 6 files following the DQN pattern with minor extensions:
- **types.rs** replaces statistics.rs (includes QAT metrics)
- **model.rs** added for trait abstraction
- **tests.rs** separated due to size
- **trainer.rs** remains largest but manageable (~2,000 lines vs 2,915)

This split improves:
✅ **Modularity**: Clear separation of concerns
✅ **Maintainability**: Smaller, focused files
✅ **Testability**: Dedicated test file
✅ **Compilation Speed**: Smaller compilation units
✅ **Code Navigation**: Easier to find specific components

**Estimated total lines after split**: 2,915 lines (no change, just reorganized)

**Ready for implementation**: ✅ Yes, plan is complete and low-risk.