foxhunt/docs/archive/ml_models/HYPERPARAMETER_TUNING_IMPLEMENTATION.md

Hyperparameter Tuning Implementation for ML Training Service

Summary

This document describes the complete implementation of hyperparameter tuning handlers for the Foxhunt HFT ML Training Service. The implementation follows the requirements precisely:

1. start_tuning_job: Generates unique job_id (UUID), spawns Optuna controller subprocess, returns immediately (async execution)
2. get_tuning_job_status: Queries job status from internal state/disk (written by Optuna), returns current trial, best params, metrics
3. stop_tuning_job: Sends SIGTERM to Optuna subprocess, waits for graceful shutdown (5s timeout), returns final status

Architecture Overview

┌─────────────────────────────────────────────────────────────┐
│                  ML Training Service (gRPC)                  │
├─────────────────────────────────────────────────────────────┤
│                                                               │
│  ┌───────────────────────────────────────────────┐          │
│  │   MLTrainingServiceImpl                       │          │
│  │   (existing service)                          │          │
│  │   - start_training()                          │          │
│  │   - train_model() [INTERNAL]                  │          │
│  │   ...                                         │          │
│  └───────────────────────────────────────────────┘          │
│                         │                                    │
│                         │ delegates to                       │
│                         ▼                                    │
│  ┌───────────────────────────────────────────────┐          │
│  │   TuningHandlers (new)                        │          │
│  │   - start_tuning_job()                        │          │
│  │   - get_tuning_job_status()                   │          │
│  │   - stop_tuning_job()                         │          │
│  └───────────────────────────────────────────────┘          │
│                         │                                    │
│                         │ uses                               │
│                         ▼                                    │
│  ┌───────────────────────────────────────────────┐          │
│  │   TuningManager (new)                         │          │
│  │   - Job state: Arc<RwLock<HashMap>>           │          │
│  │   - Process handles                           │          │
│  │   - Spawns/monitors Optuna subprocess         │          │
│  └───────────────────────────────────────────────┘          │
│                         │                                    │
└─────────────────────────┼────────────────────────────────────┘
                          │
                          │ spawns
                          ▼
                  ┌──────────────────┐
                  │  Optuna Process  │
                  │  (Python)        │
                  │  hyperparameter_ │
                  │  tuner.py        │
                  └──────────────────┘
                          │
                          │ writes status to
                          ▼
                  ┌──────────────────┐
                  │  File System     │
                  │  {working_dir}/  │
                  │  {job_id}/       │
                  │  status.json     │
                  └──────────────────┘

Files Created

1. /home/jgrusewski/Work/foxhunt/services/ml_training_service/src/tuning_manager.rs

Purpose: Core hyperparameter tuning orchestration logic

Key Components:

• TuningManager: Manages tuning jobs and subprocess lifecycle
• TuningJob: Job metadata (id, model_type, status, trials, best_params, etc.)
• ProcessHandle: Subprocess handle for running Optuna process
• Status tracking via in-memory HashMap + disk persistence

Key Methods:

pub async fn start_tuning_job(
    &self,
    model_type: String,
    num_trials: u32,
    config_path: String,
    description: String,
    tags: HashMap<String, String>,
) -> Result<Uuid>

pub async fn get_tuning_job_status(&self, job_id: Uuid) -> Result<TuningJob>

pub async fn stop_tuning_job(&self, job_id: Uuid, reason: String) -> Result<()>

Subprocess Management:

• Spawns Python Optuna process with command:

  python3 hyperparameter_tuner.py \
      --job-id <UUID> \
      --model-type <MODEL> \
      --num-trials <N> \
      --config-path <PATH> \
      --output-dir <WORKING_DIR>/<JOB_ID>
  

• SIGTERM graceful shutdown: Uses nix::sys::signal::kill() on Unix
• 5-second timeout: tokio::time::timeout(Duration::from_secs(5), ...)
• Monitors process completion via async polling loop (5s interval)

State Management:

jobs: Arc<RwLock<HashMap<Uuid, TuningJob>>>     // In-memory state
processes: Arc<RwLock<HashMap<Uuid, ProcessHandle>>>  // Running processes

Disk Persistence:

• Optuna subprocess writes {working_dir}/{job_id}/status.json
• Format: JSON-serialized TuningJob struct
• Manager polls this file every 5 seconds
• Used for status queries and recovery

2. /home/jgrusewski/Work/foxhunt/services/ml_training_service/src/grpc_tuning_handlers.rs

Purpose: gRPC handler implementations for tuning operations

Key Components:

• TuningHandlers: Wrapper around TuningManager for gRPC handlers
• Proto message conversions (internal → protobuf)
• Input validation

Handlers:

start_tuning_job

pub async fn start_tuning_job(
    &self,
    request: Request<StartTuningJobRequest>,
) -> Result<Response<StartTuningJobResponse>, Status>

