foxhunt/WAVE1_AGENT3_GRPC_API_DESIGN.md

# WAVE1_AGENT3: gRPC API Design for Multi-Model, Multi-Asset Training

**Agent**: WAVE1_AGENT3
**Date**: 2025-10-22
**Status**: Design Complete
**Objective**: Review and enhance gRPC API for multi-model, multi-asset training (4 models × 4 assets = 16 jobs)

---

## Executive Summary

This document proposes enhancements to the ML Training Service gRPC API to support efficient multi-model, multi-asset training while maintaining backward compatibility. The recommended approach uses a **job hierarchy pattern with multiplexed streaming** to provide clear client control, efficient progress updates, and production-grade observability.

**Key Decisions**:
- **API Pattern**: Job hierarchy with `oneof` (refined Option 3)
- **Streaming**: Single multiplexed stream with `child_job_id` routing
- **Hierarchy**: Parent (batch) → Child jobs (flat, one per model/asset pair)
- **Backward Compatibility**: Fully preserved via `oneof` pattern

---

## 1. Current API Analysis

### 1.1 Existing Proto Definition

**File**: `/home/jgrusewski/Work/foxhunt/services/ml_training_service/proto/ml_training.proto`

**Current Single-Model Training**:
```proto
message StartTrainingRequest {
  string model_type = 1;                // Single model: "TLOB", "MAMBA_2", "DQN", "PPO", "LIQUID", "TFT"
  DataSource data_source = 2;           // Training data source configuration
  Hyperparameters hyperparameters = 3;  // Model-specific training parameters
  bool use_gpu = 4;                     // Whether to use GPU acceleration
  string description = 5;               // Optional job description
  map<string, string> tags = 6;         // Optional categorization tags
}

message StartTrainingResponse {
  string job_id = 1;
  TrainingStatus status = 2;
  string message = 3;
}
```

**Current Batch Tuning** (Hyperparameter optimization only):
```proto
message BatchStartTuningJobsRequest {
  repeated string model_types = 1;         // List of models to tune
  uint32 trials_per_model = 2;             // Number of trials for each model
  string config_path = 3;                  // Path to tuning configuration file
  DataSource data_source = 4;              // Training data source for all models
  bool use_gpu = 5;                        // Whether to use GPU acceleration
  bool auto_export_yaml = 6;               // Automatically export best params to YAML
  string yaml_export_path = 7;             // Custom YAML export path
  string description = 8;                  // Optional batch job description
  map<string, string> tags = 9;            // Optional categorization tags
}
```

**Current Implementation Status**:
- ✅ Single-model training: **IMPLEMENTED** (lines 63-516 in `ml_training.proto`)
- ✅ Batch hyperparameter tuning: **PROTO DEFINED** (lines 441-516)
- ❌ Batch hyperparameter tuning: **NOT IMPLEMENTED** (service.rs shows `Status::unimplemented`)
- ❌ Multi-asset training: **NOT SUPPORTED** (no asset field in proto)
- ❌ Multi-model regular training: **NOT SUPPORTED** (only single `model_type` string)

**Streaming Status**:
```proto
rpc SubscribeToTrainingStatus(SubscribeToTrainingStatusRequest) returns (stream TrainingStatusUpdate);

message SubscribeToTrainingStatusRequest {
  string job_id = 1;  // Subscribe to single job only
}

message TrainingStatusUpdate {
  string job_id = 1;
  TrainingStatus status = 2;
  float progress_percentage = 3;
  uint32 current_epoch = 4;
  uint32 total_epochs = 5;
  map<string, float> metrics = 6;
  string message = 7;
  int64 timestamp = 8;
  FinancialMetrics financial_metrics = 9;
  ResourceUsage resource_usage = 10;
}
```

**Limitation**: Client must open N separate streams for N jobs, leading to connection overhead.

---

## 2. Design Analysis: Options Comparison

### 2.1 Option 1: Explicit Multi-Model in Proto

**Proposed Changes**:
```proto
message StartTrainingRequest {
  repeated string model_types = 1;  // NEW: Train multiple models
  repeated string assets = 2;        // NEW: Train on multiple assets
  DataSource data_source = 3;        // Existing field
  Hyperparameters hyperparameters = 4;
  bool use_gpu = 5;
  string description = 6;
  map<string, string> tags = 7;
}
```

**Analysis**:
| Aspect | Assessment |
|--------|------------|
| **Flexibility** | ⭐⭐⭐⭐⭐ Maximum client control over model/asset combinations |
| **Simplicity** | ⭐⭐ Complex validation (N×M matrix), ambiguous job hierarchy |
| **Backward Compatibility** | ❌ **BREAKS EXISTING API** - changes field semantics |
| **Validation Complexity** | ❌ Server must validate all (model, asset) combinations |
| **Job Hierarchy** | ❌ Implicit - unclear parent/child relationship |

