Merge pull request #584 from Dstack-TEE/worktree-vmm-discovery

feat(vmm): local VMM instance discovery

Author: kvinwang

vmm/src/discovery.rs

@@ -0,0 +1,137 @@
+// SPDX-FileCopyrightText: © 2025 Phala Network <dstack@phala.network>
+//
+// SPDX-License-Identifier: Apache-2.0
+
+//! VMM instance discovery via registration files.
+//!
+//! On startup each `dstack-vmm` process writes a JSON file to a well-known
+//! directory so that CLI tools can discover all running instances on the host.
+
+use anyhow::{Context, Result};
+use serde::{Deserialize, Serialize};
+use std::path::{Path, PathBuf};
+use tracing::{info, warn};
+use uuid::Uuid;
+
+/// Well-known directory where VMM instances register themselves.
+const DISCOVERY_DIR: &str = "/run/dstack-vmm";
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct VmmInstanceInfo {
+    /// Unique identifier for this VMM instance (random UUID).
+    pub id: String,
+    /// Process ID.
+    pub pid: u32,
+    /// Address the external API listens on (e.g. "unix:./vmm.sock" or "0.0.0.0:9080").
+    pub address: String,
+    /// Working directory of the VMM process.
+    pub working_dir: String,
+    /// Path to the configuration file (if provided via -c).
+    pub config_file: Option<String>,
+    /// Path where VM images are stored.
+    pub image_path: String,
+    /// Path where VM runtime data is stored.
+    pub run_path: String,
+    /// Node name from configuration.
+    pub node_name: String,
+    /// VMM version string.
+    pub version: String,
+    /// Unix timestamp (seconds) when the instance started.
+    pub started_at: u64,
+}
+
+/// Handle that manages the lifecycle of a discovery registration file.
+/// The file is removed when this handle is dropped.
+pub struct DiscoveryRegistration {
+    path: PathBuf,
+}
+
+impl DiscoveryRegistration {
+    /// Register a new VMM instance. Creates the discovery directory if needed
+    /// and writes the instance info JSON file.
+    pub fn register(
+        listen_address: &str,
+        config_file: Option<&str>,
+        image_path: &Path,
+        run_path: &Path,
+        node_name: &str,
+        version: &str,
+    ) -> Result<Self> {
+        let dir = Path::new(DISCOVERY_DIR);
+        fs_err::create_dir_all(dir).context("failed to create discovery directory")?;
+
+        let id = Uuid::new_v4().to_string();
+        let path = dir.join(format!("{id}.json"));
+
+        let cwd = std::env::current_dir().context("failed to get current directory")?;
+
+        let info = VmmInstanceInfo {
+            id: id.clone(),
+            pid: std::process::id(),
+            address: listen_address.to_string(),
+            working_dir: cwd.to_string_lossy().to_string(),
+            config_file: config_file.map(|s| s.to_string()),
+            image_path: image_path.to_string_lossy().to_string(),
+            run_path: run_path.to_string_lossy().to_string(),
+            node_name: node_name.to_string(),
+            version: version.to_string(),
+            started_at: std::time::SystemTime::now()
+                .duration_since(std::time::UNIX_EPOCH)
+                .unwrap_or_default()
+                .as_secs(),
+        };
+
+        let json =
+            serde_json::to_string_pretty(&info).context("failed to serialize instance info")?;
+        fs_err::write(&path, &json).context("failed to write discovery file")?;
+
+        info!("registered VMM instance {id} at {}", path.display());
+        Ok(Self { path })
+    }
+}
+
+impl Drop for DiscoveryRegistration {
+    fn drop(&mut self) {
+        match std::fs::remove_file(&self.path) {
+            Ok(()) => info!("unregistered VMM instance at {}", self.path.display()),
+            Err(e) => warn!(
+                "failed to remove discovery file {}: {e}",
+                self.path.display()
+            ),
+        }
+    }
+}
+
+/// Clean up stale discovery files from dead processes.
+pub fn cleanup_stale_registrations() {
+    let dir = Path::new(DISCOVERY_DIR);
+    let entries = match std::fs::read_dir(dir) {
+        Ok(e) => e,
+        Err(_) => return,
+    };
+
+    for entry in entries.flatten() {
+        let path = entry.path();
+        if path.extension().and_then(|e| e.to_str()) != Some("json") {
+            continue;
+        }
+        let content = match std::fs::read_to_string(&path) {
+            Ok(c) => c,
+            Err(_) => continue,
+        };
+        let info: VmmInstanceInfo = match serde_json::from_str(&content) {
+            Ok(i) => i,
+            Err(_) => continue,
+        };
+        // Check if the process is still alive
+        let alive = Path::new(&format!("/proc/{}", info.pid)).exists();
+        if !alive {
+            info!(
+                "removing stale discovery file for pid {} at {}",
+                info.pid,
+                path.display()
+            );
+            let _ = std::fs::remove_file(&path);
+        }
+    }
+}

vmm/src/main.rs

@@ -23,6 +23,7 @@ use tracing::{error, info};
 
 mod app;
 mod config;
+mod discovery;
 mod guest_api_service;
 mod host_api_service;
 mod main_routes;
@@ -179,6 +180,31 @@ async fn main() -> Result<()> {
         }
     }
 
+    // Register this VMM instance for local discovery
+    discovery::cleanup_stale_registrations();
+    let listen_address = {
+        // Use Rocket's Endpoint type to parse the address exactly as Rocket would,
+        // then override the port with the figment's port value (matching Rocket's behavior).
+        let endpoint: rocket::listener::Endpoint =
+            figment.extract_inner("address").unwrap_or_default();
+        match endpoint.tcp() {
+            Some(addr) => {
+                let port: u16 = figment.extract_inner("port").unwrap_or(addr.port());
+                format!("{}:{port}", addr.ip())
+            }
+            None => endpoint.to_string(),
+        }
+    };
+    let _discovery_reg = discovery::DiscoveryRegistration::register(
+        &listen_address,
+        args.config.as_deref(),
+        &config.image_path,
+        &config.run_path,
+        &config.node_name,
+        &app_version(),
+    )
+    .context("failed to register VMM instance for discovery")?;
+
     let api_auth = ApiToken::new(config.auth.tokens.clone(), config.auth.enabled);
     let supervisor = {
         let cfg = &config.supervisor;