flodl-cli 0.8.0

libtorch manager and GPU diagnostic tool for Rust deep learning
Documentation
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
1275
1276
1277
1278
1279
1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305
1306
1307
1308
1309
1310
1311
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329
1330
1331
1332
1333
1334
1335
1336
1337
1338
1339
1340
1341
1342
1343
1344
1345
1346
1347
1348
1349
1350
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
1366
1367
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377
1378
1379
1380
1381
1382
1383
1384
1385
1386
1387
1388
1389
1390
1391
1392
1393
1394
1395
1396
1397
1398
1399
1400
1401
1402
1403
1404
1405
1406
1407
1408
1409
1410
1411
1412
1413
1414
1415
1416
1417
1418
1419
1420
1421
1422
1423
1424
1425
1426
1427
1428
1429
1430
1431
1432
1433
1434
1435
1436
1437
1438
1439
1440
1441
1442
1443
1444
1445
1446
1447
1448
1449
1450
1451
1452
1453
1454
1455
1456
1457
1458
1459
1460
1461
1462
1463
1464
1465
1466
1467
1468
1469
1470
1471
1472
1473
1474
1475
1476
1477
1478
1479
1480
1481
1482
1483
1484
1485
1486
1487
1488
1489
1490
1491
1492
1493
1494
1495
1496
1497
1498
1499
1500
1501
1502
1503
1504
1505
1506
1507
1508
1509
1510
1511
1512
1513
1514
1515
1516
1517
1518
1519
1520
1521
1522
1523
1524
1525
1526
1527
1528
1529
1530
1531
1532
1533
1534
1535
1536
1537
1538
1539
1540
1541
1542
1543
1544
1545
1546
1547
1548
1549
1550
1551
1552
1553
1554
1555
1556
1557
1558
1559
1560
1561
1562
1563
1564
1565
1566
1567
1568
1569
1570
1571
1572
1573
1574
1575
1576
1577
1578
1579
1580
1581
1582
1583
1584
1585
1586
1587
1588
1589
1590
1591
1592
1593
1594
1595
1596
1597
1598
1599
1600
1601
1602
1603
//! `fdl join` — self-deploy this box as a dial-in worker.
//!
//! The worker-side walk-in for a discovery window (see docs/ddp/01-reference.md):
//! dials the controller's mux port, offers this host's GPUs, and — once
//! admitted — the training binary takes over in agent role (it joins,
//! then spawns and supervises this host's relay and rank children; see
//! flodl's `distributed::launcher::agent`). fdl-cli stays protocol-blind:
//! its whole job is orchestration —
//!
//!   1. resolve settings (flags over the fdl.yml `join:` block),
//!   2. prepare this box ([`crate::prepare`]): the GPU gate, the dataset
//!      source root, the node-local directories the data plane writes —
//!      all of it BEFORE the dial, because admission starts a window
//!      deadline,
//!   3. optionally bring up an ssh `-L` forward of the controller port
//!      (the guardrailed-sshd trust path: reachability = authentication),
//!   4. synthesize the agent bootstrap spec into the binary's
//!      environment (`FLODL_INTERNAL_AGENT_JSON`, hex-encoded JSON — the
//!      same envelope cluster fan-out ships),
//!   5. run + supervise the binary, and in `--persist` mode re-dial
//!      with backoff when it exits (the systemd / golden-image loop).
//!
//! The spec (which may carry the pre-shared session token) rides the
//! child's ENVIRONMENT, never argv — owner-readable via
//! `/proc/<pid>/environ` instead of world-readable via `ps`, the same
//! salt hygiene as the launcher's fan-out.

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;

/// Agent bootstrap env var — must match flodl's
/// `distributed::launcher::ENV_AGENT_JSON` (the field names in the hex
/// JSON must match `AgentSpec`; locked by `agent_spec_shape_is_the_wire_contract`).
const ENV_AGENT_JSON: &str = "FLODL_INTERNAL_AGENT_JSON";

