[None][fix] Make SleepConfig picklable by replacing closure lambda in defaultdict

[hhzhang16]

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.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 52310273-2eee-4119-a4aa-c15fa2e12a7b

📥 Commits

Reviewing files that changed from the base of the PR and between f8572ab and 76b368f.

📒 Files selected for processing (2)

• tensorrt_llm/llmapi/llm_args.py
• tests/unittest/llmapi/test_llm_args.py

---

📝 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.
Description check	✅ Passed	The PR description covers the issue, the solution, test coverage, and includes the required checklist. It clearly explains the problem (closure lambda not picklable during MPI serialization), the fix (dataclass replacement), and testing approach.
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.

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

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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