• Validates input (model_type, num_trials, config_path)
• Generates UUID via TuningManager::start_tuning_job()
• Returns immediately with TuningRunning status
• Subprocess spawned asynchronously in background

get_tuning_job_status

pub async fn get_tuning_job_status(
    &self,
    request: Request<GetTuningJobStatusRequest>,
) -> Result<Response<GetTuningJobStatusResponse>, Status>

• Parses job_id UUID
• Queries TuningManager::get_tuning_job_status()
• Returns:
  • current_trial, total_trials
  • best_params: HashMap<String, f32>
  • best_metrics: HashMap<String, f32> (sharpe_ratio, training_loss, etc.)
  • trial_history: Vec<TrialResult>
  • timestamps (started_at, updated_at)

stop_tuning_job

pub async fn stop_tuning_job(
    &self,
    request: Request<StopTuningJobRequest>,
) -> Result<Response<StopTuningJobResponse>, Status>

• Sends SIGTERM to Optuna process
• Waits up to 5 seconds for graceful shutdown
• Updates job status to Stopped
• Returns final status with completion message

3. Updated Files

/home/jgrusewski/Work/foxhunt/services/ml_training_service/src/lib.rs

Added module exports:

pub mod grpc_tuning_handlers;
pub mod tuning_manager;

/home/jgrusewski/Work/foxhunt/services/ml_training_service/Cargo.toml

Added Unix signal handling dependency:

[target.'cfg(unix)'.dependencies]
nix = { version = "0.29", features = ["signal"] }

Integration with Existing Service

Required Changes to service.rs

Add TuningHandlers to MLTrainingServiceImpl:

// In service.rs, update the struct:
pub struct MLTrainingServiceImpl {
    orchestrator: Arc<TrainingOrchestrator>,
    tuning_handlers: Arc<TuningHandlers>,  // ADD THIS
    config: MLConfig,
}

impl MLTrainingServiceImpl {
    pub fn new(
        orchestrator: Arc<TrainingOrchestrator>,
        tuning_manager: Arc<TuningManager>,  // ADD THIS PARAMETER
        config: MLConfig
    ) -> Self {
        let tuning_handlers = Arc::new(TuningHandlers::new(tuning_manager));
        Self {
            orchestrator,
            tuning_handlers,  // INITIALIZE
            config,
        }
    }
}

Add trait implementations (add to impl MlTrainingService for MLTrainingServiceImpl):

#[tonic::async_trait]
impl MlTrainingService for MLTrainingServiceImpl {
    // ... existing methods ...

    /// Start hyperparameter tuning job
    async fn start_tuning_job(
        &self,
        request: Request<StartTuningJobRequest>,
    ) -> Result<Response<StartTuningJobResponse>, Status> {
        self.tuning_handlers.start_tuning_job(request).await
    }

    /// Get tuning job status
    async fn get_tuning_job_status(
        &self,
        request: Request<GetTuningJobStatusRequest>,
    ) -> Result<Response<GetTuningJobStatusResponse>, Status> {
        self.tuning_handlers.get_tuning_job_status(request).await
    }

    /// Stop tuning job
    async fn stop_tuning_job(
        &self,
        request: Request<StopTuningJobRequest>,
    ) -> Result<Response<StopTuningJobResponse>, Status> {
        self.tuning_handlers.stop_tuning_job(request).await
    }

    // train_model() already exists - used internally by Optuna subprocess
}

Required Changes to main.rs

Initialize TuningManager in the serve() function:

// In main.rs, after storage initialization:

// Initialize TuningManager
let tuner_script_path = std::env::var("TUNER_SCRIPT_PATH")
    .unwrap_or_else(|_| "/opt/foxhunt/scripts/hyperparameter_tuner.py".to_string());
let tuning_working_dir = std::env::var("TUNING_WORKING_DIR")
    .unwrap_or_else(|_| "/var/lib/foxhunt/tuning_jobs".to_string());

let tuning_manager = Arc::new(TuningManager::new(
    tuner_script_path,
    tuning_working_dir,
));

info!("TuningManager initialized");

// Create gRPC service with tuning_manager
let training_service = MLTrainingServiceImpl::new(
    Arc::clone(&orchestrator),
    Arc::clone(&tuning_manager),  // ADD THIS
    ml_config.clone()
);

Proto Definitions (Already Complete)

The proto definitions in /home/jgrusewski/Work/foxhunt/services/ml_training_service/proto/ml_training.proto are already complete with:

Messages

• StartTuningJobRequest / StartTuningJobResponse
• GetTuningJobStatusRequest / GetTuningJobStatusResponse
• StopTuningJobRequest / StopTuningJobResponse
• TrainModelRequest / TrainModelResponse (internal, for Optuna)
• TrialResult

Enums

