# Hyperparameter Tuning Integration Checklist

## Overview

This checklist guides integration of the Optuna hyperparameter tuner (`hyperparameter_tuner.py`) with the Rust ML Training Service.

## Files Created

- ✅ `hyperparameter_tuner.py` (609 lines) - Main Optuna controller
- ✅ `tuning_config.yaml` (4.7KB) - Search space definitions for 6 model types
- ✅ `requirements-tuner.txt` - Python dependencies
- ✅ `HYPERPARAMETER_TUNING.md` (21KB) - Comprehensive documentation
- ✅ `scripts/generate_python_proto.sh` - Proto stub generation
- ✅ `scripts/example_tuning_job.sh` - Example usage script
- ✅ `tests/test_hyperparameter_tuner.py` - Unit tests

## Integration Steps

### 1. Install Python Dependencies

```bash
cd services/ml_training_service
pip3 install -r requirements-tuner.txt
```

**Dependencies**:
- `optuna>=3.0.0` - Hyperparameter optimization
- `grpcio>=1.50.0` - gRPC client
- `PyYAML>=6.0` - Config parsing
- `nvidia-ml-py3>=7.352.0` - GPU monitoring

### 2. Generate Python gRPC Stubs

```bash
cd services/ml_training_service
./scripts/generate_python_proto.sh
```

**Output**:
- `proto/ml_training_pb2.py`
- `proto/ml_training_pb2_grpc.py`

### 3. Verify Proto Definitions

Ensure these proto messages exist in `proto/ml_training.proto`:

- ✅ `TrainModelRequest` (line 179)
- ✅ `TrainModelResponse` (line 187)
- ✅ `StartTuningJobRequest` (line 132)
- ✅ `GetTuningJobStatusRequest` (line 149)
- ✅ `StopTuningJobRequest` (line 167)

**Enums**:
- ✅ `TuningJobStatus` (line 221)
- ✅ `TrialState` (line 231)

### 4. Implement Rust Service Methods

#### 4.1 StartTuningJob

```rust
// services/ml_training_service/src/service.rs

use tokio::process::Command;
use std::collections::HashMap;

pub async fn start_tuning_job(
    &self,
    request: StartTuningJobRequest,
) -> Result<StartTuningJobResponse, Status> {
    // Generate job ID
    let job_id = Uuid::new_v4().to_string();
    let storage_path = format!("/minio/studies/study_{}.log", job_id);

    // Serialize data source to JSON
    let data_source_json = serde_json::to_string(&request.data_source)
        .map_err(|e| Status::internal(format!("Failed to serialize data source: {}", e)))?;

    // Build Python subprocess command
    let mut cmd = Command::new("python3");
    cmd.current_dir(env::current_dir()?)
        .arg("hyperparameter_tuner.py")
        .arg("--job-id").arg(&job_id)
        .arg("--model-type").arg(&request.model_type)
        .arg("--num-trials").arg(request.num_trials.to_string())
        .arg("--config").arg(&request.config_path)
        .arg("--data-source-json").arg(&data_source_json)
        .arg("--storage-path").arg(&storage_path)
        .arg("--grpc-host").arg("localhost")
        .arg("--grpc-port").arg("50054");

    if request.use_gpu {
        cmd.arg("--use-gpu");
    }

    // Spawn subprocess (non-blocking)
    let child = cmd.spawn()
        .map_err(|e| Status::internal(format!("Failed to spawn tuner: {}", e)))?;

    // Store child process handle for monitoring
    self.tuning_jobs.write().await.insert(job_id.clone(), child);

    // Persist job metadata to database
    sqlx::query!(
        "INSERT INTO tuning_jobs (id, model_type, num_trials, status, storage_path, created_at)
         VALUES ($1, $2, $3, $4, $5, NOW())",
        job_id,
        request.model_type,
        request.num_trials as i32,
        "RUNNING",
        storage_path
    )
    .execute(&self.db_pool)
    .await
    .map_err(|e| Status::internal(format!("Database error: {}", e)))?;

    Ok(StartTuningJobResponse {
        job_id,
        status: TuningJobStatus::TuningRunning.into(),
        message: "Tuning job started successfully".to_string(),
    })
}
```

