1use 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
17pub fn config_to_args(resolved: &ResolvedConfig) -> Vec<String> {
21 let mut args = Vec::new();
22
23 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 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 let o = &resolved.output;
64 push_opt(&mut args, "--output", &o.dir);
65 push_num(&mut args, "--monitor", &o.monitor);
66
67 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
126fn inside_docker() -> bool {
130 Path::new("/.dockerenv").exists()
131}
132
133const DEFAULT_CONTAINER_PROJECT_ROOT: &str = "/workspace";
138
139static COMPOSE_MOUNT_CACHE: std::sync::OnceLock<
146 std::collections::HashMap<String, String>,
147> = std::sync::OnceLock::new();
148
149fn container_project_root(project_root: &Path, service: &str) -> String {
165 let cache = COMPOSE_MOUNT_CACHE
166 .get_or_init(|| parse_compose_project_mounts(project_root));
167 cache
168 .get(service)
169 .cloned()
170 .unwrap_or_else(|| DEFAULT_CONTAINER_PROJECT_ROOT.to_string())
171}
172
173fn parse_compose_project_mounts(
182 project_root: &Path,
183) -> std::collections::HashMap<String, String> {
184 let compose_path = project_root.join("docker-compose.yml");
185 let text = match std::fs::read_to_string(&compose_path) {
186 Ok(t) => t,
187 Err(_) => return std::collections::HashMap::new(),
188 };
189 let doc: serde_yaml_ng::Value = match serde_yaml_ng::from_str(&text) {
190 Ok(d) => d,
191 Err(_) => return std::collections::HashMap::new(),
192 };
193 let mut out = std::collections::HashMap::new();
194 let services = match doc.get("services").and_then(|v| v.as_mapping()) {
195 Some(s) => s,
196 None => return out,
197 };
198 for (name, svc) in services {
199 let svc_name = match name.as_str() {
200 Some(s) => s,
201 None => continue,
202 };
203 let volumes = match svc.get("volumes").and_then(|v| v.as_sequence()) {
204 Some(v) => v,
205 None => continue,
206 };
207 if let Some(container_path) = find_project_mount(volumes) {
208 let cleaned = container_path.trim_end_matches('/').to_string();
211 let cleaned = if cleaned.is_empty() {
212 "/".to_string()
213 } else {
214 cleaned
215 };
216 out.insert(svc_name.to_string(), cleaned);
217 }
218 }
219 out
220}
221
222fn find_project_mount(volumes: &[serde_yaml_ng::Value]) -> Option<String> {
226 for entry in volumes {
227 if let Some(s) = entry.as_str() {
228 let mut parts = s.splitn(3, ':');
233 let host = parts.next()?;
234 let container = parts.next()?;
235 if host == "." || host == "./" {
236 return Some(container.to_string());
237 }
238 } else if let Some(m) = entry.as_mapping() {
239 let source = m
241 .get(serde_yaml_ng::Value::String("source".into()))
242 .and_then(|v| v.as_str());
243 let target = m
244 .get(serde_yaml_ng::Value::String("target".into()))
245 .and_then(|v| v.as_str());
246 if matches!(source, Some(".") | Some("./")) {
247 if let Some(t) = target {
248 return Some(t.to_string());
249 }
250 }
251 }
252 }
253 None
254}
255
256fn libtorch_env(project_root: &Path) -> Result<Vec<(String, String)>, String> {
262 let mut env = Vec::new();
263
264 env.push((
266 "LIBTORCH_CPU_PATH".into(),
267 "./libtorch/precompiled/cpu".into(),
268 ));
269
270 if let Some((info, host_path)) = resolve_libtorch(project_root)? {
271 env.push(("LIBTORCH_HOST_PATH".into(), host_path));
272
273 if let Some(cuda) = &info.cuda_version {
275 if cuda != "none" {
276 let cuda_version = if cuda.matches('.').count() < 2 {
277 format!("{cuda}.0")
278 } else {
279 cuda.clone()
280 };
281 let cuda_tag = cuda_version
282 .splitn(3, '.')
283 .take(2)
284 .collect::<Vec<_>>()
285 .join(".");
286 env.push(("CUDA_VERSION".into(), cuda_version));
287 env.push(("CUDA_TAG".into(), cuda_tag));
288 }
289 }
290 }
291
292 Ok(env)
293}
294
295fn resolve_libtorch(
313 project_root: &Path,
314) -> Result<Option<(libtorch::detect::LibtorchInfo, String)>, String> {
315 if let Some(resolved) = resolve_libtorch_from_overlay(project_root)? {
316 return Ok(Some(resolved));
317 }
318 let Some(info) = libtorch::detect::read_active(project_root) else {
319 return Ok(None);
320 };
321 let host_path = format!("./libtorch/{}", info.path);
322 Ok(Some((info, host_path)))
323}
324
325fn resolve_libtorch_from_overlay(
338 project_root: &Path,
339) -> Result<Option<(libtorch::detect::LibtorchInfo, String)>, String> {
340 let Ok(env_name) = std::env::var("FDL_ENV") else {
341 return Ok(None);
342 };
343 let env_name = env_name.trim();
344 if env_name.is_empty() {
345 return Ok(None);
346 }
347 let base_path = config::find_config_in(project_root).ok_or_else(|| {
350 format!(
351 "FDL_ENV={env_name} is set but no fdl config file exists in {}",
352 project_root.display()
353 )
354 })?;
355 let cfg = config::load_project_with_env(&base_path, Some(env_name))
356 .map_err(|e| format!("FDL_ENV={env_name}: cannot resolve the overlay: {e}"))?;
357 let Some(cluster) = cfg.cluster else {
358 return Ok(None);
359 };
360 let host_name = crate::cluster::resolve_local_hostname();
361 let Some(entry) = cluster.workers.iter().find(|w| w.host == host_name) else {
362 return Ok(None);
363 };
364 let Some(arch) = entry.arch.as_ref() else {
365 return Ok(None);
366 };
367 let variant_dir = std::path::PathBuf::from(&entry.path)
368 .join("libtorch")
369 .join(arch);
370 Ok(resolve_libtorch_at(&variant_dir))
371}
372
373pub(crate) fn resolve_libtorch_at(
383 path: &Path,
384) -> Option<(libtorch::detect::LibtorchInfo, String)> {
385 if path.is_file()
386 && path.file_name()
387 .and_then(|n| n.to_str())
388 .is_some_and(|n| n.starts_with(".active"))
389 {
390 let libtorch_root = path.parent()?;
391 let info = libtorch::detect::read_active_from(path, libtorch_root)?;
392 let host_path = libtorch_root.join(&info.path).display().to_string();
393 return Some((info, host_path));
394 }
395 if path.join(".active").exists() {
396 let info = libtorch::detect::read_active_from(
397 &path.join(".active"),
398 path,
399 )?;
400 let host_path = path.join(&info.path).display().to_string();
401 return Some((info, host_path));
402 }
403 if path.join("lib").is_dir() {
404 let info = libtorch::detect::libtorch_info_from_dir(
405 path.display().to_string(),
406 path,
407 );
408 let host_path = path.display().to_string();
409 return Some((info, host_path));
410 }
411 None
412}
413
414fn spawn_docker_shell(command: &str, project_root: &Path) -> ExitCode {
420 let env_vars = match libtorch_env(project_root) {
421 Ok(v) => v,
422 Err(e) => {
423 eprintln!("fdl: {e}");
424 return ExitCode::FAILURE;
425 }
426 };
427
428 let mut cmd = std::process::Command::new("sh");
429 cmd.args(["-c", command])
430 .current_dir(project_root)
431 .env(
436 "HOSTNAME",
437 crate::cluster::resolve_local_hostname(),
438 )
439 .stdout(Stdio::inherit())
440 .stderr(Stdio::inherit())
441 .stdin(Stdio::inherit());
442
443 for (key, val) in &env_vars {
444 cmd.env(key, val);
445 }
446
447 match cmd.status() {
448 Ok(s) if s.success() => ExitCode::SUCCESS,
449 Ok(s) => ExitCode::from(s.code().unwrap_or(1) as u8),
450 Err(e) => {
451 cli_error!("{e}");
452 ExitCode::FAILURE
453 }
454 }
455}
456
457pub(crate) use crate::util::shell::posix_quote;
460
461fn split_append_dashdash(s: &str) -> (String, String) {
472 let s = s.trim();
473 if s == "--" {
474 return (String::new(), String::new());
475 }
476 if let Some(rest) = s.strip_prefix("-- ") {
477 return (String::new(), rest.trim().to_string());
478 }
479 if let Some(prefix) = s.strip_suffix(" --") {
480 return (prefix.trim().to_string(), String::new());
481 }
482 if let Some(idx) = s.find(" -- ") {
483 let before = &s[..idx];
484 let after = &s[idx + 4..];
485 return (before.trim().to_string(), after.trim().to_string());
486 }
487 (s.to_string(), String::new())
488}
489
490fn split_user_args_dashdash(args: &[String]) -> (&[String], Option<&[String]>) {
496 match args.iter().position(|a| a == "--") {
497 Some(idx) => (&args[..idx], Some(&args[idx + 1..])),
498 None => (args, None),
499 }
500}
501
502pub(crate) fn compose_run_command(
517 run: &str,
518 user_args: &[String],
519 append: Option<&str>,
520) -> String {
521 let suffix = append.map(str::trim).filter(|s| !s.is_empty());
522 let (append_pre, append_post, append_has_dashdash) = match suffix {
523 Some(s) => {
524 let (pre, post) = split_append_dashdash(s);
525 let has = s == "--"
530 || s.starts_with("-- ")
531 || s.ends_with(" --")
532 || s.contains(" -- ");
533 (pre, post, has)
534 }
535 None => (String::new(), String::new(), false),
536 };
537 let (user_pre, user_post_opt) = split_user_args_dashdash(user_args);
538
539 let mut out = String::from(run.trim());
540 if !append_pre.is_empty() {
541 out.push(' ');
542 out.push_str(&append_pre);
543 }
544 for a in user_pre {
545 out.push(' ');
546 out.push_str(&posix_quote(a));
547 }
548
549 let needs_separator = append_has_dashdash
550 || !append_post.is_empty()
551 || user_post_opt.is_some_and(|p| !p.is_empty());
552 if needs_separator {
553 out.push_str(" --");
554 if !append_post.is_empty() {
555 out.push(' ');
556 out.push_str(&append_post);
557 }
558 if let Some(post) = user_post_opt {
559 for a in post {
560 out.push(' ');
561 out.push_str(&posix_quote(a));
562 }
563 }
564 }
565 out
566}
567
568fn testing_cluster_env_arg() -> String {
593 match std::env::var("FLODL_TESTING_CLUSTER_JSON") {
594 Ok(_) => " -e FLODL_TESTING_CLUSTER_JSON".to_string(),
595 Err(_) => String::new(),
596 }
597}
598
599pub fn exec_script(
600 command: &str,
601 append: Option<&str>,
602 user_args: &[String],
603 docker_service: Option<&str>,
604 cwd: &Path,
605) -> ExitCode {
606 let inner_cmd = compose_run_command(command, user_args, append);
607
608 match docker_service {
609 Some(service) if !inside_docker() => {
610 let overlay = crate::cluster::cluster_compose_overlay_arg(cwd);
614 let testing_env_arg = testing_cluster_env_arg();
615 let docker_cmd = format!(
616 "docker compose{overlay} run --rm{testing_env_arg} {service} bash -c {}",
617 posix_quote(&inner_cmd)
618 );
619 spawn_docker_shell(&docker_cmd, cwd)
620 }
621 _ => {
622 let (shell, flag) = if cfg!(target_os = "windows") {
623 ("cmd", "/C")
624 } else {
625 ("sh", "-c")
626 };
627
628 match std::process::Command::new(shell)
629 .args([flag, inner_cmd.as_str()])
630 .current_dir(cwd)
631 .stdout(Stdio::inherit())
632 .stderr(Stdio::inherit())
633 .stdin(Stdio::inherit())
634 .status()
635 {
636 Ok(s) if s.success() => ExitCode::SUCCESS,
637 Ok(s) => ExitCode::from(s.code().unwrap_or(1) as u8),
638 Err(e) => {
639 cli_error!("{e}");
640 ExitCode::FAILURE
641 }
642 }
643 }
644 }
645}
646
647pub fn exec_command(
654 cmd_config: &CommandConfig,
655 preset_name: Option<&str>,
656 extra_args: &[String],
657 cmd_dir: &Path,
658 project_root: &Path,
659) -> ExitCode {
660 let entry = match &cmd_config.entry {
661 Some(e) => e.as_str(),
662 None => {
663 eprintln!(
664 "error: no entry point defined in {}/fdl.yaml",
665 cmd_dir.display()
666 );
667 return ExitCode::FAILURE;
668 }
669 };
670
671 if let Some(schema) = &cmd_config.schema {
680 if let Err(e) = config::validate_tail(extra_args, schema) {
681 cli_error!("{e}");
682 return ExitCode::FAILURE;
683 }
684 }
685
686 let resolved = match preset_name {
688 Some(name) => match cmd_config.commands.get(name) {
689 Some(preset) => {
690 if let Some(schema) = &cmd_config.schema {
694 if let Err(e) = config::validate_preset_for_exec(name, preset, schema) {
695 cli_error!("{e}");
696 return ExitCode::FAILURE;
697 }
698 }
699 config::merge_preset(cmd_config, preset)
700 }
701 None => {
702 cli_error!("unknown command '{name}'");
703 eprintln!();
704 print_command_help(cmd_config, "");
705 return ExitCode::FAILURE;
706 }
707 },
708 None => config::defaults_only(cmd_config),
709 };
710
711 let mut args = config_to_args(&resolved);
713
714 args.extend(extra_args.iter().cloned());
716
717 let use_docker = cmd_config.docker.is_some() && !inside_docker();
719
720 if use_docker {
721 let service = cmd_config.docker.as_deref().unwrap();
722 let workdir = cmd_dir
723 .strip_prefix(project_root)
724 .unwrap_or(cmd_dir)
725 .to_string_lossy();
726
727 let container_root = container_project_root(project_root, service);
736 let args_str = shell_join(&args);
737 let inner = if workdir.is_empty() || workdir == "." {
738 format!("{entry} {args_str}")
739 } else {
740 format!(
741 "cd {} && {entry} {args_str}",
742 posix_quote(&format!("{container_root}/{workdir}"))
743 )
744 };
745
746 if preset_name.is_some() {
747 eprintln!("fdl: [{service}] {inner}");
748 }
749
750 let overlay = crate::cluster::cluster_compose_overlay_arg(project_root);
758 let testing_env_arg = testing_cluster_env_arg();
759 let docker_cmd = format!(
768 "docker compose{overlay} run --rm -e {}{testing_env_arg} {service} bash -c {}",
769 posix_quote(&format!("FDL_PROJECT_ROOT={container_root}")),
770 posix_quote(&inner),
771 );
772 spawn_docker_shell(&docker_cmd, project_root)
773 } else {
774 let parts: Vec<&str> = entry.split_whitespace().collect();
776 if parts.is_empty() {
777 cli_error!("empty entry point");
778 return ExitCode::FAILURE;
779 }
780 let program = parts[0];
781 let entry_args = &parts[1..];
782
783 if preset_name.is_some() {
784 let preview: Vec<&str> = args.iter().map(|s| s.as_str()).collect();
785 eprintln!("fdl: {entry} {}", preview.join(" "));
786 }
787
788 match std::process::Command::new(program)
789 .args(entry_args)
790 .args(&args)
791 .current_dir(cmd_dir)
792 .stdout(Stdio::inherit())
793 .stderr(Stdio::inherit())
794 .stdin(Stdio::inherit())
795 .status()
796 {
797 Ok(s) if s.success() => ExitCode::SUCCESS,
798 Ok(s) => ExitCode::from(s.code().unwrap_or(1) as u8),
799 Err(e) => {
800 cli_error!("failed to execute '{program}': {e}");
801 ExitCode::FAILURE
802 }
803 }
804 }
805}
806
807fn shell_join(args: &[String]) -> String {
815 args.iter()
816 .map(|a| posix_quote(a))
817 .collect::<Vec<_>>()
818 .join(" ")
819}
820
821pub fn print_run_help(
827 name: &str,
828 description: Option<&str>,
829 run: &str,
830 append: Option<&str>,
831 docker: Option<&str>,
832) {
833 if let Some(desc) = description {
834 eprintln!("{} {desc}", style::bold(name));
835 } else {
836 eprintln!("{}", style::bold(name));
837 }
838 eprintln!();
839 eprintln!("{}:", style::yellow("Usage"));
840 eprintln!(" fdl {name} [-- <args>... [-- <runner-args>...]]");
841 eprintln!();
842 eprintln!("{}:", style::yellow("Runs"));
843 let composed = match append.map(str::trim).filter(|s| !s.is_empty()) {
844 Some(suffix) => {
845 let (pre, post) = split_append_dashdash(suffix);
846 let left = if pre.is_empty() {
850 "[<args>]".to_string()
851 } else {
852 format!("{pre} [<args>]")
853 };
854 let right = if post.is_empty() {
855 "[<runner-args>]".to_string()
856 } else {
857 format!("{post} [<runner-args>]")
858 };
859 format!("{run} {left} -- {right}")
860 }
861 None => format!("{run} [<args>] [-- <runner-args>]"),
862 };
863 if let Some(svc) = docker {
864 eprintln!(
865 " {} {svc} -c {composed:?}",
866 style::dim("docker compose run --rm")
867 );
868 } else {
869 eprintln!(" {composed}");
870 }
871 eprintln!();
872 eprintln!(
873 "{} the first `--` separates fdl args from the run script; a second `--` splits cargo-side args from runner-side args.",
874 style::dim("Note:"),
875 );
876 eprintln!(
877 "{} `append:` is split on its own `--` and merged half-and-half; pass `--no-append` to drop it entirely.",
878 style::dim("Note:"),
879 );
880}
881
882pub fn print_command_help(cmd_config: &CommandConfig, name: &str) {
885 let (presets, sub_cmds) = split_commands_by_kind(&cmd_config.commands);
886 let preset_slot = cmd_config.arg_name.as_deref().unwrap_or("preset");
887
888 let width = help_width();
890
891 print_title(cmd_config, name);
892 print_usage_line(cmd_config, name, &presets, &sub_cmds, preset_slot);
893 print_arguments_section(cmd_config, &presets, preset_slot, width);
894 print_sub_commands_section(&sub_cmds);
895 print_schema_commands_section(cmd_config, name);
896 print_options_section(cmd_config, width);
897 print_entry_section(cmd_config);
898 print_defaults_section(cmd_config);
899}
900
901fn print_title(cmd_config: &CommandConfig, name: &str) {
902 if let Some(desc) = &cmd_config.description {
903 eprintln!("{} {desc}", style::bold(name));
904 } else {
905 eprintln!("{}", style::bold(name));
906 }
907}
908
909fn print_usage_line(
910 cmd_config: &CommandConfig,
911 name: &str,
912 presets: &CommandGroup,
913 sub_cmds: &CommandGroup,
914 preset_slot: &str,
915) {
916 let usage_tail = build_usage_tail(
919 cmd_config.schema.as_ref(),
920 !presets.is_empty(),
921 !sub_cmds.is_empty(),
922 preset_slot,
923 );
924 eprintln!();
925 eprintln!("{}:", style::yellow("Usage"));
926 eprintln!(" fdl {name}{usage_tail}");
927}
928
929fn print_arguments_section(
930 cmd_config: &CommandConfig,
931 presets: &CommandGroup,
932 preset_slot: &str,
933 width: usize,
934) {
935 let has_schema_args = cmd_config
941 .schema
942 .as_ref()
943 .is_some_and(|s| !s.args.is_empty());
944 if !has_schema_args && presets.is_empty() {
945 return;
946 }
947 eprintln!();
948 eprintln!("{}:", style::yellow("Arguments"));
949 let avail = width.saturating_sub(4);
950 if let Some(schema) = &cmd_config.schema {
951 for a in &schema.args {
952 for line in format_arg(a, avail) {
953 eprintln!(" {line}");
954 }
955 }
956 }
957 if !presets.is_empty() {
958 let slot_label = format!("[<{preset_slot}>]");
959 eprintln!(
960 " {} Named preset, one of:",
961 style::green(&format!("{:<20}", slot_label))
962 );
963 for (pname, spec) in presets {
964 let desc = spec.description.as_deref().unwrap_or("-");
965 eprintln!(
966 " {} {}",
967 style::green(&format!("{:<18}", pname)),
968 desc
969 );
970 }
971 }
972}
973
974fn print_sub_commands_section(sub_cmds: &CommandGroup) {
975 if sub_cmds.is_empty() {
978 return;
979 }
980 eprintln!();
981 eprintln!("{}:", style::yellow("Commands"));
982 for (sub_name, sub_spec) in sub_cmds {
983 let desc = sub_spec.description.as_deref().unwrap_or("-");
984 eprintln!(
985 " {} {}",
986 style::green(&format!("{:<20}", sub_name)),
987 desc
988 );
989 }
990}
991
992fn print_schema_commands_section(cmd_config: &CommandConfig, name: &str) {
998 let Some(schema) = &cmd_config.schema else {
999 return;
1000 };
1001 if schema.commands.is_empty() {
1002 return;
1003 }
1004 eprintln!();
1005 eprintln!("{}:", style::yellow("Commands"));
1006 for (sub_name, sub_schema) in &schema.commands {
1007 let desc = sub_schema.description.as_deref().unwrap_or("-");
1008 eprintln!(
1009 " {} {}",
1010 style::green(&format!("{:<20}", sub_name)),
1011 desc
1012 );
1013 }
1014 eprintln!();
1015 eprintln!(
1016 " Run {} for a subcommand's options.",
1017 style::dim(&format!("fdl {name} <command> --help"))
1018 );
1019}
1020
1021fn print_options_section(cmd_config: &CommandConfig, width: usize) {
1022 let Some(schema) = &cmd_config.schema else {
1025 return;
1026 };
1027 if schema.options.is_empty() {
1028 return;
1029 }
1030 eprintln!();
1031 eprintln!("{}:", style::yellow("Options"));
1032 let avail = width.saturating_sub(4);
1033 for (long, spec) in &schema.options {
1034 for line in format_option(long, spec, avail) {
1035 eprintln!(" {line}");
1036 }
1037 }
1038}
1039
1040fn print_entry_section(cmd_config: &CommandConfig) {
1041 let Some(entry) = &cmd_config.entry else {
1042 return;
1043 };
1044 eprintln!();
1045 eprintln!("{}:", style::yellow("Entry"));
1046 eprintln!(" {entry}");
1047 if let Some(service) = &cmd_config.docker {
1048 eprintln!(
1049 " {}",
1050 style::dim(&format!("[docker: {service}]"))
1051 );
1052 }
1053 eprintln!();
1054 eprintln!(
1055 " Any extra {} are forwarded to the entry point.",
1056 style::dim("[options]")
1057 );
1058}
1059
1060fn print_defaults_section(cmd_config: &CommandConfig) {
1061 if cmd_config.ddp.is_none() && cmd_config.training.is_none() {
1062 return;
1063 }
1064 eprintln!();
1065 eprintln!("{}:", style::yellow("Defaults"));
1066 if let Some(d) = &cmd_config.ddp {
1067 if let Some(mode) = &d.mode {
1068 eprintln!(" {} {mode}", style::dim("ddp.mode"));
1069 }
1070 if let Some(anchor) = &d.anchor {
1071 eprintln!(" {} {}", style::dim("ddp.anchor"), value_to_string(anchor));
1072 }
1073 }
1074 if let Some(t) = &cmd_config.training {
1075 if let Some(e) = t.epochs {
1076 eprintln!(" {} {e}", style::dim("training.epochs"));
1077 }
1078 if let Some(bs) = t.batch_size {
1079 eprintln!(" {} {bs}", style::dim("training.batch_size"));
1080 }
1081 if let Some(lr) = t.lr {
1082 eprintln!(" {} {lr}", style::dim("training.lr"));
1083 }
1084 if let Some(seed) = t.seed {
1085 eprintln!(" {} {seed}", style::dim("training.seed"));
1086 }
1087 }
1088}
1089
1090pub fn print_preset_help(cmd_config: &CommandConfig, cmd_name: &str, preset_name: &str) {
1092 let preset = match cmd_config.commands.get(preset_name) {
1093 Some(s) => s,
1094 None => {
1095 eprintln!("unknown command: {preset_name}");
1096 return;
1097 }
1098 };
1099
1100 let desc = preset.description.as_deref().unwrap_or("(no description)");
1102 eprintln!(
1103 "{} {} {}",
1104 style::bold(cmd_name),
1105 style::green(preset_name),
1106 desc
1107 );
1108
1109 eprintln!();
1110 eprintln!("{}:", style::yellow("Usage"));
1111 eprintln!(
1112 " fdl {cmd_name} {preset_name} {}",
1113 style::dim("[extra options]")
1114 );
1115
1116 let resolved = config::merge_preset(cmd_config, preset);
1118
1119 eprintln!();
1120 eprintln!("{}:", style::yellow("Effective config"));
1121
1122 let d = &resolved.ddp;
1124 print_config_field("ddp.mode", &d.mode);
1125 print_config_value("ddp.anchor", &d.anchor);
1126 print_config_field("ddp.max_anchor", &d.max_anchor);
1127 print_config_field("ddp.overhead_target", &d.overhead_target);
1128 print_config_field("ddp.divergence_threshold", &d.divergence_threshold);
1129 print_config_value("ddp.max_batch_diff", &d.max_batch_diff);
1130 print_config_field("ddp.max_grad_norm", &d.max_grad_norm);
1131 if d.timeline == Some(true) {
1132 eprintln!(" {} true", style::dim("ddp.timeline"));
1133 }
1134
1135 let t = &resolved.training;
1137 print_config_field("training.epochs", &t.epochs);
1138 print_config_field("training.batch_size", &t.batch_size);
1139 print_config_field("training.batches_per_epoch", &t.batches_per_epoch);
1140 print_config_field("training.lr", &t.lr);
1141 print_config_field("training.seed", &t.seed);
1142
1143 let o = &resolved.output;
1145 print_config_field("output.dir", &o.dir);
1146 print_config_field("output.monitor", &o.monitor);
1147
1148 if !resolved.options.is_empty() {
1150 eprintln!();
1151 eprintln!("{}:", style::yellow("Options"));
1152 for (key, val) in &resolved.options {
1153 eprintln!(
1154 " {} {}",
1155 style::green(&format!("--{}", key.replace('_', "-"))),
1156 value_to_string(val)
1157 );
1158 }
1159 }
1160
1161 if let Some(entry) = &cmd_config.entry {
1163 let args = config_to_args(&resolved);
1164 let args_str = args.join(" ");
1165 let docker_info = cmd_config
1166 .docker
1167 .as_ref()
1168 .map(|s| format!("[{s}] ", ))
1169 .unwrap_or_default();
1170
1171 eprintln!();
1172 eprintln!("{}:", style::yellow("Effective command"));
1173 eprintln!(
1174 " {}{}{}",
1175 style::dim(&docker_info),
1176 entry,
1177 if args_str.is_empty() {
1178 String::new()
1179 } else {
1180 format!(" {args_str}")
1181 }
1182 );
1183 }
1184
1185 eprintln!();
1186 eprintln!(
1187 "Extra {} after the command name are appended to the entry.",
1188 style::dim("[options]")
1189 );
1190}
1191
1192fn print_config_field<T: std::fmt::Display>(label: &str, val: &Option<T>) {
1193 if let Some(v) = val {
1194 eprintln!(" {} {v}", style::dim(label));
1195 }
1196}
1197
1198fn print_config_value(label: &str, val: &Option<serde_json::Value>) {
1199 if let Some(v) = val {
1200 if !v.is_null() {
1201 eprintln!(" {} {}", style::dim(label), value_to_string(v));
1202 }
1203 }
1204}
1205
1206pub fn print_project_help(
1208 project: &config::ProjectConfig,
1209 project_root: &Path,
1210 active_env: Option<&str>,
1211) {
1212 let visible_builtins = builtins::visible_top_level();
1213 if let Some(desc) = &project.description {
1214 eprintln!("{} {}", style::bold("fdl"), desc);
1215 } else {
1216 eprintln!("{} {}", style::bold("fdl"), env!("CARGO_PKG_VERSION"));
1217 }
1218
1219 eprintln!();
1220 eprintln!("{}:", style::yellow("Usage"));
1221 eprintln!(
1222 " fdl {} {}",
1223 style::dim("<command>"),
1224 style::dim("[options]")
1225 );
1226
1227 eprintln!();
1228 eprintln!("{}:", style::yellow("Options"));
1229 eprintln!(
1230 " {} Show this help",
1231 style::green(&format!("{:<18}", "-h, --help"))
1232 );
1233 eprintln!(
1234 " {} Show version",
1235 style::green(&format!("{:<18}", "-V, --version"))
1236 );
1237 eprintln!(
1238 " {} Use fdl.<name>.yml overlay (also: --env <name>, FDL_ENV=<name>)",
1239 style::green(&format!("{:<18}", "@<name>"))
1240 );
1241 eprintln!(
1242 " {} Verbose output",
1243 style::green(&format!("{:<18}", "-v"))
1244 );
1245 eprintln!(
1246 " {} Debug output",
1247 style::green(&format!("{:<18}", "-vv"))
1248 );
1249 eprintln!(
1250 " {} Trace output (maximum detail)",
1251 style::green(&format!("{:<18}", "-vvv"))
1252 );
1253 eprintln!(
1254 " {} Suppress non-error output",
1255 style::green(&format!("{:<18}", "-q, --quiet"))
1256 );
1257 eprintln!(
1258 " {} Force ANSI color (bypass TTY / NO_COLOR detection)",
1259 style::green(&format!("{:<18}", "--ansi"))
1260 );
1261 eprintln!(
1262 " {} Disable ANSI color output",
1263 style::green(&format!("{:<18}", "--no-ansi"))
1264 );
1265 eprintln!(
1266 " {} Drop a run command's `append:` suffix",
1267 style::green(&format!("{:<18}", "--no-append"))
1268 );
1269 eprintln!(
1270 " {} Skip the cluster pre-flight build",
1271 style::green(&format!("{:<18}", "--no-prebuild"))
1272 );
1273
1274 eprintln!();
1276 eprintln!("{}:", style::yellow("Built-in"));
1277 for (name, desc) in &visible_builtins {
1278 eprintln!(" {} {desc}", style::green(&format!("{:<18}", name)));
1279 }
1280
1281 if !project.commands.is_empty() {
1288 eprintln!();
1289 eprintln!("{}:", style::yellow("Commands"));
1290 for (name, spec) in &project.commands {
1291 let desc: String = match spec.description.clone() {
1292 Some(d) => d,
1293 None => {
1294 let is_path_kind = spec.run.is_none();
1298 if is_path_kind {
1299 let child_dir = spec.resolve_path(name, project_root);
1300 config::load_command_with_env(&child_dir, active_env)
1301 .ok()
1302 .and_then(|c| c.description)
1303 .unwrap_or_else(|| "(sub-command)".into())
1304 } else {
1305 spec.run
1306 .as_deref()
1307 .unwrap_or("(command)")
1308 .to_string()
1309 }
1310 }
1311 };
1312 eprintln!(" {} {desc}", style::green(&format!("{:<18}", name)));
1313 }
1314 }
1315
1316 if let Some(base_config) = config::find_config(project_root) {
1318 let envs = crate::overlay::list_envs(&base_config);
1319 if !envs.is_empty() {
1320 eprintln!();
1321 eprintln!("{}:", style::yellow("Environments"));
1322 for e in &envs {
1323 let active_marker = if Some(e.as_str()) == active_env {
1324 style::green(" (active)")
1325 } else {
1326 String::new()
1327 };
1328 eprintln!(
1329 " {} Overlay from fdl.{}.yml{active_marker}",
1330 style::green(&format!("{:<18}", format!("@{e}"))),
1331 e
1332 );
1333 }
1334 eprintln!();
1335 eprintln!(
1336 "Use {} to run a command with an environment overlay.",
1337 style::dim("fdl @<env> <command>")
1338 );
1339 }
1340 }
1341
1342 eprintln!();
1343 eprintln!(
1344 "Use {} for more information on a command.",
1345 style::dim("fdl <command> -h")
1346 );
1347}
1348
1349fn build_usage_tail(
1358 schema: Option<&Schema>,
1359 has_presets: bool,
1360 has_sub_commands: bool,
1361 preset_slot: &str,
1362) -> String {
1363 let mut parts = String::new();
1364 let slot = match (has_presets, has_sub_commands) {
1365 (true, false) => Some(format!("[<{preset_slot}>]")),
1366 (false, true) => Some("[<command>]".to_string()),
1367 (true, true) => Some(format!("[<{preset_slot}>|<command>]")),
1368 (false, false) => None,
1369 };
1370 if let Some(s) = slot {
1371 parts.push(' ');
1372 parts.push_str(&style::dim(&s));
1373 }
1374 if let Some(s) = schema {
1375 for a in &s.args {
1376 parts.push(' ');
1377 parts.push_str(&format_arg_usage(a));
1378 }
1379 }
1380 parts.push(' ');
1381 parts.push_str(&style::dim("[options]"));
1382 parts
1383}
1384
1385type CommandGroup = Vec<(String, crate::config::CommandSpec)>;
1386
1387fn split_commands_by_kind(
1392 commands: &BTreeMap<String, crate::config::CommandSpec>,
1393) -> (CommandGroup, CommandGroup) {
1394 use crate::config::CommandKind;
1395 let mut presets = Vec::new();
1396 let mut sub_cmds = Vec::new();
1397 for (k, v) in commands {
1398 match v.kind() {
1399 Ok(CommandKind::Preset) => presets.push((k.clone(), v.clone())),
1400 _ => sub_cmds.push((k.clone(), v.clone())),
1401 }
1402 }
1403 (presets, sub_cmds)
1404}
1405
1406fn format_arg_usage(a: &ArgSpec) -> String {
1407 let suffix = if a.variadic { "..." } else { "" };
1408 let core = format!("<{}>{suffix}", a.name);
1409 if a.required && a.default.is_none() {
1410 style::green(&core)
1411 } else {
1412 style::dim(&format!("[{core}]"))
1413 }
1414}
1415
1416const ARG_COL: usize = 22;
1420const OPT_COL: usize = 30;
1423
1424fn format_arg(a: &ArgSpec, avail_width: usize) -> Vec<String> {
1425 let left = format_arg_usage(a);
1426 let visible = visible_width(&left);
1427 let segs = desc_segments(a.description.as_deref(), &a.default, &a.choices, &a.ty);
1428 format_row(&left, visible, ARG_COL, &segs, avail_width)
1429}
1430
1431fn format_option(long: &str, spec: &OptionSpec, avail_width: usize) -> Vec<String> {
1436 let flag = match &spec.short {
1437 Some(s) => format!("-{s}, --{long}"),
1438 None => format!(" --{long}"),
1439 };
1440 let placeholder = option_placeholder(&spec.ty);
1441 let (left, visible) = if placeholder.is_empty() {
1442 (style::green(&flag), flag.chars().count())
1443 } else {
1444 (
1445 style::green(&format!("{flag} {placeholder}")),
1446 flag.chars().count() + 1 + placeholder.chars().count(),
1447 )
1448 };
1449 let segs = desc_segments(spec.description.as_deref(), &spec.default, &spec.choices, &spec.ty);
1450 let mut out = format_row(&left, visible, OPT_COL, &segs, avail_width);
1451 if let Some(env) = &spec.env {
1452 out.push(format!(
1453 "{}{}",
1454 " ".repeat(OPT_COL),
1455 style::dim(&format!("[env: {env}]"))
1456 ));
1457 }
1458 out
1459}
1460
1461struct Seg {
1465 text: String,
1466 styled: String,
1467}
1468
1469impl Seg {
1470 fn plain(s: &str) -> Seg {
1471 Seg { text: s.to_string(), styled: s.to_string() }
1472 }
1473 fn dim(s: &str) -> Seg {
1474 Seg { text: s.to_string(), styled: style::dim(s) }
1475 }
1476}
1477
1478fn desc_segments(
1483 description: Option<&str>,
1484 default: &Option<serde_json::Value>,
1485 choices: &Option<Vec<serde_json::Value>>,
1486 ty: &str,
1487) -> Vec<Seg> {
1488 let mut segs: Vec<Seg> = description
1489 .unwrap_or("-")
1490 .split_whitespace()
1491 .map(Seg::plain)
1492 .collect();
1493 if let Some(d) = default {
1494 let is_empty_list = matches!(d, serde_json::Value::Array(a) if a.is_empty());
1496 let is_false = matches!(d, serde_json::Value::Bool(false));
1497 if !d.is_null() && !is_false && !is_empty_list {
1498 segs.push(Seg::dim(&format!("[default: {}]", format_value(d))));
1499 }
1500 }
1501 if let Some(choices) = choices {
1502 if !choices.is_empty() {
1503 let list = choices.iter().map(format_value).collect::<Vec<_>>().join(", ");
1504 segs.push(Seg::dim(&format!("[possible: {list}]")));
1505 }
1506 }
1507 if ty.starts_with("list[") {
1509 segs.push(Seg::dim("(repeat or comma-separate)"));
1510 }
1511 segs
1512}
1513
1514fn wrap_segments(segs: &[Seg], width: usize) -> Vec<String> {
1519 let width = width.max(1);
1520 let mut lines: Vec<String> = Vec::new();
1521 let mut cur = String::new();
1522 let mut cur_w = 0usize;
1523 for seg in segs {
1524 let w = seg.text.chars().count();
1525 if cur_w == 0 {
1526 cur.push_str(&seg.styled);
1527 cur_w = w;
1528 } else if cur_w + 1 + w <= width {
1529 cur.push(' ');
1530 cur.push_str(&seg.styled);
1531 cur_w += 1 + w;
1532 } else {
1533 lines.push(std::mem::take(&mut cur));
1534 cur.push_str(&seg.styled);
1535 cur_w = w;
1536 }
1537 }
1538 if !cur.is_empty() {
1539 lines.push(cur);
1540 }
1541 if lines.is_empty() {
1542 lines.push(String::new());
1543 }
1544 lines
1545}
1546
1547fn format_row(
1554 label: &str,
1555 label_visible: usize,
1556 desc_col: usize,
1557 segs: &[Seg],
1558 avail_width: usize,
1559) -> Vec<String> {
1560 const MIN_DESC: usize = 20;
1562 let desc_width = avail_width.saturating_sub(desc_col).max(MIN_DESC);
1563 let desc_lines = wrap_segments(segs, desc_width);
1564 let pad = " ".repeat(desc_col);
1565 let mut out: Vec<String> = Vec::with_capacity(desc_lines.len() + 1);
1566 if label_visible < desc_col {
1567 let gap = " ".repeat(desc_col - label_visible);
1568 out.push(format!("{label}{gap}{}", desc_lines[0]));
1569 } else {
1570 out.push(label.to_string());
1572 out.push(format!("{pad}{}", desc_lines[0]));
1573 }
1574 for line in &desc_lines[1..] {
1575 out.push(format!("{pad}{line}"));
1576 }
1577 for line in &mut out {
1579 while line.ends_with(' ') {
1580 line.pop();
1581 }
1582 }
1583 out
1584}
1585
1586fn option_placeholder(ty: &str) -> &'static str {
1587 match ty {
1588 "bool" => "",
1589 "int" => "<N>",
1590 "float" => "<F>",
1591 "path" => "<PATH>",
1592 "list[path]" => "<PATH>...",
1593 t if t.starts_with("list[") => "<VALUE>...",
1594 _ => "<VALUE>",
1595 }
1596}
1597
1598fn format_value(v: &serde_json::Value) -> String {
1599 match v {
1600 serde_json::Value::String(s) => s.clone(),
1601 other => other.to_string(),
1602 }
1603}
1604
1605fn visible_width(s: &str) -> usize {
1608 strip_ansi(s).chars().count()
1611}
1612
1613fn help_width() -> usize {
1619 const DEFAULT: usize = 100;
1620 const MIN: usize = 60;
1621 const MAX: usize = 120;
1622 let raw = std::env::var("COLUMNS")
1623 .ok()
1624 .and_then(|s| s.trim().parse::<usize>().ok())
1625 .filter(|&c| c > 0)
1626 .or_else(term_cols)
1627 .unwrap_or(DEFAULT);
1628 raw.clamp(MIN, MAX)
1629}
1630
1631#[cfg(unix)]
1637fn term_cols() -> Option<usize> {
1638 use std::io::IsTerminal;
1639 use std::os::unix::io::AsRawFd;
1640
1641 let stderr = std::io::stderr();
1642 if !stderr.is_terminal() {
1643 return None;
1644 }
1645 #[repr(C)]
1646 struct Winsize {
1647 row: u16,
1648 col: u16,
1649 xpixel: u16,
1650 ypixel: u16,
1651 }
1652 #[cfg(any(target_os = "linux", target_os = "android"))]
1655 const TIOCGWINSZ: std::os::raw::c_ulong = 0x5413;
1656 #[cfg(not(any(target_os = "linux", target_os = "android")))]
1657 const TIOCGWINSZ: std::os::raw::c_ulong = 0x4008_7468;
1658 unsafe extern "C" {
1659 fn ioctl(fd: std::os::raw::c_int, request: std::os::raw::c_ulong, ...) -> std::os::raw::c_int;
1660 }
1661 let mut ws = Winsize { row: 0, col: 0, xpixel: 0, ypixel: 0 };
1662 let rc = unsafe { ioctl(stderr.as_raw_fd(), TIOCGWINSZ, &mut ws as *mut Winsize) };
1665 (rc == 0 && ws.col > 0).then_some(ws.col as usize)
1666}
1667
1668#[cfg(not(unix))]
1669fn term_cols() -> Option<usize> {
1670 None
1671}
1672
1673fn strip_ansi(s: &str) -> String {
1674 let mut out = String::with_capacity(s.len());
1675 let mut chars = s.chars().peekable();
1676 while let Some(c) = chars.next() {
1677 if c == '\x1b' && chars.peek() == Some(&'[') {
1678 chars.next();
1679 for c in chars.by_ref() {
1680 if c.is_ascii_alphabetic() {
1681 break;
1682 }
1683 }
1684 } else {
1685 out.push(c);
1686 }
1687 }
1688 out
1689}
1690
1691#[cfg(test)]
1692mod tests {
1693 use super::*;
1694 use crate::util::test_env::env_lock;
1695
1696 fn plain_segs(words: &[&str]) -> Vec<Seg> {
1699 words.iter().map(|w| Seg::plain(w)).collect()
1700 }
1701
1702 #[test]
1703 fn wrap_segments_packs_words_within_width() {
1704 let segs = plain_segs(&["alpha", "beta", "gamma", "delta"]);
1705 let lines = wrap_segments(&segs, 12);
1707 assert_eq!(lines, vec!["alpha beta", "gamma delta"]);
1708 for line in &lines {
1709 assert!(line.chars().count() <= 12);
1710 }
1711 }
1712
1713 #[test]
1714 fn wrap_segments_oversized_segment_gets_its_own_line() {
1715 let segs = plain_segs(&["short", "supercalifragilistic", "tail"]);
1716 let lines = wrap_segments(&segs, 10);
1717 assert_eq!(lines, vec!["short", "supercalifragilistic", "tail"]);
1719 }
1720
1721 #[test]
1722 fn format_row_aligns_continuation_under_description_column() {
1723 let segs = plain_segs(&["aaaa", "bbbb", "cccc", "dddd", "eeee"]);
1725 let rows = format_row("--x", 3, 8, &segs, 28);
1726 assert_eq!(rows[0], "--x aaaa bbbb cccc dddd"); assert_eq!(rows[1], format!("{}eeee", " ".repeat(8)));
1729 for row in &rows[1..] {
1731 assert!(row.starts_with(&" ".repeat(8)));
1732 assert!(!row.starts_with(&" ".repeat(9)));
1733 }
1734 }
1735
1736 #[test]
1737 fn format_row_overflowing_label_drops_description_below() {
1738 let segs = plain_segs(&["desc"]);
1739 let rows = format_row("--a-very-long-flag", 18, 8, &segs, 40);
1741 assert_eq!(rows[0], "--a-very-long-flag");
1742 assert_eq!(rows[1], format!("{}desc", " ".repeat(8)));
1743 }
1744
1745 #[test]
1746 fn help_width_honors_columns_env_within_clamp() {
1747 let _lock = env_lock();
1748 let prev = std::env::var("COLUMNS").ok();
1749 unsafe { std::env::set_var("COLUMNS", "90") };
1751 assert_eq!(help_width(), 90);
1752 unsafe { std::env::set_var("COLUMNS", "9999") };
1753 assert_eq!(help_width(), 120); unsafe { std::env::set_var("COLUMNS", "10") };
1755 assert_eq!(help_width(), 60); match prev {
1757 Some(v) => unsafe { std::env::set_var("COLUMNS", v) },
1758 None => unsafe { std::env::remove_var("COLUMNS") },
1759 }
1760 }
1761
1762 fn unique_tmp_dir(tag: &str) -> std::path::PathBuf {
1763 static SEQ: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1764 let d = std::env::temp_dir().join(format!(
1765 "fdl-run-test-{tag}-{}-{}",
1766 std::process::id(),
1767 SEQ.fetch_add(1, std::sync::atomic::Ordering::Relaxed),
1768 ));
1769 std::fs::create_dir_all(&d).unwrap();
1770 d
1771 }
1772
1773 #[test]
1774 fn overlay_libtorch_finds_fdl_yaml_spelling() {
1775 let _guard = env_lock();
1779 let dir = unique_tmp_dir("yaml-spelling");
1780 std::fs::write(dir.join("fdl.yaml"), "description: base\n").unwrap();
1781 std::fs::write(
1782 dir.join("fdl.testenv.yaml"),
1783 "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",
1784 )
1785 .unwrap();
1786 unsafe { std::env::set_var("FDL_ENV", "testenv") };
1787 let resolved = resolve_libtorch_from_overlay(&dir);
1790 unsafe { std::env::remove_var("FDL_ENV") };
1791 std::fs::remove_dir_all(&dir).ok();
1792 assert!(matches!(resolved, Ok(None)), "{resolved:?}");
1793 }
1794
1795 #[test]
1796 fn overlay_libtorch_load_failure_is_loud() {
1797 let _guard = env_lock();
1800 let dir = unique_tmp_dir("broken-overlay");
1801 std::fs::write(dir.join("fdl.yml"), "description: base\n").unwrap();
1802 unsafe { std::env::set_var("FDL_ENV", "missing-env") };
1804 let resolved = resolve_libtorch_from_overlay(&dir);
1805 unsafe { std::env::remove_var("FDL_ENV") };
1806 std::fs::remove_dir_all(&dir).ok();
1807 let err = match resolved {
1808 Ok(v) => panic!("expected Err, got Ok({v:?})"),
1809 Err(e) => e,
1810 };
1811 assert!(err.contains("missing-env"), "{err}");
1812 }
1813
1814 #[test]
1815 fn posix_quote_passes_safe_strings_through() {
1816 assert_eq!(posix_quote("hello"), "hello");
1817 assert_eq!(posix_quote("-p"), "-p");
1818 assert_eq!(posix_quote("flodl-hf"), "flodl-hf");
1819 assert_eq!(posix_quote("a/b.c"), "a/b.c");
1820 assert_eq!(posix_quote("KEY=val"), "KEY=val");
1821 }
1822
1823 #[test]
1824 fn posix_quote_wraps_unsafe_strings() {
1825 assert_eq!(posix_quote(""), "''");
1826 assert_eq!(posix_quote("foo bar"), "'foo bar'");
1827 assert_eq!(posix_quote("a$b"), "'a$b'");
1828 assert_eq!(posix_quote("a\"b"), "'a\"b'");
1829 }
1830
1831 #[test]
1832 fn posix_quote_escapes_embedded_single_quotes() {
1833 assert_eq!(posix_quote("it's"), "'it'\\''s'");
1834 assert_eq!(posix_quote("'"), "''\\'''");
1835 }
1836
1837 #[test]
1838 fn posix_quote_round_trips_shell_join_output() {
1839 let args: Vec<String> = ["--tag", "$HOME", "a\"b"]
1847 .iter()
1848 .map(|s| s.to_string())
1849 .collect();
1850 let inner = format!("cd /workspace/bench && train {}", shell_join(&args));
1851 assert_eq!(
1852 posix_quote(&inner),
1853 "'cd /workspace/bench && train --tag '\\''$HOME'\\'' '\\''a\"b'\\'''"
1854 );
1855 }
1856
1857 #[test]
1858 fn shell_join_quotes_shell_metacharacters() {
1859 let args: Vec<String> = ["--model", "mlp", "--tag", "$HOME", "*.py", "a;b"]
1863 .iter()
1864 .map(|s| s.to_string())
1865 .collect();
1866 assert_eq!(shell_join(&args), "--model mlp --tag '$HOME' '*.py' 'a;b'");
1867 }
1868
1869 #[test]
1870 fn compose_run_command_no_extras_passes_run_through() {
1871 assert_eq!(compose_run_command("echo hello", &[], None), "echo hello");
1872 }
1873
1874 #[test]
1875 fn compose_run_command_inserts_user_args_between_run_and_append() {
1876 let user = vec!["-p".to_string(), "flodl-hf".to_string()];
1877 let out = compose_run_command("cargo test live", &user, Some("-- --nocapture --ignored"));
1878 assert_eq!(out, "cargo test live -p flodl-hf -- --nocapture --ignored");
1879 }
1880
1881 #[test]
1882 fn compose_run_command_quotes_user_args_with_spaces() {
1883 let user = vec!["--name".to_string(), "with space".to_string()];
1884 let out = compose_run_command("cmd", &user, None);
1885 assert_eq!(out, "cmd --name 'with space'");
1886 }
1887
1888 #[test]
1889 fn compose_run_command_omits_empty_append() {
1890 let out = compose_run_command("cmd", &["arg".to_string()], Some(""));
1891 assert_eq!(out, "cmd arg");
1892 let out2 = compose_run_command("cmd", &["arg".to_string()], Some(" "));
1893 assert_eq!(out2, "cmd arg");
1894 }
1895
1896 #[test]
1897 fn compose_run_command_user_double_dash_threads_runner_args() {
1898 let user = vec![
1899 "-p".to_string(),
1900 "foo".to_string(),
1901 "--".to_string(),
1902 "--ignored".to_string(),
1903 ];
1904 let out = compose_run_command("cargo test", &user, Some("-- --nocapture"));
1905 assert_eq!(out, "cargo test -p foo -- --nocapture --ignored");
1906 }
1907
1908 #[test]
1909 fn compose_run_command_user_double_dash_without_append() {
1910 let user = vec![
1911 "-p".to_string(),
1912 "foo".to_string(),
1913 "--".to_string(),
1914 "--ignored".to_string(),
1915 ];
1916 let out = compose_run_command("cargo test", &user, None);
1917 assert_eq!(out, "cargo test -p foo -- --ignored");
1918 }
1919
1920 #[test]
1921 fn compose_run_command_append_with_pre_and_post_halves() {
1922 let out = compose_run_command("cmd", &[], Some("--foo -- --bar"));
1923 assert_eq!(out, "cmd --foo -- --bar");
1924 }
1925
1926 #[test]
1927 fn compose_run_command_append_pre_only_no_separator() {
1928 let user = vec!["--ansi".to_string()];
1931 let out = compose_run_command("cmd", &user, Some("--no-ansi"));
1932 assert_eq!(out, "cmd --no-ansi --ansi");
1933 }
1934
1935 #[test]
1936 fn compose_run_command_user_only_double_dash_emits_separator() {
1937 let user = vec!["--".to_string(), "--list".to_string()];
1938 let out = compose_run_command("cargo test", &user, None);
1939 assert_eq!(out, "cargo test -- --list");
1940 }
1941
1942 #[test]
1943 fn compose_run_command_append_full_split_with_user_both_sides() {
1944 let user = vec![
1945 "-p".to_string(),
1946 "foo".to_string(),
1947 "--".to_string(),
1948 "--ignored".to_string(),
1949 ];
1950 let out = compose_run_command(
1951 "cargo test",
1952 &user,
1953 Some("--release -- --nocapture"),
1954 );
1955 assert_eq!(
1956 out,
1957 "cargo test --release -p foo -- --nocapture --ignored"
1958 );
1959 }
1960
1961 #[test]
1962 fn split_append_dashdash_handles_edges() {
1963 assert_eq!(
1964 split_append_dashdash("-- --nocapture"),
1965 (String::new(), "--nocapture".to_string())
1966 );
1967 assert_eq!(
1968 split_append_dashdash("--foo -- --bar"),
1969 ("--foo".to_string(), "--bar".to_string())
1970 );
1971 assert_eq!(
1972 split_append_dashdash("--foo --"),
1973 ("--foo".to_string(), String::new())
1974 );
1975 assert_eq!(
1976 split_append_dashdash("--"),
1977 (String::new(), String::new())
1978 );
1979 assert_eq!(
1980 split_append_dashdash("--foo"),
1981 ("--foo".to_string(), String::new())
1982 );
1983 assert_eq!(
1984 split_append_dashdash(""),
1985 (String::new(), String::new())
1986 );
1987 }
1988
1989 use std::sync::atomic::{AtomicU64, Ordering};
1998 use std::time::{SystemTime, UNIX_EPOCH};
1999
2000 static SCRATCH_SEQ: AtomicU64 = AtomicU64::new(0);
2001
2002 struct Scratch(std::path::PathBuf);
2003 impl Scratch {
2004 fn new() -> Self {
2005 let nanos = SystemTime::now().duration_since(UNIX_EPOCH)
2006 .map(|d| d.as_nanos()).unwrap_or(0);
2007 let seq = SCRATCH_SEQ.fetch_add(1, Ordering::Relaxed);
2008 let dir = std::env::temp_dir()
2009 .join(format!("fdl-resolve-libtorch-{}-{}", nanos, seq));
2010 std::fs::create_dir_all(&dir).unwrap();
2011 Self(dir)
2012 }
2013 fn path(&self) -> &std::path::Path { &self.0 }
2014 }
2015 impl Drop for Scratch {
2016 fn drop(&mut self) {
2017 let _ = std::fs::remove_dir_all(&self.0);
2018 }
2019 }
2020
2021 fn populate_lt_root(root: &std::path::Path) {
2022 for (sub, torch, archs) in [
2023 ("precompiled/v1", "1.0", "0.0"),
2024 ("builds/v2", "2.0", "1.0"),
2025 ] {
2026 let d = root.join(sub);
2027 std::fs::create_dir_all(d.join("lib")).unwrap();
2028 std::fs::write(
2029 d.join(".arch"),
2030 format!("torch={torch}\ncuda=1.0\narchs={archs}\nsource=test\n"),
2031 ).unwrap();
2032 }
2033 }
2034
2035 #[test]
2036 fn resolve_libtorch_at_pointer_file() {
2037 let s = Scratch::new();
2038 let lt = s.path().join("libtorch");
2039 populate_lt_root(<);
2040 let pointer = lt.join(".active.alt");
2041 std::fs::write(&pointer, "precompiled/v1\n").unwrap();
2042
2043 let (info, host_path) = resolve_libtorch_at(&pointer)
2044 .expect("pointer file resolves");
2045 assert_eq!(info.path, "precompiled/v1");
2046 assert_eq!(info.torch_version.as_deref(), Some("1.0"));
2047 assert_eq!(host_path, lt.join("precompiled/v1").display().to_string());
2048 }
2049
2050 #[test]
2051 fn resolve_libtorch_at_libtorch_root_dir() {
2052 let s = Scratch::new();
2053 let lt = s.path().join("libtorch");
2054 populate_lt_root(<);
2055 std::fs::write(lt.join(".active"), "builds/v2\n").unwrap();
2056
2057 let (info, host_path) = resolve_libtorch_at(<)
2058 .expect("libtorch-root dir resolves");
2059 assert_eq!(info.path, "builds/v2");
2060 assert_eq!(info.torch_version.as_deref(), Some("2.0"));
2061 assert_eq!(host_path, lt.join("builds/v2").display().to_string());
2062 }
2063
2064 #[test]
2065 fn resolve_libtorch_at_direct_variant_dir() {
2066 let s = Scratch::new();
2067 let variant = s.path().join("standalone-libtorch");
2068 std::fs::create_dir_all(variant.join("lib")).unwrap();
2069 std::fs::write(
2070 variant.join(".arch"),
2071 "torch=3.0\ncuda=2.0\narchs=1.0\nsource=test\n",
2072 ).unwrap();
2073
2074 let (info, host_path) = resolve_libtorch_at(&variant)
2075 .expect("direct variant dir resolves");
2076 assert_eq!(info.path, variant.display().to_string());
2077 assert_eq!(info.torch_version.as_deref(), Some("3.0"));
2078 assert_eq!(host_path, variant.display().to_string());
2079 }
2080
2081 #[test]
2082 fn resolve_libtorch_at_bogus_path_returns_none() {
2083 let s = Scratch::new();
2084 let bogus = s.path().join("no-lib-no-active-no-pointer");
2085 std::fs::create_dir_all(&bogus).unwrap();
2086 assert!(resolve_libtorch_at(&bogus).is_none(),
2087 "dir without lib/, .active, or pointer-shape filename → None");
2088 }
2089}