[Refactor] msgspec standardisation for data entry key names and improved type checks

tests/core/sched/test_chunk_scheduling_coordinator.py

@@ -222,7 +222,7 @@ def test_generation_mode(self):
         self.assertEqual(req.prompt_token_ids, [10, 20, 30])
         self.assertEqual(req.num_computed_tokens, 0)
         self.assertIsNone(req.additional_information)
-        self.assertEqual(req._omni_initial_model_buffer, {"left_context_size": 25})
+        self.assertEqual(req._omni_initial_model_buffer, {"meta": {"left_context_size": 25}})
 
 
 class TestChunkCoordinatorPostprocess(unittest.TestCase):

tests/distributed/omni_connectors/test_chunk_transfer_adapter.py

@@ -12,7 +12,7 @@
 from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler
 from vllm.v1.request import RequestStatus
 
-from vllm_omni.data_entry_keys import OmniPayload
+from vllm_omni.data_entry_keys import CodesStruct, MetaStruct, OmniPayload, OmniPayloadStruct
 from vllm_omni.distributed.omni_connectors.transfer_adapter.base import OmniTransferAdapterBase
 from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter import (
     OmniChunkTransferAdapter,
@@ -140,11 +140,54 @@ def test_save_async(build_adapter):
     assert task["is_finished"] is False
 
 
+def test_send_single_request_struct_without_meta_does_not_crash(build_adapter, monkeypatch):
+    """Producer may return a struct with ``meta=None`` (e.g. payload that
+    carries only ``embed`` or ``codes``). The sender's ``meta is not None``
+    guard handles this without AttributeError; ``finished_flag`` is None and
+    the cleanup path is not triggered.
+    """
+    adapter, _ = build_adapter(stage_id=1)
+    request = _req("req-no-meta", RequestStatus.WAITING, external_req_id="ext-no-meta")
+
+    adapter.custom_process_next_stage_input_func = lambda **kwargs: OmniPayloadStruct(
+        codes=CodesStruct(audio=torch.tensor([1, 2], dtype=torch.long)),
+    )
+    cleanup_calls = []
+    monkeypatch.setattr(adapter, "cleanup", lambda *a, **kw: cleanup_calls.append((a, kw)))
+
+    adapter._send_single_request({"pooling_output": None, "request": request, "is_finished": False})
+
+    assert cleanup_calls == []  # no terminal cleanup; meta.finished is unobservable
+
+
+def test_send_single_request_empty_struct_goes_on_wire(build_adapter, monkeypatch):
+    """Pin the contract: an explicitly empty ``OmniPayloadStruct()`` passes
+    the ``payload_data is None`` check and gets sent. To skip a chunk, the
+    producer must return ``None``, not an empty struct. (Filtering empty
+    structs at the adapter would require introspecting all struct fields on
+    every send and was rejected for cost vs. value.)
+    """
+    adapter, connector = build_adapter(stage_id=1)
+    request = _req("req-empty", RequestStatus.WAITING, external_req_id="ext-empty")
+
+    adapter.custom_process_next_stage_input_func = lambda **kwargs: OmniPayloadStruct()
+    monkeypatch.setattr(adapter, "cleanup", lambda *a, **kw: None)
+
+    adapter._send_single_request({"pooling_output": None, "request": request, "is_finished": False})
+
+    assert connector.put.called
+    sent_payload = connector.put.call_args.kwargs["data"]
+    assert isinstance(sent_payload, OmniPayloadStruct)
+    assert sent_payload.meta is None  # confirms it's the empty struct on the wire
+
+
 def test_send_single_request_cleans_up_after_finished_payload(build_adapter, monkeypatch):
     adapter, _ = build_adapter(stage_id=1)
     request = _req("req-finished", RequestStatus.FINISHED_STOPPED, external_req_id="ext-finished")
 
-    adapter.custom_process_next_stage_input_func = lambda **kwargs: {"x": [1], "finished": True}
+    adapter.custom_process_next_stage_input_func = lambda **kwargs: OmniPayloadStruct(
+        meta=MetaStruct(finished=torch.tensor(True, dtype=torch.bool))
+    )
     cleanup_calls = []
     monkeypatch.setattr(adapter, "cleanup", lambda *a, **kw: cleanup_calls.append((a, kw)))
 
@@ -604,3 +647,42 @@ def _super_finish(_self, request_ids, finished_status):
         OmniARScheduler.finish_requests(sched, ["r1"], RequestStatus.FINISHED_ABORTED)
 
     assert order == ["adapter", "super"]
+
+
+def test_wire_round_trip_struct_to_dict_contract():
+    """Pin the wire contract: encoding ``OmniPayloadStruct`` and decoding it
+    yields a dict equivalent to ``to_dict(struct)``.
+
+    The chunk-adapter sender uses struct attribute access while the receiver
+    uses dict-key access. This works only because ``OmniMsgpackDecoder`` has
+    no target type and decodes structs back to plain dicts. If this test
+    breaks, the receiver's dict access will silently drop fields or KeyError.
+    """
+    from vllm_omni.data_entry_keys import CodesStruct, to_dict
+    from vllm_omni.distributed.omni_connectors.utils.serialization import (
+        OmniMsgpackDecoder,
+        OmniMsgpackEncoder,
+    )
+
+    struct = OmniPayloadStruct(
+        meta=MetaStruct(
+            finished=torch.tensor(True, dtype=torch.bool),
+            left_context_size=12,
+        ),
+        codes=CodesStruct(audio=torch.tensor([1, 2, 3], dtype=torch.int64)),
+    )
+
+    encoded = OmniMsgpackEncoder().encode(struct)
+    decoded = OmniMsgpackDecoder().decode(encoded)
+
+    assert isinstance(decoded, dict)
+    assert isinstance(decoded["meta"], dict)
+    assert isinstance(decoded["meta"]["finished"], torch.Tensor)
+    assert bool(decoded["meta"]["finished"].item()) is True
+    assert decoded["meta"]["left_context_size"] == 12
+    assert torch.equal(decoded["codes"]["audio"], torch.tensor([1, 2, 3], dtype=torch.int64))
+
+    expected = to_dict(struct)
+    assert set(decoded.keys()) == set(expected.keys())
+    assert set(decoded["meta"].keys()) == set(expected["meta"].keys())
+    assert set(decoded["codes"].keys()) == set(expected["codes"].keys())

tests/model_executor/models/cosyvoice3/test_cosyvoice3_model_helpers.py

@@ -164,11 +164,12 @@ def test_forward_prefers_token_offset_when_present():
 
     runtime_info = [
         {
-            "speech_token": torch.tensor([[1, 2, 3]], dtype=torch.long),
-            "speech_feat": torch.tensor([[[0.1, 0.2], [0.3, 0.4]]], dtype=torch.float32),
-            "embedding": torch.tensor([[0.5, 0.6]], dtype=torch.float32),
-            "token_offset": 2,
-            "left_context_size": 1,
+            "embed": {
+                "speech_token": torch.tensor([[1, 2, 3]], dtype=torch.long),
+                "speech_feat": torch.tensor([[[0.1, 0.2], [0.3, 0.4]]], dtype=torch.float32),
+                "embedding": torch.tensor([[0.5, 0.6]], dtype=torch.float32),
+            },
+            "meta": {"left_context_size": 2},
         }
     ]
 
@@ -193,10 +194,12 @@ def test_forward_falls_back_to_left_context_size_for_backward_compat():
 
     runtime_info = [
         {
-            "speech_token": torch.tensor([[1, 2, 3]], dtype=torch.long),
-            "speech_feat": torch.tensor([[[0.1, 0.2], [0.3, 0.4]]], dtype=torch.float32),
-            "embedding": torch.tensor([[0.5, 0.6]], dtype=torch.float32),
-            "left_context_size": 2,
+            "embed": {
+                "speech_token": torch.tensor([[1, 2, 3]], dtype=torch.long),
+                "speech_feat": torch.tensor([[[0.1, 0.2], [0.3, 0.4]]], dtype=torch.float32),
+                "embedding": torch.tensor([[0.5, 0.6]], dtype=torch.float32),
+            },
+            "meta": {"left_context_size": 2},
         }
     ]
 
@@ -214,10 +217,12 @@ def test_forward_ignores_single_request_padded_tail_tokens():
     model = _make_code2wav_model(with_stride_cfg=True)
     runtime_info = [
         {
-            "speech_token": torch.tensor([[1, 2, 3]], dtype=torch.long),
-            "speech_feat": torch.tensor([[[0.1, 0.2], [0.3, 0.4]]], dtype=torch.float32),
-            "embedding": torch.tensor([[0.5, 0.6]], dtype=torch.float32),
-            "token_offset": 0,
+            "embed": {
+                "speech_token": torch.tensor([[1, 2, 3]], dtype=torch.long),
+                "speech_feat": torch.tensor([[[0.1, 0.2], [0.3, 0.4]]], dtype=torch.float32),
+                "embedding": torch.tensor([[0.5, 0.6]], dtype=torch.float32),
+            },
+            "meta": {"left_context_size": 0},
         }
     ]
 
@@ -238,10 +243,12 @@ def test_forward_uses_non_stream_decode_without_chunk_metadata():
 
     runtime_info = [
         {
-            "speech_token": torch.tensor([[1, 2, 3]], dtype=torch.long),
-            "speech_feat": torch.tensor([[[0.1, 0.2], [0.3, 0.4]]], dtype=torch.float32),
-            "embedding": torch.tensor([[0.5, 0.6]], dtype=torch.float32),
-            "prefix_ids": [101, 102],
+            "embed": {
+                "speech_token": torch.tensor([[1, 2, 3]], dtype=torch.long),
+                "speech_feat": torch.tensor([[[0.1, 0.2], [0.3, 0.4]]], dtype=torch.float32),
+                "embedding": torch.tensor([[0.5, 0.6]], dtype=torch.float32),
+            },
+            "ids": {"prompt": [101, 102]},
             "generated_len": 3,
         }
     ]
@@ -275,12 +282,16 @@ def test_forward_reuses_streaming_cache_state_between_chunks():
     )
     runtime_info = [
         {
-            "req_id": ["rid-stream"],
-            "speech_token": torch.tensor([[1, 2, 3]], dtype=torch.long),
-            "speech_feat": torch.tensor([[[0.1, 0.2], [0.3, 0.4]]], dtype=torch.float32),
-            "embedding": torch.tensor([[0.5, 0.6]], dtype=torch.float32),
-            "token_offset": 0,
-            "stream_finished": torch.tensor(False),
+            "embed": {
+                "speech_token": torch.tensor([[1, 2, 3]], dtype=torch.long),
+                "speech_feat": torch.tensor([[[0.1, 0.2], [0.3, 0.4]]], dtype=torch.float32),
+                "embedding": torch.tensor([[0.5, 0.6]], dtype=torch.float32),
+            },
+            "meta": {
+                "req_id": ["rid-stream"],
+                "stream_finished": torch.tensor(False),
+                "left_context_size": 0,
+            },
         }
     ]
 
