[Refactor] msgspec standardisation for data entry key names and improved type checks (#3149)

Signed-off-by: Divyansh Singhvi <divyanshsinghvi@gmail.com>
Signed-off-by: dsinghvi <divyanshsinghvi@gmail.com>

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),