foxhunt/services/ml_training_service/TRIAL_EXECUTOR_USAGE.md

Trial Executor Usage Guide

Overview

The TrialExecutor provides a fixed-size thread pool for concurrent execution of Optuna hyperparameter tuning trials with automatic GPU resource management. It ensures efficient GPU utilization while preventing resource contention.

Architecture

┌─────────────────────────────────────────────────────────┐
│                    Optuna Subprocess                     │
│              (Python hyperparameter_tuner.py)            │
└───────────────────────┬─────────────────────────────────┘
                        │ Submit trial
                        ▼
┌─────────────────────────────────────────────────────────┐
│                   TrialExecutor                          │
│                   (Fixed Pool)                           │
│  ┌────────────┐  ┌────────────┐  ┌────────────┐       │
│  │ Worker 0   │  │ Worker 1   │  │ Worker N   │       │
│  │ GPU 0      │  │ GPU 1      │  │ GPU N      │       │
│  └─────┬──────┘  └─────┬──────┘  └─────┬──────┘       │
└────────┼────────────────┼────────────────┼──────────────┘
         │                │                │
         ▼                ▼                ▼
    ┌────────────────────────────────────────────┐
    │     ML Training Service (gRPC)             │
    │          TrainModel RPC                    │
    └────────────────────────────────────────────┘

Key Components

1. Fixed Pool Size: Number of workers = Number of GPUs
2. Worker Thread: Tokio task that processes trials sequentially
3. GPU Assignment: Each worker owns a dedicated GPU (via CUDA_VISIBLE_DEVICES)
4. Trial Queue: MPSC channel for FIFO trial scheduling
5. Resource Tracking: Per-worker statistics and monitoring

GPU Detection

The executor automatically detects available GPUs:

// Automatic GPU detection (reads CUDA_VISIBLE_DEVICES)
let executor = TrialExecutor::new("http://localhost:50054".to_string()).await?;
// Pool size = GPU count (defaults to 1 if detection fails)

// Manual pool size override
let executor = TrialExecutor::with_pool_size("http://localhost:50054".to_string(), 2).await?;

GPU Detection Logic

1. Check CUDA_VISIBLE_DEVICES environment variable
2. Parse comma-separated device IDs (e.g., "0,1,2" → 3 GPUs)
3. Default to 1 GPU if not set or parsing fails

Example Environment Variables:

# Single GPU (RTX 3050 Ti)
export CUDA_VISIBLE_DEVICES=0

# Multi-GPU system
export CUDA_VISIBLE_DEVICES=0,1,2,3

# Let executor detect automatically
unset CUDA_VISIBLE_DEVICES  # Defaults to 1

Basic Usage

1. Initialize Executor

use ml_training_service::trial_executor::TrialExecutor;
use anyhow::Result;

#[tokio::main]
async fn main() -> Result<()> {
    // Create executor with automatic GPU detection
    let executor = TrialExecutor::new("http://localhost:50054".to_string()).await?;

    println!("Pool size: {}", executor.pool_size());

    // ... use executor ...

    // Graceful shutdown (60-second timeout)
    executor.shutdown().await?;
    Ok(())
}

2. Submit Trials

use std::collections::HashMap;
use ml_training_service::service::proto::DataSource;

// Prepare trial parameters
let mut hyperparameters = HashMap::new();
hyperparameters.insert("learning_rate".to_string(), 0.001);
hyperparameters.insert("batch_size".to_string(), 64.0);
hyperparameters.insert("hidden_size".to_string(), 128.0);

let data_source = DataSource {
    // ... data source configuration ...
};

// Submit trial (returns oneshot receiver for result)
let result_rx = executor
    .submit_trial(
        "trial_001".to_string(),
        "TLOB".to_string(),
        hyperparameters,
        data_source,
        true, // use_gpu
    )
    .await?;

// Wait for result
let result = result_rx.await??;
println!("Trial complete: success={}, sharpe_ratio={:.4}",
         result.success, result.sharpe_ratio);

3. Monitor Pool Statistics

// Get current pool statistics
let stats = executor.get_stats().await;

println!("Pool size: {}", stats.pool_size);
println!("Active trials: {}", stats.active_trials);
println!("Total completed: {}", stats.total_completed);
println!("Total failed: {}", stats.total_failed);

// Per-worker statistics
for worker in &stats.worker_stats {
    println!("Worker {} (GPU {}): completed={}, failed={}, oom_errors={}, busy={}",
             worker.worker_id, worker.gpu_id, worker.trials_completed,
             worker.trials_failed, worker.oom_errors, worker.is_busy);
}

Integration with Optuna

The trial executor is designed to be called from Optuna's objective function:

