cortenforge_tools/
services.rs

1use std::env;
2use std::fs;
3use std::io;
4use std::io::BufRead;
5use std::path::{Path, PathBuf};
6use std::process::{Child, Command};
7
8use serde::Deserialize;
9#[cfg(any(feature = "tui", feature = "scheduler"))]
10use sysinfo::{Pid, ProcessRefreshKind, ProcessesToUpdate, RefreshKind, System};
11use thiserror::Error;
12
13use crate::ToolConfig;
14
15#[derive(Debug, Clone, Deserialize)]
16pub struct RunManifestSummary {
17    pub schema_version: Option<u32>,
18    pub seed: Option<u64>,
19    pub output_root: Option<PathBuf>,
20    pub run_dir: Option<PathBuf>,
21    pub started_at_unix: Option<f64>,
22    pub max_frames: Option<u32>,
23}
24
25#[derive(Debug, Clone)]
26pub struct RunInfo {
27    pub path: PathBuf,
28    pub manifest: Option<RunManifestSummary>,
29    pub label_count: usize,
30    pub image_count: usize,
31    pub overlay_count: usize,
32}
33
34#[derive(Debug, Error)]
35pub enum ServiceError {
36    #[error("io error: {0}")]
37    Io(#[from] io::Error),
38    #[error("json error: {0}")]
39    Json(#[from] serde_json::Error),
40}
41
42/// List run directories under a root, with manifest and counts if available.
43pub fn list_runs(root: &Path) -> Result<Vec<RunInfo>, ServiceError> {
44    let mut out = Vec::new();
45    for entry in fs::read_dir(root)? {
46        let entry = entry?;
47        let path = entry.path();
48        if !path.is_dir() {
49            continue;
50        }
51        let name_ok = path
52            .file_name()
53            .and_then(|s| s.to_str())
54            .map(|s| s.starts_with("run_"))
55            .unwrap_or(false);
56        if !name_ok {
57            continue;
58        }
59        let manifest_path = path.join("run_manifest.json");
60        let manifest = if manifest_path.exists() {
61            let data = fs::read(&manifest_path)?;
62            serde_json::from_slice::<RunManifestSummary>(&data).ok()
63        } else {
64            None
65        };
66        let counts = count_artifacts(&path);
67        out.push(RunInfo {
68            path: path.clone(),
69            manifest,
70            label_count: counts.0,
71            image_count: counts.1,
72            overlay_count: counts.2,
73        });
74    }
75    out.sort_by(|a, b| a.path.cmp(&b.path));
76    Ok(out)
77}
78
79fn count_artifacts(run_dir: &Path) -> (usize, usize, usize) {
80    let labels = run_dir.join("labels");
81    let images = run_dir.join("images");
82    let overlays = run_dir.join("overlays");
83    let count_ext = |dir: &Path, ext: &str| -> usize {
84        fs::read_dir(dir)
85            .into_iter()
86            .flatten()
87            .filter_map(|e| e.ok())
88            .filter(|e| e.path().extension().and_then(|s| s.to_str()) == Some(ext))
89            .count()
90    };
91    (
92        count_ext(&labels, "json"),
93        count_ext(&images, "png"),
94        count_ext(&overlays, "png"),
95    )
96}
97
98#[derive(Debug, Clone)]
99pub struct ServiceCommand {
100    pub program: PathBuf,
101    pub args: Vec<String>,
102}
103
104pub fn spawn(cmd: &ServiceCommand) -> io::Result<Child> {
105    Command::new(&cmd.program).args(&cmd.args).spawn()
106}
107
108fn bin_path(bin: &str) -> io::Result<PathBuf> {
109    let mut exe = env::current_exe()?;
110    exe.pop(); // drop current binary name
111    exe.push(bin);
112    Ok(exe)
113}
114
115#[derive(Debug, Clone)]
116pub struct DatagenOptions {
117    pub output_root: PathBuf,
118    pub seed: Option<u64>,
119    pub max_frames: Option<u32>,
120    pub headless: bool,
121    pub prune_empty: bool,
122    pub prune_output_root: Option<PathBuf>,
123}
124
125/// Build a command to launch datagen (headless by default).
126pub fn datagen_command(opts: &DatagenOptions) -> io::Result<ServiceCommand> {
127    let cfg = ToolConfig::load();
128    datagen_command_with_config(&cfg, opts)
129}
130
131pub fn datagen_command_with_config(
132    cfg: &ToolConfig,
133    opts: &DatagenOptions,
134) -> io::Result<ServiceCommand> {
135    let bin = if cfg.sim_bin.is_absolute() {
136        cfg.sim_bin.clone()
137    } else {
138        bin_path(cfg.sim_bin.to_string_lossy().as_ref())?
139    };
140    let mut args = Vec::new();
141    if cfg.datagen_args.is_empty() {
142        args.push("--mode".into());
143        args.push("datagen".into());
144    } else {
145        args.extend(cfg.datagen_args.iter().cloned());
146    }
147    if let Some(seed) = opts.seed {
148        args.push("--seed".into());
149        args.push(seed.to_string());
150    }
151    if opts.headless {
152        args.push("--headless".into());
153    }
154    args.push("--output-root".into());
155    args.push(opts.output_root.display().to_string());
156    if let Some(max) = opts.max_frames {
157        args.push("--max-frames".into());
158        args.push(max.to_string());
159    }
160    if opts.headless && !opts.prune_empty && opts.prune_output_root.is_none() {
161        // no-op
162    }
163    if opts.prune_empty {
164        args.push("--prune-empty".into());
165        if let Some(root) = &opts.prune_output_root {
166            args.push("--prune-output-root".into());
167            args.push(root.display().to_string());
168        }
169    }
170    Ok(ServiceCommand { program: bin, args })
171}
172
173#[derive(Debug, Clone)]
174pub struct TrainOptions {
175    pub input_root: PathBuf,
176    pub val_ratio: f32,
177    pub batch_size: usize,
178    pub epochs: usize,
179    pub seed: Option<u64>,
180    pub drop_last: bool,
181    pub real_val_dir: Option<PathBuf>,
182    pub status_file: Option<PathBuf>,
183}
184
185/// Build a command to launch the training binary with common options.
186pub fn train_command(opts: &TrainOptions) -> io::Result<ServiceCommand> {
187    let cfg = ToolConfig::load();
188    train_command_with_config(&cfg, opts)
189}
190
191pub fn train_command_with_config(
192    cfg: &ToolConfig,
193    opts: &TrainOptions,
194) -> io::Result<ServiceCommand> {
195    let bin = if cfg.train_bin.is_absolute() {
196        cfg.train_bin.clone()
197    } else {
198        bin_path(cfg.train_bin.to_string_lossy().as_ref())?
199    };
200    let mut args = Vec::new();
201    if cfg.training_args.is_empty() {
202        args.extend([
203            "--input-root".into(),
204            opts.input_root.display().to_string(),
205            "--val-ratio".into(),
206            opts.val_ratio.to_string(),
207            "--batch-size".into(),
208            opts.batch_size.to_string(),
209            "--epochs".into(),
210            opts.epochs.to_string(),
211        ]);
212    } else {
213        args.extend(cfg.training_args.iter().cloned());
214    }
215    if let Some(seed) = opts.seed {
216        args.push("--seed".into());
217        args.push(seed.to_string());
218    }
219    if opts.drop_last {
220        args.push("--drop-last".into());
221    }
222    if let Some(val_dir) = &opts.real_val_dir {
223        args.push("--real-val-dir".into());
224        args.push(val_dir.display().to_string());
225    }
226    if let Some(status) = &opts.status_file {
227        args.push("--status-file".into());
228        args.push(status.display().to_string());
229    }
230    Ok(ServiceCommand { program: bin, args })
231}
232
233/// Read JSONL metrics (e.g., from `--metrics-out`) and return parsed entries.
234pub fn read_metrics(
235    path: &Path,
236    limit: Option<usize>,
237) -> Result<Vec<serde_json::Value>, ServiceError> {
238    let file = fs::File::open(path)?;
239    let reader = io::BufReader::new(file);
240    let mut rows: Vec<serde_json::Value> = reader
241        .lines()
242        .map_while(Result::ok)
243        .filter_map(|line| serde_json::from_str(&line).ok())
244        .collect();
245    if let Some(n) = limit {
246        if rows.len() > n {
247            rows.drain(0..rows.len().saturating_sub(n));
248        }
249    }
250    Ok(rows)
251}
252
253/// Read the last N lines of a text log.
254pub fn read_log_tail(path: &Path, limit: usize) -> Result<Vec<String>, ServiceError> {
255    let file = fs::File::open(path)?;
256    let reader = io::BufReader::new(file);
257    let mut lines: Vec<String> = reader.lines().map_while(Result::ok).collect();
258    if lines.len() > limit {
259        lines.drain(0..lines.len().saturating_sub(limit));
260    }
261    Ok(lines)
262}
263
264#[cfg(any(feature = "tui", feature = "scheduler"))]
265pub fn is_process_running(pid: u32) -> bool {
266    let mut sys = System::new_with_specifics(
267        RefreshKind::nothing().with_processes(ProcessRefreshKind::everything()),
268    );
269    let pid = Pid::from_u32(pid);
270    sys.refresh_processes(ProcessesToUpdate::Some(&[pid]), false);
271    sys.process(pid).is_some()
272}
273
274#[cfg(any(feature = "tui", feature = "scheduler"))]
275pub fn read_status(path: &Path) -> Option<serde_json::Value> {
276    let data = std::fs::read(path).ok()?;
277    serde_json::from_slice(&data).ok()
278}
cortenforge_tools/services.rs

cortenforge_tools/
services.rs
