feat: Make vllm sleep level customizable

nemo_rl/models/generation/vllm/config.py

@@ -26,6 +26,8 @@ class VllmSpecificArgs(TypedDict):
     # Additional arguments for vLLM inserted by nemo rl based on the context of when vllm is used
     skip_tokenizer_init: bool
     async_engine: bool
+    # vLLM sleep level used after generation. Defaults to 1 when unset.
+    sleep_level: NotRequired[Literal[0, 1, 2]]
     load_format: NotRequired[str]
     precision: NotRequired[str]
     kv_cache_dtype: Literal["auto", "fp8", "fp8_e4m3"]

nemo_rl/models/generation/vllm/vllm_worker.py

@@ -135,6 +135,17 @@ def __init__(
                           the vLLM worker subprocess (e.g. for quantization configs).
         """
         self.cfg = config
+        sleep_level = self.cfg["vllm_cfg"].get("sleep_level", 1)
+        if (
+            not isinstance(sleep_level, int)
+            # bool is a subclass of int, but YAML booleans should not be valid sleep levels.
+            or isinstance(sleep_level, bool)
+            or sleep_level not in (0, 1, 2)
+        ):
+            raise ValueError(
+                f"vllm_cfg.sleep_level must be 0, 1, or 2, got {sleep_level!r}"
+            )
+        self.sleep_level = sleep_level
         self.model_name = self.cfg["model_name"]
         self.tensor_parallel_size = self.cfg["vllm_cfg"]["tensor_parallel_size"]
         self.pipeline_parallel_size = self.cfg["vllm_cfg"]["pipeline_parallel_size"]
@@ -1003,13 +1014,13 @@ def sleep(self):
         # stays in sync with the receiver cache that vLLM clears internally
         # during sleep.  Without this, the sender thinks images are already
         # cached on the receiver and sends data=None, causing an assertion
-        # error.  We only clear the renderer (sender) cache here — the
+        # error.  We only clear the renderer (sender) cache here; the
         # receiver and worker-level caches are reset by sleep() internally.
         if hasattr(self.llm, "renderer") and hasattr(
             self.llm.renderer, "clear_mm_cache"
         ):
             self.llm.renderer.clear_mm_cache()
-        self.llm.sleep(level=1)
+        self.llm.sleep(level=self.sleep_level)
 
         gc.collect()
         torch.cuda.empty_cache()

nemo_rl/models/generation/vllm/vllm_worker_async.py

@@ -1148,7 +1148,7 @@ async def sleep_async(self):
         # the receiver and sends data=None, causing an assertion error.
         if hasattr(self.llm, "reset_mm_cache"):
             await self.llm.reset_mm_cache()
-        await self.llm.sleep(level=1)
+        await self.llm.sleep(level=self.sleep_level)
 
         gc.collect()
         torch.cuda.empty_cache()

tests/unit/models/generation/test_vllm_worker.py

@@ -0,0 +1,58 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any, cast
+
+import pytest
+
+from nemo_rl.models.generation.vllm.config import VllmConfig
+from nemo_rl.models.generation.vllm.vllm_worker import BaseVllmGenerationWorker
+
+
+def _worker_with_vllm_cfg(vllm_cfg: dict[str, Any]) -> BaseVllmGenerationWorker:
+    config = {
+        "model_name": "dummy-model",
+        "vllm_cfg": {
+            "tensor_parallel_size": 1,
+            "pipeline_parallel_size": 1,
+            "expert_parallel_size": 1,
+            "gpu_memory_utilization": 0.5,
+            "max_model_len": 128,
+            "skip_tokenizer_init": True,
+            "async_engine": False,
+            "precision": "bfloat16",
+            "kv_cache_dtype": "auto",
+            **vllm_cfg,
+        },
+    }
+    return BaseVllmGenerationWorker(cast(VllmConfig, config), bundle_indices=None)
+
+
+def test_vllm_sleep_level_defaults_to_level_1():
+    worker = _worker_with_vllm_cfg({})
+
+    assert worker.sleep_level == 1
+
+
+@pytest.mark.parametrize("sleep_level", [0, 1, 2])
+def test_vllm_sleep_level_accepts_supported_levels(sleep_level):
+    worker = _worker_with_vllm_cfg({"sleep_level": sleep_level})
+
+    assert worker.sleep_level == sleep_level
+
+
+@pytest.mark.parametrize("sleep_level", [-1, 3, "2", True])
+def test_vllm_sleep_level_rejects_unsupported_levels(sleep_level):
+    with pytest.raises(ValueError, match=r"vllm_cfg\.sleep_level must be 0, 1, or 2"):
+        _worker_with_vllm_cfg({"sleep_level": sleep_level})