/// Exit code for a failure retrying cannot fix: no usable GPU, a spec
/// that does not parse, a directory that cannot be created, a missing
/// binary. Distinct from 1 (a transient failure, one-shot) so a
/// fleet can act on the difference without parsing stderr:
///
/// ```ini
/// # /etc/systemd/system/flodl-join.service
/// Restart=always
/// RestartPreventExitStatus=2   # stop hot-looping a misprovisioned box
/// FailureAction=poweroff       # ... and self-deprovision it
/// ```
///
/// fdl deliberately does not power a box off itself: the decision belongs
/// to whatever owns the instance's lifecycle, and 2 is how it hears about it.
pub const EXIT_PERMANENT: i32 = 2;

/// How long the ssh forward gets to come up (auth + local bind).
const TUNNEL_READY_BUDGET: Duration = Duration::from_secs(20);

/// `--persist` re-dial backoff: floor, cap, and the attempt duration
/// past which the backoff resets to the floor (the agent ran a real
/// stint, so the next failure is a fresh incident, not a hot loop).
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);

/// Run `fdl join`. `bin_tail` is everything after the command line's
/// standalone `--`: the training binary's own arguments, forwarded
/// verbatim (rank children re-enter the binary with them). `None` =
/// no `--` was given (the config block's `args:` applies); a present
/// but empty tail is an explicit "no arguments".
///
/// Exit code: the agent's own exit code (one-shot); [`EXIT_PERMANENT`]
/// for a failure retrying cannot fix; 1 for a transient one. `--persist`
/// re-dials through transient failures and agent exits, and returns only
/// on a permanent one.
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)"
            )),
        );
    }

    // A binary named as a path must exist NOW — a persist loop retrying
    // a missing path forever helps nobody. A binary built from source
    // cannot be checked here: it does not exist until the attempt has
    // fetched and compiled the tree.
    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;
    }

    // Local active libtorch (honors FDL_LIBTORCH_CASE), anchored on the
    // project root the config walk found: its lib/ rides
    // LD_LIBRARY_PATH on the child, and its variant label rides the
    // join hello. Absent (fdl running outside a project) the child env
    // is left untouched — the binary may carry an rpath or the caller's
    // environment already provides the libs.
    let libtorch = resolve_local_libtorch(project_root.as_deref());

    // Model-sig probe cache, living exactly as long as the persist loop:
    // one slot, keyed by a digest of the probe recipe (binary identity +
    // args). An idle fleet re-dials every BACKOFF_MAX forever, and
    // without this every re-dial would rebuild the model on CPU — or,
    // for a binary that predates the probe, pay the full probe timeout
    // — for an unchanged binary.
    let mut sig_cache: Option<(u64, Option<String>)> = None;

    let mut backoff = BACKOFF_MIN;
    loop {
        let started = Instant::now();
        // What happened, phrased for the re-dial line. Every branch that
        // is not re-dialable returns from here.
        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() {
                    // The whole point of the class: a box that cannot be
                    // fixed by waiting must stop, not hot-loop.
                    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);
    }
}

// ---------------------------------------------------------------------------
// Settings resolution
// ---------------------------------------------------------------------------

/// The fully resolved join recipe: flags merged over the fdl.yml
/// `join:` block, every default applied.
#[derive(Debug)]
struct Effective {
    /// Controller mux address. Under `ssh` this is the address as seen
    /// FROM the ssh host (the `-L` forward's far end).
    controller_host: String,
    controller_port: u16,
    /// True when neither flags nor config named a controller and the
    /// loopback convention default applied (worth a stderr note).
    controller_defaulted: bool,
    /// Tunnel hop; `None` = direct dial.
    ssh: Option<SshConfig>,
    /// Pre-shared session credential (hex). `None` = open admission.
    token: Option<String>,
    /// How this box gets its training binary: a path to run as given, or
    /// a source to build. Exactly one, checked at resolution.
    bin: BinSource,
    /// libtorch variant to acquire; `None` keeps this box's active one.
    libtorch_spec: Option<String>,
    /// Logical host name in the join hello.
    host: String,
    /// Explicit CUDA device ids; `None` = all GPUs on this host.
    devices: Option<Vec<u8>>,
    persist: bool,
    /// The binary's own arguments.
    bin_args: Vec<String>,
    /// Dataset source root on this box (the mountpoint when
    /// `data_source` is set); `None` ships nothing to the ranks.
    data_path: Option<String>,
    /// Transport that establishes the source root, `<scheme>://<target>`.
    data_source: Option<String>,
    /// Integrated-GPU host-RAM share of this box; `None` ships nothing
    /// (the envelope's cluster-scope default, if any, then stands).
    gpu_ram_share: Option<f64>,
    /// Probe the binary for its model signature before each dial
    /// (default true; `--no-sig-probe` / `join.sig_probe: false`).
    sig_probe: bool,
}