**Verdict**: ❌ **NOT RECOMMENDED** - Breaks backward compatibility and introduces validation complexity.

---

### 2.2 Option 2: Implicit Orchestrator Logic

**Proposed Changes**:
```proto
// Keep existing StartTrainingRequest unchanged
// Add new orchestrator endpoint
message StartBatchTrainingRequest {
  BatchTrainingStrategy strategy = 1;  // ALL_MODELS_ALL_ASSETS, SPECIFIC_PAIRS, etc.
  repeated string assets = 2;
  DataSource data_source_template = 3;
  map<string, string> tags = 4;
}

enum BatchTrainingStrategy {
  ALL_MODELS_ALL_ASSETS = 0;    // Train all 4 models on all 4 assets (16 jobs)
  SPECIFIC_MODELS = 1;           // Client provides model list
  ASSET_SPECIFIC = 2;            // Different models per asset
}
```

**Analysis**:
| Aspect | Assessment |
|--------|------------|
| **Flexibility** | ⭐⭐ Limited - orchestrator decides combinations |
| **Simplicity** | ⭐⭐⭐⭐⭐ Cleanest client API, simple validation |
| **Backward Compatibility** | ✅ Preserves existing API |
| **Validation Complexity** | ✅ Server-side strategy enum validation |
| **Job Hierarchy** | ✅ Clear parent/child separation |

**Trade-offs**:
- ✅ Pros: Simplest client experience, clean separation
- ❌ Cons: Inflexible - what if we need 3 models on 2 assets? Requires enum updates for new patterns

**Verdict**: ⭐⭐⭐ **ACCEPTABLE** - Good for fixed use cases, but lacks flexibility for HFT experimentation.

---

### 2.3 Option 3: Job Hierarchy (Parent/Child) - **RECOMMENDED**

**Proposed Changes**:
```proto
message StartTrainingRequest {
  // Common fields that apply to both single and batch jobs
  DataSource data_source = 1;
  bool use_gpu = 2;
  string description = 3;
  map<string, string> tags = 4;

  oneof job_type {
    SingleModelJob single_job = 5;
    BatchModelJob batch_job = 6;
  }
}

// Encapsulates single training job (existing API, wrapped for clarity)
message SingleModelJob {
  string model_type = 1;              // "TLOB", "MAMBA_2", "DQN", "PPO", "TFT"
  string asset = 2;                   // "ES.FUT", "NQ.FUT", "6E.FUT", "ZN.FUT"
  Hyperparameters hyperparameters = 3;
}

// Defines batch of training jobs
message BatchModelJob {
  repeated string model_types = 1;    // Models to train (e.g., ["DQN", "PPO", "MAMBA_2", "TFT"])
  repeated string assets = 2;          // Assets to train on (e.g., ["ES.FUT", "NQ.FUT", "6E.FUT", "ZN.FUT"])
  Hyperparameters common_hyperparameters = 3; // Optional: shared hyperparameters
  string client_request_id = 4;       // Optional: for idempotency
}

message StartTrainingResponse {
  string parent_job_id = 1;           // ID for the overall batch or single job
  repeated string child_job_ids = 2;  // IDs for individual (model, asset) jobs (empty for single job)
  TrainingStatus initial_status = 3;
  string message = 4;
}
```

**Analysis**:
| Aspect | Assessment |
|--------|------------|
| **Flexibility** | ⭐⭐⭐⭐⭐ Full client control via repeated fields |
| **Simplicity** | ⭐⭐⭐⭐ Clear intent, standard gRPC `oneof` pattern |
| **Backward Compatibility** | ✅ **FULLY PRESERVED** - existing clients use `single_job` |
| **Validation Complexity** | ⭐⭐⭐⭐ Server validates combinations, clear structure |
| **Job Hierarchy** | ✅ **EXPLICIT** - `parent_job_id` links all child jobs |

**Key Benefits**:
1. **Backward Compatibility**: `oneof` is the standard gRPC pattern for API evolution
2. **Clear Intent**: Explicitly differentiates single vs. batch requests
3. **Client Flexibility**: Clients specify exact model/asset combinations
4. **Server Clarity**: Server decomposes batch into 16 child jobs (4 models × 4 assets)
5. **Observability**: `parent_job_id` provides natural grouping for monitoring

**Verdict**: ✅ **RECOMMENDED** - Best balance of flexibility, clarity, and backward compatibility.

