superfluid_worker/
process.rs

1//! The worker as its own process.
2
3use std::os::unix::io::FromRawFd;
4use std::os::unix::net::UnixStream;
5use std::path::PathBuf;
6
7use superfluid_proto::linkw;
8
9#[derive(Debug, Clone, Default)]
10pub struct WorkerArgs {
11    pub frames_fd: Option<i32>,
12    pub fd_channel_fd: Option<i32>,
13    pub model: Option<PathBuf>,
14    pub max_context: i32,
15    pub max_batch: u32,
16    pub kv_bits: i32,
17    pub venv: Option<PathBuf>,
18    pub extra: Vec<(String, String)>,
19}
20
21impl WorkerArgs {
22    pub fn parse(args: &[String]) -> Result<WorkerArgs, String> {
23        let mut a = WorkerArgs { max_context: 4096, max_batch: 8, ..Default::default() };
24        let mut it = args.iter();
25        while let Some(flag) = it.next() {
26            if !flag.starts_with("--") {
27                return Err(format!("unexpected argument {flag:?}"));
28            }
29            let value = it.next().ok_or_else(|| format!("{flag} needs a value"))?;
30            let num = |what: &str| format!("{flag}: {what} expected, got {value:?}");
31            match flag.as_str() {
32                "--frames-fd" => a.frames_fd = Some(value.parse().map_err(|_| num("a descriptor"))?),
33                "--fd-channel-fd" => a.fd_channel_fd = Some(value.parse().map_err(|_| num("a descriptor"))?),
34                "--model" => a.model = Some(PathBuf::from(value)),
35                "--max-context" => a.max_context = value.parse().map_err(|_| num("a token count"))?,
36                "--max-batch" => a.max_batch = value.parse().map_err(|_| num("a lane count"))?,
37                "--kv-bits" => a.kv_bits = value.parse().map_err(|_| num("a bit width"))?,
38                "--venv" => a.venv = Some(PathBuf::from(value)),
39                other => a.extra.push((other.to_string(), value.clone())),
40            }
41        }
42        Ok(a)
43    }
44
45    pub fn extra(&self, flag: &str) -> Option<&str> {
46        self.extra.iter().rev().find(|(k, _)| k == flag).map(|(_, v)| v.as_str())
47    }
48
49    pub fn refuse_unread(&self, reads: &[&str]) -> Result<(), String> {
50        let unread = |flag: &str, set: bool| set && !reads.contains(&flag);
51        if unread("--kv-bits", self.kv_bits != 0) {
52            return Err("--kv-bits: this runtime has no such knob".into());
53        }
54        if unread("--venv", self.venv.is_some()) {
55            return Err("--venv: this runtime has no Python environment".into());
56        }
57        match self.extra.iter().find(|(k, _)| !reads.contains(&k.as_str())) {
58            Some((k, _)) => Err(format!("unknown flag {k}")),
59            None => Ok(()),
60        }
61    }
62
63    pub fn streams(&mut self) -> Result<(UnixStream, UnixStream), String> {
64        let (Some(frames), Some(fds)) = (self.frames_fd.take(), self.fd_channel_fd.take()) else {
65            return Err("--frames-fd and --fd-channel-fd are required to serve".into());
66        };
67        // SAFETY: the parent dup2'd the Link W socketpair ends onto exactly
68        // these descriptors before exec; this process owns them from here,
69        // and `take` leaves no second owner.
70        Ok(unsafe { (UnixStream::from_raw_fd(frames), UnixStream::from_raw_fd(fds)) })
71    }
72}
73
74pub fn ring_specs_for(ctx: u32, vocab: u32) -> Vec<linkw::RingSpec> {
75    vec![
76        linkw::RingSpec {
77            ring_id: superfluid_shm::TOKEN_RING_IN,
78            kind: linkw::RingKind::Tokens,
79            slot_bytes: superfluid_shm::SharedRing::slot_bytes_for(ctx.max(512) * 4),
80            slots: 64,
81        },
82        linkw::RingSpec {
83            ring_id: superfluid_shm::TOKEN_RING_OUT,
84            kind: linkw::RingKind::Tokens,
85            slot_bytes: 16 + 1024 * 4,
86            slots: crate::server::TOKEN_OUT_SLOTS,
87        },
88        linkw::RingSpec {
89            ring_id: superfluid_shm::LOGITS_RING,
90            kind: linkw::RingKind::Logits,
91            slot_bytes: superfluid_shm::SharedRing::slot_bytes_for(vocab * 4),
92            slots: 8,
93        },
94    ]
95}
96
97#[cfg(test)]
98mod tests {
99    use super::*;
100
101    #[test]
102    fn an_odd_context_or_vocab_still_makes_rings_that_attach() {
103        for spec in ring_specs_for(4097, 50257) {
104            let ring = superfluid_shm::SharedRing::create(spec.ring_id, spec.slots, spec.slot_bytes, superfluid_shm::RingRole::Writer);
105            assert!(ring.is_ok(), "{spec:?}: {:?}", ring.err());
106        }
107        assert_eq!(ring_specs_for(4097, 50257)[0].slot_bytes, 16 + 4098 * 4);
108    }
109
110    fn args(s: &[&str]) -> Vec<String> {
111        s.iter().map(|x| x.to_string()).collect()
112    }
113
114    #[test]
115    fn the_serve_line_parses_and_keeps_what_it_does_not_own() {
116        let a = WorkerArgs::parse(&args(&[
117            "--frames-fd", "3", "--fd-channel-fd", "4", "--model", "/m.gguf", "--max-context", "8192",
118            "--max-batch", "4", "--engine", "native", "--venv", "/env",
119        ]))
120        .unwrap();
121        assert_eq!((a.frames_fd, a.fd_channel_fd), (Some(3), Some(4)));
122        assert_eq!(a.model.as_deref(), Some(std::path::Path::new("/m.gguf")));
123        assert_eq!((a.max_context, a.max_batch, a.kv_bits), (8192, 4, 0));
124        assert_eq!(a.extra("--engine"), Some("native"));
125        assert_eq!(a.venv.as_deref(), Some(std::path::Path::new("/env")));
126    }
127
128    #[test]
129    fn a_bad_line_says_what_is_wrong() {
130        assert_eq!(WorkerArgs::parse(&args(&["--max-batch"])).unwrap_err(), "--max-batch needs a value");
131        assert!(WorkerArgs::parse(&args(&["--max-batch", "many"])).unwrap_err().starts_with("--max-batch: a lane count"));
132        assert!(WorkerArgs::parse(&args(&["serve"])).unwrap_err().contains("unexpected argument"));
133        assert!(WorkerArgs::parse(&[]).unwrap().streams().is_err());
134    }
135
136    #[test]
137    fn a_flag_the_runtime_does_not_read_is_refused() {
138        let a = WorkerArgs::parse(&args(&["--kv-bits", "8", "--venv", "/env", "--engine", "mlx"])).unwrap();
139        assert_eq!(a.refuse_unread(&["--kv-bits", "--venv", "--engine"]), Ok(()));
140        assert_eq!(a.refuse_unread(&["--venv", "--engine"]).unwrap_err(), "--kv-bits: this runtime has no such knob");
141        assert_eq!(a.refuse_unread(&["--kv-bits", "--engine"]).unwrap_err(), "--venv: this runtime has no Python environment");
142        assert_eq!(a.refuse_unread(&["--kv-bits", "--venv"]).unwrap_err(), "unknown flag --engine");
143        assert_eq!(WorkerArgs::parse(&args(&["--model", "/m"])).unwrap().refuse_unread(&[]), Ok(()));
144    }
145}
superfluid_worker/process.rs

superfluid_worker/
process.rs