#### 4.2 GetTuningJobStatus

```rust
pub async fn get_tuning_job_status(
    &self,
    request: GetTuningJobStatusRequest,
) -> Result<GetTuningJobStatusResponse, Status> {
    let job_id = &request.job_id;

    // Query database for job metadata
    let job = sqlx::query!(
        "SELECT model_type, num_trials, status, storage_path, created_at, updated_at
         FROM tuning_jobs WHERE id = $1",
        job_id
    )
    .fetch_optional(&self.db_pool)
    .await
    .map_err(|e| Status::internal(format!("Database error: {}", e)))?
    .ok_or_else(|| Status::not_found("Tuning job not found"))?;

    // Load Optuna study from storage
    let storage_path = job.storage_path;
    let (best_params, best_metrics, trial_history, current_trial) =
        self.load_optuna_study(&storage_path, job_id).await?;

    Ok(GetTuningJobStatusResponse {
        job_id: job_id.clone(),
        status: TuningJobStatus::from_str(&job.status)?.into(),
        current_trial,
        total_trials: job.num_trials as u32,
        best_params,
        best_metrics,
        trial_history,
        message: format!("Trial {}/{}", current_trial, job.num_trials),
        started_at: job.created_at.unwrap().timestamp(),
        updated_at: job.updated_at.map(|t| t.timestamp()).unwrap_or(0),
    })
}

async fn load_optuna_study(
    &self,
    storage_path: &str,
    job_id: &str,
) -> Result<(HashMap<String, f32>, HashMap<String, f32>, Vec<TrialResult>, u32), Status> {
    // Call Python script to read study
    let output = Command::new("python3")
        .arg("-c")
        .arg(format!(r#"
import json
import optuna
from optuna.storages import JournalStorage, JournalFileStorage

storage = JournalStorage(JournalFileStorage('{}'))
study = optuna.load_study(study_name='study_{}', storage=storage)

result = {{
    'best_params': study.best_params,
    'best_value': study.best_value,
    'trials': [
        {{
            'number': t.number,
            'params': t.params,
            'value': t.value,
            'state': str(t.state),
            'datetime_start': t.datetime_start.timestamp() if t.datetime_start else 0,
            'datetime_complete': t.datetime_complete.timestamp() if t.datetime_complete else 0,
        }}
        for t in study.trials
    ]
}}

print(json.dumps(result))
"#, storage_path, job_id))
        .output()
        .await
        .map_err(|e| Status::internal(format!("Failed to load study: {}", e)))?;

    let json_str = String::from_utf8_lossy(&output.stdout);
    let data: serde_json::Value = serde_json::from_str(&json_str)
        .map_err(|e| Status::internal(format!("Failed to parse study JSON: {}", e)))?;

    // Extract best parameters
    let best_params: HashMap<String, f32> = data["best_params"]
        .as_object()
        .unwrap()
        .iter()
        .map(|(k, v)| (k.clone(), v.as_f64().unwrap() as f32))
        .collect();

    let best_metrics = HashMap::from([
        ("sharpe_ratio".to_string(), data["best_value"].as_f64().unwrap() as f32)
    ]);

    // Convert trial history
    let trial_history: Vec<TrialResult> = data["trials"]
        .as_array()
        .unwrap()
        .iter()
        .map(|t| TrialResult {
            trial_number: t["number"].as_u64().unwrap() as u32,
            params: t["params"]
                .as_object()
                .unwrap()
                .iter()
                .map(|(k, v)| (k.clone(), v.as_f64().unwrap() as f32))
                .collect(),
            objective_value: t["value"].as_f64().unwrap_or(-999.0) as f32,
            state: TrialState::from_str(t["state"].as_str().unwrap()).unwrap().into(),
            started_at: t["datetime_start"].as_i64().unwrap_or(0),
            completed_at: t["datetime_complete"].as_i64().unwrap_or(0),
            metrics: HashMap::new(),
        })
        .collect();

    let current_trial = trial_history.len() as u32;

    Ok((best_params, best_metrics, trial_history, current_trial))
}
```

