eval_magic/pipeline/
slots.rs1use std::fs;
11use std::path::{Path, PathBuf};
12
13#[derive(Debug, Clone, PartialEq, Eq)]
15pub struct RunSlot {
16 pub run_index: Option<u32>,
19 pub dir: PathBuf,
20}
21
22pub fn run_key(eval_id: &str, condition: &str, run_index: Option<u32>) -> String {
25 match run_index {
26 Some(k) => format!("{eval_id}:{condition}:r{k}"),
27 None => format!("{eval_id}:{condition}"),
28 }
29}
30
31pub fn run_slots(cond_dir: &Path) -> Vec<RunSlot> {
36 let mut slots: Vec<RunSlot> = fs::read_dir(cond_dir)
37 .into_iter()
38 .flatten()
39 .flatten()
40 .filter(|e| e.path().is_dir())
41 .filter_map(|e| {
42 let name = e.file_name().to_string_lossy().into_owned();
43 let index: u32 = name.strip_prefix("run-")?.parse().ok()?;
44 Some(RunSlot {
45 run_index: Some(index),
46 dir: e.path(),
47 })
48 })
49 .collect();
50
51 if slots.is_empty() {
52 return vec![RunSlot {
53 run_index: None,
54 dir: cond_dir.to_path_buf(),
55 }];
56 }
57 slots.sort_by_key(|s| s.run_index);
58 slots
59}
60
61#[cfg(test)]
62mod tests {
63 use super::*;
64 use tempfile::TempDir;
65
66 fn legacy_slot(cond_dir: &Path) -> Vec<RunSlot> {
67 vec![RunSlot {
68 run_index: None,
69 dir: cond_dir.to_path_buf(),
70 }]
71 }
72
73 #[test]
74 fn flat_cond_dir_yields_single_legacy_slot() {
75 let root = TempDir::new().unwrap();
76 let cond_dir = root.path().join("with_skill");
77 fs::create_dir_all(&cond_dir).unwrap();
78 fs::write(cond_dir.join("run.json"), "{}").unwrap();
79
80 assert_eq!(run_slots(&cond_dir), legacy_slot(&cond_dir));
81 }
82
83 #[test]
84 fn nonexistent_cond_dir_yields_single_legacy_slot() {
85 let root = TempDir::new().unwrap();
86 let cond_dir = root.path().join("missing");
87
88 assert_eq!(run_slots(&cond_dir), legacy_slot(&cond_dir));
89 }
90
91 #[test]
92 fn nested_run_dirs_yield_indexed_slots_sorted_numerically() {
93 let root = TempDir::new().unwrap();
94 let cond_dir = root.path().join("with_skill");
95 for k in [10, 2, 1] {
96 fs::create_dir_all(cond_dir.join(format!("run-{k}"))).unwrap();
97 }
98
99 let slots = run_slots(&cond_dir);
100 assert_eq!(
101 slots
102 .iter()
103 .map(|s| s.run_index)
104 .collect::<Vec<Option<u32>>>(),
105 vec![Some(1), Some(2), Some(10)]
106 );
107 assert_eq!(slots[2].dir, cond_dir.join("run-10"));
108 }
109
110 #[test]
111 fn entries_that_are_not_run_dirs_are_ignored() {
112 let root = TempDir::new().unwrap();
113 let cond_dir = root.path().join("with_skill");
114 fs::create_dir_all(cond_dir.join("run-1")).unwrap();
115 fs::create_dir_all(cond_dir.join("outputs")).unwrap();
116 fs::create_dir_all(cond_dir.join("run-zero")).unwrap();
117 fs::write(cond_dir.join("run-2"), "a file, not a dir").unwrap();
118
119 let slots = run_slots(&cond_dir);
120 assert_eq!(
121 slots,
122 vec![RunSlot {
123 run_index: Some(1),
124 dir: cond_dir.join("run-1"),
125 }]
126 );
127 }
128
129 #[test]
130 fn cond_dir_with_only_non_run_entries_is_a_legacy_slot() {
131 let root = TempDir::new().unwrap();
132 let cond_dir = root.path().join("with_skill");
133 fs::create_dir_all(cond_dir.join("outputs")).unwrap();
134 fs::write(cond_dir.join("grading.json"), "{}").unwrap();
135
136 assert_eq!(run_slots(&cond_dir), legacy_slot(&cond_dir));
137 }
138}