somatize_runtime/tracking/
local_tracker.rs

1//! Local run-directory tracker.
2
3use super::JsonlEventSink;
4use chrono::Utc;
5use somatize_core::error::{Result, SomaError};
6use somatize_core::study::Study;
7use somatize_core::tracking::{
8    EventSink, GitInfo, RunKind, RunManifest, RunState, RunStatus, Tracker,
9};
10use std::fs;
11use std::path::{Path, PathBuf};
12use std::process::Command;
13use std::sync::Arc;
14
15const EVENTS_FILE: &str = "events.jsonl";
16const METRICS_FILE: &str = "metrics.jsonl";
17const MANIFEST_FILE: &str = "manifest.json";
18const STATUS_FILE: &str = "status.json";
19const STUDY_FILE: &str = "study.json";
20const FLUSH_EVERY: usize = 20;
21
22/// A run directory under `<root>/runs/<run_id>/`.
23///
24/// `create` starts a new run (manifest + running status + fresh logs);
25/// `open` re-attaches to an existing directory for resume, continuing
26/// the event sequence where it left off.
27pub struct LocalTracker {
28    run_id: String,
29    dir: PathBuf,
30    sink: Arc<JsonlEventSink>,
31}
32
33impl LocalTracker {
34    /// Create a new run directory under `root` (e.g. `.soma`) and write
35    /// its manifest (environment fields best-effort) and running status.
36    pub fn create(root: impl AsRef<Path>, kind: RunKind, name: &str) -> Result<Self> {
37        let run_id = new_run_id(kind);
38        let dir = root.as_ref().join("runs").join(&run_id);
39        fs::create_dir_all(&dir)?;
40
41        let mut manifest = RunManifest::new(&run_id, kind, name);
42        manifest.soma_version = Some(env!("CARGO_PKG_VERSION").to_string());
43        manifest.hostname = hostname();
44        manifest.git = collect_git_info(Path::new("."));
45        manifest.argv = std::env::args().collect();
46        manifest.entrypoint = manifest.argv.first().cloned();
47        manifest.cwd = std::env::current_dir()
48            .ok()
49            .map(|p| p.display().to_string());
50        if kind == RunKind::Study {
51            manifest.study_path = Some(STUDY_FILE.to_string());
52        }
53        atomic_write_json(&dir.join(MANIFEST_FILE), &manifest)?;
54        atomic_write_json(&dir.join(STATUS_FILE), &RunStatus::running())?;
55
56        let sink = JsonlEventSink::create(
57            &dir.join(EVENTS_FILE),
58            Some(&dir.join(METRICS_FILE)),
59            FLUSH_EVERY,
60        )?;
61        Ok(Self {
62            run_id,
63            dir,
64            sink: Arc::new(sink),
65        })
66    }
67
68    /// Re-attach to an existing run directory (resume). The event
69    /// sequence continues after the last recorded line, and the status
70    /// flips back to running.
71    pub fn open(run_dir: impl AsRef<Path>) -> Result<Self> {
72        let dir = run_dir.as_ref().to_path_buf();
73        let manifest = load_manifest(&dir)?;
74        let start_seq = repair_and_count_lines(&dir.join(EVENTS_FILE))?;
75        let sink = JsonlEventSink::append(
76            &dir.join(EVENTS_FILE),
77            Some(&dir.join(METRICS_FILE)),
78            FLUSH_EVERY,
79            start_seq,
80        )?;
81        atomic_write_json(&dir.join(STATUS_FILE), &RunStatus::running())?;
82        Ok(Self {
83            run_id: manifest.run_id,
84            dir,
85            sink: Arc::new(sink),
86        })
87    }
88}
89
90impl Tracker for LocalTracker {
91    fn run_id(&self) -> &str {
92        &self.run_id
93    }
94
95    fn run_dir(&self) -> &Path {
96        &self.dir
97    }
98
99    fn sink(&self) -> Arc<dyn EventSink> {
100        self.sink.clone()
101    }
102
103    fn save_manifest(&self, manifest: &RunManifest) -> Result<()> {
104        atomic_write_json(&self.dir.join(MANIFEST_FILE), manifest)
105    }
106
107    fn save_artifact(&self, rel_path: &str, bytes: &[u8]) -> Result<()> {
108        let path = self.dir.join(rel_path);
109        if let Some(parent) = path.parent() {
110            fs::create_dir_all(parent)?;
111        }
112        fs::write(path, bytes)?;
113        Ok(())
114    }
115
116    fn save_study(&self, study: &Study) -> Result<()> {
117        atomic_write_json(&self.dir.join(STUDY_FILE), study)
118    }
119
120    fn heartbeat(&self) -> Result<()> {
121        let mut status = load_status(&self.dir)?;
122        let now = Utc::now();
123        status.heartbeat_at = Some(now);
124        status.updated_at = now;
125        atomic_write_json(&self.dir.join(STATUS_FILE), &status)
126    }
127
128    fn finalize(&self, state: RunState) -> Result<()> {
129        self.sink.flush();
130        let now = Utc::now();
131        let status = RunStatus {
132            state,
133            updated_at: now,
134            heartbeat_at: Some(now),
135            finished_at: Some(now),
136        };
137        atomic_write_json(&self.dir.join(STATUS_FILE), &status)
138    }
139}
140
141/// Read a run's manifest.
142pub fn load_manifest(run_dir: &Path) -> Result<RunManifest> {
143    let bytes = fs::read(run_dir.join(MANIFEST_FILE))?;
144    serde_json::from_slice(&bytes).map_err(|e| SomaError::Serialization(e.to_string()))
145}
146
147/// Read a run's status file.
148pub fn load_status(run_dir: &Path) -> Result<RunStatus> {
149    let bytes = fs::read(run_dir.join(STATUS_FILE))?;
150    serde_json::from_slice(&bytes).map_err(|e| SomaError::Serialization(e.to_string()))
151}
152
153/// Best-effort git context via subprocess; all-`None` outside a repo.
154pub fn collect_git_info(dir: &Path) -> GitInfo {
155    let run = |args: &[&str]| -> Option<String> {
156        let out = Command::new("git")
157            .args(args)
158            .current_dir(dir)
159            .output()
160            .ok()?;
161        if !out.status.success() {
162            return None;
163        }
164        let s = String::from_utf8_lossy(&out.stdout).trim().to_string();
165        (!s.is_empty()).then_some(s)
166    };
167    GitInfo {
168        sha: run(&["rev-parse", "HEAD"]),
169        branch: run(&["rev-parse", "--abbrev-ref", "HEAD"]),
170        dirty: run(&["status", "--porcelain"]).map(|s| !s.is_empty()).or({
171            // `git status` succeeded with empty output → clean repo.
172            run(&["rev-parse", "HEAD"]).map(|_| false)
173        }),
174    }
175}
176
177fn hostname() -> Option<String> {
178    std::env::var("HOSTNAME")
179        .ok()
180        .filter(|h| !h.is_empty())
181        .or_else(|| {
182            fs::read_to_string("/etc/hostname")
183                .ok()
184                .map(|s| s.trim().to_string())
185                .filter(|h| !h.is_empty())
186        })
187        .or_else(|| {
188            let out = Command::new("hostname").output().ok()?;
189            let s = String::from_utf8_lossy(&out.stdout).trim().to_string();
190            (!s.is_empty()).then_some(s)
191        })
192}
193
194/// Sortable, collision-resistant run id: `<prefix>_<utc-compact>_<4hex>`.
195fn new_run_id(kind: RunKind) -> String {
196    let prefix = match kind {
197        RunKind::Study => "study",
198        RunKind::Trial => "trial",
199        _ => "run",
200    };
201    let nanos = std::time::SystemTime::now()
202        .duration_since(std::time::UNIX_EPOCH)
203        .unwrap_or_default()
204        .as_nanos();
205    format!(
206        "{prefix}_{}_{:04x}",
207        Utc::now().format("%Y%m%dT%H%M%S"),
208        (nanos & 0xffff) as u16
209    )
210}
211
212/// Write JSON via tmp file + rename so readers never observe a partial
213/// file and a crash mid-write preserves the previous version.
214fn atomic_write_json<T: serde::Serialize>(path: &Path, value: &T) -> Result<()> {
215    let json =
216        serde_json::to_vec_pretty(value).map_err(|e| SomaError::Serialization(e.to_string()))?;
217    let tmp = path.with_extension("json.tmp");
218    fs::write(&tmp, &json)?;
219    fs::rename(&tmp, path)?;
220    Ok(())
221}
222
223/// Prepare the events log for appending: a crash mid-write leaves a
224/// torn trailing line with no `\n`; appending after it would concatenate
225/// the next event onto garbage. Truncate the torn tail (that event was
226/// never durably recorded) and return the number of complete lines —
227/// the next sequence number.
228fn repair_and_count_lines(path: &Path) -> Result<u64> {
229    let content = match fs::read(path) {
230        Ok(c) => c,
231        Err(_) => return Ok(0), // no events yet — fresh log
232    };
233    let newlines = content.iter().filter(|b| **b == b'\n').count() as u64;
234    if content.is_empty() || content.last() == Some(&b'\n') {
235        return Ok(newlines);
236    }
237    let keep = content
238        .iter()
239        .rposition(|b| *b == b'\n')
240        .map(|i| i + 1)
241        .unwrap_or(0);
242    let file = fs::OpenOptions::new().write(true).open(path)?;
243    file.set_len(keep as u64)?;
244    Ok(newlines)
245}
somatize_runtime/tracking/local_tracker.rs

somatize_runtime/tracking/
local_tracker.rs