#### 4.3 StopTuningJob

```rust
pub async fn stop_tuning_job(
    &self,
    request: StopTuningJobRequest,
) -> Result<StopTuningJobResponse, Status> {
    let job_id = &request.job_id;

    // Remove from active jobs
    let mut child = self.tuning_jobs.write().await.remove(job_id)
        .ok_or_else(|| Status::not_found("Tuning job not found or already stopped"))?;

    // Send SIGTERM for graceful shutdown
    child.kill().await
        .map_err(|e| Status::internal(format!("Failed to kill process: {}", e)))?;

    // Wait for exit (max 30 seconds)
    let status = tokio::time::timeout(
        Duration::from_secs(30),
        child.wait()
    ).await;

    match status {
        Ok(Ok(exit_status)) => {
            log::info!("Tuning job {} exited: {}", job_id, exit_status);
        }
        Ok(Err(e)) => {
            log::error!("Failed to wait for tuning job {}: {}", job_id, e);
        }
        Err(_) => {
            log::warn!("Tuning job {} did not exit in 30s, force killed", job_id);
        }
    }

    // Update database
    sqlx::query!(
        "UPDATE tuning_jobs SET status = 'STOPPED', updated_at = NOW() WHERE id = $1",
        job_id
    )
    .execute(&self.db_pool)
    .await
    .map_err(|e| Status::internal(format!("Database error: {}", e)))?;

    Ok(StopTuningJobResponse {
        success: true,
        message: "Tuning job stopped successfully".to_string(),
        final_status: TuningJobStatus::TuningStopped.into(),
    })
}
```

#### 4.4 TrainModel (Internal)

```rust
pub async fn train_model(
    &self,
    request: TrainModelRequest,
) -> Result<TrainModelResponse, Status> {
    let start_time = std::time::Instant::now();

    log::info!(
        "TrainModel called by trial {}: model={}, params={:?}",
        request.trial_id,
        request.model_type,
        request.hyperparameters
    );

    // Convert map<string, float> to typed hyperparameters
    let hyperparameters = self.convert_hyperparameters(
        &request.model_type,
        &request.hyperparameters
    )?;

    // Execute training (reuse existing training logic)
    let result = self.execute_training(
        &request.model_type,
        hyperparameters,
        request.data_source,
        request.use_gpu
    ).await;

    let duration = start_time.elapsed().as_secs();

    match result {
        Ok(metrics) => {
            Ok(TrainModelResponse {
                success: true,
                sharpe_ratio: metrics.financial_metrics.sharpe_ratio,
                training_loss: metrics.final_loss,
                validation_metrics: metrics.validation_metrics,
                error_message: String::new(),
                training_duration_seconds: duration as i64,
            })
        }
        Err(e) => {
            log::error!("Training failed for trial {}: {}", request.trial_id, e);
            Ok(TrainModelResponse {
                success: false,
                sharpe_ratio: 0.0,
                training_loss: f32::INFINITY,
                validation_metrics: HashMap::new(),
                error_message: e.to_string(),
                training_duration_seconds: duration as i64,
            })
        }
    }
}
```

### 5. Database Schema

Add tuning jobs table:

```sql
-- migrations/XXX_create_tuning_jobs_table.sql

CREATE TABLE tuning_jobs (
    id TEXT PRIMARY KEY,
    model_type TEXT NOT NULL,
    num_trials INTEGER NOT NULL,
    status TEXT NOT NULL,
    storage_path TEXT NOT NULL,
    best_params JSONB,
    best_sharpe_ratio REAL,
    created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
    started_at TIMESTAMPTZ,
    completed_at TIMESTAMPTZ,
    updated_at TIMESTAMPTZ,
    description TEXT,
    tags JSONB,
    error_message TEXT
);

CREATE INDEX idx_tuning_jobs_status ON tuning_jobs(status);
CREATE INDEX idx_tuning_jobs_model_type ON tuning_jobs(model_type);
CREATE INDEX idx_tuning_jobs_created_at ON tuning_jobs(created_at DESC);
```

### 6. MinIO Configuration

