fix(client): use canonical TRT-LLM transfer catalogs

[chienchunhung]

Summary

• publish TRT-LLM weights under canonical parameter paths instead of runtime-only next_attn / next_layer_layernorm aliases
• require exact source/target tensor-name catalogs before initializing NIXL writes
• fail closed when aliases have no canonical backing path or either catalog is incomplete
• canonicalize every TRT-LLM publication/target path and close per-rank log handlers on all exits

This fixes partial post-transform installation observed while qualifying NVIDIA/TensorRT-LLM#15641: 64 of 135 Llama tensors were advertised under donor-only alias paths, leaving the receiver partially initialized.

Test plan

• PYTHONPATH=modelexpress_client/python python3 -m pytest modelexpress_client/python/tests/test_trtllm_live_transfer.py -q (5 passed)
• ruff and Python compilation for the changed files
• TinyLlama BF16 TP=1: donor, full receiver, and no-shards receiver match HF baseline
• TinyLlama BF16 TP=2: ranks 0/1 each publish 135 tensors; full and no-shards receivers match TP=2 baseline
• TP=2 donor / TP=1 receiver mismatch rejects P2P before target NIXL initialization and completes full HF fallback

Follow-ups

• add publisher heartbeats to keep long-lived TRT-LLM source metadata fresh under the default server reaper timeout
• replace private TRT-LLM integration helpers with a versioned public publication/transfer API

Summary by CodeRabbit

• Bug Fixes

  • Improved live model weight transfers by excluding duplicate runtime aliases from published parameter catalogs.
  • Transfers now fail fast when source and target parameter catalogs do not match exactly.
  • Preserved support for compatible dtype conversion and fallback loading.

• Tests

  • Added coverage for aliased parameters, exact catalog matches, and incomplete catalog failures.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[coderabbitai]

Walkthrough

TRT-LLM live weight transfer now canonicalizes aliased parameters, publishes tensor_source metadata, and fails fast when source and target parameter catalogs differ. Tests cover alias exclusion, exact matches, and incomplete catalogs.

Changes

TRT-LLM Catalog Validation

Layer / File(s)	Summary
Canonical parameter publication modelexpress_client/python/modelexpress/trtllm_live_transfer.py	Canonical parameter helpers exclude runtime-only aliases, publication uses one name per storage, and worker metadata writes tensor_source only.
Exact catalog loader validation modelexpress_client/python/modelexpress/trtllm_live_transfer.py, modelexpress_client/python/tests/test_trtllm_live_transfer.py	MxLiveWeightLoader requires exact source-target catalog equality before transfer planning, while tests cover canonical aliases and missing-key failures.

Estimated code review effort: 3 (Moderate) | ~25 minutes

Poem

I’m a rabbit with tensors in tow,
Canonical names now neatly grow.
Aliases hop out of the chart,
Mismatched catalogs quickly depart.
tensor_source leads the way—
Clean transfers bloom today!

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly reflects the main change: canonical TRT-LLM transfer catalogs for client-side weight transfer.

---

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

[coderabbitai]

Actionable comments posted: 1

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

modelexpress_client/python/modelexpress/trtllm_live_transfer.py (1)

297-300: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick win

Use canonical parameter names in both catalogs.

publish_model_params() is canonicalized, but these paths still build catalogs from the default named_parameters(). That lets alias names like next_attn survive into the published or target descriptor set, so the new exact-match check can reject a compatible transfer.

• modelexpress_client/python/modelexpress/trtllm_live_transfer.py#L297-L300: build the worker descriptor catalog from _canonical_named_parameters(torch_model) before device filtering.
• modelexpress_client/python/modelexpress/trtllm_live_transfer.py#L399-L399: build target_params from _canonical_named_parameters(model) so validation compares canonical catalogs.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@modelexpress_client/python/modelexpress/trtllm_live_transfer.py` around lines
297 - 300, Update the worker descriptor catalog construction near lines 297-300
to use _canonical_named_parameters(torch_model) before device filtering, and
update target_params near line 399 to use _canonical_named_parameters(model).
Apply both changes in trtllm_live_transfer.py so published and validation
catalogs consistently use canonical parameter names.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@modelexpress_client/python/modelexpress/trtllm_live_transfer.py`:
- Line 399: Update load_weights() so the remainder of its execution, including
_require_exact_catalog_match, is wrapped in a finally-based cleanup path that
removes fh from the modelexpress logger and closes it. Ensure cleanup runs on
both catalog-mismatch exceptions and normal completion, preventing retained file
descriptors and duplicate log output.

---

Outside diff comments:
In `@modelexpress_client/python/modelexpress/trtllm_live_transfer.py`:
- Around line 297-300: Update the worker descriptor catalog construction near
lines 297-300 to use _canonical_named_parameters(torch_model) before device
filtering, and update target_params near line 399 to use
_canonical_named_parameters(model). Apply both changes in
trtllm_live_transfer.py so published and validation catalogs consistently use
canonical parameter names.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 3cb4a2ec-d68f-4bcc-914f-281c0147aa8c

📥 Commits

Reviewing files that changed from the base of the PR and between 9b57602 and b4148e3.

📒 Files selected for processing (2)

• modelexpress_client/python/modelexpress/trtllm_live_transfer.py
• modelexpress_client/python/tests/test_trtllm_live_transfer.py

[zhengluo-nv]

/ok to test b4148e3

[zhengluo-nv]

/ok to test 9d3bc31

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[zhengluo-nv]

Addressed the remaining CodeRabbit review findings in 9d3bc317:

• publish_from_worker() and the target catalog now both use _canonical_named_parameters(), so runtime-only aliases cannot survive on either side of exact catalog validation.
• The per-rank log handler is removed and closed in finally on every exit path, with a failure-path regression test.

Validation: focused tests 5 passed, Ruff and Python compilation passed, and the live Build worker image (trt-llm) plus P2P test (trt-llm) jobs both succeeded.