[None][fix] Make SleepConfig picklable by replacing closure lambda in defaultdict by hhzhang16 · Pull Request #13918 · NVIDIA/TensorRT-LLM

hhzhang16 · 2026-05-08T19:57:35Z

Summary by CodeRabbit

Bug Fixes
- Improved SleepConfig serialization to properly support pickle-based object persistence and restoration.
Tests
- Added test coverage verifying SleepConfig can be successfully pickled and unpickled while maintaining configuration integrity, including default values and custom overrides.

Description

SleepConfig._make_defaulted_restore_modes() returns a defaultdict whose default_factory is a closure lambda (lambda: default_mode). Because defaultdict stores default_factory as an instance attribute, it is serialised when llm_args is pickled during MPI worker initialisation (mpi_session.py distributes configuration to each rank via pickle). A closure lambda has no module-level name and raises AttributeError at pickle time, making any SleepConfig-enabled multi-rank deployment fail before the engine starts.

Fix: introduce _SleepConfigDefaultFactory, a private @DataClass that holds the concrete RestoreMode value and returns it via __call__. It is a drop-in replacement for the lambda and is fully picklable. No behaviour change — _validate_restore_modes calls v.default_factory() and continues to work correctly.

Test Coverage

test_SleepConfig_is_picklable:

Default SleepConfig survives pickle — the core regression check
default_factory still works after round-trip — a missing key on the restored defaultdict returns a RestoreMode rather than raising KeyError, confirming _SleepConfigDefaultFactory was properly reconstructed by pickle
Custom overrides survive pickle — ensures the explicit key→value mapping is also preserved correctly

PR Checklist

Please review the following before submitting your PR:

PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

Signed-off-by: Hannah Zhang <hannahz@nvidia.com>

coderabbitai · 2026-05-08T20:04:44Z

📝 Walkthrough

Walkthrough

This PR introduces pickling support for SleepConfig by replacing an unpicklable lambda default factory with a picklable @dataclass. A new factory class captures the default mode, updates the restore_modes initialization logic, and adds a test verifying successful pickle round-trip behavior including defaultdict fallback handling.

Changes

SleepConfig Pickling Support

Layer / File(s)	Summary
Picklable Factory Definition `tensorrt_llm/llmapi/llm_args.py`	New internal `_SleepConfigDefaultFactory` dataclass with `__call__` method to replace unpicklable `lambda: default_mode`.
Restore Modes Integration `tensorrt_llm/llmapi/llm_args.py`	`_make_defaulted_restore_modes` now uses `_SleepConfigDefaultFactory(default_mode)` as the defaultdict factory while preserving normalization logic.
Pickling Verification `tests/unittest/llmapi/test_llm_args.py`	New test `test_SleepConfig_is_picklable` verifies default and custom `restore_modes` survive pickle round-trip and defaultdict fallback works after deserialization.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 20.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and concisely describes the main change: making SleepConfig picklable by replacing a closure lambda in defaultdict with a dataclass-based factory.
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	The PR description comprehensively explains the problem (closure lambda in defaultdict causing pickle failures in MPI), the solution (picklable dataclass replacement), and test coverage details.

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

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

Tip

💬 Introducing Slack Agent: The best way for teams to turn conversations into code.

Slack Agent is built on CodeRabbit's deep understanding of your code, so your team can collaborate across the entire SDLC without losing context.

Generate code and open pull requests
Plan features and break down work
Investigate incidents and troubleshoot customer tickets together
Automate recurring tasks and respond to alerts with triggers
Summarize progress and report instantly

Built for teams:

Shared memory across your entire org—no repeating context
Per-thread sandboxes to safely plan and execute work
Governance built-in—scoped access, auditability, and budget controls

One agent for your entire SDLC. Right inside Slack.

👉 Get started

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

…managed deployment Append a new section to §06 distinguishing two failover models: - Shadow failover (existing): pre-warmed standby, sub-5s activation, requires §07 compile cache + new executor state machine. - Restart-after-death (new section): no standby, external supervisor restarts trtllm-serve, replacement zero-copy materializes from surviving GMS daemon via PR NVIDIA#13926's RO branch. Works on main today (no new TRT-LLM code). Document the supporting gaps: - TRT-LLM has no built-in worker restart machinery; restart-after-death requires an external supervisor (systemd / K8s / Dynamo). - GMS daemon must be a node-level service, not a worker subprocess. - Daemon commit-survival on writer disconnect is the load-bearing upstream assumption (GMS-7 open question). - Sleep/wake (NVIDIA#13918 + NVIDIA#14052) doesn't compose with GMS weights — three implementation tiers documented; Tier-A validator recommended for NVIDIA#13926. Add a self-managed deployment recipe: systemd unit for the GMS daemon, trtllm-serve unit with --gms-mode=auto for the worker, verification steps, and caveats (single-rank only, daemon lifetime discipline, no auto sleep/wake for weights, compile-cache cold-restart cost pre-§07). README TOC entry updated to point at the new section. Signed-off-by: Chien-Chun Hung <chienchunh@nvidia.com>

chienchunhung · 2026-05-12T20:03:28Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-05-12T20:09:47Z