Ensure MinIO mount point exists:

```bash
# Create studies directory in MinIO
docker exec -it foxhunt-minio-1 mkdir -p /data/studies

# Or via Docker Compose volume
# docker-compose.yml
services:
  ml_training_service:
    volumes:
      - minio_data:/minio
```

### 7. Testing

#### Unit Tests

```bash
cd services/ml_training_service
python3 -m pytest tests/test_hyperparameter_tuner.py -v
```

#### Integration Test

```bash
# 1. Start ML Training Service
cargo run -p ml_training_service

# 2. Run example tuning job
./scripts/example_tuning_job.sh

# 3. Verify study created
ls -lh /tmp/study_*.log
```

#### E2E Test via gRPC

```bash
# Use grpcurl or TLI
grpcurl -plaintext \
  -d '{
    "model_type": "TLOB",
    "num_trials": 10,
    "config_path": "tuning_config.yaml",
    "data_source": {"file_path": "/tmp/data.parquet"},
    "use_gpu": false
  }' \
  localhost:50054 ml_training.MLTrainingService/StartTuningJob
```

## Verification Checklist

- [ ] Python dependencies installed
- [ ] gRPC stubs generated
- [ ] Proto definitions complete
- [ ] Rust service methods implemented
- [ ] Database migration applied
- [ ] MinIO mount configured
- [ ] Unit tests passing
- [ ] Integration test successful
- [ ] E2E test via gRPC working

## Key Requirements Validation

✅ **Optuna 3.0+ with JournalStorage**: `requirements-tuner.txt` specifies `optuna>=3.0.0`

✅ **Sequential trials (n_jobs=1)**: Line 432 in `hyperparameter_tuner.py`:
```python
self.study.optimize(
    self.objective,
    n_trials=self.num_trials,
    n_jobs=1,  # CRITICAL: Sequential execution for 4GB VRAM
    ...
)
```

✅ **MedianPruner**: Lines 285-290 configure MedianPruner from `tuning_config.yaml`

✅ **pynvml for GPU monitoring**: Lines 48-53 initialize pynvml, `GPUMonitor` class uses direct library calls (NOT subprocess)

✅ **Graceful SIGTERM shutdown**: Lines 63-69 implement signal handlers

✅ **Read search spaces from YAML**: Line 317 loads `tuning_config.yaml`, lines 342-396 sample parameters

✅ **TrainModel gRPC calls**: Lines 165-228 implement `GRPCModelTrainer.train_model()`

✅ **Sharpe ratio as objective**: Lines 421-439 return `sharpe_ratio` from `TrainModelResponse`

✅ **MinIO persistence**: Line 281 creates `JournalFileStorage` with user-provided path

## Production Readiness

### Performance

- **Sequential trials**: Safe for 4GB VRAM (RTX 3050 Ti)
- **Crash recovery**: JournalStorage persists after each trial
- **GPU monitoring**: Pre-trial memory checks prevent OOM

### Reliability

- **Signal handling**: Graceful shutdown on SIGTERM/SIGINT
- **Error handling**: Failed trials return -999.0 (continue optimization)
- **gRPC timeouts**: 3600s (1 hour) for long training jobs

### Observability

- **Structured logging**: All major events logged with timestamps
- **Trial history**: Complete record in `GetTuningJobStatusResponse`
- **GPU metrics**: Memory usage logged after each trial

## Troubleshooting

See `HYPERPARAMETER_TUNING.md` Section "Troubleshooting" for common issues and solutions.

## Next Steps

1. ✅ Create hyperparameter tuner (COMPLETE)
2. ⏭️ Implement Rust service methods (use code above)
3. ⏭️ Apply database migration
4. ⏭️ Run integration tests
5. ⏭️ Test with real ML models
6. ⏭️ Benchmark 50-trial optimization
7. ⏭️ Document performance in production

## References

- `HYPERPARAMETER_TUNING.md` - Comprehensive documentation (21KB)
- `hyperparameter_tuner.py` - Main implementation (609 lines)
- `tuning_config.yaml` - Search space definitions
- `tests/test_hyperparameter_tuner.py` - Unit tests