---

## 3. Streaming Progress Design

### 3.1 Streaming Options Analysis

For 16 simultaneous training jobs (4 models × 4 assets):

**Option A: Single Stream per Parent Job with Hierarchical Updates**
```proto
message TrainingProgressUpdate {
  string parent_job_id = 1;
  repeated ChildJobUpdate child_updates = 2;
}
```
❌ **Drawback**: Complex nested message parsing on client side, inefficient for partial updates.

**Option B: Separate Stream per Child Job**
```proto
// Client opens 16 separate streams
rpc SubscribeToTrainingStatus(job_id) returns (stream TrainingStatusUpdate);
```
❌ **Drawback**: 16 open gRPC connections = high overhead, connection management complexity.

**Option C: Multiplexed Stream with `child_job_id` Routing** ⭐ **RECOMMENDED**
```proto
rpc StreamTrainingProgress(StreamTrainingProgressRequest) returns (stream TrainingProgressUpdate);

message StreamTrainingProgressRequest {
  oneof subscription_target {
    string parent_job_id = 1;          // Subscribe to all child jobs in batch
    repeated string child_job_ids = 2; // Subscribe to specific child jobs
  }
}

message TrainingProgressUpdate {
  string child_job_id = 1;             // REQUIRED: Unique ID for (model, asset) job
  string parent_job_id = 2;            // Optional: Batch ID this child belongs to
  string model_type = 3;               // For context/filtering (e.g., "DQN")
  string asset = 4;                    // For context/filtering (e.g., "ES.FUT")
  JobStatus status = 5;                // PENDING, RUNNING, COMPLETED, FAILED, CANCELLED
  float progress_percentage = 6;       // 0.0 to 100.0
  uint32 current_epoch = 7;
  uint32 total_epochs = 8;
  map<string, float> metrics = 9;      // loss, sharpe_ratio, etc.
  string message = 10;
  int64 timestamp = 11;
  FinancialMetrics financial_metrics = 12;
  ResourceUsage resource_usage = 13;
  string error_message = 14;           // Populated if status = FAILED
}

enum JobStatus {
  UNKNOWN = 0;
  PENDING = 1;
  RUNNING = 2;
  COMPLETED = 3;
  FAILED = 4;
  CANCELLED = 5;
}
```

### 3.2 Multiplexed Streaming Benefits

**Efficiency**:
- **1 network connection** for all 16 jobs (vs. 16 separate connections)
- Reduced TCP overhead, lower latency, better resource utilization

**Clarity**:
- Each message explicitly identifies `child_job_id` (e.g., `"DQN_ES.FUT"`)
- Client routes messages to UI components based on `child_job_id`

**Scalability**:
- Handles 100+ jobs without linear connection growth
- Server broadcasts updates via single stream

**HFT-Specific Advantages**:
- **Granular real-time updates**: Critical for HFT monitoring
- **Low latency**: Single connection = lower overhead
- **Observability**: `parent_job_id` + `child_job_id` enable hierarchical tracking

**Example Client Usage** (Rust):
```rust
// Subscribe to all child jobs in batch
let request = StreamTrainingProgressRequest {
    subscription_target: Some(SubscriptionTarget::ParentJobId("batch_001".to_string())),
};

let mut stream = client.stream_training_progress(request).await?.into_inner();

while let Some(update) = stream.message().await? {
    match update.child_job_id.as_str() {
        "DQN_ES.FUT" => update_dqn_es_dashboard(&update),
        "PPO_NQ.FUT" => update_ppo_nq_dashboard(&update),
        _ => log::info!("Update for {}: {}%", update.child_job_id, update.progress_percentage),
    }
}
```

---

## 4. Job Hierarchy Design

### 4.1 Hierarchy Options

**Option A: Parent → Asset → Model** (3-level hierarchy)
```
Batch Request (parent_job_id: "batch_001")
├── ES.FUT (asset_job_id: "ES.FUT_001")
│   ├── DQN (child_job_id: "DQN_ES.FUT")
│   ├── PPO (child_job_id: "PPO_ES.FUT")
│   ├── MAMBA-2 (child_job_id: "MAMBA2_ES.FUT")
│   └── TFT (child_job_id: "TFT_ES.FUT")
├── NQ.FUT (asset_job_id: "NQ.FUT_001")
│   ├── ...
```
❌ **Drawback**: Unnecessary complexity unless asset-level operations exist (e.g., asset-specific resource pools, asset-level cancellation).

