refactor: make mm_data_mooncake msgpack-native (drop PickleWrapper)

python/sglang/srt/disaggregation/encode_receiver.py

@@ -30,7 +30,11 @@
     get_mooncake_transfer_engine,
 )
 from sglang.srt.environ import envs
-from sglang.srt.managers.io_struct import GenerateReqInput, TokenizedGenerateReqInput
+from sglang.srt.managers.io_struct import (
+    GenerateReqInput,
+    MooncakeMMUrlItem,
+    TokenizedGenerateReqInput,
+)
 from sglang.srt.managers.multimodal_processor import get_mm_processor, import_processors
 from sglang.srt.managers.schedule_batch import Modality, Req
 from sglang.srt.server_args import ServerArgs
@@ -945,7 +949,7 @@ async def _send_encode_and_rdma_request(self):
         for modality in modalities:
             assigned_nums = self.num_items_assigned[modality]
             num_parts = modality_num_parts[modality]
-            mm_data_modality = [d for d in mm_data_all if d["modality"] == modality]
+            mm_data_modality = [d for d in mm_data_all if d.modality == modality]
             cum_num_items = 0
             cum_idx = 0
             for idx, assigned_num in enumerate(assigned_nums):
@@ -957,7 +961,7 @@ async def _send_encode_and_rdma_request(self):
                     {
                         "encoder_idx": idx,
                         "mm_items": [
-                            d["url"]
+                            d.url
                             for d in mm_data_modality[
                                 cum_num_items : cum_num_items + assigned_num
                             ]
@@ -1568,7 +1572,7 @@ async def recv_mm_data(
                 self.context, zmq.PULL, host=self.host
             )
             mm_data = self._extract_url_data(request_obj)
-            modalities = [m.get("modality") for m in mm_data]
+            modalities = [m.modality for m in mm_data]
             logger.info(
                 f"[{req_id}] Sending encode request to E, "
                 f"modalities={modalities}, num_items={len(mm_data)}"
@@ -1916,7 +1920,7 @@ def _assign_items_by_modality(
         Assign multimodal items across encoders by modality with cross-modality load balancing.
 
         Args:
-            mm_data: List of multimodal data items, each with a "modality" key
+            mm_data: List of multimodal data items, each with a modality
             encoder_num: Number of encoders
             random_shuffle: Whether to shuffle the encoder indices
 
@@ -1928,14 +1932,14 @@ def _assign_items_by_modality(
         if random_shuffle:
             random.shuffle(encode_idx)
         # Get unique modalities with order preserved
-        modalities = list(dict.fromkeys(mm_item.get("modality") for mm_item in mm_data))
+        modalities = list(dict.fromkeys(mm_item.modality for mm_item in mm_data))
         # Use OrderedDict to explicitly maintain modality order
         num_items_assigned = OrderedDict()
         current_offset = 0
 
         for modality in modalities:
             mm_data_modality = [
-                mm_item for mm_item in mm_data if mm_item.get("modality") == modality
+                mm_item for mm_item in mm_data if mm_item.modality == modality
             ]
             num_items = len(mm_data_modality)
             if num_items == 0:
@@ -1955,7 +1959,7 @@ def _assign_items_by_modality(
 
         return num_items_assigned
 
-    def _extract_url_data(self, request_obj) -> List[Dict]:
+    def _extract_url_data(self, request_obj) -> List[MooncakeMMUrlItem]:
         def flatten_mm_items(items):
             if not isinstance(items, list):
                 return [items]
@@ -1987,10 +1991,10 @@ def to_raw_url(mm_item):
                 mm_items = flatten_mm_items(mm_items)
                 for mm_item in mm_items:
                     mm_data.append(
-                        {
-                            "url": to_raw_url(mm_item),
-                            "modality": modality,
-                        }
+                        MooncakeMMUrlItem(
+                            url=to_raw_url(mm_item),
+                            modality=modality,
+                        )
                     )
         return mm_data
 
@@ -2080,7 +2084,7 @@ async def encode(
         effective_urls = encode_urls if encode_urls is not None else self.encode_urls
 
         # get unique modalities with order preserved
-        modalities = [mm_item.get("modality") for mm_item in mm_data]
+        modalities = [mm_item.modality for mm_item in mm_data]
         modalities = list(dict.fromkeys(modalities))
         encode_requests = []
 
@@ -2098,7 +2102,7 @@ async def encode(
         for modality in modalities:
             num_items_assigned_modality = num_items_assigned.get(modality)
             mm_data_modality = [
-                mm_item for mm_item in mm_data if mm_item.get("modality") == modality
+                mm_item for mm_item in mm_data if mm_item.modality == modality
             ]
 
             num_parts = modality_num_parts[modality]
@@ -2114,7 +2118,7 @@ async def encode(
                         "encoder_idx": idx,
                         "encoder_url": effective_urls[idx],
                         "mm_items": [
-                            mm_item.get("url")
+                            mm_item.url
                             for mm_item in mm_data_modality[
                                 cum_num_items : cum_num_items + assigned_num
                             ]
@@ -2241,18 +2245,16 @@ async def encode(
 
         effective_urls = encode_urls if encode_urls is not None else self.encode_urls
 
-        # gRPC currently only supports image; flatten new dict formats to simple lists
-        if mm_data and isinstance(mm_data[0], dict):
+        # gRPC currently only supports image; flatten typed multimodal items to simple lists
+        if mm_data and isinstance(mm_data[0], MooncakeMMUrlItem):
             non_image = [
-                item.get("modality")
-                for item in mm_data
-                if item.get("modality") != Modality.IMAGE
+                item.modality for item in mm_data if item.modality != Modality.IMAGE
             ]
             if non_image:
                 raise NotImplementedError(
                     f"gRPC encode only supports IMAGE modality, got: {non_image}"
                 )
-            img_data = [item.get("url") for item in mm_data]
+            img_data = [item.url for item in mm_data]
         else:
             img_data = mm_data
         if isinstance(num_items_assigned, dict):

python/sglang/srt/managers/io_struct.py

@@ -148,6 +148,22 @@ class SessionParams(msgspec.Struct, kw_only=True, array_like=True):
 ]
 
 
+class MooncakeMMUrlItem(msgspec.Struct, array_like=True):
+    """One multimodal item captured at tokenizer dispatch for mooncake encode routing.
+
+    ``url`` is the raw payload from MMReceiverBase._extract_url_data: the URL
+    string (ImageData.url, or dict["url"] when present) in the common case, or
+    the original raw multimodal item (e.g. a dict for processor/precomputed
+    formats, or bytes) when no URL is available. Typed as Any because msgspec
+    cannot express a union mixing str and bytes; the field still round-trips
+    over msgpack since every passthrough payload is a native msgpack type.
+    ``modality`` tags which encoder pool handles it.
+    """
+
+    url: Any
+    modality: Modality
+
+
 @dataclass
 class GenerateReqInput:
     # Request ID(s). If omitted, generated during normalization. For batch
@@ -845,10 +861,9 @@ class TokenizedGenerateReqInput(BaseReq, kw_only=True):
 
     need_wait_for_mm_inputs: Optional[bool] = None
     num_items_assigned: Optional[Dict[Modality, List[int]]] = None
-    # Pickled Optional[List[{"url": MultimodalDataInputItem, "modality": Modality}]]
-    # from MMReceiverBase._extract_url_data. "url" is ImageData.url,
-    # dict["url"] when present, or the original raw multimodal item.
-    mm_data_mooncake: Optional[PickleWrapper] = None
+    # Multimodal URL snapshot from MMReceiverBase._extract_url_data for mooncake.
+    # Each item carries the raw URL/payload and modality captured at tokenizer dispatch.
+    mm_data_mooncake: Optional[List[MooncakeMMUrlItem]] = None
     # Encoder URL snapshot frozen at tokenizer-side dispatch time so that
     # encoder_idx assignments stay consistent in the scheduler subprocess.
     # Internal IPC only.
@@ -863,12 +878,10 @@ class TokenizedGenerateReqInput(BaseReq, kw_only=True):
 
     def wrap_pickle_fields(self):
         self.mm_inputs = wrap_as_pickle(self.mm_inputs)
-        self.mm_data_mooncake = wrap_as_pickle(self.mm_data_mooncake)
         self.time_stats = wrap_as_pickle(self.time_stats)
 
     def unwrap_pickle_fields(self):
         self.mm_inputs = unwrap_from_pickle(self.mm_inputs)
-        self.mm_data_mooncake = unwrap_from_pickle(self.mm_data_mooncake)
         self.time_stats = unwrap_from_pickle(self.time_stats)
 
 

test/registered/unit/managers/test_io_struct_msgpack.py

@@ -0,0 +1,104 @@
+import unittest
+
+from sglang.srt.managers.io_struct import (
+    MooncakeMMUrlItem,
+    TokenizedGenerateReqInput,
+    msgpack_decode,
+    msgpack_encode,
+)
+from sglang.srt.managers.schedule_batch import Modality
+from sglang.srt.sampling.sampling_params import SamplingParams
+from sglang.test.ci.ci_register import (
+    register_amd_ci,
+    register_cpu_ci,
+    register_cuda_ci,
+)
+
+register_cuda_ci(est_time=8, stage="base-b", runner_config="1-gpu-large")
+register_amd_ci(est_time=8, suite="stage-b-test-1-gpu-small-amd")
+register_cpu_ci(est_time=8, suite="base-c-test-cpu")
+
+
+class TestTokenizedGenerateReqInputMsgpack(unittest.TestCase):
+    def _make_request(self, mm_data_mooncake):
+        return TokenizedGenerateReqInput(
+            input_text="",
+            input_ids=None,
+            input_embeds=None,
+            mm_inputs=None,
+            token_type_ids=None,
+            sampling_params=SamplingParams(),
+            return_logprob=False,
+            logprob_start_len=0,
+            top_logprobs_num=0,
+            token_ids_logprob=None,
+            stream=False,
+            mm_data_mooncake=mm_data_mooncake,
+        )
+
+    def _round_trip(self, req):
+        req.wrap_pickle_fields()
+        decoded = msgpack_decode(msgpack_encode(req))
+        decoded.unwrap_pickle_fields()
+        return decoded
+
+    def test_mm_data_mooncake_round_trips_populated(self):
+        items = [
+            MooncakeMMUrlItem(url="http://x/a.jpg", modality=Modality.IMAGE),
+            MooncakeMMUrlItem(url="http://x/b.mp4", modality=Modality.VIDEO),
+        ]
+
+        decoded = self._round_trip(self._make_request(items))
+
+        self.assertEqual(decoded.mm_data_mooncake, items)
+        self.assertEqual(decoded.mm_data_mooncake[0].url, "http://x/a.jpg")
+        self.assertEqual(decoded.mm_data_mooncake[0].modality, Modality.IMAGE)
+        self.assertEqual(decoded.mm_data_mooncake[1].url, "http://x/b.mp4")
+        self.assertEqual(decoded.mm_data_mooncake[1].modality, Modality.VIDEO)
+
+    def test_mm_data_mooncake_round_trips_dict_payload(self):
+        # _extract_url_data passes raw dict multimodal items through unchanged
+        # when they carry no "url" key, so the field must survive non-string
+        # payloads over msgpack.
+        items = [
+            MooncakeMMUrlItem(
+                url={"image": "data:base64", "detail": "high"},
+                modality=Modality.IMAGE,
+            ),
+        ]
+
+        decoded = self._round_trip(self._make_request(items))
+
+        self.assertEqual(
+            decoded.mm_data_mooncake[0].url, {"image": "data:base64", "detail": "high"}
+        )
+        self.assertEqual(decoded.mm_data_mooncake[0].modality, Modality.IMAGE)
+
+    def test_mm_data_mooncake_round_trips_none(self):
+        decoded = self._round_trip(self._make_request(None))
+
+        self.assertIsNone(decoded.mm_data_mooncake)
+
+    def test_mm_data_mooncake_round_trips_empty_list(self):
+        decoded = self._round_trip(self._make_request([]))
+
+        self.assertEqual(decoded.mm_data_mooncake, [])
+
+    def test_mm_data_mooncake_round_trips_all_modalities(self):
+        for modality in (Modality.IMAGE, Modality.VIDEO, Modality.AUDIO):
+            with self.subTest(modality=modality):
+                items = [
+                    MooncakeMMUrlItem(
+                        url=f"http://x/{modality.name.lower()}",
+                        modality=modality,
+                    )
+                ]
+
+                decoded = self._round_trip(self._make_request(items))
+
+                self.assertEqual(decoded.mm_data_mooncake[0].url, items[0].url)
+                self.assertEqual(decoded.mm_data_mooncake[0].modality, modality)
+
+
+if __name__ == "__main__":
+    unittest.main()