diff --git a/ml/src/ppo/circuit_breaker.rs b/ml/src/ppo/circuit_breaker.rs index 9ecd806c0..c427e5d80 100644 --- a/ml/src/ppo/circuit_breaker.rs +++ b/ml/src/ppo/circuit_breaker.rs @@ -1,275 +1,6 @@ -//! Simplified Circuit Breaker for PPO Training +//! Circuit Breaker for PPO Training //! -//! Adapted from DQN's CircuitBreaker for PPO-specific failure management. -//! Provides dynamic throttling to prevent runaway losses during training. -//! Unlike the full risk crate CircuitBreaker, this is designed for single-process -//! ML training without requiring Redis coordination or broker services. +//! Re-exports the shared circuit breaker implementation from the DQN module. +//! Both DQN and PPO use the same circuit breaker pattern for training throttling. -use std::sync::atomic::{AtomicU32, AtomicU64, Ordering}; -use std::sync::Arc; -use std::time::{Duration, Instant}; - -use parking_lot::RwLock; -use tracing::{debug, info, warn}; - -/// Circuit breaker state -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum CircuitState { - /// Circuit is closed - training allowed - Closed, - /// Circuit is open - training blocked (cooldown period) - Open, - /// Circuit is half-open - limited training to test recovery - HalfOpen, -} - -/// Configuration for the circuit breaker -#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] -pub struct CircuitBreakerConfig { - /// Number of consecutive failures before opening - pub failure_threshold: usize, - /// Number of consecutive successes needed to close from half-open - pub success_threshold: usize, - /// Cooldown duration when circuit opens (in seconds) - #[serde(with = "duration_serde")] - pub timeout_duration: Duration, - /// Maximum number of test calls in half-open state - pub half_open_max_calls: usize, -} - -// Helper module for Duration serialization -mod duration_serde { - use serde::{Deserialize, Deserializer, Serialize, Serializer}; - use std::time::Duration; - - pub(super) fn serialize(duration: &Duration, serializer: S) -> Result - where - S: Serializer, - { - duration.as_secs().serialize(serializer) - } - - pub(super) fn deserialize<'de, D>(deserializer: D) -> Result - where - D: Deserializer<'de>, - { - let secs = u64::deserialize(deserializer)?; - Ok(Duration::from_secs(secs)) - } -} - -impl Default for CircuitBreakerConfig { - fn default() -> Self { - Self { - failure_threshold: 5, - success_threshold: 3, - timeout_duration: Duration::from_secs(60), - half_open_max_calls: 2, - } - } -} - -/// Simplified circuit breaker for PPO training -#[derive(Debug)] -pub struct CircuitBreaker { - config: CircuitBreakerConfig, - state: Arc>, - consecutive_failures: AtomicU32, - consecutive_successes: AtomicU32, - half_open_calls: AtomicU32, - open_timestamp: Arc>>, - total_failures: AtomicU64, - total_successes: AtomicU64, -} - -impl CircuitBreaker { - /// Create new circuit breaker with configuration - pub fn new(config: CircuitBreakerConfig) -> Self { - info!( - "Initializing PPO Circuit Breaker - failure_threshold={}, success_threshold={}, timeout={}s", - config.failure_threshold, - config.success_threshold, - config.timeout_duration.as_secs() - ); - - Self { - config, - state: Arc::new(RwLock::new(CircuitState::Closed)), - consecutive_failures: AtomicU32::new(0), - consecutive_successes: AtomicU32::new(0), - half_open_calls: AtomicU32::new(0), - open_timestamp: Arc::new(RwLock::new(None)), - total_failures: AtomicU64::new(0), - total_successes: AtomicU64::new(0), - } - } - - /// Check if training request should be allowed through - pub fn allow_request(&self) -> bool { - let current_state = *self.state.read(); - - match current_state { - CircuitState::Closed => true, - CircuitState::Open => { - // Check if timeout has elapsed - if let Some(open_time) = *self.open_timestamp.read() { - if open_time.elapsed() >= self.config.timeout_duration { - // Transition to half-open - info!( - "PPO Circuit breaker transitioning to HALF-OPEN after {}s cooldown", - self.config.timeout_duration.as_secs() - ); - *self.state.write() = CircuitState::HalfOpen; - self.half_open_calls.store(0, Ordering::SeqCst); - true - } else { - debug!( - "PPO Circuit breaker OPEN - blocking request ({}s remaining)", - (self.config.timeout_duration - open_time.elapsed()).as_secs() - ); - false - } - } else { - // Shouldn't happen, but allow request - warn!("PPO Circuit breaker OPEN but no timestamp - allowing request"); - true - } - } - CircuitState::HalfOpen => { - // Allow limited test calls - let calls = self.half_open_calls.fetch_add(1, Ordering::SeqCst); - if calls < self.config.half_open_max_calls as u32 { - debug!( - "PPO Circuit breaker HALF-OPEN - allowing test call {}/{}", - calls + 1, - self.config.half_open_max_calls - ); - true - } else { - debug!("PPO Circuit breaker HALF-OPEN - max test calls reached"); - false - } - } - } - } - - /// Record a successful training step - pub fn record_success(&self) { - self.total_successes.fetch_add(1, Ordering::Relaxed); - self.consecutive_failures.store(0, Ordering::SeqCst); - let successes = self.consecutive_successes.fetch_add(1, Ordering::SeqCst) + 1; - - let current_state = *self.state.read(); - - match current_state { - CircuitState::Closed => { - // Already closed, nothing to do - debug!( - "PPO Circuit breaker CLOSED - success recorded ({} consecutive)", - successes - ); - } - CircuitState::HalfOpen => { - if successes >= self.config.success_threshold as u32 { - // Transition to closed - info!( - "PPO Circuit breaker transitioning to CLOSED after {} consecutive successes", - successes - ); - *self.state.write() = CircuitState::Closed; - self.consecutive_successes.store(0, Ordering::SeqCst); - *self.open_timestamp.write() = None; - } else { - debug!( - "PPO Circuit breaker HALF-OPEN - success {}/{}", - successes, self.config.success_threshold - ); - } - } - CircuitState::Open => { - // Shouldn't happen, but reset if we somehow got a success - warn!("PPO Circuit breaker OPEN but received success - resetting"); - *self.state.write() = CircuitState::Closed; - self.consecutive_successes.store(0, Ordering::SeqCst); - *self.open_timestamp.write() = None; - } - } - } - - /// Record a failed training step (e.g., NaN loss, gradient explosion) - pub fn record_failure(&self) { - self.total_failures.fetch_add(1, Ordering::Relaxed); - self.consecutive_successes.store(0, Ordering::SeqCst); - let failures = self.consecutive_failures.fetch_add(1, Ordering::SeqCst) + 1; - - let current_state = *self.state.read(); - - match current_state { - CircuitState::Closed => { - if failures >= self.config.failure_threshold as u32 { - // Transition to open - warn!( - "PPO Circuit breaker transitioning to OPEN after {} consecutive failures", - failures - ); - *self.state.write() = CircuitState::Open; - *self.open_timestamp.write() = Some(Instant::now()); - self.consecutive_failures.store(0, Ordering::SeqCst); - } else { - debug!( - "PPO Circuit breaker CLOSED - failure {}/{}", - failures, self.config.failure_threshold - ); - } - } - CircuitState::HalfOpen => { - // Any failure in half-open goes back to open - warn!("PPO Circuit breaker HALF-OPEN - failure detected, returning to OPEN"); - *self.state.write() = CircuitState::Open; - *self.open_timestamp.write() = Some(Instant::now()); - self.consecutive_failures.store(0, Ordering::SeqCst); - self.half_open_calls.store(0, Ordering::SeqCst); - } - CircuitState::Open => { - // Already open, just track - debug!("PPO Circuit breaker OPEN - additional failure recorded"); - } - } - } - - /// Get current circuit state - pub fn current_state(&self) -> CircuitState { - *self.state.read() - } - - /// Get statistics - pub fn stats(&self) -> CircuitBreakerStats { - CircuitBreakerStats { - state: self.current_state(), - consecutive_failures: self.consecutive_failures.load(Ordering::Relaxed), - consecutive_successes: self.consecutive_successes.load(Ordering::Relaxed), - total_failures: self.total_failures.load(Ordering::Relaxed), - total_successes: self.total_successes.load(Ordering::Relaxed), - } - } - - /// Reset the circuit breaker to closed state - pub fn reset(&self) { - info!("Manually resetting PPO circuit breaker to CLOSED state"); - *self.state.write() = CircuitState::Closed; - self.consecutive_failures.store(0, Ordering::SeqCst); - self.consecutive_successes.store(0, Ordering::SeqCst); - self.half_open_calls.store(0, Ordering::SeqCst); - *self.open_timestamp.write() = None; - } -} - -/// Circuit breaker statistics -#[derive(Debug, Clone)] -pub struct CircuitBreakerStats { - pub state: CircuitState, - pub consecutive_failures: u32, - pub consecutive_successes: u32, - pub total_failures: u64, - pub total_successes: u64, -} +pub use crate::dqn::circuit_breaker::{CircuitBreaker, CircuitBreakerConfig, CircuitBreakerStats, CircuitState};