somatize_core/
tracking.rs

1//! Experiment tracking types: run manifests, status, event envelopes,
2//! and the [`EventSink`]/[`Tracker`] traits.
3//!
4//! A *run* is the unit of tracking — one training session, study, or
5//! fit — materialized as a directory of append-only logs plus small
6//! atomic JSON files. This module holds only the schema and trait
7//! contracts; the file-writing implementation lives in `soma-runtime`
8//! (`LocalTracker`), and a future remote backend implements the same
9//! [`Tracker`] trait.
10
11use crate::error::Result;
12use crate::event::Event;
13use crate::study::Study;
14use chrono::{DateTime, Utc};
15use serde::{Deserialize, Serialize};
16use std::collections::HashMap;
17use std::path::Path;
18use std::sync::Arc;
19
20/// Version of the on-disk run schema (manifest + logs layout).
21pub const RUN_SCHEMA_VERSION: u32 = 1;
22
23/// What kind of work a run tracks.
24#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
25#[serde(rename_all = "snake_case")]
26#[non_exhaustive]
27pub enum RunKind {
28    /// A `GraphSession::fit` over a static pipeline.
29    Fit,
30    /// A native training loop (materialize/forward/backward/step).
31    Train,
32    /// A hyperparameter study.
33    Study,
34    /// A single trial within a study.
35    Trial,
36    /// Anything else — also the fallback when deserializing a kind
37    /// written by a newer soma, so old readers never fail on new kinds.
38    #[serde(other)]
39    Other,
40}
41
42/// Lifecycle state of a run.
43#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
44#[serde(rename_all = "snake_case")]
45#[non_exhaustive]
46pub enum RunState {
47    /// The run's process claims to be alive; trust it only while
48    /// [`RunStatus::heartbeat_at`] is fresh.
49    Running,
50    /// Finished successfully.
51    Completed,
52    /// Finished with an error.
53    Failed,
54}
55
56/// Best-effort git context captured at run start.
57#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
58pub struct GitInfo {
59    /// Commit hash of `HEAD`.
60    #[serde(default)]
61    pub sha: Option<String>,
62    /// Checked-out branch name, `None` on a detached head.
63    #[serde(default)]
64    pub branch: Option<String>,
65    /// Whether the working tree had uncommitted changes — a dirty run
66    /// is one the recorded `sha` cannot fully reproduce.
67    #[serde(default)]
68    pub dirty: Option<bool>,
69}
70
71/// Compact description of the graph a run executed, with pointers to
72/// the full topology files inside the run directory.
73#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
74pub struct GraphSummaryInfo {
75    /// Number of nodes in the executed graph.
76    pub n_nodes: usize,
77    /// Node ids, in the graph's insertion order.
78    pub node_ids: Vec<String>,
79    /// Relative path to the serialized graph (e.g. `graph.json`).
80    #[serde(default)]
81    pub graph_path: Option<String>,
82    /// Relative path to the mermaid rendering (e.g. `graph.mmd`).
83    #[serde(default)]
84    pub mermaid_path: Option<String>,
85}
86
87/// Run manifest — written once, atomically, at run start.
88///
89/// Mutable lifecycle state (running/completed/failed, heartbeat) lives
90/// in the separate [`RunStatus`] file so the manifest never needs
91/// rewriting after creation.
92#[derive(Debug, Clone, Serialize, Deserialize)]
93pub struct RunManifest {
94    /// On-disk layout version this run was written with
95    /// (see [`RUN_SCHEMA_VERSION`]).
96    pub schema_version: u32,
97    /// Unique run identifier — also the run directory's name.
98    pub run_id: String,
99    /// What kind of work this run tracks.
100    pub kind: RunKind,
101    /// Human-readable run name (not required to be unique).
102    pub name: String,
103    /// When the run started.
104    pub created_at: DateTime<Utc>,
105    /// Version of soma that wrote this run.
106    #[serde(default)]
107    pub soma_version: Option<String>,
108    /// Python interpreter version, for runs started from the bindings.
109    #[serde(default)]
110    pub python_version: Option<String>,
111    /// Host the run executed on.
112    #[serde(default)]
113    pub hostname: Option<String>,
114    /// Best-effort git context captured at run start.
115    #[serde(default)]
116    pub git: GitInfo,
117    /// Script or module that started the run.
118    #[serde(default)]
119    pub entrypoint: Option<String>,
120    /// Command-line arguments of the launching process.
121    #[serde(default)]
122    pub argv: Vec<String>,
123    /// Working directory the run was started from.
124    #[serde(default)]
125    pub cwd: Option<String>,
126    /// Named seeds, e.g. `{"torch": 42}`.
127    #[serde(default)]
128    pub seeds: HashMap<String, i64>,
129    /// Hyperparameters the caller declared for this run — the knobs
130    /// that live outside the graph (learning rate, batch size, …) and
131    /// so cannot be recovered from a filter's config hash. What makes
132    /// a `ParamChanged` derivation possible at all.
133    #[serde(default)]
134    pub params: HashMap<String, serde_json::Value>,
135    /// What the person starting this run expected, and why. Recorded at
136    /// the start rather than the end on purpose: a hypothesis written
137    /// after seeing the result is a conclusion.
138    #[serde(default)]
139    pub hypothesis: Option<String>,
140    /// Free-form labels for filtering run listings.
141    #[serde(default)]
142    pub tags: Vec<String>,
143    /// Free-form notes attached at run start.
144    #[serde(default)]
145    pub notes: Option<String>,
146    /// Run this one derives from — the edge the experiment pool's
147    /// lineage is built on. Set explicitly, never inferred.
148    #[serde(default)]
149    pub parent_run_id: Option<String>,
150    /// Compact description of the executed graph, absent for
151    /// graph-less runs (e.g. a study run).
152    #[serde(default)]
153    pub graph: Option<GraphSummaryInfo>,
154    /// Relative path to `study.json` for study runs.
155    #[serde(default)]
156    pub study_path: Option<String>,
157}
158
159impl RunManifest {
160    /// Minimal manifest; callers fill in environment fields.
161    pub fn new(run_id: impl Into<String>, kind: RunKind, name: impl Into<String>) -> Self {
162        Self {
163            schema_version: RUN_SCHEMA_VERSION,
164            run_id: run_id.into(),
165            kind,
166            name: name.into(),
167            created_at: Utc::now(),
168            soma_version: None,
169            python_version: None,
170            hostname: None,
171            git: GitInfo::default(),
172            entrypoint: None,
173            argv: Vec::new(),
174            cwd: None,
175            seeds: HashMap::new(),
176            params: HashMap::new(),
177            hypothesis: None,
178            tags: Vec::new(),
179            notes: None,
180            parent_run_id: None,
181            graph: None,
182            study_path: None,
183        }
184    }
185}
186
187/// Mutable run status — small, atomically rewritten (`status.json`).
188///
189/// The heartbeat lets a reader distinguish a live run from a crashed
190/// one without any protocol: `state == Running` with a stale
191/// `heartbeat_at` means the process died.
192#[derive(Debug, Clone, Serialize, Deserialize)]
193pub struct RunStatus {
194    /// Current lifecycle state.
195    pub state: RunState,
196    /// When this status file was last rewritten, for any reason.
197    pub updated_at: DateTime<Utc>,
198    /// Last liveness ping. Stale while `state` is
199    /// [`RunState::Running`] means the process died.
200    #[serde(default)]
201    pub heartbeat_at: Option<DateTime<Utc>>,
202    /// When the run reached a terminal state, `None` while running.
203    #[serde(default)]
204    pub finished_at: Option<DateTime<Utc>>,
205}
206
207impl RunStatus {
208    /// Fresh status for a run that just started: state
209    /// [`RunState::Running`] with the heartbeat stamped now.
210    pub fn running() -> Self {
211        let now = Utc::now();
212        Self {
213            state: RunState::Running,
214            updated_at: now,
215            heartbeat_at: Some(now),
216            finished_at: None,
217        }
218    }
219}
220
221/// One line of `events.jsonl`: a monotonic sequence number and wall
222/// timestamp wrapped around the event, with the event's own
223/// `event_type` tag flattened into the same object.
224#[derive(Debug, Clone, Serialize, Deserialize)]
225pub struct EventEnvelope {
226    /// Monotonic sequence number within the run — the order events were
227    /// recorded in, which timestamps alone cannot guarantee.
228    pub seq: u64,
229    /// Wall-clock time the event was recorded.
230    pub ts: DateTime<Utc>,
231    /// The event itself, flattened into the envelope's JSON object.
232    #[serde(flatten)]
233    pub event: Event,
234}
235
236/// A lossless, ordered consumer of events.
237///
238/// Unlike broadcast subscribers (which may lag and drop), sinks are
239/// invoked synchronously from `EventBus::emit` and must never lose an
240/// event. Implementations should buffer writes and must not panic;
241/// I/O errors are to be swallowed (optionally logged), never surfaced
242/// into the training loop.
243pub trait EventSink: Send + Sync {
244    /// Record one event. Called synchronously on the emitting thread.
245    fn record(&self, event: &Event);
246
247    /// Flush any buffered state to durable storage.
248    fn flush(&self) {}
249}
250
251/// A tracking backend bound to one run.
252///
253/// The local implementation writes a run directory; a remote backend
254/// can implement the same contract over HTTP.
255pub trait Tracker: Send + Sync {
256    /// Identifier of the run this tracker is bound to.
257    fn run_id(&self) -> &str;
258
259    /// Root directory of the run (for file-based backends).
260    fn run_dir(&self) -> &Path;
261
262    /// The sink that persists events for this run.
263    fn sink(&self) -> Arc<dyn EventSink>;
264
265    /// Atomically write the manifest.
266    fn save_manifest(&self, manifest: &RunManifest) -> Result<()>;
267
268    /// Write an artifact at a path relative to the run directory,
269    /// creating parent directories as needed.
270    fn save_artifact(&self, rel_path: &str, bytes: &[u8]) -> Result<()>;
271
272    /// Atomically write `study.json` (tmp + rename — readers never see
273    /// a partial study, and a crash mid-write preserves the previous
274    /// complete version).
275    fn save_study(&self, study: &Study) -> Result<()>;
276
277    /// Refresh `heartbeat_at` in the status file.
278    fn heartbeat(&self) -> Result<()>;
279
280    /// Set the terminal state, stamp `finished_at`, and flush the sink.
281    fn finalize(&self, state: RunState) -> Result<()>;
282}
283
284#[cfg(test)]
285mod tests {
286    use super::*;
287    use crate::event::MetricRecord;
288
289    #[test]
290    fn manifest_roundtrip_and_defaults() {
291        let mut m = RunManifest::new("run_x", RunKind::Train, "baseline");
292        m.tags = vec!["mos".into()];
293        m.seeds.insert("torch".into(), 42);
294        let json = serde_json::to_string(&m).unwrap();
295        let back: RunManifest = serde_json::from_str(&json).unwrap();
296        assert_eq!(back.run_id, "run_x");
297        assert_eq!(back.schema_version, RUN_SCHEMA_VERSION);
298        assert_eq!(back.seeds["torch"], 42);
299
300        // Old manifests without the optional fields still load.
301        let minimal = serde_json::json!({
302            "schema_version": 1,
303            "run_id": "r",
304            "kind": "fit",
305            "name": "n",
306            "created_at": "2026-07-26T10:00:00Z",
307        });
308        let back: RunManifest = serde_json::from_value(minimal).unwrap();
309        assert!(back.git.sha.is_none());
310        assert!(back.argv.is_empty());
311    }
312
313    #[test]
314    fn envelope_flattens_event_type() {
315        let env = EventEnvelope {
316            seq: 7,
317            ts: Utc::now(),
318            event: Event::MetricReported {
319                run_id: "r1".into(),
320                metric: MetricRecord {
321                    name: "val_f1".into(),
322                    value: 0.9,
323                    step: 3,
324                    timestamp: Utc::now(),
325                },
326                node_id: None,
327                trial_id: None,
328            },
329        };
330        let json = serde_json::to_value(&env).unwrap();
331        assert_eq!(json["seq"], 7);
332        assert_eq!(json["event_type"], "MetricReported");
333        assert_eq!(json["metric"]["name"], "val_f1");
334        let back: EventEnvelope = serde_json::from_value(json).unwrap();
335        assert_eq!(back.seq, 7);
336        assert!(matches!(back.event, Event::MetricReported { .. }));
337    }
338
339    #[test]
340    fn run_status_serde() {
341        let s = RunStatus::running();
342        let json = serde_json::to_string(&s).unwrap();
343        assert!(json.contains("\"running\""));
344        let back: RunStatus = serde_json::from_str(&json).unwrap();
345        assert_eq!(back.state, RunState::Running);
346        assert!(back.finished_at.is_none());
347    }
348
349    #[test]
350    fn run_status_terminal_states_roundtrip() {
351        for state in [RunState::Completed, RunState::Failed] {
352            let now = Utc::now();
353            let s = RunStatus {
354                state,
355                updated_at: now,
356                heartbeat_at: Some(now),
357                finished_at: Some(now),
358            };
359            let back: RunStatus =
360                serde_json::from_str(&serde_json::to_string(&s).unwrap()).unwrap();
361            assert_eq!(back.state, state);
362            assert_eq!(back.finished_at, Some(now));
363        }
364        // Back-compat: a status without the optional timestamps loads.
365        let minimal = serde_json::json!({
366            "state": "completed",
367            "updated_at": "2026-07-26T10:00:00Z",
368        });
369        let back: RunStatus = serde_json::from_value(minimal).unwrap();
370        assert_eq!(back.state, RunState::Completed);
371        assert!(back.heartbeat_at.is_none());
372        assert!(back.finished_at.is_none());
373    }
374
375    #[test]
376    fn unknown_run_kind_falls_back_to_other() {
377        // A manifest written by a future soma with a new kind must not
378        // break `LocalTracker::open` on this version.
379        let manifest = serde_json::json!({
380            "schema_version": 2,
381            "run_id": "r",
382            "kind": "evaluation",
383            "name": "n",
384            "created_at": "2026-07-26T10:00:00Z",
385            "some_future_field": {"nested": true},
386        });
387        let back: RunManifest = serde_json::from_value(manifest).unwrap();
388        assert_eq!(back.kind, RunKind::Other);
389        // A reader can detect the newer schema explicitly.
390        assert!(back.schema_version > RUN_SCHEMA_VERSION);
391    }
392
393    #[test]
394    fn envelope_roundtrips_one_event_per_level() {
395        let now = Utc::now();
396        let metric = MetricRecord {
397            name: "f1".into(),
398            value: 0.5,
399            step: 1,
400            timestamp: now,
401        };
402        let events = vec![
403            Event::RunFailed {
404                run_id: "r".into(),
405                error: "boom".into(),
406            },
407            Event::TrialMetric {
408                study_id: "s".into(),
409                trial_id: "t".into(),
410                metric: metric.clone(),
411            },
412            Event::StudyProgress {
413                study_id: "s".into(),
414                completed: 1,
415                total: 4,
416                best_value: 0.5,
417            },
418            Event::MemberExploited {
419                study_id: "s".into(),
420                generation: 1,
421                replaced_id: "a".into(),
422                donor_id: "b".into(),
423            },
424            Event::HealthFlag {
425                run_id: "r".into(),
426                node_id: "n".into(),
427                step: 3,
428                flag: "LEAKAGE".into(),
429                detail: "cka=0.99".into(),
430            },
431        ];
432        for (i, event) in events.into_iter().enumerate() {
433            let env = EventEnvelope {
434                seq: i as u64,
435                ts: now,
436                event,
437            };
438            let json = serde_json::to_value(&env).unwrap();
439            // The envelope's own fields never collide with payloads.
440            assert_eq!(json["seq"], i as u64);
441            assert!(json["event_type"].is_string());
442            let back: EventEnvelope = serde_json::from_value(json).unwrap();
443            assert_eq!(back.seq, i as u64);
444            assert_eq!(back.ts, now);
445        }
446    }
447
448    #[test]
449    fn git_info_and_graph_summary_serde() {
450        let git = GitInfo {
451            sha: Some("abc123".into()),
452            branch: Some("main".into()),
453            dirty: Some(true),
454        };
455        let back: GitInfo = serde_json::from_str(&serde_json::to_string(&git).unwrap()).unwrap();
456        assert_eq!(back, git);
457        assert_eq!(GitInfo::default(), GitInfo::default());
458        assert!(GitInfo::default().sha.is_none());
459
460        let summary = GraphSummaryInfo {
461            n_nodes: 2,
462            node_ids: vec!["a".into(), "b".into()],
463            graph_path: Some("graph.json".into()),
464            mermaid_path: None,
465        };
466        let back: GraphSummaryInfo =
467            serde_json::from_str(&serde_json::to_string(&summary).unwrap()).unwrap();
468        assert_eq!(back, summary);
469        // Back-compat: path fields are optional.
470        let minimal: GraphSummaryInfo =
471            serde_json::from_value(serde_json::json!({"n_nodes": 1, "node_ids": ["x"]})).unwrap();
472        assert_eq!(minimal.n_nodes, 1);
473        assert!(minimal.graph_path.is_none());
474        assert_eq!(GraphSummaryInfo::default().n_nodes, 0);
475    }
476}
somatize_core/tracking.rs

somatize_core/
tracking.rs