@@ -321,12 +332,16 @@ def test_forward_clears_streaming_cache_on_terminal_chunk():
     )
     runtime_info = [
         {
-            "req_id": ["rid-stream"],
-            "speech_token": torch.tensor([[1, 2, 3]], dtype=torch.long),
-            "speech_feat": torch.tensor([[[0.1, 0.2], [0.3, 0.4]]], dtype=torch.float32),
-            "embedding": torch.tensor([[0.5, 0.6]], dtype=torch.float32),
-            "token_offset": 0,
-            "stream_finished": torch.tensor(False),
+            "embed": {
+                "speech_token": torch.tensor([[1, 2, 3]], dtype=torch.long),
+                "speech_feat": torch.tensor([[[0.1, 0.2], [0.3, 0.4]]], dtype=torch.float32),
+                "embedding": torch.tensor([[0.5, 0.6]], dtype=torch.float32),
+            },
+            "meta": {
+                "req_id": ["rid-stream"],
+                "stream_finished": torch.tensor(False),
+                "left_context_size": 0,
+            },
         }
     ]
 
@@ -338,7 +353,7 @@ def test_forward_clears_streaming_cache_on_terminal_chunk():
     )
     assert "rid-stream" in model._stream_vocoder_cache_by_req
 
-    runtime_info[0]["stream_finished"] = torch.tensor(True)
+    runtime_info[0]["meta"]["stream_finished"] = torch.tensor(True)
     out = model.forward(
         input_ids=torch.tensor([0, 1, 2], dtype=torch.long),
         positions=torch.tensor([0, 1, 2], dtype=torch.long),

tests/model_executor/stage_input_processors/test_cosyvoice3_stage_input_processors.py

@@ -64,9 +64,11 @@ def test_talker2code2wav_async_chunk_final_payload_uses_absolute_token_offset():
         external_req_id="rid-0",
         output_token_ids=[1, 2, 6562, 3],
         additional_information={
-            "speech_token": [torch.tensor([[11, 12, 13]])],
-            "speech_feat": [torch.tensor([[[0.1, 0.2], [0.3, 0.4]]])],
-            "embedding": [torch.tensor([[0.5, 0.6]])],
+            "embed": {
+                "speech_token": [torch.tensor([[11, 12, 13]])],
+                "speech_feat": [torch.tensor([[[0.1, 0.2], [0.3, 0.4]]])],
+                "embedding": [torch.tensor([[0.5, 0.6]])],
+            },
         },
         is_finished=lambda: True,
     )
