diff --git a/dstack/vmm/src/app.rs b/dstack/vmm/src/app.rs index 7a9a32b81..1a6e78a74 100644 --- a/dstack/vmm/src/app.rs +++ b/dstack/vmm/src/app.rs @@ -385,6 +385,19 @@ impl App { } pub async fn start_vm(&self, id: &str) -> Result<()> { + self.start_vm_with_restart_policy(id, true).await + } + + async fn start_vm_with_restart_policy( + &self, + id: &str, + reset_restart_policy: bool, + ) -> Result<()> { + if reset_restart_policy { + if let Some(vm) = self.lock().get_mut(id) { + vm.state.auto_restart.reset(); + } + } { let state = self.lock(); if let Some(vm) = state.get(id) { @@ -464,6 +477,9 @@ impl App { } pub async fn stop_vm(&self, id: &str) -> Result<()> { + if let Some(vm) = self.lock().get_mut(id) { + vm.state.auto_restart.reset(); + } self.set_started(id, false)?; self.stop_vm_process(id).await?; Ok(()) @@ -1165,25 +1181,86 @@ impl App { .filter(|v| v.state.status.is_running()) .map(|v| v.config.id.clone()) .collect::>(); - let exited_vms = self - .lock() - .iter_vms() - .filter(|vm| { + let now = std::time::Instant::now(); + let mut restart_vms = Vec::new(); + { + let mut state = self.lock(); + for vm in state.vms.values_mut() { + let id = &vm.config.manifest.id; if vm.state.removing { - return false; + vm.state.auto_restart.reset(); + continue; } - let Ok(workdir) = self.work_dir(&vm.config.manifest.id) else { - warn!(id = %vm.config.manifest.id, "skipping restart: invalid VM id"); - return false; + if running_vms.contains(id) { + if vm + .state + .auto_restart + .observe_running(now, self.config.cvm.auto_restart.reset_window) + { + info!( + id, + "automatic restart retry budget reset after healthy window" + ); + } + continue; + } + let Ok(workdir) = self.work_dir(id) else { + warn!(id, "skipping restart: invalid VM id"); + vm.state.auto_restart.reset(); + continue; }; let started = workdir.started().unwrap_or(false); - started && !running_vms.contains(&vm.config.manifest.id) - }) - .map(|vm| vm.config.manifest.id.clone()) - .collect::>(); - for id in exited_vms { - info!("Restarting VM {id}"); - self.start_vm(&id).await?; + if !started { + vm.state.auto_restart.reset(); + continue; + } + match vm + .state + .auto_restart + .observe_exited(now, &self.config.cvm.auto_restart) + { + AutoRestartDecision::Scheduled { delay_secs } => { + info!(id, delay_secs, "automatic restart scheduled"); + } + AutoRestartDecision::Restart { + attempt, + next_delay_secs, + } => { + info!(id, attempt, next_delay_secs, "automatic restart attempt"); + restart_vms.push(id.clone()); + } + AutoRestartDecision::Exhausted { attempts } => { + warn!(id, attempts, "automatic restart retry limit exhausted"); + vm.state.events.push_back(pb::GuestEvent { + event: "vmm.auto_restart.exhausted".into(), + body: format!( + "Automatic restart stopped after {attempts} failed attempts" + ), + timestamp: SystemTime::now() + .duration_since(SystemTime::UNIX_EPOCH) + .unwrap_or_default() + .as_millis() as u64, + }); + while vm.state.events.len() > self.config.event_buffer_size { + vm.state.events.pop_front(); + } + } + AutoRestartDecision::Wait => {} + } + } + } + for id in restart_vms { + // A manual stop may have landed after the restart decision was made. + let Ok(workdir) = self.work_dir(&id) else { + warn!(id, "skipping restart: invalid VM id"); + continue; + }; + if !workdir.started().unwrap_or(false) { + continue; + } + if let Err(error) = self.start_vm_with_restart_policy(&id, false).await { + warn!(id, %error, "automatic restart attempt failed"); + } } Ok(()) } @@ -1616,6 +1693,77 @@ mod tests { fn hex_of(byte: u8, len: usize) -> String { hex::encode(vec![byte; len]) } + fn restart_config() -> crate::config::AutoRestartConfig { + crate::config::AutoRestartConfig { + enabled: true, + interval: 1, + max_retries: 3, + initial_backoff: 2, + max_backoff: 5, + reset_window: 10, + } + } + + #[test] + fn auto_restart_policy_backs_off_caps_and_exhausts_once() { + let config = restart_config(); + let start = std::time::Instant::now(); + let mut state = AutoRestartState::default(); + assert_eq!( + state.observe_exited(start, &config), + AutoRestartDecision::Scheduled { delay_secs: 2 } + ); + assert_eq!( + state.observe_exited(start + std::time::Duration::from_secs(1), &config), + AutoRestartDecision::Wait + ); + assert_eq!( + state.observe_exited(start + std::time::Duration::from_secs(2), &config), + AutoRestartDecision::Restart { + attempt: 1, + next_delay_secs: 4 + } + ); + assert_eq!( + state.observe_exited(start + std::time::Duration::from_secs(6), &config), + AutoRestartDecision::Restart { + attempt: 2, + next_delay_secs: 5 + } + ); + assert_eq!( + state.observe_exited(start + std::time::Duration::from_secs(11), &config), + AutoRestartDecision::Restart { + attempt: 3, + next_delay_secs: 5 + } + ); + assert_eq!( + state.observe_exited(start + std::time::Duration::from_secs(12), &config), + AutoRestartDecision::Exhausted { attempts: 3 } + ); + assert_eq!( + state.observe_exited(start + std::time::Duration::from_secs(20), &config), + AutoRestartDecision::Wait + ); + } + + #[test] + fn auto_restart_policy_resets_only_after_healthy_window() { + let config = restart_config(); + let start = std::time::Instant::now(); + let mut state = AutoRestartState::default(); + state.observe_exited(start, &config); + state.observe_exited(start + std::time::Duration::from_secs(2), &config); + assert!(!state.observe_running(start + std::time::Duration::from_secs(3), 10)); + assert!(!state.observe_running(start + std::time::Duration::from_secs(12), 10)); + assert!(state.observe_running(start + std::time::Duration::from_secs(13), 10)); + assert_eq!(state.attempts, 0); + assert_eq!( + state.observe_exited(start + std::time::Duration::from_secs(14), &config), + AutoRestartDecision::Scheduled { delay_secs: 2 } + ); + } #[test] fn simulator_config_is_written_separately_with_measurement_inputs() -> Result<()> { @@ -2351,6 +2499,77 @@ pub struct VmState { state: VmStateMut, } +/// Per-process retry bookkeeping; intentionally reset whenever the VMM restarts. +#[derive(Debug, Clone, Default)] +struct AutoRestartState { + attempts: u32, + next_retry: Option, + healthy_since: Option, + exhausted_reported: bool, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum AutoRestartDecision { + Wait, + Scheduled { delay_secs: u64 }, + Restart { attempt: u32, next_delay_secs: u64 }, + Exhausted { attempts: u32 }, +} + +impl AutoRestartState { + fn reset(&mut self) { + *self = Self::default(); + } + + fn observe_running(&mut self, now: std::time::Instant, reset_window: u64) -> bool { + let healthy_since = self.healthy_since.get_or_insert(now); + if self.attempts > 0 + && now.duration_since(*healthy_since) >= std::time::Duration::from_secs(reset_window) + { + self.reset(); + return true; + } + false + } + + fn observe_exited( + &mut self, + now: std::time::Instant, + config: &crate::config::AutoRestartConfig, + ) -> AutoRestartDecision { + self.healthy_since = None; + if self.attempts >= config.max_retries { + if self.exhausted_reported { + return AutoRestartDecision::Wait; + } + self.exhausted_reported = true; + return AutoRestartDecision::Exhausted { + attempts: self.attempts, + }; + } + let Some(next_retry) = self.next_retry else { + self.next_retry = Some(now + std::time::Duration::from_secs(config.initial_backoff)); + return AutoRestartDecision::Scheduled { + delay_secs: config.initial_backoff, + }; + }; + if now < next_retry { + return AutoRestartDecision::Wait; + } + self.attempts += 1; + let shift = self.attempts.min(63); + let delay_secs = config + .initial_backoff + .saturating_mul(1u64 << shift) + .min(config.max_backoff); + self.next_retry = Some(now + std::time::Duration::from_secs(delay_secs)); + AutoRestartDecision::Restart { + attempt: self.attempts, + next_delay_secs: delay_secs, + } + } +} + #[derive(Debug, Clone, Default)] struct VmStateMut { boot_progress: String, @@ -2359,6 +2578,7 @@ struct VmStateMut { runtime_networks: Vec, devices: GpuConfig, events: VecDeque, + auto_restart: AutoRestartState, /// True when the VM is being removed (cleanup in progress). removing: bool, } diff --git a/dstack/vmm/src/config.rs b/dstack/vmm/src/config.rs index 3fdd1d9d9..62458ce5f 100644 --- a/dstack/vmm/src/config.rs +++ b/dstack/vmm/src/config.rs @@ -166,7 +166,33 @@ pub struct PortMappingConfig { #[derive(Debug, Clone, Deserialize)] pub struct AutoRestartConfig { pub enabled: bool, + /// How often the supervisor state is sampled. pub interval: u64, + /// Maximum consecutive automatic restart attempts before intervention. + pub max_retries: u32, + /// Delay before the first retry. Later retries use exponential backoff. + pub initial_backoff: u64, + /// Upper bound for the exponential retry delay. + pub max_backoff: u64, + /// Continuous healthy runtime required to reset the retry budget. + pub reset_window: u64, +} + +impl AutoRestartConfig { + pub fn validate(&self) -> Result<()> { + if self.enabled { + if self.interval == 0 { + bail!("cvm.auto_restart.interval must be greater than zero when enabled"); + } + if self.initial_backoff == 0 { + bail!("cvm.auto_restart.initial_backoff must be greater than zero when enabled"); + } + if self.initial_backoff > self.max_backoff { + bail!("cvm.auto_restart.initial_backoff must not exceed max_backoff"); + } + } + Ok(()) + } } impl PortMappingConfig { @@ -814,6 +840,43 @@ impl Config { mod tests { use super::*; + #[test] + fn auto_restart_config_rejects_hot_loop_and_inverted_backoff() { + let mut config = AutoRestartConfig { + enabled: true, + interval: 0, + max_retries: 3, + initial_backoff: 2, + max_backoff: 5, + reset_window: 10, + }; + assert!(config + .validate() + .unwrap_err() + .to_string() + .contains("interval")); + config.interval = 1; + config.initial_backoff = 0; + assert!(config + .validate() + .unwrap_err() + .to_string() + .contains("initial_backoff")); + config.initial_backoff = 6; + assert!(config + .validate() + .unwrap_err() + .to_string() + .contains("max_backoff")); + config.max_backoff = 6; + assert!(config.validate().is_ok()); + + config.enabled = false; + config.interval = 0; + config.initial_backoff = 7; + assert!(config.validate().is_ok()); + } + #[test] fn test_parse_qemu_version_debian_format() { let output = "QEMU emulator version 8.2.2 (Debian 2:8.2.2+ds-0ubuntu1.4+tdx1.0)\nCopyright (c) 2003-2023 Fabrice Bellard and the QEMU Project developers"; diff --git a/dstack/vmm/src/main.rs b/dstack/vmm/src/main.rs index b16dbc2e8..a572cfe9f 100644 --- a/dstack/vmm/src/main.rs +++ b/dstack/vmm/src/main.rs @@ -153,11 +153,11 @@ async fn auto_restart_task(app: App) { let mut interval = tokio::time::interval(Duration::from_secs(app.config.cvm.auto_restart.interval)); loop { + interval.tick().await; info!("Checking for exited VMs"); if let Err(err) = app.try_restart_exited_vms().await { error!("Failed to restart exited VMs: {err:?}"); } - interval.tick().await; } } @@ -185,6 +185,11 @@ async fn main() -> Result<()> { .host_api .validate() .context("Invalid host_api configuration")?; + config + .cvm + .auto_restart + .validate() + .context("Invalid cvm.auto_restart configuration")?; // Handle commands match args.command.unwrap_or_default() { diff --git a/dstack/vmm/vmm.toml b/dstack/vmm/vmm.toml index b405e5e9f..36fe2d47f 100644 --- a/dstack/vmm/vmm.toml +++ b/dstack/vmm/vmm.toml @@ -119,6 +119,10 @@ range = [ [cvm.auto_restart] enabled = true interval = 20 +max_retries = 5 +initial_backoff = 5 +max_backoff = 300 +reset_window = 300 [cvm.gpu] enabled = false