mpi-rs 0.1.0

//! The pure-Rust MPI process launcher, shared by the `mpiexec` and `mpirun`
//! binaries.
//!
//! [`run`] spawns N copies of a program, assigns each a rank, and runs a tiny
//! PMI-style rendezvous so every process learns the network address of every
//! peer before `mpi::initialize()` returns.
//!
//! Usage:
//!
//! ```text
//! mpiexec -n <N> <program> [args...]
//! mpiexec --np <N> <program> [args...]
//!
//! # multi-host (some ranks launched on remote hosts via ssh):
//! mpiexec -n <N> --host h1,h2,... --remote-bin /path/on/remote <local-program> [args...]
//! ```
//!
//! Hosts are assigned to ranks round-robin. A host of `localhost`/`127.0.0.1`
//! (or an empty host list) runs locally; any other host is launched with
//! `ssh <host> ...`. When any remote host is used the rendezvous binds a
//! routable address and children auto-detect their own LAN IP.

use std::io::{Read, Write};
use std::net::{SocketAddr, TcpListener, TcpStream, UdpSocket};
use std::process::{Command, Stdio};

fn usage() -> ! {
    eprintln!("usage: mpiexec -n <N> [--host h1,h2,...] [--remote-bin PATH] <program> [args...]");
    std::process::exit(2);
}

/// Read a single newline-terminated line directly off the stream (no buffering
/// beyond the line, so the stream can be reused for writing afterwards).
fn read_line(stream: &mut TcpStream) -> std::io::Result<String> {
    let mut buf = Vec::new();
    let mut byte = [0u8; 1];
    loop {
        let n = stream.read(&mut byte)?;
        if n == 0 || byte[0] == b'\n' {
            break;
        }
        buf.push(byte[0]);
    }
    Ok(String::from_utf8_lossy(&buf).into_owned())
}

/// Best-effort detection of this machine's primary outbound IP address.
fn detect_local_ip() -> Option<String> {
    let sock = UdpSocket::bind("0.0.0.0:0").ok()?;
    // Connecting a UDP socket sends no packets; it just selects the source
    // address the OS would use, i.e. our primary interface IP.
    sock.connect("8.8.8.8:80").ok()?;
    Some(sock.local_addr().ok()?.ip().to_string())
}

fn is_local_host(host: &str) -> bool {
    matches!(host, "" | "localhost" | "127.0.0.1" | "::1")
}

