Skip to main content

flodl_cli/
join.rs

1//! `fdl join` — self-deploy this box as a dial-in worker.
2//!
3//! The worker-side walk-in for a discovery window (see docs/ddp/01-reference.md):
4//! dials the controller's mux port, offers this host's GPUs, and — once
5//! admitted — the training binary takes over in agent role (it joins,
6//! then spawns and supervises this host's relay and rank children; see
7//! flodl's `distributed::launcher::agent`). fdl-cli stays protocol-blind:
8//! its whole job is orchestration —
9//!
10//!   1. resolve settings (flags over the fdl.yml `join:` block),
11//!   2. prepare this box ([`crate::prepare`]): the GPU gate, the dataset
12//!      source root, the node-local directories the data plane writes —
13//!      all of it BEFORE the dial, because admission starts a window
14//!      deadline,
15//!   3. optionally bring up an ssh `-L` forward of the controller port
16//!      (the guardrailed-sshd trust path: reachability = authentication),
17//!   4. synthesize the agent bootstrap spec into the binary's
18//!      environment (`FLODL_INTERNAL_AGENT_JSON`, hex-encoded JSON — the
19//!      same envelope cluster fan-out ships),
20//!   5. run + supervise the binary, and in `--persist` mode re-dial
21//!      with backoff when it exits (the systemd / golden-image loop).
22//!
23//! The spec (which may carry the pre-shared session token) rides the
24//! child's ENVIRONMENT, never argv — owner-readable via
25//! `/proc/<pid>/environ` instead of world-readable via `ps`, the same
26//! salt hygiene as the launcher's fan-out.
27
28use std::net::{TcpListener, TcpStream};
29use std::path::{Path, PathBuf};
30use std::process::{Child, Command, Stdio};
31use std::time::{Duration, Instant};
32
33use crate::builtins::JoinArgs;
34use crate::config::{self, DEFAULT_CONTROLLER_PORT, SshConfig, WorkerJoin, WorkerSource};
35use crate::context::Context;
36use crate::prepare::{self, DataSpec, Fail, PrepareSpec, Prepared, SourceSpec};
37use crate::style;
38
39/// Agent bootstrap env var — must match flodl's
40/// `distributed::launcher::ENV_AGENT_JSON` (the field names in the hex
41/// JSON must match `AgentSpec`; locked by `agent_spec_shape_is_the_wire_contract`).
42const ENV_AGENT_JSON: &str = "FLODL_INTERNAL_AGENT_JSON";
43
44/// Exit code for a failure retrying cannot fix: no usable GPU, a spec
45/// that does not parse, a directory that cannot be created, a missing
46/// binary. Distinct from 1 (a transient failure, one-shot) so a
47/// fleet can act on the difference without parsing stderr:
48///
49/// ```ini
50/// # /etc/systemd/system/flodl-join.service
51/// Restart=always
52/// RestartPreventExitStatus=2   # stop hot-looping a misprovisioned box
53/// FailureAction=poweroff       # ... and self-deprovision it
54/// ```
55///
56/// fdl deliberately does not power a box off itself: the decision belongs
57/// to whatever owns the instance's lifecycle, and 2 is how it hears about it.
58pub const EXIT_PERMANENT: i32 = 2;
59
60/// How long the ssh forward gets to come up (auth + local bind).
61const TUNNEL_READY_BUDGET: Duration = Duration::from_secs(20);
62
63/// `--persist` re-dial backoff: floor, cap, and the attempt duration
64/// past which the backoff resets to the floor (the agent ran a real
65/// stint, so the next failure is a fresh incident, not a hot loop).
66const BACKOFF_MIN: Duration = Duration::from_secs(5);
67const BACKOFF_MAX: Duration = Duration::from_secs(60);
68const BACKOFF_RESET_AFTER: Duration = Duration::from_secs(120);
69
70/// Run `fdl join`. `bin_tail` is everything after the command line's
71/// standalone `--`: the training binary's own arguments, forwarded
72/// verbatim (rank children re-enter the binary with them). `None` =
73/// no `--` was given (the config block's `args:` applies); a present
74/// but empty tail is an explicit "no arguments".
75///
76/// Exit code: the agent's own exit code (one-shot); [`EXIT_PERMANENT`]
77/// for a failure retrying cannot fix; 1 for a transient one. `--persist`
78/// re-dials through transient failures and agent exits, and returns only
79/// on a permanent one.
80pub fn run(cli: &JoinArgs, bin_tail: Option<&[String]>) -> i32 {
81    let (block, project_root) = match load_join_block() {
82        Ok(pair) => pair,
83        Err(e) => {
84            crate::cli_error!("{e}");
85            return EXIT_PERMANENT;
86        }
87    };
88    let eff = match resolve_effective(
89        cli,
90        bin_tail,
91        block,
92        &crate::cluster::resolve_local_hostname(),
93    ) {
94        Ok(eff) => eff,
95        Err(e) => {
96            crate::cli_error!("{e}");
97            return EXIT_PERMANENT;
98        }
99    };
100    if eff.controller_defaulted {
101        eprintln!(
102            "{}",
103            style::dim(&format!(
104                "fdl join: no controller configured; dialing \
105                 127.0.0.1:{DEFAULT_CONTROLLER_PORT} (pass an address or set \
106                 `join.controller` in fdl.yml)"
107            )),
108        );
109    }
110
111    // A binary named as a path must exist NOW — a persist loop retrying
112    // a missing path forever helps nobody. A binary built from source
113    // cannot be checked here: it does not exist until the attempt has
114    // fetched and compiled the tree.
115    if let BinSource::Given(path) = &eff.bin
116        && !Path::new(path).is_file()
117    {
118        crate::cli_error!(
119            "training binary not found: {path} — build it first and \
120                 point `--bin` (or fdl.yml `join.bin`) at it, or hand this \
121                 box a `--source` to build",
122        );
123        return EXIT_PERMANENT;
124    }
125
126    // Local active libtorch (honors FDL_LIBTORCH_CASE), anchored on the
127    // project root the config walk found: its lib/ rides
128    // LD_LIBRARY_PATH on the child, and its variant label rides the
129    // join hello. Absent (fdl running outside a project) the child env
130    // is left untouched — the binary may carry an rpath or the caller's
131    // environment already provides the libs.
132    let libtorch = resolve_local_libtorch(project_root.as_deref());
133
134    // Model-sig probe cache, living exactly as long as the persist loop:
135    // one slot, keyed by a digest of the probe recipe (binary identity +
136    // args). An idle fleet re-dials every BACKOFF_MAX forever, and
137    // without this every re-dial would rebuild the model on CPU — or,
138    // for a binary that predates the probe, pay the full probe timeout
139    // — for an unchanged binary.
140    let mut sig_cache: Option<(u64, Option<String>)> = None;
141
142    let mut backoff = BACKOFF_MIN;
143    loop {
144        let started = Instant::now();
145        // What happened, phrased for the re-dial line. Every branch that
146        // is not re-dialable returns from here.
147        let outcome = match attempt(&eff, libtorch.as_ref(), &mut sig_cache) {
148            Ok(code) => {
149                if !eff.persist {
150                    return code;
151                }
152                format!("agent exited with code {code}")
153            }
154            Err(fail) => {
155                crate::cli_error!("{}", fail.message());
156                if fail.is_permanent() {
157                    // The whole point of the class: a box that cannot be
158                    // fixed by waiting must stop, not hot-loop.
159                    eprintln!(
160                        "{}",
161                        style::dim(&format!(
162                            "fdl join: not re-dialing — retrying cannot \
163                             fix this (exit {EXIT_PERMANENT})"
164                        )),
165                    );
166                    return EXIT_PERMANENT;
167                }
168                if !eff.persist {
169                    return 1;
170                }
171                "attempt failed".to_string()
172            }
173        };
174        if started.elapsed() > BACKOFF_RESET_AFTER {
175            backoff = BACKOFF_MIN;
176        }
177        eprintln!(
178            "fdl join: {outcome} after {}s; re-dialing in {}s (--persist)",
179            started.elapsed().as_secs(),
180            backoff.as_secs(),
181        );
182        std::thread::sleep(backoff);
183        backoff = (backoff * 2).min(BACKOFF_MAX);
184    }
185}
186
187// ---------------------------------------------------------------------------
188// Settings resolution
189// ---------------------------------------------------------------------------
190
191/// The fully resolved join recipe: flags merged over the fdl.yml
192/// `join:` block, every default applied.
193#[derive(Debug)]
194struct Effective {
195    /// Controller mux address. Under `ssh` this is the address as seen
196    /// FROM the ssh host (the `-L` forward's far end).
197    controller_host: String,
198    controller_port: u16,
199    /// True when neither flags nor config named a controller and the
200    /// loopback convention default applied (worth a stderr note).
201    controller_defaulted: bool,
202    /// Tunnel hop; `None` = direct dial.
203    ssh: Option<SshConfig>,
204    /// Pre-shared session credential (hex). `None` = open admission.
205    token: Option<String>,
206    /// How this box gets its training binary: a path to run as given, or
207    /// a source to build. Exactly one, checked at resolution.
208    bin: BinSource,
209    /// libtorch variant to acquire; `None` keeps this box's active one.
210    libtorch_spec: Option<String>,
211    /// Logical host name in the join hello.
212    host: String,
213    /// Explicit CUDA device ids; `None` = all GPUs on this host.
214    devices: Option<Vec<u8>>,
215    persist: bool,
216    /// The binary's own arguments.
217    bin_args: Vec<String>,
218    /// Dataset source root on this box (the mountpoint when
219    /// `data_source` is set); `None` ships nothing to the ranks.
220    data_path: Option<String>,
221    /// Transport that establishes the source root, `<scheme>://<target>`.
222    data_source: Option<String>,
223    /// Integrated-GPU host-RAM share of this box; `None` ships nothing
224    /// (the envelope's cluster-scope default, if any, then stands).
225    gpu_ram_share: Option<f64>,
226    /// Probe the binary for its model signature before each dial
227    /// (default true; `--no-sig-probe` / `join.sig_probe: false`).
228    sig_probe: bool,
229}
230
231/// Where this box's training binary comes from. `source:` and `bin:` are
232/// mutually exclusive because they answer the same question, and keeping
233/// the given-binary case a separate variant rather than a third kind of
234/// source spec is what keeps an artifact-versus-source distinction out of
235/// the source grammar.
236#[derive(Debug, PartialEq, Eq)]
237enum BinSource {
238    /// A path on this box, run as given.
239    Given(String),
240    /// Fetched and built here.
241    Build(WorkerSource),
242}
243
244impl Effective {
245    /// Everything [`crate::prepare`] has to settle. The tunnel block
246    /// rides along to both artifact specs: the data host and the source
247    /// host are the controller box in the shape this exists for, so its
248    /// key and options apply to them too.
249    fn prepare_spec<'a>(
250        &'a self,
251        active_libtorch: Option<&'a (PathBuf, String)>,
252    ) -> PrepareSpec<'a> {
253        PrepareSpec {
254            data: DataSpec {
255                path: self.data_path.as_deref(),
256                source: self.data_source.as_deref(),
257                ssh: self.ssh.as_ref(),
258            },
259            libtorch: self.libtorch_spec.as_deref(),
260            active_libtorch,
261            devices: self.devices.as_deref(),
262            source: match &self.bin {
263                BinSource::Given(_) => None,
264                BinSource::Build(s) => Some(SourceSpec {
265                    from: &s.from,
266                    cwd: s.cwd.as_deref(),
267                    build: s.build.as_deref(),
268                    bin: s.bin.as_deref(),
269                    ssh: self.ssh.as_ref(),
270                }),
271            },
272        }
273    }
274}
275
276/// Merge flags over the config block. Pure — all I/O (hostname, config
277/// load) happens in the callers so this stays table-testable.
278fn resolve_effective(
279    cli: &JoinArgs,
280    bin_tail: Option<&[String]>,
281    block: Option<WorkerJoin>,
282    local_hostname: &str,
283) -> Result<Effective, String> {
284    let block = block.unwrap_or_default();
285
286    if cli.identity.is_some() && cli.ssh.is_none() && block.ssh.is_none() {
287        return Err("--identity is the tunnel's key file — it needs an ssh hop \
288             (`--ssh` or fdl.yml `join.ssh`)"
289            .to_string());
290    }
291
292    // Tunnel hop: `--ssh [user@]host[:port]` replaces the block's
293    // target/user/port but inherits its identity_file/options (not
294    // expressible in the compact form); `--identity` wins last. A block
295    // without `target:` is an authoring error — ssh needs a host.
296    let ssh = match (&cli.ssh, block.ssh) {
297        (Some(spec), b) => {
298            let mut cfg = parse_ssh_spec(spec)?;
299            if let Some(b) = b {
300                cfg.identity_file = b.identity_file;
301                cfg.options = b.options;
302            }
303            Some(cfg)
304        }
305        (None, Some(b)) => {
306            if b.target.is_none() {
307                return Err("fdl.yml join.ssh needs a `target:` (the tunnel host)".to_string());
308            }
309            Some(b)
310        }
311        (None, None) => None,
312    };
313    let mut ssh = ssh;
314    if let (Some(cfg), Some(id)) = (ssh.as_mut(), &cli.identity) {
315        cfg.identity_file = Some(id.clone());
316    }
317
318    // Controller: flag > block > loopback convention. Through a tunnel
319    // the loopback default is THE convention (guardrailed sshd on the
320    // controller box), no note needed; a bare loopback default deserves
321    // one.
322    let named = cli.controller.as_ref().or(block.controller.as_ref());
323    let controller_defaulted = named.is_none() && ssh.is_none();
324    let (controller_host, controller_port) = match named {
325        Some(spec) => parse_host_port(spec)?,
326        None => ("127.0.0.1".to_string(), DEFAULT_CONTROLLER_PORT),
327    };
328
329    // The binary: a path to run, or a source to build. Flags win over
330    // the block on each side, and naming both ways is an authoring error
331    // rather than a precedence puzzle — a box that builds its own binary
332    // and is also handed one has no defensible answer.
333    let bin_path = cli.bin.clone().or(block.bin);
334    let source = match (cli.source.clone(), block.source) {
335        (Some(from), b) => Some(WorkerSource {
336            from,
337            // The compact `--source` flag carries only the transport, so
338            // the rest keeps coming from the block unless its own flag
339            // overrides it (same shape as `--ssh`).
340            cwd: cli
341                .source_cwd
342                .clone()
343                .or_else(|| b.as_ref().and_then(|b| b.cwd.clone())),
344            build: cli
345                .source_build
346                .clone()
347                .or_else(|| b.as_ref().and_then(|b| b.build.clone())),
348            // No artifact anywhere is legal: a published tree carries a
349            // run manifest that names it, and that manifest is the
350            // authority when it is there.
351            bin: cli
352                .source_bin
353                .clone()
354                .or_else(|| b.as_ref().and_then(|b| b.bin.clone())),
355        }),
356        (None, Some(mut b)) => {
357            if let Some(cwd) = cli.source_cwd.clone() {
358                b.cwd = Some(cwd);
359            }
360            if let Some(build) = cli.source_build.clone() {
361                b.build = Some(build);
362            }
363            if let Some(bin) = cli.source_bin.clone() {
364                b.bin = Some(bin);
365            }
366            Some(b)
367        }
368        (None, None) => None,
369    };
370    // A source flag with no source to attach to is an authoring error,
371    // not something to drop on the floor: the operator meant it to change
372    // the run.
373    if source.is_none() {
374        for (flag, set) in [
375            ("--source-cwd", cli.source_cwd.is_some()),
376            ("--source-build", cli.source_build.is_some()),
377            ("--source-bin", cli.source_bin.is_some()),
378        ] {
379            if set {
380                return Err(format!(
381                    "{flag} has no source to apply to — pass `--source \
382                     <spec>` too, or set `join.source` in fdl.yml"
383                ));
384            }
385        }
386    }
387
388    let bin = match (bin_path, source) {
389        (Some(_), Some(_)) => {
390            return Err("`bin:` and `source:` both name this box's training binary \
391                 — keep the one you mean. `bin:` runs a binary as given; \
392                 `source:` fetches and builds one here"
393                .to_string());
394        }
395        (Some(path), None) => BinSource::Given(path),
396        (None, Some(source)) => BinSource::Build(source),
397        (None, None) => {
398            return Err("no training binary configured — pass `--bin <path>` (run \
399                 it as given) or `--source <spec>` (build it here), or set \
400                 `join.bin` / `join.source` in fdl.yml. The binary is the \
401                 protocol: it dials, joins, and runs this host's ranks"
402                .to_string());
403        }
404    };
405
406    let devices = match &cli.devices {
407        Some(spec) => parse_devices(spec)?,
408        None => block.devices,
409    };
410
411    // A `--` tail — even an empty one — REPLACES the block's args: the
412    // args must match the run, so "explicitly none" must be sayable.
413    let bin_args = match bin_tail {
414        Some(tail) => tail.to_vec(),
415        None => block.args,
416    };
417
418    Ok(Effective {
419        controller_host,
420        controller_port,
421        controller_defaulted,
422        ssh,
423        token: cli.token.clone().or(block.token),
424        bin,
425        host: cli
426            .host
427            .clone()
428            .or(block.host)
429            .unwrap_or_else(|| local_hostname.to_string()),
430        devices,
431        persist: cli.persist || block.persist,
432        bin_args,
433        libtorch_spec: cli.libtorch.clone().or(block.libtorch),
434        data_path: cli.data_path.clone().or(block.data_path),
435        data_source: cli.data_source.clone().or(block.data_source),
436        gpu_ram_share: cli.gpu_ram_share.or(block.gpu_ram_share),
437        // The flag only disables; `sig_probe: false` in yml is the
438        // standing form of the same choice. Default on.
439        sig_probe: !cli.no_sig_probe && block.sig_probe.unwrap_or(true),
440    })
441}
442
443/// Load the top-level `join:` block from the PROJECT config (base
444/// fdl.yml merged with the active env overlay when one is selected),
445/// plus the directory it lives in (the project root — where libtorch/
446/// is anchored). The walk steps over command-level fdl.ymls
447/// ([`config::find_project_config`]): `fdl join` typically runs from
448/// the command dir the training binary expects as cwd (e.g.
449/// `ddp-bench/`), whose own fdl.yml is a command config that neither
450/// carries a `join:` block nor marks the libtorch root. `Ok(None)`
451/// root/block when there is no project at all — flags carry everything
452/// then; a present-but-broken project config is a loud error, not a
453/// silent fallback (the operator may be relying on `join.bin`).
454fn load_join_block() -> Result<(Option<WorkerJoin>, Option<PathBuf>), String> {
455    let cwd =
456        std::env::current_dir().map_err(|e| format!("cannot read the current directory: {e}"))?;
457    let Some(config_path) = config::find_project_config(&cwd) else {
458        return Ok((None, None));
459    };
460    let env_name = std::env::var("FDL_ENV")
461        .ok()
462        .filter(|s| !s.trim().is_empty());
463    let project = config::load_project_with_env(&config_path, env_name.as_deref())
464        .map_err(|e| format!("cannot load {}: {e}", config_path.display()))?;
465    let root = config_path.parent().map(Path::to_path_buf);
466    Ok((project.join, root))
467}
468
469/// Parse `host[:port]`, default port [`DEFAULT_CONTROLLER_PORT`] —
470/// same convention as `fdl status --addr`.
471fn parse_host_port(spec: &str) -> Result<(String, u16), String> {
472    match spec.rsplit_once(':') {
473        Some((host, port)) => {
474            let port = port.parse::<u16>().map_err(|_| {
475                format!("invalid controller address `{spec}` — expected host[:port]")
476            })?;
477            if host.is_empty() {
478                return Err(format!(
479                    "invalid controller address `{spec}` — expected host[:port]"
480                ));
481            }
482            Ok((host.to_string(), port))
483        }
484        None => Ok((spec.to_string(), DEFAULT_CONTROLLER_PORT)),
485    }
486}
487
488/// Parse the compact tunnel spec `[user@]host[:port]` into an
489/// [`SshConfig`] (target/user/port only; identity/options come from
490/// the config block or `--identity`).
491fn parse_ssh_spec(spec: &str) -> Result<SshConfig, String> {
492    let (user, rest) = match spec.split_once('@') {
493        Some((u, r)) if !u.is_empty() => (Some(u.to_string()), r),
494        Some(_) => {
495            return Err(format!("invalid --ssh `{spec}` — empty user before `@`"));
496        }
497        None => (None, spec),
498    };
499    let (host, port) = match rest.rsplit_once(':') {
500        Some((h, p)) => {
501            let port = p
502                .parse::<u16>()
503                .map_err(|_| format!("invalid --ssh `{spec}` — expected [user@]host[:port]"))?;
504            (h, Some(port))
505        }
506        None => (rest, None),
507    };
508    if host.is_empty() {
509        return Err(format!(
510            "invalid --ssh `{spec}` — expected [user@]host[:port]"
511        ));
512    }
513    Ok(SshConfig {
514        target: Some(host.to_string()),
515        port,
516        user,
517        identity_file: None,
518        options: Vec::new(),
519    })
520}
521
522/// Parse `--devices`: comma-separated CUDA ids, or `all` for
523/// every GPU on this host (= unset).
524fn parse_devices(spec: &str) -> Result<Option<Vec<u8>>, String> {
525    if spec.trim().eq_ignore_ascii_case("all") {
526        return Ok(None);
527    }
528    spec.split(',')
529        .map(|s| {
530            s.trim()
531                .parse::<u8>()
532                .map_err(|_| format!("invalid --devices `{spec}` — expected e.g. `0,1` or `all`"))
533        })
534        .collect::<Result<Vec<u8>, String>>()
535        .map(Some)
536}
537
538// ---------------------------------------------------------------------------
539// Agent spec synthesis
540// ---------------------------------------------------------------------------
541
542/// Build the hex-encoded JSON payload for [`ENV_AGENT_JSON`]. Field
543/// names ARE the wire contract with flodl's `AgentSpec` deserializer;
544/// optional fields are omitted (serde defaults fill them).
545///
546/// The prepared data path travels this way rather than in the join hello
547/// because the controller has nothing to say about it: it never
548/// configured this host, and only this box knows where its source root
549/// actually ended up. flodl's agent inserts it into the envelope its
550/// rank children read, beside the same-shaped rewrite it already does
551/// for the controller address.
552fn agent_spec_hex(
553    eff: &Effective,
554    dial: (&str, u16),
555    libtorch_label: &str,
556    prepared: &Prepared,
557    model_sig_hex: Option<&str>,
558) -> String {
559    let mut spec = serde_json::json!({
560        "host": eff.host,
561        "controller_host": dial.0,
562        "controller_port": dial.1,
563        "libtorch": libtorch_label,
564    });
565    if let Some(token) = &eff.token {
566        spec["salt_hex"] = serde_json::json!(token);
567    }
568    if let Some(devices) = &eff.devices {
569        spec["local_devices"] = serde_json::json!(devices);
570    }
571    if let Some(data) = &prepared.data_path {
572        spec["data_path"] = serde_json::json!(data.display().to_string());
573    }
574    if let Some(run) = &prepared.run_id {
575        spec["run_id"] = serde_json::json!(run);
576    }
577    // Host-hardware truth, same as data_path: the agent writes it into
578    // the envelope's host block, overriding any cluster-scope default
579    // the controller stamped. Omitted when undeclared, so that default
580    // stands.
581    if let Some(share) = eff.gpu_ram_share {
582        spec["gpu_ram_share"] = serde_json::json!(share);
583    }
584    // Probed from the binary itself, so admission can refuse a box
585    // building a different model while it still costs only this box's
586    // own attempt. Omitted when the probe was skipped or failed.
587    if let Some(sig) = model_sig_hex {
588        spec["model_sig_hex"] = serde_json::json!(sig);
589    }
590    hex_encode(spec.to_string().as_bytes())
591}
592
593/// Lowercase hex — flodl's `cluster::hex_decode` counterpart.
594fn hex_encode(bytes: &[u8]) -> String {
595    let mut s = String::with_capacity(bytes.len() * 2);
596    for b in bytes {
597        s.push_str(&format!("{b:02x}"));
598    }
599    s
600}
601
602/// Active libtorch of this box, anchored on the project root when the
603/// config walk found one (a command-dir cwd's `Context::resolve` would
604/// stop a level too low); the plain context fallback covers project-less
605/// setups (`~/.flodl`).
606fn resolve_local_libtorch(project_root: Option<&Path>) -> Option<(PathBuf, String)> {
607    let root = match project_root {
608        Some(r) => r.to_path_buf(),
609        None => Context::resolve().root,
610    };
611    crate::libtorch::detect::active_variant(&root)
612}
613
614/// `LD_LIBRARY_PATH` for the training binary, with this box's inherited
615/// value appended.
616///
617/// The ordering inside is not ours to choose: on ROCm the system runtime
618/// must precede libtorch's own lib dir, because libtorch-rocm bundles the
619/// whole userspace ROCm stack and a bundle that disagrees with the host's
620/// amdkfd driver segfaults the rank at its first GPU op. Prepending
621/// unconditionally, which is what this did before, is the segfault
622/// configuration on an AMD box.
623fn child_ld_library_path(libtorch_dir: &Path, variant: &str) -> String {
624    let lib = libtorch_dir.join("lib").display().to_string();
625    let vendor = crate::libtorch::detect::variant_vendor(variant);
626    let value = crate::libtorch::detect::ld_library_path_value(
627        vendor,
628        &lib,
629        &crate::libtorch::detect::local_rocm_lib_dir(),
630    );
631    match std::env::var("LD_LIBRARY_PATH") {
632        Ok(cur) if !cur.is_empty() => format!("{value}:{cur}"),
633        _ => value,
634    }
635}
636
637// ---------------------------------------------------------------------------
638// One attempt: tunnel up (optional), agent run, teardown
639// ---------------------------------------------------------------------------
640
641/// One full join attempt: prepare, tunnel, dial, supervise. Returns the
642/// agent's exit code, or the classed reason preparation/orchestration
643/// stopped.
644///
645/// Preparation comes first and every attempt re-runs it: admission
646/// starts a window deadline, so a mount established after the dial burns
647/// it, and re-running is how `--persist` becomes a provisioning loop.
648///
649/// The tunnel — when one is configured — lives exactly as long as the
650/// attempt: rebuilt fresh each re-dial, so a half-dead forward can never
651/// outlive the run it served.
652fn attempt(
653    eff: &Effective,
654    active_libtorch: Option<&(PathBuf, String)>,
655    sig_cache: &mut Option<(u64, Option<String>)>,
656) -> Result<i32, Fail> {
657    let mut notes = Vec::new();
658    let prepared = prepare::prepare(&eff.prepare_spec(active_libtorch), &mut notes);
659    prepare::print_notes("join", &notes);
660    let prepared = prepared?;
661
662    // A built binary runs in the project directory inside the fetched
663    // tree, which is where its own relative paths resolve; a given one
664    // keeps fdl's cwd, exactly as before.
665    let (bin, bin_cwd) = match (&eff.bin, &prepared.bin) {
666        (BinSource::Build(_), Some(built)) => (built.bin.clone(), Some(built.cwd.clone())),
667        (BinSource::Given(path), None) => (PathBuf::from(path), None),
668        // Neither combination is reachable — preparation returns a binary
669        // exactly when it was given a source — so say so rather than
670        // quietly preferring one and hiding a wiring inversion.
671        (kind, built) => {
672            return Err(Fail::Permanent(format!(
673                "internal: preparation and the resolved binary disagree \
674                 ({}, built={})",
675                match kind {
676                    BinSource::Given(_) => "a path was given",
677                    BinSource::Build(_) => "a source was given",
678                },
679                built.is_some(),
680            )));
681        }
682    };
683
684    // The run's arguments belong to the run. A published manifest replaces
685    // whatever this box carried, because rank children re-enter the binary
686    // with them: a standing fleet holding its own copy would train the new
687    // run with the previous one's hyperparameters. Saying so out loud is
688    // the difference between authority and a silent substitution.
689    let args: &[String] = match &prepared.args {
690        Some(published) => {
691            if !eff.bin_args.is_empty() && published != &eff.bin_args {
692                eprintln!(
693                    "{}",
694                    style::dim(&format!(
695                        "fdl join: the published run's arguments replace \
696                         this box's ({} -> {})",
697                        eff.bin_args.join(" "),
698                        published.join(" "),
699                    )),
700                );
701            }
702            published
703        }
704        None => &eff.bin_args,
705    };
706
707    // Model-signature probe, before the tunnel like everything else in
708    // preparation: admission starts a window deadline, and the probe
709    // runs the binary's whole main up to `Trainer::run`. Cached by
710    // recipe digest across re-dials — the OUTCOME is cached, a failed
711    // probe included, so an unchanged binary pays the probe (or its
712    // timeout) once, not once per backoff tick. A rebuild changes the
713    // mtime and a re-publish changes the args, so staleness invalidates
714    // itself; the libtorch variant is deliberately not in the recipe
715    // (the manifest — names, shapes, dtypes — is device-independent).
716    let model_sig_hex = if eff.sig_probe {
717        match probe_recipe_digest(&bin, args) {
718            Some(digest) => match sig_cache {
719                Some((key, cached)) if *key == digest => cached.clone(),
720                _ => {
721                    let sig =
722                        model_sig_probe(&bin, bin_cwd.as_deref(), args, prepared.libtorch.as_ref());
723                    *sig_cache = Some((digest, sig.clone()));
724                    sig
725                }
726            },
727            // The binary un-stat-able between resolution and here is a
728            // race with a rebuild: probe uncached, next attempt keys.
729            None => model_sig_probe(&bin, bin_cwd.as_deref(), args, prepared.libtorch.as_ref()),
730        }
731    } else {
732        None
733    };
734
735    let mut tunnel: Option<Child> = None;
736    let dial: (String, u16) = match &eff.ssh {
737        Some(ssh) => {
738            let local_port = pick_local_port().map_err(Fail::Transient)?;
739            let argv =
740                build_tunnel_argv(ssh, local_port, &eff.controller_host, eff.controller_port);
741            eprintln!(
742                "fdl join: opening tunnel {} -> {}:{} (local port {local_port})",
743                ssh.target.as_deref().unwrap_or("?"),
744                eff.controller_host,
745                eff.controller_port,
746            );
747            let mut child = Command::new(&argv[0])
748                .args(&argv[1..])
749                .stdin(Stdio::null())
750                .spawn()
751                .map_err(|e| {
752                    // No ssh on the box is a provisioning fact, not a
753                    // passing condition.
754                    Fail::Permanent(format!("spawn ssh tunnel: {e}"))
755                })?;
756            // Auth failure and an unreachable host are both possible
757            // here and ssh does not let us tell them apart, so this
758            // stays re-dialable: a wrong key hits the backoff cap and
759            // keeps saying so, once a minute, loudly.
760            if let Err(e) = wait_tunnel_ready(&mut child, local_port) {
761                let _ = child.kill();
762                let _ = child.wait();
763                return Err(Fail::Transient(e));
764            }
765            tunnel = Some(child);
766            ("127.0.0.1".to_string(), local_port)
767        }
768        None => (eff.controller_host.clone(), eff.controller_port),
769    };
770
771    // Preparation is the authority on libtorch: it either acquired a
772    // variant or carried the active one through, and the label it settles
773    // on is what the join hello announces.
774    let libtorch_label = prepared
775        .libtorch
776        .as_ref()
777        .map(|(_, l)| l.as_str())
778        .unwrap_or("");
779    let spec_hex = agent_spec_hex(
780        eff,
781        (&dial.0, dial.1),
782        libtorch_label,
783        &prepared,
784        model_sig_hex.as_deref(),
785    );
786
787    let mut cmd = Command::new(&bin);
788    cmd.args(args)
789        .env(ENV_AGENT_JSON, &spec_hex)
790        // Children report under the logical roster name even when it
791        // differs from `hostname` — same override fan-out applies.
792        .env(crate::cluster::ENV_HOST_OVERRIDE, &eff.host)
793        .stdin(Stdio::null());
794    if let Some(cwd) = &bin_cwd {
795        cmd.current_dir(cwd);
796    }
797    if let Some((dir, variant)) = &prepared.libtorch {
798        cmd.env("LD_LIBRARY_PATH", child_ld_library_path(dir, variant));
799    }
800
801    // A given path was checked before the loop and a built one was just
802    // written, so a spawn failure here is the file itself: not
803    // executable, wrong architecture, bad interpreter.
804    let status = cmd
805        .status()
806        .map_err(|e| Fail::Permanent(format!("run {}: {e}", bin.display())));
807    if let Some(mut t) = tunnel.take() {
808        let _ = t.kill();
809        let _ = t.wait();
810    }
811    Ok(status?.code().unwrap_or(1))
812}
813
814/// Digest of the probe recipe — everything the probe's answer depends
815/// on: the binary's identity (path, mtime, size) and the arguments the
816/// run enters it with. One u64 via std's SipHash rather than a
817/// cryptographic hash: flodl-cli is zero-dep, the key lives one process
818/// and guards a cache on a box that is trusted (not proven), and a
819/// collision's worst case — a stale signature in the hello — lands on
820/// the formation-time backstop. `None` when the binary cannot be
821/// stat'ed (a race with a rebuild): probe uncached, key next attempt.
822fn probe_recipe_digest(bin: &Path, args: &[String]) -> Option<u64> {
823    use std::hash::{Hash, Hasher};
824    let meta = std::fs::metadata(bin).ok()?;
825    let mut h = std::collections::hash_map::DefaultHasher::new();
826    bin.hash(&mut h);
827    meta.len().hash(&mut h);
828    meta.modified()
829        .ok()?
830        .duration_since(std::time::UNIX_EPOCH)
831        .ok()?
832        .as_nanos()
833        .hash(&mut h);
834    args.hash(&mut h);
835    Some(h.finish())
836}
837
838/// How many trailing stdout lines the probe keeps as evidence when no
839/// signature appears.
840const PROBE_TAIL_LINES: usize = 8;
841
842/// Probe-run marker env (flodl's launcher contract: with it set,
843/// `Trainer::run` builds the model on CPU, prints the signature line
844/// and exits — before auto-promote, before any cluster role).
845const ENV_MODEL_SIG_PROBE: &str = "FLODL_INTERNAL_MODEL_SIG_PROBE";
846
847/// Stdout line the probe scans for (main-body prints above it are
848/// harmless — the prefix is the protocol, not the whole stream).
849const MODEL_SIG_LINE: &str = "flodl-model-sig: ";
850
851/// Ceiling on the probe re-run. A current flodl answers in the time its
852/// main takes to reach `Trainer::run`; a binary built against a flodl
853/// that predates the probe ignores the env and runs its whole main,
854/// which is exactly what this bounds.
855const MODEL_SIG_PROBE_TIMEOUT: Duration = Duration::from_secs(120);
856
857/// Re-run the training binary as a model-signature probe and return the
858/// 64-hex signature it prints.
859///
860/// Best-effort BY DESIGN: every failure degrades to `None` — the hello
861/// then gates nothing and the formation-time handshake check stays the
862/// backstop — but each failure mode says so, because one of them
863/// (a non-zero exit) predicts the rank children failing the same way
864/// after formation, with the same binary and the same arguments.
865fn model_sig_probe(
866    bin: &Path,
867    cwd: Option<&Path>,
868    args: &[String],
869    libtorch: Option<&(PathBuf, String)>,
870) -> Option<String> {
871    eprintln!(
872        "{}",
873        style::dim(
874            "fdl join: probing the binary for its model signature \
875                    (--no-sig-probe skips this)"
876        ),
877    );
878    let mut cmd = Command::new(bin);
879    cmd.args(args)
880        .env(ENV_MODEL_SIG_PROBE, "1")
881        .stdin(Stdio::null())
882        .stdout(Stdio::piped());
883    if let Some(dir) = cwd {
884        cmd.current_dir(dir);
885    }
886    if let Some((dir, variant)) = libtorch {
887        cmd.env("LD_LIBRARY_PATH", child_ld_library_path(dir, variant));
888    }
889    let mut child = match cmd.spawn() {
890        Ok(c) => c,
891        Err(e) => {
892            eprintln!(
893                "fdl join: model-sig probe could not run {}: {e}; joining \
894                 without a signature",
895                bin.display(),
896            );
897            return None;
898        }
899    };
900    let stdout = child.stdout.take().expect("stdout was piped");
901    let reader = std::thread::spawn(move || {
902        use std::io::{BufRead, BufReader};
903        let mut sig = None;
904        // Last few lines kept as evidence: when no signature turns up,
905        // what the binary actually said beats anything we could infer
906        // about why. Bounded so a chatty binary cannot grow this.
907        let mut tail: std::collections::VecDeque<String> = std::collections::VecDeque::new();
908        for line in BufReader::new(stdout).lines() {
909            let Ok(line) = line else { break };
910            if let Some(rest) = line.strip_prefix(MODEL_SIG_LINE) {
911                sig = Some(rest.trim().to_string());
912            }
913            if tail.len() == PROBE_TAIL_LINES {
914                tail.pop_front();
915            }
916            tail.push_back(line);
917        }
918        (sig, tail)
919    });
920    let deadline = Instant::now() + MODEL_SIG_PROBE_TIMEOUT;
921    let status = loop {
922        match child.try_wait() {
923            Ok(Some(st)) => break Some(st),
924            Ok(None) if Instant::now() >= deadline => {
925                let _ = child.kill();
926                let _ = child.wait();
927                break None;
928            }
929            Ok(None) => std::thread::sleep(Duration::from_millis(50)),
930            Err(_) => {
931                let _ = child.kill();
932                let _ = child.wait();
933                break None;
934            }
935        }
936    };
937    let (sig, tail) = reader.join().unwrap_or_default();
938    let sig = sig.filter(|s| s.len() == 64 && s.bytes().all(|b| b.is_ascii_hexdigit()));
939    match (&status, &sig) {
940        (Some(st), Some(_)) if st.success() => sig,
941        (None, _) => {
942            eprintln!(
943                "fdl join: model-sig probe killed after {}s — a binary built \
944                 against a flodl that predates the probe runs its whole main \
945                 here; joining without a signature (`--no-sig-probe` or \
946                 `join.sig_probe: false` silences this)",
947                MODEL_SIG_PROBE_TIMEOUT.as_secs(),
948            );
949            None
950        }
951        (Some(st), _) if !st.success() => {
952            // The loud one: rank children re-enter this binary with these
953            // arguments, so this failure is what the cohort would see
954            // AFTER formation.
955            eprintln!(
956                "fdl join: WARNING: the training binary exited with {} under \
957                 the model-sig probe — rank children re-enter it with the \
958                 same arguments after admission, so if this failure is real \
959                 it takes the cohort's formation with it. Check: {} {}",
960                st.code().map_or("a signal".to_string(), |c| c.to_string()),
961                bin.display(),
962                args.join(" "),
963            );
964            None
965        }
966        _ => {
967            // Exit 0 and no signature has two very different causes: a
968            // binary older than the probe contract, or one that failed
969            // before reaching `Trainer::run` while still exiting 0.
970            // Guessing the first was wrong often enough to matter (a
971            // read-only project dir defeats a binary that creates an
972            // output directory in main, which is the walk-in's normal
973            // condition), so quote what it said and let the reader
974            // judge.
975            eprintln!(
976                "fdl join: model-sig probe exited 0 without printing a \
977                 signature; joining without one — the formation-time check \
978                 still applies. Either the binary predates the probe, or it \
979                 failed before reaching the trainer (check its output above, \
980                 and that this box can write wherever it writes).",
981            );
982            if !tail.is_empty() {
983                eprintln!("fdl join: last lines of the probe's output:");
984                for line in &tail {
985                    eprintln!("    {line}");
986                }
987            }
988            None
989        }
990    }
991}
992
993/// Reserve a loopback port for the tunnel's local end: bind :0, read
994/// the assignment, release. The tiny bind-to-ssh race is absorbed by
995/// the retry loop around each attempt.
996fn pick_local_port() -> Result<u16, String> {
997    let listener =
998        TcpListener::bind("127.0.0.1:0").map_err(|e| format!("reserve local tunnel port: {e}"))?;
999    let port = listener
1000        .local_addr()
1001        .map_err(|e| format!("reserve local tunnel port: {e}"))?
1002        .port();
1003    Ok(port)
1004}
1005
1006/// Assemble the tunnel command: `ssh -N -T` + user options first (they
1007/// win — OpenSSH takes the first value it sees per key) + flodl's
1008/// non-interactive defaults + the `-L` forward. Returned as argv for
1009/// testability.
1010fn build_tunnel_argv(
1011    ssh: &SshConfig,
1012    local_port: u16,
1013    controller_host: &str,
1014    controller_port: u16,
1015) -> Vec<String> {
1016    let mut argv: Vec<String> = vec!["ssh".into(), "-N".into(), "-T".into()];
1017    if let Some(warning) = crate::cluster::batchmode_override_warning(
1018        &ssh.options,
1019        ssh.target.as_deref().unwrap_or("?"),
1020    ) {
1021        eprintln!("{warning}");
1022    }
1023    for opt in &ssh.options {
1024        argv.push("-o".into());
1025        argv.push(opt.clone());
1026    }
1027    if let Some(port) = ssh.port {
1028        argv.push("-p".into());
1029        argv.push(port.to_string());
1030    }
1031    if let Some(user) = ssh.user.as_deref() {
1032        argv.push("-l".into());
1033        argv.push(user.to_string());
1034    }
1035    if let Some(id) = ssh.identity_file.as_deref() {
1036        argv.push("-i".into());
1037        argv.push(id.to_string());
1038    }
1039    // BatchMode: never hang on a prompt (a passphrase prompt inside a
1040    // systemd unit wedges forever). ExitOnForwardFailure: a forward the
1041    // remote refuses (permitopen mismatch) must kill ssh, not leave a
1042    // tunnel that black-holes the dial. ServerAlive: a silently dead
1043    // link tears the agent down instead of hanging the run.
1044    argv.push("-o".into());
1045    argv.push("BatchMode=yes".into());
1046    argv.push("-o".into());
1047    argv.push("ExitOnForwardFailure=yes".into());
1048    argv.push("-o".into());
1049    argv.push("ServerAliveInterval=30".into());
1050    argv.push("-L".into());
1051    argv.push(format!(
1052        "127.0.0.1:{local_port}:{controller_host}:{controller_port}"
1053    ));
1054    argv.push(ssh.target.clone().unwrap_or_default());
1055    argv
1056}
1057
1058/// Block until ssh's local forward accepts (auth done, listener bound)
1059/// or the budget runs out. A probe connection that reaches the far mux
1060/// and immediately EOFs is by-design harmless (the dispatcher drops
1061/// pre-magic EOFs and keeps serving). An early ssh exit is the loud
1062/// path: auth or forward failure, with ssh's own stderr right above.
1063fn wait_tunnel_ready(child: &mut Child, local_port: u16) -> Result<(), String> {
1064    let deadline = Instant::now() + TUNNEL_READY_BUDGET;
1065    let addr = std::net::SocketAddr::from(([127, 0, 0, 1], local_port));
1066    loop {
1067        if let Ok(Some(status)) = child.try_wait() {
1068            return Err(format!(
1069                "ssh tunnel exited ({status}) before the forward came up — \
1070                 see its output above (auth failure, or the remote refused \
1071                 the forward)"
1072            ));
1073        }
1074        if let Ok(probe) = TcpStream::connect_timeout(&addr, Duration::from_millis(500)) {
1075            drop(probe);
1076            return Ok(());
1077        }
1078        if Instant::now() >= deadline {
1079            return Err(format!(
1080                "ssh tunnel did not come up within {}s (local port \
1081                 {local_port} never accepted)",
1082                TUNNEL_READY_BUDGET.as_secs(),
1083            ));
1084        }
1085        std::thread::sleep(Duration::from_millis(200));
1086    }
1087}
1088
1089#[cfg(test)]
1090mod tests {
1091    use super::*;
1092
1093    fn no_flags() -> JoinArgs {
1094        JoinArgs {
1095            controller: None,
1096            ssh: None,
1097            identity: None,
1098            token: None,
1099            bin: None,
1100            source: None,
1101            source_cwd: None,
1102            source_build: None,
1103            source_bin: None,
1104            libtorch: None,
1105            host: None,
1106            devices: None,
1107            persist: false,
1108            data_path: None,
1109            data_source: None,
1110            gpu_ram_share: None,
1111            no_sig_probe: false,
1112        }
1113    }
1114
1115    fn full_block() -> WorkerJoin {
1116        WorkerJoin {
1117            controller: Some("10.0.0.9:9000".into()),
1118            ssh: Some(SshConfig {
1119                target: Some("bastion".into()),
1120                port: Some(2222),
1121                user: Some("join-user".into()),
1122                identity_file: Some("/etc/flodl/join_key".into()),
1123                options: vec!["StrictHostKeyChecking=accept-new".into()],
1124            }),
1125            token: Some("aa".repeat(16)),
1126            bin: Some("target/release/train".into()),
1127            source: None,
1128            libtorch: Some("auto".into()),
1129            host: Some("worker-7".into()),
1130            devices: Some(vec![0, 1]),
1131            persist: true,
1132            args: vec!["--model".into(), "lenet".into()],
1133            data_path: Some("/flodl/data".into()),
1134            data_source: Some("sshfs://flodl@ctrl:/srv/data".into()),
1135            gpu_ram_share: Some(0.5),
1136            sig_probe: None,
1137        }
1138    }
1139
1140    /// A block that builds its binary instead of naming one.
1141    fn source_block() -> WorkerJoin {
1142        WorkerJoin {
1143            source: Some(WorkerSource {
1144                from: "rsync://exa:/home/op/rdl".into(),
1145                cwd: Some("ddp-bench".into()),
1146                build: Some("cargo build --release --bin ddp-bench".into()),
1147                bin: Some("target/release/ddp-bench".into()),
1148            }),
1149            bin: None,
1150            ..full_block()
1151        }
1152    }
1153
1154    #[test]
1155    fn flags_win_over_the_config_block() {
1156        let cli = JoinArgs {
1157            controller: Some("exa".into()),
1158            ssh: Some("op@front:22".into()),
1159            identity: Some("/tmp/id".into()),
1160            token: Some("bb".repeat(16)),
1161            bin: Some("other/bin".into()),
1162            libtorch: Some("cu128".into()),
1163            host: Some("pascal".into()),
1164            devices: Some("2".into()),
1165            persist: false,
1166            data_path: Some("/mnt/corpus".into()),
1167            data_source: Some("sshfs://exa/mnt/corpus".into()),
1168            ..no_flags()
1169        };
1170        let tail: Vec<String> = vec!["--epochs".into(), "3".into()];
1171        let eff = resolve_effective(&cli, Some(&tail), Some(full_block()), "localbox").unwrap();
1172        assert_eq!(eff.controller_host, "exa");
1173        assert_eq!(eff.controller_port, DEFAULT_CONTROLLER_PORT);
1174        assert!(!eff.controller_defaulted);
1175        let ssh = eff.ssh.as_ref().unwrap();
1176        assert_eq!(ssh.target.as_deref(), Some("front"));
1177        assert_eq!(ssh.user.as_deref(), Some("op"));
1178        assert_eq!(ssh.port, Some(22));
1179        // Compact --ssh keeps the block's options; --identity wins last.
1180        assert_eq!(ssh.identity_file.as_deref(), Some("/tmp/id"));
1181        assert_eq!(
1182            ssh.options,
1183            vec!["StrictHostKeyChecking=accept-new".to_string()]
1184        );
1185        assert_eq!(eff.token.as_deref(), Some("bb".repeat(16).as_str()));
1186        assert_eq!(eff.bin, BinSource::Given("other/bin".into()));
1187        assert_eq!(eff.libtorch_spec.as_deref(), Some("cu128"));
1188        assert_eq!(eff.host, "pascal");
1189        assert_eq!(eff.devices, Some(vec![2]));
1190        // persist: block `true` sticks (the flag can only turn it on).
1191        assert!(eff.persist);
1192        // A `--` tail replaces the block's args.
1193        assert_eq!(eff.bin_args, vec!["--epochs".to_string(), "3".into()]);
1194        assert_eq!(eff.data_path.as_deref(), Some("/mnt/corpus"));
1195        assert_eq!(eff.data_source.as_deref(), Some("sshfs://exa/mnt/corpus"));
1196    }
1197
1198    #[test]
1199    fn block_fills_everything_the_flags_left_unset() {
1200        let eff = resolve_effective(&no_flags(), None, Some(full_block()), "localbox").unwrap();
1201        assert_eq!(eff.controller_host, "10.0.0.9");
1202        assert_eq!(eff.controller_port, 9000);
1203        let ssh = eff.ssh.as_ref().unwrap();
1204        assert_eq!(ssh.target.as_deref(), Some("bastion"));
1205        assert_eq!(ssh.identity_file.as_deref(), Some("/etc/flodl/join_key"));
1206        assert_eq!(eff.bin, BinSource::Given("target/release/train".into()));
1207        assert_eq!(eff.libtorch_spec.as_deref(), Some("auto"));
1208        assert_eq!(eff.host, "worker-7");
1209        assert_eq!(eff.devices, Some(vec![0, 1]));
1210        assert!(eff.persist);
1211        assert_eq!(eff.bin_args, vec!["--model".to_string(), "lenet".into()]);
1212        assert_eq!(eff.data_path.as_deref(), Some("/flodl/data"));
1213        assert_eq!(
1214            eff.data_source.as_deref(),
1215            Some("sshfs://flodl@ctrl:/srv/data"),
1216        );
1217        // The tunnel block's key and options carry to the data mount:
1218        // same box, same key (which that key must permit — see
1219        // `prepare::DataSpec::ssh`).
1220        let spec = eff.prepare_spec(None);
1221        assert_eq!(
1222            spec.data.ssh.and_then(|s| s.identity_file.as_deref()),
1223            Some("/etc/flodl/join_key"),
1224        );
1225    }
1226
1227    #[test]
1228    fn a_source_block_becomes_a_source_spec_carrying_the_same_key() {
1229        let eff = resolve_effective(&no_flags(), None, Some(source_block()), "localbox").unwrap();
1230        let spec = eff.prepare_spec(None);
1231        let source = spec.source.expect("a source block yields a source spec");
1232        assert_eq!(source.from, "rsync://exa:/home/op/rdl");
1233        assert_eq!(source.cwd, Some("ddp-bench"));
1234        assert_eq!(source.bin, Some("target/release/ddp-bench"));
1235        // The pull runs over the same hop the tunnel uses.
1236        assert_eq!(
1237            source.ssh.and_then(|s| s.identity_file.as_deref()),
1238            Some("/etc/flodl/join_key"),
1239        );
1240    }
1241
1242    #[test]
1243    fn naming_both_a_binary_and_a_source_is_a_loud_error() {
1244        // Not a precedence puzzle: a box handed both has no defensible
1245        // answer, so it must be told rather than guessed at.
1246        let block = WorkerJoin {
1247            bin: Some("target/release/train".into()),
1248            ..source_block()
1249        };
1250        let err = resolve_effective(&no_flags(), None, Some(block), "x").unwrap_err();
1251        assert!(err.contains("both name"), "got: {err}");
1252    }
1253
1254    #[test]
1255    fn a_source_flag_keeps_the_blocks_other_source_fields() {
1256        // Same shape as the compact `--ssh`: the flag carries the
1257        // transport, the block still answers for the rest.
1258        let cli = JoinArgs {
1259            source: Some("file:///mnt/rdl".into()),
1260            ..no_flags()
1261        };
1262        let eff = resolve_effective(&cli, None, Some(source_block()), "x").unwrap();
1263        assert_eq!(
1264            eff.bin,
1265            BinSource::Build(WorkerSource {
1266                from: "file:///mnt/rdl".into(),
1267                cwd: Some("ddp-bench".into()),
1268                build: Some("cargo build --release --bin ddp-bench".into()),
1269                bin: Some("target/release/ddp-bench".into()),
1270            }),
1271        );
1272    }
1273
1274    #[test]
1275    fn a_source_with_no_artifact_is_legal_because_a_manifest_may_name_it() {
1276        // The controller's published tree carries a run manifest, and that
1277        // manifest is the authority. Refusing here would make every worker
1278        // config repeat what the publish already said.
1279        let cli = JoinArgs {
1280            source: Some("file:///mnt/rdl".into()),
1281            ..no_flags()
1282        };
1283        let eff = resolve_effective(&cli, None, None, "x").unwrap();
1284        assert_eq!(
1285            eff.bin,
1286            BinSource::Build(WorkerSource {
1287                from: "file:///mnt/rdl".into(),
1288                cwd: None,
1289                build: None,
1290                bin: None,
1291            }),
1292        );
1293    }
1294
1295    #[test]
1296    fn a_source_detail_flag_with_no_source_is_a_loud_error() {
1297        // Silently dropping it would leave the operator with a run that
1298        // ignored what they typed, which is the failure `--`-forwarded
1299        // options already taught this CLI once.
1300        let cli = JoinArgs {
1301            bin: Some("t/bin".into()),
1302            source_cwd: Some("ddp-bench".into()),
1303            ..no_flags()
1304        };
1305        let err = resolve_effective(&cli, None, None, "x").unwrap_err();
1306        assert!(err.contains("--source-cwd"), "got: {err}");
1307        assert!(err.contains("no source"), "got: {err}");
1308    }
1309
1310    #[test]
1311    fn defaults_are_loopback_hostname_and_all_devices() {
1312        let cli = JoinArgs {
1313            bin: Some("t/bin".into()),
1314            ..no_flags()
1315        };
1316        let eff = resolve_effective(&cli, None, None, "localbox").unwrap();
1317        assert_eq!(eff.controller_host, "127.0.0.1");
1318        assert_eq!(eff.controller_port, DEFAULT_CONTROLLER_PORT);
1319        assert!(eff.controller_defaulted);
1320        assert!(eff.ssh.is_none());
1321        assert!(eff.token.is_none());
1322        assert_eq!(eff.host, "localbox");
1323        assert_eq!(eff.devices, None);
1324        assert!(!eff.persist);
1325        assert!(eff.bin_args.is_empty());
1326        // No data fields: prepare checks nothing and ships nothing, so
1327        // the training binary keeps its own default.
1328        assert!(eff.data_path.is_none());
1329        assert!(eff.data_source.is_none());
1330    }
1331
1332    #[test]
1333    fn an_explicit_empty_tail_clears_the_block_args() {
1334        // `fdl join --` = "this run takes no arguments" — it must
1335        // replace the block's list, not fall back to it.
1336        let eff =
1337            resolve_effective(&no_flags(), Some(&[]), Some(full_block()), "localbox").unwrap();
1338        assert!(eff.bin_args.is_empty());
1339    }
1340
1341    #[test]
1342    fn identity_without_an_ssh_hop_is_a_loud_error() {
1343        let cli = JoinArgs {
1344            identity: Some("/tmp/id".into()),
1345            bin: Some("t/bin".into()),
1346            ..no_flags()
1347        };
1348        let err = resolve_effective(&cli, None, None, "x").unwrap_err();
1349        assert!(err.contains("ssh hop"), "got: {err}");
1350    }
1351
1352    #[test]
1353    fn missing_bin_is_a_loud_error() {
1354        let err = resolve_effective(&no_flags(), None, None, "x").unwrap_err();
1355        assert!(err.contains("--bin"), "got: {err}");
1356        assert!(err.contains("join.bin"), "got: {err}");
1357    }
1358
1359    #[test]
1360    fn ssh_implies_the_loopback_controller_without_a_note() {
1361        let cli = JoinArgs {
1362            ssh: Some("join@ctrl".into()),
1363            bin: Some("t/bin".into()),
1364            ..no_flags()
1365        };
1366        let eff = resolve_effective(&cli, None, None, "x").unwrap();
1367        assert_eq!(eff.controller_host, "127.0.0.1");
1368        assert_eq!(eff.controller_port, DEFAULT_CONTROLLER_PORT);
1369        assert!(
1370            !eff.controller_defaulted,
1371            "tunnel loopback is the convention"
1372        );
1373    }
1374
1375    #[test]
1376    fn block_ssh_without_target_is_a_loud_error() {
1377        let block = WorkerJoin {
1378            ssh: Some(SshConfig::default()),
1379            bin: Some("t/bin".into()),
1380            ..WorkerJoin::default()
1381        };
1382        let err = resolve_effective(&no_flags(), None, Some(block), "x").unwrap_err();
1383        assert!(err.contains("target"), "got: {err}");
1384    }
1385
1386    #[test]
1387    fn spec_parsers_cover_their_shapes() {
1388        assert_eq!(
1389            parse_host_port("exa").unwrap(),
1390            ("exa".to_string(), DEFAULT_CONTROLLER_PORT),
1391        );
1392        assert_eq!(
1393            parse_host_port("exa:9000").unwrap(),
1394            ("exa".to_string(), 9000)
1395        );
1396        assert!(parse_host_port(":9000").is_err());
1397        assert!(parse_host_port("exa:banana").is_err());
1398
1399        let ssh = parse_ssh_spec("join@ctrl:2222").unwrap();
1400        assert_eq!(ssh.target.as_deref(), Some("ctrl"));
1401        assert_eq!(ssh.user.as_deref(), Some("join"));
1402        assert_eq!(ssh.port, Some(2222));
1403        let bare = parse_ssh_spec("ctrl").unwrap();
1404        assert_eq!(bare.target.as_deref(), Some("ctrl"));
1405        assert_eq!(bare.user, None);
1406        assert_eq!(bare.port, None);
1407        assert!(parse_ssh_spec("@ctrl").is_err());
1408        assert!(parse_ssh_spec("join@").is_err());
1409        assert!(parse_ssh_spec("ctrl:pear").is_err());
1410
1411        assert_eq!(parse_devices("0,1").unwrap(), Some(vec![0, 1]));
1412        assert_eq!(parse_devices(" 2 ").unwrap(), Some(vec![2]));
1413        assert_eq!(parse_devices("all").unwrap(), None);
1414        assert!(parse_devices("0,x").is_err());
1415    }
1416
1417    /// The JSON field names are flodl's `AgentSpec` wire contract —
1418    /// this test IS the cross-crate compatibility lock (flodl-cli is
1419    /// The cache key binds exactly what the probe's answer depends on:
1420    /// same binary + same args is a hit; touched binary, different
1421    /// args, or a missing file is not.
1422    #[test]
1423    fn probe_recipe_digest_binds_binary_identity_and_args() {
1424        let dir = std::env::temp_dir().join(format!("fdl-sig-digest-test-{}", std::process::id()));
1425        std::fs::create_dir_all(&dir).unwrap();
1426        let bin = dir.join("train");
1427        std::fs::write(&bin, b"v1").unwrap();
1428        let args = vec!["--model".to_string(), "lenet".to_string()];
1429        let base = probe_recipe_digest(&bin, &args).unwrap();
1430        assert_eq!(probe_recipe_digest(&bin, &args).unwrap(), base);
1431        assert_ne!(
1432            probe_recipe_digest(&bin, &["--model".to_string(), "resnet".to_string()]).unwrap(),
1433            base,
1434            "args are part of the recipe (a re-publish must re-probe)"
1435        );
1436        // A rebuild: same path, new content — size or mtime moves.
1437        std::fs::write(&bin, b"v2 longer").unwrap();
1438        assert_ne!(
1439            probe_recipe_digest(&bin, &args).unwrap(),
1440            base,
1441            "a rebuilt binary must re-probe"
1442        );
1443        assert_eq!(probe_recipe_digest(&dir.join("absent"), &args), None);
1444        let _ = std::fs::remove_dir_all(&dir);
1445    }
1446
1447    /// The probe's child-process contract, driven with shell-script
1448    /// stand-ins for the training binary: the prefixed line is found
1449    /// among main-body noise, and every failure mode (no line, bad
1450    /// line, non-zero exit) degrades to `None` rather than erroring —
1451    /// the hello then gates nothing and formation stays the backstop.
1452    /// (The timeout path is deliberately not exercised: it is a 120s
1453    /// wait by construction.)
1454    #[cfg(unix)]
1455    #[test]
1456    fn model_sig_probe_parses_the_line_and_absorbs_failures() {
1457        // Each case is `/bin/sh -c <body>`, not a script this test writes
1458        // and then execs. Writing an executable and running it from a
1459        // multithreaded process is a race: a `Command` spawned anywhere
1460        // else in the binary during the window where the write fd is open
1461        // inherits that fd across the fork, and the exec then fails with
1462        // ETXTBSY. Reproduced at about 1 run in 60 -- the probe returned
1463        // None and the message said "Text file busy", which reads like a
1464        // parsing bug and is not one. /bin/sh is never opened for writing.
1465        let sh = PathBuf::from("/bin/sh");
1466        let run = |body: String| model_sig_probe(&sh, None, &["-c".to_string(), body], None);
1467        let sig = "ab".repeat(32);
1468        assert_eq!(
1469            run(format!("echo main noise; echo '{MODEL_SIG_LINE}{sig}'")),
1470            Some(sig),
1471        );
1472        assert_eq!(run("exit 0".to_string()), None);
1473        assert_eq!(run("exit 3".to_string()), None);
1474        assert_eq!(run(format!("echo '{MODEL_SIG_LINE}not-hex-at-all'")), None,);
1475    }
1476
1477    /// zero-dep on flodl by design, so the shape is asserted literally;
1478    /// flodl's `agent_spec_round_trips_through_hex` holds the other end).
1479    #[test]
1480    fn agent_spec_shape_is_the_wire_contract() {
1481        let cli = JoinArgs {
1482            token: Some("ab".repeat(16)),
1483            bin: Some("t/bin".into()),
1484            host: Some("pascal".into()),
1485            devices: Some("0,1".into()),
1486            gpu_ram_share: Some(0.5),
1487            ..no_flags()
1488        };
1489        let eff = resolve_effective(&cli, None, None, "x").unwrap();
1490        let prepared = Prepared {
1491            data_path: Some(PathBuf::from("/flodl/data")),
1492            run_id: Some("a1b2c3d4e5f60718".to_string()),
1493            ..Prepared::default()
1494        };
1495        let hex = agent_spec_hex(
1496            &eff,
1497            ("127.0.0.1", 40123),
1498            "builds/sm61-sm120",
1499            &prepared,
1500            Some(&"cd".repeat(32)),
1501        );
1502        let bytes: Vec<u8> = (0..hex.len())
1503            .step_by(2)
1504            .map(|i| u8::from_str_radix(&hex[i..i + 2], 16).unwrap())
1505            .collect();
1506        let spec: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
1507        assert_eq!(spec["host"], "pascal");
1508        assert_eq!(spec["controller_host"], "127.0.0.1");
1509        assert_eq!(spec["controller_port"], 40123);
1510        assert_eq!(spec["salt_hex"], "ab".repeat(16));
1511        assert_eq!(spec["local_devices"], serde_json::json!([0, 1]));
1512        assert_eq!(spec["libtorch"], "builds/sm61-sm120");
1513        assert_eq!(spec["data_path"], "/flodl/data");
1514        assert_eq!(spec["run_id"], "a1b2c3d4e5f60718");
1515        assert_eq!(spec["gpu_ram_share"], 0.5);
1516        assert_eq!(spec["model_sig_hex"], "cd".repeat(32));
1517        // Optional fields are OMITTED when unset, never null — flodl's
1518        // serde defaults own the fallbacks.
1519        let open = {
1520            let cli = JoinArgs {
1521                bin: Some("t/bin".into()),
1522                ..no_flags()
1523            };
1524            let eff = resolve_effective(&cli, None, None, "cloud-1").unwrap();
1525            agent_spec_hex(&eff, ("10.0.0.1", 1337), "", &Prepared::default(), None)
1526        };
1527        let bytes: Vec<u8> = (0..open.len())
1528            .step_by(2)
1529            .map(|i| u8::from_str_radix(&open[i..i + 2], 16).unwrap())
1530            .collect();
1531        let spec: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
1532        assert!(spec.get("salt_hex").is_none());
1533        assert!(spec.get("local_devices").is_none());
1534        assert!(spec.get("dataset_sig_hex").is_none());
1535        // A box that declares no source root must ship no key at all:
1536        // an empty string here would point every rank at the process cwd.
1537        assert!(spec.get("data_path").is_none());
1538        // Same rule for the run id: a `--bin` box carries none, and an
1539        // absent key is what gates nothing at admission.
1540        assert!(spec.get("run_id").is_none());
1541        // And for the RAM share: an absent key is what lets the
1542        // envelope's cluster-scope default stand.
1543        assert!(spec.get("gpu_ram_share").is_none());
1544    }
1545
1546    #[test]
1547    fn tunnel_argv_orders_user_options_before_the_defaults() {
1548        let ssh = SshConfig {
1549            target: Some("ctrl".into()),
1550            port: Some(2222),
1551            user: Some("join-user".into()),
1552            identity_file: Some("/etc/flodl/join_key".into()),
1553            options: vec!["ServerAliveInterval=5".into()],
1554        };
1555        let argv = build_tunnel_argv(&ssh, 40123, "127.0.0.1", 1337);
1556        assert_eq!(argv[0], "ssh");
1557        assert!(argv.contains(&"-N".to_string()));
1558        assert!(argv.contains(&"BatchMode=yes".to_string()));
1559        assert!(argv.contains(&"ExitOnForwardFailure=yes".to_string()));
1560        // First -o value wins in OpenSSH: the user's override must
1561        // appear before flodl's default of the same key.
1562        let user_pos = argv
1563            .iter()
1564            .position(|a| a == "ServerAliveInterval=5")
1565            .unwrap();
1566        let default_pos = argv
1567            .iter()
1568            .position(|a| a == "ServerAliveInterval=30")
1569            .unwrap();
1570        assert!(user_pos < default_pos);
1571        assert!(argv.contains(&"127.0.0.1:40123:127.0.0.1:1337".to_string()));
1572        assert_eq!(argv.last().map(String::as_str), Some("ctrl"));
1573        let p = argv.iter().position(|a| a == "-p").unwrap();
1574        assert_eq!(argv[p + 1], "2222");
1575        let l = argv.iter().position(|a| a == "-l").unwrap();
1576        assert_eq!(argv[l + 1], "join-user");
1577        let i = argv.iter().position(|a| a == "-i").unwrap();
1578        assert_eq!(argv[i + 1], "/etc/flodl/join_key");
1579    }
1580
1581    #[test]
1582    fn wait_tunnel_ready_sees_a_live_listener_and_a_dead_child() {
1583        // A child that exits immediately stands in for a failed ssh.
1584        let mut dead = Command::new("true").spawn().unwrap();
1585        std::thread::sleep(Duration::from_millis(50));
1586        let err = wait_tunnel_ready(&mut dead, 1).unwrap_err();
1587        assert!(err.contains("before the forward came up"), "got: {err}");
1588
1589        // A live listener on the reserved port = ready, child untouched.
1590        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
1591        let port = listener.local_addr().unwrap().port();
1592        let mut slow = Command::new("sleep").arg("5").spawn().unwrap();
1593        assert!(wait_tunnel_ready(&mut slow, port).is_ok());
1594        let _ = slow.kill();
1595        let _ = slow.wait();
1596    }
1597
1598    #[test]
1599    fn hex_encode_is_lowercase_bytewise() {
1600        assert_eq!(hex_encode(b"\x00\xff\x10"), "00ff10");
1601        assert_eq!(hex_encode(b"{}"), "7b7d");
1602    }
1603}