superfluid_daemon/
batches.rs

1use std::path::PathBuf;
2use std::sync::atomic::{AtomicU64, Ordering};
3use std::sync::Arc;
4
5const BATCH_POLICY: crate::inference::RequestPolicy = crate::inference::RequestPolicy {
6    limits: BATCH_LIMITS,
7    sampling: crate::inference::SamplingOverrides {
8        temperature: None,
9        top_p: None,
10        top_k: None,
11        min_p: None,
12        repeat_penalty: None,
13        spec_max_temperature: None,
14        fallback: None,
15    },
16    qos: crate::inference::RequestQos::AGENT,
17};
18
19const BATCH_LIMITS: crate::inference::TokenLimits = crate::inference::TokenLimits {
20    default_max_tokens: Some(crate::inference::DEFAULT_MAX_TOKENS),
21    max_context: crate::inference::DEFAULT_MAX_CONTEXT,
22};
23
24use serde::{Deserialize, Serialize};
25
26use crate::Daemon;
27
28#[derive(Debug, Clone, Serialize, Deserialize)]
29pub struct RequestCounts {
30    pub total: u64,
31    pub completed: u64,
32    pub failed: u64,
33}
34
35#[derive(Debug, Clone, Serialize, Deserialize)]
36pub struct Batch {
37    pub id: String,
38    #[serde(default = "batch_object")]
39    pub object: String,
40    pub endpoint: String,
41    pub input_file_id: String,
42    pub completion_window: String,
43    pub status: String,
44    pub output_file_id: Option<String>,
45    #[serde(default)]
46    pub error_file_id: String,
47    pub created_at: u64,
48    #[serde(default)]
49    pub completed_at: u64,
50    #[serde(default)]
51    pub seq: u64,
52    pub request_counts: RequestCounts,
53    #[serde(default, skip_serializing_if = "Option::is_none")]
54    pub owner: Option<String>,
55}
56
57fn batch_object() -> String {
58    "batch".to_string()
59}
60fn now_secs() -> u64 {
61    std::time::SystemTime::now()
62        .duration_since(std::time::UNIX_EPOCH)
63        .map(|d| d.as_secs())
64        .unwrap_or(0)
65}
66
67/// Batches run on at most this many threads; the rest wait in `validating`, in order.
68pub const MAX_RUNNING_BATCHES: usize = 4;
69pub const MAX_QUEUED_BATCHES: usize = 256;
70
71type Job = Box<dyn FnOnce() + Send>;
72
73#[derive(Default)]
74struct Runner {
75    running: usize,
76    queued: std::collections::VecDeque<Job>,
77}
78
79#[derive(Debug)]
80pub struct BatchesBusy;
81
82static STATUS: std::sync::Mutex<()> = std::sync::Mutex::new(());
83
84pub struct BatchStore {
85    dir: PathBuf,
86    runner: std::sync::Mutex<Runner>,
87}
88
89impl BatchStore {
90    pub fn new(dir: PathBuf) -> BatchStore {
91        let _ = std::fs::create_dir_all(&dir);
92        BatchStore { dir, runner: std::sync::Mutex::new(Runner::default()) }
93    }
94
95    /// Creates a batch and runs `job` with its id on a batch worker, or queues it behind the
96    /// running ones. Refuses, creating nothing, when the queue is full.
97    pub fn start(
98        self: &Arc<Self>,
99        endpoint: &str,
100        input_file_id: &str,
101        completion_window: &str,
102        total: u64,
103        owner: Option<&str>,
104        job: impl FnOnce(String) + Send + 'static,
105    ) -> Result<Batch, BatchesBusy> {
106        let mut r = self.runner.lock().unwrap_or_else(|e| e.into_inner());
107        if r.queued.len() >= MAX_QUEUED_BATCHES {
108            return Err(BatchesBusy);
109        }
110        let batch = self.create(endpoint, input_file_id, completion_window, total, owner);
111        let id = batch.id.clone();
112        let job: Job = Box::new(move || job(id));
113        if r.running >= MAX_RUNNING_BATCHES {
114            r.queued.push_back(job);
115            return Ok(batch);
116        }
117        r.running += 1;
118        drop(r);
119        let store = Arc::clone(self);
120        let spawned = std::thread::Builder::new()
121            .name("superfluid-batch".into())
122            .spawn(move || store.work(job));
123        if spawned.is_err() {
124            self.runner.lock().unwrap_or_else(|e| e.into_inner()).running -= 1;
125            let mut b = batch.clone();
126            b.status = "failed".to_string();
127            b.completed_at = now_secs();
128            self.write(&b);
129        }
130        Ok(batch)
131    }
132
133    fn work(&self, mut job: Job) {
134        loop {
135            if std::panic::catch_unwind(std::panic::AssertUnwindSafe(job)).is_err() {
136                eprintln!("superfluid: a batch worker panicked; its batch is left as it was");
137            }
138            let mut r = self.runner.lock().unwrap_or_else(|e| e.into_inner());
139            match r.queued.pop_front() {
140                Some(next) => job = next,
141                None => {
142                    r.running -= 1;
143                    return;
144                }
145            }
146        }
147    }
148
149    fn new_id() -> String {
150        static N: AtomicU64 = AtomicU64::new(0);
151        let nanos = std::time::SystemTime::now()
152            .duration_since(std::time::UNIX_EPOCH)
153            .map(|d| d.as_nanos() as u64)
154            .unwrap_or(0);
155        format!("batch-{:x}{:x}", nanos, N.fetch_add(1, Ordering::Relaxed))
156    }
157
158    fn valid_id(id: &str) -> bool {
159        id.starts_with("batch-") && id.len() > 6 && id[6..].chars().all(|c| c.is_ascii_hexdigit())
160    }
161    fn path(&self, id: &str) -> PathBuf {
162        self.dir.join(format!("{id}.json"))
163    }
164
165    fn write(&self, b: &Batch) {
166        if let Ok(j) = serde_json::to_vec(b) {
167            let path = self.path(&b.id);
168            let tmp = path.with_extension(format!("json.tmp.{:?}", std::thread::current().id()).replace(['(', ')'], ""));
169            if std::fs::write(&tmp, j).is_ok() && std::fs::rename(&tmp, &path).is_err() {
170                let _ = std::fs::remove_file(&tmp);
171            }
172        }
173    }
174
175    pub fn create(
176        &self,
177        endpoint: &str,
178        input_file_id: &str,
179        completion_window: &str,
180        total: u64,
181        owner: Option<&str>,
182    ) -> Batch {
183        let b = Batch {
184            id: Self::new_id(),
185            object: batch_object(),
186            endpoint: endpoint.to_string(),
187            input_file_id: input_file_id.to_string(),
188            completion_window: completion_window.to_string(),
189            status: "validating".to_string(),
190            output_file_id: None,
191            error_file_id: String::new(),
192            created_at: now_secs(),
193            completed_at: 0,
194            seq: {
195                static SEQ: AtomicU64 = AtomicU64::new(1);
196                SEQ.fetch_add(1, Ordering::Relaxed)
197            },
198            request_counts: RequestCounts { total, completed: 0, failed: 0 },
199            owner: owner.map(str::to_string),
200        };
201        self.write(&b);
202        b
203    }
204
205    pub fn get(&self, id: &str) -> Option<Batch> {
206        if !Self::valid_id(id) {
207            return None;
208        }
209        serde_json::from_slice(&std::fs::read(self.path(id)).ok()?).ok()
210    }
211
212    pub fn list(&self) -> Vec<Batch> {
213        let mut out = Vec::new();
214        if let Ok(rd) = std::fs::read_dir(&self.dir) {
215            for e in rd.flatten() {
216                if e.path().extension().and_then(|x| x.to_str()) == Some("json") {
217                    if let Ok(b) = std::fs::read(e.path()) {
218                        if let Ok(m) = serde_json::from_slice::<Batch>(&b) {
219                            out.push(m);
220                        }
221                    }
222                }
223            }
224        }
225        out.sort_by_key(|b| std::cmp::Reverse(b.created_at));
226        out
227    }
228
229    pub fn cancel(&self, id: &str) -> Option<(Batch, bool)> {
230        let _one = STATUS.lock().unwrap_or_else(|e| e.into_inner());
231        let mut b = self.get(id)?;
232        let stops = b.status == "in_progress" || b.status == "validating";
233        if b.status == "validating" {
234            b.status = "cancelled".to_string();
235            b.completed_at = now_secs();
236            self.write(&b);
237        } else if stops {
238            b.status = "cancelling".to_string();
239            self.write(&b);
240        }
241        Some((b, stops))
242    }
243
244    fn begin(&self, id: &str) -> Option<Batch> {
245        let _one = STATUS.lock().unwrap_or_else(|e| e.into_inner());
246        let mut b = self.get(id).filter(|b| b.status == "validating")?;
247        b.status = "in_progress".to_string();
248        self.write(&b);
249        Some(b)
250    }
251}
252
253#[derive(Deserialize)]
254struct BatchLine {
255    #[serde(default)]
256    custom_id: String,
257    #[serde(default)]
258    url: String,
259    body: serde_json::Value,
260}
261
262pub fn process(
263    daemon: Arc<Daemon>,
264    store: Arc<BatchStore>,
265    model_name: String,
266    batch_id: String,
267    qos: crate::inference::RequestQos,
268    pace: Option<Arc<crate::keypolicy::KeyPolicy>>,
269) {
270    let Some(mut batch) = store.begin(&batch_id) else {
271        return;
272    };
273
274    let input = match daemon.files().content(&batch.input_file_id) {
275        Some(bytes) => bytes,
276        None => {
277            batch.status = "failed".to_string();
278            batch.completed_at = now_secs();
279            store.write(&batch);
280            return;
281        }
282    };
283    let text = String::from_utf8_lossy(&input);
284    let mut out_lines = String::new();
285    let (mut completed, mut failed) = (0u64, 0u64);
286
287    for (i, line) in text.lines().enumerate() {
288        let line = line.trim();
289        if line.is_empty() {
290            continue;
291        }
292        if store.get(&batch_id).map(|b| b.status) == Some("cancelling".to_string()) {
293            batch.status = "cancelled".to_string();
294            batch.completed_at = now_secs();
295            batch.request_counts = RequestCounts { total: batch.request_counts.total, completed, failed };
296            store.write(&batch);
297            return;
298        }
299        if let Some(key) = &pace {
300            let cancelled = || store.get(&batch_id).map(|b| b.status) == Some("cancelling".to_string());
301            if !key.pace(cancelled) {
302                batch.status = "cancelled".to_string();
303                batch.completed_at = now_secs();
304                batch.request_counts = RequestCounts { total: batch.request_counts.total, completed, failed };
305                store.write(&batch);
306                return;
307            }
308        }
309        let parsed: Result<BatchLine, _> = serde_json::from_str(line);
310        let (status_code, body_val): (u16, serde_json::Value) = match parsed {
311            Ok(bl) => run_line(&daemon, &model_name, &bl.url, &bl.body, qos),
312            Err(e) => (400, serde_json::json!({"error": {"message": format!("bad input line: {e}")}})),
313        };
314        let custom_id = serde_json::from_str::<BatchLine>(line)
315            .map(|b| b.custom_id)
316            .unwrap_or_else(|_| format!("line-{i}"));
317        if (200..300).contains(&status_code) {
318            completed += 1;
319        } else {
320            failed += 1;
321        }
322        let row = serde_json::json!({
323            "id": format!("batch_req_{i}"),
324            "custom_id": custom_id,
325            "response": {"status_code": status_code, "body": body_val},
326            "error": serde_json::Value::Null,
327        });
328        out_lines.push_str(&row.to_string());
329        out_lines.push('\n');
330    }
331
332    let output_file_id = daemon
333        .files()
334        .put("batch_output.jsonl", "batch_output", out_lines.as_bytes(), batch.owner.as_deref())
335        .ok()
336        .map(|m| m.id);
337    batch.output_file_id = output_file_id;
338    batch.request_counts = RequestCounts { total: batch.request_counts.total, completed, failed };
339    batch.status = "completed".to_string();
340    batch.completed_at = now_secs();
341    store.write(&batch);
342}
343
344fn run_line(
345    daemon: &Daemon,
346    model_name: &str,
347    url: &str,
348    body: &serde_json::Value,
349    qos: crate::inference::RequestQos,
350) -> (u16, serde_json::Value) {
351    match url {
352        "/v1/chat/completions" => match serde_json::from_value::<crate::inference::ChatRequest>(body.clone()) {
353            Ok(req) => match crate::openai::chat_once(
354                daemon,
355                &req,
356                model_name,
357                crate::inference::RequestPolicy { qos, ..BATCH_POLICY },
358            ) {
359                Ok(v) => (200, v),
360                Err(e) => (500, serde_json::json!({"error": {"message": e.to_string()}})),
361            },
362            Err(e) => (400, serde_json::json!({"error": {"message": format!("bad chat body: {e}")}})),
363        },
364        "/v1/embeddings" => {
365            let input = body.get("input").and_then(|v| v.as_str()).unwrap_or("");
366            match daemon.embed_as(input, Some(qos.class)) {
367                Ok(v) => (
368                    200,
369                    serde_json::json!({"object": "list", "data": [{"object":"embedding","index":0,"embedding": v}], "model": model_name}),
370                ),
371                Err(e) => (500, serde_json::json!({"error": {"message": e.to_string()}})),
372            }
373        }
374        other => (400, serde_json::json!({"error": {"message": format!("unsupported batch url: {other}")}})),
375    }
376}
377
378pub fn batches_dir_for(sessions_dir: &std::path::Path) -> PathBuf {
379    sessions_dir.join("batches")
380}
381
382#[cfg(test)]
383mod tests {
384    use super::*;
385
386    #[test]
387    fn batches_run_on_a_bounded_pool_and_a_full_queue_refuses() {
388        let dir = std::env::temp_dir().join(format!("superfluid-batch-pool-{}", std::process::id()));
389        let store = Arc::new(BatchStore::new(dir.clone()));
390        let (release, gate) = std::sync::mpsc::channel::<()>();
391        let gate = Arc::new(std::sync::Mutex::new(gate));
392        let live = Arc::new(AtomicU64::new(0));
393        let peak = Arc::new(AtomicU64::new(0));
394        let done = Arc::new(AtomicU64::new(0));
395        let total = MAX_RUNNING_BATCHES + MAX_QUEUED_BATCHES;
396        for _ in 0..total {
397            let (gate, live, peak, done) = (Arc::clone(&gate), Arc::clone(&live), Arc::clone(&peak), Arc::clone(&done));
398            store
399                .start("/v1/chat/completions", "file-x", "24h", 1, None, move |_id| {
400                    peak.fetch_max(live.fetch_add(1, Ordering::SeqCst) + 1, Ordering::SeqCst);
401                    let _ = gate.lock().unwrap().recv();
402                    live.fetch_sub(1, Ordering::SeqCst);
403                    done.fetch_add(1, Ordering::SeqCst);
404                })
405                .expect("room");
406        }
407        let refused = store.start("/v1/chat/completions", "file-x", "24h", 1, None, |_| {});
408        assert!(refused.is_err(), "a full queue refuses");
409        assert_eq!(store.list().len(), total, "a refused batch is not created");
410        for _ in 0..total {
411            release.send(()).unwrap();
412        }
413        for _ in 0..400 {
414            if done.load(Ordering::SeqCst) == total as u64 {
415                break;
416            }
417            std::thread::sleep(std::time::Duration::from_millis(5));
418        }
419        assert_eq!(done.load(Ordering::SeqCst), total as u64);
420        assert!(peak.load(Ordering::SeqCst) <= MAX_RUNNING_BATCHES as u64, "peak {}", peak.load(Ordering::SeqCst));
421        assert!(store.start("/v1/chat/completions", "file-x", "24h", 1, None, |_| {}).is_ok(), "room again");
422        let _ = std::fs::remove_dir_all(dir);
423    }
424
425    #[test]
426    fn a_batch_cancelled_while_queued_never_starts() {
427        let dir = std::env::temp_dir().join(format!("superfluid-batch-queued-cancel-{}", std::process::id()));
428        let store = BatchStore::new(dir.clone());
429        let b = store.create("/v1/chat/completions", "file-x", "24h", 1, None);
430        let (cancelled, stops) = store.cancel(&b.id).unwrap();
431        assert!(stops);
432        assert_eq!(cancelled.status, "cancelled");
433        assert!(store.begin(&b.id).is_none(), "a worker reaching it later does not run it");
434        assert_eq!(store.get(&b.id).unwrap().status, "cancelled");
435        let _ = std::fs::remove_dir_all(dir);
436    }
437}
superfluid_daemon/batches.rs

superfluid_daemon/
batches.rs
