fix(vmm): bound automatic restart retries

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::<BTreeSet<_>>();
-        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::<Vec<_>>();
-        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<std::time::Instant>,
+    healthy_since: Option<std::time::Instant>,
+    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<Networking>,
     devices: GpuConfig,
     events: VecDeque<pb::GuestEvent>,
+    auto_restart: AutoRestartState,
     /// True when the VM is being removed (cleanup in progress).
     removing: bool,
 }

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";

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() {