/// Where this box's training binary comes from. `source:` and `bin:` are
/// mutually exclusive because they answer the same question, and keeping
/// the given-binary case a separate variant rather than a third kind of
/// source spec is what keeps an artifact-versus-source distinction out of
/// the source grammar.
#[derive(Debug, PartialEq, Eq)]
enum BinSource {
    /// A path on this box, run as given.
    Given(String),
    /// Fetched and built here.
    Build(WorkerSource),
}

impl Effective {
    /// Everything [`crate::prepare`] has to settle. The tunnel block
    /// rides along to both artifact specs: the data host and the source
    /// host are the controller box in the shape this exists for, so its
    /// key and options apply to them too.
    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(),
                }),
            },
        }
    }
}

/// Merge flags over the config block. Pure — all I/O (hostname, config
/// load) happens in the callers so this stays table-testable.
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());
    }

    // Tunnel hop: `--ssh [user@]host[:port]` replaces the block's
    // target/user/port but inherits its identity_file/options (not
    // expressible in the compact form); `--identity` wins last. A block
    // without `target:` is an authoring error — ssh needs a host.
    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());
    }

    // Controller: flag > block > loopback convention. Through a tunnel
    // the loopback default is THE convention (guardrailed sshd on the
    // controller box), no note needed; a bare loopback default deserves
    // one.
    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),
    };

    // The binary: a path to run, or a source to build. Flags win over
    // the block on each side, and naming both ways is an authoring error
    // rather than a precedence puzzle — a box that builds its own binary
    // and is also handed one has no defensible answer.
    let bin_path = cli.bin.clone().or(block.bin);
    let source = match (cli.source.clone(), block.source) {
        (Some(from), b) => Some(WorkerSource {
            from,
            // The compact `--source` flag carries only the transport, so
            // the rest keeps coming from the block unless its own flag
            // overrides it (same shape as `--ssh`).
            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())),
            // No artifact anywhere is legal: a published tree carries a
            // run manifest that names it, and that manifest is the
            // authority when it is there.
            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,
    };
    // A source flag with no source to attach to is an authoring error,
    // not something to drop on the floor: the operator meant it to change
    // the run.
    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,
    };

    // A `--` tail — even an empty one — REPLACES the block's args: the
    // args must match the run, so "explicitly none" must be sayable.
    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),
        // The flag only disables; `sig_probe: false` in yml is the
        // standing form of the same choice. Default on.
        sig_probe: !cli.no_sig_probe && block.sig_probe.unwrap_or(true),
    })
}

/// Load the top-level `join:` block from the PROJECT config (base
/// fdl.yml merged with the active env overlay when one is selected),
/// plus the directory it lives in (the project root — where libtorch/
/// is anchored). The walk steps over command-level fdl.ymls
/// ([`config::find_project_config`]): `fdl join` typically runs from
/// the command dir the training binary expects as cwd (e.g.
/// `ddp-bench/`), whose own fdl.yml is a command config that neither
/// carries a `join:` block nor marks the libtorch root. `Ok(None)`
/// root/block when there is no project at all — flags carry everything
/// then; a present-but-broken project config is a loud error, not a
/// silent fallback (the operator may be relying on `join.bin`).
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))
}

/// Parse `host[:port]`, default port [`DEFAULT_CONTROLLER_PORT`] —
/// same convention as `fdl status --addr`.
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)),
    }
}

/// Parse the compact tunnel spec `[user@]host[:port]` into an
/// [`SshConfig`] (target/user/port only; identity/options come from
/// the config block or `--identity`).
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(),
    })
}

/// Parse `--devices`: comma-separated CUDA ids, or `all` for
/// every GPU on this host (= unset).
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)
}

