Merge pull request Dstack-TEE#146 from Dstack-TEE/auto-restart

vmm: Auto restart exited VMs

Author: kvinwang

vmm/src/app.rs

@@ -12,7 +12,7 @@ use guest_api::client::DefaultClient as GuestClient;
 use id_pool::IdPool;
 use ra_rpc::client::RaClient;
 use serde::{Deserialize, Serialize};
-use std::collections::{HashMap, HashSet};
+use std::collections::{BTreeSet, HashMap, HashSet};
 use std::net::IpAddr;
 use std::path::{Path, PathBuf};
 use std::sync::{Arc, Mutex, MutexGuard};
@@ -658,6 +658,33 @@ impl App {
             })
             .collect())
     }
+
+    pub(crate) async fn try_restart_exited_vms(&self) -> Result<()> {
+        let running_vms = self
+            .supervisor
+            .list()
+            .await
+            .context("Failed to list VMs")?
+            .iter()
+            .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 workdir = self.work_dir(&vm.config.manifest.id);
+                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?;
+        }
+        Ok(())
+    }
 }
 
 fn paginate<T>(items: Vec<T>, page: u32, page_size: u32) -> impl Iterator<Item = T> {

vmm/src/config.rs

@@ -62,6 +62,12 @@ pub struct PortMappingConfig {
     pub range: Vec<PortRange>,
 }
 
+#[derive(Debug, Clone, Deserialize)]
+pub struct AutoRestartConfig {
+    pub enabled: bool,
+    pub interval: u64,
+}
+
 impl PortMappingConfig {
     pub fn is_allowed(&self, protocol: &str, port: u16) -> bool {
         if !self.enabled {
@@ -111,6 +117,9 @@ pub struct CvmConfig {
     /// The tmp CA key
     #[serde(default)]
     pub tmp_ca_key: String,
+
+    /// Auto restart configuration
+    pub auto_restart: AutoRestartConfig,
 }
 
 #[derive(Debug, Clone, Deserialize)]

vmm/src/main.rs

@@ -1,4 +1,4 @@
-use std::path::Path;
+use std::{path::Path, time::Duration};
 
 use anyhow::{anyhow, Context, Result};
 use app::App;
@@ -16,6 +16,7 @@ use rocket::{
 use rocket_apitoken::ApiToken;
 use rocket_vsock_listener::VsockListener;
 use supervisor_client::SupervisorClient;
+use tracing::{error, info};
 
 mod app;
 mod config;
@@ -102,6 +103,22 @@ async fn run_host_api(app: App, figment: Figment) -> Result<()> {
     Ok(())
 }
 
+async fn auto_restart_task(app: App) {
+    if !app.config.cvm.auto_restart.enabled {
+        info!("Auto restart CVMs is disabled");
+        return;
+    }
+    let mut interval =
+        tokio::time::interval(Duration::from_secs(app.config.cvm.auto_restart.interval));
+    loop {
+        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;
+    }
+}
+
 #[rocket::main]
 async fn main() -> Result<()> {
     {
@@ -130,6 +147,7 @@ async fn main() -> Result<()> {
     };
     let state = app::App::new(config, supervisor);
     state.reload_vms().await.context("Failed to reload VMs")?;
+    tokio::spawn(auto_restart_task(state.clone()));
 
     tokio::select! {
         result = run_external_api(state.clone(), figment.clone(), api_auth) => {

vmm/vmm.toml

@@ -38,6 +38,10 @@ range = [
     { protocol = "tcp", from = 1, to = 20000 },
 ]
 
+[cvm.auto_restart]
+enabled = true
+interval = 20
+
 [cvm.gpu]
 enabled = false
 # The product IDs of the GPUs to discover