feat: Make vllm sleep level customizable

[Kipok]

What does this PR do ?

Makes vllm sleep level customizable. E.g. sleep level 2 is useful for cases when CPU memory is limited

Summary by CodeRabbit

• New Features

  • vLLM generation now supports configurable sleep levels (0, 1, or 2) via configuration, defaulting to level 1 when unset.

• Tests

  • Added unit tests validating sleep level configuration behavior and error handling.

[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

This PR introduces configurable sleep levels for vLLM workers. A new optional sleep_level configuration field (values 0, 1, or 2) is added to the vLLM config schema. A validation method is implemented in BaseVllmGenerationWorker to enforce type and value constraints, then applied to both synchronous and asynchronous sleep calls. Unit tests verify default behavior, accepted values, and rejection of invalid inputs.

Changes

Sleep Level Configuration and Implementation

Layer / File(s)	Summary
Sleep level configuration schema nemo_rl/models/generation/vllm/config.py	VllmSpecificArgs adds optional sleep_level: NotRequired[Literal[0, 1, 2]] field with default behavior documented.
Sleep level validation and worker integration nemo_rl/models/generation/vllm/vllm_worker.py, nemo_rl/models/generation/vllm/vllm_worker_async.py	BaseVllmGenerationWorker._sleep_level() validates config value (rejects non-int and out-of-range values), raising ValueError on invalid input. Both VllmGenerationWorker.sleep() and VllmAsyncGenerationWorker.sleep_async use the validated level instead of hardcoded 1.
Sleep level validation tests tests/unit/models/generation/test_vllm_worker.py	Test module with helper to construct workers with custom configs. Tests verify default return of 1, acceptance of 0, 1, 2, and rejection of invalid values (-1, 3, string "2", boolean True) with descriptive error message.

🎯 2 (Simple) | ⏱️ ~10 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 2

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 44.44% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.
Test Results For Major Changes	⚠️ Warning	PR introduces a configurable sleep feature with performance/memory implications (sleep level 2 for CPU memory) but lacks test results and performance metrics documentation in the PR description.	Document test results and add performance/memory impact analysis in PR description, particularly for sleep level 2.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
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.
Title check	✅ Passed	The title accurately summarizes the main change: making vLLM sleep level customizable across configuration and implementation.

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

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch igitman/vllm-sleep-control

---

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

[coderabbitai]

Actionable comments posted: 1

🤖 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 `@tests/unit/models/generation/test_vllm_worker.py`:
- Line 46: The regex in the pytest.raises match is using unescaped dots and
should be made a raw string with the dots escaped to avoid over-matching; update
the assertion in the test_vllm_worker test (the pytest.raises line that
currently matches "vllm_cfg.sleep_level must be 0, 1, or 2") to use a raw string
and escape the literal dots (e.g., change the match to a raw pattern where '.'
is written as '\.').

🪄 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: 3d403493-38f5-4cce-8ded-c73aae6bdb1c

📥 Commits

Reviewing files that changed from the base of the PR and between 30108c3 and 7832ac9.

📒 Files selected for processing (4)

• nemo_rl/models/generation/vllm/config.py
• nemo_rl/models/generation/vllm/vllm_worker.py
• nemo_rl/models/generation/vllm/vllm_worker_async.py
• tests/unit/models/generation/test_vllm_worker.py

[terrykong]

@yuki-97 @parthmannan to review

@Kipok out of curiosity, what scenario was this helpful in? if host mem was around 900GB for GB but VRAM was ~200GB with only fraction of that being model weights, i would imagine the savings here to not be very impactful unless there was something else very large eating up host memory.

[Kipok]

I haven't looked too much into what was eating up cpu memory, but somehow my experiments using nemotron-super worked on one internal cluster, but didn't work on another one because of cpu out-of-memory. I think there was something like 50gb difference or so in total available memory per node but that was enough to push it over the limit. Setting sleep level to 2 helped to resolve it with no impact on speed. But again I don't fully understand what was taking the memory and maybe there is a better fix, but this seems to work and I thought could be useful to others as well

[terrykong]

FWIW, the change is okay with me, just trying to understand the symptom

@joyang-nv @guyueh1 could this have potentially been explained with the numa issue on some clusters?

[Kipok]

thanks @yuki-97, addressed the comments

[yuki-97]

LGTM, thanks @Kipok

[guyueh1]

@Kipok I think this will cause issues to validation; in grpo.py, when doing validation, we first do refit then do validation then sleep the engine; then in the next step we wakeup the engine and directly start generation because we think we have done refit already. If we use level=2 here, then the sleep will discard the weights and the next generation after validation will use dummy weights. I implemented a way to only set sleep_level=2 in proper places in #2495 in parallel to your PR

[Kipok]

sure, sounds good, let me close this one then!