Python Optuna Integration

# hyperparameter_tuner.py

import optuna
import grpc
from ml_training_service_pb2 import TrainModelRequest
from ml_training_service_pb2_grpc import MlTrainingServiceStub

def objective(trial: optuna.Trial) -> float:
    """Optuna objective function that calls ML Training Service via gRPC."""

    # Sample hyperparameters
    hyperparameters = {
        "learning_rate": trial.suggest_float("learning_rate", 1e-5, 1e-2, log=True),
        "batch_size": trial.suggest_int("batch_size", 16, 128, step=16),
        "hidden_size": trial.suggest_int("hidden_size", 64, 512, step=64),
        "dropout": trial.suggest_float("dropout", 0.1, 0.5),
    }

    # Connect to ML Training Service
    channel = grpc.insecure_channel("localhost:50054")
    stub = MlTrainingServiceStub(channel)

    # Call TrainModel RPC (routed through TrialExecutor)
    request = TrainModelRequest(
        trial_id=str(trial.number),
        model_type="TLOB",
        hyperparameters=hyperparameters,
        use_gpu=True,
        # ... data source ...
    )

    response = stub.TrainModel(request)

    if not response.success:
        raise optuna.TrialPruned(f"Training failed: {response.error_message}")

    # Return Sharpe ratio as objective (maximize)
    return response.sharpe_ratio

# Create study
study = optuna.create_study(
    direction="maximize",
    study_name="tlob_hyperparameter_tuning",
)

# Run optimization
study.optimize(objective, n_trials=100)

print(f"Best trial: {study.best_trial.number}")
print(f"Best Sharpe ratio: {study.best_value:.4f}")
print(f"Best params: {study.best_params}")

Rust Service Integration

use ml_training_service::trial_executor::TrialExecutor;
use ml_training_service::tuning_manager::TuningManager;
use std::sync::Arc;

pub struct MLTrainingServiceImpl {
    trial_executor: Arc<TrialExecutor>,
    tuning_manager: Arc<TuningManager>,
}

impl MLTrainingServiceImpl {
    pub async fn new(grpc_endpoint: String) -> Result<Self> {
        // Initialize trial executor
        let trial_executor = Arc::new(
            TrialExecutor::new(grpc_endpoint).await?
        );

        // Initialize tuning manager
        let tuning_manager = Arc::new(
            TuningManager::new(
                "scripts/hyperparameter_tuner.py".to_string(),
                "/tmp/tuning_jobs".to_string(),
            )
        );

        Ok(Self {
            trial_executor,
            tuning_manager,
        })
    }

    // TrainModel RPC handler (called by Optuna)
    async fn train_model(
        &self,
        request: TrainModelRequest,
    ) -> Result<TrainModelResponse> {
        // Submit trial to executor
        let result_rx = self.trial_executor
            .submit_trial(
                request.trial_id,
                request.model_type,
                request.hyperparameters,
                request.data_source.unwrap(),
                request.use_gpu,
            )
            .await?;

        // Wait for result
        let result = result_rx.await??;

        Ok(TrainModelResponse {
            success: result.success,
            sharpe_ratio: result.sharpe_ratio,
            training_loss: result.training_loss,
            validation_metrics: result.validation_metrics,
            error_message: result.error_message,
            training_duration_seconds: result.training_duration_seconds,
        })
    }
}

Resource Management

Worker Lifecycle

1. Initialization: Worker spawned with dedicated GPU ID
2. Waiting: Worker blocks on trial queue (1-second timeout)
3. Execution: Worker runs trial with CUDA_VISIBLE_DEVICES set
4. Completion: Worker updates statistics and returns result
5. Shutdown: Worker finishes current trial and exits

GPU Isolation

Each worker sets CUDA_VISIBLE_DEVICES for its subprocess:

// Worker 0 → GPU 0
std::env::set_var("CUDA_VISIBLE_DEVICES", "0");

// Worker 1 → GPU 1
std::env::set_var("CUDA_VISIBLE_DEVICES", "1");

Note: This only affects subprocesses spawned by the TrainModel gRPC call, not the current process.

Graceful Shutdown

// Shutdown with 60-second timeout
executor.shutdown().await?;

// Shutdown behavior:
// 1. Cancel shutdown token (signals workers to stop)
// 2. Wait for current trials to finish (max 60s)
// 3. Force kill workers after timeout

Error Handling

OOM (Out of Memory) Errors

The executor tracks OOM errors per worker:

let stats = executor.get_stats().await;
for worker in &stats.worker_stats {
    if worker.oom_errors > 0 {
        warn!("Worker {} experienced {} OOM errors",
              worker.worker_id, worker.oom_errors);
    }
}

