use std::net::{TcpListener, TcpStream};
use std::path::{Path, PathBuf};
use std::process::{Child, Command, Stdio};
use std::time::{Duration, Instant};
use crate::builtins::JoinArgs;
use crate::config::{self, DEFAULT_CONTROLLER_PORT, SshConfig, WorkerJoin, WorkerSource};
use crate::context::Context;
use crate::prepare::{self, DataSpec, Fail, PrepareSpec, Prepared, SourceSpec};
use crate::style;
const ENV_AGENT_JSON: &str = "FLODL_INTERNAL_AGENT_JSON";
pub const EXIT_PERMANENT: i32 = 2;
const TUNNEL_READY_BUDGET: Duration = Duration::from_secs(20);
const BACKOFF_MIN: Duration = Duration::from_secs(5);
const BACKOFF_MAX: Duration = Duration::from_secs(60);
const BACKOFF_RESET_AFTER: Duration = Duration::from_secs(120);
pub fn run(cli: &JoinArgs, bin_tail: Option<&[String]>) -> i32 {
let (block, project_root) = match load_join_block() {
Ok(pair) => pair,
Err(e) => {
crate::cli_error!("{e}");
return EXIT_PERMANENT;
}
};
let eff = match resolve_effective(
cli,
bin_tail,
block,
&crate::cluster::resolve_local_hostname(),
) {
Ok(eff) => eff,
Err(e) => {
crate::cli_error!("{e}");
return EXIT_PERMANENT;
}
};
if eff.controller_defaulted {
eprintln!(
"{}",
style::dim(&format!(
"fdl join: no controller configured; dialing \
127.0.0.1:{DEFAULT_CONTROLLER_PORT} (pass an address or set \
`join.controller` in fdl.yml)"
)),
);
}
if let BinSource::Given(path) = &eff.bin
&& !Path::new(path).is_file()
{
crate::cli_error!(
"training binary not found: {path} — build it first and \
point `--bin` (or fdl.yml `join.bin`) at it, or hand this \
box a `--source` to build",
);
return EXIT_PERMANENT;
}
let libtorch = resolve_local_libtorch(project_root.as_deref());
let mut sig_cache: Option<(u64, Option<String>)> = None;
let mut backoff = BACKOFF_MIN;
loop {
let started = Instant::now();
let outcome = match attempt(&eff, libtorch.as_ref(), &mut sig_cache) {
Ok(code) => {
if !eff.persist {
return code;
}
format!("agent exited with code {code}")
}
Err(fail) => {
crate::cli_error!("{}", fail.message());
if fail.is_permanent() {
eprintln!(
"{}",
style::dim(&format!(
"fdl join: not re-dialing — retrying cannot \
fix this (exit {EXIT_PERMANENT})"
)),
);
return EXIT_PERMANENT;
}
if !eff.persist {
return 1;
}
"attempt failed".to_string()
}
};
if started.elapsed() > BACKOFF_RESET_AFTER {
backoff = BACKOFF_MIN;
}
eprintln!(
"fdl join: {outcome} after {}s; re-dialing in {}s (--persist)",
started.elapsed().as_secs(),
backoff.as_secs(),
);
std::thread::sleep(backoff);
backoff = (backoff * 2).min(BACKOFF_MAX);
}
}
#[derive(Debug)]
struct Effective {
controller_host: String,
controller_port: u16,
controller_defaulted: bool,
ssh: Option<SshConfig>,
token: Option<String>,
bin: BinSource,
libtorch_spec: Option<String>,
host: String,
devices: Option<Vec<u8>>,
persist: bool,
bin_args: Vec<String>,
data_path: Option<String>,
data_source: Option<String>,
gpu_ram_share: Option<f64>,
sig_probe: bool,
}
#[derive(Debug, PartialEq, Eq)]
enum BinSource {
Given(String),
Build(WorkerSource),
}
impl Effective {
fn prepare_spec<'a>(
&'a self,
active_libtorch: Option<&'a (PathBuf, String)>,
) -> PrepareSpec<'a> {
PrepareSpec {
data: DataSpec {
path: self.data_path.as_deref(),
source: self.data_source.as_deref(),
ssh: self.ssh.as_ref(),
},
libtorch: self.libtorch_spec.as_deref(),
active_libtorch,
devices: self.devices.as_deref(),
source: match &self.bin {
BinSource::Given(_) => None,
BinSource::Build(s) => Some(SourceSpec {
from: &s.from,
cwd: s.cwd.as_deref(),
build: s.build.as_deref(),
bin: s.bin.as_deref(),
ssh: self.ssh.as_ref(),
}),
},
}
}
}
fn resolve_effective(
cli: &JoinArgs,
bin_tail: Option<&[String]>,
block: Option<WorkerJoin>,
local_hostname: &str,
) -> Result<Effective, String> {
let block = block.unwrap_or_default();
if cli.identity.is_some() && cli.ssh.is_none() && block.ssh.is_none() {
return Err("--identity is the tunnel's key file — it needs an ssh hop \
(`--ssh` or fdl.yml `join.ssh`)"
.to_string());
}
let ssh = match (&cli.ssh, block.ssh) {
(Some(spec), b) => {
let mut cfg = parse_ssh_spec(spec)?;
if let Some(b) = b {
cfg.identity_file = b.identity_file;
cfg.options = b.options;
}
Some(cfg)
}
(None, Some(b)) => {
if b.target.is_none() {
return Err("fdl.yml join.ssh needs a `target:` (the tunnel host)".to_string());
}
Some(b)
}
(None, None) => None,
};
let mut ssh = ssh;
if let (Some(cfg), Some(id)) = (ssh.as_mut(), &cli.identity) {
cfg.identity_file = Some(id.clone());
}
let named = cli.controller.as_ref().or(block.controller.as_ref());
let controller_defaulted = named.is_none() && ssh.is_none();
let (controller_host, controller_port) = match named {
Some(spec) => parse_host_port(spec)?,
None => ("127.0.0.1".to_string(), DEFAULT_CONTROLLER_PORT),
};
let bin_path = cli.bin.clone().or(block.bin);
let source = match (cli.source.clone(), block.source) {
(Some(from), b) => Some(WorkerSource {
from,
cwd: cli
.source_cwd
.clone()
.or_else(|| b.as_ref().and_then(|b| b.cwd.clone())),
build: cli
.source_build
.clone()
.or_else(|| b.as_ref().and_then(|b| b.build.clone())),
bin: cli
.source_bin
.clone()
.or_else(|| b.as_ref().and_then(|b| b.bin.clone())),
}),
(None, Some(mut b)) => {
if let Some(cwd) = cli.source_cwd.clone() {
b.cwd = Some(cwd);
}
if let Some(build) = cli.source_build.clone() {
b.build = Some(build);
}
if let Some(bin) = cli.source_bin.clone() {
b.bin = Some(bin);
}
Some(b)
}
(None, None) => None,
};
if source.is_none() {
for (flag, set) in [
("--source-cwd", cli.source_cwd.is_some()),
("--source-build", cli.source_build.is_some()),
("--source-bin", cli.source_bin.is_some()),
] {
if set {
return Err(format!(
"{flag} has no source to apply to — pass `--source \
<spec>` too, or set `join.source` in fdl.yml"
));
}
}
}
let bin = match (bin_path, source) {
(Some(_), Some(_)) => {
return Err("`bin:` and `source:` both name this box's training binary \
— keep the one you mean. `bin:` runs a binary as given; \
`source:` fetches and builds one here"
.to_string());
}
(Some(path), None) => BinSource::Given(path),
(None, Some(source)) => BinSource::Build(source),
(None, None) => {
return Err("no training binary configured — pass `--bin <path>` (run \
it as given) or `--source <spec>` (build it here), or set \
`join.bin` / `join.source` in fdl.yml. The binary is the \
protocol: it dials, joins, and runs this host's ranks"
.to_string());
}
};
let devices = match &cli.devices {
Some(spec) => parse_devices(spec)?,
None => block.devices,
};
let bin_args = match bin_tail {
Some(tail) => tail.to_vec(),
None => block.args,
};
Ok(Effective {
controller_host,
controller_port,
controller_defaulted,
ssh,
token: cli.token.clone().or(block.token),
bin,
host: cli
.host
.clone()
.or(block.host)
.unwrap_or_else(|| local_hostname.to_string()),
devices,
persist: cli.persist || block.persist,
bin_args,
libtorch_spec: cli.libtorch.clone().or(block.libtorch),
data_path: cli.data_path.clone().or(block.data_path),
data_source: cli.data_source.clone().or(block.data_source),
gpu_ram_share: cli.gpu_ram_share.or(block.gpu_ram_share),
sig_probe: !cli.no_sig_probe && block.sig_probe.unwrap_or(true),
})
}
fn load_join_block() -> Result<(Option<WorkerJoin>, Option<PathBuf>), String> {
let cwd =
std::env::current_dir().map_err(|e| format!("cannot read the current directory: {e}"))?;
let Some(config_path) = config::find_project_config(&cwd) else {
return Ok((None, None));
};
let env_name = std::env::var("FDL_ENV")
.ok()
.filter(|s| !s.trim().is_empty());
let project = config::load_project_with_env(&config_path, env_name.as_deref())
.map_err(|e| format!("cannot load {}: {e}", config_path.display()))?;
let root = config_path.parent().map(Path::to_path_buf);
Ok((project.join, root))
}
fn parse_host_port(spec: &str) -> Result<(String, u16), String> {
match spec.rsplit_once(':') {
Some((host, port)) => {
let port = port.parse::<u16>().map_err(|_| {
format!("invalid controller address `{spec}` — expected host[:port]")
})?;
if host.is_empty() {
return Err(format!(
"invalid controller address `{spec}` — expected host[:port]"
));
}
Ok((host.to_string(), port))
}
None => Ok((spec.to_string(), DEFAULT_CONTROLLER_PORT)),
}
}
fn parse_ssh_spec(spec: &str) -> Result<SshConfig, String> {
let (user, rest) = match spec.split_once('@') {
Some((u, r)) if !u.is_empty() => (Some(u.to_string()), r),
Some(_) => {
return Err(format!("invalid --ssh `{spec}` — empty user before `@`"));
}
None => (None, spec),
};
let (host, port) = match rest.rsplit_once(':') {
Some((h, p)) => {
let port = p
.parse::<u16>()
.map_err(|_| format!("invalid --ssh `{spec}` — expected [user@]host[:port]"))?;
(h, Some(port))
}
None => (rest, None),
};
if host.is_empty() {
return Err(format!(
"invalid --ssh `{spec}` — expected [user@]host[:port]"
));
}
Ok(SshConfig {
target: Some(host.to_string()),
port,
user,
identity_file: None,
options: Vec::new(),
})
}
fn parse_devices(spec: &str) -> Result<Option<Vec<u8>>, String> {
if spec.trim().eq_ignore_ascii_case("all") {
return Ok(None);
}
spec.split(',')
.map(|s| {
s.trim()
.parse::<u8>()
.map_err(|_| format!("invalid --devices `{spec}` — expected e.g. `0,1` or `all`"))
})
.collect::<Result<Vec<u8>, String>>()
.map(Some)
}
fn agent_spec_hex(
eff: &Effective,
dial: (&str, u16),
libtorch_label: &str,
prepared: &Prepared,
model_sig_hex: Option<&str>,
) -> String {
let mut spec = serde_json::json!({
"host": eff.host,
"controller_host": dial.0,
"controller_port": dial.1,
"libtorch": libtorch_label,
});
if let Some(token) = &eff.token {
spec["salt_hex"] = serde_json::json!(token);
}
if let Some(devices) = &eff.devices {
spec["local_devices"] = serde_json::json!(devices);
}
if let Some(data) = &prepared.data_path {
spec["data_path"] = serde_json::json!(data.display().to_string());
}
if let Some(run) = &prepared.run_id {
spec["run_id"] = serde_json::json!(run);
}
if let Some(share) = eff.gpu_ram_share {
spec["gpu_ram_share"] = serde_json::json!(share);
}
if let Some(sig) = model_sig_hex {
spec["model_sig_hex"] = serde_json::json!(sig);
}
hex_encode(spec.to_string().as_bytes())
}
fn hex_encode(bytes: &[u8]) -> String {
let mut s = String::with_capacity(bytes.len() * 2);
for b in bytes {
s.push_str(&format!("{b:02x}"));
}
s
}
fn resolve_local_libtorch(project_root: Option<&Path>) -> Option<(PathBuf, String)> {
let root = match project_root {
Some(r) => r.to_path_buf(),
None => Context::resolve().root,
};
crate::libtorch::detect::active_variant(&root)
}
fn child_ld_library_path(libtorch_dir: &Path, variant: &str) -> String {
let lib = libtorch_dir.join("lib").display().to_string();
let vendor = crate::libtorch::detect::variant_vendor(variant);
let value = crate::libtorch::detect::ld_library_path_value(
vendor,
&lib,
&crate::libtorch::detect::local_rocm_lib_dir(),
);
match std::env::var("LD_LIBRARY_PATH") {
Ok(cur) if !cur.is_empty() => format!("{value}:{cur}"),
_ => value,
}
}
fn attempt(
eff: &Effective,
active_libtorch: Option<&(PathBuf, String)>,
sig_cache: &mut Option<(u64, Option<String>)>,
) -> Result<i32, Fail> {
let mut notes = Vec::new();
let prepared = prepare::prepare(&eff.prepare_spec(active_libtorch), &mut notes);
prepare::print_notes("join", ¬es);
let prepared = prepared?;
let (bin, bin_cwd) = match (&eff.bin, &prepared.bin) {
(BinSource::Build(_), Some(built)) => (built.bin.clone(), Some(built.cwd.clone())),
(BinSource::Given(path), None) => (PathBuf::from(path), None),
(kind, built) => {
return Err(Fail::Permanent(format!(
"internal: preparation and the resolved binary disagree \
({}, built={})",
match kind {
BinSource::Given(_) => "a path was given",
BinSource::Build(_) => "a source was given",
},
built.is_some(),
)));
}
};
let args: &[String] = match &prepared.args {
Some(published) => {
if !eff.bin_args.is_empty() && published != &eff.bin_args {
eprintln!(
"{}",
style::dim(&format!(
"fdl join: the published run's arguments replace \
this box's ({} -> {})",
eff.bin_args.join(" "),
published.join(" "),
)),
);
}
published
}
None => &eff.bin_args,
};
let model_sig_hex = if eff.sig_probe {
match probe_recipe_digest(&bin, args) {
Some(digest) => match sig_cache {
Some((key, cached)) if *key == digest => cached.clone(),
_ => {
let sig =
model_sig_probe(&bin, bin_cwd.as_deref(), args, prepared.libtorch.as_ref());
*sig_cache = Some((digest, sig.clone()));
sig
}
},
None => model_sig_probe(&bin, bin_cwd.as_deref(), args, prepared.libtorch.as_ref()),
}
} else {
None
};
let mut tunnel: Option<Child> = None;
let dial: (String, u16) = match &eff.ssh {
Some(ssh) => {
let local_port = pick_local_port().map_err(Fail::Transient)?;
let argv =
build_tunnel_argv(ssh, local_port, &eff.controller_host, eff.controller_port);
eprintln!(
"fdl join: opening tunnel {} -> {}:{} (local port {local_port})",
ssh.target.as_deref().unwrap_or("?"),
eff.controller_host,
eff.controller_port,
);
let mut child = Command::new(&argv[0])
.args(&argv[1..])
.stdin(Stdio::null())
.spawn()
.map_err(|e| {
Fail::Permanent(format!("spawn ssh tunnel: {e}"))
})?;
if let Err(e) = wait_tunnel_ready(&mut child, local_port) {
let _ = child.kill();
let _ = child.wait();
return Err(Fail::Transient(e));
}
tunnel = Some(child);
("127.0.0.1".to_string(), local_port)
}
None => (eff.controller_host.clone(), eff.controller_port),
};
let libtorch_label = prepared
.libtorch
.as_ref()
.map(|(_, l)| l.as_str())
.unwrap_or("");
let spec_hex = agent_spec_hex(
eff,
(&dial.0, dial.1),
libtorch_label,
&prepared,
model_sig_hex.as_deref(),
);
let mut cmd = Command::new(&bin);
cmd.args(args)
.env(ENV_AGENT_JSON, &spec_hex)
.env(crate::cluster::ENV_HOST_OVERRIDE, &eff.host)
.stdin(Stdio::null());
if let Some(cwd) = &bin_cwd {
cmd.current_dir(cwd);
}
if let Some((dir, variant)) = &prepared.libtorch {
cmd.env("LD_LIBRARY_PATH", child_ld_library_path(dir, variant));
}
let status = cmd
.status()
.map_err(|e| Fail::Permanent(format!("run {}: {e}", bin.display())));
if let Some(mut t) = tunnel.take() {
let _ = t.kill();
let _ = t.wait();
}
Ok(status?.code().unwrap_or(1))
}
fn probe_recipe_digest(bin: &Path, args: &[String]) -> Option<u64> {
use std::hash::{Hash, Hasher};
let meta = std::fs::metadata(bin).ok()?;
let mut h = std::collections::hash_map::DefaultHasher::new();
bin.hash(&mut h);
meta.len().hash(&mut h);
meta.modified()
.ok()?
.duration_since(std::time::UNIX_EPOCH)
.ok()?
.as_nanos()
.hash(&mut h);
args.hash(&mut h);
Some(h.finish())
}
const PROBE_TAIL_LINES: usize = 8;
const ENV_MODEL_SIG_PROBE: &str = "FLODL_INTERNAL_MODEL_SIG_PROBE";
const MODEL_SIG_LINE: &str = "flodl-model-sig: ";
const MODEL_SIG_PROBE_TIMEOUT: Duration = Duration::from_secs(120);
fn model_sig_probe(
bin: &Path,
cwd: Option<&Path>,
args: &[String],
libtorch: Option<&(PathBuf, String)>,
) -> Option<String> {
eprintln!(
"{}",
style::dim(
"fdl join: probing the binary for its model signature \
(--no-sig-probe skips this)"
),
);
let mut cmd = Command::new(bin);
cmd.args(args)
.env(ENV_MODEL_SIG_PROBE, "1")
.stdin(Stdio::null())
.stdout(Stdio::piped());
if let Some(dir) = cwd {
cmd.current_dir(dir);
}
if let Some((dir, variant)) = libtorch {
cmd.env("LD_LIBRARY_PATH", child_ld_library_path(dir, variant));
}
let mut child = match cmd.spawn() {
Ok(c) => c,
Err(e) => {
eprintln!(
"fdl join: model-sig probe could not run {}: {e}; joining \
without a signature",
bin.display(),
);
return None;
}
};
let stdout = child.stdout.take().expect("stdout was piped");
let reader = std::thread::spawn(move || {
use std::io::{BufRead, BufReader};
let mut sig = None;
let mut tail: std::collections::VecDeque<String> = std::collections::VecDeque::new();
for line in BufReader::new(stdout).lines() {
let Ok(line) = line else { break };
if let Some(rest) = line.strip_prefix(MODEL_SIG_LINE) {
sig = Some(rest.trim().to_string());
}
if tail.len() == PROBE_TAIL_LINES {
tail.pop_front();
}
tail.push_back(line);
}
(sig, tail)
});
let deadline = Instant::now() + MODEL_SIG_PROBE_TIMEOUT;
let status = loop {
match child.try_wait() {
Ok(Some(st)) => break Some(st),
Ok(None) if Instant::now() >= deadline => {
let _ = child.kill();
let _ = child.wait();
break None;
}
Ok(None) => std::thread::sleep(Duration::from_millis(50)),
Err(_) => {
let _ = child.kill();
let _ = child.wait();
break None;
}
}
};
let (sig, tail) = reader.join().unwrap_or_default();
let sig = sig.filter(|s| s.len() == 64 && s.bytes().all(|b| b.is_ascii_hexdigit()));
match (&status, &sig) {
(Some(st), Some(_)) if st.success() => sig,
(None, _) => {
eprintln!(
"fdl join: model-sig probe killed after {}s — a binary built \
against a flodl that predates the probe runs its whole main \
here; joining without a signature (`--no-sig-probe` or \
`join.sig_probe: false` silences this)",
MODEL_SIG_PROBE_TIMEOUT.as_secs(),
);
None
}
(Some(st), _) if !st.success() => {
eprintln!(
"fdl join: WARNING: the training binary exited with {} under \
the model-sig probe — rank children re-enter it with the \
same arguments after admission, so if this failure is real \
it takes the cohort's formation with it. Check: {} {}",
st.code().map_or("a signal".to_string(), |c| c.to_string()),
bin.display(),
args.join(" "),
);
None
}
_ => {
eprintln!(
"fdl join: model-sig probe exited 0 without printing a \
signature; joining without one — the formation-time check \
still applies. Either the binary predates the probe, or it \
failed before reaching the trainer (check its output above, \
and that this box can write wherever it writes).",
);
if !tail.is_empty() {
eprintln!("fdl join: last lines of the probe's output:");
for line in &tail {
eprintln!(" {line}");
}
}
None
}
}
}
fn pick_local_port() -> Result<u16, String> {
let listener =
TcpListener::bind("127.0.0.1:0").map_err(|e| format!("reserve local tunnel port: {e}"))?;
let port = listener
.local_addr()
.map_err(|e| format!("reserve local tunnel port: {e}"))?
.port();
Ok(port)
}
fn build_tunnel_argv(
ssh: &SshConfig,
local_port: u16,
controller_host: &str,
controller_port: u16,
) -> Vec<String> {
let mut argv: Vec<String> = vec!["ssh".into(), "-N".into(), "-T".into()];
if let Some(warning) = crate::cluster::batchmode_override_warning(
&ssh.options,
ssh.target.as_deref().unwrap_or("?"),
) {
eprintln!("{warning}");
}
for opt in &ssh.options {
argv.push("-o".into());
argv.push(opt.clone());
}
if let Some(port) = ssh.port {
argv.push("-p".into());
argv.push(port.to_string());
}
if let Some(user) = ssh.user.as_deref() {
argv.push("-l".into());
argv.push(user.to_string());
}
if let Some(id) = ssh.identity_file.as_deref() {
argv.push("-i".into());
argv.push(id.to_string());
}
argv.push("-o".into());
argv.push("BatchMode=yes".into());
argv.push("-o".into());
argv.push("ExitOnForwardFailure=yes".into());
argv.push("-o".into());
argv.push("ServerAliveInterval=30".into());
argv.push("-L".into());
argv.push(format!(
"127.0.0.1:{local_port}:{controller_host}:{controller_port}"
));
argv.push(ssh.target.clone().unwrap_or_default());
argv
}
fn wait_tunnel_ready(child: &mut Child, local_port: u16) -> Result<(), String> {
let deadline = Instant::now() + TUNNEL_READY_BUDGET;
let addr = std::net::SocketAddr::from(([127, 0, 0, 1], local_port));
loop {
if let Ok(Some(status)) = child.try_wait() {
return Err(format!(
"ssh tunnel exited ({status}) before the forward came up — \
see its output above (auth failure, or the remote refused \
the forward)"
));
}
if let Ok(probe) = TcpStream::connect_timeout(&addr, Duration::from_millis(500)) {
drop(probe);
return Ok(());
}
if Instant::now() >= deadline {
return Err(format!(
"ssh tunnel did not come up within {}s (local port \
{local_port} never accepted)",
TUNNEL_READY_BUDGET.as_secs(),
));
}
std::thread::sleep(Duration::from_millis(200));
}
}
#[cfg(test)]
mod tests {
use super::*;
fn no_flags() -> JoinArgs {
JoinArgs {
controller: None,
ssh: None,
identity: None,
token: None,
bin: None,
source: None,
source_cwd: None,
source_build: None,
source_bin: None,
libtorch: None,
host: None,
devices: None,
persist: false,
data_path: None,
data_source: None,
gpu_ram_share: None,
no_sig_probe: false,
}
}
fn full_block() -> WorkerJoin {
WorkerJoin {
controller: Some("10.0.0.9:9000".into()),
ssh: Some(SshConfig {
target: Some("bastion".into()),
port: Some(2222),
user: Some("join-user".into()),
identity_file: Some("/etc/flodl/join_key".into()),
options: vec!["StrictHostKeyChecking=accept-new".into()],
}),
token: Some("aa".repeat(16)),
bin: Some("target/release/train".into()),
source: None,
libtorch: Some("auto".into()),
host: Some("worker-7".into()),
devices: Some(vec![0, 1]),
persist: true,
args: vec!["--model".into(), "lenet".into()],
data_path: Some("/flodl/data".into()),
data_source: Some("sshfs://flodl@ctrl:/srv/data".into()),
gpu_ram_share: Some(0.5),
sig_probe: None,
}
}
fn source_block() -> WorkerJoin {
WorkerJoin {
source: Some(WorkerSource {
from: "rsync://exa:/home/op/rdl".into(),
cwd: Some("ddp-bench".into()),
build: Some("cargo build --release --bin ddp-bench".into()),
bin: Some("target/release/ddp-bench".into()),
}),
bin: None,
..full_block()
}
}
#[test]
fn flags_win_over_the_config_block() {
let cli = JoinArgs {
controller: Some("exa".into()),
ssh: Some("op@front:22".into()),
identity: Some("/tmp/id".into()),
token: Some("bb".repeat(16)),
bin: Some("other/bin".into()),
libtorch: Some("cu128".into()),
host: Some("pascal".into()),
devices: Some("2".into()),
persist: false,
data_path: Some("/mnt/corpus".into()),
data_source: Some("sshfs://exa/mnt/corpus".into()),
..no_flags()
};
let tail: Vec<String> = vec!["--epochs".into(), "3".into()];
let eff = resolve_effective(&cli, Some(&tail), Some(full_block()), "localbox").unwrap();
assert_eq!(eff.controller_host, "exa");
assert_eq!(eff.controller_port, DEFAULT_CONTROLLER_PORT);
assert!(!eff.controller_defaulted);
let ssh = eff.ssh.as_ref().unwrap();
assert_eq!(ssh.target.as_deref(), Some("front"));
assert_eq!(ssh.user.as_deref(), Some("op"));
assert_eq!(ssh.port, Some(22));
assert_eq!(ssh.identity_file.as_deref(), Some("/tmp/id"));
assert_eq!(
ssh.options,
vec!["StrictHostKeyChecking=accept-new".to_string()]
);
assert_eq!(eff.token.as_deref(), Some("bb".repeat(16).as_str()));
assert_eq!(eff.bin, BinSource::Given("other/bin".into()));
assert_eq!(eff.libtorch_spec.as_deref(), Some("cu128"));
assert_eq!(eff.host, "pascal");
assert_eq!(eff.devices, Some(vec![2]));
assert!(eff.persist);
assert_eq!(eff.bin_args, vec!["--epochs".to_string(), "3".into()]);
assert_eq!(eff.data_path.as_deref(), Some("/mnt/corpus"));
assert_eq!(eff.data_source.as_deref(), Some("sshfs://exa/mnt/corpus"));
}
#[test]
fn block_fills_everything_the_flags_left_unset() {
let eff = resolve_effective(&no_flags(), None, Some(full_block()), "localbox").unwrap();
assert_eq!(eff.controller_host, "10.0.0.9");
assert_eq!(eff.controller_port, 9000);
let ssh = eff.ssh.as_ref().unwrap();
assert_eq!(ssh.target.as_deref(), Some("bastion"));
assert_eq!(ssh.identity_file.as_deref(), Some("/etc/flodl/join_key"));
assert_eq!(eff.bin, BinSource::Given("target/release/train".into()));
assert_eq!(eff.libtorch_spec.as_deref(), Some("auto"));
assert_eq!(eff.host, "worker-7");
assert_eq!(eff.devices, Some(vec![0, 1]));
assert!(eff.persist);
assert_eq!(eff.bin_args, vec!["--model".to_string(), "lenet".into()]);
assert_eq!(eff.data_path.as_deref(), Some("/flodl/data"));
assert_eq!(
eff.data_source.as_deref(),
Some("sshfs://flodl@ctrl:/srv/data"),
);
let spec = eff.prepare_spec(None);
assert_eq!(
spec.data.ssh.and_then(|s| s.identity_file.as_deref()),
Some("/etc/flodl/join_key"),
);
}
#[test]
fn a_source_block_becomes_a_source_spec_carrying_the_same_key() {
let eff = resolve_effective(&no_flags(), None, Some(source_block()), "localbox").unwrap();
let spec = eff.prepare_spec(None);
let source = spec.source.expect("a source block yields a source spec");
assert_eq!(source.from, "rsync://exa:/home/op/rdl");
assert_eq!(source.cwd, Some("ddp-bench"));
assert_eq!(source.bin, Some("target/release/ddp-bench"));
assert_eq!(
source.ssh.and_then(|s| s.identity_file.as_deref()),
Some("/etc/flodl/join_key"),
);
}
#[test]
fn naming_both_a_binary_and_a_source_is_a_loud_error() {
let block = WorkerJoin {
bin: Some("target/release/train".into()),
..source_block()
};
let err = resolve_effective(&no_flags(), None, Some(block), "x").unwrap_err();
assert!(err.contains("both name"), "got: {err}");
}
#[test]
fn a_source_flag_keeps_the_blocks_other_source_fields() {
let cli = JoinArgs {
source: Some("file:///mnt/rdl".into()),
..no_flags()
};
let eff = resolve_effective(&cli, None, Some(source_block()), "x").unwrap();
assert_eq!(
eff.bin,
BinSource::Build(WorkerSource {
from: "file:///mnt/rdl".into(),
cwd: Some("ddp-bench".into()),
build: Some("cargo build --release --bin ddp-bench".into()),
bin: Some("target/release/ddp-bench".into()),
}),
);
}
#[test]
fn a_source_with_no_artifact_is_legal_because_a_manifest_may_name_it() {
let cli = JoinArgs {
source: Some("file:///mnt/rdl".into()),
..no_flags()
};
let eff = resolve_effective(&cli, None, None, "x").unwrap();
assert_eq!(
eff.bin,
BinSource::Build(WorkerSource {
from: "file:///mnt/rdl".into(),
cwd: None,
build: None,
bin: None,
}),
);
}
#[test]
fn a_source_detail_flag_with_no_source_is_a_loud_error() {
let cli = JoinArgs {
bin: Some("t/bin".into()),
source_cwd: Some("ddp-bench".into()),
..no_flags()
};
let err = resolve_effective(&cli, None, None, "x").unwrap_err();
assert!(err.contains("--source-cwd"), "got: {err}");
assert!(err.contains("no source"), "got: {err}");
}
#[test]
fn defaults_are_loopback_hostname_and_all_devices() {
let cli = JoinArgs {
bin: Some("t/bin".into()),
..no_flags()
};
let eff = resolve_effective(&cli, None, None, "localbox").unwrap();
assert_eq!(eff.controller_host, "127.0.0.1");
assert_eq!(eff.controller_port, DEFAULT_CONTROLLER_PORT);
assert!(eff.controller_defaulted);
assert!(eff.ssh.is_none());
assert!(eff.token.is_none());
assert_eq!(eff.host, "localbox");
assert_eq!(eff.devices, None);
assert!(!eff.persist);
assert!(eff.bin_args.is_empty());
assert!(eff.data_path.is_none());
assert!(eff.data_source.is_none());
}
#[test]
fn an_explicit_empty_tail_clears_the_block_args() {
let eff =
resolve_effective(&no_flags(), Some(&[]), Some(full_block()), "localbox").unwrap();
assert!(eff.bin_args.is_empty());
}
#[test]
fn identity_without_an_ssh_hop_is_a_loud_error() {
let cli = JoinArgs {
identity: Some("/tmp/id".into()),
bin: Some("t/bin".into()),
..no_flags()
};
let err = resolve_effective(&cli, None, None, "x").unwrap_err();
assert!(err.contains("ssh hop"), "got: {err}");
}
#[test]
fn missing_bin_is_a_loud_error() {
let err = resolve_effective(&no_flags(), None, None, "x").unwrap_err();
assert!(err.contains("--bin"), "got: {err}");
assert!(err.contains("join.bin"), "got: {err}");
}
#[test]
fn ssh_implies_the_loopback_controller_without_a_note() {
let cli = JoinArgs {
ssh: Some("join@ctrl".into()),
bin: Some("t/bin".into()),
..no_flags()
};
let eff = resolve_effective(&cli, None, None, "x").unwrap();
assert_eq!(eff.controller_host, "127.0.0.1");
assert_eq!(eff.controller_port, DEFAULT_CONTROLLER_PORT);
assert!(
!eff.controller_defaulted,
"tunnel loopback is the convention"
);
}
#[test]
fn block_ssh_without_target_is_a_loud_error() {
let block = WorkerJoin {
ssh: Some(SshConfig::default()),
bin: Some("t/bin".into()),
..WorkerJoin::default()
};
let err = resolve_effective(&no_flags(), None, Some(block), "x").unwrap_err();
assert!(err.contains("target"), "got: {err}");
}
#[test]
fn spec_parsers_cover_their_shapes() {
assert_eq!(
parse_host_port("exa").unwrap(),
("exa".to_string(), DEFAULT_CONTROLLER_PORT),
);
assert_eq!(
parse_host_port("exa:9000").unwrap(),
("exa".to_string(), 9000)
);
assert!(parse_host_port(":9000").is_err());
assert!(parse_host_port("exa:banana").is_err());
let ssh = parse_ssh_spec("join@ctrl:2222").unwrap();
assert_eq!(ssh.target.as_deref(), Some("ctrl"));
assert_eq!(ssh.user.as_deref(), Some("join"));
assert_eq!(ssh.port, Some(2222));
let bare = parse_ssh_spec("ctrl").unwrap();
assert_eq!(bare.target.as_deref(), Some("ctrl"));
assert_eq!(bare.user, None);
assert_eq!(bare.port, None);
assert!(parse_ssh_spec("@ctrl").is_err());
assert!(parse_ssh_spec("join@").is_err());
assert!(parse_ssh_spec("ctrl:pear").is_err());
assert_eq!(parse_devices("0,1").unwrap(), Some(vec![0, 1]));
assert_eq!(parse_devices(" 2 ").unwrap(), Some(vec![2]));
assert_eq!(parse_devices("all").unwrap(), None);
assert!(parse_devices("0,x").is_err());
}
#[test]
fn probe_recipe_digest_binds_binary_identity_and_args() {
let dir = std::env::temp_dir().join(format!("fdl-sig-digest-test-{}", std::process::id()));
std::fs::create_dir_all(&dir).unwrap();
let bin = dir.join("train");
std::fs::write(&bin, b"v1").unwrap();
let args = vec!["--model".to_string(), "lenet".to_string()];
let base = probe_recipe_digest(&bin, &args).unwrap();
assert_eq!(probe_recipe_digest(&bin, &args).unwrap(), base);
assert_ne!(
probe_recipe_digest(&bin, &["--model".to_string(), "resnet".to_string()]).unwrap(),
base,
"args are part of the recipe (a re-publish must re-probe)"
);
std::fs::write(&bin, b"v2 longer").unwrap();
assert_ne!(
probe_recipe_digest(&bin, &args).unwrap(),
base,
"a rebuilt binary must re-probe"
);
assert_eq!(probe_recipe_digest(&dir.join("absent"), &args), None);
let _ = std::fs::remove_dir_all(&dir);
}
#[cfg(unix)]
#[test]
fn model_sig_probe_parses_the_line_and_absorbs_failures() {
let sh = PathBuf::from("/bin/sh");
let run = |body: String| model_sig_probe(&sh, None, &["-c".to_string(), body], None);
let sig = "ab".repeat(32);
assert_eq!(
run(format!("echo main noise; echo '{MODEL_SIG_LINE}{sig}'")),
Some(sig),
);
assert_eq!(run("exit 0".to_string()), None);
assert_eq!(run("exit 3".to_string()), None);
assert_eq!(run(format!("echo '{MODEL_SIG_LINE}not-hex-at-all'")), None,);
}
#[test]
fn agent_spec_shape_is_the_wire_contract() {
let cli = JoinArgs {
token: Some("ab".repeat(16)),
bin: Some("t/bin".into()),
host: Some("pascal".into()),
devices: Some("0,1".into()),
gpu_ram_share: Some(0.5),
..no_flags()
};
let eff = resolve_effective(&cli, None, None, "x").unwrap();
let prepared = Prepared {
data_path: Some(PathBuf::from("/flodl/data")),
run_id: Some("a1b2c3d4e5f60718".to_string()),
..Prepared::default()
};
let hex = agent_spec_hex(
&eff,
("127.0.0.1", 40123),
"builds/sm61-sm120",
&prepared,
Some(&"cd".repeat(32)),
);
let bytes: Vec<u8> = (0..hex.len())
.step_by(2)
.map(|i| u8::from_str_radix(&hex[i..i + 2], 16).unwrap())
.collect();
let spec: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
assert_eq!(spec["host"], "pascal");
assert_eq!(spec["controller_host"], "127.0.0.1");
assert_eq!(spec["controller_port"], 40123);
assert_eq!(spec["salt_hex"], "ab".repeat(16));
assert_eq!(spec["local_devices"], serde_json::json!([0, 1]));
assert_eq!(spec["libtorch"], "builds/sm61-sm120");
assert_eq!(spec["data_path"], "/flodl/data");
assert_eq!(spec["run_id"], "a1b2c3d4e5f60718");
assert_eq!(spec["gpu_ram_share"], 0.5);
assert_eq!(spec["model_sig_hex"], "cd".repeat(32));
let open = {
let cli = JoinArgs {
bin: Some("t/bin".into()),
..no_flags()
};
let eff = resolve_effective(&cli, None, None, "cloud-1").unwrap();
agent_spec_hex(&eff, ("10.0.0.1", 1337), "", &Prepared::default(), None)
};
let bytes: Vec<u8> = (0..open.len())
.step_by(2)
.map(|i| u8::from_str_radix(&open[i..i + 2], 16).unwrap())
.collect();
let spec: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
assert!(spec.get("salt_hex").is_none());
assert!(spec.get("local_devices").is_none());
assert!(spec.get("dataset_sig_hex").is_none());
assert!(spec.get("data_path").is_none());
assert!(spec.get("run_id").is_none());
assert!(spec.get("gpu_ram_share").is_none());
}
#[test]
fn tunnel_argv_orders_user_options_before_the_defaults() {
let ssh = SshConfig {
target: Some("ctrl".into()),
port: Some(2222),
user: Some("join-user".into()),
identity_file: Some("/etc/flodl/join_key".into()),
options: vec!["ServerAliveInterval=5".into()],
};
let argv = build_tunnel_argv(&ssh, 40123, "127.0.0.1", 1337);
assert_eq!(argv[0], "ssh");
assert!(argv.contains(&"-N".to_string()));
assert!(argv.contains(&"BatchMode=yes".to_string()));
assert!(argv.contains(&"ExitOnForwardFailure=yes".to_string()));
let user_pos = argv
.iter()
.position(|a| a == "ServerAliveInterval=5")
.unwrap();
let default_pos = argv
.iter()
.position(|a| a == "ServerAliveInterval=30")
.unwrap();
assert!(user_pos < default_pos);
assert!(argv.contains(&"127.0.0.1:40123:127.0.0.1:1337".to_string()));
assert_eq!(argv.last().map(String::as_str), Some("ctrl"));
let p = argv.iter().position(|a| a == "-p").unwrap();
assert_eq!(argv[p + 1], "2222");
let l = argv.iter().position(|a| a == "-l").unwrap();
assert_eq!(argv[l + 1], "join-user");
let i = argv.iter().position(|a| a == "-i").unwrap();
assert_eq!(argv[i + 1], "/etc/flodl/join_key");
}
#[test]
fn wait_tunnel_ready_sees_a_live_listener_and_a_dead_child() {
let mut dead = Command::new("true").spawn().unwrap();
std::thread::sleep(Duration::from_millis(50));
let err = wait_tunnel_ready(&mut dead, 1).unwrap_err();
assert!(err.contains("before the forward came up"), "got: {err}");
let listener = TcpListener::bind("127.0.0.1:0").unwrap();
let port = listener.local_addr().unwrap().port();
let mut slow = Command::new("sleep").arg("5").spawn().unwrap();
assert!(wait_tunnel_ready(&mut slow, port).is_ok());
let _ = slow.kill();
let _ = slow.wait();
}
#[test]
fn hex_encode_is_lowercase_bytewise() {
assert_eq!(hex_encode(b"\x00\xff\x10"), "00ff10");
assert_eq!(hex_encode(b"{}"), "7b7d");
}
}