Skip to main content

flodl_cli/
run.rs

1//! Command resolution and execution.
2//!
3//! Merges structured config sections into CLI arguments, resolves named
4//! command presets, and spawns the target process (directly or through
5//! Docker when a `docker:` service is declared).
6
7use std::collections::BTreeMap;
8use std::path::Path;
9use std::process::{ExitCode, Stdio};
10
11use crate::builtins;
12use crate::cli_error;
13use crate::config::{self, ArgSpec, CommandConfig, OptionSpec, ResolvedConfig, Schema};
14use crate::libtorch;
15use crate::style;
16
17// ── Config to CLI args ──────────────────────────────────────────────────
18
19/// Translate a resolved config into CLI arguments for the entry point.
20pub fn config_to_args(resolved: &ResolvedConfig) -> Vec<String> {
21    let mut args = Vec::new();
22
23    // DDP section
24    let d = &resolved.ddp;
25    push_opt(&mut args, "--mode", &d.mode);
26    push_opt(&mut args, "--policy", &d.policy);
27    push_opt(&mut args, "--backend", &d.backend);
28    push_value(&mut args, "--anchor", &d.anchor);
29    push_num(&mut args, "--max-anchor", &d.max_anchor);
30    push_float(&mut args, "--overhead-target", &d.overhead_target);
31    push_float(&mut args, "--divergence-threshold", &d.divergence_threshold);
32    push_value(&mut args, "--max-batch-diff", &d.max_batch_diff);
33    push_float(&mut args, "--max-grad-norm", &d.max_grad_norm);
34    push_num(&mut args, "--snapshot-timeout", &d.snapshot_timeout);
35    push_num(&mut args, "--checkpoint-every", &d.checkpoint_every);
36    push_value(&mut args, "--progressive", &d.progressive);
37    if let Some(hint) = &d.speed_hint {
38        args.push("--speed-hint".into());
39        args.push(format!("{}:{}", hint.slow_rank, hint.ratio));
40    }
41    if let Some(ratios) = &d.partition_ratios {
42        let s: Vec<String> = ratios.iter().map(|r| format!("{r}")).collect();
43        args.push("--partition-ratios".into());
44        args.push(s.join(","));
45    }
46    if let Some(ratio) = d.lr_scale_ratio {
47        args.push("--lr-scale-ratio".into());
48        args.push(format!("{ratio}"));
49    }
50    if d.timeline == Some(true) {
51        args.push("--timeline".into());
52    }
53
54    // Training section
55    let t = &resolved.training;
56    push_num(&mut args, "--epochs", &t.epochs);
57    push_num(&mut args, "--batch-size", &t.batch_size);
58    push_num(&mut args, "--batches", &t.batches_per_epoch);
59    push_float(&mut args, "--lr", &t.lr);
60    push_num(&mut args, "--seed", &t.seed);
61
62    // Output section
63    let o = &resolved.output;
64    push_opt(&mut args, "--output", &o.dir);
65    push_num(&mut args, "--monitor", &o.monitor);
66
67    // Pass-through options
68    for (key, val) in &resolved.options {
69        let flag = format!("--{}", key.replace('_', "-"));
70        match val {
71            serde_json::Value::Bool(true) => args.push(flag),
72            serde_json::Value::Bool(false) => {}
73            serde_json::Value::Null => {}
74            other => {
75                args.push(flag);
76                args.push(value_to_string(other));
77            }
78        }
79    }
80
81    args
82}
83
84fn push_opt(args: &mut Vec<String>, flag: &str, val: &Option<String>) {
85    if let Some(v) = val {
86        args.push(flag.into());
87        args.push(v.clone());
88    }
89}
90
91fn push_num<T: std::fmt::Display>(args: &mut Vec<String>, flag: &str, val: &Option<T>) {
92    if let Some(v) = val {
93        args.push(flag.into());
94        args.push(v.to_string());
95    }
96}
97
98fn push_float(args: &mut Vec<String>, flag: &str, val: &Option<f64>) {
99    if let Some(v) = val {
100        args.push(flag.into());
101        args.push(format!("{v}"));
102    }
103}
104
105fn push_value(args: &mut Vec<String>, flag: &str, val: &Option<serde_json::Value>) {
106    if let Some(v) = val {
107        match v {
108            serde_json::Value::Null => {}
109            other => {
110                args.push(flag.into());
111                args.push(value_to_string(other));
112            }
113        }
114    }
115}
116
117fn value_to_string(v: &serde_json::Value) -> String {
118    match v {
119        serde_json::Value::String(s) => s.clone(),
120        serde_json::Value::Number(n) => n.to_string(),
121        serde_json::Value::Bool(b) => b.to_string(),
122        other => other.to_string(),
123    }
124}
125
126// ── Docker detection ────────────────────────────────────────────────────
127
128/// Check if we're already running inside a Docker container.
129fn inside_docker() -> bool {
130    Path::new("/.dockerenv").exists()
131}
132
133/// Default container path we assume the host project root is mounted
134/// at when `docker-compose.yml` is missing or can't be parsed.
135/// Matches the convention `fdl init` writes into every generated
136/// compose service (`.:/workspace`).
137const DEFAULT_CONTAINER_PROJECT_ROOT: &str = "/workspace";
138
139/// Per-process cache of the container-side project-root path, keyed by
140/// docker-compose service name. Populated lazily on first lookup and
141/// reused for the life of the `fdl` invocation. `docker-compose.yml` is
142/// user-edited and version-controlled, so re-parsing once per
143/// invocation is cheap — the cache only avoids re-parsing *within* a
144/// single invocation.
145static COMPOSE_MOUNT_CACHE: std::sync::OnceLock<std::collections::HashMap<String, String>> =
146    std::sync::OnceLock::new();
147
148/// Resolve the absolute container path where the host project root is
149/// mounted inside `service`.
150///
151/// Reads `docker-compose.yml` at `project_root` once per process and
152/// caches the `service → container_path` mapping. Falls back to
153/// [`DEFAULT_CONTAINER_PROJECT_ROOT`] when the compose file is missing,
154/// unparseable, or doesn't declare a matching bind mount for `.`.
155///
156/// This is what lets `exec_command` generate `cd <container-path>`
157/// prefixes that work regardless of the service's `working_dir:` —
158/// e.g. flodl-hf's `hf-parity` service uses
159/// `working_dir: /workspace/flodl-hf` so Python parity scripts can
160/// `import` sibling helpers, and a naive relative `cd flodl-hf/convert`
161/// would resolve to the non-existent
162/// `/workspace/flodl-hf/flodl-hf/convert`.
163fn container_project_root(project_root: &Path, service: &str) -> String {
164    let cache = COMPOSE_MOUNT_CACHE.get_or_init(|| parse_compose_project_mounts(project_root));
165    cache
166        .get(service)
167        .cloned()
168        .unwrap_or_else(|| DEFAULT_CONTAINER_PROJECT_ROOT.to_string())
169}
170
171/// Parse `<project_root>/docker-compose.yml` and return a map of
172/// `service → container-mount-path` for every service that bind-mounts
173/// the project root (host `.` or `./`).
174///
175/// Handles both short-form (`".:/workspace"`) and long-form
176/// (`{ type: bind, source: ., target: /workspace }`) volume entries.
177/// Read errors, parse errors, and missing sections all yield an empty
178/// map — callers fall back to the convention.
179fn parse_compose_project_mounts(project_root: &Path) -> std::collections::HashMap<String, String> {
180    let compose_path = project_root.join("docker-compose.yml");
181    let text = match std::fs::read_to_string(&compose_path) {
182        Ok(t) => t,
183        Err(_) => return std::collections::HashMap::new(),
184    };
185    let doc: serde_yaml_ng::Value = match serde_yaml_ng::from_str(&text) {
186        Ok(d) => d,
187        Err(_) => return std::collections::HashMap::new(),
188    };
189    let mut out = std::collections::HashMap::new();
190    let services = match doc.get("services").and_then(|v| v.as_mapping()) {
191        Some(s) => s,
192        None => return out,
193    };
194    for (name, svc) in services {
195        let svc_name = match name.as_str() {
196            Some(s) => s,
197            None => continue,
198        };
199        let volumes = match svc.get("volumes").and_then(|v| v.as_sequence()) {
200            Some(v) => v,
201            None => continue,
202        };
203        if let Some(container_path) = find_project_mount(volumes) {
204            // Strip a trailing `/` so later `format!("{root}/{workdir}")`
205            // never produces `//` in the middle of a path.
206            let cleaned = container_path.trim_end_matches('/').to_string();
207            let cleaned = if cleaned.is_empty() {
208                "/".to_string()
209            } else {
210                cleaned
211            };
212            out.insert(svc_name.to_string(), cleaned);
213        }
214    }
215    out
216}
217
218/// Inside a service's `volumes:` sequence, find the entry that
219/// bind-mounts the project root (host path `.` or `./`) and return the
220/// container-side target path.
221fn find_project_mount(volumes: &[serde_yaml_ng::Value]) -> Option<String> {
222    for entry in volumes {
223        if let Some(s) = entry.as_str() {
224            // Short form: "host:container[:options]". Docker-compose's
225            // short-form parser only splits on the first two `:` on
226            // POSIX, but fdl-generated hosts are always `.` so naive
227            // split-and-check works fine here.
228            let mut parts = s.splitn(3, ':');
229            let host = parts.next()?;
230            let container = parts.next()?;
231            if host == "." || host == "./" {
232                return Some(container.to_string());
233            }
234        } else if let Some(m) = entry.as_mapping() {
235            // Long form: { type: bind, source: ., target: /workspace }.
236            let source = m
237                .get(serde_yaml_ng::Value::String("source".into()))
238                .and_then(|v| v.as_str());
239            let target = m
240                .get(serde_yaml_ng::Value::String("target".into()))
241                .and_then(|v| v.as_str());
242            if matches!(source, Some(".") | Some("./"))
243                && let Some(t) = target
244            {
245                return Some(t.to_string());
246            }
247        }
248    }
249    None
250}
251
252/// Resolve libtorch env vars from the project root, matching the Makefile logic:
253///   LIBTORCH_HOST_PATH = ./libtorch/<active_variant>          (standalone)
254///                      = <host.path>/libtorch/<host.arch>     (overlay)
255///   LIBTORCH_CPU_PATH  = ./libtorch/precompiled/cpu
256///   CUDA_VERSION, CUDA_TAG from .arch metadata
257///   FDL_GPU_FEATURE   = the cargo feature the active variant needs
258fn libtorch_env(project_root: &Path) -> Result<Vec<(String, String)>, String> {
259    let mut env = Vec::new();
260
261    // CPU path is always the same.
262    env.push((
263        "LIBTORCH_CPU_PATH".into(),
264        "./libtorch/precompiled/cpu".into(),
265    ));
266
267    if let Some((info, host_path)) = resolve_libtorch(project_root)? {
268        env.push(("LIBTORCH_HOST_PATH".into(), host_path.clone()));
269
270        // Native commands: fill what the docker services get from their
271        // compose env, so the scaffold's printed next steps
272        // (`./fdl libtorch download --cpu` then `./fdl build` then
273        // `./fdl run`) are true without hand exports. `LIBTORCH_PATH`
274        // is what flodl-sys's build.rs consumes (fill-when-absent: the
275        // inherited value is build.rs's documented manual override and
276        // keeps the last word); `LD_LIBRARY_PATH` is what lets the
277        // linked binary load at runtime, vendor-ordered — on ROCm the
278        // system runtime must precede libtorch's bundled copy.
279        let abs_path = if Path::new(&host_path).is_relative() {
280            project_root.join(&host_path).display().to_string()
281        } else {
282            host_path
283        };
284        if std::env::var_os("LIBTORCH_PATH").is_none() {
285            env.push(("LIBTORCH_PATH".into(), abs_path.clone()));
286        }
287        let lib = format!("{abs_path}/lib");
288        let inherited = std::env::var("LD_LIBRARY_PATH").unwrap_or_default();
289        if !inherited.split(':').any(|p| p == lib) {
290            let value = crate::libtorch::detect::ld_library_path_value(
291                crate::libtorch::detect::variant_vendor(&info.path),
292                &lib,
293                &crate::libtorch::detect::local_rocm_lib_dir(),
294            );
295            let composed = if inherited.is_empty() {
296                value
297            } else {
298                format!("{value}:{inherited}")
299            };
300            env.push(("LD_LIBRARY_PATH".into(), composed));
301        }
302
303        // macOS: bake the runtime search path into the BINARY, because
304        // neither of the two obvious fixes works here.
305        //
306        // Its loader ignores `LD_LIBRARY_PATH` entirely, so everything
307        // above is inert there and a scaffolded project builds and then
308        // dies with `dyld: Library not loaded: @rpath/libtorch.dylib`.
309        // The apparent fix, exporting `DYLD_LIBRARY_PATH` beside it, does
310        // not survive: run lines execute through `sh -c`, /bin/sh is SIP
311        // restricted, and dyld purges every `DYLD_*` from a restricted
312        // process before the shell ever execs the binary. The other
313        // apparent fix, an rpath from `flodl-sys`'s build script, does not
314        // reach here either -- `cargo:rustc-link-arg` does not propagate
315        // to dependents (the same property that forces the libtorch_cuda
316        // dlopen).
317        //
318        // RUSTFLAGS is neither: SIP has no opinion on it, and it applies
319        // to the final link of whatever cargo is building. The binary
320        // then locates libtorch on its own, which also means it still
321        // runs when launched directly rather than through `fdl run`.
322        //
323        // Fill-when-absent, like LIBTORCH_PATH above: a caller's own
324        // RUSTFLAGS is theirs, and appending to it would silently change
325        // a build they had configured. Linux is deliberately untouched --
326        // LD_LIBRARY_PATH already works there, and setting RUSTFLAGS
327        // would invalidate every existing cargo cache for no gain.
328        if cfg!(target_os = "macos") && std::env::var_os("RUSTFLAGS").is_none() {
329            env.push(("RUSTFLAGS".into(), format!("-C link-arg=-Wl,-rpath,{lib}")));
330        }
331
332        // The cargo feature this variant needs, so a `run:` line can say
333        // `--features "$FDL_GPU_FEATURE"` instead of hardcoding a vendor.
334        // Run lines execute under `bash -c` / `sh -c`, so the expansion
335        // is the shell's; this just has to be in the child's env.
336        //
337        // Defaults to `cuda` when no variant resolves, which reproduces
338        // exactly what the hardcoded `--features cuda` did before: a GPU
339        // command with no GPU libtorch fails the same way it always has,
340        // rather than failing differently in a way nobody recognises.
341        // `source::build_env` (join/publish recipes) deliberately answers
342        // "" for the same case instead — there a CPU variant is the
343        // publish gate's advertised cheap mode, not a misconfiguration.
344        env.push((
345            "FDL_GPU_FEATURE".into(),
346            match crate::libtorch::detect::variant_vendor(&info.path) {
347                Some(v) => v.cargo_feature().to_string(),
348                None => "cuda".to_string(),
349            },
350        ));
351
352        // CUDA version from .arch metadata.
353        if let Some(cuda) = &info.cuda_version
354            && cuda != "none"
355        {
356            let cuda_version = if cuda.matches('.').count() < 2 {
357                format!("{cuda}.0")
358            } else {
359                cuda.clone()
360            };
361            let cuda_tag = cuda_version
362                .splitn(3, '.')
363                .take(2)
364                .collect::<Vec<_>>()
365                .join(".");
366            env.push(("CUDA_VERSION".into(), cuda_version));
367            env.push(("CUDA_TAG".into(), cuda_tag));
368        }
369    }
370
371    Ok(env)
372}
373
374/// Resolve `(LibtorchInfo, host_path)` for `libtorch_env`.
375///
376/// Priority:
377///   1. Cluster overlay's per-host `arch:` (when `FDL_ENV` is
378///      set, the merged config has a `cluster:` block, AND the current
379///      hostname matches an entry). Resolved via the convention
380///      `<host.path>/libtorch/<arch>`, so each host in a shared-checkout
381///      heterogeneous rig picks its own libtorch without flipping the
382///      global `.active`.
383///   2. `project_root/libtorch/.active` (or `.active.<case>` via the
384///      `FDL_LIBTORCH_CASE` env var). Standalone single-host default.
385///
386/// Returns `Ok(None)` only when neither path resolves — the caller (env
387/// builder) then omits `LIBTORCH_HOST_PATH`, which surfaces as a
388/// libtorch-missing error from the downstream cargo/Docker invocation.
389/// `Err` when an active `FDL_ENV` overlay fails to load (see
390/// [`resolve_libtorch_from_overlay`]).
391fn resolve_libtorch(
392    project_root: &Path,
393) -> Result<Option<(libtorch::detect::LibtorchInfo, String)>, String> {
394    if let Some(resolved) = resolve_libtorch_from_overlay(project_root)? {
395        return Ok(Some(resolved));
396    }
397    let Some(info) = libtorch::detect::read_active(project_root) else {
398        return Ok(None);
399    };
400    let host_path = format!("./libtorch/{}", info.path);
401    Ok(Some((info, host_path)))
402}
403
404/// Try to resolve libtorch from the active cluster overlay's current-
405/// host entry. `Ok(None)` when the overlay legitimately doesn't apply
406/// (no `FDL_ENV`, no `cluster:` block, the current host isn't listed,
407/// or the entry's `arch:` is unset). `Err` when `FDL_ENV` is set but the
408/// config cannot be loaded — the user asked for that env, so silently
409/// falling back to `.active` could select the wrong libtorch.
410///
411/// Convention: libtorch lives at `<host.path>/libtorch/<host.arch>`
412/// on every host. The controller's view uses `<host.path>` directly
413/// here because this function is the controller-side (local) path
414/// resolver — when fdl runs locally as the current host, that host's
415/// own `path:` IS the controller's view.
416fn resolve_libtorch_from_overlay(
417    project_root: &Path,
418) -> Result<Option<(libtorch::detect::LibtorchInfo, String)>, String> {
419    let Ok(env_name) = std::env::var("FDL_ENV") else {
420        return Ok(None);
421    };
422    let env_name = env_name.trim();
423    if env_name.is_empty() {
424        return Ok(None);
425    }
426    // Same discovery set as `find_config` (fdl.yaml / fdl.yml / fdl.json) —
427    // a hardcoded fdl.yml here silently skipped fdl.yaml projects.
428    let base_path = config::find_config_in(project_root).ok_or_else(|| {
429        format!(
430            "FDL_ENV={env_name} is set but no fdl config file exists in {}",
431            project_root.display()
432        )
433    })?;
434    let cfg = config::load_project_with_env(&base_path, Some(env_name))
435        .map_err(|e| format!("FDL_ENV={env_name}: cannot resolve the overlay: {e}"))?;
436    let Some(cluster) = cfg.cluster else {
437        return Ok(None);
438    };
439    let host_name = crate::cluster::resolve_local_hostname();
440    let Some(entry) = cluster.workers.iter().find(|w| w.host == host_name) else {
441        return Ok(None);
442    };
443    let Some(arch) = entry.arch.as_ref() else {
444        return Ok(None);
445    };
446    let variant_dir = std::path::PathBuf::from(&entry.path)
447        .join("libtorch")
448        .join(arch);
449    Ok(resolve_libtorch_at(&variant_dir))
450}
451
452/// Resolve a libtorch variant dir (the per-host `arch:` applied as
453/// `<path>/libtorch/<arch>`) into
454/// `(LibtorchInfo, absolute host path for Docker bind mount)`. Accepts
455/// the same three shapes as `probe::check_libtorch_at`:
456///   1. Pointer file `.active*` — read pointer, resolve variant against
457///      the file's parent dir.
458///   2. Directory containing `.active` — read its `.active`.
459///   3. Direct variant dir (has `lib/`) — use as-is, parse `.arch` if
460///      present.
461pub(crate) fn resolve_libtorch_at(path: &Path) -> Option<(libtorch::detect::LibtorchInfo, String)> {
462    if path.is_file()
463        && path
464            .file_name()
465            .and_then(|n| n.to_str())
466            .is_some_and(|n| n.starts_with(".active"))
467    {
468        let libtorch_root = path.parent()?;
469        let info = libtorch::detect::read_active_from(path, libtorch_root)?;
470        let host_path = libtorch_root.join(&info.path).display().to_string();
471        return Some((info, host_path));
472    }
473    if path.join(".active").exists() {
474        let info = libtorch::detect::read_active_from(&path.join(".active"), path)?;
475        let host_path = path.join(&info.path).display().to_string();
476        return Some((info, host_path));
477    }
478    if path.join("lib").is_dir() {
479        let info = libtorch::detect::libtorch_info_from_dir(path.display().to_string(), path);
480        let host_path = path.display().to_string();
481        return Some((info, host_path));
482    }
483    None
484}
485
486/// Spawn a shell command with libtorch env vars set.
487///
488/// `FLODL_VERBOSITY` is forwarded to Docker containers via the
489/// `environment:` section in docker-compose.yml (bare variable name
490/// passes the host value through when set, ignored otherwise).
491fn spawn_docker_shell(command: &str, project_root: &Path) -> ExitCode {
492    let env_vars = match libtorch_env(project_root) {
493        Ok(v) => v,
494        Err(e) => {
495            eprintln!("fdl: {e}");
496            return ExitCode::FAILURE;
497        }
498    };
499
500    let mut cmd = std::process::Command::new("sh");
501    cmd.args(["-c", command])
502        .current_dir(project_root)
503        // Export HOSTNAME so docker-compose's `hostname: ${HOSTNAME}`
504        // interpolation resolves to the host's hostname. bash sets
505        // HOSTNAME as a shell built-in but doesn't export it; docker
506        // compose only reads exported env vars.
507        .env("HOSTNAME", crate::cluster::resolve_local_hostname())
508        .stdout(Stdio::inherit())
509        .stderr(Stdio::inherit())
510        .stdin(Stdio::inherit());
511
512    for (key, val) in &env_vars {
513        cmd.env(key, val);
514    }
515
516    match cmd.status() {
517        Ok(s) if s.success() => ExitCode::SUCCESS,
518        Ok(s) => ExitCode::from(s.code().unwrap_or(1) as u8),
519        Err(e) => {
520            cli_error!("{e}");
521            ExitCode::FAILURE
522        }
523    }
524}
525
526// ── Run-kind execution ──────────────────────────────────────────────────
527
528pub(crate) use crate::util::shell::posix_quote;
529
530/// Split `s` on the first whitespace-bounded `--` token, returning the
531/// halves with that token removed. Trim each half. When no such token is
532/// found, the whole string is returned as the "before" half and the
533/// "after" half is empty.
534///
535/// Whitespace-bounded means the `--` must be a standalone token: a
536/// leading `--`, a trailing `--`, a sole `--`, or a ` -- ` between
537/// other tokens. A bare `--foo` (no separator) is not a match. Quoted
538/// content in `s` passes through verbatim — split scanning happens on
539/// the raw string, not its shell-tokenised form.
540fn split_append_dashdash(s: &str) -> (String, String) {
541    let s = s.trim();
542    if s == "--" {
543        return (String::new(), String::new());
544    }
545    if let Some(rest) = s.strip_prefix("-- ") {
546        return (String::new(), rest.trim().to_string());
547    }
548    if let Some(prefix) = s.strip_suffix(" --") {
549        return (prefix.trim().to_string(), String::new());
550    }
551    if let Some(idx) = s.find(" -- ") {
552        let before = &s[..idx];
553        let after = &s[idx + 4..];
554        return (before.trim().to_string(), after.trim().to_string());
555    }
556    (s.to_string(), String::new())
557}
558
559/// Split `args` on the first standalone `--` token, returning the
560/// halves with that token removed. Returns `(before, Some(after))` when
561/// a separator is present, `(args, None)` otherwise. The `None` case
562/// keeps callers from emitting a stray `--` when the user did not pass
563/// runner-side args.
564fn split_user_args_dashdash(args: &[String]) -> (&[String], Option<&[String]>) {
565    match args.iter().position(|a| a == "--") {
566        Some(idx) => (&args[..idx], Some(&args[idx + 1..])),
567        None => (args, None),
568    }
569}
570
571/// Compose the final shell command from `run` + `append` + `user_args`.
572///
573/// Layout: `run [append-pre] [user-pre] -- [append-post] [user-post]`,
574/// where `append-pre` / `append-post` are halves of the yml `append:`
575/// field split on its first standalone `--`, and `user-pre` /
576/// `user-post` are halves of the CLI tokens that followed fdl's own
577/// first `--`, split again on a second `--`. The `--` separator only
578/// emits when at least one of the post halves is non-empty (or the
579/// `append` half explicitly carries one — preserving the legacy
580/// `append: -- --nocapture` shape).
581///
582/// User args go after append on each side: `append` seeds defaults,
583/// and last-wins for cargo-style flags lets the user override without
584/// ceremony (e.g. `append: --no-ansi` + CLI `--ansi`).
585pub(crate) fn compose_run_command(run: &str, user_args: &[String], append: Option<&str>) -> String {
586    let suffix = append.map(str::trim).filter(|s| !s.is_empty());
587    let (append_pre, append_post, append_has_dashdash) = match suffix {
588        Some(s) => {
589            let (pre, post) = split_append_dashdash(s);
590            // Distinguish "append had a `--`" (legacy `-- --nocapture`
591            // with empty pre, non-empty post) from "append had no `--`"
592            // (post defaults empty). Without this we'd swallow the
593            // separator on legacy yml.
594            let has = s == "--" || s.starts_with("-- ") || s.ends_with(" --") || s.contains(" -- ");
595            (pre, post, has)
596        }
597        None => (String::new(), String::new(), false),
598    };
599    let (user_pre, user_post_opt) = split_user_args_dashdash(user_args);
600
601    let mut out = String::from(run.trim());
602    if !append_pre.is_empty() {
603        out.push(' ');
604        out.push_str(&append_pre);
605    }
606    for a in user_pre {
607        out.push(' ');
608        out.push_str(&posix_quote(a));
609    }
610
611    let needs_separator = append_has_dashdash
612        || !append_post.is_empty()
613        || user_post_opt.is_some_and(|p| !p.is_empty());
614    if needs_separator {
615        out.push_str(" --");
616        if !append_post.is_empty() {
617            out.push(' ');
618            out.push_str(&append_post);
619        }
620        if let Some(post) = user_post_opt {
621            for a in post {
622                out.push(' ');
623                out.push_str(&posix_quote(a));
624            }
625        }
626    }
627    out
628}
629
630/// Run an inline `run:` script, optionally wrapped in Docker.
631///
632/// `user_args` (from CLI tokens after `--`) are POSIX-quoted and spliced
633/// between `command` and `append`, so a script like `cargo test live`
634/// with `append: -- --nocapture --ignored` still receives its libtest
635/// flags after a user-supplied `-p flodl-hf`.
636/// Forward the testing-cluster envelope into a docker-compose run
637/// invocation. When `fdl @cluster-test-{nccl,cpu} <cmd>` activates an
638/// overlay with a `cluster:` block, the dispatcher sets
639/// `FLODL_TESTING_CLUSTER_JSON` in fdl-cli's own env (see
640/// `dispatch_config` in main.rs). This helper checks that variable and
641/// returns a bare ` -e NAME` fragment (docker passes the value through
642/// from the environment the `sh -c` child inherits from this process)
643/// so the inner cargo process can see it; without it, the env var dies
644/// at the docker boundary and `discover_test_cluster()` inside the
645/// container silently falls back to local autodetect.
646///
647/// The value never appears on the command line, so this stays correct
648/// even if the envelope encoding changes (today it is hex, which would
649/// be shell-safe inline; the bare form does not depend on that).
650/// Source of truth for the env-var name lives in
651/// `flodl::distributed::testing::ENV_TESTING_CLUSTER_JSON`; mirrored
652/// here as a literal because flodl-cli is decoupled from the flodl
653/// library crate by policy (it must build without libtorch).
654fn testing_cluster_env_arg() -> String {
655    let mut out = String::new();
656    for name in TESTING_ENV_VARS {
657        if std::env::var(name).is_ok() {
658            out.push_str(" -e ");
659            out.push_str(name);
660        }
661    }
662    out
663}
664
665/// Testing env vars forwarded into a docker-compose run when present.
666///
667/// `FLODL_TESTING_CLUSTER_JSON` injects a cluster topology;
668/// `FLODL_TESTING_GPU_JSON` injects a spoofed GPU survey (source of
669/// truth: `flodl_hw::ENV_TESTING_GPU_JSON`). Both die at the docker
670/// boundary without an explicit `-e`, and both fail *silently* when
671/// they do -- the container falls back to real detection and the test
672/// quietly measures the host's actual hardware instead of the described
673/// one. Names are literals here because flodl-cli is decoupled from the
674/// flodl library crate by policy; `flodl-hw` is a dependency, so the GPU
675/// one is asserted against its constant in the tests below.
676const TESTING_ENV_VARS: &[&str] = &["FLODL_TESTING_CLUSTER_JSON", "FLODL_TESTING_GPU_JSON"];
677
678/// The logical `docker:` value meaning "whichever GPU container matches
679/// the active libtorch variant".
680pub const LOGICAL_GPU_SERVICE: &str = "gpu";
681
682/// Resolve a `docker:` value to a concrete docker-compose service.
683///
684/// Only `gpu` is logical; every other value passes through untouched, so
685/// `docker: cuda` and `docker: rocm` remain explicit escapes for anyone
686/// who wants to pin a container regardless of the active variant.
687///
688/// Why the container cannot simply BE vendor-neutral the way the cargo
689/// feature is: both vendors build from one source tree, so a feature is
690/// derivable, but a CUDA image (nvidia/cuda base, nvidia runtime,
691/// `/dev/nvidia*`) and a ROCm image (rocm/dev-ubuntu, `/dev/kfd` +
692/// `/dev/dri`, render group) are genuinely different artifacts. A
693/// service that declared both device sets would fail to start on a host
694/// missing either. So the service must be *selected*, and this is where.
695///
696/// The AMD arm falls back to the CUDA container **with a message** when
697/// no `rocm` service is defined. This repo and `fdl init`'s scaffolds
698/// both define one now, so the fallback is for a project whose
699/// `docker-compose.yml` predates it (or was hand-trimmed). It stays a
700/// message rather than an error because landing in the CUDA container
701/// produces the *better* diagnostic: flodl-sys then reports exactly
702/// which part of its ROCm path is missing, where compose's bare "no such
703/// service: rocm" would say less.
704pub fn resolve_docker_service(name: &str, project_root: &Path) -> String {
705    if name != LOGICAL_GPU_SERVICE {
706        return name.to_string();
707    }
708    let vendor = resolve_libtorch(project_root)
709        .ok()
710        .flatten()
711        .and_then(|(info, _)| crate::libtorch::detect::variant_vendor(&info.path));
712
713    // CONVENTION: the compose service name IS the cargo feature name
714    // (`cuda`, `rocm`). One table instead of two, and `GpuVendor` is
715    // `#[non_exhaustive]`, so a vendor added upstream resolves to a
716    // sensibly-named service here without touching this function.
717    let Some(vendor) = vendor else {
718        // No GPU variant resolved: nothing to select on. `cuda` is the
719        // historical default and its own failure is the informative one.
720        return "cuda".to_string();
721    };
722    let service = vendor.cargo_feature();
723    if service == "cuda" || compose_has_service(service, project_root) {
724        return service.to_string();
725    }
726    eprintln!(
727        "fdl: the active libtorch variant targets {vendor}, but no `{service}` \
728         docker-compose service is defined; falling back to the `cuda` container. \
729         Define a `{service}` service, or set `docker:` explicitly on the command \
730         to silence this."
731    );
732    "cuda".to_string()
733}
734
735/// Whether the merged compose config defines `service`.
736///
737/// One subprocess, and only on the AMD path -- an NVIDIA host never pays
738/// for it.
739fn compose_has_service(service: &str, project_root: &Path) -> bool {
740    std::process::Command::new("docker")
741        .args(["compose", "config", "--services"])
742        .current_dir(project_root)
743        .output()
744        .is_ok_and(|o| {
745            o.status.success()
746                && String::from_utf8_lossy(&o.stdout)
747                    .lines()
748                    .any(|l| l.trim() == service)
749        })
750}
751
752pub fn exec_script(
753    command: &str,
754    append: Option<&str>,
755    user_args: &[String],
756    docker_service: Option<&str>,
757    cwd: &Path,
758) -> ExitCode {
759    let inner_cmd = compose_run_command(command, user_args, append);
760
761    match docker_service {
762        Some(service) if !inside_docker() => {
763            // Quote the whole composed command for the outer
764            // `bash -c` so user args containing shell metacharacters
765            // don't escape the inner shell.
766            let overlay = crate::cluster::cluster_compose_overlay_arg(cwd);
767            let testing_env_arg = testing_cluster_env_arg();
768            let service = resolve_docker_service(service, cwd);
769            let docker_cmd = format!(
770                "docker compose{overlay} run --rm{testing_env_arg} {service} bash -c {}",
771                posix_quote(&inner_cmd)
772            );
773            spawn_docker_shell(&docker_cmd, cwd)
774        }
775        _ => {
776            let (shell, flag) = if cfg!(target_os = "windows") {
777                ("cmd", "/C")
778            } else {
779                ("sh", "-c")
780            };
781
782            // The same libtorch env the docker path gets (via the
783            // compose `environment:` passthrough): a native `run:` line
784            // saying `--features "$FDL_GPU_FEATURE"` was silently
785            // getting an empty variable, which broke the scaffolded
786            // gpu-* commands in `fdl init --native` projects — for both
787            // vendors.
788            let env_vars = match libtorch_env(cwd) {
789                Ok(v) => v,
790                Err(e) => {
791                    eprintln!("fdl: {e}");
792                    return ExitCode::FAILURE;
793                }
794            };
795            let mut cmd = std::process::Command::new(shell);
796            cmd.args([flag, inner_cmd.as_str()]);
797            for (key, val) in &env_vars {
798                cmd.env(key, val);
799            }
800            match cmd
801                .current_dir(cwd)
802                .stdout(Stdio::inherit())
803                .stderr(Stdio::inherit())
804                .stdin(Stdio::inherit())
805                .status()
806            {
807                Ok(s) if s.success() => ExitCode::SUCCESS,
808                Ok(s) => ExitCode::from(s.code().unwrap_or(1) as u8),
809                Err(e) => {
810                    cli_error!("{e}");
811                    ExitCode::FAILURE
812                }
813            }
814        }
815    }
816}
817
818// ── Command execution ───────────────────────────────────────────────────
819
820/// Execute a sub-command, optionally with a named preset (inline command).
821///
822/// `project_root` is needed to resolve Docker compose context and
823/// compute the relative workdir for containerized execution.
824pub fn exec_command(
825    cmd_config: &CommandConfig,
826    preset_name: Option<&str>,
827    extra_args: &[String],
828    cmd_dir: &Path,
829    project_root: &Path,
830) -> ExitCode {
831    let entry = match &cmd_config.entry {
832        Some(e) => e.as_str(),
833        None => {
834            eprintln!(
835                "error: no entry point defined in {}/fdl.yaml",
836                cmd_dir.display()
837            );
838            return ExitCode::FAILURE;
839        }
840    };
841
842    // Tail validation pre-flight. Runs whenever a schema is present:
843    // - `choices:` on declared options → always enforced.
844    // - Unknown flags → rejected only when `schema.strict` is set
845    //   (lenient mode tolerates pass-through flags the binary may
846    //   consume directly).
847    // fdl-generated args (from the structured ddp/training/output
848    // blocks) are intentionally skipped — those are the binary's
849    // surface, not the user's.
850    if let Some(schema) = &cmd_config.schema
851        && let Err(e) = config::validate_tail(extra_args, schema)
852    {
853        cli_error!("{e}");
854        return ExitCode::FAILURE;
855    }
856
857    // Resolve config: preset overrides merged with root defaults.
858    let resolved = match preset_name {
859        Some(name) => match cmd_config.commands.get(name) {
860            Some(preset) => {
861                // Validate *this* preset only (choices + strict unknowns).
862                // Whole-map validation is deferred so a broken sibling
863                // preset doesn't block a correct one from running.
864                if let Some(schema) = &cmd_config.schema
865                    && let Err(e) = config::validate_preset_for_exec(name, preset, schema)
866                {
867                    cli_error!("{e}");
868                    return ExitCode::FAILURE;
869                }
870                config::merge_preset(cmd_config, preset)
871            }
872            None => {
873                cli_error!("unknown command '{name}'");
874                eprintln!();
875                print_command_help(cmd_config, "");
876                return ExitCode::FAILURE;
877            }
878        },
879        None => config::defaults_only(cmd_config),
880    };
881
882    // Build argument list from config.
883    let mut args = config_to_args(&resolved);
884
885    // Append extra CLI args (these override config-derived args).
886    args.extend(extra_args.iter().cloned());
887
888    // Docker wrapping or direct execution.
889    let use_docker = cmd_config.docker.is_some() && !inside_docker();
890
891    if use_docker {
892        let service = resolve_docker_service(cmd_config.docker.as_deref().unwrap(), project_root);
893        let workdir = cmd_dir
894            .strip_prefix(project_root)
895            .unwrap_or(cmd_dir)
896            .to_string_lossy();
897
898        // Build the inner command: cd <container-root>/<workdir> && <entry> <args>
899        //
900        // `<container-root>` is discovered from docker-compose.yml's
901        // bind mount for the project (host `.` → container target) via
902        // [`container_project_root`], so the generated `cd` works
903        // regardless of the service's own `working_dir:`. Falls back
904        // to `/workspace` (the fdl init convention) when the compose
905        // file is missing or silent on this service.
906        let container_root = container_project_root(project_root, &service);
907        let args_str = shell_join(&args);
908        let inner = if workdir.is_empty() || workdir == "." {
909            format!("{entry} {args_str}")
910        } else {
911            format!(
912                "cd {} && {entry} {args_str}",
913                posix_quote(&format!("{container_root}/{workdir}"))
914            )
915        };
916
917        if preset_name.is_some() {
918            eprintln!("fdl: [{service}] {inner}");
919        }
920
921        // Surface the container-side workspace root to the inner
922        // process so entry binaries can re-anchor argv path arguments
923        // independently of the per-task `cd <root>/<workdir>` we
924        // injected above. Mirrors the `HF_HOME` pattern: `fdl` owns
925        // the env, the binary just reads it. Without this, a user
926        // typing `flodl-hf/tests/.exports/bert` from the host repo
927        // root resolves against the wrong cwd inside the container.
928        let overlay = crate::cluster::cluster_compose_overlay_arg(project_root);
929        let testing_env_arg = testing_cluster_env_arg();
930        // Quote the whole composed command for the outer `sh -c`,
931        // exactly like `exec_script`. A double-quoted wrapper would
932        // let the outer shell expand `$`/backticks inside it (host-side,
933        // defeating shell_join's quoting) and break on any `"` in an
934        // argument.
935        // FDL_PROJECT_ROOT is quoted too: this whole string goes through
936        // `sh -c`, and a container root with a space would otherwise
937        // splice the env value across arguments.
938        let docker_cmd = format!(
939            "docker compose{overlay} run --rm -e {}{testing_env_arg} {service} bash -c {}",
940            posix_quote(&format!("FDL_PROJECT_ROOT={container_root}")),
941            posix_quote(&inner),
942        );
943        spawn_docker_shell(&docker_cmd, project_root)
944    } else {
945        // Direct execution (inside container or no docker configured).
946        let parts: Vec<&str> = entry.split_whitespace().collect();
947        if parts.is_empty() {
948            cli_error!("empty entry point");
949            return ExitCode::FAILURE;
950        }
951        let program = parts[0];
952        let entry_args = &parts[1..];
953
954        if preset_name.is_some() {
955            let preview: Vec<&str> = args.iter().map(|s| s.as_str()).collect();
956            eprintln!("fdl: {entry} {}", preview.join(" "));
957        }
958
959        // Same libtorch env as the docker path — see exec_script's
960        // native arm for why (a native entry saying
961        // `--features "$FDL_GPU_FEATURE"` got nothing before).
962        let env_vars = match libtorch_env(project_root) {
963            Ok(v) => v,
964            Err(e) => {
965                eprintln!("fdl: {e}");
966                return ExitCode::FAILURE;
967            }
968        };
969        let mut cmd = std::process::Command::new(program);
970        cmd.args(entry_args).args(&args);
971        for (key, val) in &env_vars {
972            cmd.env(key, val);
973        }
974        match cmd
975            .current_dir(cmd_dir)
976            .stdout(Stdio::inherit())
977            .stderr(Stdio::inherit())
978            .stdin(Stdio::inherit())
979            .status()
980        {
981            Ok(s) if s.success() => ExitCode::SUCCESS,
982            Ok(s) => ExitCode::from(s.code().unwrap_or(1) as u8),
983            Err(e) => {
984                cli_error!("failed to execute '{program}': {e}");
985                ExitCode::FAILURE
986            }
987        }
988    }
989}
990
991/// Join args into a single shell-safe string for the docker `bash -c "…"`
992/// path. Each token is POSIX-quoted via [`posix_quote`], so shell
993/// metacharacters in a value (`$`, backticks, globs, `;`, `|`, redirections,
994/// …) are passed literally rather than expanded or interpreted by the inner
995/// shell. The previous predicate only quoted tokens containing a space, `"`,
996/// or empty — leaving `$HOME`, `` `cmd` ``, `*.py`, `a;b` to be interpreted.
997/// Mirrors `exec_script`, which quotes through the same helper.
998fn shell_join(args: &[String]) -> String {
999    args.iter()
1000        .map(|a| posix_quote(a))
1001        .collect::<Vec<_>>()
1002        .join(" ")
1003}
1004
1005// ── Help output ─────────────────────────────────────────────────────────
1006
1007/// Print help for a `run:`-kind command. Shows the inline script,
1008/// any `append:` suffix, the Docker service (if any), and the `--`
1009/// forwarding contract.
1010pub fn print_run_help(
1011    name: &str,
1012    description: Option<&str>,
1013    run: &str,
1014    append: Option<&str>,
1015    docker: Option<&str>,
1016) {
1017    if let Some(desc) = description {
1018        eprintln!("{} {desc}", style::bold(name));
1019    } else {
1020        eprintln!("{}", style::bold(name));
1021    }
1022    eprintln!();
1023    eprintln!("{}:", style::yellow("Usage"));
1024    eprintln!("    fdl {name} [-- <args>... [-- <runner-args>...]]");
1025    eprintln!();
1026    eprintln!("{}:", style::yellow("Runs"));
1027    let composed = match append.map(str::trim).filter(|s| !s.is_empty()) {
1028        Some(suffix) => {
1029            let (pre, post) = split_append_dashdash(suffix);
1030            // Show whichever halves the yml actually populated, with
1031            // the user slots interleaved at the right side per the
1032            // merge contract.
1033            let left = if pre.is_empty() {
1034                "[<args>]".to_string()
1035            } else {
1036                format!("{pre} [<args>]")
1037            };
1038            let right = if post.is_empty() {
1039                "[<runner-args>]".to_string()
1040            } else {
1041                format!("{post} [<runner-args>]")
1042            };
1043            format!("{run} {left} -- {right}")
1044        }
1045        None => format!("{run} [<args>] [-- <runner-args>]"),
1046    };
1047    if let Some(svc) = docker {
1048        eprintln!(
1049            "    {} {svc} -c {composed:?}",
1050            style::dim("docker compose run --rm")
1051        );
1052    } else {
1053        eprintln!("    {composed}");
1054    }
1055    eprintln!();
1056    eprintln!(
1057        "{} the first `--` separates fdl args from the run script; a second `--` splits cargo-side args from runner-side args.",
1058        style::dim("Note:"),
1059    );
1060    eprintln!(
1061        "{} `append:` is split on its own `--` and merged half-and-half; pass `--no-append` to drop it entirely.",
1062        style::dim("Note:"),
1063    );
1064}
1065
1066/// Print help for a sub-command (its arguments, nested commands, and
1067/// entry). Orchestrates the per-section helpers below.
1068pub fn print_command_help(cmd_config: &CommandConfig, name: &str) {
1069    let (presets, sub_cmds) = split_commands_by_kind(&cmd_config.commands);
1070    let preset_slot = cmd_config.arg_name.as_deref().unwrap_or("preset");
1071
1072    // Wrap descriptions to the terminal width (or COLUMNS / a sane default).
1073    let width = help_width();
1074
1075    print_title(cmd_config, name);
1076    print_usage_line(cmd_config, name, &presets, &sub_cmds, preset_slot);
1077    print_arguments_section(cmd_config, &presets, preset_slot, width);
1078    print_sub_commands_section(&sub_cmds);
1079    print_schema_commands_section(cmd_config, name);
1080    print_options_section(cmd_config, width);
1081    print_entry_section(cmd_config);
1082    print_defaults_section(cmd_config);
1083}
1084
1085fn print_title(cmd_config: &CommandConfig, name: &str) {
1086    if let Some(desc) = &cmd_config.description {
1087        eprintln!("{} {desc}", style::bold(name));
1088    } else {
1089        eprintln!("{}", style::bold(name));
1090    }
1091}
1092
1093fn print_usage_line(
1094    cmd_config: &CommandConfig,
1095    name: &str,
1096    presets: &CommandGroup,
1097    sub_cmds: &CommandGroup,
1098    preset_slot: &str,
1099) {
1100    // The first-positional slot reflects what is actually accepted here:
1101    // preset name, sub-command name, or either.
1102    let usage_tail = build_usage_tail(
1103        cmd_config.schema.as_ref(),
1104        !presets.is_empty(),
1105        !sub_cmds.is_empty(),
1106        preset_slot,
1107    );
1108    eprintln!();
1109    eprintln!("{}:", style::yellow("Usage"));
1110    eprintln!("    fdl {name}{usage_tail}");
1111}
1112
1113fn print_arguments_section(
1114    cmd_config: &CommandConfig,
1115    presets: &CommandGroup,
1116    preset_slot: &str,
1117    width: usize,
1118) {
1119    // Schema-declared positionals (typed slots on the entry binary) and
1120    // the preset slot (dispatched by fdl before the binary sees argv)
1121    // both land in the first-positional position, so they share one
1122    // section. Schema args render first; the preset slot with its
1123    // value list follows.
1124    let has_schema_args = cmd_config
1125        .schema
1126        .as_ref()
1127        .is_some_and(|s| !s.args.is_empty());
1128    if !has_schema_args && presets.is_empty() {
1129        return;
1130    }
1131    eprintln!();
1132    eprintln!("{}:", style::yellow("Arguments"));
1133    let avail = width.saturating_sub(4);
1134    if let Some(schema) = &cmd_config.schema {
1135        for a in &schema.args {
1136            for line in format_arg(a, avail) {
1137                eprintln!("    {line}");
1138            }
1139        }
1140    }
1141    if !presets.is_empty() {
1142        let slot_label = format!("[<{preset_slot}>]");
1143        eprintln!(
1144            "    {}  Named preset, one of:",
1145            style::green(&format!("{:<20}", slot_label))
1146        );
1147        for (pname, spec) in presets {
1148            let desc = spec.description.as_deref().unwrap_or("-");
1149            eprintln!(
1150                "      {}  {}",
1151                style::green(&format!("{:<18}", pname)),
1152                desc
1153            );
1154        }
1155    }
1156}
1157
1158fn print_sub_commands_section(sub_cmds: &CommandGroup) {
1159    // Run/Path kinds only — true sub-commands with their own behavior
1160    // (an inline script or a nested fdl.yml).
1161    if sub_cmds.is_empty() {
1162        return;
1163    }
1164    eprintln!();
1165    eprintln!("{}:", style::yellow("Commands"));
1166    for (sub_name, sub_spec) in sub_cmds {
1167        let desc = sub_spec.description.as_deref().unwrap_or("-");
1168        eprintln!(
1169            "    {}  {}",
1170            style::green(&format!("{:<20}", sub_name)),
1171            desc
1172        );
1173    }
1174}
1175
1176/// List the entry binary's own subcommands when its schema is a tree
1177/// (a variant-shaped `#[derive(FdlArgs)]` CLI). Distinct from
1178/// [`print_sub_commands_section`], which lists fdl.yml-level Run/Path
1179/// commands — these come from the binary's `--fdl-schema` output, and each
1180/// has its own flag set (drill in with `fdl <name> <subcommand> --help`).
1181fn print_schema_commands_section(cmd_config: &CommandConfig, name: &str) {
1182    let Some(schema) = &cmd_config.schema else {
1183        return;
1184    };
1185    if schema.commands.is_empty() {
1186        return;
1187    }
1188    eprintln!();
1189    eprintln!("{}:", style::yellow("Commands"));
1190    for (sub_name, sub_schema) in &schema.commands {
1191        let desc = sub_schema.description.as_deref().unwrap_or("-");
1192        eprintln!(
1193            "    {}  {}",
1194            style::green(&format!("{:<20}", sub_name)),
1195            desc
1196        );
1197    }
1198    eprintln!();
1199    eprintln!(
1200        "    Run {} for a subcommand's options.",
1201        style::dim(&format!("fdl {name} <command> --help"))
1202    );
1203}
1204
1205fn print_options_section(cmd_config: &CommandConfig, width: usize) {
1206    // Schema-driven options. Renders only when a schema block is present
1207    // in fdl.yaml; the "Defaults" section covers ddp/training/output.
1208    let Some(schema) = &cmd_config.schema else {
1209        return;
1210    };
1211    if schema.options.is_empty() {
1212        return;
1213    }
1214    eprintln!();
1215    eprintln!("{}:", style::yellow("Options"));
1216    let avail = width.saturating_sub(4);
1217    for (long, spec) in &schema.options {
1218        for line in format_option(long, spec, avail) {
1219            eprintln!("    {line}");
1220        }
1221    }
1222}
1223
1224fn print_entry_section(cmd_config: &CommandConfig) {
1225    let Some(entry) = &cmd_config.entry else {
1226        return;
1227    };
1228    eprintln!();
1229    eprintln!("{}:", style::yellow("Entry"));
1230    eprintln!("    {entry}");
1231    if let Some(service) = &cmd_config.docker {
1232        eprintln!("     {}", style::dim(&format!("[docker: {service}]")));
1233    }
1234    eprintln!();
1235    eprintln!(
1236        "    Any extra {} are forwarded to the entry point.",
1237        style::dim("[options]")
1238    );
1239}
1240
1241fn print_defaults_section(cmd_config: &CommandConfig) {
1242    if cmd_config.ddp.is_none() && cmd_config.training.is_none() {
1243        return;
1244    }
1245    eprintln!();
1246    eprintln!("{}:", style::yellow("Defaults"));
1247    if let Some(d) = &cmd_config.ddp {
1248        if let Some(mode) = &d.mode {
1249            eprintln!("    {}  {mode}", style::dim("ddp.mode"));
1250        }
1251        if let Some(anchor) = &d.anchor {
1252            eprintln!(
1253                "    {}  {}",
1254                style::dim("ddp.anchor"),
1255                value_to_string(anchor)
1256            );
1257        }
1258    }
1259    if let Some(t) = &cmd_config.training {
1260        if let Some(e) = t.epochs {
1261            eprintln!("    {}  {e}", style::dim("training.epochs"));
1262        }
1263        if let Some(bs) = t.batch_size {
1264            eprintln!("    {}  {bs}", style::dim("training.batch_size"));
1265        }
1266        if let Some(lr) = t.lr {
1267            eprintln!("    {}  {lr}", style::dim("training.lr"));
1268        }
1269        if let Some(seed) = t.seed {
1270            eprintln!("    {}  {seed}", style::dim("training.seed"));
1271        }
1272    }
1273}
1274
1275/// Print help for a named preset command nested inside a sub-command.
1276pub fn print_preset_help(cmd_config: &CommandConfig, cmd_name: &str, preset_name: &str) {
1277    let preset = match cmd_config.commands.get(preset_name) {
1278        Some(s) => s,
1279        None => {
1280            eprintln!("unknown command: {preset_name}");
1281            return;
1282        }
1283    };
1284
1285    // Title.
1286    let desc = preset.description.as_deref().unwrap_or("(no description)");
1287    eprintln!(
1288        "{} {} {}",
1289        style::bold(cmd_name),
1290        style::green(preset_name),
1291        desc
1292    );
1293
1294    eprintln!();
1295    eprintln!("{}:", style::yellow("Usage"));
1296    eprintln!(
1297        "    fdl {cmd_name} {preset_name} {}",
1298        style::dim("[extra options]")
1299    );
1300
1301    // Show the merged config that this preset produces.
1302    let resolved = config::merge_preset(cmd_config, preset);
1303
1304    eprintln!();
1305    eprintln!("{}:", style::yellow("Effective config"));
1306
1307    // DDP fields.
1308    let d = &resolved.ddp;
1309    print_config_field("ddp.mode", &d.mode);
1310    print_config_value("ddp.anchor", &d.anchor);
1311    print_config_field("ddp.max_anchor", &d.max_anchor);
1312    print_config_field("ddp.overhead_target", &d.overhead_target);
1313    print_config_field("ddp.divergence_threshold", &d.divergence_threshold);
1314    print_config_value("ddp.max_batch_diff", &d.max_batch_diff);
1315    print_config_field("ddp.max_grad_norm", &d.max_grad_norm);
1316    if d.timeline == Some(true) {
1317        eprintln!("    {}  true", style::dim("ddp.timeline"));
1318    }
1319
1320    // Training fields.
1321    let t = &resolved.training;
1322    print_config_field("training.epochs", &t.epochs);
1323    print_config_field("training.batch_size", &t.batch_size);
1324    print_config_field("training.batches_per_epoch", &t.batches_per_epoch);
1325    print_config_field("training.lr", &t.lr);
1326    print_config_field("training.seed", &t.seed);
1327
1328    // Output fields.
1329    let o = &resolved.output;
1330    print_config_field("output.dir", &o.dir);
1331    print_config_field("output.monitor", &o.monitor);
1332
1333    // Pass-through options.
1334    if !resolved.options.is_empty() {
1335        eprintln!();
1336        eprintln!("{}:", style::yellow("Options"));
1337        for (key, val) in &resolved.options {
1338            eprintln!(
1339                "    {}  {}",
1340                style::green(&format!("--{}", key.replace('_', "-"))),
1341                value_to_string(val)
1342            );
1343        }
1344    }
1345
1346    // Show the effective command.
1347    if let Some(entry) = &cmd_config.entry {
1348        let args = config_to_args(&resolved);
1349        let args_str = args.join(" ");
1350        let docker_info = cmd_config
1351            .docker
1352            .as_ref()
1353            .map(|s| format!("[{s}] ",))
1354            .unwrap_or_default();
1355
1356        eprintln!();
1357        eprintln!("{}:", style::yellow("Effective command"));
1358        eprintln!(
1359            "    {}{}{}",
1360            style::dim(&docker_info),
1361            entry,
1362            if args_str.is_empty() {
1363                String::new()
1364            } else {
1365                format!(" {args_str}")
1366            }
1367        );
1368    }
1369
1370    eprintln!();
1371    eprintln!(
1372        "Extra {} after the command name are appended to the entry.",
1373        style::dim("[options]")
1374    );
1375}
1376
1377fn print_config_field<T: std::fmt::Display>(label: &str, val: &Option<T>) {
1378    if let Some(v) = val {
1379        eprintln!("    {}  {v}", style::dim(label));
1380    }
1381}
1382
1383fn print_config_value(label: &str, val: &Option<serde_json::Value>) {
1384    if let Some(v) = val
1385        && !v.is_null()
1386    {
1387        eprintln!("    {}  {}", style::dim(label), value_to_string(v));
1388    }
1389}
1390
1391/// Print the project help with scripts and commands.
1392pub fn print_project_help(
1393    project: &config::ProjectConfig,
1394    project_root: &Path,
1395    active_env: Option<&str>,
1396) {
1397    let visible_builtins = builtins::visible_top_level();
1398    if let Some(desc) = &project.description {
1399        eprintln!("{} {}", style::bold("fdl"), desc);
1400    } else {
1401        eprintln!("{} {}", style::bold("fdl"), env!("CARGO_PKG_VERSION"));
1402    }
1403
1404    eprintln!();
1405    eprintln!("{}:", style::yellow("Usage"));
1406    eprintln!(
1407        "    fdl {} {}",
1408        style::dim("<command>"),
1409        style::dim("[options]")
1410    );
1411
1412    eprintln!();
1413    eprintln!("{}:", style::yellow("Options"));
1414    eprintln!(
1415        "    {}  Show this help",
1416        style::green(&format!("{:<18}", "-h, --help"))
1417    );
1418    eprintln!(
1419        "    {}  Show version",
1420        style::green(&format!("{:<18}", "-V, --version"))
1421    );
1422    eprintln!(
1423        "    {}  Use fdl.<name>.yml overlay (also: --env <name>, FDL_ENV=<name>)",
1424        style::green(&format!("{:<18}", "@<name>"))
1425    );
1426    eprintln!(
1427        "    {}  Scope visible GPUs, e.g. 0,1 or all (any position)",
1428        style::green(&format!("{:<18}", "--gpus <spec>"))
1429    );
1430    eprintln!(
1431        "    {}  Verbose output",
1432        style::green(&format!("{:<18}", "-v"))
1433    );
1434    eprintln!(
1435        "    {}  Debug output",
1436        style::green(&format!("{:<18}", "-vv"))
1437    );
1438    eprintln!(
1439        "    {}  Trace output (maximum detail)",
1440        style::green(&format!("{:<18}", "-vvv"))
1441    );
1442    eprintln!(
1443        "    {}  Suppress non-error output",
1444        style::green(&format!("{:<18}", "-q, --quiet"))
1445    );
1446    eprintln!(
1447        "    {}  Force ANSI color (bypass TTY / NO_COLOR detection)",
1448        style::green(&format!("{:<18}", "--ansi"))
1449    );
1450    eprintln!(
1451        "    {}  Disable ANSI color output",
1452        style::green(&format!("{:<18}", "--no-ansi"))
1453    );
1454    eprintln!(
1455        "    {}  Drop a run command's `append:` suffix",
1456        style::green(&format!("{:<18}", "--no-append"))
1457    );
1458    eprintln!(
1459        "    {}  Skip the cluster pre-flight build",
1460        style::green(&format!("{:<18}", "--no-prebuild"))
1461    );
1462
1463    // Built-in commands.
1464    eprintln!();
1465    eprintln!("{}:", style::yellow("Built-in"));
1466    for (name, desc) in &visible_builtins {
1467        eprintln!("    {}  {desc}", style::green(&format!("{:<18}", name)));
1468    }
1469
1470    // Commands: unified section. Each entry in `project.commands` is one
1471    // of: an inline `run:` script, a `path:` (or convention-default)
1472    // pointer to a child fdl.yml, or — at nested levels only — an inline
1473    // preset. Descriptions come from the `CommandSpec`; for `path:`
1474    // commands missing their own description, fall back to loading the
1475    // child fdl.yml's description.
1476    if !project.commands.is_empty() {
1477        eprintln!();
1478        eprintln!("{}:", style::yellow("Commands"));
1479        for (name, spec) in &project.commands {
1480            let desc: String = match spec.description.clone() {
1481                Some(d) => d,
1482                None => {
1483                    // For path-kind entries, fall back to the child config's
1484                    // own description so `commands: { ddp-bench: }` still
1485                    // shows a useful blurb.
1486                    let is_path_kind = spec.run.is_none();
1487                    if is_path_kind {
1488                        let child_dir = spec.resolve_path(name, project_root);
1489                        config::load_command_with_env(&child_dir, active_env)
1490                            .ok()
1491                            .and_then(|c| c.description)
1492                            .unwrap_or_else(|| "(sub-command)".into())
1493                    } else {
1494                        spec.run.as_deref().unwrap_or("(command)").to_string()
1495                    }
1496                }
1497            };
1498            eprintln!("    {}  {desc}", style::green(&format!("{:<18}", name)));
1499        }
1500    }
1501
1502    // Available environments (sibling fdl.<env>.yml files at project root).
1503    if let Some(base_config) = config::find_config(project_root) {
1504        let envs = crate::overlay::list_envs(&base_config);
1505        if !envs.is_empty() {
1506            eprintln!();
1507            eprintln!("{}:", style::yellow("Environments"));
1508            for e in &envs {
1509                let active_marker = if Some(e.as_str()) == active_env {
1510                    style::green(" (active)")
1511                } else {
1512                    String::new()
1513                };
1514                eprintln!(
1515                    "    {}  Overlay from fdl.{}.yml{active_marker}",
1516                    style::green(&format!("{:<18}", format!("@{e}"))),
1517                    e
1518                );
1519            }
1520            eprintln!();
1521            eprintln!(
1522                "Use {} to run a command with an environment overlay.",
1523                style::dim("fdl @<env> <command>")
1524            );
1525        }
1526    }
1527
1528    eprintln!();
1529    eprintln!(
1530        "Use {} for more information on a command.",
1531        style::dim("fdl <command> -h")
1532    );
1533}
1534
1535// ── Schema-driven help helpers ──────────────────────────────────────────
1536
1537/// Build the part of `fdl <cmd>...` after the command name: positionals
1538/// rendered as `<name>` (required) or `[<name>]` (optional), plus a slot
1539/// for the first-positional picker — `[<preset>]` when only presets exist,
1540/// `[<command>]` when only sub-commands exist, `[<preset>|<command>]` when
1541/// both — and `[options]`. The preset placeholder is customisable per
1542/// sub-command via `arg-name:`.
1543fn build_usage_tail(
1544    schema: Option<&Schema>,
1545    has_presets: bool,
1546    has_sub_commands: bool,
1547    preset_slot: &str,
1548) -> String {
1549    let mut parts = String::new();
1550    let slot = match (has_presets, has_sub_commands) {
1551        (true, false) => Some(format!("[<{preset_slot}>]")),
1552        (false, true) => Some("[<command>]".to_string()),
1553        (true, true) => Some(format!("[<{preset_slot}>|<command>]")),
1554        (false, false) => None,
1555    };
1556    if let Some(s) = slot {
1557        parts.push(' ');
1558        parts.push_str(&style::dim(&s));
1559    }
1560    if let Some(s) = schema {
1561        for a in &s.args {
1562            parts.push(' ');
1563            parts.push_str(&format_arg_usage(a));
1564        }
1565    }
1566    parts.push(' ');
1567    parts.push_str(&style::dim("[options]"));
1568    parts
1569}
1570
1571type CommandGroup = Vec<(String, crate::config::CommandSpec)>;
1572
1573/// Partition a `commands:` map into (presets, sub-commands) by resolved
1574/// `CommandKind`. Entries whose `kind()` errors (both run and path set)
1575/// are treated as sub-commands so they still render somewhere — the
1576/// error surfaces when the user tries to dispatch them.
1577fn split_commands_by_kind(
1578    commands: &BTreeMap<String, crate::config::CommandSpec>,
1579) -> (CommandGroup, CommandGroup) {
1580    use crate::config::CommandKind;
1581    let mut presets = Vec::new();
1582    let mut sub_cmds = Vec::new();
1583    for (k, v) in commands {
1584        match v.kind() {
1585            Ok(CommandKind::Preset) => presets.push((k.clone(), v.clone())),
1586            _ => sub_cmds.push((k.clone(), v.clone())),
1587        }
1588    }
1589    (presets, sub_cmds)
1590}
1591
1592fn format_arg_usage(a: &ArgSpec) -> String {
1593    let suffix = if a.variadic { "..." } else { "" };
1594    let core = format!("<{}>{suffix}", a.name);
1595    if a.required && a.default.is_none() {
1596        style::green(&core)
1597    } else {
1598        style::dim(&format!("[{core}]"))
1599    }
1600}
1601
1602/// Label-column width for the `Arguments` section (chars, after the
1603/// section's own left indent). Descriptions wrap into the column to the
1604/// right of this.
1605const ARG_COL: usize = 22;
1606/// Label-column width for the `Options` section (option flags run wider
1607/// than positional args, so they get a roomier column).
1608const OPT_COL: usize = 30;
1609
1610fn format_arg(a: &ArgSpec, avail_width: usize) -> Vec<String> {
1611    let left = format_arg_usage(a);
1612    let visible = visible_width(&left);
1613    let segs = desc_segments(a.description.as_deref(), &a.default, &a.choices, &a.ty);
1614    format_row(&left, visible, ARG_COL, &segs, avail_width)
1615}
1616
1617/// Format an option row into one or more display lines: the flag (with its
1618/// value placeholder) in the label column, the description word-wrapped
1619/// into an aligned column to its right. A flag wider than the column drops
1620/// its description to the next line rather than crowding it.
1621fn format_option(long: &str, spec: &OptionSpec, avail_width: usize) -> Vec<String> {
1622    let flag = match &spec.short {
1623        Some(s) => format!("-{s}, --{long}"),
1624        None => format!("    --{long}"),
1625    };
1626    let placeholder = option_placeholder(&spec.ty);
1627    let (left, visible) = if placeholder.is_empty() {
1628        (style::green(&flag), flag.chars().count())
1629    } else {
1630        (
1631            style::green(&format!("{flag} {placeholder}")),
1632            flag.chars().count() + 1 + placeholder.chars().count(),
1633        )
1634    };
1635    let segs = desc_segments(
1636        spec.description.as_deref(),
1637        &spec.default,
1638        &spec.choices,
1639        &spec.ty,
1640    );
1641    let mut out = format_row(&left, visible, OPT_COL, &segs, avail_width);
1642    if let Some(env) = &spec.env {
1643        out.push(format!(
1644            "{}{}",
1645            " ".repeat(OPT_COL),
1646            style::dim(&format!("[env: {env}]"))
1647        ));
1648    }
1649    out
1650}
1651
1652/// A word/segment for description wrapping: `text` is the visible content
1653/// (what counts toward the wrap width), `styled` is what actually gets
1654/// printed (may carry ANSI escapes, which have zero visible width).
1655struct Seg {
1656    text: String,
1657    styled: String,
1658}
1659
1660impl Seg {
1661    fn plain(s: &str) -> Seg {
1662        Seg {
1663            text: s.to_string(),
1664            styled: s.to_string(),
1665        }
1666    }
1667    fn dim(s: &str) -> Seg {
1668        Seg {
1669            text: s.to_string(),
1670            styled: style::dim(s),
1671        }
1672    }
1673}
1674
1675/// Break a description (plus its `[default:]` / `[possible:]` / list-type
1676/// annotations) into wrap units. Free-text words split on whitespace so
1677/// they reflow; each annotation stays whole (a `[possible: a, b, c]` list
1678/// reads better unbroken than reflowed mid-item).
1679fn desc_segments(
1680    description: Option<&str>,
1681    default: &Option<serde_json::Value>,
1682    choices: &Option<Vec<serde_json::Value>>,
1683    ty: &str,
1684) -> Vec<Seg> {
1685    let mut segs: Vec<Seg> = description
1686        .unwrap_or("-")
1687        .split_whitespace()
1688        .map(Seg::plain)
1689        .collect();
1690    if let Some(d) = default {
1691        // Skip noisy defaults: bool false, empty list, null.
1692        let is_empty_list = matches!(d, serde_json::Value::Array(a) if a.is_empty());
1693        let is_false = matches!(d, serde_json::Value::Bool(false));
1694        if !d.is_null() && !is_false && !is_empty_list {
1695            segs.push(Seg::dim(&format!("[default: {}]", format_value(d))));
1696        }
1697    }
1698    if let Some(choices) = choices
1699        && !choices.is_empty()
1700    {
1701        let list = choices
1702            .iter()
1703            .map(format_value)
1704            .collect::<Vec<_>>()
1705            .join(", ");
1706        segs.push(Seg::dim(&format!("[possible: {list}]")));
1707    }
1708    // Annotate list types so users know about repeat/comma semantics.
1709    if ty.starts_with("list[") {
1710        segs.push(Seg::dim("(repeat or comma-separate)"));
1711    }
1712    segs
1713}
1714
1715/// Greedily pack segments into lines no wider than `width` visible chars,
1716/// one space between words. A single segment wider than `width` (e.g. a
1717/// long unbreakable annotation) gets its own overflowing line rather than
1718/// being split.
1719fn wrap_segments(segs: &[Seg], width: usize) -> Vec<String> {
1720    let width = width.max(1);
1721    let mut lines: Vec<String> = Vec::new();
1722    let mut cur = String::new();
1723    let mut cur_w = 0usize;
1724    for seg in segs {
1725        let w = seg.text.chars().count();
1726        if cur_w == 0 {
1727            cur.push_str(&seg.styled);
1728            cur_w = w;
1729        } else if cur_w + 1 + w <= width {
1730            cur.push(' ');
1731            cur.push_str(&seg.styled);
1732            cur_w += 1 + w;
1733        } else {
1734            lines.push(std::mem::take(&mut cur));
1735            cur.push_str(&seg.styled);
1736            cur_w = w;
1737        }
1738    }
1739    if !cur.is_empty() {
1740        lines.push(cur);
1741    }
1742    if lines.is_empty() {
1743        lines.push(String::new());
1744    }
1745    lines
1746}
1747
1748/// Lay out a two-column help row: a `label` of visible width
1749/// `label_visible` in a column `desc_col` chars wide, then `segs` wrapped
1750/// into the description column to its right. `avail_width` is the printable
1751/// width the section has after its own left indent. Continuation lines
1752/// align under the description column; a label too wide for its column
1753/// drops the description to the next line.
1754fn format_row(
1755    label: &str,
1756    label_visible: usize,
1757    desc_col: usize,
1758    segs: &[Seg],
1759    avail_width: usize,
1760) -> Vec<String> {
1761    // Floor so a narrow terminal still leaves a usable description column.
1762    const MIN_DESC: usize = 20;
1763    let desc_width = avail_width.saturating_sub(desc_col).max(MIN_DESC);
1764    let desc_lines = wrap_segments(segs, desc_width);
1765    let pad = " ".repeat(desc_col);
1766    let mut out: Vec<String> = Vec::with_capacity(desc_lines.len() + 1);
1767    if label_visible < desc_col {
1768        let gap = " ".repeat(desc_col - label_visible);
1769        out.push(format!("{label}{gap}{}", desc_lines[0]));
1770    } else {
1771        // Label overflows its column: give it its own line, description below.
1772        out.push(label.to_string());
1773        out.push(format!("{pad}{}", desc_lines[0]));
1774    }
1775    for line in &desc_lines[1..] {
1776        out.push(format!("{pad}{line}"));
1777    }
1778    // Drop trailing padding (e.g. a "-" placeholder leaves a padded blank).
1779    for line in &mut out {
1780        while line.ends_with(' ') {
1781            line.pop();
1782        }
1783    }
1784    out
1785}
1786
1787fn option_placeholder(ty: &str) -> &'static str {
1788    match ty {
1789        "bool" => "",
1790        "int" => "<N>",
1791        "float" => "<F>",
1792        "path" => "<PATH>",
1793        "list[path]" => "<PATH>...",
1794        t if t.starts_with("list[") => "<VALUE>...",
1795        _ => "<VALUE>",
1796    }
1797}
1798
1799fn format_value(v: &serde_json::Value) -> String {
1800    match v {
1801        serde_json::Value::String(s) => s.clone(),
1802        other => other.to_string(),
1803    }
1804}
1805
1806/// Rough visible width helper: styled strings wrap their visible content
1807/// in ANSI escapes, so we use the unstyled inputs we started from.
1808fn visible_width(s: &str) -> usize {
1809    // The inputs we pass here come from pre-styling helpers that already
1810    // know the raw length. Strip ANSI to be safe.
1811    strip_ansi(s).chars().count()
1812}
1813
1814/// Printable width to wrap help output to. Precedence: an explicit
1815/// `COLUMNS` env var (lets CI and pipelines pin it), else the controlling
1816/// terminal's width when stderr is a TTY, else a readable default. Clamped
1817/// so ultra-wide terminals don't stretch descriptions past comfortable
1818/// reading length and narrow ones stay usable.
1819fn help_width() -> usize {
1820    const DEFAULT: usize = 100;
1821    const MIN: usize = 60;
1822    const MAX: usize = 120;
1823    let raw = std::env::var("COLUMNS")
1824        .ok()
1825        .and_then(|s| s.trim().parse::<usize>().ok())
1826        .filter(|&c| c > 0)
1827        .or_else(term_cols)
1828        .unwrap_or(DEFAULT);
1829    raw.clamp(MIN, MAX)
1830}
1831
1832/// The controlling terminal's column count via `TIOCGWINSZ`, or `None`
1833/// when stderr is not a TTY (help is piped/redirected) or the query fails.
1834/// Dep-free: the minimal-deps policy precludes a terminal-size crate, so we
1835/// declare the one `ioctl` we need. Non-unix targets have no equivalent
1836/// here and fall back to the default width.
1837#[cfg(unix)]
1838fn term_cols() -> Option<usize> {
1839    use std::io::IsTerminal;
1840    use std::os::unix::io::AsRawFd;
1841
1842    let stderr = std::io::stderr();
1843    if !stderr.is_terminal() {
1844        return None;
1845    }
1846    #[repr(C)]
1847    struct Winsize {
1848        row: u16,
1849        col: u16,
1850        xpixel: u16,
1851        ypixel: u16,
1852    }
1853    // TIOCGWINSZ request code: 0x5413 on Linux/Android (incl. WSL), the
1854    // packed 0x4008_7468 on the BSDs/macOS.
1855    #[cfg(any(target_os = "linux", target_os = "android"))]
1856    const TIOCGWINSZ: std::os::raw::c_ulong = 0x5413;
1857    #[cfg(not(any(target_os = "linux", target_os = "android")))]
1858    const TIOCGWINSZ: std::os::raw::c_ulong = 0x4008_7468;
1859    unsafe extern "C" {
1860        fn ioctl(
1861            fd: std::os::raw::c_int,
1862            request: std::os::raw::c_ulong,
1863            ...
1864        ) -> std::os::raw::c_int;
1865    }
1866    let mut ws = Winsize {
1867        row: 0,
1868        col: 0,
1869        xpixel: 0,
1870        ypixel: 0,
1871    };
1872    // SAFETY: `ioctl(TIOCGWINSZ, &Winsize)` writes the window size into the
1873    // struct; the fd is stderr, verified above to be a terminal.
1874    let rc = unsafe { ioctl(stderr.as_raw_fd(), TIOCGWINSZ, &mut ws as *mut Winsize) };
1875    (rc == 0 && ws.col > 0).then_some(ws.col as usize)
1876}
1877
1878#[cfg(not(unix))]
1879fn term_cols() -> Option<usize> {
1880    None
1881}
1882
1883fn strip_ansi(s: &str) -> String {
1884    let mut out = String::with_capacity(s.len());
1885    let mut chars = s.chars().peekable();
1886    while let Some(c) = chars.next() {
1887        if c == '\x1b' && chars.peek() == Some(&'[') {
1888            chars.next();
1889            for c in chars.by_ref() {
1890                if c.is_ascii_alphabetic() {
1891                    break;
1892                }
1893            }
1894        } else {
1895            out.push(c);
1896        }
1897    }
1898    out
1899}
1900
1901#[cfg(test)]
1902mod tests {
1903    use super::*;
1904    use crate::util::test_env::env_lock;
1905
1906    // Help-column wrapping. Plain segments keep these independent of the
1907    // color state (no `style::*` calls), so they need no style lock.
1908    fn plain_segs(words: &[&str]) -> Vec<Seg> {
1909        words.iter().map(|w| Seg::plain(w)).collect()
1910    }
1911
1912    #[test]
1913    fn wrap_segments_packs_words_within_width() {
1914        let segs = plain_segs(&["alpha", "beta", "gamma", "delta"]);
1915        // width 12: "alpha beta" (10) fits, +" gamma" (16) does not.
1916        let lines = wrap_segments(&segs, 12);
1917        assert_eq!(lines, vec!["alpha beta", "gamma delta"]);
1918        for line in &lines {
1919            assert!(line.chars().count() <= 12);
1920        }
1921    }
1922
1923    #[test]
1924    fn wrap_segments_oversized_segment_gets_its_own_line() {
1925        let segs = plain_segs(&["short", "supercalifragilistic", "tail"]);
1926        let lines = wrap_segments(&segs, 10);
1927        // The oversized word overflows alone rather than being split.
1928        assert_eq!(lines, vec!["short", "supercalifragilistic", "tail"]);
1929    }
1930
1931    #[test]
1932    fn format_row_aligns_continuation_under_description_column() {
1933        // desc_col 8, avail 28 → desc width 20 (the MIN_DESC floor).
1934        let segs = plain_segs(&["aaaa", "bbbb", "cccc", "dddd", "eeee"]);
1935        let rows = format_row("--x", 3, 8, &segs, 28);
1936        // "aaaa bbbb cccc dddd" = 19 ≤ 20 fits; " eeee" would be 24, wraps.
1937        assert_eq!(rows[0], "--x     aaaa bbbb cccc dddd"); // 3 + 5 spaces = column 8
1938        assert_eq!(rows[1], format!("{}eeee", " ".repeat(8)));
1939        // Continuation lines are indented exactly to the column.
1940        for row in &rows[1..] {
1941            assert!(row.starts_with(&" ".repeat(8)));
1942            assert!(!row.starts_with(&" ".repeat(9)));
1943        }
1944    }
1945
1946    #[test]
1947    fn format_row_overflowing_label_drops_description_below() {
1948        let segs = plain_segs(&["desc"]);
1949        // Label wider than the 8-char column → own line, desc aligned below.
1950        let rows = format_row("--a-very-long-flag", 18, 8, &segs, 40);
1951        assert_eq!(rows[0], "--a-very-long-flag");
1952        assert_eq!(rows[1], format!("{}desc", " ".repeat(8)));
1953    }
1954
1955    #[test]
1956    fn help_width_honors_columns_env_within_clamp() {
1957        let _lock = env_lock();
1958        let prev = std::env::var("COLUMNS").ok();
1959        // SAFETY: guarded by the process-wide env lock.
1960        unsafe { std::env::set_var("COLUMNS", "90") };
1961        assert_eq!(help_width(), 90);
1962        unsafe { std::env::set_var("COLUMNS", "9999") };
1963        assert_eq!(help_width(), 120); // clamped to MAX
1964        unsafe { std::env::set_var("COLUMNS", "10") };
1965        assert_eq!(help_width(), 60); // clamped to MIN
1966        match prev {
1967            Some(v) => unsafe { std::env::set_var("COLUMNS", v) },
1968            None => unsafe { std::env::remove_var("COLUMNS") },
1969        }
1970    }
1971
1972    fn unique_tmp_dir(tag: &str) -> std::path::PathBuf {
1973        static SEQ: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1974        let d = std::env::temp_dir().join(format!(
1975            "fdl-run-test-{tag}-{}-{}",
1976            std::process::id(),
1977            SEQ.fetch_add(1, std::sync::atomic::Ordering::Relaxed),
1978        ));
1979        std::fs::create_dir_all(&d).unwrap();
1980        d
1981    }
1982
1983    #[test]
1984    fn overlay_libtorch_finds_fdl_yaml_spelling() {
1985        // The resolver previously hardcoded `fdl.yml` while config
1986        // discovery accepts fdl.yaml / fdl.yml / fdl.json — an fdl.yaml
1987        // project under FDL_ENV silently fell back to `.active`.
1988        let _guard = env_lock();
1989        let dir = unique_tmp_dir("yaml-spelling");
1990        std::fs::write(dir.join("fdl.yaml"), "description: base\n").unwrap();
1991        std::fs::write(
1992            dir.join("fdl.testenv.yaml"),
1993            "cluster:\n  controller:\n    host: 127.0.0.1\n    port: 29500\n    path: /opt/flodl\n  workers:\n    - host: not-this-host\n      local_devices: [0]\n      nccl_socket_ifname: lo\n      path: /opt/flodl\n",
1994        )
1995        .unwrap();
1996        unsafe { std::env::set_var("FDL_ENV", "testenv") };
1997        // Loads through fdl.yaml + overlay; current host isn't listed →
1998        // legitimate Ok(None), NOT an Err and NOT a hardcoded-name miss.
1999        let resolved = resolve_libtorch_from_overlay(&dir);
2000        unsafe { std::env::remove_var("FDL_ENV") };
2001        std::fs::remove_dir_all(&dir).ok();
2002        assert!(matches!(resolved, Ok(None)), "{resolved:?}");
2003    }
2004
2005    #[test]
2006    fn overlay_libtorch_load_failure_is_loud() {
2007        // A broken overlay under FDL_ENV must error, not silently fall
2008        // back to `.active` (wrong libtorch on heterogeneous rigs).
2009        let _guard = env_lock();
2010        let dir = unique_tmp_dir("broken-overlay");
2011        std::fs::write(dir.join("fdl.yml"), "description: base\n").unwrap();
2012        // FDL_ENV names an overlay that doesn't exist -> load error.
2013        unsafe { std::env::set_var("FDL_ENV", "missing-env") };
2014        let resolved = resolve_libtorch_from_overlay(&dir);
2015        unsafe { std::env::remove_var("FDL_ENV") };
2016        std::fs::remove_dir_all(&dir).ok();
2017        let err = match resolved {
2018            Ok(v) => panic!("expected Err, got Ok({v:?})"),
2019            Err(e) => e,
2020        };
2021        assert!(err.contains("missing-env"), "{err}");
2022    }
2023
2024    #[test]
2025    fn posix_quote_passes_safe_strings_through() {
2026        assert_eq!(posix_quote("hello"), "hello");
2027        assert_eq!(posix_quote("-p"), "-p");
2028        assert_eq!(posix_quote("flodl-hf"), "flodl-hf");
2029        assert_eq!(posix_quote("a/b.c"), "a/b.c");
2030        assert_eq!(posix_quote("KEY=val"), "KEY=val");
2031    }
2032
2033    #[test]
2034    fn posix_quote_wraps_unsafe_strings() {
2035        assert_eq!(posix_quote(""), "''");
2036        assert_eq!(posix_quote("foo bar"), "'foo bar'");
2037        assert_eq!(posix_quote("a$b"), "'a$b'");
2038        assert_eq!(posix_quote("a\"b"), "'a\"b'");
2039    }
2040
2041    #[test]
2042    fn posix_quote_escapes_embedded_single_quotes() {
2043        assert_eq!(posix_quote("it's"), "'it'\\''s'");
2044        assert_eq!(posix_quote("'"), "''\\'''");
2045    }
2046
2047    #[test]
2048    fn posix_quote_round_trips_shell_join_output() {
2049        // The docker exec paths nest quoting: shell_join single-quotes
2050        // each arg, then the whole `cd … && entry args` command is
2051        // posix_quote'd again for the outer `sh -c`. The embedded
2052        // single quotes must re-escape as '\'' so the inner bash sees
2053        // the args byte-for-byte. A double-quoted wrapper ("{inner}")
2054        // instead lets the OUTER shell expand $/backticks and breaks
2055        // on any `"` in an arg.
2056        let args: Vec<String> = ["--tag", "$HOME", "a\"b"]
2057            .iter()
2058            .map(|s| s.to_string())
2059            .collect();
2060        let inner = format!("cd /workspace/bench && train {}", shell_join(&args));
2061        assert_eq!(
2062            posix_quote(&inner),
2063            "'cd /workspace/bench && train --tag '\\''$HOME'\\'' '\\''a\"b'\\'''"
2064        );
2065    }
2066
2067    #[test]
2068    fn shell_join_quotes_shell_metacharacters() {
2069        // M23: safe tokens pass through bare and join with spaces; values with
2070        // shell metacharacters ($, glob, ;) are single-quoted so the inner
2071        // `bash -c "…"` treats them literally instead of expanding them.
2072        let args: Vec<String> = ["--model", "mlp", "--tag", "$HOME", "*.py", "a;b"]
2073            .iter()
2074            .map(|s| s.to_string())
2075            .collect();
2076        assert_eq!(shell_join(&args), "--model mlp --tag '$HOME' '*.py' 'a;b'");
2077    }
2078
2079    #[test]
2080    fn compose_run_command_no_extras_passes_run_through() {
2081        assert_eq!(compose_run_command("echo hello", &[], None), "echo hello");
2082    }
2083
2084    #[test]
2085    fn compose_run_command_inserts_user_args_between_run_and_append() {
2086        let user = vec!["-p".to_string(), "flodl-hf".to_string()];
2087        let out = compose_run_command("cargo test live", &user, Some("-- --nocapture --ignored"));
2088        assert_eq!(out, "cargo test live -p flodl-hf -- --nocapture --ignored");
2089    }
2090
2091    #[test]
2092    fn compose_run_command_quotes_user_args_with_spaces() {
2093        let user = vec!["--name".to_string(), "with space".to_string()];
2094        let out = compose_run_command("cmd", &user, None);
2095        assert_eq!(out, "cmd --name 'with space'");
2096    }
2097
2098    #[test]
2099    fn compose_run_command_omits_empty_append() {
2100        let out = compose_run_command("cmd", &["arg".to_string()], Some(""));
2101        assert_eq!(out, "cmd arg");
2102        let out2 = compose_run_command("cmd", &["arg".to_string()], Some("   "));
2103        assert_eq!(out2, "cmd arg");
2104    }
2105
2106    #[test]
2107    fn compose_run_command_user_double_dash_threads_runner_args() {
2108        let user = vec![
2109            "-p".to_string(),
2110            "foo".to_string(),
2111            "--".to_string(),
2112            "--ignored".to_string(),
2113        ];
2114        let out = compose_run_command("cargo test", &user, Some("-- --nocapture"));
2115        assert_eq!(out, "cargo test -p foo -- --nocapture --ignored");
2116    }
2117
2118    #[test]
2119    fn compose_run_command_user_double_dash_without_append() {
2120        let user = vec![
2121            "-p".to_string(),
2122            "foo".to_string(),
2123            "--".to_string(),
2124            "--ignored".to_string(),
2125        ];
2126        let out = compose_run_command("cargo test", &user, None);
2127        assert_eq!(out, "cargo test -p foo -- --ignored");
2128    }
2129
2130    #[test]
2131    fn compose_run_command_append_with_pre_and_post_halves() {
2132        let out = compose_run_command("cmd", &[], Some("--foo -- --bar"));
2133        assert_eq!(out, "cmd --foo -- --bar");
2134    }
2135
2136    #[test]
2137    fn compose_run_command_append_pre_only_no_separator() {
2138        // append carries a default flag with no `--` token; user supplies
2139        // an override. Defaults seed first, user wins via last-flag-wins.
2140        let user = vec!["--ansi".to_string()];
2141        let out = compose_run_command("cmd", &user, Some("--no-ansi"));
2142        assert_eq!(out, "cmd --no-ansi --ansi");
2143    }
2144
2145    #[test]
2146    fn compose_run_command_user_only_double_dash_emits_separator() {
2147        let user = vec!["--".to_string(), "--list".to_string()];
2148        let out = compose_run_command("cargo test", &user, None);
2149        assert_eq!(out, "cargo test -- --list");
2150    }
2151
2152    #[test]
2153    fn compose_run_command_append_full_split_with_user_both_sides() {
2154        let user = vec![
2155            "-p".to_string(),
2156            "foo".to_string(),
2157            "--".to_string(),
2158            "--ignored".to_string(),
2159        ];
2160        let out = compose_run_command("cargo test", &user, Some("--release -- --nocapture"));
2161        assert_eq!(out, "cargo test --release -p foo -- --nocapture --ignored");
2162    }
2163
2164    #[test]
2165    fn split_append_dashdash_handles_edges() {
2166        assert_eq!(
2167            split_append_dashdash("-- --nocapture"),
2168            (String::new(), "--nocapture".to_string())
2169        );
2170        assert_eq!(
2171            split_append_dashdash("--foo -- --bar"),
2172            ("--foo".to_string(), "--bar".to_string())
2173        );
2174        assert_eq!(
2175            split_append_dashdash("--foo --"),
2176            ("--foo".to_string(), String::new())
2177        );
2178        assert_eq!(split_append_dashdash("--"), (String::new(), String::new()));
2179        assert_eq!(
2180            split_append_dashdash("--foo"),
2181            ("--foo".to_string(), String::new())
2182        );
2183        assert_eq!(split_append_dashdash(""), (String::new(), String::new()));
2184    }
2185
2186    // ── resolve_libtorch_at: 3-shape libtorch variant resolution ────────
2187    //
2188    // Each test builds a synthetic libtorch dir under a per-test scratch
2189    // path (the minimal-deps policy precludes pulling in `tempfile`) and feeds the path
2190    // through `resolve_libtorch_at`. Variant names (`precompiled/v1`,
2191    // `builds/v2`) are deliberately abstract — the resolver is structural,
2192    // not rig-aware.
2193
2194    use std::sync::atomic::{AtomicU64, Ordering};
2195    use std::time::{SystemTime, UNIX_EPOCH};
2196
2197    static SCRATCH_SEQ: AtomicU64 = AtomicU64::new(0);
2198
2199    struct Scratch(std::path::PathBuf);
2200    impl Scratch {
2201        fn new() -> Self {
2202            let nanos = SystemTime::now()
2203                .duration_since(UNIX_EPOCH)
2204                .map(|d| d.as_nanos())
2205                .unwrap_or(0);
2206            let seq = SCRATCH_SEQ.fetch_add(1, Ordering::Relaxed);
2207            let dir = std::env::temp_dir().join(format!("fdl-resolve-libtorch-{}-{}", nanos, seq));
2208            std::fs::create_dir_all(&dir).unwrap();
2209            Self(dir)
2210        }
2211        fn path(&self) -> &std::path::Path {
2212            &self.0
2213        }
2214    }
2215    impl Drop for Scratch {
2216        fn drop(&mut self) {
2217            let _ = std::fs::remove_dir_all(&self.0);
2218        }
2219    }
2220
2221    fn populate_lt_root(root: &std::path::Path) {
2222        for (sub, torch, archs) in [
2223            ("precompiled/v1", "1.0", "0.0"),
2224            ("builds/v2", "2.0", "1.0"),
2225        ] {
2226            let d = root.join(sub);
2227            std::fs::create_dir_all(d.join("lib")).unwrap();
2228            std::fs::write(
2229                d.join(".arch"),
2230                format!("torch={torch}\ncuda=1.0\narchs={archs}\nsource=test\n"),
2231            )
2232            .unwrap();
2233        }
2234    }
2235
2236    #[test]
2237    fn resolve_libtorch_at_pointer_file() {
2238        let s = Scratch::new();
2239        let lt = s.path().join("libtorch");
2240        populate_lt_root(&lt);
2241        let pointer = lt.join(".active.alt");
2242        std::fs::write(&pointer, "precompiled/v1\n").unwrap();
2243
2244        let (info, host_path) = resolve_libtorch_at(&pointer).expect("pointer file resolves");
2245        assert_eq!(info.path, "precompiled/v1");
2246        assert_eq!(info.torch_version.as_deref(), Some("1.0"));
2247        assert_eq!(host_path, lt.join("precompiled/v1").display().to_string());
2248    }
2249
2250    #[test]
2251    fn resolve_libtorch_at_libtorch_root_dir() {
2252        let s = Scratch::new();
2253        let lt = s.path().join("libtorch");
2254        populate_lt_root(&lt);
2255        std::fs::write(lt.join(".active"), "builds/v2\n").unwrap();
2256
2257        let (info, host_path) = resolve_libtorch_at(&lt).expect("libtorch-root dir resolves");
2258        assert_eq!(info.path, "builds/v2");
2259        assert_eq!(info.torch_version.as_deref(), Some("2.0"));
2260        assert_eq!(host_path, lt.join("builds/v2").display().to_string());
2261    }
2262
2263    #[test]
2264    fn resolve_libtorch_at_direct_variant_dir() {
2265        let s = Scratch::new();
2266        let variant = s.path().join("standalone-libtorch");
2267        std::fs::create_dir_all(variant.join("lib")).unwrap();
2268        std::fs::write(
2269            variant.join(".arch"),
2270            "torch=3.0\ncuda=2.0\narchs=1.0\nsource=test\n",
2271        )
2272        .unwrap();
2273
2274        let (info, host_path) = resolve_libtorch_at(&variant).expect("direct variant dir resolves");
2275        assert_eq!(info.path, variant.display().to_string());
2276        assert_eq!(info.torch_version.as_deref(), Some("3.0"));
2277        assert_eq!(host_path, variant.display().to_string());
2278    }
2279
2280    #[test]
2281    fn resolve_libtorch_at_bogus_path_returns_none() {
2282        let s = Scratch::new();
2283        let bogus = s.path().join("no-lib-no-active-no-pointer");
2284        std::fs::create_dir_all(&bogus).unwrap();
2285        assert!(
2286            resolve_libtorch_at(&bogus).is_none(),
2287            "dir without lib/, .active, or pointer-shape filename → None"
2288        );
2289    }
2290
2291    #[test]
2292    fn resolve_docker_service_passes_explicit_names_through() {
2293        // Only `gpu` is logical. An explicit pin must survive untouched,
2294        // including one naming a service that does not exist yet --
2295        // resolving it would defeat the point of pinning.
2296        let root = Path::new("/nonexistent");
2297        for name in ["cuda", "rocm", "dev", "bench", "something-custom"] {
2298            assert_eq!(resolve_docker_service(name, root), name, "{name}");
2299        }
2300    }
2301
2302    #[test]
2303    fn resolve_docker_service_falls_back_when_no_variant_resolves() {
2304        // No libtorch under this root, so nothing to select on: `cuda`
2305        // is the historical default and its own failure is the
2306        // informative one.
2307        assert_eq!(
2308            resolve_docker_service(LOGICAL_GPU_SERVICE, Path::new("/nonexistent")),
2309            "cuda",
2310        );
2311    }
2312
2313    #[test]
2314    fn testing_env_var_names_match_their_source_of_truth() {
2315        // The names are literals here (flodl-cli is decoupled from the
2316        // flodl library by policy). flodl-hw IS a dependency, so at
2317        // least that one can be pinned to its constant rather than
2318        // trusted. A rename there would otherwise silently stop the
2319        // forward, and the container would fall back to real hardware
2320        // while the test still passed.
2321        assert!(
2322            TESTING_ENV_VARS.contains(&flodl_hw::ENV_TESTING_GPU_JSON),
2323            "flodl_hw::ENV_TESTING_GPU_JSON = {:?} is not forwarded into docker; \
2324             TESTING_ENV_VARS = {TESTING_ENV_VARS:?}",
2325            flodl_hw::ENV_TESTING_GPU_JSON,
2326        );
2327    }
2328}