// ---------------------------------------------------------------------------
// Agent spec synthesis
// ---------------------------------------------------------------------------

/// Build the hex-encoded JSON payload for [`ENV_AGENT_JSON`]. Field
/// names ARE the wire contract with flodl's `AgentSpec` deserializer;
/// optional fields are omitted (serde defaults fill them).
///
/// The prepared data path travels this way rather than in the join hello
/// because the controller has nothing to say about it: it never
/// configured this host, and only this box knows where its source root
/// actually ended up. flodl's agent inserts it into the envelope its
/// rank children read, beside the same-shaped rewrite it already does
/// for the controller address.
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);
    }
    // Host-hardware truth, same as data_path: the agent writes it into
    // the envelope's host block, overriding any cluster-scope default
    // the controller stamped. Omitted when undeclared, so that default
    // stands.
    if let Some(share) = eff.gpu_ram_share {
        spec["gpu_ram_share"] = serde_json::json!(share);
    }
    // Probed from the binary itself, so admission can refuse a box
    // building a different model while it still costs only this box's
    // own attempt. Omitted when the probe was skipped or failed.
    if let Some(sig) = model_sig_hex {
        spec["model_sig_hex"] = serde_json::json!(sig);
    }
    hex_encode(spec.to_string().as_bytes())
}

/// Lowercase hex — flodl's `cluster::hex_decode` counterpart.
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
}

/// Active libtorch of this box, anchored on the project root when the
/// config walk found one (a command-dir cwd's `Context::resolve` would
/// stop a level too low); the plain context fallback covers project-less
/// setups (`~/.flodl`).
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)
}

/// `LD_LIBRARY_PATH` for the training binary, with this box's inherited
/// value appended.
///
/// The ordering inside is not ours to choose: on ROCm the system runtime
/// must precede libtorch's own lib dir, because libtorch-rocm bundles the
/// whole userspace ROCm stack and a bundle that disagrees with the host's
/// amdkfd driver segfaults the rank at its first GPU op. Prepending
/// unconditionally, which is what this did before, is the segfault
/// configuration on an AMD box.
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,
    }
}

// ---------------------------------------------------------------------------
// One attempt: tunnel up (optional), agent run, teardown
// ---------------------------------------------------------------------------

/// One full join attempt: prepare, tunnel, dial, supervise. Returns the
/// agent's exit code, or the classed reason preparation/orchestration
/// stopped.
///
/// Preparation comes first and every attempt re-runs it: admission
/// starts a window deadline, so a mount established after the dial burns
/// it, and re-running is how `--persist` becomes a provisioning loop.
///
/// The tunnel — when one is configured — lives exactly as long as the
/// attempt: rebuilt fresh each re-dial, so a half-dead forward can never
/// outlive the run it served.
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", &notes);
    let prepared = prepared?;

    // A built binary runs in the project directory inside the fetched
    // tree, which is where its own relative paths resolve; a given one
    // keeps fdl's cwd, exactly as before.
    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),
        // Neither combination is reachable — preparation returns a binary
        // exactly when it was given a source — so say so rather than
        // quietly preferring one and hiding a wiring inversion.
        (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(),
            )));
        }
    };

    // The run's arguments belong to the run. A published manifest replaces
    // whatever this box carried, because rank children re-enter the binary
    // with them: a standing fleet holding its own copy would train the new
    // run with the previous one's hyperparameters. Saying so out loud is
    // the difference between authority and a silent substitution.
    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,
    };

    // Model-signature probe, before the tunnel like everything else in
    // preparation: admission starts a window deadline, and the probe
    // runs the binary's whole main up to `Trainer::run`. Cached by
    // recipe digest across re-dials — the OUTCOME is cached, a failed
    // probe included, so an unchanged binary pays the probe (or its
    // timeout) once, not once per backoff tick. A rebuild changes the
    // mtime and a re-publish changes the args, so staleness invalidates
    // itself; the libtorch variant is deliberately not in the recipe
    // (the manifest — names, shapes, dtypes — is device-independent).
    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
                }
            },
            // The binary un-stat-able between resolution and here is a
            // race with a rebuild: probe uncached, next attempt keys.
            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| {
                    // No ssh on the box is a provisioning fact, not a
                    // passing condition.
                    Fail::Permanent(format!("spawn ssh tunnel: {e}"))
                })?;
            // Auth failure and an unreachable host are both possible
            // here and ssh does not let us tell them apart, so this
            // stays re-dialable: a wrong key hits the backoff cap and
            // keeps saying so, once a minute, loudly.
            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),
    };

    // Preparation is the authority on libtorch: it either acquired a
    // variant or carried the active one through, and the label it settles
    // on is what the join hello announces.
    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)
        // Children report under the logical roster name even when it
        // differs from `hostname` — same override fan-out applies.
        .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));
    }

    // A given path was checked before the loop and a built one was just
    // written, so a spawn failure here is the file itself: not
    // executable, wrong architecture, bad interpreter.
    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))
}

