fix: serialize disagg first_gen_log_probs int keys for Rust transport

[nv-yna]

Summary

• Fix disaggregated logprobs CI failure (test_deployment[disaggregated_logprobs-2])
• TRT-LLM PR #11727 adds first_gen_log_probs to DisaggregatedParams with integer token-ID dict keys (e.g. {4710: Logprob(...)})
• After dataclasses.asdict(), these int keys break Rust pythonize 0.23 depythonize (dict_key_not_string error) → cascades into Disconnected: Stream ended before generation completed
• Add serialize_first_gen_log_probs / deserialize_first_gen_log_probs to DisaggregatedParamsCodec using TRT-LLM own list-of-dicts format ([{"token_id": id, "logprob": float, "rank": int}])
• Add unit tests for the serialization round-trip

Test plan

• Reproduced original failure on H100x2 cluster
• test_deployment[disaggregated_logprobs-2] passes with fix (all logprobs payloads validated — 300 tokens with logprobs each)
• test_deployment[disaggregated-2] passes (no regression on non-logprobs disagg)
• Unit tests for serialization round-trip (no GPU required)
• CI multi-gpu tests pass

Fixes: DYN-2265 / NVBugs 5926823

Summary by CodeRabbit

• Bug Fixes

  • Improved token log probability handling in distributed inference scenarios for better data format compatibility.

• Tests

  • Added comprehensive test coverage for log probability serialization and deserialization.

[github-actions]

👋 Hi nv-yna! Thank you for contributing to ai-dynamo/dynamo.

Just a reminder: The NVIDIA Test Github Validation CI runs an essential subset of the testing framework to quickly catch errors.Your PR reviewers may elect to test the changes comprehensively before approving your changes.

🚀

[coderabbitai]

Walkthrough

This pull request adds serialization and deserialization logic for first_gen_log_probs in the disaggregated request handling pipeline. New codec methods transform log probabilities between TRT-LLM's internal dictionary format and a transport-compatible list format, with handler integration and comprehensive test coverage.

Changes

Cohort / File(s)	Summary
Codec Implementation components/src/dynamo/trtllm/utils/disagg_utils.py	Added Logprob import and two static methods: serialize_first_gen_log_probs() converts internal dict-of-dicts to list-of-lists format; deserialize_first_gen_log_probs() converts back to dict mapping token_id to Logprob instances. Both handle missing keys gracefully.
Handler Integration components/src/dynamo/trtllm/request_handlers/handler_base.py	Integrated codec methods into request handling pipeline: deserialization in _decode_disaggregated_params_from_prefill() and serialization in _encode_and_pack_disaggregated_params().
Test Coverage tests/serve/test_disagg_logprobs_serialization.py	New test module with six test cases validating round-trip serialization/deserialization: None passthrough, missing field handling, single/multi-token roundtrips, rank preservation, and empty list scenarios. Includes JSON safety verification.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Poem

A hop through the logs, from dict to a list,
The codec transforms what transport won't miss—
Token by token, with rank shining bright,
The rabbits rejoice at serialization done right! 🐰✨

🚥 Pre-merge checks | ✅ 3

✅ Passed checks (3 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and specifically describes the main fix: serializing disaggregated first_gen_log_probs integer keys for Rust transport, which directly addresses the CI failure detailed in the description.
Description check	✅ Passed	The description provides a clear summary, details the root cause and solution, includes specific test validation, and identifies related issues. However, it does not follow the repository's template structure with explicit 'Overview', 'Details', 'Where should the reviewer start', and 'Related Issues' sections.
Docstring Coverage	✅ Passed	Docstring coverage is 80.00% which is sufficient. The required threshold is 80.00%.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

Tip

Try Coding Plans. Let us write the prompt for your AI agent so you can ship faster (with fewer bugs).
Share your feedback on Discord.

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 2

🧹 Nitpick comments (1)

tests/serve/test_disagg_logprobs_serialization.py (1)

104-107: Add strict=True to zip() for safer iteration.

The static analysis flagged this zip() call. Since original and recovered should have equal lengths after a successful round-trip, adding strict=True would catch any unexpected length mismatch.

🔧 Proposed fix

-        for orig_pos, rec_pos in zip(original, recovered):
+        for orig_pos, rec_pos in zip(original, recovered, strict=True):

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tests/serve/test_disagg_logprobs_serialization.py` around lines 104 - 107,
The zip over original and recovered should be strict to catch length mismatches;
update the loop that iterates with "for orig_pos, rec_pos in zip(original,
recovered):" to use "zip(original, recovered, strict=True)" so any unexpected
unequal lengths after the round-trip fail fast; keep the loop body (asserting
rec_pos[tid].logprob and rank) unchanged.

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@tests/serve/test_disagg_logprobs_serialization.py`:
- Around line 36-38: The test class TestDisaggLogprobsSerializationRoundtrip
currently has only two markers; add the required GPU and type markers by
annotating the class with `@pytest.mark.gpu_0` and `@pytest.mark.unit` (in addition
to the existing `@pytest.mark.pre_merge` and `@pytest.mark.trtllm`) so the class has
a scheduling marker, a GPU marker, and a type marker per guidelines.
- Around line 30-33: The test imports optional dependencies Logprob and
DisaggregatedParamsCodec directly; wrap those imports in a try/except
ImportError at module import time and call pytest.skip("missing optional
dependency: tensorrt_llm", allow_module_level=True) inside the except so the
test module is not collected when tensorrt_llm (and its Logprob) or disagg utils
are absent; update the import block referencing Logprob and
DisaggregatedParamsCodec accordingly to perform the guarded import and early
skip.

---

Nitpick comments:
In `@tests/serve/test_disagg_logprobs_serialization.py`:
- Around line 104-107: The zip over original and recovered should be strict to
catch length mismatches; update the loop that iterates with "for orig_pos,
rec_pos in zip(original, recovered):" to use "zip(original, recovered,
strict=True)" so any unexpected unequal lengths after the round-trip fail fast;
keep the loop body (asserting rec_pos[tid].logprob and rank) unchanged.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 8e1547d0-90eb-41c6-99b3-2453c325b748

📥 Commits

Reviewing files that changed from the base of the PR and between 5e51d6d and 2ea183a.

📒 Files selected for processing (3)

• components/src/dynamo/trtllm/request_handlers/handler_base.py
• components/src/dynamo/trtllm/utils/disagg_utils.py
• tests/serve/test_disagg_logprobs_serialization.py

[indrajit96]

LGTM!
Some minor comments for test