# 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**:
```rust
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:
  ```bash
  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**:
```rust
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`
```rust
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`
```rust
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`
```rust
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:
```rust
pub mod grpc_tuning_handlers;
pub mod tuning_manager;
```

#### `/home/jgrusewski/Work/foxhunt/services/ml_training_service/Cargo.toml`
Added Unix signal handling dependency:
```toml
[target.'cfg(unix)'.dependencies]
nix = { version = "0.29", features = ["signal"] }
```

## Integration with Existing Service

### Required Changes to `service.rs`

Add `TuningHandlers` to `MLTrainingServiceImpl`:

```rust
// 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`):

```rust
#[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:

```rust
// 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**:
   ```python
   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**:
   ```python
   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:
   ```python
   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:
   ```python
   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**:
   ```python
   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)

```rust
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

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

```bash
# 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