PR_Github #48023 [ run ] triggered by Bot. Commit: 76b368f Link to invocation

tensorrt-cicd · 2026-05-13T05:33:19Z

PR_Github #48023 [ run ] completed with state SUCCESS. Commit: 76b368f
/LLM/main/L0_MergeRequest_PR pipeline #37859 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

Signed-off-by: Hannah Zhang <hannahz@nvidia.com>

chienchunhung · 2026-05-13T16:15:48Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-05-13T16:21:10Z

PR_Github #48214 [ run ] triggered by Bot. Commit: 6e276f2 Link to invocation

Signed-off-by: Hannah Zhang <hannahz@nvidia.com>

tensorrt-cicd · 2026-05-14T10:37:01Z

PR_Github #48214 [ run ] completed with state SUCCESS. Commit: 6e276f2
/LLM/main/L0_MergeRequest_PR pipeline #38032 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

chienchunhung · 2026-05-14T17:29:35Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-05-14T17:35:33Z

PR_Github #48410 [ run ] triggered by Bot. Commit: 0e8f3be Link to invocation

tensorrt-cicd · 2026-05-15T02:29:45Z

PR_Github #48410 [ run ] completed with state SUCCESS. Commit: 0e8f3be
/LLM/main/L0_MergeRequest_PR pipeline #38212 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

chienchunhung · 2026-05-15T16:30:21Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-05-15T16:35:42Z

PR_Github #48606 [ run ] triggered by Bot. Commit: 0e8f3be Link to invocation

tensorrt-cicd · 2026-05-15T22:27:12Z

PR_Github #48606 [ run ] completed with state SUCCESS. Commit: 0e8f3be
/LLM/main/L0_MergeRequest_PR pipeline #38389 completed with status: 'SUCCESS'

CI Report

Link to invocation

hhzhang16 added 2 commits May 8, 2026 15:48

feat: make SLeepConfig picklable for MPI worker distribution

94cfb18

Signed-off-by: Hannah Zhang <hannahz@nvidia.com>

test: add pickle test for SleepConfig

76b368f

Signed-off-by: Hannah Zhang <hannahz@nvidia.com>

hhzhang16 requested a review from a team as a code owner May 8, 2026 19:57

hhzhang16 requested a review from hchings May 8, 2026 19:57

github-actions Bot assigned hhzhang16 May 8, 2026

chienchunhung self-requested a review May 8, 2026 20:06

svc-trtllm-gh-bot added the Community want to contribute PRs initiated from Community label May 8, 2026

galletas1712 mentioned this pull request May 11, 2026

DRAFT proof: Kimi K2.5 + Eagle3 shadow-failover stack (do not merge) #13394

Draft

chienchunhung approved these changes May 12, 2026

View reviewed changes

Comment thread tests/unittest/llmapi/test_llm_args.py Outdated

Comment thread tests/unittest/llmapi/test_llm_args.py Outdated

hhzhang16 force-pushed the hannahz/dep-934-fix-sleepconfig-pickling branch from 3a2498d to 76b368f Compare May 12, 2026 20:44

chienchunhung mentioned this pull request May 12, 2026

[None][feat] add single-rank MPI sleep/wakeup and rank-0 collective_rpc shim #14052

Merged

1 task

fix: fix yapf violation and address reviewer comments

6e276f2

Signed-off-by: Hannah Zhang <hannahz@nvidia.com>

lint: fix formatting

e7ce604

Signed-off-by: Hannah Zhang <hannahz@nvidia.com>

pcastonguay approved these changes May 13, 2026

View reviewed changes

hchings reviewed May 13, 2026

View reviewed changes

Comment thread tests/unittest/llmapi/test_llm_args.py

test: reviewer comments

0e8f3be

Signed-off-by: Hannah Zhang <hannahz@nvidia.com>

chienchunhung enabled auto-merge (squash) May 15, 2026 19:20

chienchunhung merged commit 731eb9b into NVIDIA:main May 15, 2026
6 checks passed

Uh oh!

Conversation

hhzhang16 commented May 8, 2026 • edited by chienchunhung Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary by CodeRabbit

Description

Test Coverage

PR Checklist

GitHub Bot Help

Uh oh!

coderabbitai Bot commented May 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Estimated code review effort

❌ Failed checks (1 warning)

Uh oh!

Uh oh!

Uh oh!

chienchunhung commented May 12, 2026

Uh oh!

tensorrt-cicd commented May 12, 2026

Uh oh!

tensorrt-cicd commented May 13, 2026

Uh oh!

chienchunhung commented May 13, 2026

Uh oh!

tensorrt-cicd commented May 13, 2026

Uh oh!

Uh oh!

tensorrt-cicd commented May 14, 2026

Uh oh!

chienchunhung commented May 14, 2026

Uh oh!

tensorrt-cicd commented May 14, 2026

Uh oh!

tensorrt-cicd commented May 15, 2026

Uh oh!

chienchunhung commented May 15, 2026

Uh oh!

tensorrt-cicd commented May 15, 2026

Uh oh!

tensorrt-cicd commented May 15, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

hhzhang16 commented May 8, 2026 •

edited by chienchunhung

Loading

coderabbitai Bot commented May 8, 2026 •

edited

Loading