**Option B: Parent → Child Jobs (Flat)** ⭐ **RECOMMENDED**
```
Batch Request (parent_job_id: "batch_001")
├── DQN_ES.FUT (child_job_id: "01")
├── DQN_NQ.FUT (child_job_id: "02")
├── DQN_6E.FUT (child_job_id: "03")
├── DQN_ZN.FUT (child_job_id: "04")
├── PPO_ES.FUT (child_job_id: "05")
├── PPO_NQ.FUT (child_job_id: "06")
├── ...
├── TFT_ZN.FUT (child_job_id: "16")
```

### 4.2 Recommended Hierarchy Details

**Parent Job**:
- **ID**: `parent_job_id` (UUID, e.g., `"3e4a891c-7f2b-4d5e-9c1a-8f3b2d4e5a6c"`)
- **Purpose**: Logical grouping for all 16 child jobs
- **Lifetime**: Exists until all child jobs complete/fail
- **Status**: Computed from child statuses (e.g., RUNNING if any child is RUNNING)

**Child Job**:
- **ID**: `child_job_id` (UUID, e.g., `"a1b2c3d4-e5f6-7890-abcd-ef1234567890"`)
- **Metadata**: `model_type` ("DQN"), `asset` ("ES.FUT")
- **Parent Link**: References `parent_job_id`
- **Independence**: Runs independently, failures isolated

**Database Schema** (PostgreSQL):
```sql
CREATE TABLE training_jobs (
    job_id UUID PRIMARY KEY,
    parent_job_id UUID,  -- NULL for single jobs, references parent for batch
    model_type VARCHAR(20) NOT NULL,
    asset VARCHAR(20) NOT NULL,
    status VARCHAR(20) NOT NULL,
    progress_percentage REAL DEFAULT 0.0,
    created_at TIMESTAMPTZ DEFAULT NOW(),
    started_at TIMESTAMPTZ,
    completed_at TIMESTAMPTZ,
    error_message TEXT,
    -- Index for efficient queries
    INDEX idx_parent_job (parent_job_id),
    INDEX idx_status (status),
    INDEX idx_created_at (created_at DESC)
);
```

**Benefits of Flat Hierarchy**:
1. **Simplicity**: Easier to query, monitor, and debug
2. **Independence**: Each (model, asset) pair is self-contained
3. **Performance**: No nested queries needed for status retrieval
4. **HFT-Appropriate**: Clear 1:1 mapping to training artifacts

---

## 5. Backward Compatibility Strategy

### 5.1 Migration Path

**Phase 1: Proto Update** (Breaking change controlled via `oneof`)
```proto
message StartTrainingRequest {
  // OLD CLIENT (existing behavior, no code changes):
  // Implicitly uses single_job with model_type="DQN", asset="ES.FUT"

  // NEW CLIENT (opt-in to batch API):
  // Explicitly uses batch_job with model_types=["DQN","PPO"], assets=["ES.FUT","NQ.FUT"]

  oneof job_type {
    SingleModelJob single_job = 5;
    BatchModelJob batch_job = 6;
  }
}
```

**Phase 2: Server Implementation**
```rust
async fn start_training(&self, request: Request<StartTrainingRequest>) -> Result<Response<StartTrainingResponse>, Status> {
    let req = request.into_inner();

    match req.job_type {
        Some(JobType::SingleJob(single)) => {
            // Existing single-model training logic
            let job_id = self.orchestrator.submit_training_job(single).await?;
            Ok(Response::new(StartTrainingResponse {
                parent_job_id: job_id.clone(),
                child_job_ids: vec![],  // Empty for single job
                initial_status: TrainingStatus::Pending,
                message: format!("Training job {} submitted", job_id),
            }))
        }
        Some(JobType::BatchJob(batch)) => {
            // NEW: Batch training logic
            let parent_id = Uuid::new_v4().to_string();
            let mut child_ids = Vec::new();

            for model in &batch.model_types {
                for asset in &batch.assets {
                    let child_id = self.orchestrator.submit_training_job_with_parent(
                        model, asset, &parent_id
                    ).await?;
                    child_ids.push(child_id);
                }
            }

            Ok(Response::new(StartTrainingResponse {
                parent_job_id: parent_id,
                child_job_ids: child_ids,
                initial_status: TrainingStatus::Pending,
                message: format!("Batch training job submitted with {} child jobs", child_ids.len()),
            }))
        }
        None => Err(Status::invalid_argument("job_type must be specified")),
    }
}
```

**Phase 3: Client Migration** (Gradual rollout)
1. **Week 1**: Deploy new server with backward compatibility
2. **Week 2-4**: Existing clients continue using `single_job` (no changes needed)
3. **Week 5+**: New clients adopt `batch_job` for multi-model training

### 5.2 Compatibility Testing

