Enforce strict data types for record.py. Bump default worker init timeout to 60

Author: nv-alicheng

scripts/zmq_pubsub_demo.py

@@ -153,13 +153,11 @@ async def publish_test_events(publisher) -> None:
             event_type=SampleEventType.RECV_FIRST,
             timestamp_ns=10010,
             sample_uuid=uuid1,
-            data={"ttft_ms": 10.0},
         ),
         EventRecord(
             event_type=SampleEventType.RECV_FIRST,
             timestamp_ns=10190,
             sample_uuid=uuid2,
-            data={"ttft_ms": 187.0},
         ),
         EventRecord(
             event_type=SampleEventType.RECV_NON_FIRST,
@@ -190,7 +188,7 @@ async def publish_test_events(publisher) -> None:
             event_type=SampleEventType.COMPLETE,
             timestamp_ns=10211,
             sample_uuid=uuid1,
-            data={"tokens": 50},
+            data="Hello world",
         ),
         EventRecord(
             event_type=SampleEventType.RECV_NON_FIRST,
@@ -211,7 +209,7 @@ async def publish_test_events(publisher) -> None:
             event_type=SampleEventType.COMPLETE,
             timestamp_ns=10219,
             sample_uuid=uuid2,
-            data={"tokens": 75},
+            data="Sample output for uuid2",
         ),
     ]
 

src/inference_endpoint/async_utils/services/event_logger/sql_writer.py

@@ -17,7 +17,7 @@
 
 from pathlib import Path
 
-import orjson
+import msgspec
 from inference_endpoint.core.record import EventRecord
 from sqlalchemy import BigInteger, Integer, LargeBinary, String, create_engine
 from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, sessionmaker
@@ -55,17 +55,13 @@ class EventRowModel(Base):
 
 
 def _record_to_row(record: EventRecord) -> EventRowModel:
-    """Convert an EventRecord to an EventRowModel using EventType topic strings."""
-    data_bytes = b""
-    if record.data:
-        data_bytes = orjson.dumps(record.data)
     # event_type.topic is set by EventTypeMeta on each enum member
     topic = record.event_type.topic  # type: ignore[attr-defined]
     return EventRowModel(
         sample_uuid=record.sample_uuid,
         event_type=topic,
         timestamp_ns=record.timestamp_ns,
-        data=data_bytes,
+        data=msgspec.json.encode(record.data),
     )
 
 

src/inference_endpoint/async_utils/transport/protocol.py

@@ -35,7 +35,7 @@
     decode_event_record,
     encode_event_record,
 )
-from inference_endpoint.core.types import Query, QueryResult, StreamChunk
+from inference_endpoint.core.types import ErrorData, Query, QueryResult, StreamChunk
 
 if TYPE_CHECKING:
     from inference_endpoint.async_utils.transport.zmq.context import ManagedZMQContext
@@ -341,14 +341,12 @@ def _on_readable(self) -> None:
                 try:
                     event_record = decode_event_record(payload)
                 except msgspec.DecodeError as e:
-                    # Record an error instead
-                    # TODO: Make `data` field more rigidly typed
                     event_record = EventRecord(
                         event_type=ErrorEventType.GENERIC,
-                        data={
-                            "error_type": "msgspec.DecodeError",
-                            "error_message": str(e),
-                        },
+                        data=ErrorData(
+                            error_type="msgspec.DecodeError",
+                            error_message=str(e),
+                        ),
                     )
                 records.append(event_record)
         except StopIteration:

src/inference_endpoint/core/record.py

@@ -19,6 +19,8 @@
 
 import msgspec
 
+from .types import OUTPUT_TYPE, ErrorData
+
 TOPIC_FRAME_SIZE: Final[int] = 40
 """int: Fixed bytesize for the encoded topic string. PUB messages will be prefixed by a
 topic string corresponding to the EventType. This topic will be null-padded to this fixed
@@ -151,7 +153,7 @@ class EventRecord(msgspec.Struct, kw_only=True):  # type: ignore[call-arg]
     event_type: EventType
     timestamp_ns: int = msgspec.field(default_factory=time.monotonic_ns)
     sample_uuid: str = ""
-    data: dict[str, Any] = msgspec.field(default_factory=dict)
+    data: OUTPUT_TYPE | ErrorData | None = None
 
 
 _ENCODER = msgspec.msgpack.Encoder(enc_hook=EventType.encode_hook)

src/inference_endpoint/core/types.py

@@ -51,9 +51,6 @@ class QueryStatus(Enum):
 OUTPUT_ELEM_TYPE = str | tuple[str, ...]
 """Type for a single output or reasoning value: string (non-streaming) or tuple of strings (streaming)."""
 
-_OUTPUT_DICT_TYPE = dict[str, str | list[str]]
-_OUTPUT_RESULT_TYPE = str | tuple[str, ...] | _OUTPUT_DICT_TYPE | None
-
 
 class TextModelOutput(msgspec.Struct, tag=True, kw_only=True):  # type: ignore[call-arg]
     """Structured output from a text model.