/// Digest of the probe recipe — everything the probe's answer depends
/// on: the binary's identity (path, mtime, size) and the arguments the
/// run enters it with. One u64 via std's SipHash rather than a
/// cryptographic hash: flodl-cli is zero-dep, the key lives one process
/// and guards a cache on a box that is trusted (not proven), and a
/// collision's worst case — a stale signature in the hello — lands on
/// the formation-time backstop. `None` when the binary cannot be
/// stat'ed (a race with a rebuild): probe uncached, key next attempt.
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())
}

/// How many trailing stdout lines the probe keeps as evidence when no
/// signature appears.
const PROBE_TAIL_LINES: usize = 8;

/// Probe-run marker env (flodl's launcher contract: with it set,
/// `Trainer::run` builds the model on CPU, prints the signature line
/// and exits — before auto-promote, before any cluster role).
const ENV_MODEL_SIG_PROBE: &str = "FLODL_INTERNAL_MODEL_SIG_PROBE";

/// Stdout line the probe scans for (main-body prints above it are
/// harmless — the prefix is the protocol, not the whole stream).
const MODEL_SIG_LINE: &str = "flodl-model-sig: ";

/// Ceiling on the probe re-run. A current flodl answers in the time its
/// main takes to reach `Trainer::run`; a binary built against a flodl
/// that predates the probe ignores the env and runs its whole main,
/// which is exactly what this bounds.
const MODEL_SIG_PROBE_TIMEOUT: Duration = Duration::from_secs(120);

/// Re-run the training binary as a model-signature probe and return the
/// 64-hex signature it prints.
///
/// Best-effort BY DESIGN: every failure degrades to `None` — the hello
/// then gates nothing and the formation-time handshake check stays the
/// backstop — but each failure mode says so, because one of them
/// (a non-zero exit) predicts the rank children failing the same way
/// after formation, with the same binary and the same arguments.
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;
        // Last few lines kept as evidence: when no signature turns up,
        // what the binary actually said beats anything we could infer
        // about why. Bounded so a chatty binary cannot grow this.
        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() => {
            // The loud one: rank children re-enter this binary with these
            // arguments, so this failure is what the cohort would see
            // AFTER formation.
            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
        }
        _ => {
            // Exit 0 and no signature has two very different causes: a
            // binary older than the probe contract, or one that failed
            // before reaching `Trainer::run` while still exiting 0.
            // Guessing the first was wrong often enough to matter (a
            // read-only project dir defeats a binary that creates an
            // output directory in main, which is the walk-in's normal
            // condition), so quote what it said and let the reader
            // judge.
            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
        }
    }
}

/// Reserve a loopback port for the tunnel's local end: bind :0, read
/// the assignment, release. The tiny bind-to-ssh race is absorbed by
/// the retry loop around each attempt.
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)
}

/// Assemble the tunnel command: `ssh -N -T` + user options first (they
/// win — OpenSSH takes the first value it sees per key) + flodl's
/// non-interactive defaults + the `-L` forward. Returned as argv for
/// testability.
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());
    }
    // BatchMode: never hang on a prompt (a passphrase prompt inside a
    // systemd unit wedges forever). ExitOnForwardFailure: a forward the
    // remote refuses (permitopen mismatch) must kill ssh, not leave a
    // tunnel that black-holes the dial. ServerAlive: a silently dead
    // link tears the agent down instead of hanging the run.
    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
}

