Save inference parity implementation and evaluation harness
This commit is contained in:
Georg Bauer
@@ -0,0 +1,359 @@
|
||||
//! 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);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user