From 4e746f60cac2bdf013e4c8f26bc0b29d04f1a765 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Tue, 4 Aug 2026 02:10:30 -0700 Subject: [PATCH] feat(supervisor): expose concurrent UDS auto-start --- dstack/Cargo.lock | 1 + dstack/supervisor/client/Cargo.toml | 1 + dstack/supervisor/client/src/lib.rs | 55 ++++++++++++++++++++++++++-- dstack/supervisor/client/src/main.rs | 50 ++++++++++++++++++++++++- dstack/supervisor/src/web_api.rs | 10 +++-- 5 files changed, 109 insertions(+), 8 deletions(-) diff --git a/dstack/Cargo.lock b/dstack/Cargo.lock index 63766de8f..7d45c4b89 100644 --- a/dstack/Cargo.lock +++ b/dstack/Cargo.lock @@ -7405,6 +7405,7 @@ dependencies = [ "hyper", "hyper-util", "hyperlocal", + "libc", "log", "serde", "serde_json", diff --git a/dstack/supervisor/client/Cargo.toml b/dstack/supervisor/client/Cargo.toml index 99db01f7e..e1a691b05 100644 --- a/dstack/supervisor/client/Cargo.toml +++ b/dstack/supervisor/client/Cargo.toml @@ -27,6 +27,7 @@ serde.workspace = true http-body-util.workspace = true tracing-subscriber.workspace = true log.workspace = true +libc.workspace = true fs-err.workspace = true futures.workspace = true diff --git a/dstack/supervisor/client/src/lib.rs b/dstack/supervisor/client/src/lib.rs index 53257e5c9..f377549ec 100644 --- a/dstack/supervisor/client/src/lib.rs +++ b/dstack/supervisor/client/src/lib.rs @@ -11,6 +11,37 @@ use supervisor::{ProcessConfig, ProcessInfo, Response}; pub use supervisor; +#[cfg(unix)] +fn acquire_uds_start_lock(uds: &Path) -> Result { + use std::fs::OpenOptions; + use std::os::unix::fs::{MetadataExt as _, OpenOptionsExt as _}; + + let lock_path = uds.with_extension("lock"); + let lock = OpenOptions::new() + .read(true) + .write(true) + .create(true) + .mode(0o600) + .custom_flags(libc::O_CLOEXEC | libc::O_NOFOLLOW) + .open(&lock_path) + .with_context(|| { + format!( + "Failed to open Supervisor start lock {}", + lock_path.display() + ) + })?; + let metadata = lock.metadata()?; + let effective_uid = unsafe { libc::geteuid() }; + if metadata.uid() != effective_uid || metadata.mode() & 0o077 != 0 { + anyhow::bail!("Supervisor start lock is not owner-only"); + } + let result = unsafe { libc::flock(std::os::fd::AsRawFd::as_raw_fd(&lock), libc::LOCK_EX) }; + if result != 0 { + return Err(std::io::Error::last_os_error()).context("Failed to lock Supervisor startup"); + } + Ok(lock) +} + #[derive(Debug, Clone)] pub struct SupervisorClient { base_url: Arc, @@ -41,6 +72,12 @@ impl SupervisorClient { anyhow::bail!("Failed to connect to supervisor at {uri}"); } info!("Failed to connect to supervisor at {uri}, trying to start supervisor"); + let _start_lock = acquire_uds_start_lock(uds.as_ref())?; + // Another caller may have completed startup while this caller waited. + if client.probe(Duration::from_millis(500)).await.is_ok() { + info!("Connected to supervisor at {uri} after waiting for startup lock"); + return Ok(client); + } // if the uds exists, remove it if std::path::Path::new(uds.as_ref()).exists() { fs_err::remove_file(uds.as_ref())?; @@ -51,7 +88,8 @@ impl SupervisorClient { let log_file = log_file.as_ref().to_path_buf(); std::thread::spawn(move || { // start supervisor - let result = std::process::Command::new(supervisor_path) + let mut command = std::process::Command::new(supervisor_path); + command .arg("--uds") .arg(uds) .arg("--pid-file") @@ -59,8 +97,19 @@ impl SupervisorClient { .arg("--log-file") .arg(log_file) .args(if detached { &["--detach"][..] } else { &[] }) - .env("RUST_LOG", "info,rocket=warn") - .output(); + .env("RUST_LOG", "info,rocket=warn"); + #[cfg(unix)] + { + use std::os::unix::process::CommandExt as _; + + unsafe { + command.pre_exec(|| { + libc::umask(0o077); + Ok(()) + }); + } + } + let result = command.output(); let output = match result { Ok(output) => output, Err(err) => { diff --git a/dstack/supervisor/client/src/main.rs b/dstack/supervisor/client/src/main.rs index 4f50793e5..6c1e1e4c8 100644 --- a/dstack/supervisor/client/src/main.rs +++ b/dstack/supervisor/client/src/main.rs @@ -13,6 +13,26 @@ struct Cli { #[arg(long, default_value = "unix:/var/run/supervisor.sock")] base_url: String, + /// Start a missing Supervisor before connecting to a trusted Unix socket. + #[arg(long, requires_all = ["supervisor_path", "pid_file", "log_file"])] + auto_start: bool, + + /// Supervisor executable used only with --auto-start. + #[arg(long)] + supervisor_path: Option, + + /// PID file passed to an auto-started Supervisor. + #[arg(long)] + pid_file: Option, + + /// Log file passed to an auto-started Supervisor. + #[arg(long)] + log_file: Option, + + /// Detach an auto-started Supervisor process. + #[arg(long, requires = "auto_start")] + detached: bool, + #[command(subcommand)] command: Commands, } @@ -50,11 +70,37 @@ async fn main() -> Result<()> { { use tracing_subscriber::{fmt, EnvFilter}; let filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info")); - fmt().with_env_filter(filter).with_ansi(false).init(); + fmt() + .with_env_filter(filter) + .with_ansi(false) + .with_writer(std::io::stderr) + .init(); } let cli = Cli::parse(); - let client = SupervisorClient::new(&cli.base_url); + let client = if cli.auto_start { + let uds = cli + .base_url + .strip_prefix("unix:") + .ok_or_else(|| anyhow::anyhow!("--auto-start requires a unix: base URL"))?; + SupervisorClient::start_and_connect_uds( + cli.supervisor_path + .as_deref() + .ok_or_else(|| anyhow::anyhow!("missing --supervisor-path"))?, + uds, + cli.pid_file + .as_deref() + .ok_or_else(|| anyhow::anyhow!("missing --pid-file"))?, + cli.log_file + .as_deref() + .ok_or_else(|| anyhow::anyhow!("missing --log-file"))?, + cli.detached, + true, + ) + .await? + } else { + SupervisorClient::new(&cli.base_url) + }; match cli.command { Commands::Deploy { id, command, args } => { diff --git a/dstack/supervisor/src/web_api.rs b/dstack/supervisor/src/web_api.rs index 2521ee20b..dbe0841dc 100644 --- a/dstack/supervisor/src/web_api.rs +++ b/dstack/supervisor/src/web_api.rs @@ -6,7 +6,7 @@ use anyhow::{anyhow, Result}; use or_panic::ResultOrPanic; use rocket::figment::Figment; use rocket::serde::json::Json; -use rocket::{delete, get, post, routes, Build, Rocket, State}; +use rocket::{delete, get, post, routes, Build, Rocket, Shutdown, State}; use serde::{Deserialize, Serialize}; use tokio::signal; use tracing::info; @@ -81,8 +81,12 @@ fn clear(supervisor: &State) -> Json> { } #[post("/shutdown")] -async fn shutdown(supervisor: &State) -> Json> { - to_json(perform_shutdown(supervisor, false).await) +async fn shutdown(supervisor: &State, shutdown: Shutdown) -> Json> { + let result = supervisor.shutdown().await; + if result.is_ok() { + shutdown.notify(); + } + to_json(result) } async fn perform_shutdown(supervisor: &Supervisor, force: bool) -> Result<()> {