• TuningJobStatus: TUNING_PENDING, TUNING_RUNNING, TUNING_COMPLETED, TUNING_FAILED, TUNING_STOPPED
• TrialState: TRIAL_RUNNING, TRIAL_COMPLETE, TRIAL_PRUNED, TRIAL_FAILED

Python Optuna Integration

Expected Python Script Structure

The Optuna subprocess (hyperparameter_tuner.py) should:

1. Parse command-line arguments:

   parser.add_argument('--job-id', required=True)
   parser.add_argument('--model-type', required=True)
   parser.add_argument('--num-trials', type=int, required=True)
   parser.add_argument('--config-path', required=True)
   parser.add_argument('--output-dir', required=True)
   

2. Create Optuna study:

   study = optuna.create_study(
       direction="maximize",  # Maximize Sharpe ratio
       study_name=f"{model_type}_{job_id}",
   )
   

3. Define objective function that calls back to gRPC:

   def objective(trial):
       # Sample hyperparameters from Optuna
       learning_rate = trial.suggest_float("learning_rate", 1e-5, 1e-2, log=True)
       batch_size = trial.suggest_categorical("batch_size", [32, 64, 128, 256])
       # ...
   
       # Call ML training service's train_model() gRPC method
       response = stub.TrainModel(TrainModelRequest(
           model_type=model_type,
           hyperparameters={
               "learning_rate": learning_rate,
               "batch_size": float(batch_size),
               # ...
           },
           trial_id=str(trial.number),
           use_gpu=True,
       ))
   
       return response.sharpe_ratio  # Optimization objective
   

4. Write status to disk after each trial:

   status = {
       "id": job_id,
       "model_type": model_type,
       "status": "Running",
       "num_trials": num_trials,
       "current_trial": trial.number,
       "best_params": study.best_params,
       "best_metrics": {
           "sharpe_ratio": study.best_value,
           "training_loss": study.best_trial.user_attrs.get("training_loss"),
       },
       "trial_history": [
           {
               "trial_number": t.number,
               "params": t.params,
               "objective_value": t.value,
               "state": str(t.state),
               ...
           }
           for t in study.trials
       ],
       ...
   }
   
   with open(f"{output_dir}/status.json", "w") as f:
       json.dump(status, f)
   

5. Handle SIGTERM gracefully:

   import signal
   
   def signal_handler(sig, frame):
       # Write final status
       status["status"] = "Stopped"
       with open(f"{output_dir}/status.json", "w") as f:
           json.dump(status, f)
       sys.exit(0)
   
   signal.signal(signal.SIGTERM, signal_handler)
   

Usage Example

Client Code (gRPC)

use tonic::Request;
use proto::{
    ml_training_service_client::MlTrainingServiceClient,
    StartTuningJobRequest, GetTuningJobStatusRequest,
};

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let mut client = MlTrainingServiceClient::connect("http://localhost:50054").await?;

    // 1. Start tuning job
    let response = client.start_tuning_job(Request::new(StartTuningJobRequest {
        model_type: "TLOB".to_string(),
        num_trials: 50,
        config_path: "/etc/foxhunt/tuning_configs/tlob_search_space.yaml".to_string(),
        use_gpu: true,
        description: "TLOB hyperparameter optimization for BTC/USD".to_string(),
        tags: {
            let mut tags = std::collections::HashMap::new();
            tags.insert("symbol".to_string(), "BTC/USD".to_string());
            tags.insert("env".to_string(), "production".to_string());
            tags
        },
        ..Default::default()
    })).await?;

    let job_id = response.into_inner().job_id;
    println!("Started tuning job: {}", job_id);

    // 2. Poll for status
    loop {
        tokio::time::sleep(tokio::time::Duration::from_secs(10)).await;

        let status_response = client.get_tuning_job_status(Request::new(
            GetTuningJobStatusRequest {
                job_id: job_id.clone(),
            }
        )).await?;

        let status = status_response.into_inner();
        println!(
            "Trial {}/{} - Best Sharpe: {:.4}",
            status.current_trial,
            status.total_trials,
            status.best_metrics.get("sharpe_ratio").unwrap_or(&0.0)
        );

        if status.status == TuningJobStatus::TuningCompleted as i32 {
            println!("Tuning complete!");
            println!("Best params: {:?}", status.best_params);
            break;
        }
    }

    Ok(())
}

Testing

Unit Tests

All modules include comprehensive unit tests:

tuning_manager.rs:

• test_tuning_job_creation() - Job initialization
• test_trial_result_creation() - Trial result structure
• test_tuning_manager_creation() - Manager creation
• test_get_nonexistent_job() - Error handling

grpc_tuning_handlers.rs:

• test_status_conversion() - TuningJobStatus conversions
• test_trial_state_conversion() - TrialState conversions
• test_tuning_job_validation() - Input validation

