refactor: generalize ZMQ pub/sub over message type via MessageCodec

[nv-alicheng]

Summary

• Replaces EventRecord-specific publisher/subscriber classes with generic ZmqMessagePublisher[T] / ZmqMessageSubscriber[T] parameterized by a MessageCodec[T] Protocol.
• EventRecordCodec preserves the existing wire format and decode-error wrapping behavior, so all current callers (EventPublisherService, EventLoggerService, MetricsAggregatorService) keep their semantics — they just thread the codec through super().__init__.
• Exposes sndhwm / linger / conflate / rcvhwm knobs so future callers can choose drop-old vs delivery-guarantee semantics per message type.

Why

The next step is a MetricsPublisher that publishes MetricsSnapshot for a live TUI consumer. Building it directly on top of the EventRecord-specific pub/sub would duplicate the batching / pending-queue / async-writer machinery, or shoehorn snapshots through an EventRecord-shaped wrapper. Generalizing first is one rename pass + a codec extraction; afterward the metrics publisher is a thin composition.

Design context lives at .cursor_artifacts/zmq_metrics_publisher.md (not part of this PR; just for reference).

Behavior preservation

• MessageSubscriber._on_readable catches msgspec.DecodeError only — same as before.
• EventRecordCodec.on_decode_error returns the same ErrorEventType.GENERIC wrapper that the old _on_readable constructed inline.
• All existing socket options (SNDHWM=0, LINGER=-1, IMMEDIATE=1, RCVHWM=0) are the new parameter defaults.

Test plan

• tests/unit/transport/test_zmq_pool_transport.py — updated to use ZmqMessagePublisher(EventRecordCodec(), ...)
• tests/unit/async_utils/test_event_publisher.py — CollectingEventSubscriber updated
• All 154 tests pass in tests/unit/transport/, tests/unit/async_utils/test_event_publisher.py, tests/unit/async_utils/services/event_logger/, tests/unit/async_utils/services/metrics_aggregator/ (run in inference_endpoint_dev container)
• pre-commit run clean (ruff, ruff-format, mypy, license headers, prettier)
• Reviewer to confirm no behavior change in production paths beyond the explicit ones called out above

🤖 Generated with Claude Code

[github-actions]

MLCommons CLA bot All contributors have signed the MLCommons CLA ✍️ ✅

[gemini-code-assist]

Code Review

This pull request generalizes the pub/sub transport layer by introducing a generic MessageCodec protocol and refactoring the existing EventRecord-specific publisher and subscriber into generic MessagePublisher[T] and MessageSubscriber[T] classes. The ZMQ implementations were updated to support these generics and now include additional socket configuration options such as conflate and high-water marks. Feedback was provided to broaden the exception handling in the subscriber's read loop to ensure the implementation remains truly generic and adheres to the codec protocol contract.

[gemini-code-assist]

The MessageSubscriber should catch Exception instead of msgspec.DecodeError to remain truly generic and adhere to the MessageCodec protocol contract. The protocol defines on_decode_error as accepting any Exception, and different codec implementations might raise various error types (e.g., ValueError, TypeError, or library-specific exceptions like json.JSONDecodeError) during decoding. Catching only msgspec.DecodeError here couples the base protocol implementation to a specific library and could lead to unhandled exceptions crashing the reader loop if a non-msgspec codec is used.

Suggested change

	except msgspec.DecodeError as e:
	except Exception as e:

[arekay-nv]

Review Council — Multi-AI Code Review

Reviewed by Codex + Claude | Depth: standard | HEAD: faaff7b

Found 1 issue across 1 file: 0 critical, 0 high, 1 medium, 0 low.

Codex returned no actionable regressions: "The refactor from EventRecord-specific pub/sub classes to generic codec-driven message transport appears internally consistent." Claude flagged one test-coverage gap on the new on_decode_error codec method (see inline comment).

Note (not posted as inline): a few docs/**/DESIGN.md files still reference the old symbol names (encode_event_record, decode_event_record, ZmqEventRecordPublisher, ZmqEventRecordSubscriber). Worth a follow-up doc sweep.

[arekay-nv]

[Review Council — Claude] medium · testing

EventRecordCodec.on_decode_error (in core/record.py:187) is brand-new public-API logic introduced by this PR with two distinct branches:

def on_decode_error(self, payload: bytes, exc: Exception) -> EventRecord:
    if not isinstance(exc, msgspec.DecodeError):
        raise exc
    return EventRecord(
        event_type=ErrorEventType.GENERIC,
        data=ErrorData(error_type=type(exc).__name__, error_message=str(exc)),
    )

grep -r on_decode_error tests/ src/ returns zero hits outside the dispatch site in protocol.py. Neither the wrap-on-DecodeError branch nor the re-raise-on-other-exception branch is exercised. This is the only piece of behavior that meaningfully differs between the old inline _on_readable decode-error handling and the new codec abstraction (the previous implementation unconditionally wrapped any msgspec.DecodeError; the new code re-raises non-DecodeError exceptions, which is a real behavioral change).

More importantly, this is the cold path that fires when something else has already gone wrong — exactly the path you don't want untested. The re-raise branch in particular escapes _on_readable (which catches only StopIteration), so a non-DecodeError propagates out of the asyncio reader callback and silently de-registers the subscriber.

Suggested addition in TestEventRecordRoundTrip (or a new TestEventRecordCodecOnDecodeError):

1. Garbage payload returns a wrapped EventRecord(event_type=ErrorEventType.GENERIC, data=ErrorData(...)) whose error_message matches str(exc).
2. on_decode_error(b"...", ValueError("not a decode error")) re-raises ValueError (verifies the not isinstance(exc, msgspec.DecodeError) branch).

[arekay-nv]

Review Council — Multi-AI Code Review

Reviewed by Codex + Claude | Depth: standard | HEAD: faaff7b

Found 1 issue across 1 file:

• 0 critical/high
• 1 medium
• 0 low

#	File	Line	Severity	Category	Reviewer(s)	Summary
1	tests/unit/core/test_record.py	32	medium	testing	Claude	New EventRecordCodec.on_decode_error has 2 distinct branches (wrap on msgspec.DecodeError, re-raise on anything else) — neither exercised; grep shows zero test references; this is the only behavioral change introduced by the codec refactor

Codex verdict

The refactor from EventRecord-specific pub/sub classes to generic codec-driven message transport appears internally consistent. I did not find an actionable regression in the changed code paths or any missed in-repo call sites that would break existing behavior.

Cross-reviewer agreement: only one finding overall, from Claude. Codex independently concluded the refactor is clean.

Follow-up (not posted inline)

A few design docs reference the old symbol names — likely a small doc sweep:

• docs/load_generator/DESIGN.md
• docs/async_utils/DESIGN.md
• docs/async_utils/services/DESIGN.md
• docs/async_utils/services/event_logger/DESIGN.md

Stale references include encode_event_record, decode_event_record, ZmqEventRecordPublisher, ZmqEventRecordSubscriber.

Inline review: #pullrequestreview-4231671106

🤖 Posted by /review-council skill.

[arekay-nv]

#300 🎉 !