xil_workers/
lib.rs

1//! Client side of the newline-delimited JSON protocol spoken by the Python
2//! ML workers (`chatterbox_turbo_worker.py`, `whisper_worker.py`,
3//! `mmaudio_worker.py`). The workers stay Python, in their own venvs; Rust
4//! only starts them, writes one JSON request per line and reads one JSON
5//! response per line.
6//!
7//! Startup differs by caller, so both reads are exposed: `stem-verify`
8//! parses the very first line as the ready message, while the SFX and TTS
9//! clients skip model-loading noise until a `{"ready": true}` line appears.
10
11use std::io::{self, BufRead, BufReader, Write};
12use std::path::{Path, PathBuf};
13use std::process::{Child, ChildStdin, ChildStdout, Command, Stdio};
14
15use serde_json::Value;
16use xil_core::pyjson::{dumps, Style};
17
18/// A running worker process.
19pub struct Worker {
20    child: Child,
21    stdin: Option<ChildStdin>,
22    stdout: BufReader<ChildStdout>,
23}
24
25impl Worker {
26    /// `subprocess.Popen([python, script, *args], stdin=PIPE, stdout=PIPE)`
27    /// with stderr inherited, as every Python caller leaves it.
28    pub fn spawn(python: &Path, script: &Path, args: &[String]) -> io::Result<Worker> {
29        let mut child = Command::new(python)
30            .arg(script)
31            .args(args)
32            .stdin(Stdio::piped())
33            .stdout(Stdio::piped())
34            .stderr(Stdio::inherit())
35            .spawn()?;
36        let stdin = child.stdin.take();
37        let stdout = BufReader::new(child.stdout.take().expect("piped stdout"));
38        Ok(Worker {
39            child,
40            stdin,
41            stdout,
42        })
43    }
44
45    /// `proc.stdout.readline()`: the next line including its newline, or
46    /// an empty string at end of stream.
47    pub fn read_line(&mut self) -> io::Result<String> {
48        let mut buf = Vec::new();
49        self.stdout.read_until(b'\n', &mut buf)?;
50        Ok(String::from_utf8_lossy(&buf).into_owned())
51    }
52
53    /// Read lines until one parses as JSON with a truthy `ready`, returning
54    /// it. Non-JSON lines and other messages go to `on_noise`. `None` means
55    /// the worker exited first.
56    pub fn wait_ready(&mut self, mut on_noise: impl FnMut(&str)) -> io::Result<Option<Value>> {
57        loop {
58            let raw = self.read_line()?;
59            if raw.is_empty() {
60                return Ok(None);
61            }
62            let line = raw.trim();
63            if line.is_empty() {
64                continue;
65            }
66            match serde_json::from_str::<Value>(line) {
67                Ok(msg) if is_truthy(msg.get("ready")) => return Ok(Some(msg)),
68                _ => on_noise(line),
69            }
70        }
71    }
72
73    /// Write `json.dumps(request) + "\n"` and flush.
74    pub fn send(&mut self, request: &Value) -> io::Result<()> {
75        let stdin = self
76            .stdin
77            .as_mut()
78            .ok_or_else(|| io::Error::new(io::ErrorKind::BrokenPipe, "worker stdin closed"))?;
79        let mut line = dumps(request, Style::COMPACT);
80        line.push('\n');
81        stdin.write_all(line.as_bytes())?;
82        stdin.flush()
83    }
84
85    /// Close stdin (the worker's loop ends) and wait for it to exit.
86    pub fn close(mut self) -> io::Result<()> {
87        drop(self.stdin.take());
88        self.child.wait().map(|_| ())
89    }
90}
91
92/// Python truthiness for the `ready` field.
93pub fn is_truthy(v: Option<&Value>) -> bool {
94    match v {
95        None | Some(Value::Null) | Some(Value::Bool(false)) => false,
96        Some(Value::Number(n)) => n.as_f64() != Some(0.0),
97        Some(Value::String(s)) => !s.is_empty(),
98        Some(Value::Array(a)) => !a.is_empty(),
99        Some(Value::Object(o)) => !o.is_empty(),
100        Some(Value::Bool(true)) => true,
101    }
102}
103
104/// `<package_dir>/<name>` — where a worker script lives next to the Python
105/// package that ships it.
106pub fn worker_script(package_dir: &Path, name: &str) -> PathBuf {
107    package_dir.join(name)
108}
109
110#[cfg(all(test, unix))]
111mod tests {
112    use super::*;
113    use std::os::unix::fs::PermissionsExt;
114
115    /// A shell "python" that ignores the script argument, prints noise and
116    /// a ready line, then echoes each request's `n` back doubled.
117    fn fake_python(dir: &Path) -> PathBuf {
118        let p = dir.join("fake-python");
119        std::fs::write(
120            &p,
121            "#!/bin/sh\necho 'loading weights...'\necho '{\"ready\": true, \"device\": \"cpu\"}'\nwhile read line; do echo '{\"done\": true}'; done\n",
122        )
123        .unwrap();
124        std::fs::set_permissions(&p, std::fs::Permissions::from_mode(0o755)).unwrap();
125        p
126    }
127
128    #[test]
129    fn handshake_request_close() {
130        let tmp = tempfile::tempdir().unwrap();
131        let py = fake_python(tmp.path());
132        let mut w = Worker::spawn(&py, Path::new("worker.py"), &["cpu".into()]).unwrap();
133        let mut noise = Vec::new();
134        let ready = w
135            .wait_ready(|l| noise.push(l.to_string()))
136            .unwrap()
137            .unwrap();
138        assert_eq!(ready["device"], "cpu");
139        assert_eq!(noise, vec!["loading weights..."]);
140        w.send(&serde_json::json!({"prompt": "door"})).unwrap();
141        let resp: Value = serde_json::from_str(&w.read_line().unwrap()).unwrap();
142        assert_eq!(resp["done"], true);
143        w.close().unwrap();
144    }
145}
xil_workers/lib.rs

xil_workers/
lib.rs
