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:

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):

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:

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:

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:

// 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:

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

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

// 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

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):

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

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)

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

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

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:

// 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

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:

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:

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:

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

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):

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):

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)

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)

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

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

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