src/inference_endpoint/endpoint_client/config.py

@@ -71,7 +71,7 @@ class HTTPClientConfig:
     cpu_affinity: AffinityPlan | None = None
 
     # Worker lifecycle timeouts
-    worker_initialization_timeout: float = 40.0  # init
+    worker_initialization_timeout: float = 60.0  # init
     worker_graceful_shutdown_wait: float = 0.5  # post-run
     worker_force_kill_timeout: float = 0.5  # post-run
 

tests/unit/async_utils/test_event_publisher.py

@@ -151,7 +151,7 @@ async def test_publish_sends_data_on_ipc_socket(
         record = EventRecord(
             event_type=SessionEventType.STARTED,
             sample_uuid="",
-            data={"manual_socket": True},
+            data="manual_socket_test",
         )
         event_publisher_service.publish(record)
         # Yield so the publisher's event loop can drain the send buffer
@@ -168,7 +168,7 @@ async def test_publish_sends_data_on_ipc_socket(
         assert topic_bytes == b"session.started"
         rec = decode_event_record(bytes(payload))
         assert rec.event_type.value == SessionEventType.STARTED.value
-        assert rec.data == {"manual_socket": True}
+        assert rec.data == "manual_socket_test"
         # Socket is closed by ManagedZMQContext.cleanup() in ev_pub_zmq_context fixture teardown.
 
     def test_singleton_returns_same_instance(
@@ -190,15 +190,15 @@ async def test_publish_session_event_received_by_subscriber(
         record = EventRecord(
             event_type=SessionEventType.STARTED,
             sample_uuid="",
-            data={"key": "value"},
+            data="value",
         )
         event_publisher_service.publish(record)
         await asyncio.sleep(0.05)  # Let publisher drain send buffer
         await asyncio.wait_for(received_event.wait(), timeout=_WAIT_RECORDS_TIMEOUT)
         assert len(collecting_subscriber.received) == 1
         rec = collecting_subscriber.received[0]
         assert rec.event_type.value == SessionEventType.STARTED.value
-        assert rec.data == {"key": "value"}
+        assert rec.data == "value"
 
     @pytest.mark.asyncio
     async def test_publish_sample_event_received_by_subscriber(
@@ -210,7 +210,7 @@ async def test_publish_sample_event_received_by_subscriber(
         record = EventRecord(
             event_type=SampleEventType.COMPLETE,
             sample_uuid="sample-1",
-            data={"latency_ns": 42},
+            data="sample output",
         )
         event_publisher_service.publish(record)
         await asyncio.sleep(0.05)  # Let publisher drain send buffer
@@ -219,7 +219,7 @@ async def test_publish_sample_event_received_by_subscriber(
         rec = collecting_subscriber.received[0]
         assert rec.event_type.value == SampleEventType.COMPLETE.value
         assert rec.sample_uuid == "sample-1"
-        assert rec.data == {"latency_ns": 42}
+        assert rec.data == "sample output"
 
     @pytest.mark.asyncio
     async def test_multiple_events_received_in_order(
@@ -232,12 +232,11 @@ async def test_multiple_events_received_in_order(
             record = EventRecord(
                 event_type=SampleEventType.ISSUED,
                 sample_uuid=f"sample-{i}",
-                data={"seq": i},
             )
             event_publisher_service.publish(record)
             await asyncio.sleep(0.05)  # Small delay between events (demo-style)
         await asyncio.wait_for(received_event.wait(), timeout=_WAIT_RECORDS_TIMEOUT)
         assert len(collecting_subscriber.received) == 3
         for i in range(3):
             assert collecting_subscriber.received[i].sample_uuid == f"sample-{i}"
-            assert collecting_subscriber.received[i].data.get("seq") == i
+            assert collecting_subscriber.received[i].data is None

tests/unit/core/test_record.py

@@ -17,7 +17,6 @@
 
 import time
 
-import msgspec
 import pytest
 from inference_endpoint.core.record import (
     TOPIC_FRAME_SIZE,
@@ -29,6 +28,7 @@
     decode_event_record,
     encode_event_record,
 )
+from inference_endpoint.core.types import ErrorData, TextModelOutput
 
 
 class TestEventType:
@@ -64,8 +64,7 @@ def test_construction_with_only_event_type_uses_defaults(self):
         after = time.monotonic_ns()
         assert before <= record.timestamp_ns <= after
         assert record.sample_uuid == ""
-        assert record.data == {}
-        assert isinstance(record.data, dict)
+        assert record.data is None
 
 
 class TestEncodeEventRecord:
@@ -76,17 +75,16 @@ def test_returns_tuple_of_topic_bytes_padded_and_payload_bytes_with_valid_msgpac
         record = EventRecord(
             event_type=SampleEventType.ISSUED,
             sample_uuid="test-uuid",
-            data={"key": "value"},
+            data="test-output",
         )
         topic_bytes, payload = encode_event_record(record)
         assert isinstance(topic_bytes, bytes)
         assert len(topic_bytes) == TOPIC_FRAME_SIZE
         assert topic_bytes.rstrip(b"\x00") == b"sample.issued"
         assert isinstance(payload, bytes)
-        decoded = msgspec.msgpack.decode(payload)
-        assert isinstance(decoded, dict)
-        assert decoded.get("sample_uuid") == "test-uuid"
-        assert decoded.get("data") == {"key": "value"}
+        decoded = decode_event_record(payload)
+        assert decoded.sample_uuid == "test-uuid"
+        assert decoded.data == "test-output"
 
     def test_topic_bytes_padded_matches_event_type_for_session_sample_error(self):
         """Topic is null-padded to TOPIC_FRAME_SIZE for single-frame ZMQ sends."""
@@ -106,37 +104,55 @@ def test_session_event_round_trips_with_all_fields(self):
         record = EventRecord(
             event_type=SessionEventType.STARTED,
             sample_uuid="sess-1",
-            data={"session_id": "abc"},
         )
         _, payload = encode_event_record(record)
         decoded = decode_event_record(payload)
         assert decoded.event_type.topic == SessionEventType.STARTED.topic
         assert decoded.sample_uuid == "sess-1"
-        assert decoded.data == {"session_id": "abc"}
+        assert decoded.data is None
         assert isinstance(decoded.timestamp_ns, int)
         assert decoded.timestamp_ns == record.timestamp_ns
 
-    def test_sample_event_round_trips(self):
+    def test_sample_event_round_trips_with_output(self):
         record = EventRecord(
             event_type=SampleEventType.COMPLETE,
             sample_uuid="sample-42",
-            data={"latency_ns": 1000},
+            data="output text",
         )
         _, payload = encode_event_record(record)
         decoded = decode_event_record(payload)
         assert decoded.event_type.topic == SampleEventType.COMPLETE.topic
         assert decoded.sample_uuid == "sample-42"
-        assert decoded.data == {"latency_ns": 1000}
+        assert decoded.data == "output text"
 
-    def test_error_event_round_trips_with_defaults(self):
+    def test_sample_event_round_trips_with_text_model_output(self):
+        record = EventRecord(
+            event_type=SampleEventType.COMPLETE,
+            sample_uuid="sample-42",
+            data=TextModelOutput(output="out", reasoning="reason"),
+        )
+        _, payload = encode_event_record(record)
+        decoded = decode_event_record(payload)
+        assert decoded.event_type.topic == SampleEventType.COMPLETE.topic
+        assert decoded.sample_uuid == "sample-42"
+        assert isinstance(decoded.data, TextModelOutput)
+        assert decoded.data.output == "out"
+        assert decoded.data.reasoning == "reason"
+
+    def test_error_event_round_trips_with_error_data(self):
         record = EventRecord(
             event_type=ErrorEventType.LOADGEN,
-            data={"message": "error details"},
+            data=ErrorData(
+                error_type="LoadgenError",
+                error_message="error details",
+            ),
         )
         _, payload = encode_event_record(record)
         decoded = decode_event_record(payload)
         assert decoded.event_type.topic == ErrorEventType.LOADGEN.topic
-        assert decoded.data == {"message": "error details"}
+        assert isinstance(decoded.data, ErrorData)
+        assert decoded.data.error_type == "LoadgenError"
+        assert decoded.data.error_message == "error details"
         assert decoded.sample_uuid == ""
 
     def test_record_with_only_event_type_round_trips_with_defaults(self):
@@ -145,7 +161,7 @@ def test_record_with_only_event_type_round_trips_with_defaults(self):
         decoded = decode_event_record(payload)
         assert decoded.event_type.topic == SessionEventType.ENDED.topic
         assert decoded.sample_uuid == ""
-        assert decoded.data == {}
+        assert decoded.data is None
         assert decoded.timestamp_ns > 0
 
     def test_explicit_timestamp_ns_preserved_round_trip(self):
@@ -157,12 +173,3 @@ def test_explicit_timestamp_ns_preserved_round_trip(self):
         _, payload = encode_event_record(record)
         decoded = decode_event_record(payload)
         assert decoded.timestamp_ns == ts
-
-    def test_nested_and_list_data_round_trips(self):
-        record = EventRecord(
-            event_type=SampleEventType.TRANSPORT_RECV,
-            data={"nested": {"a": 1}, "list": [1, 2, 3]},
-        )
-        _, payload = encode_event_record(record)
-        decoded = decode_event_record(payload)
-        assert decoded.data == {"nested": {"a": 1}, "list": [1, 2, 3]}