Worker Crash Recovery

Workers are not automatically restarted on crash. The pool remains at reduced capacity until manual intervention.

Future Enhancement: Add worker restart logic:

// TODO: Implement worker crash recovery
// - Detect worker exit via JoinHandle
// - Spawn replacement worker
// - Maintain pool_size workers

Trial Timeout

Trials timeout after 30 minutes:

// Built into execute_trial_with_gpu
match timeout(Duration::from_secs(1800), client.train_model(request)).await {
    Ok(Ok(resp)) => /* success */,
    Ok(Err(e)) => /* gRPC error */,
    Err(_) => /* timeout after 30 minutes */,
}

Performance Tuning

Pool Size Recommendations

Hardware	Pool Size	Rationale
Single GPU (RTX 3050 Ti)	1	Sequential trials (validated in research)
2 GPUs	2	Parallel trials, 1 per GPU
4 GPUs	4	Parallel trials, 1 per GPU
8 GPUs	8	Parallel trials, 1 per GPU

Rule of Thumb: pool_size = num_gpus for maximum throughput without contention.

Memory Management

• Trial Isolation: Each trial runs in separate subprocess with dedicated GPU
• No Sharing: Workers never share GPU memory
• OOM Detection: Tracked per worker, enables batch size reduction strategies

Monitoring

Add Prometheus metrics (future enhancement):

// TODO: Add Prometheus metrics
// - trial_executor_pool_size (gauge)
// - trial_executor_active_trials (gauge)
// - trial_executor_trials_completed_total (counter)
// - trial_executor_trials_failed_total (counter)
// - trial_executor_oom_errors_total (counter)
// - trial_executor_trial_duration_seconds (histogram)

Testing

Unit Tests

# Run trial executor tests
cargo test -p ml_training_service --test trial_executor_test

# Run with logging
RUST_LOG=debug cargo test -p ml_training_service --test trial_executor_test -- --nocapture

Integration Tests

# Start ML Training Service
cargo run -p ml_training_service

# Run Optuna hyperparameter tuning
python scripts/hyperparameter_tuner.py --model-type TLOB --num-trials 10

Load Testing

# Simulate 100 concurrent trials
python scripts/optuna_load_test.py --num-trials 100 --num-workers 4

Configuration

Environment Variables

# GPU configuration
export CUDA_VISIBLE_DEVICES=0,1,2,3  # Multi-GPU

# gRPC endpoint
export ML_TRAINING_SERVICE_ENDPOINT=http://localhost:50054

# Logging
export RUST_LOG=info,ml_training_service::trial_executor=debug

Runtime Configuration

// Custom pool size (override GPU detection)
let executor = TrialExecutor::with_pool_size(
    "http://localhost:50054".to_string(),
    2  // 2 workers even if 4 GPUs detected
).await?;

Troubleshooting

Issue: Pool size is 1 but I have multiple GPUs

Solution: Set CUDA_VISIBLE_DEVICES environment variable:

export CUDA_VISIBLE_DEVICES=0,1,2,3

Issue: Workers stuck in busy state

Symptoms: worker.is_busy = true for extended period

Debugging:

# Check if gRPC service is responding
grpcurl -plaintext localhost:50054 list

# Check GPU utilization
nvidia-smi

Issue: High OOM error rate

Solution: Reduce batch size in hyperparameters:

# Optuna objective function
hyperparameters = {
    "batch_size": trial.suggest_int("batch_size", 8, 32, step=8),  # Reduced range
}

Issue: Shutdown timeout

Symptoms: Shutdown takes 60+ seconds

Cause: Workers still running trials

Solution: Wait for active trials to complete before shutting down:

// Check active trials
let stats = executor.get_stats().await;
if stats.active_trials > 0 {
    warn!("Waiting for {} active trials to complete", stats.active_trials);
    tokio::time::sleep(Duration::from_secs(10)).await;
}

executor.shutdown().await?;

Future Enhancements

1. Worker Crash Recovery: Automatically restart crashed workers
2. Batch Size Reduction: Retry OOM trials with reduced batch size
3. Prometheus Metrics: Real-time monitoring and alerting
4. Trial Prioritization: Priority queue instead of FIFO
5. Multi-Node Distribution: Distribute trials across multiple machines
6. GPU Memory Profiling: Track per-trial memory usage
7. Adaptive Timeout: Dynamic timeout based on model type and data size

References

• Research Validation: Single GPU (RTX 3050 Ti) validated for sequential trial execution
• gRPC Protocol: services/ml_training_service/proto/ml_training.proto
• Optuna Documentation: https://optuna.readthedocs.io/
• CUDA Documentation: https://docs.nvidia.com/cuda/