Integration Testing

# 1. Start ML training service
cargo run -p ml_training_service -- serve --dev

# 2. Create test tuning config (YAML)
cat > /tmp/test_search_space.yaml <<EOF
model_type: TLOB
search_space:
  learning_rate: [1e-5, 1e-2, log]
  batch_size: [32, 64, 128, 256]
  hidden_dim: [128, 256, 512]
  dropout_rate: [0.1, 0.3]
objectives:
  primary: sharpe_ratio
  minimize: training_loss
EOF

# 3. Test gRPC methods
grpcurl -plaintext \
  -d '{"model_type":"TLOB","num_trials":10,"config_path":"/tmp/test_search_space.yaml"}' \
  localhost:50054 \
  ml_training.MLTrainingService/StartTuningJob

# 4. Query status (use job_id from step 3)
grpcurl -plaintext \
  -d '{"job_id":"<UUID>"}' \
  localhost:50054 \
  ml_training.MLTrainingService/GetTuningJobStatus

# 5. Stop job
grpcurl -plaintext \
  -d '{"job_id":"<UUID>","reason":"Manual stop for testing"}' \
  localhost:50054 \
  ml_training.MLTrainingService/StopTuningJob

Environment Variables

Add to .env or docker-compose.yml:

# Tuning configuration
TUNER_SCRIPT_PATH=/opt/foxhunt/scripts/hyperparameter_tuner.py
TUNING_WORKING_DIR=/var/lib/foxhunt/tuning_jobs

Key Design Decisions

1. Async Execution: start_tuning_job returns immediately, Optuna runs in background subprocess
2. State Persistence: Dual storage (in-memory HashMap + disk JSON) for reliability
3. Process Isolation: Each tuning job = separate Python process (fail-isolation)
4. Graceful Shutdown: SIGTERM with 5s timeout (not SIGKILL)
5. Status Polling: Manager polls subprocess status file every 5s
6. UUID Job IDs: Generated by Rust service (not Python script)
7. Metrics Storage: Best params stored as HashMap<String, f32> for flexibility
8. Trial History: Full history persisted for post-analysis

Production Considerations

Error Handling

• Process spawn failures → job status = Failed
• SIGTERM timeout → warning logged, process may still be running
• Status file parsing errors → fall back to in-memory state
• Invalid job_id → gRPC Status::not_found

Resource Management

• No limit on concurrent tuning jobs (add semaphore if needed)
• Each job creates dedicated output directory
• Old job directories should be cleaned up periodically

Monitoring

• Log all tuning job lifecycle events (start, stop, complete, fail)
• Track subprocess PIDs for debugging
• Monitor disk usage in TUNING_WORKING_DIR

Security

• Validate config_path to prevent directory traversal
• Run Python subprocess with restricted permissions
• Consider sandboxing Optuna process (containers, seccomp)

Dependencies Added

• nix 0.29 (Unix only): For SIGTERM signal handling
• serde/serde_json: Already in workspace (for JSON persistence)
• tokio: Already in workspace (for async subprocess)
• uuid: Already in workspace (for job IDs)

File Locations Summary

services/ml_training_service/
├── src/
│   ├── tuning_manager.rs               # NEW: Core tuning logic
│   ├── grpc_tuning_handlers.rs         # NEW: gRPC handlers
│   ├── service.rs                      # MODIFY: Add tuning methods
│   ├── main.rs                         # MODIFY: Initialize TuningManager
│   └── lib.rs                          # MODIFIED: Export new modules
├── Cargo.toml                          # MODIFIED: Add nix dependency
└── proto/
    └── ml_training.proto               # UNCHANGED: Already complete

/opt/foxhunt/scripts/
└── hyperparameter_tuner.py             # TO BE IMPLEMENTED (Python Optuna)

/var/lib/foxhunt/tuning_jobs/
└── {job_id}/
    └── status.json                     # Written by Optuna subprocess

Next Steps

1. Implement Python Optuna script (hyperparameter_tuner.py)
2. Update service.rs with integration code (see above)
3. Update main.rs with TuningManager initialization
4. Add tuning config templates (YAML search space definitions)
5. Test end-to-end workflow with real model training
6. Add Prometheus metrics for tuning job monitoring
7. Implement cleanup script for old tuning job directories

Status

Implementation: ✅ Complete (Rust side)

• tuning_manager.rs: Fully implemented with tests
• grpc_tuning_handlers.rs: Fully implemented with tests
• lib.rs: Module exports added
• Cargo.toml: Dependencies added

Integration: ⚠️ Requires Changes

• service.rs: Add 3 trait methods + struct field
• main.rs: Initialize TuningManager
• Proto messages: Already complete

Python Optuna Script: ❌ Not Started

• Needs implementation based on spec above

Testing: ⚠️ Unit tests complete, E2E pending Python script