Files
foxhunt/services/ml_training_service/tests/integration_tuning_test.rs
jgrusewski 3817b06f19 feat(ml_training_service): add JobSpawner field to MLTrainingServiceImpl
Wire JobSpawner into the gRPC service struct so that training jobs can
be persisted to PostgreSQL before being dispatched to K8s. This is the
first step toward a durable job queue that survives pod restarts.

- Add `job_spawner: Option<Arc<JobSpawner>>` to MLTrainingServiceImpl
- Extend `new()` constructor to accept the spawner parameter
- Add `DatabaseManager::pg_pool()` accessor for cheap PgPool cloning
- Construct JobSpawner in main.rs and pass `Some(job_spawner)` to service
- Add compile-time test verifying the field exists

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-27 22:29:10 +01:00

845 lines
29 KiB
Rust

//! Integration Tests for Hyperparameter Tuning
//!
//! Comprehensive end-to-end tests for the Optuna-based hyperparameter tuning subsystem,
//! covering:
//! - Single trial E2E flow (StartTuningJob → completion → checkpoint/study saved)
//! - Trial pruning (MedianPruner early stopping)
//! - Concurrent trials (sequential execution for single GPU)
//! - Error handling (invalid model, missing data, OOM)
//! - Progress streaming (StreamTuningProgress subscription)
//! - Crash recovery (service restart mid-job)
use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::Arc;
use std::time::Duration;
use tempfile::TempDir;
use tokio::time::timeout;
use tonic::{Request, Status};
use uuid::Uuid;
use ml_training_service::{
database::DatabaseManager,
orchestrator::TrainingOrchestrator,
service::{
proto::{
ml_training_service_server::MlTrainingService, DataSource, GetTuningJobStatusRequest,
StartTuningJobRequest, StopTuningJobRequest, StreamProgressRequest, TrainModelRequest,
TrialState as ProtoTrialState, TuningJobStatus as ProtoTuningJobStatus,
},
MLTrainingServiceImpl,
},
storage::{ModelStorageManager, StorageConfig},
tuning_manager::{TuningJobStatus, TuningManager},
};
/// Test helper to create database configuration
async fn create_test_database_config() -> config::database::DatabaseConfig {
let database_url = std::env::var("DATABASE_URL").unwrap_or_else(|_| {
"postgresql://foxhunt:foxhunt_dev_password@localhost:5432/foxhunt".to_string()
});
config::database::DatabaseConfig {
url: database_url.clone(),
max_connections: 5,
min_connections: 1,
connect_timeout: Duration::from_secs(10),
query_timeout: Duration::from_secs(30),
enable_query_logging: false,
application_name: Some("ml_training_tuning_test".to_string()),
pool: config::PoolConfig {
min_connections: 1,
max_connections: 5,
acquire_timeout_secs: 5,
max_lifetime_secs: 3600,
idle_timeout_secs: 300,
test_before_acquire: true,
database_url: database_url.clone(),
health_check_enabled: true,
health_check_interval_secs: 60,
},
transaction: config::TransactionConfig::default(),
}
}
/// Setup test service with temporary storage
async fn setup_test_service() -> (Arc<TuningManager>, Arc<MLTrainingServiceImpl>, TempDir) {
let temp_dir = TempDir::new().expect("Failed to create temp dir");
// Setup tuning manager
let tuner_script = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
.join("hyperparameter_tuner.py")
.to_string_lossy()
.to_string();
let working_dir = temp_dir
.path()
.join("tuning_jobs")
.to_string_lossy()
.to_string();
let tuning_manager = Arc::new(TuningManager::new(tuner_script, working_dir));
// Setup orchestrator
let ml_config = config::MLConfig::default();
let db_config = create_test_database_config().await;
let database = Arc::new(
DatabaseManager::new(&db_config)
.await
.expect("Failed to create database"),
);
let storage_config = StorageConfig {
storage_type: "local".to_string(),
local_base_path: Some(temp_dir.path().join("models")),
enable_compression: false,
};
let storage = Arc::new(
ModelStorageManager::new(storage_config)
.await
.expect("Failed to create storage"),
);
let orchestrator = Arc::new(
TrainingOrchestrator::new(ml_config.clone(), database, storage)
.await
.expect("Failed to create orchestrator"),
);
let tuning_mgr = Arc::new(tuning_manager.clone());
let promotion_mgr = Arc::new(ml_training_service::promotion_manager::PromotionManager::new());
let service = Arc::new(MLTrainingServiceImpl::new(
orchestrator,
tuning_mgr,
ml_config,
promotion_mgr,
None, // no K8s dispatcher in tests
None, // no job spawner in tests
));
(tuning_manager, service, temp_dir)
}
/// Create real tuning configuration file (NO MOCKS)
async fn create_real_tuning_config(path: &PathBuf) -> anyhow::Result<()> {
mod test_helpers;
test_helpers::create_real_tuning_config(path.as_path()).await
}
/// Create real training data for testing (NO MOCKS)
async fn create_real_training_data(path: &PathBuf) -> anyhow::Result<()> {
mod test_helpers;
test_helpers::create_real_training_data(path.as_path()).await
}
// ============================================================================
// Test 1: Single Trial E2E Flow
// ============================================================================
#[tokio::test]
#[ignore = "Run with: cargo test --test integration_tuning_test -- --ignored"]
async fn test_single_trial_e2e_flow() {
let (tuning_manager, _service, temp_dir) = setup_test_service().await;
// Create mock config and data
let config_path = temp_dir.path().join("tuning_config.yaml");
let data_path = temp_dir.path().join("training_data.parquet");
create_real_tuning_config(&config_path).await.unwrap();
create_real_training_data(&data_path).await.unwrap();
// Start tuning job with 1 trial
let mut tags = HashMap::new();
tags.insert("test".to_string(), "single_trial".to_string());
let job_id = tuning_manager
.start_tuning_job(
"TLOB".to_string(),
1, // Single trial
config_path.to_string_lossy().to_string(),
"Test single trial E2E".to_string(),
tags,
)
.await
.expect("Failed to start tuning job");
// Wait for job completion (timeout: 5 minutes)
let completion_timeout = Duration::from_secs(300);
let start_time = tokio::time::Instant::now();
loop {
if start_time.elapsed() > completion_timeout {
panic!("Job did not complete within 5 minutes");
}
tokio::time::sleep(Duration::from_secs(5)).await;
// Check job status
let job_status = tuning_manager
.get_tuning_job_status(job_id)
.await
.expect("Failed to get job status");
match job_status.status {
TuningJobStatus::Completed => {
// Verify job completed successfully
assert_eq!(job_status.current_trial, 1, "Should have completed 1 trial");
assert!(
!job_status.best_params.is_empty(),
"Should have best params"
);
assert!(
job_status.best_metrics.contains_key("sharpe_ratio"),
"Should have Sharpe ratio"
);
assert!(
!job_status.trial_history.is_empty(),
"Should have trial history"
);
// Verify checkpoint saved (would be in MinIO in production)
let checkpoint_dir = temp_dir.path().join("models").join(job_id.to_string());
// In mock test, checkpoint may not exist - verify manager tracked it
println!("✓ Job completed: {:?}", job_status);
break;
},
TuningJobStatus::Failed => {
panic!(
"Job failed unexpectedly: {:?}",
job_status.error_message.unwrap_or_default()
);
},
TuningJobStatus::Running | TuningJobStatus::Pending => {
// Continue waiting
continue;
},
TuningJobStatus::Stopped => {
panic!("Job was stopped unexpectedly");
},
}
}
}
// ============================================================================
// Test 2: Trial Pruning
// ============================================================================
#[tokio::test]
#[ignore = "Run with: cargo test --test integration_tuning_test -- --ignored"]
async fn test_trial_pruning() {
let (tuning_manager, _service, temp_dir) = setup_test_service().await;
// Create mock config with aggressive pruning
let config_path = temp_dir.path().join("pruning_config.yaml");
let pruning_config = r#"
model_type: "TLOB"
search_space:
learning_rate: [0.0001, 0.01]
batch_size: [32, 64]
objective:
metric: "sharpe_ratio"
direction: "maximize"
pruner:
type: "median"
n_startup_trials: 2 # Low threshold for faster pruning
n_warmup_steps: 5
sampler:
type: "tpe"
"#;
tokio::fs::write(&config_path, pruning_config)
.await
.unwrap();
let data_path = temp_dir.path().join("training_data.parquet");
create_real_training_data(&data_path).await.unwrap();
// Start tuning job with 10 trials
let job_id = tuning_manager
.start_tuning_job(
"TLOB".to_string(),
10,
config_path.to_string_lossy().to_string(),
"Test trial pruning".to_string(),
HashMap::new(),
)
.await
.expect("Failed to start tuning job");
// Wait for completion and verify pruned trials
let completion_timeout = Duration::from_secs(600); // 10 minutes
let start_time = tokio::time::Instant::now();
loop {
if start_time.elapsed() > completion_timeout {
panic!("Job did not complete within 10 minutes");
}
tokio::time::sleep(Duration::from_secs(10)).await;
let job_status = tuning_manager
.get_tuning_job_status(job_id)
.await
.expect("Failed to get job status");
match job_status.status {
TuningJobStatus::Completed | TuningJobStatus::Failed => {
// Verify at least one trial was pruned
let pruned_count = job_status
.trial_history
.iter()
.filter(|t| {
matches!(
t.state,
ml_training_service::tuning_manager::TrialState::Pruned
)
})
.count();
println!(
"✓ Pruned {} trials out of {}",
pruned_count,
job_status.trial_history.len()
);
assert!(
pruned_count > 0 || job_status.trial_history.len() < 10,
"Expected pruned trials or early termination"
);
break;
},
TuningJobStatus::Running | TuningJobStatus::Pending => continue,
TuningJobStatus::Stopped => panic!("Job was stopped unexpectedly"),
}
}
}
// ============================================================================
// Test 3: Concurrent Trials (Sequential for Single GPU)
// ============================================================================
#[tokio::test]
#[ignore = "Requires GPU"]
async fn test_concurrent_trials_sequential_execution() {
let (tuning_manager, _service, temp_dir) = setup_test_service().await;
let config_path = temp_dir.path().join("tuning_config.yaml");
let data_path = temp_dir.path().join("training_data.parquet");
create_real_tuning_config(&config_path).await.unwrap();
create_real_training_data(&data_path).await.unwrap();
// Start 3 separate tuning jobs
let mut job_ids = Vec::new();
for i in 0..3 {
let mut tags = HashMap::new();
tags.insert("batch".to_string(), "concurrent_test".to_string());
tags.insert("trial".to_string(), format!("{}", i));
let job_id = tuning_manager
.start_tuning_job(
"TLOB".to_string(),
1, // 1 trial each
config_path.to_string_lossy().to_string(),
format!("Concurrent test job {}", i),
tags,
)
.await
.expect("Failed to start tuning job");
job_ids.push(job_id);
tokio::time::sleep(Duration::from_millis(100)).await; // Stagger submissions
}
// Verify only 1 job runs at a time (sequential execution for single GPU)
let mut completion_count = 0;
let timeout_duration = Duration::from_secs(900); // 15 minutes
let start_time = tokio::time::Instant::now();
while completion_count < 3 {
if start_time.elapsed() > timeout_duration {
panic!("Not all jobs completed within 15 minutes");
}
tokio::time::sleep(Duration::from_secs(10)).await;
let mut running_count = 0;
let mut completed_count = 0;
for job_id in &job_ids {
if let Ok(status) = tuning_manager.get_tuning_job_status(*job_id).await {
match status.status {
TuningJobStatus::Running => running_count += 1,
TuningJobStatus::Completed => completed_count += 1,
_ => {},
}
}
}
// Verify sequential execution (at most 1 running at a time)
assert!(
running_count <= 1,
"Expected sequential execution, found {} jobs running simultaneously",
running_count
);
if completed_count > completion_count {
println!("✓ Completed {}/{} jobs", completed_count, 3);
completion_count = completed_count;
}
}
println!("✓ All 3 jobs completed sequentially");
}
// ============================================================================
// Test 4: Error Handling
// ============================================================================
#[tokio::test]
async fn test_error_handling_invalid_model_type() {
let (tuning_manager, _service, temp_dir) = setup_test_service().await;
let config_path = temp_dir.path().join("tuning_config.yaml");
create_real_tuning_config(&config_path).await.unwrap();
// Try to start job with invalid model type
let result = tuning_manager
.start_tuning_job(
"INVALID_MODEL".to_string(), // Invalid model type
1,
config_path.to_string_lossy().to_string(),
"Test invalid model".to_string(),
HashMap::new(),
)
.await;
// Should either fail immediately or fail during execution
match result {
Ok(job_id) => {
// If job started, verify it fails
tokio::time::sleep(Duration::from_secs(10)).await;
let status = tuning_manager.get_tuning_job_status(job_id).await.unwrap();
assert_eq!(
status.status,
TuningJobStatus::Failed,
"Job should fail with invalid model type"
);
},
Err(_) => {
// Immediate failure is also acceptable
println!("✓ Job rejected immediately with invalid model type");
},
}
}
#[tokio::test]
async fn test_error_handling_missing_training_data() {
let (tuning_manager, _service, temp_dir) = setup_test_service().await;
let config_path = temp_dir.path().join("tuning_config.yaml");
create_real_tuning_config(&config_path).await.unwrap();
// Note: NOT creating training data file
let result = tuning_manager
.start_tuning_job(
"TLOB".to_string(),
1,
config_path.to_string_lossy().to_string(),
"Test missing data".to_string(),
HashMap::new(),
)
.await;
// Should start but fail during execution
match result {
Ok(job_id) => {
tokio::time::sleep(Duration::from_secs(15)).await;
let status = tuning_manager.get_tuning_job_status(job_id).await.unwrap();
// May fail or still be running (mock mode might not detect missing data)
println!("✓ Job handled missing data: status={:?}", status.status);
},
Err(e) => {
println!("✓ Job failed to start with missing data: {}", e);
},
}
}
// ============================================================================
// Test 5: Progress Streaming
// ============================================================================
#[tokio::test]
#[ignore = "Integration test - run with --ignored"]
async fn test_progress_streaming() {
let (tuning_manager, _service, temp_dir) = setup_test_service().await;
let config_path = temp_dir.path().join("tuning_config.yaml");
let data_path = temp_dir.path().join("training_data.parquet");
create_real_tuning_config(&config_path).await.unwrap();
create_real_training_data(&data_path).await.unwrap();
// Start tuning job with 2 trials
let job_id = tuning_manager
.start_tuning_job(
"TLOB".to_string(),
2,
config_path.to_string_lossy().to_string(),
"Test progress streaming".to_string(),
HashMap::new(),
)
.await
.expect("Failed to start tuning job");
// Subscribe to progress updates
let mut progress_rx = tuning_manager.subscribe_to_progress(job_id);
// Collect progress updates
let mut received_updates = 0;
let timeout_duration = Duration::from_secs(600); // 10 minutes
loop {
match timeout(Duration::from_secs(30), progress_rx.recv()).await {
Ok(Ok(event)) => {
// Filter for this job's events
if event.job_id == job_id {
received_updates += 1;
println!(
"✓ Progress update {}: trial {}/{} (Sharpe: {:.2})",
received_updates,
event.current_trial,
event.total_trials,
event.trial_sharpe
);
// Check for job completion event
if event.update_type
== ml_training_service::tuning_manager::ProgressUpdateType::JobComplete
{
println!("✓ Job completed, stream should close soon");
break;
}
}
},
Ok(Err(e)) => {
// Channel closed (expected after job completion)
println!("✓ Progress stream closed: {:?}", e);
break;
},
Err(_) => {
// Timeout waiting for update - check job status
if let Ok(status) = tuning_manager.get_tuning_job_status(job_id).await {
if matches!(
status.status,
TuningJobStatus::Completed | TuningJobStatus::Failed
) {
println!("✓ Job completed");
break;
}
}
},
}
// Safety timeout
if tokio::time::Instant::now().elapsed() > timeout_duration {
panic!("Progress streaming test exceeded 10 minute timeout");
}
}
assert!(
received_updates >= 2,
"Expected at least 2 progress updates (trial completions), got {}",
received_updates
);
}
// ============================================================================
// Test 6: Crash Recovery
// ============================================================================
#[tokio::test]
#[ignore = "Integration test - run with --ignored"]
async fn test_crash_recovery() {
let temp_dir = TempDir::new().expect("Failed to create temp dir");
let config_path = temp_dir.path().join("tuning_config.yaml");
let data_path = temp_dir.path().join("training_data.parquet");
create_real_tuning_config(&config_path).await.unwrap();
create_real_training_data(&data_path).await.unwrap();
let tuner_script = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
.join("hyperparameter_tuner.py")
.to_string_lossy()
.to_string();
let working_dir = temp_dir
.path()
.join("tuning_jobs")
.to_string_lossy()
.to_string();
// Phase 1: Start job with 5 trials
let job_id = {
let tuning_manager = TuningManager::new(tuner_script.clone(), working_dir.clone());
let job_id = tuning_manager
.start_tuning_job(
"TLOB".to_string(),
5,
config_path.to_string_lossy().to_string(),
"Test crash recovery".to_string(),
HashMap::new(),
)
.await
.expect("Failed to start tuning job");
// Wait for trial 2 to complete
loop {
tokio::time::sleep(Duration::from_secs(10)).await;
if let Ok(status) = tuning_manager.get_tuning_job_status(job_id).await {
if status.current_trial >= 2 {
println!("✓ Completed trial 2, simulating crash...");
break;
}
}
}
// Simulate crash by stopping the job
tuning_manager
.stop_tuning_job(job_id, "Simulated crash".to_string())
.await
.ok();
job_id
};
// Drop tuning_manager to simulate crash
tokio::time::sleep(Duration::from_secs(2)).await;
// Phase 2: Restart service and resume job
let tuning_manager = TuningManager::new(tuner_script, working_dir.clone());
// Try to load job status from disk (Optuna study should be persisted)
let status_path = PathBuf::from(&working_dir)
.join(job_id.to_string())
.join("status.json");
if tokio::fs::metadata(&status_path).await.is_ok() {
println!("✓ Found persisted job status");
// Resume job (would be automatic in production with MinIO)
let resumed_job_id = tuning_manager
.start_tuning_job(
"TLOB".to_string(),
5,
config_path.to_string_lossy().to_string(),
"Resumed after crash".to_string(),
HashMap::new(),
)
.await
.expect("Failed to resume job");
// Verify job continues from trial 3
tokio::time::sleep(Duration::from_secs(30)).await;
if let Ok(status) = tuning_manager.get_tuning_job_status(resumed_job_id).await {
// In mock mode, may restart from 0, but in production with Optuna+MinIO
// it would resume from trial 3
println!(
"✓ Resumed job at trial {} (production: should resume at trial 3)",
status.current_trial
);
}
} else {
println!("⚠ Status file not persisted (expected in mock mode)");
}
}
// ============================================================================
// Test 7: TrainModel gRPC Endpoint (Called by Optuna)
// ============================================================================
#[tokio::test]
async fn test_train_model_grpc_endpoint() {
let (_tuning_manager, service, temp_dir) = setup_test_service().await;
let data_path = temp_dir.path().join("training_data.parquet");
create_real_training_data(&data_path).await.unwrap();
// Simulate Optuna calling TrainModel for a single trial
let mut hyperparameters = HashMap::new();
hyperparameters.insert("learning_rate".to_string(), 0.001);
hyperparameters.insert("batch_size".to_string(), 64.0);
hyperparameters.insert("hidden_dim".to_string(), 256.0);
let request = Request::new(TrainModelRequest {
model_type: "TLOB".to_string(),
hyperparameters,
data_source: Some(DataSource {
source: Some(
ml_training_service::service::proto::data_source::Source::FilePath(
data_path.to_string_lossy().to_string(),
),
),
start_time: 0,
end_time: 0,
}),
use_gpu: false, // Use CPU for testing
trial_id: Uuid::new_v4().to_string(),
});
let response = service.train_model(request).await;
match response {
Ok(resp) => {
let result = resp.into_inner();
println!("✓ TrainModel succeeded:");
println!(" - Sharpe ratio: {}", result.sharpe_ratio);
println!(" - Training loss: {}", result.training_loss);
println!(" - Duration: {}s", result.training_duration_seconds);
// Verify response structure
assert!(
result.sharpe_ratio.is_finite(),
"Sharpe ratio should be valid"
);
assert!(
result.training_loss >= 0.0,
"Training loss should be non-negative"
);
},
Err(e) => {
// In mock mode, may fail due to missing real data - verify error is reasonable
println!("⚠ TrainModel failed (expected in mock mode): {:?}", e);
assert!(
matches!(
e.code(),
tonic::Code::Internal | tonic::Code::InvalidArgument
),
"Error code should be reasonable"
);
},
}
}
// ============================================================================
// Test 8: Stop Tuning Job Mid-Execution
// ============================================================================
#[tokio::test]
#[ignore = "Integration test - run with --ignored"]
async fn test_stop_tuning_job() {
let (tuning_manager, _service, temp_dir) = setup_test_service().await;
let config_path = temp_dir.path().join("tuning_config.yaml");
let data_path = temp_dir.path().join("training_data.parquet");
create_real_tuning_config(&config_path).await.unwrap();
create_real_training_data(&data_path).await.unwrap();
// Start job with 10 trials
let job_id = tuning_manager
.start_tuning_job(
"TLOB".to_string(),
10,
config_path.to_string_lossy().to_string(),
"Test stop mid-execution".to_string(),
HashMap::new(),
)
.await
.expect("Failed to start tuning job");
// Wait for at least 1 trial to complete
loop {
tokio::time::sleep(Duration::from_secs(10)).await;
if let Ok(status) = tuning_manager.get_tuning_job_status(job_id).await {
if status.current_trial >= 1 {
println!("✓ Completed trial 1, stopping job...");
break;
}
}
}
// Stop the job
let stop_result = tuning_manager
.stop_tuning_job(job_id, "User requested stop".to_string())
.await;
assert!(stop_result.is_ok(), "Failed to stop job");
// Verify job status is Stopped
tokio::time::sleep(Duration::from_secs(5)).await;
let final_status = tuning_manager
.get_tuning_job_status(job_id)
.await
.expect("Failed to get job status");
assert_eq!(
final_status.status,
TuningJobStatus::Stopped,
"Job should be in Stopped state"
);
println!("✓ Job stopped after {} trials", final_status.current_trial);
}
// ============================================================================
// Test 9: GetTuningJobStatus with Nonexistent Job
// ============================================================================
#[tokio::test]
async fn test_get_status_nonexistent_job() {
let (tuning_manager, _service, _temp_dir) = setup_test_service().await;
let fake_job_id = Uuid::new_v4();
let result = tuning_manager.get_tuning_job_status(fake_job_id).await;
assert!(result.is_err(), "Should fail for nonexistent job");
println!("✓ Correctly rejected nonexistent job ID");
}
// ============================================================================
// Test 10: Validation of Tuning Parameters
// ============================================================================
#[tokio::test]
async fn test_parameter_validation() {
let (tuning_manager, _service, temp_dir) = setup_test_service().await;
let config_path = temp_dir.path().join("tuning_config.yaml");
create_real_tuning_config(&config_path).await.unwrap();
// Test 1: Zero trials
let result = tuning_manager
.start_tuning_job(
"TLOB".to_string(),
0, // Invalid: zero trials
config_path.to_string_lossy().to_string(),
"Test zero trials".to_string(),
HashMap::new(),
)
.await;
// Should be rejected (either immediately or during validation)
match result {
Ok(job_id) => {
tokio::time::sleep(Duration::from_secs(5)).await;
let status = tuning_manager.get_tuning_job_status(job_id).await;
// May fail or be invalid
println!("✓ Zero trials handled: {:?}", status);
},
Err(_) => {
println!("✓ Zero trials rejected immediately");
},
}
// Test 2: Empty model type
let result = tuning_manager
.start_tuning_job(
"".to_string(), // Invalid: empty model type
1,
config_path.to_string_lossy().to_string(),
"Test empty model".to_string(),
HashMap::new(),
)
.await;
assert!(result.is_err(), "Empty model type should be rejected");
println!("✓ Empty model type rejected");
}