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 " {} Scope visible GPUs, e.g. 0,1 or all (any position)",
1243 style::green(&format!("{:<18}", "--gpus <spec>"))
1244 );
1245 eprintln!(
1246 " {} Verbose output",
1247 style::green(&format!("{:<18}", "-v"))
1248 );
1249 eprintln!(
1250 " {} Debug output",
1251 style::green(&format!("{:<18}", "-vv"))
1252 );
1253 eprintln!(
1254 " {} Trace output (maximum detail)",
1255 style::green(&format!("{:<18}", "-vvv"))
1256 );
1257 eprintln!(
1258 " {} Suppress non-error output",
1259 style::green(&format!("{:<18}", "-q, --quiet"))
1260 );
1261 eprintln!(
1262 " {} Force ANSI color (bypass TTY / NO_COLOR detection)",
1263 style::green(&format!("{:<18}", "--ansi"))
1264 );
1265 eprintln!(
1266 " {} Disable ANSI color output",
1267 style::green(&format!("{:<18}", "--no-ansi"))
1268 );
1269 eprintln!(
1270 " {} Drop a run command's `append:` suffix",
1271 style::green(&format!("{:<18}", "--no-append"))
1272 );
1273 eprintln!(
1274 " {} Skip the cluster pre-flight build",
1275 style::green(&format!("{:<18}", "--no-prebuild"))
1276 );
1277
1278 eprintln!();
1280 eprintln!("{}:", style::yellow("Built-in"));
1281 for (name, desc) in &visible_builtins {
1282 eprintln!(" {} {desc}", style::green(&format!("{:<18}", name)));
1283 }
1284
1285 if !project.commands.is_empty() {
1292 eprintln!();
1293 eprintln!("{}:", style::yellow("Commands"));
1294 for (name, spec) in &project.commands {
1295 let desc: String = match spec.description.clone() {
1296 Some(d) => d,
1297 None => {
1298 let is_path_kind = spec.run.is_none();
1302 if is_path_kind {
1303 let child_dir = spec.resolve_path(name, project_root);
1304 config::load_command_with_env(&child_dir, active_env)
1305 .ok()
1306 .and_then(|c| c.description)
1307 .unwrap_or_else(|| "(sub-command)".into())
1308 } else {
1309 spec.run
1310 .as_deref()
1311 .unwrap_or("(command)")
1312 .to_string()
1313 }
1314 }
1315 };
1316 eprintln!(" {} {desc}", style::green(&format!("{:<18}", name)));
1317 }
1318 }
1319
1320 if let Some(base_config) = config::find_config(project_root) {
1322 let envs = crate::overlay::list_envs(&base_config);
1323 if !envs.is_empty() {
1324 eprintln!();
1325 eprintln!("{}:", style::yellow("Environments"));
1326 for e in &envs {
1327 let active_marker = if Some(e.as_str()) == active_env {
1328 style::green(" (active)")
1329 } else {
1330 String::new()
1331 };
1332 eprintln!(
1333 " {} Overlay from fdl.{}.yml{active_marker}",
1334 style::green(&format!("{:<18}", format!("@{e}"))),
1335 e
1336 );
1337 }
1338 eprintln!();
1339 eprintln!(
1340 "Use {} to run a command with an environment overlay.",
1341 style::dim("fdl @<env> <command>")
1342 );
1343 }
1344 }
1345
1346 eprintln!();
1347 eprintln!(
1348 "Use {} for more information on a command.",
1349 style::dim("fdl <command> -h")
1350 );
1351}
1352
1353fn build_usage_tail(
1362 schema: Option<&Schema>,
1363 has_presets: bool,
1364 has_sub_commands: bool,
1365 preset_slot: &str,
1366) -> String {
1367 let mut parts = String::new();
1368 let slot = match (has_presets, has_sub_commands) {
1369 (true, false) => Some(format!("[<{preset_slot}>]")),
1370 (false, true) => Some("[<command>]".to_string()),
1371 (true, true) => Some(format!("[<{preset_slot}>|<command>]")),
1372 (false, false) => None,
1373 };
1374 if let Some(s) = slot {
1375 parts.push(' ');
1376 parts.push_str(&style::dim(&s));
1377 }
1378 if let Some(s) = schema {
1379 for a in &s.args {
1380 parts.push(' ');
1381 parts.push_str(&format_arg_usage(a));
1382 }
1383 }
1384 parts.push(' ');
1385 parts.push_str(&style::dim("[options]"));
1386 parts
1387}
1388
1389type CommandGroup = Vec<(String, crate::config::CommandSpec)>;
1390
1391fn split_commands_by_kind(
1396 commands: &BTreeMap<String, crate::config::CommandSpec>,
1397) -> (CommandGroup, CommandGroup) {
1398 use crate::config::CommandKind;
1399 let mut presets = Vec::new();
1400 let mut sub_cmds = Vec::new();
1401 for (k, v) in commands {
1402 match v.kind() {
1403 Ok(CommandKind::Preset) => presets.push((k.clone(), v.clone())),
1404 _ => sub_cmds.push((k.clone(), v.clone())),
1405 }
1406 }
1407 (presets, sub_cmds)
1408}
1409
1410fn format_arg_usage(a: &ArgSpec) -> String {
1411 let suffix = if a.variadic { "..." } else { "" };
1412 let core = format!("<{}>{suffix}", a.name);
1413 if a.required && a.default.is_none() {
1414 style::green(&core)
1415 } else {
1416 style::dim(&format!("[{core}]"))
1417 }
1418}
1419
1420const ARG_COL: usize = 22;
1424const OPT_COL: usize = 30;
1427
1428fn format_arg(a: &ArgSpec, avail_width: usize) -> Vec<String> {
1429 let left = format_arg_usage(a);
1430 let visible = visible_width(&left);
1431 let segs = desc_segments(a.description.as_deref(), &a.default, &a.choices, &a.ty);
1432 format_row(&left, visible, ARG_COL, &segs, avail_width)
1433}
1434
1435fn format_option(long: &str, spec: &OptionSpec, avail_width: usize) -> Vec<String> {
1440 let flag = match &spec.short {
1441 Some(s) => format!("-{s}, --{long}"),
1442 None => format!(" --{long}"),
1443 };
1444 let placeholder = option_placeholder(&spec.ty);
1445 let (left, visible) = if placeholder.is_empty() {
1446 (style::green(&flag), flag.chars().count())
1447 } else {
1448 (
1449 style::green(&format!("{flag} {placeholder}")),
1450 flag.chars().count() + 1 + placeholder.chars().count(),
1451 )
1452 };
1453 let segs = desc_segments(spec.description.as_deref(), &spec.default, &spec.choices, &spec.ty);
1454 let mut out = format_row(&left, visible, OPT_COL, &segs, avail_width);
1455 if let Some(env) = &spec.env {
1456 out.push(format!(
1457 "{}{}",
1458 " ".repeat(OPT_COL),
1459 style::dim(&format!("[env: {env}]"))
1460 ));
1461 }
1462 out
1463}
1464
1465struct Seg {
1469 text: String,
1470 styled: String,
1471}
1472
1473impl Seg {
1474 fn plain(s: &str) -> Seg {
1475 Seg { text: s.to_string(), styled: s.to_string() }
1476 }
1477 fn dim(s: &str) -> Seg {
1478 Seg { text: s.to_string(), styled: style::dim(s) }
1479 }
1480}
1481
1482fn desc_segments(
1487 description: Option<&str>,
1488 default: &Option<serde_json::Value>,
1489 choices: &Option<Vec<serde_json::Value>>,
1490 ty: &str,
1491) -> Vec<Seg> {
1492 let mut segs: Vec<Seg> = description
1493 .unwrap_or("-")
1494 .split_whitespace()
1495 .map(Seg::plain)
1496 .collect();
1497 if let Some(d) = default {
1498 let is_empty_list = matches!(d, serde_json::Value::Array(a) if a.is_empty());
1500 let is_false = matches!(d, serde_json::Value::Bool(false));
1501 if !d.is_null() && !is_false && !is_empty_list {
1502 segs.push(Seg::dim(&format!("[default: {}]", format_value(d))));
1503 }
1504 }
1505 if let Some(choices) = choices {
1506 if !choices.is_empty() {
1507 let list = choices.iter().map(format_value).collect::<Vec<_>>().join(", ");
1508 segs.push(Seg::dim(&format!("[possible: {list}]")));
1509 }
1510 }
1511 if ty.starts_with("list[") {
1513 segs.push(Seg::dim("(repeat or comma-separate)"));
1514 }
1515 segs
1516}
1517
1518fn wrap_segments(segs: &[Seg], width: usize) -> Vec<String> {
1523 let width = width.max(1);
1524 let mut lines: Vec<String> = Vec::new();
1525 let mut cur = String::new();
1526 let mut cur_w = 0usize;
1527 for seg in segs {
1528 let w = seg.text.chars().count();
1529 if cur_w == 0 {
1530 cur.push_str(&seg.styled);
1531 cur_w = w;
1532 } else if cur_w + 1 + w <= width {
1533 cur.push(' ');
1534 cur.push_str(&seg.styled);
1535 cur_w += 1 + w;
1536 } else {
1537 lines.push(std::mem::take(&mut cur));
1538 cur.push_str(&seg.styled);
1539 cur_w = w;
1540 }
1541 }
1542 if !cur.is_empty() {
1543 lines.push(cur);
1544 }
1545 if lines.is_empty() {
1546 lines.push(String::new());
1547 }
1548 lines
1549}
1550
1551fn format_row(
1558 label: &str,
1559 label_visible: usize,
1560 desc_col: usize,
1561 segs: &[Seg],
1562 avail_width: usize,
1563) -> Vec<String> {
1564 const MIN_DESC: usize = 20;
1566 let desc_width = avail_width.saturating_sub(desc_col).max(MIN_DESC);
1567 let desc_lines = wrap_segments(segs, desc_width);
1568 let pad = " ".repeat(desc_col);
1569 let mut out: Vec<String> = Vec::with_capacity(desc_lines.len() + 1);
1570 if label_visible < desc_col {
1571 let gap = " ".repeat(desc_col - label_visible);
1572 out.push(format!("{label}{gap}{}", desc_lines[0]));
1573 } else {
1574 out.push(label.to_string());
1576 out.push(format!("{pad}{}", desc_lines[0]));
1577 }
1578 for line in &desc_lines[1..] {
1579 out.push(format!("{pad}{line}"));
1580 }
1581 for line in &mut out {
1583 while line.ends_with(' ') {
1584 line.pop();
1585 }
1586 }
1587 out
1588}
1589
1590fn option_placeholder(ty: &str) -> &'static str {
1591 match ty {
1592 "bool" => "",
1593 "int" => "<N>",
1594 "float" => "<F>",
1595 "path" => "<PATH>",
1596 "list[path]" => "<PATH>...",
1597 t if t.starts_with("list[") => "<VALUE>...",
1598 _ => "<VALUE>",
1599 }
1600}
1601
1602fn format_value(v: &serde_json::Value) -> String {
1603 match v {
1604 serde_json::Value::String(s) => s.clone(),
1605 other => other.to_string(),
1606 }
1607}
1608
1609fn visible_width(s: &str) -> usize {
1612 strip_ansi(s).chars().count()
1615}
1616
1617fn help_width() -> usize {
1623 const DEFAULT: usize = 100;
1624 const MIN: usize = 60;
1625 const MAX: usize = 120;
1626 let raw = std::env::var("COLUMNS")
1627 .ok()
1628 .and_then(|s| s.trim().parse::<usize>().ok())
1629 .filter(|&c| c > 0)
1630 .or_else(term_cols)
1631 .unwrap_or(DEFAULT);
1632 raw.clamp(MIN, MAX)
1633}
1634
1635#[cfg(unix)]
1641fn term_cols() -> Option<usize> {
1642 use std::io::IsTerminal;
1643 use std::os::unix::io::AsRawFd;
1644
1645 let stderr = std::io::stderr();
1646 if !stderr.is_terminal() {
1647 return None;
1648 }
1649 #[repr(C)]
1650 struct Winsize {
1651 row: u16,
1652 col: u16,
1653 xpixel: u16,
1654 ypixel: u16,
1655 }
1656 #[cfg(any(target_os = "linux", target_os = "android"))]
1659 const TIOCGWINSZ: std::os::raw::c_ulong = 0x5413;
1660 #[cfg(not(any(target_os = "linux", target_os = "android")))]
1661 const TIOCGWINSZ: std::os::raw::c_ulong = 0x4008_7468;
1662 unsafe extern "C" {
1663 fn ioctl(fd: std::os::raw::c_int, request: std::os::raw::c_ulong, ...) -> std::os::raw::c_int;
1664 }
1665 let mut ws = Winsize { row: 0, col: 0, xpixel: 0, ypixel: 0 };
1666 let rc = unsafe { ioctl(stderr.as_raw_fd(), TIOCGWINSZ, &mut ws as *mut Winsize) };
1669 (rc == 0 && ws.col > 0).then_some(ws.col as usize)
1670}
1671
1672#[cfg(not(unix))]
1673fn term_cols() -> Option<usize> {
1674 None
1675}
1676
1677fn strip_ansi(s: &str) -> String {
1678 let mut out = String::with_capacity(s.len());
1679 let mut chars = s.chars().peekable();
1680 while let Some(c) = chars.next() {
1681 if c == '\x1b' && chars.peek() == Some(&'[') {
1682 chars.next();
1683 for c in chars.by_ref() {
1684 if c.is_ascii_alphabetic() {
1685 break;
1686 }
1687 }
1688 } else {
1689 out.push(c);
1690 }
1691 }
1692 out
1693}
1694
1695#[cfg(test)]
1696mod tests {
1697 use super::*;
1698 use crate::util::test_env::env_lock;
1699
1700 fn plain_segs(words: &[&str]) -> Vec<Seg> {
1703 words.iter().map(|w| Seg::plain(w)).collect()
1704 }
1705
1706 #[test]
1707 fn wrap_segments_packs_words_within_width() {
1708 let segs = plain_segs(&["alpha", "beta", "gamma", "delta"]);
1709 let lines = wrap_segments(&segs, 12);
1711 assert_eq!(lines, vec!["alpha beta", "gamma delta"]);
1712 for line in &lines {
1713 assert!(line.chars().count() <= 12);
1714 }
1715 }
1716
1717 #[test]
1718 fn wrap_segments_oversized_segment_gets_its_own_line() {
1719 let segs = plain_segs(&["short", "supercalifragilistic", "tail"]);
1720 let lines = wrap_segments(&segs, 10);
1721 assert_eq!(lines, vec!["short", "supercalifragilistic", "tail"]);
1723 }
1724
1725 #[test]
1726 fn format_row_aligns_continuation_under_description_column() {
1727 let segs = plain_segs(&["aaaa", "bbbb", "cccc", "dddd", "eeee"]);
1729 let rows = format_row("--x", 3, 8, &segs, 28);
1730 assert_eq!(rows[0], "--x aaaa bbbb cccc dddd"); assert_eq!(rows[1], format!("{}eeee", " ".repeat(8)));
1733 for row in &rows[1..] {
1735 assert!(row.starts_with(&" ".repeat(8)));
1736 assert!(!row.starts_with(&" ".repeat(9)));
1737 }
1738 }
1739
1740 #[test]
1741 fn format_row_overflowing_label_drops_description_below() {
1742 let segs = plain_segs(&["desc"]);
1743 let rows = format_row("--a-very-long-flag", 18, 8, &segs, 40);
1745 assert_eq!(rows[0], "--a-very-long-flag");
1746 assert_eq!(rows[1], format!("{}desc", " ".repeat(8)));
1747 }
1748
1749 #[test]
1750 fn help_width_honors_columns_env_within_clamp() {
1751 let _lock = env_lock();
1752 let prev = std::env::var("COLUMNS").ok();
1753 unsafe { std::env::set_var("COLUMNS", "90") };
1755 assert_eq!(help_width(), 90);
1756 unsafe { std::env::set_var("COLUMNS", "9999") };
1757 assert_eq!(help_width(), 120); unsafe { std::env::set_var("COLUMNS", "10") };
1759 assert_eq!(help_width(), 60); match prev {
1761 Some(v) => unsafe { std::env::set_var("COLUMNS", v) },
1762 None => unsafe { std::env::remove_var("COLUMNS") },
1763 }
1764 }
1765
1766 fn unique_tmp_dir(tag: &str) -> std::path::PathBuf {
1767 static SEQ: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1768 let d = std::env::temp_dir().join(format!(
1769 "fdl-run-test-{tag}-{}-{}",
1770 std::process::id(),
1771 SEQ.fetch_add(1, std::sync::atomic::Ordering::Relaxed),
1772 ));
1773 std::fs::create_dir_all(&d).unwrap();
1774 d
1775 }
1776
1777 #[test]
1778 fn overlay_libtorch_finds_fdl_yaml_spelling() {
1779 let _guard = env_lock();
1783 let dir = unique_tmp_dir("yaml-spelling");
1784 std::fs::write(dir.join("fdl.yaml"), "description: base\n").unwrap();
1785 std::fs::write(
1786 dir.join("fdl.testenv.yaml"),
1787 "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",
1788 )
1789 .unwrap();
1790 unsafe { std::env::set_var("FDL_ENV", "testenv") };
1791 let resolved = resolve_libtorch_from_overlay(&dir);
1794 unsafe { std::env::remove_var("FDL_ENV") };
1795 std::fs::remove_dir_all(&dir).ok();
1796 assert!(matches!(resolved, Ok(None)), "{resolved:?}");
1797 }
1798
1799 #[test]
1800 fn overlay_libtorch_load_failure_is_loud() {
1801 let _guard = env_lock();
1804 let dir = unique_tmp_dir("broken-overlay");
1805 std::fs::write(dir.join("fdl.yml"), "description: base\n").unwrap();
1806 unsafe { std::env::set_var("FDL_ENV", "missing-env") };
1808 let resolved = resolve_libtorch_from_overlay(&dir);
1809 unsafe { std::env::remove_var("FDL_ENV") };
1810 std::fs::remove_dir_all(&dir).ok();
1811 let err = match resolved {
1812 Ok(v) => panic!("expected Err, got Ok({v:?})"),
1813 Err(e) => e,
1814 };
1815 assert!(err.contains("missing-env"), "{err}");
1816 }
1817
1818 #[test]
1819 fn posix_quote_passes_safe_strings_through() {
1820 assert_eq!(posix_quote("hello"), "hello");
1821 assert_eq!(posix_quote("-p"), "-p");
1822 assert_eq!(posix_quote("flodl-hf"), "flodl-hf");
1823 assert_eq!(posix_quote("a/b.c"), "a/b.c");
1824 assert_eq!(posix_quote("KEY=val"), "KEY=val");
1825 }
1826
1827 #[test]
1828 fn posix_quote_wraps_unsafe_strings() {
1829 assert_eq!(posix_quote(""), "''");
1830 assert_eq!(posix_quote("foo bar"), "'foo bar'");
1831 assert_eq!(posix_quote("a$b"), "'a$b'");
1832 assert_eq!(posix_quote("a\"b"), "'a\"b'");
1833 }
1834
1835 #[test]
1836 fn posix_quote_escapes_embedded_single_quotes() {
1837 assert_eq!(posix_quote("it's"), "'it'\\''s'");
1838 assert_eq!(posix_quote("'"), "''\\'''");
1839 }
1840
1841 #[test]
1842 fn posix_quote_round_trips_shell_join_output() {
1843 let args: Vec<String> = ["--tag", "$HOME", "a\"b"]
1851 .iter()
1852 .map(|s| s.to_string())
1853 .collect();
1854 let inner = format!("cd /workspace/bench && train {}", shell_join(&args));
1855 assert_eq!(
1856 posix_quote(&inner),
1857 "'cd /workspace/bench && train --tag '\\''$HOME'\\'' '\\''a\"b'\\'''"
1858 );
1859 }
1860
1861 #[test]
1862 fn shell_join_quotes_shell_metacharacters() {
1863 let args: Vec<String> = ["--model", "mlp", "--tag", "$HOME", "*.py", "a;b"]
1867 .iter()
1868 .map(|s| s.to_string())
1869 .collect();
1870 assert_eq!(shell_join(&args), "--model mlp --tag '$HOME' '*.py' 'a;b'");
1871 }
1872
1873 #[test]
1874 fn compose_run_command_no_extras_passes_run_through() {
1875 assert_eq!(compose_run_command("echo hello", &[], None), "echo hello");
1876 }
1877
1878 #[test]
1879 fn compose_run_command_inserts_user_args_between_run_and_append() {
1880 let user = vec!["-p".to_string(), "flodl-hf".to_string()];
1881 let out = compose_run_command("cargo test live", &user, Some("-- --nocapture --ignored"));
1882 assert_eq!(out, "cargo test live -p flodl-hf -- --nocapture --ignored");
1883 }
1884
1885 #[test]
1886 fn compose_run_command_quotes_user_args_with_spaces() {
1887 let user = vec!["--name".to_string(), "with space".to_string()];
1888 let out = compose_run_command("cmd", &user, None);
1889 assert_eq!(out, "cmd --name 'with space'");
1890 }
1891
1892 #[test]
1893 fn compose_run_command_omits_empty_append() {
1894 let out = compose_run_command("cmd", &["arg".to_string()], Some(""));
1895 assert_eq!(out, "cmd arg");
1896 let out2 = compose_run_command("cmd", &["arg".to_string()], Some(" "));
1897 assert_eq!(out2, "cmd arg");
1898 }
1899
1900 #[test]
1901 fn compose_run_command_user_double_dash_threads_runner_args() {
1902 let user = vec![
1903 "-p".to_string(),
1904 "foo".to_string(),
1905 "--".to_string(),
1906 "--ignored".to_string(),
1907 ];
1908 let out = compose_run_command("cargo test", &user, Some("-- --nocapture"));
1909 assert_eq!(out, "cargo test -p foo -- --nocapture --ignored");
1910 }
1911
1912 #[test]
1913 fn compose_run_command_user_double_dash_without_append() {
1914 let user = vec![
1915 "-p".to_string(),
1916 "foo".to_string(),
1917 "--".to_string(),
1918 "--ignored".to_string(),
1919 ];
1920 let out = compose_run_command("cargo test", &user, None);
1921 assert_eq!(out, "cargo test -p foo -- --ignored");
1922 }
1923
1924 #[test]
1925 fn compose_run_command_append_with_pre_and_post_halves() {
1926 let out = compose_run_command("cmd", &[], Some("--foo -- --bar"));
1927 assert_eq!(out, "cmd --foo -- --bar");
1928 }
1929
1930 #[test]
1931 fn compose_run_command_append_pre_only_no_separator() {
1932 let user = vec!["--ansi".to_string()];
1935 let out = compose_run_command("cmd", &user, Some("--no-ansi"));
1936 assert_eq!(out, "cmd --no-ansi --ansi");
1937 }
1938
1939 #[test]
1940 fn compose_run_command_user_only_double_dash_emits_separator() {
1941 let user = vec!["--".to_string(), "--list".to_string()];
1942 let out = compose_run_command("cargo test", &user, None);
1943 assert_eq!(out, "cargo test -- --list");
1944 }
1945
1946 #[test]
1947 fn compose_run_command_append_full_split_with_user_both_sides() {
1948 let user = vec![
1949 "-p".to_string(),
1950 "foo".to_string(),
1951 "--".to_string(),
1952 "--ignored".to_string(),
1953 ];
1954 let out = compose_run_command(
1955 "cargo test",
1956 &user,
1957 Some("--release -- --nocapture"),
1958 );
1959 assert_eq!(
1960 out,
1961 "cargo test --release -p foo -- --nocapture --ignored"
1962 );
1963 }
1964
1965 #[test]
1966 fn split_append_dashdash_handles_edges() {
1967 assert_eq!(
1968 split_append_dashdash("-- --nocapture"),
1969 (String::new(), "--nocapture".to_string())
1970 );
1971 assert_eq!(
1972 split_append_dashdash("--foo -- --bar"),
1973 ("--foo".to_string(), "--bar".to_string())
1974 );
1975 assert_eq!(
1976 split_append_dashdash("--foo --"),
1977 ("--foo".to_string(), String::new())
1978 );
1979 assert_eq!(
1980 split_append_dashdash("--"),
1981 (String::new(), String::new())
1982 );
1983 assert_eq!(
1984 split_append_dashdash("--foo"),
1985 ("--foo".to_string(), String::new())
1986 );
1987 assert_eq!(
1988 split_append_dashdash(""),
1989 (String::new(), String::new())
1990 );
1991 }
1992
1993 use std::sync::atomic::{AtomicU64, Ordering};
2002 use std::time::{SystemTime, UNIX_EPOCH};
2003
2004 static SCRATCH_SEQ: AtomicU64 = AtomicU64::new(0);
2005
2006 struct Scratch(std::path::PathBuf);
2007 impl Scratch {
2008 fn new() -> Self {
2009 let nanos = SystemTime::now().duration_since(UNIX_EPOCH)
2010 .map(|d| d.as_nanos()).unwrap_or(0);
2011 let seq = SCRATCH_SEQ.fetch_add(1, Ordering::Relaxed);
2012 let dir = std::env::temp_dir()
2013 .join(format!("fdl-resolve-libtorch-{}-{}", nanos, seq));
2014 std::fs::create_dir_all(&dir).unwrap();
2015 Self(dir)
2016 }
2017 fn path(&self) -> &std::path::Path { &self.0 }
2018 }
2019 impl Drop for Scratch {
2020 fn drop(&mut self) {
2021 let _ = std::fs::remove_dir_all(&self.0);
2022 }
2023 }
2024
2025 fn populate_lt_root(root: &std::path::Path) {
2026 for (sub, torch, archs) in [
2027 ("precompiled/v1", "1.0", "0.0"),
2028 ("builds/v2", "2.0", "1.0"),
2029 ] {
2030 let d = root.join(sub);
2031 std::fs::create_dir_all(d.join("lib")).unwrap();
2032 std::fs::write(
2033 d.join(".arch"),
2034 format!("torch={torch}\ncuda=1.0\narchs={archs}\nsource=test\n"),
2035 ).unwrap();
2036 }
2037 }
2038
2039 #[test]
2040 fn resolve_libtorch_at_pointer_file() {
2041 let s = Scratch::new();
2042 let lt = s.path().join("libtorch");
2043 populate_lt_root(<);
2044 let pointer = lt.join(".active.alt");
2045 std::fs::write(&pointer, "precompiled/v1\n").unwrap();
2046
2047 let (info, host_path) = resolve_libtorch_at(&pointer)
2048 .expect("pointer file resolves");
2049 assert_eq!(info.path, "precompiled/v1");
2050 assert_eq!(info.torch_version.as_deref(), Some("1.0"));
2051 assert_eq!(host_path, lt.join("precompiled/v1").display().to_string());
2052 }
2053
2054 #[test]
2055 fn resolve_libtorch_at_libtorch_root_dir() {
2056 let s = Scratch::new();
2057 let lt = s.path().join("libtorch");
2058 populate_lt_root(<);
2059 std::fs::write(lt.join(".active"), "builds/v2\n").unwrap();
2060
2061 let (info, host_path) = resolve_libtorch_at(<)
2062 .expect("libtorch-root dir resolves");
2063 assert_eq!(info.path, "builds/v2");
2064 assert_eq!(info.torch_version.as_deref(), Some("2.0"));
2065 assert_eq!(host_path, lt.join("builds/v2").display().to_string());
2066 }
2067
2068 #[test]
2069 fn resolve_libtorch_at_direct_variant_dir() {
2070 let s = Scratch::new();
2071 let variant = s.path().join("standalone-libtorch");
2072 std::fs::create_dir_all(variant.join("lib")).unwrap();
2073 std::fs::write(
2074 variant.join(".arch"),
2075 "torch=3.0\ncuda=2.0\narchs=1.0\nsource=test\n",
2076 ).unwrap();
2077
2078 let (info, host_path) = resolve_libtorch_at(&variant)
2079 .expect("direct variant dir resolves");
2080 assert_eq!(info.path, variant.display().to_string());
2081 assert_eq!(info.torch_version.as_deref(), Some("3.0"));
2082 assert_eq!(host_path, variant.display().to_string());
2083 }
2084
2085 #[test]
2086 fn resolve_libtorch_at_bogus_path_returns_none() {
2087 let s = Scratch::new();
2088 let bogus = s.path().join("no-lib-no-active-no-pointer");
2089 std::fs::create_dir_all(&bogus).unwrap();
2090 assert!(resolve_libtorch_at(&bogus).is_none(),
2091 "dir without lib/, .active, or pointer-shape filename → None");
2092 }
2093}