diff --git a/.github/workflows/pr-test-rust.yml b/.github/workflows/pr-test-rust.yml index c0e1ca10e..4924da2fc 100644 --- a/.github/workflows/pr-test-rust.yml +++ b/.github/workflows/pr-test-rust.yml @@ -583,7 +583,7 @@ jobs: secrets: inherit e2e-2gpu-chat-zmq-dp: - name: e2e-2gpu-chat-zmq-dp (vllm) + name: e2e-2gpu-chat-zmq-dp (${{ matrix.engine }}) needs: [build-wheel, detect-changes] if: >- always() @@ -593,19 +593,33 @@ jobs: || (needs.detect-changes.result == 'success' && (needs.detect-changes.outputs.common == 'true' || needs.detect-changes.outputs.chat-completions == 'true'))) - # The 1-GPU ZMQ suite with grouped workers: each vLLM ZMQ worker launches - # two DP engines on one socket set (tp=1 models, one GPU per engine), so - # the gateway's handshake, the connector's least-loaded selection, and the - # wave protocol all run against a real dp=2 group. + # The 1-GPU ZMQ suite with grouped workers: each ZMQ worker launches two + # DP engines on one socket set (tp=1 models, one GPU per engine), so the + # gateway's handshake, the connector's least-loaded selection, and + # per-rank output attribution all run against a real dp=2 group. The wave + # protocol is vLLM-only (TokenSpeed ranks run independently). + strategy: + fail-fast: false + matrix: + include: + - engine: vllm + timeout: 46 + # The TokenSpeed ZMQ lane restarts the engine per model group, and + # DP doubles model load + warmup. The ~22-minute TokenSpeed source + # build eats the front of the budget, and at 50 the first dp run was + # killed mid-suite while progressing (28 minutes of test time); 60 + # leaves the tests the same ~40 minutes the test_timeout allows. + - engine: tokenspeed + timeout: 60 uses: ./.github/workflows/e2e-gpu-job.yml with: - engine: vllm + engine: ${{ matrix.engine }} # Tier selects the TEST SET (and model downloads): this lane runs the # tier-1 chat suite; the second GPU serves the second engine of the # group, not 2-GPU-marked tests. gpu_tier: "1" runner: 2-gpu-h100 - timeout: 46 + timeout: ${{ matrix.timeout }} test_timeout: 40 test_dirs: e2e_test/chat_completions connection_mode: zmq diff --git a/bindings/python/src/smg/serve.py b/bindings/python/src/smg/serve.py index 8d8929eeb..ef5678fa2 100644 --- a/bindings/python/src/smg/serve.py +++ b/bindings/python/src/smg/serve.py @@ -61,6 +61,13 @@ def _backend_arg_int(backend_args: list[str], flag: str, default: int) -> int: return default +# The band `derive_handshake_port` folds every worker's handshake port into. +# Other launcher-derived tcp ports must stay out of it: a port in this band can +# collide with SOME worker's handshake listener for the right ipc path. +_ZMQ_HANDSHAKE_PORT_BASE = 20000 +_ZMQ_HANDSHAKE_PORT_SPAN = 10000 + + def _zmq_handshake_port(ipc_url: str) -> int: """The tcp handshake port SMG derives from an ipc:// URL. @@ -73,7 +80,7 @@ def _zmq_handshake_port(ipc_url: str) -> int: for b in path.encode(): h ^= b h = (h * 0x100000001B3) & 0xFFFFFFFFFFFFFFFF - return 20000 + (h % 10000) + return _ZMQ_HANDSHAKE_PORT_BASE + (h % _ZMQ_HANDSHAKE_PORT_SPAN) def _reject_handshake_port_collisions(ports: list[int]) -> None: @@ -360,14 +367,41 @@ def build_command( def _build_zmq_command( self, args: argparse.Namespace, backend_args: list[str], port: int ) -> list[str]: - """Launch a headless TokenSpeed scheduler that dials SMG's ZMQ handshake. + """Launch a headless TokenSpeed scheduler group that dials SMG's ZMQ handshake. SMG (the router) binds the tcp handshake + ipc data-plane sockets it - derives from the ipc:// worker URL; this engine connects in. Each worker - is a standalone engine (`--zmq-engine-index 0`); running several is - dense data parallelism as N independent ZMQ workers. + derives from the ipc:// worker URL; the engines connect in. An + engine-level ``--data-parallel-size N`` (after ``--``) launches a + grouped worker: N ranks on one socket set, each dialing with its own + identity (``--zmq-engine-index`` is the group's base). The default + stays one standalone engine per worker. + + Two ports are the launcher's to own, or co-located workers collide: + + - ``--port`` seeds TokenSpeed's whole derived control-plane port + cluster (torch.distributed store at ``port + 233``, and neighbors). + Left at the engine default, every worker on the host derives the + same cluster, and back-to-back engine restarts race the previous + process's teardown (EADDRINUSE on the distributed store). + - ``--dist-init-addr`` pins that store explicitly. TokenSpeed derives + it from ``--port`` at dp==1 but refuses to guess for dp>1; passing + the same derivation it would use keeps one port layout for both. """ rpc_port = _zmq_handshake_port(_zmq_ipc_url(port)) + # Mirrors TokenSpeed's own dp==1 derivation (ZMQ_TCP_PORT_DELTA); + # reflected below the u16 ceiling instead of wrapping into low ports. + dist_port = port + 233 if port + 233 <= 65535 else port - 233 + # Hop over the SMG handshake band: a dist port inside it can land on + # a worker's rpc listener (this worker's included — the band is a hash + # of the ipc path). The +233 branch enters the band only from below, + # so one span-wide hop exits it for good; the -233 branch starts far + # above the band. + if ( + _ZMQ_HANDSHAKE_PORT_BASE + <= dist_port + < _ZMQ_HANDSHAKE_PORT_BASE + _ZMQ_HANDSHAKE_PORT_SPAN + ): + dist_port += _ZMQ_HANDSHAKE_PORT_SPAN cmd = [ sys.executable, "-m", @@ -376,6 +410,10 @@ def _build_zmq_command( "--headless", "--model", getattr(args, "model", ""), + "--port", + str(port), + "--dist-init-addr", + f"127.0.0.1:{dist_port}", "--data-parallel-address", "127.0.0.1", "--data-parallel-rpc-port", @@ -406,6 +444,8 @@ def _build_zmq_command( [ "--model", "--headless", + "--port", + "--dist-init-addr", "--data-parallel-address", "--data-parallel-rpc-port", "--zmq-engine-index", diff --git a/bindings/python/tests/test_serve.py b/bindings/python/tests/test_serve.py index 0231b6da4..9be6a49d0 100644 --- a/bindings/python/tests/test_serve.py +++ b/bindings/python/tests/test_serve.py @@ -858,9 +858,53 @@ def test_build_zmq_command(self): assert str(expected_port) in cmd assert "--zmq-engine-index" in cmd assert "0" in cmd + # The launcher owns the engine's control-plane port layout: --port + # seeds the derived cluster per worker, --dist-init-addr pins the + # torch.distributed store at TokenSpeed's own dp==1 derivation + # (port + 233) so dp==1 and dp>1 share one layout. + assert cmd[cmd.index("--port") + 1] == "31000" + assert cmd[cmd.index("--dist-init-addr") + 1] == "127.0.0.1:31233" for arg in backend_args: assert arg in cmd + def test_build_zmq_command_passes_dp_size_through(self): + # DP is the engine's flag: the launcher forwards it untouched and the + # ranks each dial the shared socket set with their own identity. + launcher = TokenspeedWorkerLauncher() + args = argparse.Namespace(model="/tmp/model", connection_mode="zmq") + cmd = launcher.build_command(args, ["--data-parallel-size", "2"], "127.0.0.1", 31000) + + assert cmd[cmd.index("--data-parallel-size") + 1] == "2" + # dp>1 hard-requires the explicit store address the launcher always passes. + assert cmd[cmd.index("--dist-init-addr") + 1] == "127.0.0.1:31233" + + def test_build_zmq_command_reflects_dist_port_below_u16_ceiling(self): + launcher = TokenspeedWorkerLauncher() + args = argparse.Namespace(model="/tmp/model", connection_mode="zmq") + cmd = launcher.build_command(args, [], "127.0.0.1", 65500) + + assert cmd[cmd.index("--dist-init-addr") + 1] == f"127.0.0.1:{65500 - 233}" + + def test_dist_port_never_enters_the_handshake_band(self): + # The handshake port is a hash of the ipc path folded into + # 20000..=29999, so a dist port inside that band can land on some + # worker's rpc listener — which port collides depends on the socket + # dir (uid). The launcher must keep the dist port out of the band + # entirely; that also implies it never equals this worker's own + # rpc port. Sweep every worker port whose naive +233 derivation + # lands in the band, plus the band edges and the u16 reflection. + launcher = TokenspeedWorkerLauncher() + args = argparse.Namespace(model="/tmp/model", connection_mode="zmq") + ports = [*range(19767, 29767), 19766, 29767, 31000, 65500] + for port in ports: + cmd = launcher.build_command(args, [], "127.0.0.1", port) + dist_port = int(cmd[cmd.index("--dist-init-addr") + 1].rsplit(":", 1)[1]) + rpc_port = _zmq_handshake_port(_zmq_ipc_url(port)) + assert not 20000 <= dist_port <= 29999, (port, dist_port) + assert dist_port != rpc_port, (port, dist_port) + assert dist_port != port, (port, dist_port) + assert 1 <= dist_port <= 65535, (port, dist_port) + def test_build_zmq_command_filters_launcher_owned_flags(self): launcher = TokenspeedWorkerLauncher() args = argparse.Namespace(model="/tmp/model", connection_mode="zmq") diff --git a/crates/engine_zmq_client/src/connector.rs b/crates/engine_zmq_client/src/connector.rs index e1974a3e2..1f3b74cda 100644 --- a/crates/engine_zmq_client/src/connector.rs +++ b/crates/engine_zmq_client/src/connector.rs @@ -1208,6 +1208,7 @@ mod tests { cached_tokens: vec![0], output_token_logprobs_val: vec![vec![]], output_token_logprobs_idx: vec![vec![]], + engine_index: 0, }; let done = BatchTokenIDOutSlim { rids: vec!["ts-1".into()], @@ -1218,6 +1219,7 @@ mod tests { cached_tokens: vec![0], output_token_logprobs_val: vec![vec![]], output_token_logprobs_idx: vec![vec![]], + engine_index: 0, }; engine .send_output(vec![Bytes::from(encode_msgpack(&chunk).unwrap())]) diff --git a/crates/engine_zmq_client/src/protocol/tokenspeed/mod.rs b/crates/engine_zmq_client/src/protocol/tokenspeed/mod.rs index 0280d802c..2eec47a35 100644 --- a/crates/engine_zmq_client/src/protocol/tokenspeed/mod.rs +++ b/crates/engine_zmq_client/src/protocol/tokenspeed/mod.rs @@ -99,8 +99,10 @@ impl EngineProtocol for TokenSpeedProtocol { } fn data_parallel_rank(_request: &Self::Request) -> Option { - // The TokenSpeed generate request carries no DP-rank field, so requests - // route to the sole engine (single-engine ZMQ). DP fan-out is future work. + // The TokenSpeed generate request carries no DP-rank field and needs + // none: rank routing is purely by ZMQ identity — the connector selects + // a rank and sends on that rank's socket identity. `None` means "never + // pinned", so every request goes through least-loaded selection. None } @@ -139,6 +141,7 @@ impl EngineProtocol for TokenSpeedProtocol { } let payload = frames.first().map(AsRef::as_ref).unwrap_or_default(); let batch: BatchTokenIDOutSlim = decode_msgpack(payload)?; + let engine_index = batch.engine_index; let outputs = batch.into_outputs()?; let finished_request_ids = outputs .iter() @@ -146,11 +149,11 @@ impl EngineProtocol for TokenSpeedProtocol { .map(|output| output.request_id.clone()) .collect(); Ok(EngineBatch { - // Single-engine ZMQ: TokenSpeed batches carry no engine index and - // no piggybacked scheduler load. - engine_index: 0, + engine_index, outputs, finished_request_ids, + // The slim batch piggybacks no scheduler load, so DP selection + // scores TokenSpeed ranks on the gateway's own in-flight counts. load: None, wave: None, }) @@ -171,6 +174,7 @@ mod tests { cached_tokens: vec![0, 0], output_token_logprobs_val: vec![vec![], vec![]], output_token_logprobs_idx: vec![vec![], vec![]], + engine_index: 1, } } @@ -181,6 +185,9 @@ mod tests { assert_eq!(decoded.outputs.len(), 2); assert_eq!(decoded.finished_request_ids, vec!["b".to_string()]); assert!(decoded.load.is_none()); + // The batch names its producing DP rank; the connector routes in-flight + // release and scoring by it. + assert_eq!(decoded.engine_index, 1); } #[test] diff --git a/crates/engine_zmq_client/src/protocol/tokenspeed/output.rs b/crates/engine_zmq_client/src/protocol/tokenspeed/output.rs index baaabe97c..3c6810cee 100644 --- a/crates/engine_zmq_client/src/protocol/tokenspeed/output.rs +++ b/crates/engine_zmq_client/src/protocol/tokenspeed/output.rs @@ -51,11 +51,16 @@ pub struct BatchTokenIDOutSlim { /// parallel to it per request. Empty inner `Vec` when logprobs were not /// requested. pub output_token_logprobs_idx: Vec>, + /// Producing DP rank's engine index. The output PULL socket carries no + /// routing identity, so under DP the batch itself names its rank. Appended + /// tail field: an older (pre-DP) sender emits 9 elements and this defaults + /// to `0`, which is also the sole rank of a single-engine worker. + pub engine_index: u32, } impl Serialize for BatchTokenIDOutSlim { fn serialize(&self, serializer: S) -> std::result::Result { - let mut tuple = serializer.serialize_tuple(9)?; + let mut tuple = serializer.serialize_tuple(10)?; tuple.serialize_element(BATCH_TOKEN_ID_OUT_SLIM_TAG)?; tuple.serialize_element(&self.rids)?; tuple.serialize_element(&self.output_ids)?; @@ -65,6 +70,7 @@ impl Serialize for BatchTokenIDOutSlim { tuple.serialize_element(&self.cached_tokens)?; tuple.serialize_element(&self.output_token_logprobs_val)?; tuple.serialize_element(&self.output_token_logprobs_idx)?; + tuple.serialize_element(&self.engine_index)?; tuple.end() } } @@ -94,6 +100,9 @@ impl<'de> Deserialize<'de> for BatchTokenIDOutSlim { cached_tokens: next_field(&mut seq, "cached_tokens")?, output_token_logprobs_val: next_field(&mut seq, "output_token_logprobs_val")?, output_token_logprobs_idx: next_field(&mut seq, "output_token_logprobs_idx")?, + // Appended by the DP wire revision: a 9-element batch from + // an older sender means rank 0 (the only rank it can be). + engine_index: seq.next_element::()?.unwrap_or(0), }; drain_trailing(&mut seq)?; Ok(batch) @@ -178,6 +187,9 @@ impl BatchTokenIDOutSlim { cached_tokens, output_token_logprobs_val, output_token_logprobs_idx, + // Batch-level rank tag, not a per-request column; the caller reads + // it off the batch before splitting. + engine_index: _, } = self; Ok(rids @@ -218,24 +230,32 @@ mod tests { use super::*; use crate::codec::{decode_msgpack, decode_value, encode_msgpack}; - /// A slim output batch captured from the Python encoder: rids ["vec-1"], - /// output_ids [[10, 11]], finished_reasons ["length"], prompt 3 / - /// completion 2 / cached 1, logprobs [[-0.5, -0.25]] over tokens [[10, 11]]. + /// A slim output batch captured from the Python msgspec encoder: rids + /// ["vec-1"], output_ids [[10, 11]], finished_reasons ["length"], prompt 3 + /// / completion 2 / cached 1, logprobs [[-0.5, -0.25]] over tokens + /// [[10, 11]], engine_index 1 (a DP sender's rank-1 batch). const PYTHON_OUTPUT_VECTOR: &str = + "9ab34261746368546f6b656e49444f7574536c696d91a57665632d3191920a0b91a66c656e\ + 6774689103910291019192cbbfe0000000000000cbbfd000000000000091920a0b01"; + + /// The same batch as encoded before the DP wire revision: 9 elements, no + /// engine_index tail field. + const PYTHON_OUTPUT_VECTOR_PRE_DP: &str = "99b34261746368546f6b656e49444f7574536c696d91a57665632d3191920a0b91a66c656e\ 6774689103910291019192cbbfe0000000000000cbbfd000000000000091920a0b"; - fn python_output_bytes() -> Vec { - let hex: String = PYTHON_OUTPUT_VECTOR - .chars() - .filter(|c| !c.is_whitespace()) - .collect(); + fn vector_bytes(hex_vector: &str) -> Vec { + let hex: String = hex_vector.chars().filter(|c| !c.is_whitespace()).collect(); (0..hex.len()) .step_by(2) .map(|i| u8::from_str_radix(&hex[i..i + 2], 16).unwrap()) .collect() } + fn python_output_bytes() -> Vec { + vector_bytes(PYTHON_OUTPUT_VECTOR) + } + fn vector_batch() -> BatchTokenIDOutSlim { BatchTokenIDOutSlim { rids: vec!["vec-1".into()], @@ -246,11 +266,12 @@ mod tests { cached_tokens: vec![1], output_token_logprobs_val: vec![vec![-0.5, -0.25]], output_token_logprobs_idx: vec![vec![10, 11]], + engine_index: 1, } } /// The pinned cross-language vector — the exact bytes the engine sends — - /// decodes into the full 9-element (tag + 8 columns) batch. + /// decodes into the full 10-element (tag + 8 columns + rank) batch. #[test] fn python_output_vector_decodes() { let decoded: BatchTokenIDOutSlim = decode_msgpack(&python_output_bytes()).unwrap(); @@ -277,16 +298,32 @@ mod tests { } #[test] - fn batch_output_serializes_as_tagged_nine_element_array() { + fn batch_output_serializes_as_tagged_ten_element_array() { let encoded = encode_msgpack(&vector_batch()).unwrap(); let Value::Array(array) = decode_value(&encoded).unwrap() else { panic!("expected positional array"); }; - assert_eq!(array.len(), 9); + assert_eq!(array.len(), 10); assert_eq!(array[0], Value::from(BATCH_TOKEN_ID_OUT_SLIM_TAG)); assert_eq!(array[1], Value::Array(vec![Value::from("vec-1")])); // rids assert_eq!(array[3], Value::Array(vec![Value::from("length")])); // finished_reasons assert_eq!(array[6], Value::Array(vec![Value::from(1)])); // cached_tokens + assert_eq!(array[9], Value::from(1)); // engine_index + } + + /// A pre-DP sender emits 9 elements; the missing tail decodes as rank 0, + /// the only rank a single-engine worker can be. + #[test] + fn pre_dp_nine_element_batch_decodes_as_rank_zero() { + let decoded: BatchTokenIDOutSlim = + decode_msgpack(&vector_bytes(PYTHON_OUTPUT_VECTOR_PRE_DP)).unwrap(); + assert_eq!( + decoded, + BatchTokenIDOutSlim { + engine_index: 0, + ..vector_batch() + } + ); } #[test] @@ -324,6 +361,7 @@ mod tests { cached_tokens: vec![0, 1], output_token_logprobs_val: vec![vec![-0.5], vec![-1.0, -2.0]], output_token_logprobs_idx: vec![vec![10], vec![20, 21]], + engine_index: 0, }; let outputs = batch.into_outputs().unwrap(); assert_eq!(outputs.len(), 2); @@ -351,6 +389,7 @@ mod tests { cached_tokens: vec![0], output_token_logprobs_val: vec![vec![]], output_token_logprobs_idx: vec![vec![]], + engine_index: 0, }; let outputs = batch.into_outputs().unwrap(); assert!(outputs[0].output_logprobs_val.is_empty()); @@ -368,6 +407,7 @@ mod tests { cached_tokens: vec![0, 1], output_token_logprobs_val: vec![vec![], vec![]], output_token_logprobs_idx: vec![vec![], vec![]], + engine_index: 0, }; assert!(batch.into_outputs().is_err()); } @@ -384,6 +424,7 @@ mod tests { // Only one logprob column entry for two requests. output_token_logprobs_val: vec![vec![]], output_token_logprobs_idx: vec![vec![], vec![]], + engine_index: 0, }; assert!(batch.into_outputs().is_err()); } diff --git a/e2e_test/fixtures/hooks.py b/e2e_test/fixtures/hooks.py index b3bd3acae..22e06d1a5 100644 --- a/e2e_test/fixtures/hooks.py +++ b/e2e_test/fixtures/hooks.py @@ -14,7 +14,13 @@ import os import pytest -from infra import ConnectionMode, cleanup_pool, get_connection_mode_override, get_runtime +from infra import ( + ConnectionMode, + cleanup_pool, + get_connection_mode_override, + get_runtime, + get_zmq_engine_count, +) from .markers import resolve_class_marker @@ -193,6 +199,27 @@ def _filter_zmq_items(items: list[pytest.Item]) -> tuple[list[pytest.Item], list return kept, deselected +# Models whose TokenSpeed forward crashes on the 0-token idle batch a DP rank +# runs to stay in the group's collectives (flashinfer silu_and_mul cannot +# launch over an empty grid). Fixed upstream by +# https://github.com/lightseekorg/tokenspeed/pull/1077; delete this skip when +# the tokenspeed pin advances past it. +_TOKENSPEED_DP_BROKEN_MODELS = frozenset({"Qwen/Qwen3-4B-Instruct-2507"}) + + +def _filter_tokenspeed_dp_items( + items: list[pytest.Item], +) -> tuple[list[pytest.Item], list[pytest.Item]]: + """Split items into (kept, deselected) for a grouped TokenSpeed ZMQ lane.""" + kept: list[pytest.Item] = [] + deselected: list[pytest.Item] = [] + for item in items: + marker = resolve_class_marker(item, "model") + model = str(marker.args[0]) if marker is not None and marker.args else "" + (deselected if model in _TOKENSPEED_DP_BROKEN_MODELS else kept).append(item) + return kept, deselected + + def pytest_collection_modifyitems( config: pytest.Config, items: list[pytest.Item], @@ -241,6 +268,14 @@ def pytest_collection_modifyitems( config.hook.pytest_deselected(items=deselected) items[:] = kept + # Grouped TokenSpeed lane: drop the models the pinned engine cannot + # run under DP (see _TOKENSPEED_DP_BROKEN_MODELS). + if get_runtime() == "tokenspeed" and get_zmq_engine_count() > 1: + kept, deselected = _filter_tokenspeed_dp_items(items) + if deselected: + config.hook.pytest_deselected(items=deselected) + items[:] = kept + items.sort(key=_pool_sort_key) diff --git a/e2e_test/infra/__init__.py b/e2e_test/infra/__init__.py index bde475869..dc97ae978 100644 --- a/e2e_test/infra/__init__.py +++ b/e2e_test/infra/__init__.py @@ -35,6 +35,7 @@ WorkerType, get_connection_mode_override, get_runtime, + get_zmq_engine_count, is_mlx, is_sglang, is_trtllm, @@ -113,6 +114,7 @@ # Runtime helpers "get_runtime", "get_connection_mode_override", + "get_zmq_engine_count", "is_vllm", "is_sglang", "is_trtllm", diff --git a/e2e_test/infra/worker.py b/e2e_test/infra/worker.py index 7b32c9be0..9a4a953c7 100644 --- a/e2e_test/infra/worker.py +++ b/e2e_test/infra/worker.py @@ -371,9 +371,13 @@ def _build_tokenspeed_zmq_cmd(self, model_path: str, tp_size: int, spec: dict) - args = argparse.Namespace( connection_mode="zmq", model=model_path, tensor_parallel_size=tp_size ) - return TokenspeedWorkerLauncher().build_command( - args, list(spec.get("tokenspeed_args", [])), DEFAULT_HOST, self.port - ) + backend_args = list(spec.get("tokenspeed_args", [])) + # Grouped lane: an engine-level dp flag makes the launcher start that + # many ranks on this worker's socket set (see get_zmq_engine_count). + engine_count = get_zmq_engine_count() + if engine_count > 1: + backend_args += ["--data-parallel-size", str(engine_count)] + return TokenspeedWorkerLauncher().build_command(args, backend_args, DEFAULT_HOST, self.port) def _build_tokenspeed_grpc_cmd(self, model_path: str, tp_size: int, spec: dict) -> list[str]: """Build TokenSpeed gRPC server command. @@ -624,14 +628,15 @@ def start_workers( gpus_per_worker = gpus or spec.get("tp", 1) if gpus is None and mode == ConnectionMode.ZMQ: # A grouped ZMQ worker launches get_zmq_engine_count() engines, each - # tp-wide, in one process — size its GPU slice accordingly. Only the - # vLLM launcher starts engine groups; a grouped count on any other - # runtime would reserve GPUs for engines that never launch and leave - # the gateway awaiting handshakes that never come. + # tp-wide, in one process — size its GPU slice accordingly. vLLM and + # TokenSpeed launchers both start engine groups; a grouped count on + # any other runtime would reserve GPUs for engines that never launch + # and leave the gateway awaiting handshakes that never come. zmq_engine_count = get_zmq_engine_count() - if zmq_engine_count > 1 and engine != "vllm": + if zmq_engine_count > 1 and engine not in ("vllm", "tokenspeed"): raise ValueError( - f"E2E_ZMQ_ENGINE_COUNT={zmq_engine_count} is vLLM-only; got engine={engine!r}" + f"E2E_ZMQ_ENGINE_COUNT={zmq_engine_count} needs an engine-group " + f"launcher (vllm or tokenspeed); got engine={engine!r}" ) gpus_per_worker *= zmq_engine_count timeout = spec.get("startup_timeout", timeout) diff --git a/model_gateway/src/routers/grpc/zmq_client.rs b/model_gateway/src/routers/grpc/zmq_client.rs index 2a275f988..4b0dab095 100644 --- a/model_gateway/src/routers/grpc/zmq_client.rs +++ b/model_gateway/src/routers/grpc/zmq_client.rs @@ -394,17 +394,6 @@ impl ZmqEngineClient { runtime: RuntimeType, timeout: Duration, ) -> Result> { - // Single-engine scope for TokenSpeed: its wire carries no DP-rank routing - // yet (`data_parallel_rank` is always `None`), so more than one engine - // would silently send all traffic to engine 0. Reject it loudly until - // DP>1 lands. The engine count is known here (the handshake awaits it). - if matches!(runtime, RuntimeType::TokenSpeed) && engine_count > 1 { - return Err(format!( - "TokenSpeed ZMQ backend supports a single engine only (got \ - engine_count={engine_count}); DP>1 is not yet supported" - ) - .into()); - } // No silent fallback: any other runtime has no ZMQ engine adapter. // Reject before the handshake — no such engine ever dials in, so the // handshake would just block for the full timeout. @@ -2156,6 +2145,7 @@ mod tests { cached_tokens: vec![0], output_token_logprobs_val: vec![vec![-0.5]], output_token_logprobs_idx: vec![vec![10]], + engine_index: 0, }; let done = BatchTokenIDOutSlim { rids: vec!["r1".into()], @@ -2166,6 +2156,7 @@ mod tests { cached_tokens: vec![0], output_token_logprobs_val: vec![vec![-1.25]], output_token_logprobs_idx: vec![vec![11]], + engine_index: 0, }; output .send_frames(vec![bytes::Bytes::from(encode_msgpack(&chunk).unwrap())]) @@ -2819,6 +2810,7 @@ mod tests { cached_tokens: vec![0, 0], output_token_logprobs_val: vec![vec![], vec![]], output_token_logprobs_idx: vec![vec![], vec![]], + engine_index: 0, }; output .send_frames(vec![bytes::Bytes::from(encode_msgpack(&done).unwrap())]) diff --git a/model_gateway/src/workflow/steps/local/create_worker.rs b/model_gateway/src/workflow/steps/local/create_worker.rs index de4bdc14f..ac829f04d 100644 --- a/model_gateway/src/workflow/steps/local/create_worker.rs +++ b/model_gateway/src/workflow/steps/local/create_worker.rs @@ -158,23 +158,12 @@ impl StepExecutor for CreateLocalWorkerStep { } // A grouped ZMQ worker (`dp_size: N` on the spec) awaits N engines on - // one socket set. TokenSpeed's msgpack wire carries no DP-rank routing - // yet, so fail its groups at registration for the same reason as the - // runtime check above. + // one socket set. Both ZMQ runtimes route per rank: vLLM by in-request + // DP rank, TokenSpeed by per-rank socket identity with the producing + // rank named on each output batch. let zmq_engine_group = config .dp_size .filter(|&n| n > 1 && *connection_mode == ConnectionMode::Zmq); - if zmq_engine_group.is_some() && runtime_type == RuntimeType::TokenSpeed { - return Err(WorkflowError::StepFailed { - step_id: StepId::new("create_worker"), - message: format!( - "ZMQ worker {} configures dp_size={} but the TokenSpeed wire \ - does not support DP>1 yet; only vllm engine groups are supported", - config.url, - config.dp_size.unwrap_or_default() - ), - }); - } // Normalize URL let url = normalize_url(&config.url, *connection_mode); diff --git a/scripts/ci_install_tokenspeed.sh b/scripts/ci_install_tokenspeed.sh index 464851199..f2d305b6d 100755 --- a/scripts/ci_install_tokenspeed.sh +++ b/scripts/ci_install_tokenspeed.sh @@ -22,7 +22,7 @@ fi # engine-watch workflow files an issue when this drifts) rather than # floating against ``main`` — upstream has renamed APIs before and the # gRPC servicer broke until we caught up. -TOKENSPEED_REF="${TOKENSPEED_REF:-788f0b09b49237176c02bfbfe23702b7579fe223}" +TOKENSPEED_REF="${TOKENSPEED_REF:-04bc08649f3e53e19144ee86564be7c6121c99d2}" TOKENSPEED_REPO="${TOKENSPEED_REPO:-https://github.com/lightseekorg/tokenspeed.git}" TOKENSPEED_DIR="${TOKENSPEED_DIR:-/tmp/tokenspeed-src}"