From bb116bad8913b21e7f99eb22c65926a73a2b9cf2 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Mon, 23 Feb 2026 00:26:13 +0100 Subject: [PATCH] fix(trading_service): safetensors DQN loading, prediction shutdown, emergency order query - Add load_from_safetensors() to DQNAgent for weight loading via VarMap - Update RealDQNModel::from_checkpoint to try safetensors first, fall back to JSON - Replace std::mem::forget(prediction_shutdown_tx) with proper Vec-based storage that sends shutdown signal and drops senders during graceful shutdown - Wire order_manager.get_open_orders() into emergency_stop response so callers see which orders were active when the kill switch engaged Co-Authored-By: Claude Opus 4.6 --- ml/src/dqn/agent.rs | 33 +++++++++++ services/trading_service/src/main.rs | 16 ++++- .../src/services/enhanced_ml.rs | 59 ++++++++++++++++--- services/trading_service/src/services/risk.rs | 22 ++++++- 4 files changed, 118 insertions(+), 12 deletions(-) diff --git a/ml/src/dqn/agent.rs b/ml/src/dqn/agent.rs index 1e3e2b2c4..52b2d1c5a 100644 --- a/ml/src/dqn/agent.rs +++ b/ml/src/dqn/agent.rs @@ -714,6 +714,39 @@ impl DQNAgent { Ok(()) } + /// Load model weights from a safetensors file. + /// + /// This loads the Q-network weights from a `.safetensors` checkpoint, + /// then copies them to the target network via soft update. + /// Unlike `load_checkpoint` (JSON), this only restores network weights -- + /// config, metrics, and epsilon are left unchanged. + pub fn load_from_safetensors(&mut self, path: &std::path::Path) -> Result<(), MLError> { + let path_str = path.to_string_lossy(); + let safetensors_path = if !path_str.ends_with(".safetensors") { + format!("{}.safetensors", path_str) + } else { + path_str.to_string() + }; + + if !std::path::Path::new(&safetensors_path).exists() { + return Err(MLError::CheckpointError(format!( + "Safetensors checkpoint not found: {}", + safetensors_path + ))); + } + + let mut vars_clone = self.q_network.vars().clone(); + vars_clone.load(&safetensors_path).map_err(|e| { + MLError::CheckpointError(format!("Failed to load safetensors via VarMap: {}", e)) + })?; + + // Propagate loaded weights to target network + self.update_target_network()?; + + tracing::info!("DQNAgent weights loaded from safetensors: {}", safetensors_path); + Ok(()) + } + /// Get current epsilon value pub fn get_epsilon(&self) -> f64 { self.q_network.get_epsilon() diff --git a/services/trading_service/src/main.rs b/services/trading_service/src/main.rs index 54174416d..34f4fba7f 100644 --- a/services/trading_service/src/main.rs +++ b/services/trading_service/src/main.rs @@ -464,6 +464,9 @@ async fn main() -> Result<()> { PredictionGenerationLoop, PredictionLoopConfig, }; + // Shutdown senders for background loops -- dropped at end of main to signal shutdown. + let mut prediction_shutdown_handles: Vec> = Vec::new(); + if let Some(ensemble_coordinator) = service_state.ensemble_coordinator() { let prediction_config = PredictionLoopConfig::from_env(); let prediction_loop = PredictionGenerationLoop::new( @@ -488,9 +491,9 @@ async fn main() -> Result<()> { info!("✅ Background ML prediction generation loop initialized"); - // Store shutdown sender for graceful shutdown - // TODO: Add to service state for proper cleanup - std::mem::forget(prediction_shutdown_tx); // Kept alive for service lifetime + // Store shutdown sender so it is dropped during graceful shutdown, + // which closes the channel and signals the prediction loop to stop. + prediction_shutdown_handles.push(prediction_shutdown_tx); } else { warn!("⚠️ Ensemble coordinator not available - prediction generation loop disabled"); } @@ -700,6 +703,13 @@ async fn main() -> Result<()> { } } + // Signal prediction loop(s) to stop, then drop senders + for tx in &prediction_shutdown_handles { + let _ = tx.send(()); + } + drop(prediction_shutdown_handles); + info!("Background prediction loops signalled for shutdown"); + // Cleanup kill switch monitoring on shutdown info!("Stopping kill switch monitoring..."); if let Err(e) = kill_switch_system.stop_monitoring().await { diff --git a/services/trading_service/src/services/enhanced_ml.rs b/services/trading_service/src/services/enhanced_ml.rs index ec5cf2bb0..61ad33d36 100644 --- a/services/trading_service/src/services/enhanced_ml.rs +++ b/services/trading_service/src/services/enhanced_ml.rs @@ -1246,9 +1246,8 @@ impl MlService for EnhancedMLServiceImpl { /// Real DQN Model Wrapper that loads from safetensors checkpoints /// -/// This wrapper integrates the ml crate's DQN implementation with the MLModel trait -/// NOTE: Current implementation uses DQNAgent with JSON checkpoint format, not safetensors. -/// TODO: Implement safetensors loading when DQNAgent supports it. +/// This wrapper integrates the ml crate's DQN implementation with the MLModel trait. +/// Checkpoint loading tries safetensors first, then falls back to JSON format. #[derive(Debug)] struct RealDQNModel { model_id: String, @@ -1257,7 +1256,12 @@ struct RealDQNModel { } impl RealDQNModel { - /// Create new DQN model from checkpoint + /// Create new DQN model from checkpoint. + /// + /// Loading strategy: + /// 1. If `checkpoint_path` ends with `.safetensors`, load via safetensors directly. + /// 2. Otherwise, check if a `.safetensors` sibling exists and prefer it. + /// 3. Fall back to JSON checkpoint format (`DQNAgent::load_checkpoint`). pub fn from_checkpoint( model_id: String, checkpoint_path: &std::path::Path, @@ -1287,10 +1291,49 @@ impl RealDQNModel { let mut agent = DQNAgent::new(config) .map_err(|e| ml::MLError::ModelError(format!("Failed to create DQN agent: {}", e)))?; - // Load checkpoint weights (JSON format for now) - agent.load_checkpoint(checkpoint_path).map_err(|e| { - ml::MLError::ModelError(format!("Failed to load DQN checkpoint: {}", e)) - })?; + // Try safetensors first, fall back to JSON checkpoint + let safetensors_path = if checkpoint_path + .extension() + .is_some_and(|ext| ext == "safetensors") + { + Some(checkpoint_path.to_path_buf()) + } else { + let candidate = checkpoint_path.with_extension("safetensors"); + if candidate.exists() { + Some(candidate) + } else { + None + } + }; + + if let Some(st_path) = safetensors_path { + match agent.load_from_safetensors(&st_path) { + Ok(()) => { + tracing::info!( + "DQN model '{}' loaded from safetensors: {}", + model_id, + st_path.display() + ); + }, + Err(e) => { + tracing::warn!( + "Safetensors load failed for '{}', falling back to JSON: {}", + model_id, + e + ); + agent.load_checkpoint(checkpoint_path).map_err(|e2| { + ml::MLError::ModelError(format!( + "Failed to load DQN checkpoint (safetensors failed: {}, JSON failed: {})", + e, e2 + )) + })?; + }, + } + } else { + agent.load_checkpoint(checkpoint_path).map_err(|e| { + ml::MLError::ModelError(format!("Failed to load DQN checkpoint: {}", e)) + })?; + } Ok(Self { model_id, diff --git a/services/trading_service/src/services/risk.rs b/services/trading_service/src/services/risk.rs index be4aca72b..739e3d951 100644 --- a/services/trading_service/src/services/risk.rs +++ b/services/trading_service/src/services/risk.rs @@ -659,11 +659,31 @@ impl RiskService for RiskServiceImpl { info!("Emergency stop complete, global kill switch engaged"); + // Query open orders that were active at the time of emergency shutdown + let affected_orders = match self.state.order_manager.read().await.get_open_orders().await { + orders if orders.is_empty() => { + info!("No open orders at time of emergency stop"); + vec![] + }, + orders => { + let order_ids: Vec = orders + .iter() + .map(|o| o.id.to_string()) + .collect(); + warn!( + "Emergency stop affecting {} open orders: {:?}", + order_ids.len(), + order_ids + ); + order_ids + }, + }; + Ok(Response::new(EmergencyStopResponse { success: true, message: format!("Emergency stop activated: {}", req.reason), timestamp: chrono::Utc::now().timestamp(), - affected_orders: vec![], // TODO: Query order_manager for open orders at shutdown + affected_orders, })) }