360 lines
14 KiB
Rust
360 lines
14 KiB
Rust
//! External watchdog for a prebuilt Rust test or a direct reference worker.
|
|
//! Output is test progress, not proof of GPU progress. No total-runtime limit.
|
|
#[cfg(target_os = "macos")]
|
|
#[path = "../src/process_resources.rs"]
|
|
mod process_resources;
|
|
|
|
#[cfg(target_os = "macos")]
|
|
mod supervisor {
|
|
use super::process_resources::{ResourceUsage, resource_usage};
|
|
use serde_json::json;
|
|
use std::io::{self, Read, Write};
|
|
use std::os::unix::process::CommandExt;
|
|
use std::process::{Child, Command, Stdio};
|
|
use std::sync::{Arc, Mutex};
|
|
use std::thread;
|
|
use std::time::{Duration, Instant};
|
|
|
|
const SAMPLE_INTERVAL: Duration = Duration::from_millis(100);
|
|
|
|
pub(super) struct Limits {
|
|
pub(super) bytes: u64,
|
|
pub(super) start: Duration,
|
|
pub(super) idle: Duration,
|
|
}
|
|
|
|
struct Worker(Child);
|
|
|
|
impl Drop for Worker {
|
|
fn drop(&mut self) {
|
|
if !matches!(self.0.try_wait(), Ok(Some(_))) {
|
|
// This process group was created for this child alone. Also
|
|
// clean up helpers on error; memory accounting covers the test
|
|
// process, so this launcher is not a process-tree supervisor.
|
|
unsafe { libc::kill(-(self.0.id() as i32), libc::SIGKILL) };
|
|
let _ = self.0.kill();
|
|
let _ = self.0.wait();
|
|
}
|
|
}
|
|
}
|
|
|
|
#[derive(Default)]
|
|
struct Output {
|
|
last: Option<Instant>,
|
|
failed: bool,
|
|
}
|
|
|
|
// Worker-side probes are observations, not inference progress. Buffer JSON
|
|
// across pipe reads so a split heartbeat cannot keep a stuck worker alive.
|
|
// Plain output (including Rust test banners without a newline) still counts.
|
|
fn useful_output(pending: &mut Vec<u8>, bytes: &[u8]) -> bool {
|
|
pending.extend_from_slice(bytes);
|
|
let mut progressed = false;
|
|
while let Some(end) = pending.iter().position(|&b| b == b'\n') {
|
|
let line = &pending[..end];
|
|
let observation = serde_json::from_slice::<serde_json::Value>(line)
|
|
.ok()
|
|
.is_some_and(|value| {
|
|
matches!(
|
|
value["event"].as_str(),
|
|
Some(
|
|
"gpu_canary_sample"
|
|
| "gpu_canary_summary"
|
|
| "reference_canary_summary"
|
|
| "test_resource_sample"
|
|
)
|
|
)
|
|
});
|
|
progressed |= !observation && line.iter().any(|b| !b.is_ascii_whitespace());
|
|
pending.drain(..=end);
|
|
}
|
|
if pending.len() > 64 * 1024
|
|
|| pending
|
|
.iter()
|
|
.find(|b| !b.is_ascii_whitespace())
|
|
.is_some_and(|&b| b != b'{')
|
|
{
|
|
progressed = true;
|
|
pending.clear();
|
|
}
|
|
progressed
|
|
}
|
|
|
|
fn forward(
|
|
mut input: impl Read + Send + 'static,
|
|
mut output: impl Write + Send + 'static,
|
|
progress: Arc<Mutex<Output>>,
|
|
) -> thread::JoinHandle<()> {
|
|
thread::spawn(move || {
|
|
let mut bytes = [0; 4096];
|
|
let mut pending = Vec::new();
|
|
loop {
|
|
match input.read(&mut bytes) {
|
|
Ok(0) => return,
|
|
Ok(n) => {
|
|
if output
|
|
.write_all(&bytes[..n])
|
|
.and_then(|()| output.flush())
|
|
.is_err()
|
|
{
|
|
progress.lock().unwrap().failed = true;
|
|
return;
|
|
}
|
|
if useful_output(&mut pending, &bytes[..n]) {
|
|
progress.lock().unwrap().last = Some(Instant::now());
|
|
}
|
|
}
|
|
Err(error) if error.kind() == io::ErrorKind::Interrupted => continue,
|
|
Err(_) => {
|
|
progress.lock().unwrap().failed = true;
|
|
return;
|
|
}
|
|
}
|
|
}
|
|
})
|
|
}
|
|
|
|
pub(super) fn run(command: &mut Command, limits: Limits) -> Result<(), String> {
|
|
if limits.bytes == 0 || limits.start.is_zero() || limits.idle.is_zero() {
|
|
return Err("memory and progress limits must be positive".into());
|
|
}
|
|
// Fail before spawning if process accounting is unavailable.
|
|
resource_usage(std::process::id()).ok_or("process accounting unavailable")?;
|
|
let mut worker = Worker(
|
|
command
|
|
.process_group(0)
|
|
.stdin(Stdio::null())
|
|
.stdout(Stdio::piped())
|
|
.stderr(Stdio::piped())
|
|
.spawn()
|
|
.map_err(|e| format!("could not start test: {e}"))?,
|
|
);
|
|
let progress = Arc::new(Mutex::new(Output::default()));
|
|
let readers = [
|
|
forward(
|
|
worker.0.stdout.take().unwrap(),
|
|
io::stdout(),
|
|
progress.clone(),
|
|
),
|
|
forward(
|
|
worker.0.stderr.take().unwrap(),
|
|
io::stderr(),
|
|
progress.clone(),
|
|
),
|
|
];
|
|
let started = Instant::now();
|
|
let mut peak = ResourceUsage::default();
|
|
eprintln!(
|
|
"{}",
|
|
json!({"event":"test_supervisor_started", "pid":worker.0.id(),
|
|
"limit_bytes":limits.bytes, "start_seconds":limits.start.as_secs_f64(),
|
|
"idle_seconds":limits.idle.as_secs_f64(), "sample_ms":100})
|
|
);
|
|
let outcome = loop {
|
|
if let Some(status) = worker.0.try_wait().map_err(|e| e.to_string())? {
|
|
break if status.success() {
|
|
Ok(())
|
|
} else {
|
|
Err(format!("test exited with {status}"))
|
|
};
|
|
}
|
|
let now = Instant::now();
|
|
let Some(sample) = resource_usage(worker.0.id()) else {
|
|
// Exit can race the sample; missing telemetry on a live worker
|
|
// is a failure, never permission to keep running unmonitored.
|
|
if let Some(status) = worker.0.try_wait().map_err(|e| e.to_string())? {
|
|
break if status.success() {
|
|
Ok(())
|
|
} else {
|
|
Err(format!("test exited with {status}"))
|
|
};
|
|
}
|
|
break Err("process accounting lost".into());
|
|
};
|
|
peak.include(sample);
|
|
eprintln!(
|
|
"{}",
|
|
json!({"event":"test_resource_sample",
|
|
"elapsed_ms":started.elapsed().as_millis(), "usage":sample})
|
|
);
|
|
if sample.physical_bytes > limits.bytes || sample.peak_physical_bytes > limits.bytes {
|
|
break Err("memory_limit".into());
|
|
}
|
|
let output = progress.lock().map_err(|e| e.to_string())?;
|
|
if output.failed {
|
|
break Err("output monitoring failed".into());
|
|
}
|
|
let (last, timeout, reason) = output
|
|
.last
|
|
.map_or((started, limits.start, "start_timeout"), |last| {
|
|
(last, limits.idle, "continuation_timeout")
|
|
});
|
|
if now.saturating_duration_since(last) >= timeout {
|
|
break Err(reason.into());
|
|
}
|
|
drop(output);
|
|
thread::sleep(SAMPLE_INTERVAL);
|
|
};
|
|
drop(worker); // Kill and reap on every failure, including output errors.
|
|
for reader in readers {
|
|
reader.join().map_err(|_| "output reader panicked")?;
|
|
}
|
|
eprintln!(
|
|
"{}",
|
|
json!({"event":"test_resource_summary",
|
|
"elapsed_ms":started.elapsed().as_millis(), "peak":peak,
|
|
"error":outcome.as_ref().err()})
|
|
);
|
|
if progress.lock().map_err(|e| e.to_string())?.failed {
|
|
return Err("output monitoring failed".into());
|
|
}
|
|
outcome
|
|
}
|
|
|
|
pub(super) fn cli() -> Result<(), String> {
|
|
let args = std::env::args_os().skip(1).collect::<Vec<_>>();
|
|
if args.len() < 5 || (args[3] != "--command" && args.len() != 5) {
|
|
return Err("usage: test-supervisor MEMORY_MIB START_SECONDS IDLE_SECONDS TEST_BINARY TEST_FILTER\nOr: test-supervisor MEMORY_MIB START_SECONDS IDLE_SECONDS --command EXECUTABLE [ARG ...]\nMonitors the direct worker only; do not pass cargo or a spawning shell.".into());
|
|
}
|
|
let number = |i: usize| {
|
|
args[i]
|
|
.to_str()
|
|
.and_then(|s| s.parse::<u64>().ok())
|
|
.filter(|n| *n > 0)
|
|
.ok_or("limits must be positive integers")
|
|
};
|
|
let limits = Limits {
|
|
bytes: number(0)?
|
|
.checked_mul(1024 * 1024)
|
|
.ok_or("memory limit overflow")?,
|
|
start: Duration::from_secs(number(1)?),
|
|
idle: Duration::from_secs(number(2)?),
|
|
};
|
|
let mut command = worker_command(&args[3..])?;
|
|
run(&mut command, limits)
|
|
}
|
|
|
|
fn worker_command(args: &[std::ffi::OsString]) -> Result<Command, String> {
|
|
if args.len() < 2 || args[1].is_empty() {
|
|
return Err("worker executable or test filter must not be empty".into());
|
|
}
|
|
if args[0] == "--command" {
|
|
let mut command = Command::new(&args[1]);
|
|
command.args(&args[2..]);
|
|
Ok(command)
|
|
} else if args.len() == 2 {
|
|
let mut command = Command::new(&args[0]);
|
|
command
|
|
.arg(&args[1])
|
|
.args(["--include-ignored", "--test-threads=1", "--nocapture"]);
|
|
Ok(command)
|
|
} else {
|
|
Err("unexpected test arguments".into())
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn direct_worker_arguments_are_not_rust_test_arguments() {
|
|
let args = |a: &[&str]| a.iter().map(std::ffi::OsString::from).collect::<Vec<_>>();
|
|
let direct = worker_command(&args(&[
|
|
"--command",
|
|
"/usr/bin/python3",
|
|
"-u",
|
|
"reference.py",
|
|
]))
|
|
.unwrap();
|
|
assert_eq!(direct.get_program(), "/usr/bin/python3");
|
|
assert_eq!(
|
|
direct.get_args().collect::<Vec<_>>(),
|
|
["-u", "reference.py"]
|
|
);
|
|
let test = worker_command(&args(&["test-binary", "fixture"])).unwrap();
|
|
assert_eq!(
|
|
test.get_args().collect::<Vec<_>>(),
|
|
[
|
|
"fixture",
|
|
"--include-ignored",
|
|
"--test-threads=1",
|
|
"--nocapture"
|
|
]
|
|
);
|
|
assert!(worker_command(&args(&["--command", ""])).is_err());
|
|
assert!(worker_command(&args(&["test-binary", "fixture", "extra"])).is_err());
|
|
}
|
|
|
|
#[test]
|
|
fn watchdog_stops_memory_and_silence_but_not_total_runtime() {
|
|
let limits = || Limits {
|
|
bytes: 256 * 1024 * 1024,
|
|
start: Duration::from_millis(300),
|
|
idle: Duration::from_millis(300),
|
|
};
|
|
let mut sleeping = Command::new("/bin/sleep");
|
|
sleeping.arg("5");
|
|
assert_eq!(run(&mut sleeping, limits()).unwrap_err(), "start_timeout");
|
|
let mut started = Command::new("/bin/sh");
|
|
started.args(["-c", "printf ready; exec /bin/sleep 5"]);
|
|
assert_eq!(
|
|
run(&mut started, limits()).unwrap_err(),
|
|
"continuation_timeout"
|
|
);
|
|
for (prefix, expected) in [
|
|
("", "start_timeout"),
|
|
("printf 'ready\\n'; ", "continuation_timeout"),
|
|
] {
|
|
let mut heartbeat = Command::new("/bin/sh");
|
|
heartbeat.args(["-c", &format!("{prefix}while :; do printf '{{\"event\":\"gpu_canary_sample\",\"ok\":true}}\\n'; sleep 0.05; done")]);
|
|
assert_eq!(run(&mut heartbeat, limits()).unwrap_err(), expected);
|
|
}
|
|
let mut bounded = limits();
|
|
bounded.bytes = 1;
|
|
assert_eq!(run(&mut sleeping, bounded).unwrap_err(), "memory_limit");
|
|
let mut progressing = Command::new("/bin/sh");
|
|
progressing.args([
|
|
"-c",
|
|
"for i in 1 2 3 4 5 6 7 8; do printf progress; sleep 0.1; done",
|
|
]);
|
|
let before = Instant::now();
|
|
run(&mut progressing, limits()).unwrap();
|
|
assert!(before.elapsed() >= Duration::from_millis(600));
|
|
let mut failure = Command::new("/bin/sh");
|
|
failure.args(["-c", "exit 7"]);
|
|
assert!(
|
|
run(&mut failure, limits())
|
|
.unwrap_err()
|
|
.contains("exit status: 7")
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn split_probe_records_do_not_reset_progress() {
|
|
let mut pending = Vec::new();
|
|
assert!(!useful_output(&mut pending, b"{\"event\":\"gpu_canary_"));
|
|
assert!(!useful_output(&mut pending, b"sample\",\"ok\":true}\n"));
|
|
assert!(useful_output(
|
|
&mut pending,
|
|
b"{\"event\":\"mtp_tokens\",\"ids\":[42]}\n"
|
|
));
|
|
assert!(useful_output(&mut pending, b"progress"));
|
|
assert!(pending.is_empty());
|
|
}
|
|
}
|
|
}
|
|
|
|
fn main() {
|
|
#[cfg(target_os = "macos")]
|
|
if let Err(error) = supervisor::cli() {
|
|
eprintln!("Test supervisor: {error}");
|
|
std::process::exit(1);
|
|
}
|
|
#[cfg(not(target_os = "macos"))]
|
|
{
|
|
eprintln!("test-supervisor requires macOS process accounting");
|
|
std::process::exit(1);
|
|
}
|
|
}
|