kitchen_shell/
jobs.rs

1//! What runs in the background, in one place (from lumen's job registry): each job has a kind
2//! (`histology`, `model`, `export`…), a name saying what it is doing, how far it is, and a cancel
3//! flag. Only one job that drives the GPU runs at a time (a second is refused, saying what
4//! runs). The status bar lists them with a cancel button each; the work itself (off the UI
5//! thread) updates its job through its [`JobHandle`], which ends the job when dropped (finished,
6//! failed or cancelled alike).
7//!
8//! A job can leave an outcome ([`JobHandle::finish`]: what it found, or why it failed), kept for
9//! a while after it ends. A command that starts a job returns [`job_reply`] (`{"job": id}`), so
10//! whoever ran it (an agent through the MCP host) can wait for the job ([`is_running`]) and read
11//! its [`outcome`].
12
13use std::collections::VecDeque;
14use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
15use std::sync::{Arc, Mutex, OnceLock};
16
17/// A job running.
18#[derive(Debug, Clone)]
19pub struct Job {
20    /// Its id (unique in this run of the app).
21    pub id: u64,
22    /// What kind of work it is (`model`, `export`).
23    pub kind: String,
24    /// What it is doing now ("Finding cells: tile 3 of 12").
25    pub name: String,
26    /// How far (0..1), when known.
27    pub progress: Option<f32>,
28    /// Whether it drives the GPU (one such job at a time).
29    pub gpu: bool,
30    /// Set to ask it to stop.
31    pub cancel: Arc<AtomicBool>,
32}
33
34#[derive(Default)]
35struct Inner {
36    list: Vec<Job>,
37    next: u64,
38    /// The outcomes of the latest jobs that left one (newest last).
39    outcomes: VecDeque<(u64, Result<serde_json::Value, String>)>,
40}
41
42/// Outcomes kept after their jobs end.
43const KEPT_OUTCOMES: usize = 64;
44
45/// What a command that started job `id` returns: `{"job": id}`.
46pub fn job_reply(id: u64) -> serde_json::Value {
47    serde_json::json!({ "job": id })
48}
49
50/// The job a command's result names (`{"job": id, …}`), if any.
51pub fn job_of(result: &serde_json::Value) -> Option<u64> {
52    result.get("job").and_then(serde_json::Value::as_u64)
53}
54
55/// Whether the job `id` is still running.
56pub fn is_running(id: u64) -> bool {
57    inner().lock().unwrap_or_else(|e| e.into_inner()).list.iter().any(|o| o.id == id)
58}
59
60/// The outcome the job `id` left, once it ended (`None`: running, or it left none and has been
61/// forgotten).
62pub fn outcome(id: u64) -> Option<Result<serde_json::Value, String>> {
63    inner().lock().unwrap_or_else(|e| e.into_inner()).outcomes.iter().rev().find(|(i, _)| *i == id).map(|(_, o)| o.clone())
64}
65
66fn inner() -> &'static Mutex<Inner> {
67    static J: OnceLock<Mutex<Inner>> = OnceLock::new();
68    J.get_or_init(Default::default)
69}
70
71static GENERATION: AtomicU64 = AtomicU64::new(0);
72
73fn with<T>(f: impl FnOnce(&mut Inner) -> T) -> T {
74    let mut g = inner().lock().unwrap_or_else(|e| e.into_inner());
75    let out = f(&mut g);
76    GENERATION.fetch_add(1, Ordering::Relaxed);
77    out
78}
79
80/// Starts a job; refused while another GPU job runs (when this one is one too).
81pub fn start(kind: impl Into<String>, name: &str, gpu: bool) -> Result<JobHandle, String> {
82    let kind = kind.into();
83    with(|j| {
84        if gpu && let Some(other) = j.list.iter().find(|o| o.gpu) {
85            return Err(format!("{} is still running", other.name));
86        }
87        j.next += 1;
88        let cancel = Arc::new(AtomicBool::new(false));
89        crate::activity::job_started(j.next, &kind, name);
90        j.list.push(Job { id: j.next, kind, name: name.to_string(), progress: None, gpu, cancel: cancel.clone() });
91        Ok(JobHandle { id: j.next, cancel })
92    })
93}
94
95/// The jobs running, oldest first.
96pub fn running() -> Vec<Job> {
97    let g = inner().lock().unwrap_or_else(|e| e.into_inner());
98    g.list.clone()
99}
100
101/// The latest job of `kind`.
102pub fn of(kind: &str) -> Option<Job> {
103    running().into_iter().rev().find(|o| o.kind == kind)
104}
105
106/// Asks every job of `kind` to stop (when its work next checks).
107pub fn cancel(kind: &str) {
108    with(|j| j.list.iter().filter(|o| o.kind == kind).for_each(|o| o.cancel.store(true, Ordering::Relaxed)));
109}
110
111/// Asks the job `id` to stop.
112pub fn cancel_id(id: u64) {
113    with(|j| j.list.iter().filter(|o| o.id == id).for_each(|o| o.cancel.store(true, Ordering::Relaxed)));
114}
115
116/// A number that changes whenever a job starts, moves or ends (the status bar redraws then).
117pub fn generation() -> u64 {
118    GENERATION.load(Ordering::Relaxed)
119}
120
121/// A running job, held by the work doing it.
122pub struct JobHandle {
123    /// The job's id.
124    pub id: u64,
125    /// Set when someone asked it to stop.
126    pub cancel: Arc<AtomicBool>,
127}
128
129impl JobHandle {
130    /// Says what it is doing now, and how far.
131    pub fn set(&self, name: &str, progress: Option<f32>) {
132        with(|j| {
133            if let Some(o) = j.list.iter_mut().find(|o| o.id == self.id) {
134                o.name = name.to_string();
135                o.progress = progress;
136            }
137        });
138    }
139
140    /// Whether someone asked it to stop.
141    pub fn cancelled(&self) -> bool {
142        self.cancel.load(Ordering::Relaxed)
143    }
144
145    /// Leaves the job's outcome (what it found, or why it failed), read by [`outcome`] once the
146    /// job ends. The last call wins.
147    pub fn finish(&self, outcome: Result<serde_json::Value, String>) {
148        let id = self.id;
149        with(|j| {
150            j.outcomes.retain(|(i, _)| *i != id);
151            j.outcomes.push_back((id, outcome));
152            while j.outcomes.len() > KEPT_OUTCOMES {
153                j.outcomes.pop_front();
154            }
155        });
156    }
157}
158
159impl Drop for JobHandle {
160    fn drop(&mut self) {
161        let id = self.id;
162        let stopped = self.cancelled();
163        with(|j| {
164            j.list.retain(|o| o.id != id);
165            // A job that left no outcome: it was stopped, or it just ended
166            if !j.outcomes.iter().any(|(i, _)| *i == id) {
167                j.outcomes.push_back((id, if stopped { Err("stopped".to_string()) } else { Ok(serde_json::Value::Null) }));
168                while j.outcomes.len() > KEPT_OUTCOMES {
169                    j.outcomes.pop_front();
170                }
171            }
172        });
173        crate::activity::job_ended(id);
174    }
175}
176
177#[cfg(test)]
178mod tests {
179    #[test]
180    fn one_gpu_job_at_a_time_and_a_job_ends_with_its_handle() {
181        let g0 = super::generation();
182        let a = super::start("test-a", "first", true).unwrap();
183        assert!(super::start("test-b", "second", true).is_err(), "a second GPU job is refused");
184        let cpu = super::start("test-c", "on the CPU", false).unwrap();
185        a.set("first: half", Some(0.5));
186        assert_eq!(super::of("test-a").unwrap().progress, Some(0.5));
187        super::cancel("test-a");
188        assert!(a.cancelled());
189        super::cancel_id(cpu.id);
190        assert!(cpu.cancelled());
191        drop(a);
192        assert!(super::of("test-a").is_none());
193        assert!(super::start("test-b", "second", true).is_ok(), "free again");
194        drop(cpu);
195        assert!(super::generation() > g0);
196    }
197
198    #[test]
199    fn a_job_leaves_its_outcome() {
200        let a = super::start("test-outcome", "counting", false).unwrap();
201        let id = a.id;
202        assert!(super::is_running(id) && super::outcome(id).is_none());
203        assert_eq!(super::job_of(&super::job_reply(id)), Some(id));
204        a.finish(Ok(serde_json::json!({"cells": 12})));
205        drop(a);
206        assert!(!super::is_running(id));
207        assert_eq!(super::outcome(id), Some(Ok(serde_json::json!({"cells": 12}))));
208        let b = super::start("test-outcome", "stopped", false).unwrap();
209        b.cancel.store(true, std::sync::atomic::Ordering::Relaxed);
210        let id = b.id;
211        drop(b);
212        assert_eq!(super::outcome(id), Some(Err("stopped".to_string())));
213    }
214}
kitchen_shell/jobs.rs

kitchen_shell/
jobs.rs
