ts_wire/
handshake.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 startup handshake (`engine/zmq_wire.py`).
22//!
23//! The frontend binds a ROUTER handshake socket, a ROUTER input socket and a
24//! PULL output socket. Each engine (one per attention-DP rank) connects with
25//! its two-byte little-endian engine index as its identity and:
26//!
27//! 1. sends HELLO ([`ReadyMessage`]) on the handshake socket;
28//! 2. receives INIT ([`HandshakeInitMessage`]) with the input and output
29//!    addresses;
30//! 3. sends READY on the handshake socket;
31//! 4. connects its input DEALER and registers on it with
32//!    [`EngineCoreReadyResponse`];
33//! 5. connects its output PUSH.
34//!
35//! These structs use the map layout, matching the frontend's named codec.
36
37use crate::{codec::Dict, value::Value};
38
39/// HELLO status.
40pub const STATUS_HELLO: &str = "HELLO";
41/// READY status.
42pub const STATUS_READY: &str = "READY";
43
44wire_struct! {
45    /// Engine to frontend handshake status (HELLO or READY).
46    #[derive(Debug, Clone, PartialEq, Eq)]
47    pub struct ReadyMessage: map + Default {
48        /// `HELLO` or `READY`.
49        pub status: Option<String> = None,
50        /// The engine runs on the frontend's host.
51        pub local: Option<bool> = None,
52        /// The engine runs without an in-process frontend.
53        pub headless: Option<bool> = None,
54        /// Hash of the engine's parallel configuration.
55        pub parallel_config_hash: Option<String> = None,
56    }
57}
58
59wire_struct! {
60    /// Frontend-owned socket addresses sent in INIT.
61    #[derive(Debug, Clone, PartialEq, Eq)]
62    pub struct HandshakeAddresses: map + Default {
63        /// Input ROUTER addresses; engines use the first.
64        pub inputs: Vec<String> = Vec::new(),
65        /// Output PULL addresses; engines use the first.
66        pub outputs: Vec<String> = Vec::new(),
67        /// DP coordinator input, unused by TokenSpeed.
68        pub coordinator_input: Option<String> = None,
69        /// DP coordinator output, unused by TokenSpeed.
70        pub coordinator_output: Option<String> = None,
71        /// Frontend stats publisher, unused by TokenSpeed.
72        pub frontend_stats_publish_address: Option<String> = None,
73    }
74}
75
76wire_struct! {
77    /// Frontend to engine INIT, the reply to HELLO.
78    #[derive(Debug, Clone, PartialEq)]
79    pub struct HandshakeInitMessage: map {
80        /// Where to connect.
81        pub addresses: HandshakeAddresses,
82        /// Opaque to the engine.
83        pub parallel_config: Dict<Value, Value> = Dict::default(),
84    }
85}
86
87wire_struct! {
88    /// The engine's configuration, sent when it registers on the input socket.
89    ///
90    /// Field names, not order, are the contract; some names are the
91    /// frontend's vocabulary rather than TokenSpeed's.
92    #[derive(Debug, Clone, PartialEq)]
93    pub struct EngineCoreReadyResponse: map + Default {
94        /// Longest supported sequence.
95        pub max_model_len: i64 = 0,
96        /// KV pages on the device.
97        pub num_gpu_blocks: i64 = 0,
98        /// Prefix-cache matching granularity in tokens.
99        pub prefix_granularity: i64 = 0,
100        /// DP stats address, unused by TokenSpeed.
101        pub dp_stats_address: Option<String> = None,
102        /// Model dtype: `bfloat16`, `float16` or `float32`.
103        pub dtype: String = "bfloat16".to_owned(),
104        /// Multimodal encoder dtype, if any.
105        pub multimodal_encoder_dtype: Option<String> = None,
106        /// `tokenspeed-<version>`; the frontend's field name.
107        pub vllm_version: String = String::new(),
108        /// Ranks in the engine.
109        pub world_size: i64 = 1,
110        /// Attention-DP ranks.
111        pub data_parallel_size: i64 = 1,
112        /// Attention-TP ranks.
113        pub tensor_parallel_size: i64 = 1,
114        /// Pipeline stages.
115        pub pipeline_parallel_size: i64 = 1,
116        /// KV parallelism (`--kv-parallel-size`); the frontend's field name.
117        pub decode_context_parallel_size: i64 = 1,
118        /// This engine's attention-DP rank.
119        pub data_parallel_rank: i64 = 0,
120        /// Running-request limit.
121        pub max_num_seqs: i64 = 0,
122        /// Tokens per step limit; `0` is no cap.
123        pub max_num_batched_tokens: i64 = 0,
124        /// Served model name.
125        pub instance_id: String = String::new(),
126        /// KV cache capacity in tokens.
127        pub kv_cache_size_tokens: Option<i64> = None,
128        /// KV cache capacity in maximum-length requests.
129        pub kv_cache_max_concurrency: Option<f64> = None,
130        /// KV event publisher configuration.
131        pub kv_events_config: Option<Dict<Value, Value>> = None,
132    }
133}
ts_wire/handshake.rs

ts_wire/
handshake.rs