/// Entry point for the launcher binaries. Parses arguments, spawns the job
/// (locally and/or via ssh), runs the PMI rendezvous, and exits with the job's
/// aggregate exit code.
pub fn run() -> ! {
    let mut args = std::env::args().skip(1).peekable();

    let mut nprocs: Option<usize> = None;
    let mut hosts: Vec<String> = Vec::new();
    let mut remote_bin: Option<String> = None;

    while let Some(arg) = args.peek() {
        match arg.as_str() {
            "-n" | "-np" | "--n" | "--np" => {
                args.next();
                let v = args.next().unwrap_or_else(|| usage());
                nprocs = Some(v.parse().unwrap_or_else(|_| usage()));
            }
            "--host" | "--hosts" | "-host" | "-hosts" => {
                args.next();
                let v = args.next().unwrap_or_else(|| usage());
                hosts.extend(v.split(',').map(|s| s.trim().to_string()));
            }
            "--remote-bin" | "-remote-bin" => {
                args.next();
                remote_bin = Some(args.next().unwrap_or_else(|| usage()));
            }
            "-h" | "--help" => usage(),
            _ => break,
        }
    }

    let nprocs = nprocs.unwrap_or_else(|| usage());
    if nprocs == 0 {
        usage();
    }
    let program = args.next().unwrap_or_else(|| usage());
    let prog_args: Vec<String> = args.collect();

    let multihost = hosts.iter().any(|h| !is_local_host(h));

    // Bind the rendezvous. Multi-host jobs must be reachable from other nodes.
    let (bind_host, root_ip) = if multihost {
        let ip = detect_local_ip().unwrap_or_else(|| {
            eprintln!("mpiexec: could not determine local IP for multi-host launch");
            std::process::exit(1);
        });
        ("0.0.0.0", ip)
    } else {
        ("127.0.0.1", "127.0.0.1".to_string())
    };
    let listener = TcpListener::bind((bind_host, 0)).unwrap_or_else(|e| {
        eprintln!("mpiexec: failed to bind rendezvous socket: {e}");
        std::process::exit(1);
    });
    let root_port = listener.local_addr().unwrap().port();
    let root_addr = format!("{root_ip}:{root_port}");

    // Spawn the children (locally or via ssh).
    let mut children = Vec::with_capacity(nprocs);
    for rank in 0..nprocs {
        let host = if hosts.is_empty() {
            "localhost".to_string()
        } else {
            hosts[rank % hosts.len()].clone()
        };

        let spawn = if is_local_host(&host) {
            let mut cmd = Command::new(&program);
            cmd.args(&prog_args)
                .env("MPI_PMI_ROOT", &root_addr)
                .env("MPI_PMI_RANK", rank.to_string())
                .env("MPI_PMI_SIZE", nprocs.to_string())
                .env("MPI_JOBID", std::process::id().to_string())
                .stdin(Stdio::inherit())
                .stdout(Stdio::inherit())
                .stderr(Stdio::inherit());
            if multihost {
                cmd.env("MPI_MULTIHOST", "1");
            }
            cmd.spawn()
        } else {
            let bin = remote_bin.as_ref().unwrap_or_else(|| {
                eprintln!("mpiexec: --remote-bin is required for remote host '{host}'");
                std::process::exit(1);
            });
            // Build the remote command line with env vars prefixed.
            let dbg = if std::env::var("MPI_DEBUG").is_ok() {
                "MPI_DEBUG=1 "
            } else {
                ""
            };
            let jobid = std::process::id();
            let mut remote = format!(
                "{dbg}MPI_PMI_ROOT={root_addr} MPI_PMI_RANK={rank} MPI_PMI_SIZE={nprocs} MPI_JOBID={jobid} MPI_MULTIHOST=1 {bin}"
            );
            for a in &prog_args {
                remote.push(' ');
                remote.push_str(a);
            }
            Command::new("ssh")
                .arg("-o")
                .arg("BatchMode=yes")
                .arg(&host)
                .arg(&remote)
                .stdin(Stdio::inherit())
                .stdout(Stdio::inherit())
                .stderr(Stdio::inherit())
                .spawn()
        };

        match spawn {
            Ok(c) => children.push(c),
            Err(e) => {
                eprintln!("mpiexec: failed to launch rank {rank} on '{host}': {e}");
                for mut c in children {
                    let _ = c.kill();
                }
                std::process::exit(1);
            }
        }
    }

    // Rendezvous: accept a connection from every rank, learn its data-plane
    // address, then broadcast the full address table back.
    let mut conns: Vec<Option<TcpStream>> = (0..nprocs).map(|_| None).collect();
    let mut table: Vec<Option<SocketAddr>> = vec![None; nprocs];
    let mut endians: Vec<String> = vec!["le".to_string(); nprocs];

    for _ in 0..nprocs {
        let (mut stream, _peer) = listener.accept().unwrap_or_else(|e| {
            eprintln!("mpiexec: accept failed: {e}");
            std::process::exit(1);
        });
        let line = read_line(&mut stream).unwrap_or_else(|e| {
            eprintln!("mpiexec: PMI read failed: {e}");
            std::process::exit(1);
        });
        let mut parts = line.split_whitespace();
        let rank: usize = parts
            .next()
            .and_then(|s| s.parse().ok())
            .unwrap_or_else(|| {
                eprintln!("mpiexec: malformed PMI line: {line:?}");
                std::process::exit(1);
            });
        let addr: SocketAddr = parts
            .next()
            .and_then(|s| s.parse().ok())
            .unwrap_or_else(|| {
                eprintln!("mpiexec: malformed PMI address: {line:?}");
                std::process::exit(1);
            });
        if let Some(e) = parts.next() {
            endians[rank] = e.to_string();
        }
        table[rank] = Some(addr);
        conns[rank] = Some(stream);
    }

    // Serialize and broadcast the address table (rank addr endian).
    let mut blob = String::new();
    for (rank, addr) in table.iter().enumerate() {
        let addr = addr.expect("every rank should have reported an address");
        blob.push_str(&format!("{rank} {addr} {}\n", endians[rank]));
    }
    for conn in conns.iter_mut() {
        if let Some(stream) = conn.as_mut() {
            if let Err(e) = stream.write_all(blob.as_bytes()) {
                eprintln!("mpiexec: PMI write failed: {e}");
                std::process::exit(1);
            }
            let _ = stream.flush();
        }
    }
    drop(conns);
    drop(listener);

    // Wait for all children; propagate the first non-zero exit code.
    let mut exit_code = 0;
    for mut child in children {
        match child.wait() {
            Ok(status) => {
                if !status.success() && exit_code == 0 {
                    exit_code = status.code().unwrap_or(1);
                }
            }
            Err(e) => {
                eprintln!("mpiexec: error waiting for child: {e}");
                if exit_code == 0 {
                    exit_code = 1;
                }
            }
        }
    }
    std::process::exit(exit_code);
}