@@ -79,15 +81,14 @@ def test_talker2code2wav_async_chunk_final_payload_uses_absolute_token_offset():
     )
 
     assert payload is not None
-    assert payload["meta"]["finished"].item() is True
-    assert payload["codes"]["audio"] == [1, 2, 3]
-    assert payload["token_offset"] == 0
-    assert payload["left_context_size"] == 0
-    assert payload["req_id"] == ["rid-0"]
-    assert payload["stream_finished"].item() is True
-    assert "speech_token" in payload
-    assert "speech_feat" in payload
-    assert "embedding" in payload
+    assert payload.meta.finished.item() is True
+    assert payload.codes.audio.tolist() == [1, 2, 3]
+    assert payload.meta.left_context_size == 0
+    assert payload.meta.req_id == ["rid-0"]
+    assert payload.meta.stream_finished.item() is True
+    assert payload.embed.speech_token is not None
+    assert payload.embed.speech_feat is not None
+    assert payload.embed.embedding is not None
 
 
 def test_talker2code2wav_async_chunk_emits_eof_when_finished_without_valid_codes():
@@ -107,8 +108,8 @@ def test_talker2code2wav_async_chunk_emits_eof_when_finished_without_valid_codes
     )
 
     assert payload is not None
-    assert payload["codes"]["audio"] == []
-    assert payload["meta"]["finished"].item() is True
+    assert payload.codes.audio.tolist() == []
+    assert payload.meta.finished.item() is True
 
 
 def test_talker2code2wav_async_chunk_does_not_reemit_without_new_tokens():
