diff --git a/python/tokenspeed/runtime/engine/event_loop.py b/python/tokenspeed/runtime/engine/event_loop.py index 160fa1d1e0..f6aadc5ace 100644 --- a/python/tokenspeed/runtime/engine/event_loop.py +++ b/python/tokenspeed/runtime/engine/event_loop.py @@ -553,6 +553,20 @@ def __init__( ), metrics=self.metrics, ) + + # msgpack wire only: piggyback a load snapshot on every output batch. + # The socket predates the scheduler (it is built during the startup + # handshake), so the sampler late-binds here — the first point where + # scheduler, cache geometry, and output processor all exist. Same + # sources as the per-iteration Prometheus snapshot in + # _record_scheduler_iteration_metrics. + if self.server_args.zmq_msgpack and self.attn_tp_rank == 0: + self.send_to_tokenizer.load_fn = lambda: ( + len(self.output_processor.rid_to_state), + self.scheduler.waiting_size(), + self.scheduler.active_kv_pages(), + self._scheduler_cache_geometry.num_usable_pages, + ) if server_args.disaggregation_mode != "null": kv_args = get_kv_args( global_rank, diff --git a/python/tokenspeed/runtime/engine/io_struct.py b/python/tokenspeed/runtime/engine/io_struct.py index ca6cea8be6..06b7c226ee 100755 --- a/python/tokenspeed/runtime/engine/io_struct.py +++ b/python/tokenspeed/runtime/engine/io_struct.py @@ -712,10 +712,26 @@ class BatchTokenIDOutSlim(BaseBatchReq, kw_only=True): # DP the batch itself names its rank. Appended field: defaults to 0 so # older peers on either side stay compatible. engine_index: int = 0 + # Piggybacked scheduler-load snapshot from the producing rank, sampled at + # send time. The pickle-mode GetLoad poll has no msgpack transport (control + # replies are dropped on this wire), so the output batch is the only + # in-band load channel; an external frontend uses these for least-loaded + # routing across a DP group. Same sources as the per-iteration Prometheus + # snapshot: running = resident request states, waiting = scheduler queue + # depth, and the KV ratio's numerator/denominator carried as exact page + # counts. Appended fields: all default to 0 so older peers stay + # compatible; a 0 kv_total_pages means "no snapshot" to the frontend. + num_running: int = 0 + num_waiting: int = 0 + kv_active_pages: int = 0 + kv_total_pages: int = 0 @classmethod def from_full( - cls, out: BatchTokenIDOut, engine_index: int = 0 + cls, + out: BatchTokenIDOut, + engine_index: int = 0, + load: tuple[int, int, int, int] | None = None, ) -> "BatchTokenIDOutSlim": # Token source: ``out.output_ids`` — the not-yet-sent slice of each # request's generated ids. NOT ``out.decode_ids``: that is the @@ -729,8 +745,13 @@ def from_full( "BatchTokenIDOut.output_ids is None; the msgpack wire needs " "the per-request generated token ids" ) + num_running, num_waiting, kv_active_pages, kv_total_pages = load or (0, 0, 0, 0) return cls( engine_index=engine_index, + num_running=num_running, + num_waiting=num_waiting, + kv_active_pages=kv_active_pages, + kv_total_pages=kv_total_pages, rids=list(out.rids), output_ids=[list(ids) for ids in out.output_ids], finished_reasons=[_finish_type(fr) for fr in out.finished_reasons], diff --git a/python/tokenspeed/runtime/engine/zmq_msgpack.py b/python/tokenspeed/runtime/engine/zmq_msgpack.py index 4a891cd4c7..ddf50fb320 100644 --- a/python/tokenspeed/runtime/engine/zmq_msgpack.py +++ b/python/tokenspeed/runtime/engine/zmq_msgpack.py @@ -144,13 +144,25 @@ def __init__(self, socket: zmq.Socket, engine_index: int = 0) -> None: self._socket = socket self._engine_index = engine_index self._encoder = MsgpackEncoder() + # Late-bound scheduler-load sampler, () -> (num_running, num_waiting, + # kv_active_pages, kv_total_pages). The socket is created during the + # startup handshake, before the scheduler exists, so the event loop + # binds this once the scheduler is up. None means "no snapshot": + # the slim batch's load fields stay at their 0 defaults. + self.load_fn = None def send_pyobj(self, obj) -> None: if isinstance(obj, BatchTokenIDOut): # The PULL side carries no routing identity, so the batch itself # names its producing rank; the frontend attributes per-rank - # outputs and load by this index under DP. - slim = BatchTokenIDOutSlim.from_full(obj, engine_index=self._engine_index) + # outputs and load by this index under DP. The load snapshot rides + # every batch because this wire has no control-reply channel for + # the GetLoad poll (dropped below). + slim = BatchTokenIDOutSlim.from_full( + obj, + engine_index=self._engine_index, + load=self.load_fn() if self.load_fn is not None else None, + ) self._socket.send_multipart(self._encoder.encode(slim), copy=False) else: logger.warning( diff --git a/test/runtime/test_zmq_msgpack.py b/test/runtime/test_zmq_msgpack.py index 721d748497..16849e2ce7 100644 --- a/test/runtime/test_zmq_msgpack.py +++ b/test/runtime/test_zmq_msgpack.py @@ -271,6 +271,56 @@ def test_slim_out_engine_index_defaults_for_older_senders(): assert rt.engine_index == 0 +def test_slim_out_piggybacks_the_load_snapshot(): + # The msgpack wire has no control-reply channel for the GetLoad poll, so + # the output batch carries the producing rank's scheduler-load snapshot: + # (num_running, num_waiting, kv_active_pages, kv_total_pages). + slim = BatchTokenIDOutSlim.from_full(_make_batch_out(), load=(2, 5, 100, 400)) + rt = msgspec.msgpack.Decoder(BatchTokenIDOutSlim).decode(_encode_payload(slim)) + assert rt.num_running == 2 + assert rt.num_waiting == 5 + assert rt.kv_active_pages == 100 + assert rt.kv_total_pages == 400 + + +def test_slim_out_load_defaults_for_older_senders(): + # Appended fields: a 10-element array (engine_index era, pre-load) must + # still decode with the zero "no snapshot" defaults — kv_total_pages == 0 + # is the frontend's signal that no load rode this batch. + slim = BatchTokenIDOutSlim.from_full(_make_batch_out(), load=(2, 5, 100, 400)) + raw = msgspec.msgpack.decode(_encode_payload(slim)) + rt = msgspec.msgpack.Decoder(BatchTokenIDOutSlim).decode( + msgspec.msgpack.encode(raw[:10]) + ) + assert rt.num_running == 0 + assert rt.num_waiting == 0 + assert rt.kv_active_pages == 0 + assert rt.kv_total_pages == 0 + + +def test_send_socket_samples_load_fn_per_batch(): + # MsgpackSendSocket late-binds a load sampler (the socket exists before + # the scheduler); every BatchTokenIDOut send must sample it fresh. + sent = [] + + class _Sock: + def send_multipart(self, frames, copy=False): + sent.append(list(frames)) + + sender = zmq_msgpack.MsgpackSendSocket(_Sock(), engine_index=1) + samples = iter([(1, 0, 10, 400), (0, 0, 0, 400)]) + sender.load_fn = lambda: next(samples) + sender.send_pyobj(_make_batch_out()) + sender.send_pyobj(_make_batch_out()) + + decoder = msgspec.msgpack.Decoder(BatchTokenIDOutSlim) + first = decoder.decode(sent[0][0]) + second = decoder.decode(sent[1][0]) + assert (first.num_running, first.kv_active_pages) == (1, 10) + assert (second.num_running, second.kv_active_pages) == (0, 0) + assert first.engine_index == second.engine_index == 1 + + def test_slim_out_sources_output_ids_not_the_detok_window(): # ``decode_ids`` on the io_struct is the incremental-detokenization window, # which starts at the prompt tail for context; ``output_ids`` is the @@ -305,22 +355,24 @@ def test_slim_out_carries_logprob_columns(): def test_slim_out_is_tagged_positional_tuple(): - """The wire form must be a 10-element tagged array (the frontend's codec - depends on the tag and column order; engine_index is the appended tail).""" + """The wire form must be a 14-element tagged array (the frontend's codec + depends on the tag and column order; engine_index and the load snapshot + are the appended tail).""" out = _make_batch_out( output_token_logprobs_val=[[-0.5]], output_token_logprobs_idx=[[10]] ) - slim = BatchTokenIDOutSlim.from_full(out, engine_index=1) + slim = BatchTokenIDOutSlim.from_full(out, engine_index=1, load=(2, 5, 100, 400)) raw = msgspec.msgpack.decode(_encode_payload(slim)) assert isinstance(raw, list) assert raw[0] == "BatchTokenIDOutSlim" - assert len(raw) == 10 + assert len(raw) == 14 assert raw[1] == ["r1"] # rids assert raw[2] == [[10, 11]] # output_ids assert raw[3] == ["length"] # finished_reasons assert raw[7] == [[pytest.approx(-0.5)]] assert raw[8] == [[10]] - assert raw[9] == 1 # engine_index (appended) + assert raw[9] == 1 # engine_index (appended, #1046) + assert raw[10:14] == [2, 5, 100, 400] # load snapshot (appended) def test_slim_out_finish_reason_none_maps_to_empty():