somatize_worker/
ws_transport.rs

1//! WebSocket-based Transport implementation.
2//!
3//! Implements the `Transport` trait from soma-runtime, sending plans to
4//! remote workers via WebSocket and receiving results.
5
6use crate::error::{Result, WorkerError};
7use somatize_compiler::ExecutionPlan;
8use somatize_core::value::Value;
9use somatize_runtime::executor::RunMode;
10use somatize_runtime::node_catalog::NodeCatalog;
11use somatize_runtime::runner::Transport;
12use std::collections::HashMap;
13
14use crate::protocol::*;
15
16/// Transport implementation using WebSocket.
17pub struct WsTransport {
18    /// The worker's base address (`ws://host:port`); rewritten to
19    /// `http(s)://` for the bulk upload/download endpoints.
20    pub address: String,
21    /// Bearer token appended to every connection when the worker
22    /// requires authentication.
23    pub token: Option<String>,
24}
25
26/// Drive `fut` to completion from a synchronous caller.
27///
28/// On a thread of our own, with a runtime of its own, always. [`Transport`]
29/// is a synchronous trait — the effect driver calls it from
30/// `std::thread::scope` — but a caller may equally already be inside a
31/// tokio runtime: `soma.Worker.serve()` runs an axum server on a thread of
32/// the user's Python process, and the Python bindings dispatch plans from
33/// there. `block_on` inside a runtime is a *panic*, not an error.
34///
35/// This file used to hold two contradictory answers to that. `upload` and
36/// `resolve_output` paid for a thread and said why in a comment;
37/// `send_msg`, `notify` and `stream_plan` built a bare current-thread
38/// runtime and blocked on it, so they worked until they were called from
39/// the wrong place. One rule now: never assume you are outside a runtime.
40///
41/// Scoped rather than detached, so the future may borrow `self`.
42fn on_own_runtime<F, T>(fut: F) -> Result<T>
43where
44    F: std::future::Future<Output = Result<T>> + Send,
45    T: Send,
46{
47    std::thread::scope(|scope| {
48        scope
49            .spawn(|| {
50                tokio::runtime::Builder::new_current_thread()
51                    .enable_all()
52                    .build()
53                    .map_err(|e| WorkerError::Concurrency(format!("tokio: {e}")))?
54                    .block_on(fut)
55            })
56            .join()
57            .map_err(|_| WorkerError::Concurrency("transport thread panicked".into()))?
58    })
59}
60
61impl WsTransport {
62    /// A transport to the worker at `address`, authenticating with
63    /// `token` if given. Connections are opened per call, not held.
64    pub fn new(address: impl Into<String>, token: Option<String>) -> Self {
65        Self {
66            address: address.into(),
67            token,
68        }
69    }
70
71    /// The worker's HTTP address, for the bulk endpoints.
72    fn http_addr(&self) -> String {
73        self.address
74            .replace("ws://", "http://")
75            .replace("wss://", "https://")
76    }
77
78    /// Send a `CoordinatorToWorker` message and wait for the response.
79    ///
80    /// Public because a caller that builds its own plan — the Python
81    /// bindings decide which worker gets which filters, which is policy,
82    /// not transport — should not have to open its own socket to ship it.
83    /// A second `connect_async` elsewhere is a second place to get the
84    /// frame-size configuration wrong.
85    pub fn send_msg(&self, msg: &CoordinatorToWorker) -> Result<WorkerToCoordinator> {
86        on_own_runtime(async {
87            let url = if let Some(t) = &self.token {
88                format!("{}/ws?token={t}", self.address)
89            } else {
90                format!("{}/ws", self.address)
91            };
92
93            let ws_config = {
94                let mut c = tokio_tungstenite::tungstenite::protocol::WebSocketConfig::default();
95                c.max_message_size = None;
96                c.max_frame_size = None;
97                c
98            };
99
100            let (mut ws, _) =
101                tokio_tungstenite::connect_async_with_config(&url, Some(ws_config), false)
102                    .await
103                    .map_err(|e| WorkerError::Transport(format!("WS connect: {e}")))?;
104
105            use futures_util::{SinkExt, StreamExt};
106            use tokio_tungstenite::tungstenite::Message;
107
108            let json = serde_json::to_string(msg)
109                .map_err(|e| WorkerError::Encoding(format!("serialize: {e}")))?;
110
111            ws.send(Message::Text(json.into()))
112                .await
113                .map_err(|e| WorkerError::Transport(format!("WS send: {e}")))?;
114
115            while let Some(frame) = ws.next().await {
116                // Ping, Pong and Binary are not answers; a Text frame
117                // always is, one way or another.
118                let Ok(Message::Text(response)) = frame else {
119                    continue;
120                };
121                let parsed = serde_json::from_str::<WorkerToCoordinator>(&response);
122                let _ = ws.close(None).await;
123                return match parsed {
124                    // The worker said what went wrong. Surface it here
125                    // rather than at every call site.
126                    Ok(WorkerToCoordinator::Error { message }) => {
127                        Err(WorkerError::Transport(format!("remote worker: {message}")))
128                    }
129                    Ok(result) => Ok(result),
130                    // A reply this build cannot read is an error, not
131                    // something to skip. Skipping it is what made every
132                    // worker-side failure hang the caller until the socket
133                    // closed, with nothing said about why.
134                    Err(e) => Err(WorkerError::Encoding(format!(
135                        "cannot read the worker's reply: {e}. Raw: {}",
136                        response.chars().take(200).collect::<String>()
137                    ))),
138                };
139            }
140
141            Err(WorkerError::Transport(
142                "worker closed without response".into(),
143            ))
144        })
145    }
146
147    /// Send a message without waiting for an answer.
148    ///
149    /// `Shutdown` is the one that needs this: the worker is not going to
150    /// reply, so [`WsTransport::send_msg`] would block until the socket
151    /// closed.
152    pub fn notify(&self, msg: &CoordinatorToWorker) -> Result<()> {
153        let url = match &self.token {
154            Some(t) => format!("{}/ws?token={t}", self.address),
155            None => format!("{}/ws", self.address),
156        };
157        let json = serde_json::to_string(msg)
158            .map_err(|e| WorkerError::Encoding(format!("serialize: {e}")))?;
159
160        on_own_runtime(async move {
161            use futures_util::SinkExt;
162            use tokio_tungstenite::tungstenite::Message;
163
164            let (mut ws, _) = tokio_tungstenite::connect_async(&url)
165                .await
166                .map_err(|e| WorkerError::Transport(format!("WS connect: {e}")))?;
167            ws.send(Message::Text(json.into()))
168                .await
169                .map_err(|e| WorkerError::Transport(format!("WS send: {e}")))
170        })
171    }
172
173    /// Upload a value to the worker's `/upload` endpoint, for payloads too
174    /// large to travel inline in a WebSocket message.
175    pub fn upload(&self, value: &Value) -> Result<somatize_core::store::DataRef> {
176        let url = format!("{}/upload", self.http_addr());
177        let body = serde_json::to_vec(value)
178            .map_err(|e| WorkerError::Encoding(format!("serialize upload: {e}")))?;
179        let token = self.token.clone();
180
181        // Blocking HTTP on its own thread: this may be called from inside
182        // a tokio runtime, and nesting one is a panic.
183        std::thread::spawn(move || {
184            let client = reqwest::blocking::Client::new();
185            let mut req = client
186                .post(&url)
187                .header("Content-Type", "application/json")
188                .body(body);
189            if let Some(t) = &token {
190                req = req.query(&[("token", t.as_str())]);
191            }
192            let resp = req
193                .send()
194                .map_err(|e| WorkerError::Transport(format!("HTTP upload: {e}")))?;
195            if !resp.status().is_success() {
196                return Err(WorkerError::Transport(format!(
197                    "HTTP upload failed: {}",
198                    resp.status()
199                )));
200            }
201            resp.json::<somatize_core::store::DataRef>()
202                .map_err(|e| WorkerError::Encoding(format!("parse upload response: {e}")))
203        })
204        .join()
205        .map_err(|_| WorkerError::Concurrency("upload thread panicked".into()))?
206    }
207
208    /// Ship a plan and a stream of chunks over one WebSocket, collecting
209    /// results as they come back.
210    ///
211    /// The binary side of the protocol. It lived in the Python bindings,
212    /// which meant a second hand-rolled `connect_async` and a second copy
213    /// of the msgpack `StreamMessage` framing, a crate away from the enum
214    /// that defines it.
215    pub fn stream_plan(&self, plan: SerializedPlan, chunks: Vec<Value>) -> Result<Value> {
216        let stream_id = plan.plan_id.clone();
217        let total_chunks = chunks.len();
218        let url = match &self.token {
219            Some(t) => format!("{}/ws?token={t}", self.address),
220            None => format!("{}/ws", self.address),
221        };
222
223        on_own_runtime(async move {
224            use futures_util::StreamExt;
225            use tokio_tungstenite::tungstenite::Message;
226
227            let (mut ws, _) = tokio_tungstenite::connect_async(&url)
228                .await
229                .map_err(|e| WorkerError::Transport(format!("WS connect: {e}")))?;
230
231            send_frame(
232                &mut ws,
233                StreamMessage::StreamBegin {
234                    stream_id: stream_id.clone(),
235                    plan_id: stream_id.clone(),
236                    total_chunks: Some(total_chunks),
237                    plan: Box::new(plan),
238                },
239            )
240            .await?;
241
242            let mut results: Vec<Value> = Vec::new();
243
244            for (i, chunk) in chunks.into_iter().enumerate() {
245                send_frame(
246                    &mut ws,
247                    StreamMessage::ChunkData {
248                        stream_id: stream_id.clone(),
249                        chunk_index: i,
250                        value: chunk,
251                    },
252                )
253                .await?;
254
255                // Drain whatever has come back so far, so a long stream
256                // does not queue every result until the end.
257                while let Ok(Some(Ok(Message::Binary(resp)))) =
258                    tokio::time::timeout(std::time::Duration::from_millis(1), ws.next()).await
259                {
260                    if let Ok(StreamMessage::ChunkResult { value, .. }) =
261                        crate::protocol::decode_frame(&resp)
262                    {
263                        results.push(value);
264                    }
265                }
266            }
267
268            send_frame(
269                &mut ws,
270                StreamMessage::StreamEnd {
271                    stream_id: stream_id.clone(),
272                },
273            )
274            .await?;
275
276            // Whatever is left, then the barrier filters' flush.
277            let mut flushed: Option<Value> = None;
278            while let Some(Ok(Message::Binary(resp))) = ws.next().await {
279                match crate::protocol::decode_frame(&resp) {
280                    Ok(StreamMessage::ChunkResult { value, .. }) => results.push(value),
281                    Ok(StreamMessage::StreamComplete { result, .. }) => match result {
282                        PlanResult::Success { output, .. } => {
283                            let v = self.resolve_output(&output)?;
284                            if !v.is_empty() {
285                                flushed = Some(v);
286                            }
287                            break;
288                        }
289                        PlanResult::Failed { error, .. } => {
290                            return Err(WorkerError::Remote(format!("stream error: {error}")));
291                        }
292                    },
293                    _ => {}
294                }
295            }
296
297            if let Some(v) = flushed {
298                results.push(v);
299            }
300            match results.len() {
301                0 => Ok(Value::Empty),
302                1 => Ok(results.into_iter().next().unwrap()),
303                _ => Ok(somatize_runtime::executors::materialize_buffer(&results)?),
304            }
305        })
306    }
307
308    /// Resolve OutputDelivery — inline or download via HTTP.
309    pub fn resolve_output(&self, delivery: &OutputDelivery) -> Result<Value> {
310        match delivery {
311            OutputDelivery::Inline { value } => Ok(value.clone()),
312            OutputDelivery::Reference { data_ref } => {
313                let url = format!("{}/download", self.http_addr());
314                let ref_json = serde_json::to_string(data_ref)
315                    .map_err(|e| WorkerError::Encoding(format!("serialize ref: {e}")))?;
316                let token = self.token.clone();
317
318                std::thread::spawn(move || {
319                    let client = reqwest::blocking::Client::new();
320                    let mut req = client.get(&url).query(&[("ref", &ref_json)]);
321                    if let Some(t) = &token {
322                        req = req.query(&[("token", t.as_str())]);
323                    }
324                    let resp = req
325                        .send()
326                        .map_err(|e| WorkerError::Transport(format!("HTTP download: {e}")))?;
327                    let bytes = resp
328                        .bytes()
329                        .map_err(|e| WorkerError::Transport(format!("read response: {e}")))?;
330                    serde_json::from_slice(&bytes)
331                        .map_err(|e| WorkerError::Encoding(format!("deserialize: {e}")))
332                })
333                .join()
334                .map_err(|_| WorkerError::Concurrency("download thread panicked".into()))?
335            }
336        }
337    }
338}
339
340/// One msgpack frame down a socket.
341async fn send_frame<S>(ws: &mut S, msg: StreamMessage) -> Result<()>
342where
343    S: futures_util::Sink<tokio_tungstenite::tungstenite::Message> + Unpin,
344    S::Error: std::fmt::Display,
345{
346    use futures_util::SinkExt;
347    let bytes = crate::protocol::encode_frame(&msg)?;
348    ws.send(tokio_tungstenite::tungstenite::Message::Binary(
349        bytes.into(),
350    ))
351    .await
352    .map_err(|e| WorkerError::Transport(format!("WS send: {e}")))
353}
354
355/// The seam.
356///
357/// `Transport` is a `soma-runtime` trait, so these return `SomaError`
358/// while everything behind them is a typed [`WorkerError`]. A refused
359/// socket and a reply that would not decode stay distinguishable inside
360/// this crate, which is where the retry decision is made.
361/// Build the wire plan for one `execute`.
362///
363/// Split out of the trait method so it can be tested without a socket:
364/// what goes on the wire is worth an assertion, and the seed in
365/// particular used to be hardcoded `None` here.
366///
367/// `filters` stays empty, and a `NodeCatalog` cannot change that. The
368/// worker reconstructs a Python filter by unpickling
369/// `SerializedFilter::pickled_filter`, and those bytes only exist in the
370/// Python layer (`Graph.pickled_filters`); a catalog holds live
371/// `Arc<dyn Filter>` values and their states, never the pickle. So the
372/// old `TODO: serialize from NodeCatalog if needed` described something
373/// that cannot be done from here — sending entries with empty pickle
374/// bytes would be worse than sending none, since the worker would try to
375/// unpickle them. The path that CAN supply filters builds its own
376/// `SerializedPlan` in `soma-python/src/graph.rs`; this transport is for
377/// plans whose nodes the worker already has.
378fn wire_plan(
379    plan: &ExecutionPlan,
380    input: &Value,
381    mode: &RunMode,
382    seed: Option<i64>,
383) -> SerializedPlan {
384    SerializedPlan {
385        protocol_version: PROTOCOL_VERSION,
386        plan_id: somatize_core::util::timestamp_id("remote"),
387        plan: plan.clone(),
388        input: Some(InputSource::Inline {
389            value: input.clone(),
390        }),
391        filters: vec![],
392        mode: match mode {
393            RunMode::Fit { y } => ExecutionMode::Fit {
394                y: y.clone(),
395                batch_size: None,
396            },
397            RunMode::Forward => ExecutionMode::Forward,
398        },
399        seed,
400        metadata: serde_json::json!({}),
401    }
402}
403
404impl Transport for WsTransport {
405    fn execute(
406        &self,
407        plan: &ExecutionPlan,
408        _filters: &NodeCatalog,
409        input: &Value,
410        mode: &RunMode,
411        seed: Option<i64>,
412    ) -> somatize_core::error::Result<(Value, HashMap<String, Value>)> {
413        let serialized = wire_plan(plan, input, mode, seed);
414
415        let msg = CoordinatorToWorker::AssignPlan { plan: serialized };
416        match self.send_msg(&msg)? {
417            WorkerToCoordinator::PlanResult { result, .. } => match result {
418                PlanResult::Success { output, states, .. } => {
419                    let value = self.resolve_output(&output)?;
420                    Ok((value, states))
421                }
422                PlanResult::Failed { error, .. } => {
423                    Err(WorkerError::Remote(format!("remote: {error}")).into())
424                }
425            },
426            other => {
427                Err(WorkerError::Transport(format!("expected PlanResult, got: {other:?}")).into())
428            }
429        }
430    }
431
432    fn get_state(
433        &self,
434        node_ids: &[String],
435    ) -> somatize_core::error::Result<HashMap<String, Value>> {
436        let msg = CoordinatorToWorker::GetState {
437            plan_id: String::new(),
438            node_ids: node_ids.to_vec(),
439        };
440        match self.send_msg(&msg)? {
441            WorkerToCoordinator::StateResult { states, .. } => Ok(states),
442            other => {
443                Err(WorkerError::Transport(format!("expected StateResult, got: {other:?}")).into())
444            }
445        }
446    }
447
448    fn set_state(&self, states: &HashMap<String, Value>) -> somatize_core::error::Result<()> {
449        let msg = CoordinatorToWorker::SetState {
450            plan_id: String::new(),
451            states: states.clone(),
452        };
453        self.send_msg(&msg)?;
454        Ok(())
455    }
456
457    fn get_gradients(
458        &self,
459        node_ids: &[String],
460    ) -> somatize_core::error::Result<HashMap<String, Value>> {
461        let msg = CoordinatorToWorker::GetGradients {
462            plan_id: String::new(),
463            node_ids: node_ids.to_vec(),
464        };
465        match self.send_msg(&msg)? {
466            WorkerToCoordinator::GradientsResult { gradients, .. } => Ok(gradients),
467            other => Err(WorkerError::Transport(format!(
468                "expected GradientsResult, got: {other:?}"
469            ))
470            .into()),
471        }
472    }
473
474    fn apply_gradients(
475        &self,
476        gradients: &HashMap<String, Value>,
477    ) -> somatize_core::error::Result<()> {
478        let msg = CoordinatorToWorker::ApplyGradients {
479            plan_id: String::new(),
480            gradients: gradients.clone(),
481        };
482        self.send_msg(&msg)?;
483        Ok(())
484    }
485}
486
487#[cfg(test)]
488mod tests {
489    use super::*;
490
491    fn plan() -> ExecutionPlan {
492        ExecutionPlan::Execute {
493            node_id: "n".into(),
494        }
495    }
496
497    // The regression this file existed to have. `seed` was hardcoded to
498    // `None` here while `SerializedPlan::seed` documented itself as the
499    // field that stops a sweep's seeds sharing one cache line — so the
500    // remote path silently reintroduced the bug the protocol had closed.
501    #[test]
502    fn wire_plan_carries_the_run_seed() {
503        let p = wire_plan(&plan(), &Value::Empty, &RunMode::Forward, Some(7));
504        assert_eq!(p.seed, Some(7));
505    }
506
507    #[test]
508    fn wire_plan_without_a_seed_stays_unseeded() {
509        let p = wire_plan(&plan(), &Value::Empty, &RunMode::Forward, None);
510        assert_eq!(p.seed, None);
511    }
512
513    // Two seeds must not produce the same wire plan, or the worker cannot
514    // tell them apart and the cache line is shared again.
515    #[test]
516    fn different_seeds_differ_on_the_wire() {
517        let a = wire_plan(&plan(), &Value::Empty, &RunMode::Forward, Some(1));
518        let b = wire_plan(&plan(), &Value::Empty, &RunMode::Forward, Some(2));
519        assert_ne!(a.seed, b.seed);
520    }
521
522    // Not a limitation to fix later: a NodeCatalog holds live filters, not
523    // the pickle bytes the worker unpickles, so this transport cannot fill
524    // this in. Asserted so nobody "implements the TODO" by sending empty
525    // pickles, which the worker would try to unpickle.
526    #[test]
527    fn wire_plan_sends_no_filters() {
528        let p = wire_plan(&plan(), &Value::Empty, &RunMode::Forward, Some(1));
529        assert!(p.filters.is_empty());
530    }
531
532    #[test]
533    fn fit_mode_carries_labels_and_the_seed_together() {
534        let y = Value::tensor(vec![1.0], vec![1]);
535        let p = wire_plan(
536            &plan(),
537            &Value::Empty,
538            &RunMode::Fit { y: Some(y.clone()) },
539            Some(3),
540        );
541        assert_eq!(p.seed, Some(3));
542        match p.mode {
543            ExecutionMode::Fit { y: got, .. } => assert_eq!(got, Some(y)),
544            other => panic!("expected Fit, got {other:?}"),
545        }
546    }
547}
somatize_worker/ws_transport.rs

somatize_worker/
ws_transport.rs