/// Block until ssh's local forward accepts (auth done, listener bound)
/// or the budget runs out. A probe connection that reaches the far mux
/// and immediately EOFs is by-design harmless (the dispatcher drops
/// pre-magic EOFs and keeps serving). An early ssh exit is the loud
/// path: auth or forward failure, with ssh's own stderr right above.
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,
        }
    }

    /// A block that builds its binary instead of naming one.
    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));
        // Compact --ssh keeps the block's options; --identity wins last.
        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]));
        // persist: block `true` sticks (the flag can only turn it on).
        assert!(eff.persist);
        // A `--` tail replaces the block's args.
        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"),
        );
        // The tunnel block's key and options carry to the data mount:
        // same box, same key (which that key must permit — see
        // `prepare::DataSpec::ssh`).
        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"));
        // The pull runs over the same hop the tunnel uses.
        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() {
        // Not a precedence puzzle: a box handed both has no defensible
        // answer, so it must be told rather than guessed at.
        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() {
        // Same shape as the compact `--ssh`: the flag carries the
        // transport, the block still answers for the rest.
        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() {
        // The controller's published tree carries a run manifest, and that
        // manifest is the authority. Refusing here would make every worker
        // config repeat what the publish already said.
        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() {
        // Silently dropping it would leave the operator with a run that
        // ignored what they typed, which is the failure `--`-forwarded
        // options already taught this CLI once.
        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());
        // No data fields: prepare checks nothing and ships nothing, so
        // the training binary keeps its own default.
        assert!(eff.data_path.is_none());
        assert!(eff.data_source.is_none());
    }

    #[test]
    fn an_explicit_empty_tail_clears_the_block_args() {
        // `fdl join --` = "this run takes no arguments" — it must
        // replace the block's list, not fall back to it.
        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());
    }

    /// The JSON field names are flodl's `AgentSpec` wire contract —
    /// this test IS the cross-crate compatibility lock (flodl-cli is
    /// The cache key binds exactly what the probe's answer depends on:
    /// same binary + same args is a hit; touched binary, different
    /// args, or a missing file is not.
    #[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)"
        );
        // A rebuild: same path, new content — size or mtime moves.
        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);
    }

    /// The probe's child-process contract, driven with shell-script
    /// stand-ins for the training binary: the prefixed line is found
    /// among main-body noise, and every failure mode (no line, bad
    /// line, non-zero exit) degrades to `None` rather than erroring —
    /// the hello then gates nothing and formation stays the backstop.
    /// (The timeout path is deliberately not exercised: it is a 120s
    /// wait by construction.)
    #[cfg(unix)]
    #[test]
    fn model_sig_probe_parses_the_line_and_absorbs_failures() {
        // Each case is `/bin/sh -c <body>`, not a script this test writes
        // and then execs. Writing an executable and running it from a
        // multithreaded process is a race: a `Command` spawned anywhere
        // else in the binary during the window where the write fd is open
        // inherits that fd across the fork, and the exec then fails with
        // ETXTBSY. Reproduced at about 1 run in 60 -- the probe returned
        // None and the message said "Text file busy", which reads like a
        // parsing bug and is not one. /bin/sh is never opened for writing.
        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,);
    }

    /// zero-dep on flodl by design, so the shape is asserted literally;
    /// flodl's `agent_spec_round_trips_through_hex` holds the other end).
    #[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));
        // Optional fields are OMITTED when unset, never null — flodl's
        // serde defaults own the fallbacks.
        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());
        // A box that declares no source root must ship no key at all:
        // an empty string here would point every rank at the process cwd.
        assert!(spec.get("data_path").is_none());
        // Same rule for the run id: a `--bin` box carries none, and an
        // absent key is what gates nothing at admission.
        assert!(spec.get("run_id").is_none());
        // And for the RAM share: an absent key is what lets the
        // envelope's cluster-scope default stand.
        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()));
        // First -o value wins in OpenSSH: the user's override must
        // appear before flodl's default of the same key.
        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() {
        // A child that exits immediately stands in for a failed ssh.
        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}");

        // A live listener on the reserved port = ready, child untouched.
        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");
    }
}