[Core] Clean up array-like msgspec structs (#32688)
This commit is contained in:
@@ -5,7 +5,7 @@
|
|||||||
//! `msgspec.msgpack`. Two struct families are involved:
|
//! `msgspec.msgpack`. Two struct families are involved:
|
||||||
//!
|
//!
|
||||||
//! * `EventBatch` (the outer payload) — declared with
|
//! * `EventBatch` (the outer payload) — declared with
|
||||||
//! `array_like=True, omit_defaults=True, gc=False` (no tag).
|
//! `array_like=True, gc=False` (no tag).
|
||||||
//! * `KVCacheEvent` (each inner event variant) — additionally declared
|
//! * `KVCacheEvent` (each inner event variant) — additionally declared
|
||||||
//! with `tag=True`.
|
//! with `tag=True`.
|
||||||
//!
|
//!
|
||||||
@@ -17,9 +17,6 @@
|
|||||||
//! of each inner event array, so an event is
|
//! of each inner event array, so an event is
|
||||||
//! `[class_name_str, field1, field2, ...]`. The outer `EventBatch`
|
//! `[class_name_str, field1, field2, ...]`. The outer `EventBatch`
|
||||||
//! array does **not** carry a tag prefix.
|
//! array does **not** carry a tag prefix.
|
||||||
//! * `omit_defaults=True` allows trailing fields whose values equal their
|
|
||||||
//! declared defaults to be dropped from the array. The decoder therefore
|
|
||||||
//! accepts variable-length sequences for each struct shape.
|
|
||||||
//!
|
//!
|
||||||
//! This module deserializes those bytes into Rust types and exposes a single
|
//! This module deserializes those bytes into Rust types and exposes a single
|
||||||
//! [`decode_event_batch`] entry point.
|
//! [`decode_event_batch`] entry point.
|
||||||
@@ -32,7 +29,7 @@ use serde::Deserialize;
|
|||||||
/// Top-level batch payload published by SGLang.
|
/// Top-level batch payload published by SGLang.
|
||||||
///
|
///
|
||||||
/// Wire shape (`EventBatch`, `array_like`):
|
/// Wire shape (`EventBatch`, `array_like`):
|
||||||
/// `[ts: f64, events: [...], attn_dp_rank: int_or_nil_or_omitted]`.
|
/// `[ts: f64, events: [...], attn_dp_rank: int_or_nil]`.
|
||||||
/// SGLang declares `attn_dp_rank` as a Python `Optional[int]`; we decode
|
/// SGLang declares `attn_dp_rank` as a Python `Optional[int]`; we decode
|
||||||
/// it as `u32` since DP ranks are non-negative and bounded by the
|
/// it as `u32` since DP ranks are non-negative and bounded by the
|
||||||
/// publisher's `dp_size`.
|
/// publisher's `dp_size`.
|
||||||
@@ -43,7 +40,8 @@ pub struct KvEventBatch {
|
|||||||
/// Ordered list of cache events in this batch.
|
/// Ordered list of cache events in this batch.
|
||||||
pub events: Vec<KvCacheEvent>,
|
pub events: Vec<KvCacheEvent>,
|
||||||
/// Optional DP-attention rank that produced this batch. `None` if the
|
/// Optional DP-attention rank that produced this batch. `None` if the
|
||||||
/// publisher emitted nil or omitted the field via `omit_defaults`.
|
/// publisher emitted nil. The decoder also accepts an omitted field for
|
||||||
|
/// compatibility.
|
||||||
pub attn_dp_rank: Option<u32>,
|
pub attn_dp_rank: Option<u32>,
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -332,8 +330,8 @@ impl<'de> Deserialize<'de> for BoundedU32Vec {
|
|||||||
|
|
||||||
// ---------------------------------------------------------------------------
|
// ---------------------------------------------------------------------------
|
||||||
// Custom Deserialize impls — msgspec encodes these structs as msgpack arrays
|
// Custom Deserialize impls — msgspec encodes these structs as msgpack arrays
|
||||||
// (not maps), and `omit_defaults=True` means trailing optional fields may be
|
// (not maps). The visitors also accept absent trailing optional fields for
|
||||||
// absent. We therefore implement `Deserialize` by hand against `SeqAccess`.
|
// compatibility.
|
||||||
// ---------------------------------------------------------------------------
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
impl<'de> Deserialize<'de> for KvEventBatch {
|
impl<'de> Deserialize<'de> for KvEventBatch {
|
||||||
@@ -731,8 +729,7 @@ mod tests {
|
|||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn attn_dp_rank_omitted_decodes_as_none() {
|
fn attn_dp_rank_omitted_decodes_as_none() {
|
||||||
// msgspec's `omit_defaults=True` may drop attn_dp_rank entirely from
|
// Accept a payload that omits the optional trailing rank field.
|
||||||
// the wire array when it equals its default of None.
|
|
||||||
let event = build_all_blocks_cleared_bytes();
|
let event = build_all_blocks_cleared_bytes();
|
||||||
let bytes = build_batch_bytes(5.0, &[event], None, /* include_dp_field */ false);
|
let bytes = build_batch_bytes(5.0, &[event], None, /* include_dp_field */ false);
|
||||||
|
|
||||||
|
|||||||
@@ -61,7 +61,6 @@ def select_kv_publisher_dp_rank(
|
|||||||
class EventBatch(
|
class EventBatch(
|
||||||
msgspec.Struct,
|
msgspec.Struct,
|
||||||
array_like=True, # type: ignore[call-arg]
|
array_like=True, # type: ignore[call-arg]
|
||||||
omit_defaults=True, # type: ignore[call-arg]
|
|
||||||
gc=False, # type: ignore[call-arg]
|
gc=False, # type: ignore[call-arg]
|
||||||
):
|
):
|
||||||
ts: float
|
ts: float
|
||||||
@@ -72,7 +71,6 @@ class EventBatch(
|
|||||||
class KVCacheEvent(
|
class KVCacheEvent(
|
||||||
msgspec.Struct,
|
msgspec.Struct,
|
||||||
array_like=True, # type: ignore[call-arg]
|
array_like=True, # type: ignore[call-arg]
|
||||||
omit_defaults=True, # type: ignore[call-arg]
|
|
||||||
gc=False, # type: ignore[call-arg]
|
gc=False, # type: ignore[call-arg]
|
||||||
tag=True,
|
tag=True,
|
||||||
):
|
):
|
||||||
|
|||||||
@@ -72,7 +72,7 @@ def raise_if_tokenizer_required(
|
|||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
class SamplingParams(msgspec.Struct, kw_only=True, omit_defaults=True):
|
class SamplingParams(msgspec.Struct, kw_only=True, array_like=True):
|
||||||
"""
|
"""
|
||||||
The sampling parameters.
|
The sampling parameters.
|
||||||
|
|
||||||
@@ -106,10 +106,10 @@ class SamplingParams(msgspec.Struct, kw_only=True, omit_defaults=True):
|
|||||||
skip_special_tokens: bool = True
|
skip_special_tokens: bool = True
|
||||||
spaces_between_special_tokens: bool = True
|
spaces_between_special_tokens: bool = True
|
||||||
no_stop_trim: bool = False
|
no_stop_trim: bool = False
|
||||||
custom_params: Optional[Dict[str, CustomParamValue]] = None
|
|
||||||
stream_interval: Optional[int] = None
|
stream_interval: Optional[int] = None
|
||||||
logit_bias: Optional[Dict[str, float]] = None
|
logit_bias: Optional[Dict[str, float]] = None
|
||||||
sampling_seed: Optional[int] = None
|
sampling_seed: Optional[int] = None
|
||||||
|
custom_params: Optional[Dict[str, CustomParamValue]] = None
|
||||||
|
|
||||||
# --- Internal fields (populated by __post_init__ or normalize(), not API-facing) ---
|
# --- Internal fields (populated by __post_init__ or normalize(), not API-facing) ---
|
||||||
stop_strs: Optional[Union[str, List[str]]] = None # from stop
|
stop_strs: Optional[Union[str, List[str]]] = None # from stop
|
||||||
@@ -269,7 +269,7 @@ class SamplingParams(msgspec.Struct, kw_only=True, omit_defaults=True):
|
|||||||
tokenizer, self.stop_strs, self.stop_regex_strs, self.min_new_tokens
|
tokenizer, self.stop_strs, self.stop_regex_strs, self.min_new_tokens
|
||||||
)
|
)
|
||||||
|
|
||||||
# Clear API input aliases so omit_defaults=True drops them from the wire.
|
# Clear API input aliases after normalizing them into internal fields.
|
||||||
self.stop = None
|
self.stop = None
|
||||||
self.stop_regex = None
|
self.stop_regex = None
|
||||||
self.is_normalized = True
|
self.is_normalized = True
|
||||||
|
|||||||
@@ -411,11 +411,6 @@ class TestSamplingParamsMsgspecStruct(CustomTestCase):
|
|||||||
self.assertTrue(sp.spaces_between_special_tokens)
|
self.assertTrue(sp.spaces_between_special_tokens)
|
||||||
self.assertFalse(sp.no_stop_trim)
|
self.assertFalse(sp.no_stop_trim)
|
||||||
|
|
||||||
def test_msgpack_omits_default_fields(self):
|
|
||||||
encoded = msgspec.msgpack.encode(SamplingParams())
|
|
||||||
|
|
||||||
self.assertEqual(msgspec.msgpack.decode(encoded), {})
|
|
||||||
|
|
||||||
def test_msgpack_round_trip_preserves_normalized_state(self):
|
def test_msgpack_round_trip_preserves_normalized_state(self):
|
||||||
tokenizer = MagicMock()
|
tokenizer = MagicMock()
|
||||||
tokenizer.encode.side_effect = lambda s, add_special_tokens=False: {
|
tokenizer.encode.side_effect = lambda s, add_special_tokens=False: {
|
||||||
|
|||||||
Reference in New Issue
Block a user