**Test Cases**:
1. ✅ **Old client + New server**: Single-model training via `single_job`
2. ✅ **New client + New server**: Batch training via `batch_job`
3. ✅ **Old client + Old server**: No regression (existing API unchanged)
4. ✅ **Streaming**: Old clients subscribe to single `job_id`, new clients subscribe to `parent_job_id`

---

## 6. Additional Production Considerations

### 6.1 Idempotency

**Problem**: Client retries can duplicate training jobs.

**Solution**: Client-provided request ID
```proto
message BatchModelJob {
  repeated string model_types = 1;
  repeated string assets = 2;
  Hyperparameters common_hyperparameters = 3;
  string client_request_id = 4;  // NEW: For idempotency
}
```

**Server Logic**:
```rust
// Check if request_id already processed
if let Some(existing_job) = self.get_job_by_request_id(&batch.client_request_id).await? {
    return Ok(Response::new(StartTrainingResponse {
        parent_job_id: existing_job.parent_job_id,
        child_job_ids: existing_job.child_job_ids,
        initial_status: existing_job.status,
        message: "Request already processed (idempotent response)".to_string(),
    }));
}
```

### 6.2 Resource Management

**GPU Memory Budget**: 440MB total (RTX 3050 Ti has 4GB)
- DQN: ~6MB
- PPO: ~145MB
- MAMBA-2: ~164MB
- TFT-INT8: ~125MB

**Concurrency Limits**:
- **Sequential**: Train 1 model at a time (safest, 4GB headroom)
- **Parallel (2x)**: Train 2 models concurrently (e.g., DQN + PPO = 151MB)
- **Parallel (4x)**: Train all 4 models (440MB, tight but feasible)

**Recommendation**: Configurable concurrency via `max_concurrent_jobs` setting
```proto
message BatchModelJob {
  repeated string model_types = 1;
  repeated string assets = 2;
  uint32 max_concurrent_jobs = 3;  // Default: 1 (sequential), max: 4 (all parallel)
}
```

### 6.3 Granular Error Reporting

**Enhanced Progress Update**:
```proto
message TrainingProgressUpdate {
  string child_job_id = 1;
  JobStatus status = 2;

  // Error details (populated if status = FAILED)
  ErrorDetails error = 3;
}

message ErrorDetails {
  string error_code = 1;           // "OUT_OF_MEMORY", "DATA_LOAD_FAILED", etc.
  string error_message = 2;        // Human-readable description
  string stack_trace = 3;          // Full stack trace for debugging
  map<string, string> context = 4; // Additional context (e.g., epoch=5, batch_size=32)
}
```

**HFT Benefit**: Rapid diagnosis and remediation of training failures.

### 6.4 Cancellation Support

**New RPC**:
```proto
service MLTrainingService {
  // ... existing RPCs

  // Cancel a batch (cascades to all child jobs)
  rpc CancelTrainingJob(CancelTrainingJobRequest) returns (CancelTrainingJobResponse);
}

message CancelTrainingJobRequest {
  oneof target {
    string parent_job_id = 1;      // Cancel entire batch
    string child_job_id = 2;       // Cancel single child job
  }
  string reason = 3;                // Optional cancellation reason
}

message CancelTrainingJobResponse {
  bool success = 1;
  repeated string cancelled_job_ids = 2;
  string message = 3;
}
```

**Streaming Update**:
```proto
message TrainingProgressUpdate {
  JobStatus status = 5;  // Will transition to CANCELLED
  string message = 10;   // "Cancelled by user: reason"
}
```

### 6.5 Observability & Monitoring

**Metrics to Track** (Prometheus):
```
# Gauge: Active training jobs by status
foxhunt_training_jobs_active{status="running"} 16
foxhunt_training_jobs_active{status="pending"} 0

# Counter: Completed jobs by outcome
foxhunt_training_jobs_completed_total{outcome="success"} 48
foxhunt_training_jobs_completed_total{outcome="failed"} 2

# Histogram: Training duration by model
foxhunt_training_duration_seconds{model="DQN"} 15.2
foxhunt_training_duration_seconds{model="MAMBA2"} 112.5
```

**Logging**:
```
[INFO] parent_job_id=batch_001 child_job_id=01 model=DQN asset=ES.FUT status=RUNNING epoch=5/30
[ERROR] parent_job_id=batch_001 child_job_id=07 model=PPO asset=6E.FUT status=FAILED error=OUT_OF_MEMORY
```

---

## 7. Final Recommended Proto Changes

### 7.1 Complete Proto Definition

