better trtllm unittest ray.init and shutdown; add rollout config doc

Author: davidmlw

tests/workers/rollout/rollout_trtllm/test_adapter.py

@@ -142,7 +142,17 @@ def test_init_without_device_mesh(self):
 
         try:
             os.environ.setdefault("TLLM_RAY_FORCE_LOCAL_CLUSTER", "1")
-            ray.init(address="local", ignore_reinit_error=True, include_dashboard=False)
+            ray.init(
+                runtime_env={
+                    "env_vars": {
+                        "TOKENIZERS_PARALLELISM": "true",
+                        "NCCL_DEBUG": "WARN",
+                        "VLLM_LOGGING_LEVEL": "INFO",
+                        "VLLM_USE_V1": "1",
+                    }
+                },
+                ignore_reinit_error=True,
+            )
 
             config_dir = os.path.abspath("verl/verl/trainer/config")
             if not os.path.exists(config_dir):
@@ -187,5 +197,5 @@ def test_init_without_device_mesh(self):
                 os.environ.pop("RANK", None)
             else:
                 os.environ["RANK"] = prev_rank
+            print("\nShutting down Ray...")
             ray.shutdown()
-            subprocess.run(["ray", "stop"], capture_output=True)

tests/workers/rollout/rollout_trtllm/test_async_server.py

@@ -170,7 +170,17 @@ def test_async_generate(self):
         """Test TRT-LLM generate method with real model."""
         try:
             os.environ.setdefault("TLLM_RAY_FORCE_LOCAL_CLUSTER", "1")
-            ray.init(address="local", ignore_reinit_error=True, include_dashboard=False)
+            ray.init(
+                runtime_env={
+                    "env_vars": {
+                        "TOKENIZERS_PARALLELISM": "true",
+                        "NCCL_DEBUG": "WARN",
+                        "VLLM_LOGGING_LEVEL": "INFO",
+                        "VLLM_USE_V1": "1",
+                    }
+                },
+                ignore_reinit_error=True,
+            )
 
             rollout_config, model_config = self._build_rollout_config(response_length=50)
 
@@ -209,14 +219,24 @@ def test_async_generate(self):
             print(f"Log probs: {result.log_probs[:10]}...")  # Print first 10 log probs
 
         finally:
+            print("\nShutting down Ray...")
             ray.shutdown()
-            subprocess.run(["ray", "stop"], capture_output=True)
 
     def test_async_memory_management(self):
         """Test TRT-LLM async memory management (sleep) reduces memory usage."""
         try:
             os.environ.setdefault("TLLM_RAY_FORCE_LOCAL_CLUSTER", "1")
-            ray.init(address="local", ignore_reinit_error=True, include_dashboard=False)
+            ray.init(
+                runtime_env={
+                    "env_vars": {
+                        "TOKENIZERS_PARALLELISM": "true",
+                        "NCCL_DEBUG": "WARN",
+                        "VLLM_LOGGING_LEVEL": "INFO",
+                        "VLLM_USE_V1": "1",
+                    }
+                },
+                ignore_reinit_error=True,
+            )
 
             rollout_config, model_config = self._build_rollout_config(free_cache_engine=True)
 
@@ -271,5 +291,5 @@ def get_gpu_memory_mb_for_device(device_uuid: str) -> float:
             )
 
         finally:
+            print("\nShutting down Ray...")
             ray.shutdown()
-            subprocess.run(["ray", "stop"], capture_output=True)

verl/trainer/config/rollout/rollout.yaml

@@ -313,6 +313,7 @@ profiler:
   # choices: npu, torch
   tool: ${oc.select:global_profiler.tool,null}
 
+  # global tool config
   global_tool_config: ${oc.select:global_profiler.global_tool_config,null}
 
   # whether enable profile on rollout

verl/workers/rollout/trtllm_rollout/trtllm_async_server.py

@@ -178,6 +178,8 @@ async def launch_server(self):
         self.llm = await AsyncLLM(**llm_kwargs)
 
         trtllm_server = OpenAIServer(
+            # TODO: update to generator in future
+            # generator=self.llm,
             llm=self.llm,
             model=self.model_config.local_path,
             tool_parser=None,