1use std::path::Path;
16
17#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Deserialize, serde::Serialize)]
22#[serde(rename_all = "kebab-case")]
23pub enum FlowType {
24 Train,
25 Tune,
26 TuneThenTrain,
27 TrainFromTune,
28 ResumeRun,
29}
30
31impl std::fmt::Display for FlowType {
32 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
33 let s = match self {
34 FlowType::Train => "train",
35 FlowType::Tune => "tune",
36 FlowType::TuneThenTrain => "tune-then-train",
37 FlowType::TrainFromTune => "train-from-tune",
38 FlowType::ResumeRun => "resume-run",
39 };
40 write!(f, "{s}")
41 }
42}
43
44#[derive(Debug, Clone, Copy, PartialEq, Eq)]
48pub enum FieldRequirement {
49 Required,
50 Optional,
51 NotApplicable,
52}
53
54#[derive(Debug, Clone, Copy, PartialEq, Eq)]
60pub enum JobField {
61 TaskName,
62 Algorithm,
63 TotalTimesteps,
64 ResumeRunId,
65 TuneRunId,
66 NTrials,
67 TuneTimesteps,
68 EvalEps,
69}
70
71impl JobField {
72 pub fn requirement(self, flow_type: FlowType) -> FieldRequirement {
75 use FieldRequirement::*;
76 use FlowType::*;
77 use JobField::*;
78
79 match (self, flow_type) {
80 (TaskName, _) => Required,
84 (Algorithm, _) => Required,
85
86 (TotalTimesteps, Train | TuneThenTrain | TrainFromTune | ResumeRun) => Required,
89 (TotalTimesteps, Tune) => NotApplicable,
90
91 (ResumeRunId, ResumeRun) => Required,
93 (ResumeRunId, _) => NotApplicable,
94
95 (TuneRunId, TrainFromTune) => Required,
97 (TuneRunId, _) => NotApplicable,
98
99 (NTrials | TuneTimesteps | EvalEps, Tune | TuneThenTrain) => Required,
102 (NTrials | TuneTimesteps | EvalEps, _) => NotApplicable,
103 }
104 }
105}
106
107#[derive(Debug, Clone, Default)]
113pub struct JobFields {
114 pub task_name: Option<String>,
115 pub algorithm: Option<String>,
116 pub total_timesteps: Option<i64>,
117 pub resume_run_id: Option<String>,
118 pub tune_run_id: Option<String>,
119 pub n_trials: Option<i64>,
120 pub tune_timesteps: Option<i64>,
121 pub eval_eps: Option<i64>,
122}
123
124impl JobFields {
125 fn get(&self, field: JobField) -> FieldPresence<'_> {
126 match field {
127 JobField::TaskName => FieldPresence::Str(self.task_name.as_deref()),
128 JobField::Algorithm => FieldPresence::Str(self.algorithm.as_deref()),
129 JobField::TotalTimesteps => FieldPresence::Int(self.total_timesteps),
130 JobField::ResumeRunId => FieldPresence::Str(self.resume_run_id.as_deref()),
131 JobField::TuneRunId => FieldPresence::Str(self.tune_run_id.as_deref()),
132 JobField::NTrials => FieldPresence::Int(self.n_trials),
133 JobField::TuneTimesteps => FieldPresence::Int(self.tune_timesteps),
134 JobField::EvalEps => FieldPresence::Int(self.eval_eps),
135 }
136 }
137}
138
139enum FieldPresence<'a> {
140 Str(Option<&'a str>),
141 Int(Option<i64>),
142}
143
144impl FieldPresence<'_> {
145 fn is_present(&self) -> bool {
151 match self {
152 FieldPresence::Str(s) => s.is_some_and(|s| !s.trim().is_empty()),
153 FieldPresence::Int(i) => i.is_some(),
154 }
155 }
156}
157
158#[derive(Debug, Clone, PartialEq, Eq)]
160pub struct MissingField {
161 pub field: JobField,
162 pub flow_type: FlowType,
163}
164
165impl std::fmt::Display for MissingField {
166 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
167 write!(
168 f,
169 "'{:?}' is required for flow_type '{}' but was not set",
170 self.field, self.flow_type
171 )
172 }
173}
174
175const ALL_FIELDS: &[JobField] = &[
176 JobField::TaskName,
177 JobField::Algorithm,
178 JobField::TotalTimesteps,
179 JobField::ResumeRunId,
180 JobField::TuneRunId,
181 JobField::NTrials,
182 JobField::TuneTimesteps,
183 JobField::EvalEps,
184];
185
186pub fn validate_job(flow_type: FlowType, fields: &JobFields) -> Vec<MissingField> {
189 ALL_FIELDS
190 .iter()
191 .filter(|&&field| field.requirement(flow_type) == FieldRequirement::Required)
192 .filter(|&&field| !fields.get(field).is_present())
193 .map(|&field| MissingField { field, flow_type })
194 .collect()
195}
196
197#[derive(Debug, Clone, PartialEq, Eq)]
200pub enum TaskLookupError {
201 NotFound { task_name: String },
202 InvalidJson { task_name: String, message: String },
203}
204
205impl std::fmt::Display for TaskLookupError {
206 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
207 match self {
208 TaskLookupError::NotFound { task_name } => write!(
209 f,
210 "task '{task_name}' has no configs/rl/tasks/{task_name}/task.json - check the \
211 task_name against the project's task catalog"
212 ),
213 TaskLookupError::InvalidJson { task_name, message } => {
214 write!(f, "configs/rl/tasks/{task_name}/task.json is not valid JSON: {message}")
215 }
216 }
217 }
218}
219
220impl std::error::Error for TaskLookupError {}
221
222pub const RL_TASKS_DIR: &str = "configs/rl/tasks";
227
228pub fn task_exists(project_root: &Path, task_name: &str) -> Result<(), TaskLookupError> {
234 let path = project_root.join(RL_TASKS_DIR).join(task_name).join("task.json");
235 let content = std::fs::read_to_string(&path).map_err(|_| TaskLookupError::NotFound {
236 task_name: task_name.to_string(),
237 })?;
238 serde_json::from_str::<serde_json::Value>(&content).map_err(|e| TaskLookupError::InvalidJson {
239 task_name: task_name.to_string(),
240 message: e.to_string(),
241 })?;
242 Ok(())
243}
244
245#[cfg(test)]
246mod tests {
247 use super::*;
248
249 fn fields_with_task_and_algo() -> JobFields {
250 JobFields {
251 task_name: Some("find_target".to_string()),
252 algorithm: Some("ppo".to_string()),
253 ..Default::default()
254 }
255 }
256
257 #[test]
258 fn train_requires_total_timesteps() {
259 let fields = fields_with_task_and_algo();
260 let missing = validate_job(FlowType::Train, &fields);
261 assert_eq!(
262 missing,
263 vec![MissingField {
264 field: JobField::TotalTimesteps,
265 flow_type: FlowType::Train
266 }]
267 );
268 }
269
270 #[test]
271 fn train_passes_with_total_timesteps_set() {
272 let mut fields = fields_with_task_and_algo();
273 fields.total_timesteps = Some(1000);
274 assert!(validate_job(FlowType::Train, &fields).is_empty());
275 }
276
277 #[test]
278 fn tune_does_not_require_total_timesteps() {
279 let mut fields = fields_with_task_and_algo();
280 fields.n_trials = Some(10);
281 fields.tune_timesteps = Some(1000);
282 fields.eval_eps = Some(2);
283 assert!(validate_job(FlowType::Tune, &fields).is_empty());
284 }
285
286 #[test]
287 fn tune_requires_tune_specific_fields() {
288 let fields = fields_with_task_and_algo();
289 let missing = validate_job(FlowType::Tune, &fields);
290 assert_eq!(missing.len(), 3);
291 assert!(missing.iter().all(|m| m.flow_type == FlowType::Tune));
292 }
293
294 #[test]
295 fn resume_run_requires_resume_run_id_and_total_timesteps() {
296 let fields = fields_with_task_and_algo();
297 let missing = validate_job(FlowType::ResumeRun, &fields);
298 assert_eq!(missing.len(), 2);
299 assert!(missing.iter().any(|m| m.field == JobField::ResumeRunId));
300 assert!(missing.iter().any(|m| m.field == JobField::TotalTimesteps));
301 }
302
303 #[test]
304 fn train_from_tune_requires_tune_run_id() {
305 let mut fields = fields_with_task_and_algo();
306 fields.total_timesteps = Some(1000);
307 let missing = validate_job(FlowType::TrainFromTune, &fields);
308 assert_eq!(
309 missing,
310 vec![MissingField {
311 field: JobField::TuneRunId,
312 flow_type: FlowType::TrainFromTune
313 }]
314 );
315 }
316
317 #[test]
318 fn missing_task_name_and_algorithm_reported_for_every_flow_type() {
319 let missing = validate_job(FlowType::Train, &JobFields::default());
320 assert!(missing.iter().any(|m| m.field == JobField::TaskName));
321 assert!(missing.iter().any(|m| m.field == JobField::Algorithm));
322 }
323
324 #[test]
325 fn blank_string_fields_count_as_absent() {
326 let mut fields = fields_with_task_and_algo();
327 fields.resume_run_id = Some(" ".to_string());
328 let missing = validate_job(FlowType::ResumeRun, &fields);
329 assert!(missing.iter().any(|m| m.field == JobField::ResumeRunId));
330 }
331
332 #[test]
333 fn task_exists_finds_bundled_find_target_task() {
334 let dir = tempfile::tempdir().unwrap();
335 let task_dir = dir.path().join("configs/rl/tasks/find_target");
336 std::fs::create_dir_all(&task_dir).unwrap();
337 std::fs::write(task_dir.join("task.json"), r#"{"env_name": "basic_arena"}"#).unwrap();
338 assert!(task_exists(dir.path(), "find_target").is_ok());
339 }
340
341 #[test]
342 fn task_exists_reports_missing_task() {
343 let dir = tempfile::tempdir().unwrap();
344 std::fs::create_dir_all(dir.path().join("configs/rl/tasks")).unwrap();
345 let err = task_exists(dir.path(), "nonexistent").unwrap_err();
346 assert_eq!(
347 err,
348 TaskLookupError::NotFound {
349 task_name: "nonexistent".to_string()
350 }
351 );
352 }
353
354 #[test]
355 fn task_exists_reports_invalid_json() {
356 let dir = tempfile::tempdir().unwrap();
357 let task_dir = dir.path().join("configs/rl/tasks/broken");
358 std::fs::create_dir_all(&task_dir).unwrap();
359 std::fs::write(task_dir.join("task.json"), "not json").unwrap();
360 let err = task_exists(dir.path(), "broken").unwrap_err();
361 assert!(matches!(err, TaskLookupError::InvalidJson { .. }));
362 }
363}