@@ -134,8 +135,8 @@ def test_talker2code2wav_async_chunk_does_not_reemit_without_new_tokens():
     )
 
     assert payload1 is not None
-    assert payload1["codes"]["audio"] == [1, 2]
-    assert payload1["token_offset"] == 0
+    assert payload1.codes.audio.tolist() == [1, 2]
+    assert payload1.meta.left_context_size == 0
     assert payload2 is None
 
 
@@ -164,9 +165,9 @@ def test_talker2code2wav_async_chunk_waits_for_prelookahead_and_emits_cumulative
 
     assert payload_pending is None
     assert payload_ready is not None
-    assert payload_ready["codes"]["audio"] == [1, 2, 3]
-    assert payload_ready["token_offset"] == 0
-    assert payload_ready["meta"]["finished"].item() is False
+    assert payload_ready.codes.audio.tolist() == [1, 2, 3]
+    assert payload_ready.meta.left_context_size == 0
+    assert payload_ready.meta.finished.item() is False
 
 
 def test_talker2code2wav_async_chunk_final_flush_uses_previous_token_offset():
@@ -193,13 +194,13 @@ def test_talker2code2wav_async_chunk_final_flush_uses_previous_token_offset():
     )
 
     assert payload_stream is not None
-    assert payload_stream["meta"]["finished"].item() is False
-    assert payload_stream["codes"]["audio"] == [3, 4, 5]
-    assert payload_stream["token_offset"] == 0
+    assert payload_stream.meta.finished.item() is False
+    assert payload_stream.codes.audio.tolist() == [3, 4, 5]
+    assert payload_stream.meta.left_context_size == 0
     assert payload_final is not None
-    assert payload_final["meta"]["finished"].item() is True
-    assert payload_final["codes"]["audio"] == [3, 4, 5, 6]
-    assert payload_final["token_offset"] == 2
+    assert payload_final.meta.finished.item() is True
+    assert payload_final.codes.audio.tolist() == [3, 4, 5, 6]
+    assert payload_final.meta.left_context_size == 2
 
 
 def test_talker2code2wav_async_chunk_respects_prompt_token_pad_on_first_chunk():
@@ -208,7 +209,7 @@ def test_talker2code2wav_async_chunk_respects_prompt_token_pad_on_first_chunk():
         external_req_id="rid-pad",
         output_token_ids=[8, 9, 10],
         additional_information={
-            "speech_token": [torch.tensor([[1, 2, 3]])],
+            "embed": {"speech_token": [torch.tensor([[1, 2, 3]])]},
         },
         is_finished=lambda: False,
     )
@@ -229,8 +230,8 @@ def test_talker2code2wav_async_chunk_respects_prompt_token_pad_on_first_chunk():
 
     assert payload_pending is None
     assert payload_ready is not None
-    assert payload_ready["codes"]["audio"] == [8, 9, 10, 11]
-    assert payload_ready["token_offset"] == 0
+    assert payload_ready.codes.audio.tolist() == [8, 9, 10, 11]
+    assert payload_ready.meta.left_context_size == 0
 
 
 def test_talker2code2wav_async_chunk_emits_terminal_eof_without_duplicate_audio():
@@ -256,8 +257,8 @@ def test_talker2code2wav_async_chunk_emits_terminal_eof_without_duplicate_audio(
     )
 
     assert payload_stream is not None
-    assert payload_stream["meta"]["finished"].item() is False
-    assert payload_stream["codes"]["audio"] == [3, 4]
+    assert payload_stream.meta.finished.item() is False
+    assert payload_stream.codes.audio.tolist() == [3, 4]
     assert payload_final is not None
-    assert payload_final["meta"]["finished"].item() is True
-    assert payload_final["codes"]["audio"] == []
+    assert payload_final.meta.finished.item() is True
+    assert payload_final.codes.audio.tolist() == []