refactor(ml): deduplicate PPO circuit breaker (re-export from DQN)
PPO's circuit_breaker.rs was a 276-line copy-paste of DQN's version (only difference: "PPO" vs "DQN" in log messages). Since ppo/ppo.rs already imports from dqn::circuit_breaker, the PPO copy was dead code. Replace with a 6-line re-export module. Tests still pass via the canonical DQN implementation. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -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<S>(duration: &Duration, serializer: S) -> Result<S::Ok, S::Error>
|
||||
where
|
||||
S: Serializer,
|
||||
{
|
||||
duration.as_secs().serialize(serializer)
|
||||
}
|
||||
|
||||
pub(super) fn deserialize<'de, D>(deserializer: D) -> Result<Duration, D::Error>
|
||||
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<RwLock<CircuitState>>,
|
||||
consecutive_failures: AtomicU32,
|
||||
consecutive_successes: AtomicU32,
|
||||
half_open_calls: AtomicU32,
|
||||
open_timestamp: Arc<RwLock<Option<Instant>>>,
|
||||
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};
|
||||
|
||||
Reference in New Issue
Block a user