nuvec-distributed 0.6.0

//! The multi-process harness: a test re-runs its own test binary as the ranks of a ring.

#![allow(dead_code)]

use std::io::Read;
use std::net::TcpListener;
use std::process::{Command, Stdio};
use std::time::{Duration, Instant};

use nuvec_core::distributed::Group;

pub(crate) const CHILD: &str = "NUVEC_TEST_CHILD";

/// Runs `body` as `n` ranks of one ring. In the parent, starts the ranks and returns their
/// standard outputs (in rank order) once all succeed; in a rank, runs `body` and returns
/// `None`.
pub(crate) fn ranks(test: &str, n: usize, body: impl FnOnce(Group)) -> Option<Vec<String>> {
    ranks_with(test, n, &[], body)
}

/// [`ranks`] with `env` set in every rank's environment.
pub(crate) fn ranks_with(
    test: &str,
    n: usize,
    env: &[(&str, &str)],
    body: impl FnOnce(Group),
) -> Option<Vec<String>> {
    if let Ok(name) = std::env::var(CHILD) {
        if name == test {
            body(nuvec_distributed::init(true).expect("init"));
        }
        return None;
    }
    let ports: Vec<u16> = (0..n)
        .map(|_| TcpListener::bind("127.0.0.1:0").unwrap())
        .collect::<Vec<_>>()
        .iter()
        .map(|l| l.local_addr().unwrap().port())
        .collect();
    let hostfile =
        std::env::temp_dir().join(format!("nuvec-ring-{}-{test}.json", std::process::id()));
    let entries: Vec<String> = ports
        .iter()
        .map(|p| format!("[\"127.0.0.1:{p}\"]"))
        .collect();
    std::fs::write(&hostfile, format!("[{}]", entries.join(", "))).unwrap();
    let exe = std::env::current_exe().unwrap();
    let mut children: Vec<_> = (0..n)
        .map(|rank| {
            Command::new(&exe)
                .args([test, "--exact", "--nocapture", "--test-threads=1"])
                .env(CHILD, test)
                .env("MLX_HOSTFILE", &hostfile)
                .env("MLX_RANK", rank.to_string())
                .env("NUVEC_DISTRIBUTED_TIMEOUT", "20")
                .envs(env.iter().copied())
                .stdout(Stdio::piped())
                .stderr(Stdio::piped())
                .spawn()
                .unwrap()
        })
        .collect();
    let deadline = Instant::now() + Duration::from_secs(120);
    let mut outputs = Vec::new();
    for (rank, child) in children.iter_mut().enumerate() {
        let status = loop {
            if let Some(status) = child.try_wait().unwrap() {
                break Some(status);
            }
            if Instant::now() > deadline {
                break None;
            }
            std::thread::sleep(Duration::from_millis(10));
        };
        let mut out = String::new();
        let mut err = String::new();
        child
            .stdout
            .take()
            .unwrap()
            .read_to_string(&mut out)
            .unwrap();
        child
            .stderr
            .take()
            .unwrap()
            .read_to_string(&mut err)
            .unwrap();
        match status {
            Some(s) if s.success() => outputs.push(out),
            _ => {
                for c in &mut children {
                    let _ = c.kill();
                }
                panic!("{test}: rank {rank} failed ({status:?})\n{out}\n{err}");
            }
        }
    }
    let _ = std::fs::remove_file(&hostfile);
    Some(outputs)
}

/// The lines of each rank's output starting with `BITS`, which must agree.
pub(crate) fn agree(outputs: &[String]) {
    fn bits(o: &str) -> Vec<&str> {
        o.lines().filter(|l| l.starts_with("BITS")).collect()
    }
    let first = bits(&outputs[0]);
    assert!(!first.is_empty(), "no BITS lines");
    for (rank, o) in outputs.iter().enumerate() {
        assert_eq!(bits(o), first, "rank {rank} differs from rank 0");
    }
}

/// `n` free loopback addresses.
pub(crate) fn free_addresses(n: usize) -> Vec<String> {
    let listeners: Vec<TcpListener> = (0..n)
        .map(|_| TcpListener::bind("127.0.0.1:0").unwrap())
        .collect();
    listeners
        .iter()
        .map(|l| format!("127.0.0.1:{}", l.local_addr().unwrap().port()))
        .collect()
}