```proto
syntax = "proto3";
package ml_training;

service MLTrainingService {
  // Training Job Management
  rpc StartTraining(StartTrainingRequest) returns (StartTrainingResponse);
  rpc StopTraining(StopTrainingRequest) returns (StopTrainingResponse);
  rpc CancelTrainingJob(CancelTrainingJobRequest) returns (CancelTrainingJobResponse);  // NEW

  // Progress Monitoring
  rpc StreamTrainingProgress(StreamTrainingProgressRequest) returns (stream TrainingProgressUpdate);  // NEW (replaces SubscribeToTrainingStatus)

  // Job Discovery
  rpc ListTrainingJobs(ListTrainingJobsRequest) returns (ListTrainingJobsResponse);
  rpc GetTrainingJobDetails(GetTrainingJobDetailsRequest) returns (GetTrainingJobDetailsResponse);

  // ... other RPCs (ListAvailableModels, HealthCheck, Tuning RPCs)
}

// --- Training Request/Response ---

message StartTrainingRequest {
  // Common fields for both single and batch jobs
  DataSource data_source = 1;
  bool use_gpu = 2;
  string description = 3;
  map<string, string> tags = 4;

  oneof job_type {
    SingleModelJob single_job = 5;
    BatchModelJob batch_job = 6;
  }
}

message SingleModelJob {
  string model_type = 1;              // "DQN", "PPO", "MAMBA_2", "TFT"
  string asset = 2;                   // "ES.FUT", "NQ.FUT", "6E.FUT", "ZN.FUT"
  Hyperparameters hyperparameters = 3;
}

message BatchModelJob {
  repeated string model_types = 1;    // Models to train
  repeated string assets = 2;          // Assets to train on
  Hyperparameters common_hyperparameters = 3; // Shared hyperparameters
  string client_request_id = 4;       // For idempotency
  uint32 max_concurrent_jobs = 5;     // Concurrency limit (default: 1, max: 4)
}

message StartTrainingResponse {
  string parent_job_id = 1;           // Batch ID or single job ID
  repeated string child_job_ids = 2;  // Individual (model, asset) job IDs (empty for single)
  TrainingStatus initial_status = 3;
  string message = 4;
}

// --- Streaming Progress ---

message StreamTrainingProgressRequest {
  oneof subscription_target {
    string parent_job_id = 1;          // Subscribe to all child jobs in batch
    repeated string child_job_ids = 2; // Subscribe to specific child jobs
  }
}

message TrainingProgressUpdate {
  string child_job_id = 1;             // Unique ID for (model, asset) job
  string parent_job_id = 2;            // Batch ID (if part of batch)
  string model_type = 3;               // "DQN", "PPO", etc.
  string asset = 4;                    // "ES.FUT", "NQ.FUT", etc.
  JobStatus status = 5;
  float progress_percentage = 6;
  uint32 current_epoch = 7;
  uint32 total_epochs = 8;
  map<string, float> metrics = 9;      // loss, sharpe_ratio, etc.
  string message = 10;
  int64 timestamp = 11;
  FinancialMetrics financial_metrics = 12;
  ResourceUsage resource_usage = 13;
  ErrorDetails error = 14;             // Populated if status = FAILED
}

message ErrorDetails {
  string error_code = 1;               // "OUT_OF_MEMORY", "DATA_LOAD_FAILED", etc.
  string error_message = 2;
  string stack_trace = 3;
  map<string, string> context = 4;
}

// --- Cancellation ---

message CancelTrainingJobRequest {
  oneof target {
    string parent_job_id = 1;          // Cancel entire batch
    string child_job_id = 2;           // Cancel single child job
  }
  string reason = 3;
}

message CancelTrainingJobResponse {
  bool success = 1;
  repeated string cancelled_job_ids = 2;
  string message = 3;
}

// --- Enums ---

enum JobStatus {
  UNKNOWN = 0;
  PENDING = 1;
  RUNNING = 2;
  COMPLETED = 3;
  FAILED = 4;
  CANCELLED = 5;
}

// --- Existing messages (unchanged) ---
// DataSource, Hyperparameters, FinancialMetrics, ResourceUsage, etc.
```

### 7.2 Migration Checklist

**Server Implementation**:
- [ ] Update proto file with `oneof job_type`
- [ ] Implement `BatchModelJob` handler in `service.rs`
- [ ] Add job hierarchy tracking in database
- [ ] Implement multiplexed streaming in `StreamTrainingProgress`
- [ ] Add idempotency check via `client_request_id`
- [ ] Implement `CancelTrainingJob` RPC
- [ ] Add Prometheus metrics for job tracking

