04d8802c94
Remove use_double_dqn, use_dueling, use_per, use_branching,
use_distributional, use_noisy_nets, use_huber_loss, and use_cql
from DQNConfig, DQNHyperparameters, and DqnParams structs.
These features are always enabled (Rainbow DQN standard). The boolean
flags were dead code — every constructor set them to true, and the
only code paths that set them to false were in tests that disabled
features for simplicity. With the fields removed, the features are
unconditionally active, eliminating ~490 lines of dead configuration.
Key changes:
- Struct field declarations removed from 3 core config structs
- Conditional branches (if use_X { ... } else { ... }) simplified:
dueling/branching/PER network creation is now unconditional
- Checkpoint metadata hardcodes "true" for backward compatibility
- Hyperopt search space index 11 (use_branching) fixed at 1.0
- TOML/YAML config files cleaned of removed fields
- Tests that toggled these flags updated or rewritten
45 files changed, -487 net lines. Zero new test failures.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2130 lines
81 KiB
Rust
2130 lines
81 KiB
Rust
//! gRPC Service Implementation
|
|
//!
|
|
//! This module implements the MLTrainingService gRPC interface,
|
|
//! providing the external API for training job management.
|
|
|
|
use std::collections::HashMap;
|
|
use std::pin::Pin;
|
|
use std::sync::Arc;
|
|
|
|
use anyhow::Result;
|
|
use tokio::time::{timeout, Duration};
|
|
use tokio_stream::Stream;
|
|
use tonic::{Request, Response, Status};
|
|
use tracing::{debug, error, info, warn, instrument};
|
|
use uuid::Uuid;
|
|
|
|
// Import the generated protobuf types
|
|
pub mod proto {
|
|
pub mod ml_training {
|
|
tonic::include_proto!("ml_training");
|
|
}
|
|
// Re-export for convenience
|
|
pub use ml_training::*;
|
|
}
|
|
|
|
use proto::{
|
|
ml_training_service_server::MlTrainingService, DqnParams,
|
|
FinancialMetrics as ProtoFinancialMetrics, GetTrainingJobDetailsRequest,
|
|
GetTrainingJobDetailsResponse, HealthCheckRequest, HealthCheckResponse, Hyperparameters,
|
|
LiquidParams, ListAvailableModelsRequest, ListAvailableModelsResponse, ListTrainingJobsRequest,
|
|
ListTrainingJobsResponse, MambaParams, ModelDefinition, PpoParams,
|
|
ResourceUsage as ProtoResourceUsage, StartTrainingRequest, StartTrainingResponse,
|
|
StopTrainingRequest, StopTrainingResponse, SubscribeToTrainingStatusRequest, TftParams,
|
|
TlobParams, TrainingJobDetails, TrainingJobSummary, TrainingStatus as ProtoTrainingStatus,
|
|
TrainingStatusUpdate as ProtoStatusUpdate,
|
|
// Job completion & model promotion (on-demand training pipeline)
|
|
JobCompletionReport, JobCompletionAck,
|
|
ListPendingPromotionsRequest, ListPendingPromotionsResponse, PendingPromotion,
|
|
ApprovePromotionRequest, ApprovePromotionResponse,
|
|
RejectPromotionRequest, RejectPromotionResponse,
|
|
// Dashboard-facing model approval / rejection
|
|
ApproveModelRequest, ApproveModelResponse,
|
|
RejectModelRequest, RejectModelResponse,
|
|
};
|
|
|
|
use crate::batch_tuning_manager::{BatchJobStatus, BatchTuningManager};
|
|
use crate::grpc_tuning_handlers::TuningHandlers;
|
|
use crate::orchestrator::{JobStatus, TrainingOrchestrator, TrainingStatusUpdate};
|
|
use crate::tuning_manager::{TuningJobStatus, TuningManager};
|
|
use config::MLConfig;
|
|
use ml::training_pipeline::{
|
|
ModelArchitectureConfig, ProductionTrainingConfig, TrainingHyperparameters,
|
|
};
|
|
|
|
/// gRPC service implementation
|
|
pub struct MLTrainingServiceImpl {
|
|
orchestrator: Arc<TrainingOrchestrator>,
|
|
tuning_handlers: Arc<TuningHandlers>,
|
|
batch_tuning_manager: Arc<BatchTuningManager>,
|
|
/// Tracks pending model promotions and active model registry.
|
|
pub promotion_manager: Arc<crate::promotion_manager::PromotionManager>,
|
|
/// K8s job dispatcher -- None when running outside a K8s cluster.
|
|
pub k8s_dispatcher: Option<Arc<crate::k8s_dispatcher::K8sDispatcher>>,
|
|
/// Job spawner for persisting training jobs to PostgreSQL before K8s dispatch.
|
|
/// None when no database pool is available.
|
|
pub job_spawner: Option<Arc<crate::job_spawner::JobSpawner>>,
|
|
}
|
|
|
|
impl MLTrainingServiceImpl {
|
|
/// Create a new service instance
|
|
pub fn new(
|
|
orchestrator: Arc<TrainingOrchestrator>,
|
|
tuning_manager: Arc<TuningManager>,
|
|
_config: MLConfig,
|
|
promotion_manager: Arc<crate::promotion_manager::PromotionManager>,
|
|
k8s_dispatcher: Option<Arc<crate::k8s_dispatcher::K8sDispatcher>>,
|
|
job_spawner: Option<Arc<crate::job_spawner::JobSpawner>>,
|
|
) -> Self {
|
|
let tuning_handlers = Arc::new(TuningHandlers::new(Arc::clone(&tuning_manager)));
|
|
let working_dir = std::env::var("TUNING_WORKING_DIR").unwrap_or_else(|_| ".".to_string());
|
|
let batch_tuning_manager = Arc::new(BatchTuningManager::new(tuning_manager, working_dir));
|
|
Self {
|
|
orchestrator,
|
|
tuning_handlers,
|
|
batch_tuning_manager,
|
|
promotion_manager,
|
|
k8s_dispatcher,
|
|
job_spawner,
|
|
}
|
|
}
|
|
|
|
/// Convert internal job status to protobuf status
|
|
fn convert_job_status(status: &JobStatus) -> ProtoTrainingStatus {
|
|
match status {
|
|
JobStatus::Pending => ProtoTrainingStatus::Pending,
|
|
JobStatus::Running => ProtoTrainingStatus::Running,
|
|
JobStatus::Completed => ProtoTrainingStatus::Completed,
|
|
JobStatus::Failed => ProtoTrainingStatus::Failed,
|
|
JobStatus::Stopped => ProtoTrainingStatus::Stopped,
|
|
JobStatus::Paused => ProtoTrainingStatus::Paused,
|
|
}
|
|
}
|
|
|
|
/// Convert internal batch job status to protobuf batch status
|
|
fn convert_batch_status(status: &BatchJobStatus) -> proto::BatchTuningStatus {
|
|
match status {
|
|
BatchJobStatus::Pending => proto::BatchTuningStatus::BatchPending,
|
|
BatchJobStatus::Running => proto::BatchTuningStatus::BatchRunning,
|
|
BatchJobStatus::Completed => proto::BatchTuningStatus::BatchCompleted,
|
|
BatchJobStatus::PartiallyCompleted => {
|
|
proto::BatchTuningStatus::BatchPartiallyCompleted
|
|
},
|
|
BatchJobStatus::Failed => proto::BatchTuningStatus::BatchFailed,
|
|
BatchJobStatus::Stopped => proto::BatchTuningStatus::BatchStopped,
|
|
}
|
|
}
|
|
|
|
/// Convert internal tuning job status to protobuf tuning status
|
|
fn convert_tuning_job_status(status: &TuningJobStatus) -> proto::TuningJobStatus {
|
|
match status {
|
|
TuningJobStatus::Pending => proto::TuningJobStatus::TuningPending,
|
|
TuningJobStatus::Running => proto::TuningJobStatus::TuningRunning,
|
|
TuningJobStatus::Completed => proto::TuningJobStatus::TuningCompleted,
|
|
TuningJobStatus::Failed => proto::TuningJobStatus::TuningFailed,
|
|
TuningJobStatus::Stopped => proto::TuningJobStatus::TuningStopped,
|
|
}
|
|
}
|
|
|
|
/// Convert protobuf hyperparameters to internal config
|
|
fn convert_hyperparameters(
|
|
&self,
|
|
params: Option<Hyperparameters>,
|
|
) -> Result<ProductionTrainingConfig> {
|
|
let mut config = ProductionTrainingConfig::default();
|
|
|
|
if let Some(hyperparams) = params {
|
|
match hyperparams.model_params {
|
|
Some(proto::hyperparameters::ModelParams::TlobParams(tlob)) => {
|
|
config.model_config = ModelArchitectureConfig {
|
|
input_dim: 20, // Default for TLOB
|
|
hidden_dims: vec![
|
|
usize::try_from(tlob.hidden_dim).unwrap_or(256),
|
|
usize::try_from(tlob.hidden_dim).unwrap_or(256) / 2,
|
|
],
|
|
output_dim: 1,
|
|
dropout_rate: tlob.dropout_rate as f64,
|
|
activation: "relu".to_string(),
|
|
batch_norm: true,
|
|
residual_connections: false,
|
|
};
|
|
|
|
config.training_params = TrainingHyperparameters {
|
|
learning_rate: tlob.learning_rate as f64,
|
|
batch_size: usize::try_from(tlob.batch_size).unwrap_or(64),
|
|
max_epochs: usize::try_from(tlob.epochs).unwrap_or(100),
|
|
patience: 50,
|
|
validation_split: 0.2,
|
|
l2_regularization: 1e-4,
|
|
lr_decay_factor: 0.5,
|
|
lr_decay_patience: 25,
|
|
};
|
|
},
|
|
|
|
Some(proto::hyperparameters::ModelParams::MambaParams(mamba)) => {
|
|
config.model_config = ModelArchitectureConfig {
|
|
input_dim: usize::try_from(mamba.state_dim).unwrap_or(128),
|
|
hidden_dims: vec![usize::try_from(mamba.hidden_dim).unwrap_or(512)],
|
|
output_dim: 1,
|
|
dropout_rate: 0.1,
|
|
activation: "relu".to_string(),
|
|
batch_norm: true,
|
|
residual_connections: false,
|
|
};
|
|
|
|
config.training_params = TrainingHyperparameters {
|
|
learning_rate: mamba.learning_rate as f64,
|
|
batch_size: usize::try_from(mamba.batch_size).unwrap_or(32),
|
|
max_epochs: usize::try_from(mamba.epochs).unwrap_or(150),
|
|
patience: 50,
|
|
validation_split: 0.2,
|
|
l2_regularization: 1e-4,
|
|
lr_decay_factor: 0.5,
|
|
lr_decay_patience: 25,
|
|
};
|
|
},
|
|
|
|
Some(proto::hyperparameters::ModelParams::DqnParams(dqn)) => {
|
|
config.training_params = TrainingHyperparameters {
|
|
learning_rate: dqn.learning_rate as f64,
|
|
batch_size: usize::try_from(dqn.batch_size).unwrap_or(128),
|
|
max_epochs: usize::try_from(dqn.epochs).unwrap_or(200),
|
|
patience: 100,
|
|
validation_split: 0.1,
|
|
l2_regularization: 1e-5,
|
|
lr_decay_factor: 0.8,
|
|
lr_decay_patience: 50,
|
|
};
|
|
},
|
|
|
|
// Add other model parameter conversions...
|
|
_ => {
|
|
// Use default configuration
|
|
},
|
|
}
|
|
}
|
|
|
|
// Apply GPU preference
|
|
if config.performance_config.device_preference == "cpu" {
|
|
config.performance_config.device_preference = "cuda".to_string();
|
|
}
|
|
|
|
Ok(config)
|
|
}
|
|
|
|
/// Convert internal status update to protobuf
|
|
fn convert_status_update(update: TrainingStatusUpdate) -> ProtoStatusUpdate {
|
|
let financial_metrics = update.financial_metrics.map(|fm| ProtoFinancialMetrics {
|
|
simulated_return: fm.simulated_return as f32,
|
|
sharpe_ratio: fm.sharpe_ratio as f32,
|
|
max_drawdown: fm.max_drawdown as f32,
|
|
hit_rate: fm.hit_rate as f32,
|
|
avg_prediction_error_bps: fm.avg_prediction_error_bps as f32,
|
|
risk_adjusted_return: fm.risk_adjusted_return as f32,
|
|
var_5pct: 0.0, // Would be populated from actual metrics
|
|
expected_shortfall: 0.0,
|
|
});
|
|
|
|
let resource_usage = ProtoResourceUsage {
|
|
cpu_usage_percent: update.resource_usage.cpu_usage_percent,
|
|
memory_usage_gb: update.resource_usage.memory_usage_gb,
|
|
gpu_usage_percent: update.resource_usage.gpu_usage_percent.unwrap_or(0.0),
|
|
gpu_memory_usage_gb: update.resource_usage.gpu_memory_usage_gb.unwrap_or(0.0),
|
|
active_workers: update.resource_usage.active_workers,
|
|
};
|
|
|
|
ProtoStatusUpdate {
|
|
job_id: update.job_id.to_string(),
|
|
status: (Self::convert_job_status(&update.status) as i32),
|
|
progress_percentage: update.progress_percentage,
|
|
current_epoch: update.current_epoch,
|
|
total_epochs: update.total_epochs,
|
|
metrics: update
|
|
.metrics
|
|
.into_iter()
|
|
.map(|(k, v)| (k, v as f32))
|
|
.collect(),
|
|
message: update.message,
|
|
timestamp: update.timestamp.timestamp(),
|
|
financial_metrics,
|
|
resource_usage: Some(resource_usage),
|
|
}
|
|
}
|
|
}
|
|
|
|
#[tonic::async_trait]
|
|
impl MlTrainingService for MLTrainingServiceImpl {
|
|
/// Start a new training job
|
|
#[instrument(skip_all)]
|
|
async fn start_training(
|
|
&self,
|
|
request: Request<StartTrainingRequest>,
|
|
) -> Result<Response<StartTrainingResponse>, Status> {
|
|
let req = request.into_inner();
|
|
|
|
info!("Starting training job for model type: {}", req.model_type);
|
|
|
|
// Detect hyperopt mode from tags
|
|
let is_hyperopt = req.tags.get("mode").map(|m| m == "hyperopt").unwrap_or(false);
|
|
|
|
// Determine the correct binary
|
|
let binary = if is_hyperopt {
|
|
crate::k8s_dispatcher::hyperopt_binary_for_model(&req.model_type)
|
|
} else {
|
|
crate::k8s_dispatcher::training_binary_for_model(&req.model_type)
|
|
};
|
|
|
|
if is_hyperopt {
|
|
info!("Hyperopt mode: using binary {}", binary);
|
|
}
|
|
|
|
// Generate a job ID
|
|
let job_id = Uuid::new_v4();
|
|
|
|
// Extract symbol from data_source file_path (take last path component), or use default
|
|
let symbol = req
|
|
.data_source
|
|
.as_ref()
|
|
.and_then(|ds| match &ds.source {
|
|
Some(proto::data_source::Source::FilePath(p)) => {
|
|
// file_path may be "data/cache/futures-baseline/ES.FUT" — extract just the symbol
|
|
std::path::Path::new(p)
|
|
.file_name()
|
|
.and_then(|f| f.to_str())
|
|
.map(String::from)
|
|
}
|
|
_ => None,
|
|
})
|
|
.unwrap_or_else(|| "ES.FUT".to_string());
|
|
|
|
// Priority 1: Batch flow — persist job to DB, queue consumer dispatches to K8s
|
|
if let Some(ref spawner) = self.job_spawner {
|
|
let model = common::model_types::ModelType::from_str(&req.model_type).ok_or_else(
|
|
|| Status::invalid_argument(format!("Unknown model type: {}", req.model_type)),
|
|
)?;
|
|
|
|
let asset = crate::job_spawner::Asset {
|
|
symbol: symbol.clone(),
|
|
data_file: std::path::PathBuf::from(format!(
|
|
"/data/futures-baseline/{}",
|
|
symbol
|
|
)),
|
|
};
|
|
|
|
match spawner.spawn_batch(vec![asset], vec![model]).await {
|
|
Ok(batch) => {
|
|
info!(batch_id = %batch.batch_id, "Batch job created in database");
|
|
return Ok(Response::new(StartTrainingResponse {
|
|
job_id: batch.batch_id.to_string(),
|
|
status: ProtoTrainingStatus::Pending as i32,
|
|
message: format!("Batch training job queued: {}", batch.batch_id),
|
|
}));
|
|
}
|
|
Err(e) => {
|
|
warn!("JobSpawner failed, falling back to K8s/orchestrator: {}", e);
|
|
// Fall through to K8s dispatcher or in-process orchestrator
|
|
}
|
|
}
|
|
}
|
|
|
|
// Priority 2: Direct K8s dispatch
|
|
if let Some(ref dispatcher) = self.k8s_dispatcher {
|
|
let epochs = req
|
|
.hyperparameters
|
|
.as_ref()
|
|
.and_then(|hp| match &hp.model_params {
|
|
Some(proto::hyperparameters::ModelParams::TlobParams(p)) => {
|
|
Some(p.epochs)
|
|
}
|
|
Some(proto::hyperparameters::ModelParams::DqnParams(p)) => {
|
|
Some(p.epochs)
|
|
}
|
|
Some(proto::hyperparameters::ModelParams::PpoParams(p)) => {
|
|
Some(p.epochs)
|
|
}
|
|
Some(proto::hyperparameters::ModelParams::MambaParams(p)) => {
|
|
Some(p.epochs)
|
|
}
|
|
Some(proto::hyperparameters::ModelParams::LiquidParams(p)) => {
|
|
Some(p.epochs)
|
|
}
|
|
Some(proto::hyperparameters::ModelParams::TftParams(p)) => {
|
|
Some(p.epochs)
|
|
}
|
|
None => None,
|
|
})
|
|
.unwrap_or(100);
|
|
|
|
let extra_args = if is_hyperopt {
|
|
let trials = req.tags.get("trials")
|
|
.and_then(|t| t.parse::<u32>().ok())
|
|
.unwrap_or(20);
|
|
let parallel = req.tags.get("parallel")
|
|
.and_then(|p| p.parse::<u32>().ok())
|
|
.unwrap_or(0);
|
|
vec![
|
|
"--output=/output/hyperopt_results.json".to_string(),
|
|
format!("--model={}", req.model_type.to_lowercase()),
|
|
format!("--trials={}", trials),
|
|
format!("--epochs={}", epochs),
|
|
format!("--parallel={}", parallel),
|
|
]
|
|
} else {
|
|
vec![]
|
|
};
|
|
|
|
let params = crate::k8s_dispatcher::TrainingJobParams {
|
|
job_id,
|
|
model_type: req.model_type.to_lowercase(),
|
|
symbol,
|
|
epochs,
|
|
binary: binary.to_string(),
|
|
extra_args,
|
|
};
|
|
|
|
match dispatcher.dispatch(¶ms).await {
|
|
Ok(k8s_job_name) => {
|
|
info!("K8s Job dispatched: {} (job_id={})", k8s_job_name, job_id);
|
|
return Ok(Response::new(StartTrainingResponse {
|
|
job_id: job_id.to_string(),
|
|
status: (ProtoTrainingStatus::Pending as i32),
|
|
message: format!("K8s training job dispatched: {}", k8s_job_name),
|
|
}));
|
|
}
|
|
Err(e) => {
|
|
warn!("K8s dispatch failed, falling back to in-process: {}", e);
|
|
// Fall through to orchestrator
|
|
}
|
|
}
|
|
}
|
|
|
|
// Fallback: in-process orchestrator (for development or if K8s dispatch fails)
|
|
let training_config = self
|
|
.convert_hyperparameters(req.hyperparameters)
|
|
.map_err(|e| Status::invalid_argument(format!("Invalid hyperparameters: {}", e)))?;
|
|
|
|
let tags: HashMap<String, String> = req.tags.into_iter().collect();
|
|
|
|
let job_id = self
|
|
.orchestrator
|
|
.submit_job(req.model_type, training_config, req.description, tags)
|
|
.await
|
|
.map_err(|e| Status::internal(format!("Failed to submit job: {}", e)))?;
|
|
|
|
info!("Training job submitted to orchestrator: {}", job_id);
|
|
|
|
Ok(Response::new(StartTrainingResponse {
|
|
job_id: job_id.to_string(),
|
|
status: (ProtoTrainingStatus::Pending as i32),
|
|
message: "Training job submitted successfully".to_string(),
|
|
}))
|
|
}
|
|
|
|
/// Subscribe to training status updates
|
|
type SubscribeToTrainingStatusStream =
|
|
Pin<Box<dyn Stream<Item = Result<ProtoStatusUpdate, Status>> + Send>>;
|
|
|
|
#[instrument(skip_all)]
|
|
async fn subscribe_to_training_status(
|
|
&self,
|
|
request: Request<SubscribeToTrainingStatusRequest>,
|
|
) -> Result<Response<Self::SubscribeToTrainingStatusStream>, Status> {
|
|
let req = request.into_inner();
|
|
|
|
let job_id = Uuid::parse_str(&req.job_id)
|
|
.map_err(|_| Status::invalid_argument("Invalid job ID format"))?;
|
|
|
|
debug!("Subscribing to status updates for job: {}", job_id);
|
|
|
|
// Get status update receiver from orchestrator with timeout to prevent hanging
|
|
let receiver = timeout(
|
|
Duration::from_secs(5),
|
|
self.orchestrator.subscribe_to_job_status(job_id),
|
|
)
|
|
.await
|
|
.map_err(|_| Status::deadline_exceeded("Timeout subscribing to job status"))?
|
|
.map_err(|e| Status::not_found(format!("Job not found: {}", e)))?;
|
|
|
|
// Convert the broadcast receiver to a stream with proper error handling
|
|
let stream = async_stream::stream! {
|
|
let mut receiver = receiver;
|
|
loop {
|
|
match receiver.recv().await {
|
|
Ok(update) => yield Ok(Self::convert_status_update(update)),
|
|
Err(_) => break,
|
|
}
|
|
}
|
|
};
|
|
|
|
Ok(Response::new(Box::pin(stream)))
|
|
}
|
|
|
|
/// Stop a running training job
|
|
#[instrument(skip_all)]
|
|
async fn stop_training(
|
|
&self,
|
|
request: Request<StopTrainingRequest>,
|
|
) -> Result<Response<StopTrainingResponse>, Status> {
|
|
let req = request.into_inner();
|
|
|
|
let job_id = Uuid::parse_str(&req.job_id)
|
|
.map_err(|_| Status::invalid_argument("Invalid job ID format"))?;
|
|
|
|
info!("Stopping training job: {} (reason: {})", job_id, req.reason);
|
|
|
|
let success = self
|
|
.orchestrator
|
|
.stop_job(job_id, req.reason)
|
|
.await
|
|
.map_err(|e| Status::internal(format!("Failed to stop job: {}", e)))?;
|
|
|
|
let message = if success {
|
|
"Training job stopped successfully".to_string()
|
|
} else {
|
|
"Job was not running or already completed".to_string()
|
|
};
|
|
|
|
let response = StopTrainingResponse { success, message };
|
|
|
|
Ok(Response::new(response))
|
|
}
|
|
|
|
/// List available models for training
|
|
#[instrument(skip_all)]
|
|
async fn list_available_models(
|
|
&self,
|
|
_request: Request<ListAvailableModelsRequest>,
|
|
) -> Result<Response<ListAvailableModelsResponse>, Status> {
|
|
debug!("Listing available models");
|
|
|
|
// Define available models with their configurations
|
|
let models = vec![
|
|
ModelDefinition {
|
|
model_type: "TLOB".to_string(),
|
|
description:
|
|
"Time-Limit Order Book Transformer for ultra-low latency order book prediction"
|
|
.to_string(),
|
|
default_hyperparameters: Some(Hyperparameters {
|
|
model_params: Some(proto::hyperparameters::ModelParams::TlobParams(
|
|
TlobParams {
|
|
epochs: 100,
|
|
learning_rate: 0.001,
|
|
batch_size: 64,
|
|
sequence_length: 50,
|
|
hidden_dim: 256,
|
|
num_heads: 8,
|
|
num_layers: 6,
|
|
dropout_rate: 0.1,
|
|
use_positional_encoding: true,
|
|
},
|
|
)),
|
|
}),
|
|
required_features: vec![
|
|
"order_book_levels".to_string(),
|
|
"trade_flow".to_string(),
|
|
"volume_imbalance".to_string(),
|
|
],
|
|
estimated_training_time_minutes: 45,
|
|
requires_gpu: true,
|
|
},
|
|
ModelDefinition {
|
|
model_type: "MAMBA_2".to_string(),
|
|
description: "MAMBA-2 State Space Model for long-sequence financial time series"
|
|
.to_string(),
|
|
default_hyperparameters: Some(Hyperparameters {
|
|
model_params: Some(proto::hyperparameters::ModelParams::MambaParams(
|
|
MambaParams {
|
|
epochs: 150,
|
|
learning_rate: 0.0005,
|
|
batch_size: 32,
|
|
state_dim: 128,
|
|
hidden_dim: 512,
|
|
num_layers: 8,
|
|
dt_min: 0.001,
|
|
dt_max: 0.1,
|
|
use_cuda_kernels: true,
|
|
},
|
|
)),
|
|
}),
|
|
required_features: vec![
|
|
"price_series".to_string(),
|
|
"volume_series".to_string(),
|
|
"technical_indicators".to_string(),
|
|
],
|
|
estimated_training_time_minutes: 90,
|
|
requires_gpu: true,
|
|
},
|
|
ModelDefinition {
|
|
model_type: "DQN".to_string(),
|
|
description: "Deep Q-Network for reinforcement learning-based trading strategies"
|
|
.to_string(),
|
|
default_hyperparameters: Some(Hyperparameters {
|
|
model_params: Some(proto::hyperparameters::ModelParams::DqnParams(DqnParams {
|
|
epochs: 200,
|
|
learning_rate: 0.0001,
|
|
batch_size: 128,
|
|
replay_buffer_size: 100000,
|
|
epsilon_start: 1.0,
|
|
epsilon_end: 0.01,
|
|
epsilon_decay_steps: 50000,
|
|
gamma: 0.99,
|
|
target_update_frequency: 1000,
|
|
use_double_dqn: true,
|
|
use_dueling: true,
|
|
use_prioritized_replay: true,
|
|
})),
|
|
}),
|
|
required_features: vec![
|
|
"market_state".to_string(),
|
|
"portfolio_state".to_string(),
|
|
"risk_metrics".to_string(),
|
|
],
|
|
estimated_training_time_minutes: 120,
|
|
requires_gpu: true,
|
|
},
|
|
ModelDefinition {
|
|
model_type: "PPO".to_string(),
|
|
description: "Proximal Policy Optimization for continuous action space trading"
|
|
.to_string(),
|
|
default_hyperparameters: Some(Hyperparameters {
|
|
model_params: Some(proto::hyperparameters::ModelParams::PpoParams(PpoParams {
|
|
epochs: 100,
|
|
learning_rate: 0.0003,
|
|
batch_size: 64,
|
|
clip_ratio: 0.2,
|
|
value_loss_coef: 0.5,
|
|
entropy_coef: 0.01,
|
|
rollout_steps: 2048,
|
|
minibatch_size: 64,
|
|
gae_lambda: 0.95,
|
|
})),
|
|
}),
|
|
required_features: vec![
|
|
"market_state".to_string(),
|
|
"position_state".to_string(),
|
|
"risk_state".to_string(),
|
|
],
|
|
estimated_training_time_minutes: 75,
|
|
requires_gpu: true,
|
|
},
|
|
ModelDefinition {
|
|
model_type: "LIQUID".to_string(),
|
|
description: "Liquid Neural Network for adaptive market regime detection"
|
|
.to_string(),
|
|
default_hyperparameters: Some(Hyperparameters {
|
|
model_params: Some(proto::hyperparameters::ModelParams::LiquidParams(
|
|
LiquidParams {
|
|
epochs: 80,
|
|
learning_rate: 0.002,
|
|
batch_size: 48,
|
|
num_neurons: 128,
|
|
tau: 0.1,
|
|
sigma: 0.5,
|
|
use_adaptive_tau: true,
|
|
},
|
|
)),
|
|
}),
|
|
required_features: vec![
|
|
"volatility_regime".to_string(),
|
|
"liquidity_metrics".to_string(),
|
|
"market_microstructure".to_string(),
|
|
],
|
|
estimated_training_time_minutes: 60,
|
|
requires_gpu: false,
|
|
},
|
|
ModelDefinition {
|
|
model_type: "TFT".to_string(),
|
|
description: "Temporal Fusion Transformer for multi-horizon forecasting"
|
|
.to_string(),
|
|
default_hyperparameters: Some(Hyperparameters {
|
|
model_params: Some(proto::hyperparameters::ModelParams::TftParams(TftParams {
|
|
epochs: 120,
|
|
learning_rate: 0.001,
|
|
batch_size: 32,
|
|
hidden_dim: 240,
|
|
num_heads: 4,
|
|
num_layers: 3,
|
|
lookback_window: 168,
|
|
forecast_horizon: 24,
|
|
dropout_rate: 0.3,
|
|
})),
|
|
}),
|
|
required_features: vec![
|
|
"historical_prices".to_string(),
|
|
"external_regressors".to_string(),
|
|
"calendar_features".to_string(),
|
|
],
|
|
estimated_training_time_minutes: 100,
|
|
requires_gpu: true,
|
|
},
|
|
];
|
|
|
|
let response = ListAvailableModelsResponse { models };
|
|
Ok(Response::new(response))
|
|
}
|
|
|
|
/// List training jobs with filtering
|
|
#[instrument(skip_all)]
|
|
async fn list_training_jobs(
|
|
&self,
|
|
request: Request<ListTrainingJobsRequest>,
|
|
) -> Result<Response<ListTrainingJobsResponse>, Status> {
|
|
let req = request.into_inner();
|
|
|
|
debug!("Listing training jobs with filters");
|
|
|
|
// Convert status filter
|
|
let status_filter = if req.status_filter != ProtoTrainingStatus::Unknown as i32 {
|
|
Some(match req.status_filter {
|
|
x if x == ProtoTrainingStatus::Pending as i32 => JobStatus::Pending,
|
|
x if x == ProtoTrainingStatus::Running as i32 => JobStatus::Running,
|
|
x if x == ProtoTrainingStatus::Completed as i32 => JobStatus::Completed,
|
|
x if x == ProtoTrainingStatus::Failed as i32 => JobStatus::Failed,
|
|
x if x == ProtoTrainingStatus::Stopped as i32 => JobStatus::Stopped,
|
|
x if x == ProtoTrainingStatus::Paused as i32 => JobStatus::Paused,
|
|
_ => return Err(Status::invalid_argument("Invalid status filter")),
|
|
})
|
|
} else {
|
|
None
|
|
};
|
|
|
|
let model_type_filter = if req.model_type_filter.is_empty() {
|
|
None
|
|
} else {
|
|
Some(req.model_type_filter)
|
|
};
|
|
|
|
let limit = if req.page_size == 0 {
|
|
None
|
|
} else {
|
|
Some(usize::try_from(req.page_size).unwrap_or(50))
|
|
};
|
|
let offset = if req.page == 0 {
|
|
None
|
|
} else {
|
|
Some(usize::try_from((req.page - 1) * req.page_size).unwrap_or(0))
|
|
};
|
|
|
|
// Get jobs from orchestrator
|
|
let jobs = self
|
|
.orchestrator
|
|
.list_jobs(status_filter, model_type_filter, limit, offset)
|
|
.await
|
|
.map_err(|e| Status::internal(format!("Failed to list jobs: {}", e)))?;
|
|
|
|
// Convert to protobuf format
|
|
let job_summaries: Vec<TrainingJobSummary> = jobs
|
|
.into_iter()
|
|
.map(|job| TrainingJobSummary {
|
|
job_id: job.id.to_string(),
|
|
model_type: job.model_type,
|
|
status: Self::convert_job_status(&job.status) as i32,
|
|
created_at: job.created_at.timestamp(),
|
|
started_at: job.started_at.map(|dt| dt.timestamp()).unwrap_or(0),
|
|
completed_at: job.completed_at.map(|dt| dt.timestamp()).unwrap_or(0),
|
|
description: job.description,
|
|
final_loss: job.metrics.get("final_train_loss").copied().unwrap_or(0.0) as f32,
|
|
best_validation_score: job.metrics.get("final_val_loss").copied().unwrap_or(0.0)
|
|
as f32,
|
|
tags: job.tags,
|
|
})
|
|
.collect();
|
|
|
|
let response = ListTrainingJobsResponse {
|
|
jobs: job_summaries,
|
|
total_count: 0, // Would implement proper counting
|
|
page: req.page,
|
|
page_size: req.page_size,
|
|
};
|
|
|
|
Ok(Response::new(response))
|
|
}
|
|
|
|
/// Get detailed information about a training job
|
|
#[instrument(skip_all)]
|
|
async fn get_training_job_details(
|
|
&self,
|
|
request: Request<GetTrainingJobDetailsRequest>,
|
|
) -> Result<Response<GetTrainingJobDetailsResponse>, Status> {
|
|
let req = request.into_inner();
|
|
|
|
let job_id = Uuid::parse_str(&req.job_id)
|
|
.map_err(|_| Status::invalid_argument("Invalid job ID format"))?;
|
|
|
|
debug!("Getting details for training job: {}", job_id);
|
|
|
|
// Get job from orchestrator
|
|
let job = self
|
|
.orchestrator
|
|
.get_job(job_id)
|
|
.await
|
|
.map_err(|e| Status::not_found(format!("Job not found: {}", e)))?;
|
|
|
|
// Convert to detailed response
|
|
let job_details = TrainingJobDetails {
|
|
job_id: job.id.to_string(),
|
|
model_type: job.model_type,
|
|
status: Self::convert_job_status(&job.status) as i32,
|
|
created_at: job.created_at.timestamp(),
|
|
started_at: job.started_at.map(|dt| dt.timestamp()).unwrap_or(0),
|
|
completed_at: job.completed_at.map(|dt| dt.timestamp()).unwrap_or(0),
|
|
description: job.description,
|
|
hyperparameters: None, // Would serialize from job.config
|
|
data_source: None, // Would populate from job configuration
|
|
status_history: Vec::new(), // Would populate from database
|
|
final_financial_metrics: None, // Would populate from final results
|
|
model_artifact_path: job.model_artifact_path.unwrap_or_default(),
|
|
tags: job.tags,
|
|
error_message: job.error_message.unwrap_or_default(),
|
|
};
|
|
|
|
let response = GetTrainingJobDetailsResponse {
|
|
job_details: Some(job_details),
|
|
};
|
|
|
|
Ok(Response::new(response))
|
|
}
|
|
|
|
/// Health check
|
|
#[instrument(skip_all)]
|
|
async fn health_check(
|
|
&self,
|
|
_request: Request<HealthCheckRequest>,
|
|
) -> Result<Response<HealthCheckResponse>, Status> {
|
|
debug!("Health check requested");
|
|
|
|
let mut details = HashMap::new();
|
|
details.insert("service".to_string(), "ml_training_service".to_string());
|
|
details.insert("version".to_string(), "0.1.0".to_string());
|
|
details.insert("uptime".to_string(), "active".to_string());
|
|
|
|
let response = HealthCheckResponse {
|
|
healthy: true,
|
|
message: "ML Training Service is healthy".to_string(),
|
|
details,
|
|
};
|
|
|
|
Ok(Response::new(response))
|
|
}
|
|
|
|
/// Start a new hyperparameter tuning job
|
|
#[instrument(skip_all)]
|
|
async fn start_tuning_job(
|
|
&self,
|
|
request: Request<proto::StartTuningJobRequest>,
|
|
) -> Result<Response<proto::StartTuningJobResponse>, Status> {
|
|
self.tuning_handlers.start_tuning_job(request).await
|
|
}
|
|
|
|
/// Get status of a tuning job
|
|
#[instrument(skip_all)]
|
|
async fn get_tuning_job_status(
|
|
&self,
|
|
request: Request<proto::GetTuningJobStatusRequest>,
|
|
) -> Result<Response<proto::GetTuningJobStatusResponse>, Status> {
|
|
self.tuning_handlers.get_tuning_job_status(request).await
|
|
}
|
|
|
|
/// Stop a running tuning job
|
|
#[instrument(skip_all)]
|
|
async fn stop_tuning_job(
|
|
&self,
|
|
request: Request<proto::StopTuningJobRequest>,
|
|
) -> Result<Response<proto::StopTuningJobResponse>, Status> {
|
|
self.tuning_handlers.stop_tuning_job(request).await
|
|
}
|
|
|
|
/// Stream tuning progress updates
|
|
type StreamTuningProgressStream =
|
|
Pin<Box<dyn Stream<Item = Result<proto::ProgressUpdate, Status>> + Send>>;
|
|
|
|
#[instrument(skip_all)]
|
|
async fn stream_tuning_progress(
|
|
&self,
|
|
request: Request<proto::StreamProgressRequest>,
|
|
) -> Result<Response<Self::StreamTuningProgressStream>, Status> {
|
|
self.tuning_handlers.stream_tuning_progress(request).await
|
|
}
|
|
|
|
/// INTERNAL: Train a single model with specific hyperparameters (called by Optuna subprocess)
|
|
#[instrument(skip_all)]
|
|
async fn train_model(
|
|
&self,
|
|
request: Request<proto::TrainModelRequest>,
|
|
) -> Result<Response<proto::TrainModelResponse>, Status> {
|
|
let req = request.into_inner();
|
|
|
|
info!(
|
|
"INTERNAL TrainModel called by Optuna - trial_id: {}, model_type: {}",
|
|
req.trial_id, req.model_type
|
|
);
|
|
|
|
let start_time = std::time::Instant::now();
|
|
|
|
// 1. Parse hyperparameters from request
|
|
let hyperparams = req.hyperparameters;
|
|
|
|
// 2. Convert hyperparameters to training config
|
|
let training_config =
|
|
match self.convert_hyperparameters_from_map(&req.model_type, &hyperparams) {
|
|
Ok(config) => config,
|
|
Err(e) => {
|
|
return Ok(Response::new(proto::TrainModelResponse {
|
|
success: false,
|
|
sharpe_ratio: 0.0,
|
|
training_loss: f32::INFINITY,
|
|
validation_metrics: HashMap::new(),
|
|
error_message: format!("Failed to parse hyperparameters: {}", e),
|
|
training_duration_seconds: 0,
|
|
}));
|
|
},
|
|
};
|
|
|
|
// 3. Load training data (reuse existing pipeline)
|
|
let (training_data, validation_data) =
|
|
match crate::orchestrator::TrainingOrchestrator::load_training_data().await {
|
|
Ok(data) => data,
|
|
Err(e) => {
|
|
return Ok(Response::new(proto::TrainModelResponse {
|
|
success: false,
|
|
sharpe_ratio: 0.0,
|
|
training_loss: f32::INFINITY,
|
|
validation_metrics: HashMap::new(),
|
|
error_message: format!("Failed to load training data: {}", e),
|
|
training_duration_seconds: 0,
|
|
}));
|
|
},
|
|
};
|
|
|
|
// 4. Create training system and execute training
|
|
let training_system =
|
|
match ml::training_pipeline::ProductionMLTrainingSystem::new(training_config).await {
|
|
Ok(system) => system,
|
|
Err(e) => {
|
|
return Ok(Response::new(proto::TrainModelResponse {
|
|
success: false,
|
|
sharpe_ratio: 0.0,
|
|
training_loss: f32::INFINITY,
|
|
validation_metrics: HashMap::new(),
|
|
error_message: format!("Failed to create training system: {:?}", e),
|
|
training_duration_seconds: 0,
|
|
}));
|
|
},
|
|
};
|
|
|
|
// 5. Train model
|
|
let result = match training_system
|
|
.train_model(training_data, Some(validation_data.clone()))
|
|
.await
|
|
{
|
|
Ok(r) => r,
|
|
Err(e) => {
|
|
return Ok(Response::new(proto::TrainModelResponse {
|
|
success: false,
|
|
sharpe_ratio: 0.0,
|
|
training_loss: f32::INFINITY,
|
|
validation_metrics: HashMap::new(),
|
|
error_message: format!("Training failed: {:?}", e),
|
|
training_duration_seconds: 0,
|
|
}));
|
|
},
|
|
};
|
|
|
|
// 6. Calculate Sharpe ratio on validation set using backtesting
|
|
let sharpe_ratio = self
|
|
.calculate_sharpe_ratio_from_validation(&validation_data, &result)
|
|
.await;
|
|
|
|
// 7. Collect validation metrics
|
|
let mut validation_metrics = HashMap::new();
|
|
validation_metrics.insert("validation_loss".to_string(), result.final_val_loss as f32);
|
|
validation_metrics.insert("train_loss".to_string(), result.final_train_loss as f32);
|
|
|
|
// Add financial metrics if available
|
|
if let Some(last_metrics) = result.metrics_history.last() {
|
|
validation_metrics.insert(
|
|
"hit_rate".to_string(),
|
|
last_metrics.financial_metrics.hit_rate as f32,
|
|
);
|
|
validation_metrics.insert(
|
|
"avg_prediction_error_bps".to_string(),
|
|
last_metrics.financial_metrics.avg_prediction_error_bps as f32,
|
|
);
|
|
}
|
|
|
|
let duration = start_time.elapsed().as_secs() as i64;
|
|
|
|
info!(
|
|
"INTERNAL TrainModel completed - trial_id: {}, sharpe_ratio: {:.4}, training_loss: {:.6}, duration: {}s",
|
|
req.trial_id, sharpe_ratio, result.final_train_loss, duration
|
|
);
|
|
|
|
// 8. Return response with Sharpe ratio as primary objective
|
|
Ok(Response::new(proto::TrainModelResponse {
|
|
success: true,
|
|
sharpe_ratio,
|
|
training_loss: result.final_train_loss as f32,
|
|
validation_metrics,
|
|
error_message: String::new(),
|
|
training_duration_seconds: duration,
|
|
}))
|
|
}
|
|
|
|
/// Start batch tuning job for multiple models
|
|
#[instrument(skip_all)]
|
|
async fn batch_start_tuning_jobs(
|
|
&self,
|
|
request: Request<proto::BatchStartTuningJobsRequest>,
|
|
) -> Result<Response<proto::BatchStartTuningJobsResponse>, Status> {
|
|
let req = request.into_inner();
|
|
|
|
info!(
|
|
"Starting batch tuning job for models: {:?} with {} trials each",
|
|
req.model_types, req.trials_per_model
|
|
);
|
|
|
|
// Validate inputs
|
|
if req.model_types.is_empty() {
|
|
return Err(Status::invalid_argument(
|
|
"At least one model type must be specified",
|
|
));
|
|
}
|
|
|
|
if req.trials_per_model == 0 {
|
|
return Err(Status::invalid_argument(
|
|
"Number of trials per model must be > 0",
|
|
));
|
|
}
|
|
|
|
if req.config_path.is_empty() {
|
|
return Err(Status::invalid_argument(
|
|
"Configuration path cannot be empty",
|
|
));
|
|
}
|
|
|
|
// Determine data source path from the proto DataSource
|
|
let data_source_path = req.data_source.and_then(|ds| {
|
|
ds.source.map(|s| match s {
|
|
proto::data_source::Source::FilePath(p) => p,
|
|
proto::data_source::Source::HistoricalDbQuery(q) => q,
|
|
proto::data_source::Source::RealTimeStreamTopic(t) => t,
|
|
})
|
|
});
|
|
|
|
// Determine YAML export path (use custom or default)
|
|
let yaml_export_path = if req.yaml_export_path.is_empty() {
|
|
None
|
|
} else {
|
|
Some(req.yaml_export_path)
|
|
};
|
|
|
|
// Start batch tuning via manager
|
|
let batch_id = self
|
|
.batch_tuning_manager
|
|
.start_batch_tuning(
|
|
req.model_types,
|
|
req.trials_per_model,
|
|
req.config_path,
|
|
data_source_path,
|
|
req.auto_export_yaml,
|
|
yaml_export_path,
|
|
)
|
|
.await
|
|
.map_err(|e| Status::internal(format!("Failed to start batch tuning: {}", e)))?;
|
|
|
|
// Get the job to return the execution order
|
|
let job = self
|
|
.batch_tuning_manager
|
|
.get_batch_status(batch_id)
|
|
.await
|
|
.map_err(|e| {
|
|
Status::internal(format!("Batch started but failed to retrieve status: {}", e))
|
|
})?;
|
|
|
|
info!(
|
|
"Batch tuning job started: {} with execution order: {:?}",
|
|
batch_id, job.execution_order
|
|
);
|
|
|
|
let response = proto::BatchStartTuningJobsResponse {
|
|
batch_id: batch_id.to_string(),
|
|
execution_order: job.execution_order,
|
|
message: format!("Batch tuning started for {} models", job.models.len()),
|
|
status: Self::convert_batch_status(&job.status) as i32,
|
|
};
|
|
|
|
Ok(Response::new(response))
|
|
}
|
|
|
|
/// Get batch tuning job status
|
|
#[instrument(skip_all)]
|
|
async fn get_batch_tuning_status(
|
|
&self,
|
|
request: Request<proto::GetBatchTuningStatusRequest>,
|
|
) -> Result<Response<proto::GetBatchTuningStatusResponse>, Status> {
|
|
let req = request.into_inner();
|
|
|
|
let batch_id = Uuid::parse_str(&req.batch_id)
|
|
.map_err(|_| Status::invalid_argument("Invalid batch ID format"))?;
|
|
|
|
debug!("Getting batch tuning status for: {}", batch_id);
|
|
|
|
let job = self
|
|
.batch_tuning_manager
|
|
.get_batch_status(batch_id)
|
|
.await
|
|
.map_err(|e| Status::not_found(format!("Batch job not found: {}", e)))?;
|
|
|
|
// Convert model results to proto format
|
|
let results: Vec<proto::ModelTuningResult> = job
|
|
.results
|
|
.iter()
|
|
.map(|r| proto::ModelTuningResult {
|
|
model_type: r.model_type.clone(),
|
|
job_id: r.job_id.to_string(),
|
|
status: Self::convert_tuning_job_status(&r.status) as i32,
|
|
best_params: r.best_params.clone(),
|
|
best_metrics: r.best_metrics.clone(),
|
|
trials_completed: r.trials_completed,
|
|
started_at: r.started_at.timestamp(),
|
|
completed_at: r.completed_at.map(|dt| dt.timestamp()).unwrap_or(0),
|
|
error_message: r.error_message.clone().unwrap_or_default(),
|
|
})
|
|
.collect();
|
|
|
|
// Determine the current model name
|
|
let current_model = job
|
|
.execution_order
|
|
.get(job.current_model_index)
|
|
.cloned()
|
|
.unwrap_or_default();
|
|
|
|
let response = proto::GetBatchTuningStatusResponse {
|
|
batch_id: job.batch_id.to_string(),
|
|
status: Self::convert_batch_status(&job.status) as i32,
|
|
current_model_index: job.current_model_index as u32,
|
|
total_models: job.execution_order.len() as u32,
|
|
results,
|
|
current_model,
|
|
started_at: job.started_at.timestamp(),
|
|
updated_at: job.updated_at.timestamp(),
|
|
estimated_completion_time: 0, // Not yet estimated
|
|
yaml_export_path: job.yaml_export_path.clone(),
|
|
};
|
|
|
|
Ok(Response::new(response))
|
|
}
|
|
|
|
/// Stop a running batch tuning job
|
|
#[instrument(skip_all)]
|
|
async fn stop_batch_tuning_job(
|
|
&self,
|
|
request: Request<proto::StopBatchTuningJobRequest>,
|
|
) -> Result<Response<proto::StopBatchTuningJobResponse>, Status> {
|
|
let req = request.into_inner();
|
|
|
|
let batch_id = Uuid::parse_str(&req.batch_id)
|
|
.map_err(|_| Status::invalid_argument("Invalid batch ID format"))?;
|
|
|
|
info!(
|
|
"Stopping batch tuning job: {} (reason: {})",
|
|
batch_id, req.reason
|
|
);
|
|
|
|
// Stop the batch job
|
|
self.batch_tuning_manager
|
|
.stop_batch_job(batch_id, req.reason.clone())
|
|
.await
|
|
.map_err(|e| Status::internal(format!("Failed to stop batch job: {}", e)))?;
|
|
|
|
// Get final status
|
|
let job = self
|
|
.batch_tuning_manager
|
|
.get_batch_status(batch_id)
|
|
.await
|
|
.map_err(|e| Status::not_found(format!("Batch job not found after stop: {}", e)))?;
|
|
|
|
// Convert completed results
|
|
let completed_results: Vec<proto::ModelTuningResult> = job
|
|
.results
|
|
.iter()
|
|
.map(|r| proto::ModelTuningResult {
|
|
model_type: r.model_type.clone(),
|
|
job_id: r.job_id.to_string(),
|
|
status: Self::convert_tuning_job_status(&r.status) as i32,
|
|
best_params: r.best_params.clone(),
|
|
best_metrics: r.best_metrics.clone(),
|
|
trials_completed: r.trials_completed,
|
|
started_at: r.started_at.timestamp(),
|
|
completed_at: r.completed_at.map(|dt| dt.timestamp()).unwrap_or(0),
|
|
error_message: r.error_message.clone().unwrap_or_default(),
|
|
})
|
|
.collect();
|
|
|
|
let completed_count = job
|
|
.results
|
|
.iter()
|
|
.filter(|r| r.status == TuningJobStatus::Completed)
|
|
.count();
|
|
|
|
let response = proto::StopBatchTuningJobResponse {
|
|
success: true,
|
|
message: format!(
|
|
"Batch job stopped. {} of {} models completed before stop.",
|
|
completed_count,
|
|
job.execution_order.len()
|
|
),
|
|
final_status: Self::convert_batch_status(&job.status) as i32,
|
|
completed_results,
|
|
};
|
|
|
|
Ok(Response::new(response))
|
|
}
|
|
|
|
// -- Job completion & model promotion (on-demand training pipeline) --------
|
|
|
|
#[instrument(skip_all)]
|
|
async fn report_job_completion(
|
|
&self,
|
|
request: Request<JobCompletionReport>,
|
|
) -> Result<Response<JobCompletionAck>, Status> {
|
|
let report = request.into_inner();
|
|
info!(
|
|
job_id = %report.job_id,
|
|
success = report.success,
|
|
s3_path = %report.s3_path,
|
|
"received job completion report"
|
|
);
|
|
|
|
// Parse the job UUID from the report
|
|
let job_uuid = uuid::Uuid::parse_str(&report.job_id).map_err(|e| {
|
|
Status::invalid_argument(format!("invalid job_id UUID: {}", e))
|
|
})?;
|
|
|
|
// ── Failed job path ──────────────────────────────────────────────
|
|
if !report.success {
|
|
// Best-effort DB status update; log warning if it fails
|
|
if let Some(ref spawner) = self.job_spawner {
|
|
if let Err(e) = spawner.update_job_status(job_uuid, "Failed").await {
|
|
error!(job_id = %report.job_id, error = %e, "failed to update job status to Failed");
|
|
}
|
|
}
|
|
return Ok(Response::new(JobCompletionAck {
|
|
accepted: true,
|
|
promotion_status: "failed".to_string(),
|
|
}));
|
|
}
|
|
|
|
// ── Successful job path ──────────────────────────────────────────
|
|
// 1. Update DB status to Completed
|
|
if let Some(ref spawner) = self.job_spawner {
|
|
if let Err(e) = spawner.update_job_status(job_uuid, "Completed").await {
|
|
error!(job_id = %report.job_id, error = %e, "failed to update job status to Completed");
|
|
}
|
|
}
|
|
|
|
// 2. Look up model_type and symbol from the child_jobs table
|
|
let (model_type, symbol) = match self.job_spawner {
|
|
Some(ref spawner) => {
|
|
match spawner.get_job_by_id(job_uuid).await {
|
|
Ok(Some(child_job)) => {
|
|
// Symbol is stored in config_json.asset
|
|
let sym = child_job
|
|
.config_json
|
|
.get("asset")
|
|
.and_then(|v| v.as_str())
|
|
.unwrap_or("UNKNOWN")
|
|
.to_string();
|
|
(child_job.model_type, sym)
|
|
}
|
|
Ok(None) => {
|
|
warn!(job_id = %report.job_id, "child job not found in DB, using defaults");
|
|
("UNKNOWN".to_string(), "UNKNOWN".to_string())
|
|
}
|
|
Err(e) => {
|
|
warn!(job_id = %report.job_id, error = %e, "failed to look up child job");
|
|
("UNKNOWN".to_string(), "UNKNOWN".to_string())
|
|
}
|
|
}
|
|
}
|
|
None => {
|
|
// No DB — cannot resolve model_type/symbol
|
|
debug!(job_id = %report.job_id, "no JobSpawner, cannot resolve model metadata");
|
|
("UNKNOWN".to_string(), "UNKNOWN".to_string())
|
|
}
|
|
};
|
|
|
|
// 3. Register with PromotionManager
|
|
let metrics: std::collections::HashMap<String, f64> = report.metrics.into_iter().collect();
|
|
let status = self
|
|
.promotion_manager
|
|
.register_completion(&report.job_id, &model_type, &symbol, &report.s3_path, metrics)
|
|
.await;
|
|
|
|
let promotion_status = match status {
|
|
crate::promotion_manager::PromotionStatus::PendingPromotion => "pending_promotion",
|
|
crate::promotion_manager::PromotionStatus::NoImprovement => "no_improvement",
|
|
crate::promotion_manager::PromotionStatus::Registered => "registered",
|
|
crate::promotion_manager::PromotionStatus::Error(ref msg) => {
|
|
warn!(job_id = %report.job_id, error = %msg, "promotion registration error");
|
|
"error"
|
|
}
|
|
};
|
|
|
|
info!(
|
|
job_id = %report.job_id,
|
|
model_type = %model_type,
|
|
symbol = %symbol,
|
|
promotion_status = %promotion_status,
|
|
"job completion processed"
|
|
);
|
|
|
|
Ok(Response::new(JobCompletionAck {
|
|
accepted: true,
|
|
promotion_status: promotion_status.to_string(),
|
|
}))
|
|
}
|
|
|
|
#[instrument(skip_all)]
|
|
async fn list_pending_promotions(
|
|
&self,
|
|
_request: Request<ListPendingPromotionsRequest>,
|
|
) -> Result<Response<ListPendingPromotionsResponse>, Status> {
|
|
let pending = self.promotion_manager.list_pending().await;
|
|
info!(count = pending.len(), "listing pending promotions");
|
|
|
|
let promotions = pending
|
|
.into_iter()
|
|
.map(|m| PendingPromotion {
|
|
model_id: m.model_id,
|
|
model_type: m.model_type,
|
|
symbol: m.symbol,
|
|
s3_path: m.s3_path,
|
|
new_metrics: m.new_metrics.into_iter().collect(),
|
|
current_metrics: m.current_metrics.into_iter().collect(),
|
|
trained_at: m.trained_at.timestamp(),
|
|
job_id: m.job_id,
|
|
})
|
|
.collect();
|
|
|
|
Ok(Response::new(ListPendingPromotionsResponse { promotions }))
|
|
}
|
|
|
|
#[instrument(skip_all)]
|
|
async fn approve_promotion(
|
|
&self,
|
|
request: Request<ApprovePromotionRequest>,
|
|
) -> Result<Response<ApprovePromotionResponse>, Status> {
|
|
let req = request.into_inner();
|
|
if req.model_id.is_empty() {
|
|
return Err(Status::invalid_argument("model_id is required"));
|
|
}
|
|
info!(model_id = %req.model_id, "approve_promotion request received");
|
|
|
|
match self.promotion_manager.approve(&req.model_id).await {
|
|
Ok(()) => {
|
|
info!(model_id = %req.model_id, "model promotion approved and activated");
|
|
Ok(Response::new(ApprovePromotionResponse {
|
|
success: true,
|
|
message: format!("Model {} promoted to active", req.model_id),
|
|
}))
|
|
}
|
|
Err(e) => {
|
|
warn!(model_id = %req.model_id, error = %e, "approve_promotion failed");
|
|
Ok(Response::new(ApprovePromotionResponse {
|
|
success: false,
|
|
message: format!("Failed to approve promotion: {}", e),
|
|
}))
|
|
}
|
|
}
|
|
}
|
|
|
|
#[instrument(skip_all)]
|
|
async fn reject_promotion(
|
|
&self,
|
|
request: Request<RejectPromotionRequest>,
|
|
) -> Result<Response<RejectPromotionResponse>, Status> {
|
|
let req = request.into_inner();
|
|
if req.model_id.is_empty() {
|
|
return Err(Status::invalid_argument("model_id is required"));
|
|
}
|
|
let reason = if req.reason.is_empty() {
|
|
"no reason provided".to_string()
|
|
} else {
|
|
req.reason
|
|
};
|
|
info!(model_id = %req.model_id, reason = %reason, "reject_promotion request received");
|
|
|
|
match self.promotion_manager.reject(&req.model_id, &reason).await {
|
|
Ok(()) => {
|
|
info!(model_id = %req.model_id, reason = %reason, "model promotion rejected");
|
|
Ok(Response::new(RejectPromotionResponse {
|
|
success: true,
|
|
message: format!("Promotion of model {} rejected: {}", req.model_id, reason),
|
|
}))
|
|
}
|
|
Err(e) => {
|
|
warn!(model_id = %req.model_id, error = %e, "reject_promotion failed");
|
|
Ok(Response::new(RejectPromotionResponse {
|
|
success: false,
|
|
message: format!("Failed to reject promotion: {}", e),
|
|
}))
|
|
}
|
|
}
|
|
}
|
|
|
|
// -- Dashboard-facing model approval / rejection ----------------------------
|
|
|
|
#[instrument(skip_all)]
|
|
async fn approve_model(
|
|
&self,
|
|
request: Request<ApproveModelRequest>,
|
|
) -> Result<Response<ApproveModelResponse>, Status> {
|
|
let req = request.into_inner();
|
|
if req.model_id.is_empty() {
|
|
return Err(Status::invalid_argument("model_id is required"));
|
|
}
|
|
info!(model_id = %req.model_id, promoted_to = %req.promoted_to, "approve_model request received");
|
|
|
|
match self.promotion_manager.approve(&req.model_id).await {
|
|
Ok(()) => {
|
|
let target = if req.promoted_to.is_empty() {
|
|
"active".to_string()
|
|
} else {
|
|
req.promoted_to
|
|
};
|
|
info!(model_id = %req.model_id, promoted_to = %target, "model approved");
|
|
Ok(Response::new(ApproveModelResponse {
|
|
success: true,
|
|
message: format!("Model '{}' approved for '{}'", req.model_id, target),
|
|
}))
|
|
}
|
|
Err(e) => {
|
|
warn!(model_id = %req.model_id, error = %e, "approve_model failed");
|
|
Ok(Response::new(ApproveModelResponse {
|
|
success: false,
|
|
message: format!("Failed to approve model: {}", e),
|
|
}))
|
|
}
|
|
}
|
|
}
|
|
|
|
#[instrument(skip_all)]
|
|
async fn reject_model(
|
|
&self,
|
|
request: Request<RejectModelRequest>,
|
|
) -> Result<Response<RejectModelResponse>, Status> {
|
|
let req = request.into_inner();
|
|
if req.model_id.is_empty() {
|
|
return Err(Status::invalid_argument("model_id is required"));
|
|
}
|
|
let reason = if req.reason.is_empty() {
|
|
"no reason provided".to_string()
|
|
} else {
|
|
req.reason
|
|
};
|
|
info!(model_id = %req.model_id, reason = %reason, "reject_model request received");
|
|
|
|
match self.promotion_manager.reject(&req.model_id, &reason).await {
|
|
Ok(()) => {
|
|
info!(model_id = %req.model_id, reason = %reason, "model rejected");
|
|
Ok(Response::new(RejectModelResponse {
|
|
success: true,
|
|
message: format!("Model '{}' rejected: {}", req.model_id, reason),
|
|
}))
|
|
}
|
|
Err(e) => {
|
|
warn!(model_id = %req.model_id, error = %e, "reject_model failed");
|
|
Ok(Response::new(RejectModelResponse {
|
|
success: false,
|
|
message: format!("Failed to reject model: {}", e),
|
|
}))
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
impl MLTrainingServiceImpl {
|
|
/// Convert hyperparameters from flat map to ProductionTrainingConfig
|
|
fn convert_hyperparameters_from_map(
|
|
&self,
|
|
model_type: &str,
|
|
params: &HashMap<String, f32>,
|
|
) -> Result<ProductionTrainingConfig> {
|
|
let mut config = ProductionTrainingConfig::default();
|
|
|
|
// Extract common parameters
|
|
let learning_rate = params.get("learning_rate").copied().unwrap_or(0.001) as f64;
|
|
let batch_size = params.get("batch_size").copied().unwrap_or(64.0) as usize;
|
|
let epochs = params.get("epochs").copied().unwrap_or(100.0) as usize;
|
|
|
|
// Update training parameters
|
|
config.training_params.learning_rate = learning_rate;
|
|
config.training_params.batch_size = batch_size;
|
|
config.training_params.max_epochs = epochs;
|
|
|
|
// Model-specific parameter extraction
|
|
match model_type {
|
|
"TLOB" => {
|
|
let hidden_dim = params.get("hidden_dim").copied().unwrap_or(256.0) as usize;
|
|
let dropout_rate = params.get("dropout_rate").copied().unwrap_or(0.1) as f64;
|
|
|
|
config.model_config.input_dim = 20;
|
|
config.model_config.hidden_dims = vec![hidden_dim, hidden_dim / 2];
|
|
config.model_config.dropout_rate = dropout_rate;
|
|
config.model_config.output_dim = 1;
|
|
},
|
|
"MAMBA_2" => {
|
|
let state_dim = params.get("state_dim").copied().unwrap_or(128.0) as usize;
|
|
let hidden_dim = params.get("hidden_dim").copied().unwrap_or(512.0) as usize;
|
|
|
|
config.model_config.input_dim = state_dim;
|
|
config.model_config.hidden_dims = vec![hidden_dim];
|
|
config.model_config.output_dim = 1;
|
|
},
|
|
"DQN" => {
|
|
let gamma = params.get("gamma").copied().unwrap_or(0.99);
|
|
// DQN uses default config with gamma stored in tags/metadata
|
|
config.training_params.l2_regularization = 1e-5;
|
|
// Store gamma for later use
|
|
let _ = gamma; // Placeholder for future DQN-specific config
|
|
},
|
|
"PPO" => {
|
|
let clip_ratio = params.get("clip_ratio").copied().unwrap_or(0.2);
|
|
// PPO uses default config with clip_ratio stored in metadata
|
|
let _ = clip_ratio; // Placeholder for future PPO-specific config
|
|
},
|
|
"TFT" => {
|
|
let hidden_dim = params.get("hidden_dim").copied().unwrap_or(240.0) as usize;
|
|
let dropout_rate = params.get("dropout_rate").copied().unwrap_or(0.3) as f64;
|
|
|
|
config.model_config.hidden_dims = vec![hidden_dim];
|
|
config.model_config.dropout_rate = dropout_rate;
|
|
config.model_config.output_dim = 1;
|
|
},
|
|
"LIQUID" => {
|
|
let num_neurons = params.get("num_neurons").copied().unwrap_or(128.0) as usize;
|
|
config.model_config.hidden_dims = vec![num_neurons];
|
|
config.model_config.output_dim = 1;
|
|
},
|
|
_ => {
|
|
return Err(anyhow::anyhow!("Unknown model type: {}", model_type));
|
|
},
|
|
}
|
|
|
|
Ok(config)
|
|
}
|
|
|
|
/// Calculate Sharpe ratio from validation data
|
|
///
|
|
/// This performs a mini-backtest using the trained model to generate predictions
|
|
/// on the validation set and calculates financial performance metrics.
|
|
async fn calculate_sharpe_ratio_from_validation(
|
|
&self,
|
|
_validation_data: &[(ml::training_pipeline::FinancialFeatures, Vec<f64>)],
|
|
result: &ml::training_pipeline::TrainingResult,
|
|
) -> f32 {
|
|
// Extract returns from validation predictions
|
|
// In a real implementation, this would:
|
|
// 1. Use the trained model to predict on validation data
|
|
// 2. Simulate trades based on predictions
|
|
// 3. Calculate returns from simulated trades
|
|
// 4. Compute Sharpe ratio from returns
|
|
|
|
// For now, use financial metrics from training if available
|
|
if let Some(last_metrics) = result.metrics_history.last() {
|
|
// Use the financial performance metrics from training
|
|
let sharpe = last_metrics.financial_metrics.sharpe_ratio;
|
|
|
|
// Validate and clamp to reasonable range
|
|
if sharpe.is_finite() {
|
|
return sharpe.max(-5.0).min(10.0) as f32;
|
|
}
|
|
}
|
|
|
|
// Fallback: Estimate Sharpe ratio from validation loss
|
|
// Lower validation loss suggests better predictions -> higher Sharpe ratio
|
|
// This is a rough approximation: Sharpe ≈ -log(val_loss)
|
|
let val_loss = result.final_val_loss;
|
|
if val_loss > 0.0 && val_loss.is_finite() {
|
|
let estimated_sharpe = -(val_loss.ln()) / 2.0; // Normalized to ~0-3 range
|
|
estimated_sharpe.max(-2.0).min(5.0) as f32
|
|
} else {
|
|
0.0
|
|
}
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
#[allow(clippy::unwrap_used, clippy::expect_used)]
|
|
mod tests {
|
|
use super::*;
|
|
use crate::orchestrator::TrainingJob;
|
|
|
|
#[test]
|
|
fn test_status_conversion() {
|
|
assert_eq!(
|
|
MLTrainingServiceImpl::convert_job_status(&JobStatus::Pending),
|
|
ProtoTrainingStatus::Pending
|
|
);
|
|
|
|
assert_eq!(
|
|
MLTrainingServiceImpl::convert_job_status(&JobStatus::Running),
|
|
ProtoTrainingStatus::Running
|
|
);
|
|
|
|
assert_eq!(
|
|
MLTrainingServiceImpl::convert_job_status(&JobStatus::Completed),
|
|
ProtoTrainingStatus::Completed
|
|
);
|
|
|
|
assert_eq!(
|
|
MLTrainingServiceImpl::convert_job_status(&JobStatus::Failed),
|
|
ProtoTrainingStatus::Failed
|
|
);
|
|
|
|
assert_eq!(
|
|
MLTrainingServiceImpl::convert_job_status(&JobStatus::Stopped),
|
|
ProtoTrainingStatus::Stopped
|
|
);
|
|
|
|
assert_eq!(
|
|
MLTrainingServiceImpl::convert_job_status(&JobStatus::Paused),
|
|
ProtoTrainingStatus::Paused
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_training_job_creation() {
|
|
let config = ProductionTrainingConfig::default();
|
|
let mut tags = HashMap::new();
|
|
tags.insert("env".to_string(), "test".to_string());
|
|
|
|
let job = TrainingJob::new(
|
|
"TLOB".to_string(),
|
|
config,
|
|
"Test job".to_string(),
|
|
tags.clone(),
|
|
);
|
|
|
|
assert_eq!(job.model_type, "TLOB");
|
|
assert_eq!(job.status, JobStatus::Pending);
|
|
assert_eq!(job.description, "Test job");
|
|
assert_eq!(job.tags.get("env"), Some(&"test".to_string()));
|
|
assert_eq!(job.progress_percentage, 0.0);
|
|
assert_eq!(job.current_epoch, 0);
|
|
assert!(job.started_at.is_none());
|
|
assert!(job.completed_at.is_none());
|
|
}
|
|
|
|
#[test]
|
|
fn test_job_id_uniqueness() {
|
|
let config = ProductionTrainingConfig::default();
|
|
let mut job_ids = vec![];
|
|
|
|
for i in 0..100 {
|
|
let job = TrainingJob::new(
|
|
"TLOB".to_string(),
|
|
config.clone(),
|
|
format!("Job {}", i),
|
|
HashMap::new(),
|
|
);
|
|
job_ids.push(job.id);
|
|
}
|
|
|
|
// Verify all IDs are unique
|
|
let unique_ids: std::collections::HashSet<_> = job_ids.iter().collect();
|
|
assert_eq!(unique_ids.len(), 100);
|
|
}
|
|
|
|
#[test]
|
|
fn test_hyperparameter_protobuf_structure() {
|
|
let tlob_params = Hyperparameters {
|
|
model_params: Some(proto::hyperparameters::ModelParams::TlobParams(
|
|
TlobParams {
|
|
epochs: 100,
|
|
learning_rate: 0.001,
|
|
batch_size: 64,
|
|
sequence_length: 50,
|
|
hidden_dim: 256,
|
|
num_heads: 8,
|
|
num_layers: 6,
|
|
dropout_rate: 0.1,
|
|
use_positional_encoding: true,
|
|
},
|
|
)),
|
|
};
|
|
|
|
assert!(matches!(
|
|
tlob_params.model_params,
|
|
Some(proto::hyperparameters::ModelParams::TlobParams(_))
|
|
));
|
|
|
|
if let Some(proto::hyperparameters::ModelParams::TlobParams(params)) =
|
|
tlob_params.model_params
|
|
{
|
|
assert_eq!(params.epochs, 100);
|
|
assert_eq!(params.batch_size, 64);
|
|
assert_eq!(params.num_heads, 8);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_mamba_hyperparameters() {
|
|
let mamba_params = Hyperparameters {
|
|
model_params: Some(proto::hyperparameters::ModelParams::MambaParams(
|
|
MambaParams {
|
|
epochs: 150,
|
|
learning_rate: 0.0005,
|
|
batch_size: 32,
|
|
state_dim: 128,
|
|
hidden_dim: 512,
|
|
num_layers: 8,
|
|
dt_min: 0.001,
|
|
dt_max: 0.1,
|
|
use_cuda_kernels: true,
|
|
},
|
|
)),
|
|
};
|
|
|
|
if let Some(proto::hyperparameters::ModelParams::MambaParams(params)) =
|
|
mamba_params.model_params
|
|
{
|
|
assert_eq!(params.state_dim, 128);
|
|
assert_eq!(params.hidden_dim, 512);
|
|
assert!(params.use_cuda_kernels);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_dqn_hyperparameters() {
|
|
let dqn_params = Hyperparameters {
|
|
model_params: Some(proto::hyperparameters::ModelParams::DqnParams(DqnParams {
|
|
epochs: 200,
|
|
learning_rate: 0.0001,
|
|
batch_size: 128,
|
|
replay_buffer_size: 100000,
|
|
epsilon_start: 1.0,
|
|
epsilon_end: 0.01,
|
|
epsilon_decay_steps: 50000,
|
|
gamma: 0.99,
|
|
target_update_frequency: 1000,
|
|
use_double_dqn: true,
|
|
use_dueling: true,
|
|
use_prioritized_replay: true,
|
|
})),
|
|
};
|
|
|
|
if let Some(proto::hyperparameters::ModelParams::DqnParams(params)) =
|
|
dqn_params.model_params
|
|
{
|
|
assert_eq!(params.replay_buffer_size, 100000);
|
|
assert!(params.use_prioritized_replay);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_ppo_hyperparameters() {
|
|
let ppo_params = Hyperparameters {
|
|
model_params: Some(proto::hyperparameters::ModelParams::PpoParams(PpoParams {
|
|
epochs: 100,
|
|
learning_rate: 0.0003,
|
|
batch_size: 64,
|
|
clip_ratio: 0.2,
|
|
value_loss_coef: 0.5,
|
|
entropy_coef: 0.01,
|
|
rollout_steps: 2048,
|
|
minibatch_size: 64,
|
|
gae_lambda: 0.95,
|
|
})),
|
|
};
|
|
|
|
if let Some(proto::hyperparameters::ModelParams::PpoParams(params)) =
|
|
ppo_params.model_params
|
|
{
|
|
assert_eq!(params.clip_ratio, 0.2);
|
|
assert_eq!(params.rollout_steps, 2048);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_liquid_hyperparameters() {
|
|
let liquid_params = Hyperparameters {
|
|
model_params: Some(proto::hyperparameters::ModelParams::LiquidParams(
|
|
LiquidParams {
|
|
epochs: 80,
|
|
learning_rate: 0.002,
|
|
batch_size: 48,
|
|
num_neurons: 128,
|
|
tau: 0.1,
|
|
sigma: 0.5,
|
|
use_adaptive_tau: true,
|
|
},
|
|
)),
|
|
};
|
|
|
|
if let Some(proto::hyperparameters::ModelParams::LiquidParams(params)) =
|
|
liquid_params.model_params
|
|
{
|
|
assert_eq!(params.num_neurons, 128);
|
|
assert!(params.use_adaptive_tau);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_tft_hyperparameters() {
|
|
let tft_params = Hyperparameters {
|
|
model_params: Some(proto::hyperparameters::ModelParams::TftParams(TftParams {
|
|
epochs: 120,
|
|
learning_rate: 0.001,
|
|
batch_size: 32,
|
|
hidden_dim: 240,
|
|
num_heads: 4,
|
|
num_layers: 3,
|
|
lookback_window: 168,
|
|
forecast_horizon: 24,
|
|
dropout_rate: 0.3,
|
|
})),
|
|
};
|
|
|
|
if let Some(proto::hyperparameters::ModelParams::TftParams(params)) =
|
|
tft_params.model_params
|
|
{
|
|
assert_eq!(params.lookback_window, 168);
|
|
assert_eq!(params.forecast_horizon, 24);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_model_types() {
|
|
let models = vec!["TLOB", "MAMBA_2", "DQN", "PPO", "LIQUID", "TFT"];
|
|
|
|
// Verify unique count
|
|
let unique_count = models
|
|
.iter()
|
|
.collect::<std::collections::HashSet<_>>()
|
|
.len();
|
|
assert_eq!(unique_count, 6);
|
|
|
|
// Verify naming conventions
|
|
for model in models {
|
|
assert!(!model.is_empty());
|
|
assert!(model
|
|
.chars()
|
|
.all(|c| c.is_uppercase() || c.is_numeric() || c == '_'));
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_job_progress_updates() {
|
|
let config = ProductionTrainingConfig::default();
|
|
let mut job = TrainingJob::new(
|
|
"DQN".to_string(),
|
|
config,
|
|
"Progress test".to_string(),
|
|
HashMap::new(),
|
|
);
|
|
|
|
// Initial state
|
|
assert_eq!(job.progress_percentage, 0.0);
|
|
assert_eq!(job.current_epoch, 0);
|
|
|
|
// Simulate progress
|
|
job.progress_percentage = 50.0;
|
|
job.current_epoch = 50;
|
|
job.total_epochs = 100;
|
|
|
|
assert_eq!(job.progress_percentage, 50.0);
|
|
assert_eq!(job.current_epoch, 50);
|
|
assert_eq!(job.total_epochs, 100);
|
|
}
|
|
|
|
#[test]
|
|
fn test_job_metrics_tracking() {
|
|
let config = ProductionTrainingConfig::default();
|
|
let mut job = TrainingJob::new(
|
|
"TLOB".to_string(),
|
|
config,
|
|
"Metrics test".to_string(),
|
|
HashMap::new(),
|
|
);
|
|
|
|
// Add metrics
|
|
job.metrics.insert("train_loss".to_string(), 0.5);
|
|
job.metrics.insert("val_loss".to_string(), 0.6);
|
|
job.metrics.insert("accuracy".to_string(), 0.85);
|
|
|
|
assert_eq!(job.metrics.get("train_loss"), Some(&0.5));
|
|
assert_eq!(job.metrics.get("val_loss"), Some(&0.6));
|
|
assert_eq!(job.metrics.get("accuracy"), Some(&0.85));
|
|
assert_eq!(job.metrics.len(), 3);
|
|
}
|
|
|
|
#[test]
|
|
fn test_training_config_defaults() {
|
|
let config = ProductionTrainingConfig::default();
|
|
|
|
// Verify default training parameters exist
|
|
assert!(config.training_params.learning_rate > 0.0);
|
|
assert!(config.training_params.batch_size > 0);
|
|
assert!(config.training_params.max_epochs > 0);
|
|
assert!(config.training_params.validation_split >= 0.0);
|
|
assert!(config.training_params.validation_split <= 1.0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_service_impl_has_job_spawner() {
|
|
let _: fn(&MLTrainingServiceImpl) -> bool = |s| s.job_spawner.is_some();
|
|
}
|
|
|
|
#[test]
|
|
fn test_start_training_model_binary_mapping() {
|
|
use crate::k8s_dispatcher::training_binary_for_model;
|
|
assert_eq!(training_binary_for_model("tft"), "train_baseline_supervised");
|
|
assert_eq!(training_binary_for_model("dqn"), "train_baseline_rl");
|
|
assert_eq!(training_binary_for_model("ppo"), "train_baseline_rl");
|
|
assert_eq!(
|
|
training_binary_for_model("mamba2"),
|
|
"train_baseline_supervised"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_promotion_status_string_mapping() {
|
|
// Verify the string representations match what gRPC clients expect
|
|
use crate::promotion_manager::PromotionStatus;
|
|
|
|
let cases: Vec<(PromotionStatus, &str)> = vec![
|
|
(PromotionStatus::PendingPromotion, "pending_promotion"),
|
|
(PromotionStatus::NoImprovement, "no_improvement"),
|
|
(PromotionStatus::Registered, "registered"),
|
|
(PromotionStatus::Error("test".to_string()), "error"),
|
|
];
|
|
|
|
for (status, expected) in cases {
|
|
let result = match status {
|
|
PromotionStatus::PendingPromotion => "pending_promotion",
|
|
PromotionStatus::NoImprovement => "no_improvement",
|
|
PromotionStatus::Registered => "registered",
|
|
PromotionStatus::Error(_) => "error",
|
|
};
|
|
assert_eq!(result, expected);
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_report_job_completion_registers_with_promotion_manager() {
|
|
// Integration test: verify that a successful completion report
|
|
// flows through to the PromotionManager and returns the correct status
|
|
use crate::promotion_manager::PromotionManager;
|
|
|
|
let pm = PromotionManager::new();
|
|
let metrics = HashMap::from([
|
|
("best_val_loss".to_string(), 0.05),
|
|
("sharpe_ratio".to_string(), 1.5),
|
|
]);
|
|
|
|
// First registration for a type+symbol returns Registered
|
|
let status = pm
|
|
.register_completion("job-1", "DQN", "ES.FUT", "s3://b/m.bin", metrics)
|
|
.await;
|
|
assert_eq!(
|
|
status,
|
|
crate::promotion_manager::PromotionStatus::Registered
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_list_pending_promotions_maps_to_proto() {
|
|
// Verify that PendingModel fields map correctly to proto PendingPromotion
|
|
use crate::promotion_manager::PendingModel;
|
|
use chrono::Utc;
|
|
|
|
let now = Utc::now();
|
|
let model = PendingModel {
|
|
model_id: "m-123".to_string(),
|
|
model_type: "DQN".to_string(),
|
|
symbol: "ES.FUT".to_string(),
|
|
s3_path: "s3://bucket/better.bin".to_string(),
|
|
new_metrics: HashMap::from([
|
|
("best_val_loss".to_string(), 0.03),
|
|
("sharpe_ratio".to_string(), 2.5),
|
|
]),
|
|
current_metrics: HashMap::from([
|
|
("best_val_loss".to_string(), 0.10),
|
|
("sharpe_ratio".to_string(), 1.0),
|
|
]),
|
|
trained_at: now,
|
|
job_id: "job-99".to_string(),
|
|
};
|
|
|
|
// Map to proto the same way the handler does
|
|
let proto = super::proto::PendingPromotion {
|
|
model_id: model.model_id.clone(),
|
|
model_type: model.model_type.clone(),
|
|
symbol: model.symbol.clone(),
|
|
s3_path: model.s3_path.clone(),
|
|
new_metrics: model.new_metrics.clone().into_iter().collect(),
|
|
current_metrics: model.current_metrics.clone().into_iter().collect(),
|
|
trained_at: model.trained_at.timestamp(),
|
|
job_id: model.job_id.clone(),
|
|
};
|
|
|
|
assert_eq!(proto.model_id, "m-123");
|
|
assert_eq!(proto.model_type, "DQN");
|
|
assert_eq!(proto.symbol, "ES.FUT");
|
|
assert_eq!(proto.s3_path, "s3://bucket/better.bin");
|
|
assert_eq!(proto.job_id, "job-99");
|
|
assert_eq!(proto.trained_at, now.timestamp());
|
|
assert_eq!(proto.new_metrics.len(), 2);
|
|
assert_eq!(proto.current_metrics.len(), 2);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_approve_promotion_activates_pending_model() {
|
|
use crate::promotion_manager::{ActiveModel, PromotionManager};
|
|
use chrono::Utc;
|
|
|
|
let pm = Arc::new(PromotionManager::new());
|
|
|
|
// Seed an active model so register_completion returns PendingPromotion
|
|
{
|
|
let active = pm.active_models_for_test();
|
|
let mut w = active.write().await;
|
|
w.insert(
|
|
("DQN".to_string(), "ES.FUT".to_string()),
|
|
ActiveModel {
|
|
model_id: "old-model".to_string(),
|
|
s3_path: "s3://bucket/old.bin".to_string(),
|
|
metrics: HashMap::from([
|
|
("best_val_loss".to_string(), 0.10),
|
|
("sharpe_ratio".to_string(), 1.0),
|
|
]),
|
|
promoted_at: Utc::now(),
|
|
},
|
|
);
|
|
}
|
|
|
|
// Register a better model
|
|
let new_metrics = HashMap::from([
|
|
("best_val_loss".to_string(), 0.03),
|
|
("sharpe_ratio".to_string(), 2.5),
|
|
]);
|
|
let status = pm
|
|
.register_completion("job-a", "DQN", "ES.FUT", "s3://bucket/better.bin", new_metrics)
|
|
.await;
|
|
assert_eq!(
|
|
status,
|
|
crate::promotion_manager::PromotionStatus::PendingPromotion
|
|
);
|
|
|
|
let pending = pm.list_pending().await;
|
|
assert_eq!(pending.len(), 1);
|
|
let model_id = pending[0].model_id.clone();
|
|
|
|
// Approve it
|
|
pm.approve(&model_id).await.unwrap();
|
|
|
|
// Pending list should be empty
|
|
assert!(pm.list_pending().await.is_empty());
|
|
|
|
// Active model should be updated
|
|
let active = pm.active_models_for_test();
|
|
let r = active.read().await;
|
|
let am = r
|
|
.get(&("DQN".to_string(), "ES.FUT".to_string()))
|
|
.unwrap();
|
|
assert_eq!(am.model_id, model_id);
|
|
assert_eq!(am.s3_path, "s3://bucket/better.bin");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_reject_promotion_removes_pending_model() {
|
|
use crate::promotion_manager::{ActiveModel, PromotionManager};
|
|
use chrono::Utc;
|
|
|
|
let pm = Arc::new(PromotionManager::new());
|
|
|
|
// Seed an active model
|
|
{
|
|
let active = pm.active_models_for_test();
|
|
let mut w = active.write().await;
|
|
w.insert(
|
|
("PPO".to_string(), "NQ.FUT".to_string()),
|
|
ActiveModel {
|
|
model_id: "old-ppo".to_string(),
|
|
s3_path: "s3://bucket/old-ppo.bin".to_string(),
|
|
metrics: HashMap::from([
|
|
("best_val_loss".to_string(), 0.10),
|
|
("sharpe_ratio".to_string(), 1.0),
|
|
]),
|
|
promoted_at: Utc::now(),
|
|
},
|
|
);
|
|
}
|
|
|
|
// Register a better model
|
|
let new_metrics = HashMap::from([
|
|
("best_val_loss".to_string(), 0.04),
|
|
("sharpe_ratio".to_string(), 2.0),
|
|
]);
|
|
let status = pm
|
|
.register_completion("job-r", "PPO", "NQ.FUT", "s3://bucket/cand.bin", new_metrics)
|
|
.await;
|
|
assert_eq!(
|
|
status,
|
|
crate::promotion_manager::PromotionStatus::PendingPromotion
|
|
);
|
|
|
|
let pending = pm.list_pending().await;
|
|
assert_eq!(pending.len(), 1);
|
|
let model_id = pending[0].model_id.clone();
|
|
|
|
// Reject it
|
|
pm.reject(&model_id, "needs more walk-forward windows")
|
|
.await
|
|
.unwrap();
|
|
|
|
// Pending list should be empty
|
|
assert!(pm.list_pending().await.is_empty());
|
|
|
|
// Active model should NOT have changed
|
|
let active = pm.active_models_for_test();
|
|
let r = active.read().await;
|
|
let am = r
|
|
.get(&("PPO".to_string(), "NQ.FUT".to_string()))
|
|
.unwrap();
|
|
assert_eq!(am.model_id, "old-ppo");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_approve_nonexistent_model_returns_error() {
|
|
use crate::promotion_manager::PromotionManager;
|
|
|
|
let pm = PromotionManager::new();
|
|
let result = pm.approve("nonexistent-id").await;
|
|
assert!(result.is_err());
|
|
assert!(
|
|
result
|
|
.unwrap_err()
|
|
.to_string()
|
|
.contains("pending model not found")
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_reject_nonexistent_model_returns_error() {
|
|
use crate::promotion_manager::PromotionManager;
|
|
|
|
let pm = PromotionManager::new();
|
|
let result = pm.reject("nonexistent-id", "bad model").await;
|
|
assert!(result.is_err());
|
|
assert!(
|
|
result
|
|
.unwrap_err()
|
|
.to_string()
|
|
.contains("pending model not found")
|
|
);
|
|
}
|
|
}
|