ci fixes

Author: hchings

tests/special_sanity/check_device_api_usage.py

@@ -47,6 +47,7 @@
     "verl/workers/engine/veomni/transformer_impl.py",  # appear in default device_name
     "verl/workers/rollout/vllm_rollout/vllm_async_server.py",  # appear in config.cudagraph_capture_sizes
     "verl/workers/rollout/sglang_rollout/async_sglang_server.py",  # manually set CUDA_VISIBLE_DEVICES
+    "verl/workers/rollout/trtllm_rollout/trtllm_async_server.py",  # appear in config.cudagraph_capture_sizes
 ]
 
 # directory or file path must contain keyword "nccl"

tests/special_sanity/check_pr_title.py

@@ -19,7 +19,7 @@
 pr_title = os.environ.get("PR_TITLE", "").strip()
 
 # Define rules
-allowed_modules = ["fsdp", "megatron", "sglang", "vllm", "rollout", "trainer"]
+allowed_modules = ["fsdp", "megatron", "sglang", "vllm", "trtllm", "rollout", "trainer"]
 allowed_modules += ["tests", "training_utils", "recipe", "hardware", "deployment"]
 allowed_modules += ["ray", "worker", "single_controller", "misc", "docker", "ci"]
 allowed_modules += ["perf", "model", "algo", "env", "tool", "ckpt", "doc", "data", "cfg", "reward"]

verl/trainer/ppo/ray_trainer.py

@@ -929,12 +929,12 @@ def init_workers(self):
         else:
             rm_resource_pool = None
 
-            self.async_rollout_manager = AgentLoopManager(
-                config=self.config,
-                worker_group=self.actor_rollout_wg,
-                rollout_resource_pool=actor_rollout_resource_pool,
-                rm_resource_pool=rm_resource_pool,
-            )
+        self.async_rollout_manager = AgentLoopManager(
+            config=self.config,
+            worker_group=self.actor_rollout_wg,
+            rollout_resource_pool=actor_rollout_resource_pool,
+            rm_resource_pool=rm_resource_pool,
+        )
 
     def _save_checkpoint(self):
         from verl.utils.fs import local_mkdir_safe

verl/workers/rollout/trtllm_rollout/trtllm_async_server.py

@@ -25,6 +25,7 @@
 
 from verl.single_controller.ray import RayClassWithInitArgs, SubRayResourcePool
 from verl.utils.config import omega_conf_to_dataclass
+from verl.utils.device import is_cuda_available
 from verl.workers.config import HFModelConfig, RolloutConfig
 from verl.workers.rollout.replica import RolloutMode, RolloutReplica, TokenOutput
 from verl.workers.rollout.trtllm_rollout.trtllm_rollout import TRTLLMAsyncRollout
@@ -61,7 +62,7 @@ def __init__(
         bundle_indices: list[list[int]] = None,
     ):
         os.environ["TRT_LLM_DISABLE_LOAD_WEIGHTS_IN_PARALLEL"] = "1"
-        assert torch.cuda.is_available(), "TRTLLM http server should run on GPU node"
+        assert is_cuda_available, "TRTLLM http server should run on GPU node"
 
         self.config: RolloutConfig = omega_conf_to_dataclass(config)
         self.model_config: HFModelConfig = omega_conf_to_dataclass(model_config, dataclass_type=HFModelConfig)