**Client Updates**:
- [ ] Regenerate proto bindings (tonic)
- [ ] Update TLI to support batch training commands
- [ ] Add batch job monitoring in UI/dashboard
- [ ] Test backward compatibility with single-job API

**Testing**:
- [ ] Unit tests for batch job decomposition
- [ ] Integration tests for 16-job batch (4 models × 4 assets)
- [ ] Streaming tests for multiplexed updates
- [ ] Backward compatibility tests (old client + new server)
- [ ] Idempotency tests (duplicate `client_request_id`)
- [ ] Cancellation tests (parent/child)

**Documentation**:
- [ ] Update API documentation (gRPC method signatures)
- [ ] Add batch training tutorial
- [ ] Update ML_TRAINING_PARQUET_GUIDE.md with batch examples
- [ ] Create observability runbook (Grafana dashboards)

---

## 8. Comparison with Existing Batch Tuning API

**Current Batch Tuning** (Hyperparameter optimization):
```proto
message BatchStartTuningJobsRequest {
  repeated string model_types = 1;         // Multiple models
  uint32 trials_per_model = 2;             // Tuning trials
  string config_path = 3;
  DataSource data_source = 4;              // Single data source (no multi-asset)
  bool use_gpu = 5;
  bool auto_export_yaml = 6;
  string yaml_export_path = 7;
  string description = 8;
  map<string, string> tags = 9;
}
```

**Proposed Batch Training** (Regular training):
```proto
message BatchModelJob {
  repeated string model_types = 1;         // Multiple models
  repeated string assets = 2;              // NEW: Multi-asset support
  Hyperparameters common_hyperparameters = 3;
  string client_request_id = 4;           // NEW: Idempotency
  uint32 max_concurrent_jobs = 5;         // NEW: Resource control
}
```

**Key Differences**:
| Feature | Batch Tuning | Batch Training |
|---------|--------------|----------------|
| **Purpose** | Hyperparameter optimization (Optuna) | Regular model training |
| **Multi-Asset** | ❌ No (single `data_source`) | ✅ Yes (`repeated assets`) |
| **Job Hierarchy** | Flat (N tuning jobs) | Parent → Children (N×M jobs) |
| **Streaming** | `StreamTuningProgress` (trial-based) | `StreamTrainingProgress` (epoch-based) |
| **Implementation** | ❌ Not implemented (`Status::unimplemented`) | 🚧 Proposed in this doc |
| **Use Case** | Find best hyperparameters for 1 asset | Train 4 models on 4 assets (16 jobs) |

**Recommendation**: Keep both APIs separate. Batch tuning is for hyperparameter search, batch training is for multi-asset production training.

---

## 9. Conclusion

### 9.1 Summary of Recommendations

1. **API Design**: Use **Option 3 (Job Hierarchy with `oneof`)** for maximum flexibility and backward compatibility
2. **Streaming**: Implement **multiplexed streaming** with `child_job_id` routing for efficient progress updates
3. **Job Hierarchy**: Use **flat parent → child** structure (no intermediate asset level)
4. **Backward Compatibility**: Fully preserved via `oneof job_type` pattern
5. **Additional Features**: Add idempotency, cancellation, granular error reporting, and observability

### 9.2 Next Steps

**Wave 1 Agent 4** (Implementation):
1. Update `ml_training.proto` with proposed changes
2. Implement server-side batch job handler
3. Add database schema for parent/child tracking
4. Implement multiplexed streaming
5. Add unit and integration tests

**Wave 1 Agent 5** (Client Integration):
1. Regenerate gRPC client bindings
2. Update TLI with batch training commands
3. Add batch job monitoring UI
4. Test end-to-end with 16-job batch

---

## Appendix A: Example Client Usage

### A.1 Single-Model Training (Backward Compatible)

```rust
use ml_training::*;

let request = StartTrainingRequest {
    data_source: Some(DataSource {
        source: Some(data_source::Source::FilePath("test_data/ES_FUT_180d.parquet".to_string())),
        start_time: 0,
        end_time: 0,
    }),
    use_gpu: true,
    description: "DQN training on ES.FUT".to_string(),
    tags: HashMap::new(),
    job_type: Some(start_training_request::JobType::SingleJob(SingleModelJob {
        model_type: "DQN".to_string(),
        asset: "ES.FUT".to_string(),
        hyperparameters: Some(Hyperparameters {
            model_params: Some(hyperparameters::ModelParams::DqnParams(DqnParams {
                epochs: 100,
                learning_rate: 0.001,
                batch_size: 32,
                // ... other params
            })),
        }),
    })),
};

let response = client.start_training(request).await?;
println!("Job ID: {}", response.parent_job_id);
```

