Skip to main content

eval_magic/pipeline/
slots.rs

1//! Run-slot enumeration for a condition cell.
2//!
3//! A condition directory holds either a single run directly (the legacy flat
4//! layout: `eval-<id>/<cond>/run.json`) or N nested runs under 1-based
5//! `run-<k>` subdirectories (`eval-<id>/<cond>/run-2/run.json`). Stages that
6//! walk the workspace by convention enumerate slots through [`run_slots`] so
7//! both layouts — including a mix across cells within one iteration — read the
8//! same way.
9
10use std::fs;
11use std::path::{Path, PathBuf};
12
13/// One run's artifact directory within a condition cell.
14#[derive(Debug, Clone, PartialEq, Eq)]
15pub struct RunSlot {
16    /// 1-based index parsed from the `run-<k>` directory name; `None` for the
17    /// legacy flat layout where the condition directory is the single slot.
18    pub run_index: Option<u32>,
19    pub dir: PathBuf,
20}
21
22/// The lookup key for a run within the iteration: `<eval_id>:<condition>`, with
23/// a `:r<k>` suffix for indexed runs in a multi-run cell.
24pub 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
31/// Enumerate the run slots of `cond_dir`: its `run-<k>` subdirectories sorted
32/// numerically when any exist, otherwise `cond_dir` itself as the single
33/// legacy slot. Existence checks stay with the caller — a nonexistent
34/// `cond_dir` still yields the legacy slot.
35pub 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}