ts_pd/
wire.rs

1// Copyright (c) 2026 LightSeek Foundation
2//
3// Permission is hereby granted, free of charge, to any person obtaining a copy
4// of this software and associated documentation files (the "Software"), to deal
5// in the Software without restriction, including without limitation the rights
6// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
7// copies of the Software, and to permit persons to whom the Software is
8// furnished to do so, subject to the following conditions:
9//
10// The above copyright notice and this permission notice shall be included in
11// all copies or substantial portions of the Software.
12//
13// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
14// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
15// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
16// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
17// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
18// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
19// SOFTWARE.
20
21//! The P/D control messages. The frames are those the Python peers sent
22//! before this crate replaced them, byte for byte, so nodes of either
23//! version interoperate.
24
25use bytes::Bytes;
26
27use crate::{
28    net::{ascii, parse_uint},
29    status::TransferPoll,
30};
31
32/// The largest room a pre-allocation may name.
33const MAX_ROOM: u64 = (1 << 63) - 1;
34
35/// A malformed control message.
36#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
37#[error("{0}")]
38pub struct WireError(pub String);
39
40fn error(message: &str) -> WireError {
41    WireError(message.to_owned())
42}
43
44fn text(frame: &[u8], name: &str) -> Result<String, WireError> {
45    if !frame.is_ascii() {
46        return Err(WireError(format!("{name} must be ASCII")));
47    }
48    Ok(String::from_utf8_lossy(frame).into_owned())
49}
50
51/// A Decode rank's registration with a Prefill rank: where to reach it,
52/// its transfer session, and its cache.
53#[allow(clippy::too_many_arguments)]
54pub fn registration_frames(
55    host: &str,
56    port: u16,
57    session_id: &str,
58    arena_base: u64,
59    layout: &[u8],
60    tp_size: i64,
61    tp_rank: i64,
62) -> Vec<Bytes> {
63    vec![
64        Bytes::from_static(b"None"),
65        ascii(host),
66        ascii(port),
67        ascii(session_id),
68        // struct.pack("Q"): native order.
69        Bytes::copy_from_slice(&arena_base.to_ne_bytes()),
70        Bytes::copy_from_slice(layout),
71        ascii(tp_size),
72        ascii(tp_rank),
73    ]
74}
75
76/// Whether `frames` is a registration rather than a pre-allocation.
77pub fn is_registration(frames: &[Bytes]) -> bool {
78    frames
79        .first()
80        .is_some_and(|header| header.as_ref() == b"None")
81}
82
83/// A registration's reply route, readable even when its layout is rejected
84/// (`KVArgsRegisterInfo.route_header_from_zmq`): endpoint, port, session.
85pub fn registration_route(frames: &[Bytes]) -> Result<(String, u16, String), WireError> {
86    if frames.len() != 8 || !is_registration(frames) {
87        return Err(error("CachePD slab registration requires exactly 8 frames"));
88    }
89    let endpoint = text(&frames[1], "CachePD registration endpoint")?;
90    let session_id = text(&frames[3], "CachePD registration session")?;
91    if endpoint.is_empty() || endpoint.len() > 255 {
92        return Err(error("CachePD registration endpoint is invalid"));
93    }
94    if session_id.is_empty() || session_id.len() > 512 {
95        return Err(error("CachePD registration session is invalid"));
96    }
97    let port = parse_uint(&frames[2], 65535)
98        .filter(|&port| port >= 1)
99        .ok_or_else(|| error("CachePD registration port is invalid"))?;
100    Ok((endpoint, port as u16, session_id))
101}
102
103/// A Decode rank's pre-allocation of one request: its destination block
104/// manifest, empty for a rank that only joins the status agreement.
105pub fn preallocation_frames(room: u64, session_id: &str, manifest: Option<&[u8]>) -> Vec<Bytes> {
106    vec![
107        ascii(room),
108        ascii(session_id),
109        manifest.map_or_else(Bytes::new, Bytes::copy_from_slice),
110    ]
111}
112
113/// A pre-allocation's room and session (`TransferInfo.route_header_from_zmq`).
114pub fn preallocation_route(frames: &[Bytes]) -> Result<(u64, String), WireError> {
115    if frames.len() != 3 {
116        return Err(error("CachePD pre-allocation requires exactly 3 frames"));
117    }
118    let room = parse_uint(&frames[0], MAX_ROOM)
119        .ok_or_else(|| error("CachePD room must be canonical unsigned ASCII"))?;
120    let session_id = text(&frames[1], "CachePD session")?;
121    if session_id.is_empty() || session_id.len() > 512 {
122        return Err(error("CachePD request requires a Decode session"));
123    }
124    Ok((room, session_id))
125}
126
127/// A Prefill rank's report on one room.
128#[derive(Debug, Clone, PartialEq)]
129pub struct StatusMessage {
130    /// The room.
131    pub room: u64,
132    /// Its state on the sender.
133    pub status: TransferPoll,
134    /// The sender in Decode's completion-tracking rank space.
135    pub prefill_rank: i64,
136    /// The first output token, -1 if none.
137    pub bootstrap_token: i64,
138    /// The drafter's candidates for the first decode, if any.
139    pub spec_candidate_ids: Option<Vec<i32>>,
140    /// Prompt tokens the sender served from its cache.
141    pub cached_tokens: i64,
142    /// The first token's logprob, if the request returns logprobs.
143    pub bootstrap_logprob: Option<f64>,
144}
145
146impl StatusMessage {
147    /// A bare status: no first token and no usage.
148    pub fn bare(room: u64, status: TransferPoll, prefill_rank: i64) -> Self {
149        Self {
150            room,
151            status,
152            prefill_rank,
153            bootstrap_token: -1,
154            spec_candidate_ids: None,
155            cached_tokens: 0,
156            bootstrap_logprob: None,
157        }
158    }
159
160    /// The message's frames.
161    pub fn to_frames(&self) -> Vec<Bytes> {
162        let candidates = self
163            .spec_candidate_ids
164            .as_ref()
165            .map_or_else(Bytes::new, |ids| {
166                ids.iter()
167                    .flat_map(|id| id.to_le_bytes())
168                    .collect::<Vec<u8>>()
169                    .into()
170            });
171        let logprob = self.bootstrap_logprob.map_or_else(Bytes::new, |value| {
172            Bytes::copy_from_slice(&value.to_le_bytes())
173        });
174        vec![
175            ascii(self.room),
176            ascii(self.status as i64),
177            ascii(self.prefill_rank),
178            ascii(self.bootstrap_token),
179            candidates,
180            ascii(self.cached_tokens),
181            logprob,
182        ]
183    }
184
185    /// Parse a status message; the last
186    /// three frames are optional, as older senders omit them.
187    pub fn from_frames(frames: &[Bytes]) -> Result<Self, WireError> {
188        fn int(frame: &[u8], name: &str) -> Result<i64, WireError> {
189            std::str::from_utf8(frame)
190                .ok()
191                .and_then(|value| value.parse().ok())
192                .ok_or_else(|| WireError(format!("{name} is not an integer")))
193        }
194        if frames.len() < 3 {
195            return Err(error("a status message has at least 3 frames"));
196        }
197        let room =
198            u64::try_from(int(&frames[0], "room")?).map_err(|_| error("room is negative"))?;
199        let status = TransferPoll::from_wire(int(&frames[1], "status")?)
200            .ok_or_else(|| error("unknown status"))?;
201        let prefill_rank = int(&frames[2], "prefill rank")?;
202        let bootstrap_token = match frames.get(3) {
203            Some(frame) => int(frame, "bootstrap token")?,
204            None => -1,
205        };
206        let spec_candidate_ids = match frames.get(4) {
207            Some(frame) if !frame.is_empty() => {
208                if frame.len() % 4 != 0 {
209                    return Err(error("candidate ids are not int32"));
210                }
211                Some(
212                    frame
213                        .chunks_exact(4)
214                        .map(|chunk| i32::from_le_bytes(chunk.try_into().expect("4 bytes")))
215                        .collect(),
216                )
217            }
218            _ => None,
219        };
220        let cached_tokens = match frames.get(5) {
221            Some(frame) => int(frame, "cached tokens")?,
222            None => 0,
223        };
224        let bootstrap_logprob = match frames.get(6) {
225            Some(frame) if !frame.is_empty() => Some(f64::from_le_bytes(
226                frame
227                    .as_ref()
228                    .try_into()
229                    .map_err(|_| error("the logprob is not a double"))?,
230            )),
231            _ => None,
232        };
233        Ok(Self {
234            room,
235            status,
236            prefill_rank,
237            bootstrap_token,
238            spec_candidate_ids,
239            cached_tokens,
240            bootstrap_logprob,
241        })
242    }
243}
244
245/// The port a session id (`host:port`) ends in, which shards a room's
246/// transfers across workers.
247pub fn session_port(session_id: &str) -> Option<u64> {
248    session_id.rsplit(':').next()?.parse().ok()
249}
250
251#[cfg(test)]
252mod tests {
253    use super::*;
254
255    #[test]
256    fn a_status_message_round_trips() {
257        let message = StatusMessage {
258            room: 77,
259            status: TransferPoll::Success,
260            prefill_rank: 3,
261            bootstrap_token: 1234,
262            spec_candidate_ids: Some(vec![5, -1, 7]),
263            cached_tokens: 64,
264            bootstrap_logprob: Some(-0.25),
265        };
266        assert_eq!(
267            StatusMessage::from_frames(&message.to_frames()),
268            Ok(message)
269        );
270        let bare = StatusMessage::bare(9, TransferPoll::Failed, 0);
271        assert_eq!(StatusMessage::from_frames(&bare.to_frames()), Ok(bare));
272        // Older senders stop after the rank.
273        let short = [ascii(9), ascii(5), ascii(1)];
274        let parsed = StatusMessage::from_frames(&short).unwrap();
275        assert_eq!(parsed.bootstrap_token, -1);
276        assert_eq!(parsed.cached_tokens, 0);
277    }
278
279    #[test]
280    fn routes_parse_from_well_formed_frames_only() {
281        let frames = registration_frames("10.0.0.1", 4000, "10.0.0.1:5000", 0x1000, b"L", 4, 2);
282        assert_eq!(
283            registration_route(&frames),
284            Ok(("10.0.0.1".to_owned(), 4000, "10.0.0.1:5000".to_owned()))
285        );
286        assert_eq!(frames[4].as_ref(), 0x1000u64.to_ne_bytes());
287        assert!(registration_route(&frames[..7]).is_err());
288
289        let frames = preallocation_frames(42, "10.0.0.1:5000", None);
290        assert_eq!(
291            preallocation_route(&frames),
292            Ok((42, "10.0.0.1:5000".to_owned()))
293        );
294        assert!(frames[2].is_empty());
295        let bad = [ascii("042"), ascii("s"), Bytes::new()];
296        assert!(preallocation_route(&bad).is_err());
297        assert_eq!(session_port("[::1]:5000"), Some(5000));
298    }
299}
ts_pd/wire.rs

ts_pd/
wire.rs
