1use std::path::Path;
18
19#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Deserialize, serde::Serialize, schemars::JsonSchema)]
21#[serde(rename_all = "kebab-case")]
22pub enum FlowType {
23 Train,
24 Tune,
25 TuneThenTrain,
26 TrainFromTune,
27 ResumeRun,
28}
29
30impl std::fmt::Display for FlowType {
31 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
32 let s = match self {
33 FlowType::Train => "train",
34 FlowType::Tune => "tune",
35 FlowType::TuneThenTrain => "tune-then-train",
36 FlowType::TrainFromTune => "train-from-tune",
37 FlowType::ResumeRun => "resume-run",
38 };
39 write!(f, "{s}")
40 }
41}
42
43#[derive(Debug, Clone, Copy, PartialEq, Eq)]
47pub enum FieldRequirement {
48 Required,
49 Optional,
50 NotApplicable,
51}
52
53#[derive(Debug, Clone, Copy, PartialEq, Eq)]
59pub enum JobField {
60 TaskName,
61 Algorithm,
62 TotalTimesteps,
63 ResumeRunId,
64 TuneRunId,
65 NTrials,
66 TuneTimesteps,
67 EvalEps,
68}
69
70impl JobField {
71 pub fn wire_name(self) -> &'static str {
78 match self {
79 JobField::TaskName => "task_name",
80 JobField::Algorithm => "algorithm",
81 JobField::TotalTimesteps => "total_timesteps",
82 JobField::ResumeRunId => "resume_run_id",
83 JobField::TuneRunId => "tune_run_id",
84 JobField::NTrials => "n_trials",
85 JobField::TuneTimesteps => "tune_timesteps",
86 JobField::EvalEps => "eval_eps",
87 }
88 }
89
90 pub fn requirement(self, flow_type: FlowType) -> FieldRequirement {
94 use FieldRequirement::*;
95 use FlowType::*;
96 use JobField::*;
97
98 match (self, flow_type) {
99 (TaskName, _) => Required,
100 (Algorithm, _) => Required,
101
102 (TotalTimesteps, Train | TuneThenTrain | TrainFromTune | ResumeRun) => Required,
103 (TotalTimesteps, Tune) => NotApplicable,
104
105 (ResumeRunId, ResumeRun) => Required,
106 (ResumeRunId, _) => NotApplicable,
107
108 (TuneRunId, TrainFromTune) => Required,
109 (TuneRunId, _) => NotApplicable,
110
111 (NTrials | TuneTimesteps | EvalEps, Tune | TuneThenTrain) => Required,
112 (NTrials | TuneTimesteps | EvalEps, _) => NotApplicable,
113 }
114 }
115}
116
117#[derive(Debug, Clone, Default)]
123pub struct JobFields {
124 pub task_name: Option<String>,
125 pub algorithm: Option<String>,
126 pub total_timesteps: Option<i64>,
127 pub resume_run_id: Option<String>,
128 pub tune_run_id: Option<String>,
129 pub n_trials: Option<i64>,
130 pub tune_timesteps: Option<i64>,
131 pub eval_eps: Option<i64>,
132}
133
134impl JobFields {
135 fn get(&self, field: JobField) -> FieldPresence<'_> {
136 match field {
137 JobField::TaskName => FieldPresence::Str(self.task_name.as_deref()),
138 JobField::Algorithm => FieldPresence::Str(self.algorithm.as_deref()),
139 JobField::TotalTimesteps => FieldPresence::Int(self.total_timesteps),
140 JobField::ResumeRunId => FieldPresence::Str(self.resume_run_id.as_deref()),
141 JobField::TuneRunId => FieldPresence::Str(self.tune_run_id.as_deref()),
142 JobField::NTrials => FieldPresence::Int(self.n_trials),
143 JobField::TuneTimesteps => FieldPresence::Int(self.tune_timesteps),
144 JobField::EvalEps => FieldPresence::Int(self.eval_eps),
145 }
146 }
147}
148
149enum FieldPresence<'a> {
150 Str(Option<&'a str>),
151 Int(Option<i64>),
152}
153
154impl FieldPresence<'_> {
155 fn is_present(&self) -> bool {
161 match self {
162 FieldPresence::Str(s) => s.is_some_and(|s| !s.trim().is_empty()),
163 FieldPresence::Int(i) => i.is_some(),
164 }
165 }
166}
167
168#[derive(Debug, Clone, PartialEq, Eq)]
170pub struct MissingField {
171 pub field: JobField,
172 pub flow_type: FlowType,
173}
174
175impl std::fmt::Display for MissingField {
176 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
177 write!(
178 f,
179 "'{:?}' is required for flow_type '{}' but was not set",
180 self.field, self.flow_type
181 )
182 }
183}
184
185const ALL_FIELDS: &[JobField] = &[
186 JobField::TaskName,
187 JobField::Algorithm,
188 JobField::TotalTimesteps,
189 JobField::ResumeRunId,
190 JobField::TuneRunId,
191 JobField::NTrials,
192 JobField::TuneTimesteps,
193 JobField::EvalEps,
194];
195
196pub fn all_fields() -> &'static [JobField] {
200 ALL_FIELDS
201}
202
203pub fn validate_job(flow_type: FlowType, fields: &JobFields) -> Vec<MissingField> {
206 ALL_FIELDS
207 .iter()
208 .filter(|&&field| field.requirement(flow_type) == FieldRequirement::Required)
209 .filter(|&&field| !fields.get(field).is_present())
210 .map(|&field| MissingField { field, flow_type })
211 .collect()
212}
213
214#[derive(Debug, Clone, PartialEq, Eq)]
217pub enum TaskLookupError {
218 NotFound { task_name: String },
219 InvalidJson { task_name: String, message: String },
220}
221
222impl std::fmt::Display for TaskLookupError {
223 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
224 match self {
225 TaskLookupError::NotFound { task_name } => write!(
226 f,
227 "task '{task_name}' has no configs/rl/tasks/{task_name}/task.json - check the \
228 task_name against the project's task catalog"
229 ),
230 TaskLookupError::InvalidJson { task_name, message } => {
231 write!(f, "configs/rl/tasks/{task_name}/task.json is not valid JSON: {message}")
232 }
233 }
234 }
235}
236
237impl std::error::Error for TaskLookupError {}
238
239pub const RL_TASKS_DIR: &str = "configs/rl/tasks";
244
245pub fn task_exists(project_root: &Path, task_name: &str) -> Result<(), TaskLookupError> {
251 let path = project_root.join(RL_TASKS_DIR).join(task_name).join("task.json");
252 let content = std::fs::read_to_string(&path).map_err(|_| TaskLookupError::NotFound {
253 task_name: task_name.to_string(),
254 })?;
255 serde_json::from_str::<serde_json::Value>(&content).map_err(|e| TaskLookupError::InvalidJson {
256 task_name: task_name.to_string(),
257 message: e.to_string(),
258 })?;
259 Ok(())
260}
261
262pub fn list_tasks(project_root: &Path) -> std::io::Result<Vec<String>> {
268 let dir = project_root.join(RL_TASKS_DIR);
269 if !dir.exists() {
270 return Ok(Vec::new());
271 }
272
273 let mut names: Vec<String> = std::fs::read_dir(&dir)?
274 .filter_map(|entry| entry.ok())
275 .filter(|entry| entry.path().join("task.json").is_file())
276 .filter_map(|entry| entry.file_name().into_string().ok())
277 .collect();
278 names.sort();
279 Ok(names)
280}
281
282#[cfg(test)]
283mod tests {
284 use super::*;
285
286 #[test]
287 fn every_job_field_has_a_unique_wire_name() {
288 let names: Vec<&str> = ALL_FIELDS.iter().map(|f| f.wire_name()).collect();
289 let unique: std::collections::HashSet<&str> = names.iter().copied().collect();
290 assert_eq!(names.len(), unique.len());
291 }
292
293 fn fields_with_task_and_algo() -> JobFields {
294 JobFields {
295 task_name: Some("find_target".to_string()),
296 algorithm: Some("ppo".to_string()),
297 ..Default::default()
298 }
299 }
300
301 #[test]
302 fn train_requires_total_timesteps() {
303 let fields = fields_with_task_and_algo();
304 let missing = validate_job(FlowType::Train, &fields);
305 assert_eq!(
306 missing,
307 vec![MissingField {
308 field: JobField::TotalTimesteps,
309 flow_type: FlowType::Train
310 }]
311 );
312 }
313
314 #[test]
315 fn train_passes_with_total_timesteps_set() {
316 let mut fields = fields_with_task_and_algo();
317 fields.total_timesteps = Some(1000);
318 assert!(validate_job(FlowType::Train, &fields).is_empty());
319 }
320
321 #[test]
322 fn tune_does_not_require_total_timesteps() {
323 let mut fields = fields_with_task_and_algo();
324 fields.n_trials = Some(10);
325 fields.tune_timesteps = Some(1000);
326 fields.eval_eps = Some(2);
327 assert!(validate_job(FlowType::Tune, &fields).is_empty());
328 }
329
330 #[test]
331 fn tune_requires_tune_specific_fields() {
332 let fields = fields_with_task_and_algo();
333 let missing = validate_job(FlowType::Tune, &fields);
334 assert_eq!(missing.len(), 3);
335 assert!(missing.iter().all(|m| m.flow_type == FlowType::Tune));
336 }
337
338 #[test]
339 fn resume_run_requires_resume_run_id_and_total_timesteps() {
340 let fields = fields_with_task_and_algo();
341 let missing = validate_job(FlowType::ResumeRun, &fields);
342 assert_eq!(missing.len(), 2);
343 assert!(missing.iter().any(|m| m.field == JobField::ResumeRunId));
344 assert!(missing.iter().any(|m| m.field == JobField::TotalTimesteps));
345 }
346
347 #[test]
348 fn train_from_tune_requires_tune_run_id() {
349 let mut fields = fields_with_task_and_algo();
350 fields.total_timesteps = Some(1000);
351 let missing = validate_job(FlowType::TrainFromTune, &fields);
352 assert_eq!(
353 missing,
354 vec![MissingField {
355 field: JobField::TuneRunId,
356 flow_type: FlowType::TrainFromTune
357 }]
358 );
359 }
360
361 #[test]
362 fn missing_task_name_and_algorithm_reported_for_every_flow_type() {
363 let missing = validate_job(FlowType::Train, &JobFields::default());
364 assert!(missing.iter().any(|m| m.field == JobField::TaskName));
365 assert!(missing.iter().any(|m| m.field == JobField::Algorithm));
366 }
367
368 #[test]
369 fn blank_string_fields_count_as_absent() {
370 let mut fields = fields_with_task_and_algo();
371 fields.resume_run_id = Some(" ".to_string());
372 let missing = validate_job(FlowType::ResumeRun, &fields);
373 assert!(missing.iter().any(|m| m.field == JobField::ResumeRunId));
374 }
375
376 #[test]
377 fn task_exists_finds_bundled_find_target_task() {
378 let dir = tempfile::tempdir().unwrap();
379 let task_dir = dir.path().join("configs/rl/tasks/find_target");
380 std::fs::create_dir_all(&task_dir).unwrap();
381 std::fs::write(task_dir.join("task.json"), r#"{"env_name": "basic_arena"}"#).unwrap();
382 assert!(task_exists(dir.path(), "find_target").is_ok());
383 }
384
385 #[test]
386 fn task_exists_reports_missing_task() {
387 let dir = tempfile::tempdir().unwrap();
388 std::fs::create_dir_all(dir.path().join("configs/rl/tasks")).unwrap();
389 let err = task_exists(dir.path(), "nonexistent").unwrap_err();
390 assert_eq!(
391 err,
392 TaskLookupError::NotFound {
393 task_name: "nonexistent".to_string()
394 }
395 );
396 }
397
398 #[test]
399 fn list_tasks_returns_empty_when_directory_missing() {
400 let dir = tempfile::tempdir().unwrap();
401 assert_eq!(list_tasks(dir.path()).unwrap(), Vec::<String>::new());
402 }
403
404 #[test]
405 fn list_tasks_lists_directories_with_a_task_json_sorted() {
406 let dir = tempfile::tempdir().unwrap();
407 for name in ["task2", "find_target", "task1"] {
408 let task_dir = dir.path().join("configs/rl/tasks").join(name);
409 std::fs::create_dir_all(&task_dir).unwrap();
410 std::fs::write(task_dir.join("task.json"), "{}").unwrap();
411 }
412 std::fs::create_dir_all(dir.path().join("configs/rl/tasks/not-a-task")).unwrap();
413
414 assert_eq!(
415 list_tasks(dir.path()).unwrap(),
416 vec!["find_target".to_string(), "task1".to_string(), "task2".to_string()]
417 );
418 }
419
420 #[test]
421 fn task_exists_reports_invalid_json() {
422 let dir = tempfile::tempdir().unwrap();
423 let task_dir = dir.path().join("configs/rl/tasks/broken");
424 std::fs::create_dir_all(&task_dir).unwrap();
425 std::fs::write(task_dir.join("task.json"), "not json").unwrap();
426 let err = task_exists(dir.path(), "broken").unwrap_err();
427 assert!(matches!(err, TaskLookupError::InvalidJson { .. }));
428 }
429}