### A.2 Batch Training (New API)

```rust
let request = StartTrainingRequest {
    data_source: Some(DataSource {
        source: Some(data_source::Source::FilePath("test_data/{asset}_180d.parquet".to_string())),
        start_time: 0,
        end_time: 0,
    }),
    use_gpu: true,
    description: "4 models × 4 assets = 16 jobs".to_string(),
    tags: HashMap::new(),
    job_type: Some(start_training_request::JobType::BatchJob(BatchModelJob {
        model_types: vec!["DQN".to_string(), "PPO".to_string(), "MAMBA_2".to_string(), "TFT".to_string()],
        assets: vec!["ES.FUT".to_string(), "NQ.FUT".to_string(), "6E.FUT".to_string(), "ZN.FUT".to_string()],
        common_hyperparameters: Some(Hyperparameters { /* shared params */ }),
        client_request_id: Uuid::new_v4().to_string(),
        max_concurrent_jobs: 2,  // Train 2 models at a time
    })),
};

let response = client.start_training(request).await?;
println!("Batch ID: {}", response.parent_job_id);
println!("Child jobs: {:?}", response.child_job_ids);  // 16 job IDs
```

### A.3 Streaming Progress Updates

```rust
let request = StreamTrainingProgressRequest {
    subscription_target: Some(stream_training_progress_request::SubscriptionTarget::ParentJobId(
        "batch_001".to_string()
    )),
};

let mut stream = client.stream_training_progress(request).await?.into_inner();

while let Some(update) = stream.message().await? {
    println!(
        "[{}] {}/{} - Epoch {}/{} - {:.1}% - {}",
        update.child_job_id,
        update.model_type,
        update.asset,
        update.current_epoch,
        update.total_epochs,
        update.progress_percentage,
        update.message
    );

    if update.status == JobStatus::Failed as i32 {
        eprintln!("FAILED: {:?}", update.error);
    }
}
```

**Example Output**:
```
[01] DQN/ES.FUT - Epoch 5/100 - 5.0% - Training in progress
[02] DQN/NQ.FUT - Epoch 3/100 - 3.0% - Training in progress
[05] PPO/ES.FUT - Epoch 2/30 - 6.7% - Training in progress
[01] DQN/ES.FUT - Epoch 100/100 - 100.0% - Training complete
[02] DQN/NQ.FUT - Epoch 50/100 - 50.0% - Training in progress
FAILED: ErrorDetails { error_code: "OUT_OF_MEMORY", error_message: "GPU memory exceeded: 4.2GB > 4.0GB limit", ... }
```

---

## Appendix B: Database Schema

```sql
-- Parent/Batch job tracking
CREATE TABLE training_batches (
    batch_id UUID PRIMARY KEY,
    client_request_id VARCHAR(255) UNIQUE,  -- For idempotency
    description TEXT,
    created_at TIMESTAMPTZ DEFAULT NOW(),
    updated_at TIMESTAMPTZ DEFAULT NOW(),
    status VARCHAR(20) NOT NULL,  -- Computed from child statuses
    tags JSONB
);

-- Individual training jobs
CREATE TABLE training_jobs (
    job_id UUID PRIMARY KEY,
    parent_batch_id UUID REFERENCES training_batches(batch_id),  -- NULL for single jobs
    model_type VARCHAR(20) NOT NULL,
    asset VARCHAR(20) NOT NULL,
    status VARCHAR(20) NOT NULL,
    progress_percentage REAL DEFAULT 0.0,
    current_epoch INT DEFAULT 0,
    total_epochs INT NOT NULL,
    created_at TIMESTAMPTZ DEFAULT NOW(),
    started_at TIMESTAMPTZ,
    completed_at TIMESTAMPTZ,
    error_code VARCHAR(50),
    error_message TEXT,
    stack_trace TEXT,
    metrics JSONB,  -- Store training metrics as JSON

    -- Indexes for efficient queries
    INDEX idx_parent_batch (parent_batch_id),
    INDEX idx_status (status),
    INDEX idx_model_asset (model_type, asset),
    INDEX idx_created_at (created_at DESC)
);

-- Example queries
-- Get all child jobs for a batch
SELECT * FROM training_jobs WHERE parent_batch_id = 'batch_001' ORDER BY created_at;

-- Get batch status summary
SELECT
    status,
    COUNT(*) as count,
    AVG(progress_percentage) as avg_progress
FROM training_jobs
WHERE parent_batch_id = 'batch_001'
GROUP BY status;

-- Check idempotency
SELECT batch_id FROM training_batches WHERE client_request_id = 'client_req_123';
```

---

**End of Document**