foxhunt/services/ml_training_service/TUNING_INTEGRATION_CHECKLIST.md

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

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

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

// 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

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

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)

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:

-- 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:

# 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

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

Integration Test

# 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

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

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