use std::collections::HashMap;
use std::ffi::CString;
use std::io::{Read, Write};
use std::os::unix::net::UnixStream;
use std::time::Duration;
use superfluid_abi::{
op_state, Artifact, CapabilityReq, StrategyRegistration, TapSpec,
TokenRange, RecordArena, Status,
};
use superfluid_engine::{Engine, StageOutput};
use superfluid_proto::envelope::FrameClass;
use superfluid_proto::handshake::{negotiate, Hello, HelloAckW};
use superfluid_proto::linkw::{self, msg_type};
use superfluid_proto::{encode_frame, Frame, FrameDecoder};
use superfluid_shm::{
fdpass, ControlSegment, RingRole, SharedRing, SharedRings, ShmSegment, CONTROL_TAG, LOGITS_RING,
TOKEN_RING_IN, TOKEN_RING_OUT,
};
use crate::materialize::{
desc_to_msg, events_to_msg, match_result_to_msg, materialize_plan, token_ref_from_msg,
};
use crate::WorkerError;
pub const PROTO_VERSION: u16 = 8;
pub const TOKEN_OUT_SLOTS: u32 = 256;
pub fn default_ring_specs() -> Vec<linkw::RingSpec> {
vec![
linkw::RingSpec {
ring_id: TOKEN_RING_IN,
kind: linkw::RingKind::Tokens,
slot_bytes: 16 + 4096 * 4,
slots: 64,
},
linkw::RingSpec {
ring_id: TOKEN_RING_OUT,
kind: linkw::RingKind::Tokens,
slot_bytes: 16 + 1024 * 4,
slots: TOKEN_OUT_SLOTS,
},
linkw::RingSpec {
ring_id: LOGITS_RING,
kind: linkw::RingKind::Logits,
slot_bytes: 16 + 64 * 4,
slots: 32,
},
]
}
pub struct WorkerConfig {
pub engine_bundle_hash: [u8; 32],
}
enum OpBinding {
Export { buf_id: u64 },
Import,
}
pub struct WorkerServer<E: Engine> {
engine: E,
cfg: WorkerConfig,
frames: UnixStream,
fd_channel: UnixStream,
seq: u64,
rings: Option<SharedRings>,
pending_rings: HashMap<u32, SharedRing>,
pending_control: Option<ControlSegment>,
ring_specs: Vec<linkw::RingSpec>,
bufs: HashMap<u64, ShmSegment>,
pending_buf_release: std::collections::HashSet<u64>,
open_ops: HashMap<u64, OpBinding>,
op_window: usize,
decoder: FrameDecoder<fn(u16) -> bool>,
pending: std::collections::VecDeque<Frame>,
handshaken: bool,
chosen_version: u16,
last_mem: linkw::MemCountersMsg,
}
struct EmitSink<'a> {
frames: &'a mut UnixStream,
seq: &'a mut u64,
ok: bool,
}
unsafe extern "C" fn emit_tramp(
user: *mut std::ffi::c_void,
lane_tag: u64,
tokens: *const u32,
n_tokens: u32,
) {
if user.is_null() || n_tokens == 0 {
return;
}
let sink = unsafe { &mut *(user as *mut EmitSink) };
if !sink.ok {
return;
}
if tokens.is_null() {
sink.ok = false;
return;
}
let toks = unsafe { std::slice::from_raw_parts(tokens, n_tokens as usize) };
let sent = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
let msg = linkw::TickEmitMsg { lane_tag, tokens: toks.to_vec() };
postcard::to_stdvec(&msg).is_ok_and(|p| {
send_frame(sink.frames, sink.seq, msg_type::TICK_EMIT, 0, &p).is_ok()
})
}));
sink.ok = sent.unwrap_or(false);
}
fn send_frame(
frames: &mut UnixStream,
seq: &mut u64,
msg_type: u16,
correlation: u64,
payload: &[u8],
) -> Result<(), WorkerError> {
let class = if msg_type & 0x8000 != 0 {
FrameClass::Droppable
} else {
FrameClass::Required
};
let bytes = encode_frame(PROTO_VERSION, msg_type, *seq, correlation, class, payload)?;
*seq += 1;
frames.write_all(&bytes)?;
Ok(())
}
fn poll_transcribe_cancel(
frames: &mut UnixStream,
decoder: &mut FrameDecoder<fn(u16) -> bool>,
pending: &mut std::collections::VecDeque<Frame>,
) -> bool {
if frames.set_nonblocking(true).is_err() {
return false;
}
let mut buf = [0u8; 64 * 1024];
let read = frames.read(&mut buf);
if frames.set_nonblocking(false).is_err() {
return true;
}
match read {
Ok(0) => return true,
Ok(n) => decoder.feed(&buf[..n]),
Err(e)
if e.kind() == std::io::ErrorKind::WouldBlock
|| e.kind() == std::io::ErrorKind::TimedOut => {}
Err(_) => return true,
}
let mut cancel = false;
loop {
match decoder.decode_next() {
Ok(Some(frame)) => {
if frame.msg_type == msg_type::TRANSCRIBE_CANCEL {
cancel = true;
} else {
pending.push_back(frame);
}
}
Ok(None) => return cancel,
Err(_) => return true,
}
}
}
impl<E: Engine> WorkerServer<E> {
pub fn new(engine: E, cfg: WorkerConfig, frames: UnixStream, fd_channel: UnixStream) -> Self {
WorkerServer {
engine,
cfg,
frames,
fd_channel,
seq: 1,
rings: None,
pending_rings: HashMap::new(),
pending_control: None,
ring_specs: default_ring_specs(),
bufs: HashMap::new(),
pending_buf_release: std::collections::HashSet::new(),
open_ops: HashMap::new(),
op_window: 16,
decoder: FrameDecoder::new(linkw::is_known_type as fn(u16) -> bool),
pending: std::collections::VecDeque::new(),
handshaken: false,
chosen_version: 0,
last_mem: zero_mem(),
}
}
pub fn with_ring_specs(mut self, specs: Vec<linkw::RingSpec>) -> Self {
self.ring_specs = specs;
self
}
pub fn serve(mut self) -> Result<(), WorkerError> {
self.frames
.set_read_timeout(Some(Duration::from_millis(5)))?;
let mut buf = [0u8; 64 * 1024];
loop {
while let Some(frame) = self.pending.pop_front() {
self.dispatch(frame)?;
}
match self.frames.read(&mut buf) {
Ok(0) => return Ok(()),
Ok(n) => self.decoder.feed(&buf[..n]),
Err(e)
if e.kind() == std::io::ErrorKind::WouldBlock
|| e.kind() == std::io::ErrorKind::TimedOut => {}
Err(e) => return Err(e.into()),
}
while let Some(frame) = self.decoder.decode_next()? {
self.dispatch(frame)?;
}
self.engine.pump();
self.flush_op_completions()?;
}
}
fn send(&mut self, msg_type: u16, correlation: u64, payload: &[u8]) -> Result<(), WorkerError> {
send_frame(
&mut self.frames,
&mut self.seq,
msg_type,
correlation,
payload,
)
}
fn reply<T: serde::Serialize>(
&mut self,
msg_type: u16,
correlation: u64,
msg: &T,
) -> Result<(), WorkerError> {
let payload = postcard::to_stdvec(msg)?;
self.send(msg_type, correlation, &payload)
}
fn flush_op_completions(&mut self) -> Result<(), WorkerError> {
let handles: Vec<u64> = self.open_ops.keys().copied().collect();
for op in handles {
let st = match self.engine.op_poll(op) {
Ok(s) => s,
Err(_) => continue,
};
if st.state == op_state::PENDING || st.state == op_state::RUNNING {
continue;
}
let binding = self.open_ops.remove(&op).expect("present");
self.finalize_op(op, binding, st.state, st.error, st.bytes_moved)?;
}
Ok(())
}
fn finalize_op(
&mut self,
op: u64,
binding: OpBinding,
state: u8,
error: u32,
bytes_moved: u64,
) -> Result<(), WorkerError> {
if state == op_state::DONE {
if let OpBinding::Export { buf_id } = binding {
if let Some(payload) = self.engine.take_op_output(op) {
if let Some(seg) = self.bufs.get(&buf_id) {
let _ = seg.write_at(0, &payload);
}
}
}
}
self.reply(
msg_type::STATE_OP_DONE,
0,
&linkw::StateOpDoneMsg {
op_handle: op,
state,
error,
bytes_moved,
},
)?;
if let OpBinding::Export { buf_id } = binding {
if self.pending_buf_release.remove(&buf_id) {
self.bufs.remove(&buf_id);
}
}
Ok(())
}
fn dispatch(&mut self, frame: Frame) -> Result<(), WorkerError> {
let corr = frame.seq;
if !self.handshaken && frame.msg_type != msg_type::HELLO {
return Err(WorkerError::Protocol("message before handshake"));
}
if self.handshaken && frame.proto_version != self.chosen_version {
return Err(WorkerError::Protocol(
"frame version differs from the negotiated version",
));
}
match frame.msg_type {
msg_type::HELLO => {
let hello: Hello = postcard::from_bytes(&frame.payload)?;
let chosen = negotiate(hello.proto_versions, (PROTO_VERSION, PROTO_VERSION))
.ok_or(WorkerError::Protocol("no common proto version"))?;
let descs: Vec<_> = self.engine.state_spaces().iter().map(desc_to_msg).collect();
let ack = HelloAckW {
chosen_version: chosen,
engine_bundle_hash: self.cfg.engine_bundle_hash,
capabilities: self.engine.capability_descriptor().unwrap_or_default(),
state_space_descriptors: descs,
registered_strategies: Vec::new(),
limits: linkw::LinkWLimits {
max_frame: superfluid_proto::MAX_FRAME,
op_window: self.op_window as u32,
ring_specs: self.ring_specs.clone(),
},
kv_bits: self.engine.kv_bits(),
};
self.handshaken = true;
self.chosen_version = chosen;
self.reply(msg_type::HELLO_ACK_W, corr, &ack)
}
msg_type::RING_ATTACH => {
let msg: linkw::RingAttachMsg = postcard::from_bytes(&frame.payload)?;
let spec = self
.ring_specs
.iter()
.find(|s| s.ring_id == msg.ring_id)
.copied()
.ok_or(WorkerError::Protocol("attach for unknown ring id"))?;
if msg.layout != spec {
return Err(WorkerError::Protocol("ring layout disagrees with spec"));
}
let (fd, tag) = fdpass::recv_fd(&self.fd_channel)?;
if tag != msg.ring_id as u64 {
return Err(WorkerError::Protocol("fd tag does not match ring id"));
}
let seg =
ShmSegment::from_fd(fd, SharedRing::required_len(spec.slots, spec.slot_bytes))?;
let role = if msg.ring_id == TOKEN_RING_IN {
RingRole::Reader
} else {
RingRole::Writer
};
let ring = SharedRing::attach(seg, msg.ring_id, spec.slots, spec.slot_bytes, role)?;
let attach = superfluid_engine::rings::RingAttachment {
ring_id: msg.ring_id,
kind: match spec.kind {
linkw::RingKind::Tokens => superfluid_abi::ring_kind::TOKENS,
linkw::RingKind::Logits => superfluid_abi::ring_kind::LOGITS,
},
role: if msg.ring_id == TOKEN_RING_IN {
superfluid_abi::ring_role::ENGINE_READS
} else {
superfluid_abi::ring_role::ENGINE_WRITES
},
slots: spec.slots,
slot_bytes: spec.slot_bytes,
base: ring.base_ptr(),
len: ring.seg_len(),
};
self.engine
.attach_ring(&attach)
.map_err(|_| WorkerError::Protocol("engine refused ring attachment"))?;
self.pending_rings.insert(msg.ring_id, ring);
if self.rings.is_none()
&& [TOKEN_RING_IN, TOKEN_RING_OUT, LOGITS_RING]
.iter()
.all(|id| self.pending_rings.contains_key(id))
{
self.rings = Some(SharedRings {
token_in: self.pending_rings.remove(&TOKEN_RING_IN).expect("checked"),
token_out: self.pending_rings.remove(&TOKEN_RING_OUT).expect("checked"),
logits: self.pending_rings.remove(&LOGITS_RING).expect("checked"),
control: self.pending_control.take(),
});
}
self.reply(
msg_type::RING_ACK,
corr,
&linkw::RingAckMsg {
ring_id: msg.ring_id,
generation_base: 1,
},
)
}
msg_type::CONTROL_ATTACH => {
let msg: linkw::ControlAttachMsg = postcard::from_bytes(&frame.payload)?;
let (fd, tag) = fdpass::recv_fd(&self.fd_channel)?;
if tag != CONTROL_TAG {
return Err(WorkerError::Protocol("fd tag is not the control segment's"));
}
let control = ControlSegment::attach(ShmSegment::from_fd(fd, msg.bytes as usize)?)?;
match self.rings.as_mut() {
Some(rings) => rings.control = Some(control),
None => self.pending_control = Some(control),
}
self.reply(msg_type::CONTROL_ACK, corr, &linkw::ControlAckMsg { bytes: msg.bytes })
}
msg_type::TICK_SUBMIT => {
let msg: linkw::TickSubmitMsg = postcard::from_bytes(&frame.payload)?;
let Self { engine, frames, seq, rings, .. } = self;
let rings = rings
.as_mut()
.ok_or(WorkerError::Protocol("tick before ring attach"))?;
let (mut plan, arena) = materialize_plan(&msg.plan);
let mut sink = EmitSink { frames, seq, ok: true };
if msg.plan.flags & superfluid_abi::tick_flags::PARTIAL_EMITS != 0 {
plan.on_partial_emit = Some(emit_tramp);
plan.partial_emit_user = (&mut sink as *mut EmitSink).cast();
}
let result = match engine.tick(&plan, &arena, rings) {
Ok(ev) => linkw::TickResultMsg::Events(events_to_msg(ev)),
Err(s) => linkw::TickResultMsg::Rejected { status: s.raw() },
};
if let linkw::TickResultMsg::Events(ev) = &result {
self.last_mem = ev.mem;
let done: Vec<linkw::OpCompleteMsg> = ev.op_completions.clone();
for c in done {
if let Some(binding) = self.open_ops.remove(&c.op) {
self.finalize_op(c.op, binding, c.state, c.error, c.bytes_moved)?;
}
}
}
self.reply(msg_type::TICK_RESULT, corr, &result)
}
msg_type::MATCH_REQ => {
let msg: linkw::MatchReqMsg = postcard::from_bytes(&frame.payload)?;
let rings = self
.rings
.as_ref()
.ok_or(WorkerError::Protocol("match before ring attach"))?;
let span = token_ref_from_msg(&msg.span);
use superfluid_engine::Rings;
let tokens = rings
.read_tokens(&span)
.map_err(|_| WorkerError::Protocol("stale span in match"))?;
let res = self
.engine
.space_match_ref(msg.space_id, &span, &tokens, &msg.media_deps)
.map(match_result_to_msg)
.unwrap_or(linkw::MatchResMsg { spaces: Vec::new() });
self.reply(msg_type::MATCH_RES, corr, &res)
}
msg_type::SEED_ACQUIRE => {
let msg: linkw::SeedAcquireMsg = postcard::from_bytes(&frame.payload)?;
let rings = self
.rings
.as_ref()
.ok_or(WorkerError::Protocol("seed before ring attach"))?;
use superfluid_engine::Rings;
let span = token_ref_from_msg(&msg.span);
let grant = match rings.read_tokens(&span) {
Ok(tokens) => match self.engine.seed_acquire_ref(
&span,
&tokens,
msg.prefix_len,
msg.determinism_class,
) {
Ok(handle) => linkw::SeedGrantMsg::Granted {
seed_handle: handle,
ttl_ms: 1000,
ttl_ticks: self.engine.seed_lease_ticks(),
},
Err(s) => linkw::SeedGrantMsg::Refused { status: s.raw() },
},
Err(_) => linkw::SeedGrantMsg::Refused {
status: Status::RejectBadRefGeneration.raw(),
},
};
self.reply(msg_type::SEED_GRANT, corr, &grant)
}
msg_type::SEED_RELEASE => {
let msg: linkw::SeedReleaseMsg = postcard::from_bytes(&frame.payload)?;
let _ = self.engine.seed_release(msg.seed_handle);
Ok(())
}
msg_type::STATE_OP_SUBMIT => {
let msg: linkw::StateOpSubmitMsg = postcard::from_bytes(&frame.payload)?;
let reply = if self.open_ops.len() >= self.op_window {
linkw::StateOpAcceptedMsg::Rejected {
status: Status::Busy.raw(),
}
} else {
self.submit_state_op(&msg.op)
};
self.reply(msg_type::STATE_OP_ACCEPTED, corr, &reply)
}
msg_type::OP_CANCEL => {
let msg: linkw::OpCancelMsg = postcard::from_bytes(&frame.payload)?;
let _ = self.engine.op_cancel(msg.op_handle);
Ok(())
}
msg_type::STATE_SYNC_REQ => {
let msg: linkw::StateSyncReqMsg = postcard::from_bytes(&frame.payload)?;
let res = self.state_sync(&msg);
self.reply(msg_type::STATE_SYNC_RES, corr, &res)
}
msg_type::STRATEGY_REGISTER => {
let msg: linkw::StrategyRegisterMsg = postcard::from_bytes(&frame.payload)?;
let res = self.register_strategy(&msg);
self.reply(msg_type::STRATEGY_GRANT, corr, &res)
}
msg_type::BUF_ALLOC => {
let _msg: linkw::BufAllocMsg = postcard::from_bytes(&frame.payload)?;
Ok(())
}
msg_type::BUF_ATTACH => {
let msg: linkw::BufAttachMsg = postcard::from_bytes(&frame.payload)?;
let (fd, tag) = fdpass::recv_fd(&self.fd_channel)?;
if tag != msg.buf_id {
return Err(WorkerError::Protocol("fd tag does not match buf id"));
}
let seg = ShmSegment::from_fd(fd, msg.len as usize)?;
self.bufs.insert(msg.buf_id, seg);
Ok(())
}
msg_type::BUF_RELEASE => {
let msg: linkw::BufReleaseMsg = postcard::from_bytes(&frame.payload)?;
let bound = self
.open_ops
.values()
.any(|b| matches!(b, OpBinding::Export { buf_id } if *buf_id == msg.buf_id));
if bound {
self.pending_buf_release.insert(msg.buf_id);
} else {
self.bufs.remove(&msg.buf_id);
}
Ok(())
}
msg_type::PING => {
let _msg: linkw::PingMsg = postcard::from_bytes(&frame.payload)?;
let pong = linkw::PongMsg {
queue_depth: self.open_ops.len() as u32,
mem: self.last_mem,
};
self.reply(msg_type::PONG, corr, &pong)
}
msg_type::DRAIN_REQ => {
while !self.open_ops.is_empty() {
self.engine.pump();
self.flush_op_completions()?;
}
self.reply(msg_type::DRAINED, corr, &linkw::DrainedMsg)
}
msg_type::TELEMETRY => Ok(()),
msg_type::TRANSCRIBE_CANCEL => Ok(()),
_ => Err(WorkerError::Protocol("unhandled required message")),
}
}
fn submit_state_op(&mut self, op: &linkw::StateOpKind) -> linkw::StateOpAcceptedMsg {
let result: Result<(u64, OpBinding), Status> = match *op {
linkw::StateOpKind::Snapshot {
sequence,
space_id,
boundary_pos,
buf_id,
sizing_gen,
} => {
let Some(seg) = self.bufs.get(&buf_id) else {
return linkw::StateOpAcceptedMsg::Rejected {
status: Status::UnknownHandle.raw(),
};
};
let len = seg.len() as u64;
self.engine
.space_snapshot(sequence, space_id, boundary_pos, len, sizing_gen)
.map(|h| (h, OpBinding::Export { buf_id }))
}
linkw::StateOpKind::Demote {
sequence,
space_id,
range,
encoding,
buf_id,
sizing_gen,
} => {
let Some(seg) = self.bufs.get(&buf_id) else {
return linkw::StateOpAcceptedMsg::Rejected {
status: Status::UnknownHandle.raw(),
};
};
let len = seg.len() as u64;
self.engine
.space_demote(
sequence,
space_id,
TokenRange {
start: range.start,
end: range.end,
},
encoding,
len,
sizing_gen,
)
.map(|h| (h, OpBinding::Export { buf_id }))
}
linkw::StateOpKind::Restore {
sequence,
space_id,
buf_id,
} => match self.read_sealed(buf_id) {
Ok(bytes) => self
.engine
.space_restore(sequence, space_id, &bytes)
.map(|h| (h, OpBinding::Import)),
Err(s) => Err(s),
},
linkw::StateOpKind::Promote {
sequence,
space_id,
range,
buf_id,
} => match self.read_sealed(buf_id) {
Ok(bytes) => self
.engine
.space_promote(
sequence,
space_id,
TokenRange {
start: range.start,
end: range.end,
},
&bytes,
)
.map(|h| (h, OpBinding::Import)),
Err(s) => Err(s),
},
};
match result {
Ok((handle, binding)) => {
self.open_ops.insert(handle, binding);
linkw::StateOpAcceptedMsg::Accepted { op_handle: handle }
}
Err(s) => linkw::StateOpAcceptedMsg::Rejected { status: s.raw() },
}
}
fn read_sealed(&self, buf_id: u64) -> Result<Vec<u8>, Status> {
const ENVELOPE_BYTES: usize = std::mem::size_of::<superfluid_abi::StateEnvelope>();
let seg = self.bufs.get(&buf_id).ok_or(Status::UnknownHandle)?;
if seg.len() < ENVELOPE_BYTES {
return Err(Status::EnvelopeMismatch);
}
let mut head = vec![0u8; ENVELOPE_BYTES];
seg.read_at(0, &mut head)
.map_err(|_| Status::EnvelopeMismatch)?;
let off = std::mem::offset_of!(superfluid_abi::StateEnvelope, payload_len);
let payload_len = u64::from_le_bytes(head[off..off + 8].try_into().expect("8 bytes"));
let payload_len = usize::try_from(payload_len).map_err(|_| Status::EnvelopeMismatch)?;
let total = ENVELOPE_BYTES
.checked_add(payload_len)
.ok_or(Status::EnvelopeMismatch)?;
if seg.len() < total {
return Err(Status::EnvelopeMismatch);
}
let mut out = vec![0u8; total];
seg.read_at(0, &mut out)
.map_err(|_| Status::EnvelopeMismatch)?;
Ok(out)
}
fn state_sync(&mut self, msg: &linkw::StateSyncReqMsg) -> linkw::StateSyncResMsg {
let result = match *msg {
linkw::StateSyncReqMsg::Fork {
parent_sequence,
flags,
} => self.engine.seq_fork(parent_sequence, flags).map(|child| {
linkw::StateSyncOkMsg::Fork {
child_sequence: child,
}
}),
linkw::StateSyncReqMsg::Trim {
sequence,
space_id,
new_len,
} => self
.engine
.space_trim(sequence, space_id, new_len)
.map(|()| linkw::StateSyncOkMsg::Trim),
linkw::StateSyncReqMsg::ExportSize {
sequence,
space_id,
range,
encoding,
} => self
.engine
.space_export_size(
sequence,
space_id,
TokenRange {
start: range.start,
end: range.end,
},
encoding,
)
.map(
|(required_bytes, sizing_gen)| linkw::StateSyncOkMsg::ExportSize {
required_bytes,
sizing_gen,
},
),
linkw::StateSyncReqMsg::CacheEvict {
bytes_target,
ref policy,
} => self
.engine
.cache_evict(
policy.evictable_cache_classes,
policy.protected_quota_bytes,
policy.max_evict_bytes,
bytes_target,
)
.map(|bytes_freed| linkw::StateSyncOkMsg::CacheEvict { bytes_freed }),
linkw::StateSyncReqMsg::CacheEvictEntries { ref keys } => {
let raw: Vec<([u8; 32], [u8; 32])> =
keys.iter().map(|k| (k.compat_key, k.provenance_digest)).collect();
self.engine
.cache_evict_entries(&raw)
.map(|()| linkw::StateSyncOkMsg::CacheEvictEntries)
}
linkw::StateSyncReqMsg::SnapshotBoundary {
sequence,
space_id,
cap,
} => self
.engine
.space_snapshot_boundary(sequence, space_id, cap)
.map(|boundary| linkw::StateSyncOkMsg::SnapshotBoundary { boundary }),
linkw::StateSyncReqMsg::LaneSequence { lane_tag } => self
.engine
.lane_sequence(lane_tag)
.map(|sequence| linkw::StateSyncOkMsg::LaneSequence { sequence })
.ok_or(Status::UnknownHandle),
linkw::StateSyncReqMsg::CreateSequence => self
.engine
.create_sequence()
.map(|sequence| linkw::StateSyncOkMsg::CreateSequence { sequence }),
linkw::StateSyncReqMsg::FreeSequence { sequence } => self
.engine
.free_sequence(sequence)
.map(|()| linkw::StateSyncOkMsg::FreeSequence),
linkw::StateSyncReqMsg::CreateGrammarStructural { ref tag_json } => self
.engine
.grammar_create_structural(tag_json)
.map(|grammar_handle| linkw::StateSyncOkMsg::CreateGrammar { grammar_handle }),
linkw::StateSyncReqMsg::CreateGrammar { ref json_schema } => self
.engine
.grammar_create(json_schema)
.map(|grammar_handle| linkw::StateSyncOkMsg::CreateGrammar { grammar_handle }),
linkw::StateSyncReqMsg::FreeGrammar { grammar_handle } => self
.engine
.grammar_free(grammar_handle)
.map(|()| linkw::StateSyncOkMsg::FreeGrammar),
linkw::StateSyncReqMsg::CreateLogitBias { ref tokens, ref values } => {
let vals: Vec<f32> = values.iter().map(|b| f32::from_bits(*b)).collect();
self.engine
.logit_bias_create(tokens, &vals)
.map(|handle| linkw::StateSyncOkMsg::CreateLogitBias { handle })
}
linkw::StateSyncReqMsg::FreeLogitBias { handle } => self
.engine
.logit_bias_free(handle)
.map(|()| linkw::StateSyncOkMsg::FreeLogitBias),
linkw::StateSyncReqMsg::Embed { ref tokens } => self.engine.embed(tokens).map(|v| {
let mut bytes = Vec::with_capacity(v.len() * 4);
for f in v {
bytes.extend_from_slice(&f.to_le_bytes());
}
linkw::StateSyncOkMsg::Embed { embedding: bytes }
}),
linkw::StateSyncReqMsg::Transcribe {
ref audio_path,
ref language,
translate,
timestamps,
ref prompt,
stream,
} => {
let params = superfluid_engine::TranscribeParams {
language: language.clone(),
translate,
timestamps,
prompt: prompt.clone(),
};
let Self { engine, frames, seq, decoder, pending, .. } = self;
let mut sink = |start_ms: i32, end_ms: i32, text: &str| -> bool {
let msg = linkw::TranscribeSegmentMsg {
start_ms,
end_ms,
text: text.to_string(),
};
let sent = postcard::to_stdvec(&msg).is_ok_and(|p| {
send_frame(frames, seq, msg_type::TRANSCRIBE_SEGMENT, 0, &p).is_ok()
});
sent && !poll_transcribe_cancel(frames, decoder, pending)
};
if stream {
engine.transcribe(audio_path, ¶ms, Some(&mut sink))
} else {
engine.transcribe(audio_path, ¶ms, None)
}
}
.map(|t| linkw::StateSyncOkMsg::Transcribe {
text: t.text,
language: t.language,
duration_ms: t.duration_ms,
segments: t
.segments
.into_iter()
.map(|s| linkw::TranscriptSegmentMsg {
start_ms: s.start_ms,
end_ms: s.end_ms,
text: s.text,
avg_logprob_bits: s.avg_logprob.to_bits(),
no_speech_prob_bits: s.no_speech_prob.to_bits(),
compression_ratio_bits: s.compression_ratio.to_bits(),
temperature_bits: s.temperature.to_bits(),
})
.collect(),
}),
linkw::StateSyncReqMsg::LoraLoad { ref adapter_path } => self
.engine
.lora_load(adapter_path)
.map(|()| linkw::StateSyncOkMsg::LoraId { id: self.engine.lora_id() }),
linkw::StateSyncReqMsg::LoraUnload => self
.engine
.lora_unload()
.map(|()| linkw::StateSyncOkMsg::LoraId { id: self.engine.lora_id() }),
linkw::StateSyncReqMsg::LoraId => {
Ok(linkw::StateSyncOkMsg::LoraId { id: self.engine.lora_id() })
}
linkw::StateSyncReqMsg::Capabilities => Ok(linkw::StateSyncOkMsg::Capabilities {
json: self.engine.capability_descriptor().unwrap_or_default(),
}),
linkw::StateSyncReqMsg::ForwardStage {
ref tokens,
start_layer,
end_layer,
ref hidden_in,
} => self
.engine
.forward_stage(tokens, start_layer, end_layer, hidden_in.clone())
.map(|out| match out {
StageOutput::Hidden(h) => linkw::StateSyncOkMsg::ForwardStage {
hidden: Some(h),
token: None,
},
StageOutput::Token(t) => linkw::StateSyncOkMsg::ForwardStage {
hidden: None,
token: Some(t),
},
}),
linkw::StateSyncReqMsg::MediaProbe { ref image_path } => {
self.engine.media_probe(image_path).map(|i| linkw::StateSyncOkMsg::MediaProbe {
info: linkw::MediaInfoMsg {
n_tokens: i.n_tokens,
image_token_id: i.image_token_id,
boi_token_id: i.boi_token_id,
eoi_token_id: i.eoi_token_id,
preprocess_fp: i.preprocess_fp,
},
})
}
linkw::StateSyncReqMsg::MediaEncode { ref image_path } => {
self.engine
.media_encode(image_path)
.map(|(media_handle, i)| linkw::StateSyncOkMsg::MediaEncode {
media_handle,
info: linkw::MediaInfoMsg {
n_tokens: i.n_tokens,
image_token_id: i.image_token_id,
boi_token_id: i.boi_token_id,
eoi_token_id: i.eoi_token_id,
preprocess_fp: i.preprocess_fp,
},
})
}
linkw::StateSyncReqMsg::MediaRelease { media_handle } => self
.engine
.media_release(media_handle)
.map(|()| linkw::StateSyncOkMsg::MediaRelease),
linkw::StateSyncReqMsg::MediaBind {
lane_tag,
media_handle,
token_offset,
} => self
.engine
.media_bind(lane_tag, media_handle, token_offset)
.map(|()| linkw::StateSyncOkMsg::MediaBind),
linkw::StateSyncReqMsg::AdoptSequence { sequence, ref span } => {
use superfluid_engine::Rings;
let span_ref = token_ref_from_msg(span);
match self.rings.as_ref() {
None => Err(Status::RejectBadRefGeneration),
Some(rings) => match rings.read_tokens(&span_ref) {
Ok(tokens) => self
.engine
.seed_adopt_ref(sequence, &span_ref, &tokens)
.map(|seed_handle| linkw::StateSyncOkMsg::AdoptSequence { seed_handle }),
Err(_) => Err(Status::RejectBadRefGeneration),
},
}
}
linkw::StateSyncReqMsg::PublishSequence { sequence, ref span } => {
use superfluid_engine::Rings;
let span = token_ref_from_msg(span);
match self.rings.as_ref() {
None => Err(Status::RejectBadRefGeneration),
Some(rings) => match rings.read_tokens(&span) {
Ok(tokens) => self
.engine
.publish_sequence(sequence, &tokens)
.map(|()| linkw::StateSyncOkMsg::PublishSequence),
Err(_) => Err(Status::RejectBadRefGeneration),
},
}
}
linkw::StateSyncReqMsg::Unpublish { ref span } => {
use superfluid_engine::Rings;
let span = token_ref_from_msg(span);
match self.rings.as_ref() {
None => Err(Status::RejectBadRefGeneration),
Some(rings) => match rings.read_tokens(&span) {
Ok(tokens) => self.engine.unpublish(&tokens).map(|()| linkw::StateSyncOkMsg::Unpublish),
Err(_) => Err(Status::RejectBadRefGeneration),
},
}
}
};
match result {
Ok(ok) => linkw::StateSyncResMsg::Ok(ok),
Err(s) => linkw::StateSyncResMsg::Err { status: s.raw() },
}
}
fn register_strategy(&mut self, msg: &linkw::StrategyRegisterMsg) -> linkw::StrategyGrantMsg {
let strategy_id = CString::new(msg.strategy_id.as_str()).unwrap_or_default();
let impl_version = CString::new(msg.impl_version.as_str()).unwrap_or_default();
let load_paths: Vec<CString> = msg
.artifacts
.iter()
.map(|a| CString::new(a.load_path.as_str()).unwrap_or_default())
.collect();
let archs: Vec<CString> = msg
.target_archs
.iter()
.map(|a| CString::new(a.as_str()).unwrap_or_default())
.collect();
let mut arena = RecordArena::new();
let artifacts: Vec<Artifact> = msg
.artifacts
.iter()
.zip(&load_paths)
.map(|(a, p)| Artifact {
role: a.role,
_pad0: 0,
content_hash: a.content_hash,
byte_size: a.byte_size,
load_path: p.as_ptr(),
})
.collect();
let taps: Vec<TapSpec> = msg
.taps
.iter()
.map(|t| TapSpec {
layer: t.layer,
tensor: t.tensor,
dtype: t.dtype,
layout: t.layout,
})
.collect();
let capabilities: Vec<CapabilityReq> = msg
.capabilities
.iter()
.map(|c| CapabilityReq {
kind_id: c.kind_id,
_pad0: 0,
params: arena.push_records(&c.params),
})
.collect();
let arch_ptrs: Vec<u64> = archs.iter().map(|a| a.as_ptr() as u64).collect();
let reg = StrategyRegistration {
struct_size: std::mem::size_of::<StrategyRegistration>() as u64,
strategy_id: strategy_id.as_ptr(),
impl_version: impl_version.as_ptr(),
impl_hash: msg.impl_hash,
config_hash: msg.config_hash,
artifacts: arena.push_records(&artifacts),
taps: arena.push_records(&taps),
capabilities: arena.push_records(&capabilities),
target_archs: arena.push_records(&arch_ptrs),
kernel_caps_required: msg.kernel_caps_required,
_pad0: 0,
est_state_bytes: msg.est_state_bytes,
claimed_exactness: msg.claimed_exactness,
_pad1: [0; 3],
rng_contract_version: msg.rng_contract_version,
};
match self.engine.strategy_register(®) {
Ok(grant) => {
let bounds = superfluid_abi::ArenaBounds {
base: 0,
len: usize::MAX,
};
let certificates = unsafe {
superfluid_abi::array::read_array::<superfluid_abi::ExactnessCert>(
&grant.certificates,
&bounds,
)
.map(|it| {
it.map(|c| {
let ids = superfluid_abi::array::read_array::<[u8; 32]>(
&c.admissible_host_identities,
&bounds,
)
.map(|h| h.collect())
.unwrap_or_default();
linkw::ExactnessCertMsg {
cert_id: c.cert_id,
exactness: c.exactness,
sampling_modes: c.sampling_modes,
param_domain: c.param_domain,
grammar_allowed: c.grammar_allowed != 0,
logit_bias_allowed: c.logit_bias_allowed != 0,
verification_shape_class: c.verification_shape_class,
admissible_host_identities: ids,
}
})
.collect()
})
.unwrap_or_default()
};
let state_spaces = {
unsafe {
superfluid_abi::array::read_array::<superfluid_abi::StateSpaceDesc>(
&grant.state_spaces,
&bounds,
)
.map(|it| it.map(|d| desc_to_msg(&d)).collect())
.unwrap_or_default()
}
};
linkw::StrategyGrantMsg::Granted(linkw::StrategyGrantOkMsg {
strategy_slot: grant.strategy_slot,
certificates,
state_spaces,
reserved_bytes: grant.reserved_bytes,
})
}
Err(s) => linkw::StrategyGrantMsg::Refused { status: s.raw() },
}
}
}
fn zero_mem() -> linkw::MemCountersMsg {
linkw::MemCountersMsg {
allocated_bytes: 0,
host_retained_bytes: 0,
pool_blocks_total: 0,
pool_blocks_used: 0,
pool_bytes_evictable: 0,
tentative_bytes: 0,
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn the_emit_trampoline_refuses_null_pointers_from_the_engine() {
let (mut frames, mut peer) = UnixStream::pair().unwrap();
let mut seq = 1;
let toks = [4u32, 5, 6];
unsafe {
emit_tramp(std::ptr::null_mut(), 7, toks.as_ptr(), 3);
let mut sink = EmitSink { frames: &mut frames, seq: &mut seq, ok: true };
emit_tramp((&mut sink as *mut EmitSink).cast(), 7, std::ptr::null(), 3);
assert!(!sink.ok, "a null token pointer ends the partial emits");
let mut sink = EmitSink { frames: &mut frames, seq: &mut seq, ok: true };
emit_tramp((&mut sink as *mut EmitSink).cast(), 7, toks.as_ptr(), 3);
assert!(sink.ok, "a good emit is sent");
}
drop(frames);
let mut sent = Vec::new();
peer.read_to_end(&mut sent).unwrap();
assert!(!sent.is_empty(), "the good emit reached the link");
}
}