rename ServerAdapter

Author: Superjomn

verl/workers/rollout/base.py

@@ -81,7 +81,7 @@ def generate_sequences(self, prompts: DataProto) -> DataProto:
 _ROLLOUT_REGISTRY = {
     ("vllm", "async"): "verl.workers.rollout.vllm_rollout.vLLMAsyncRollout",
     ("sglang", "async"): "verl.workers.rollout.sglang_rollout.sglang_rollout.ServerAdapter",
-    ("trtllm", "async"): "verl.workers.rollout.trtllm_rollout.trtllm_rollout.TRTLLMAsyncRollout",
+    ("trtllm", "async"): "verl.workers.rollout.trtllm_rollout.trtllm_rollout.ServerAdapter",
 }
 
 

verl/workers/rollout/trtllm_rollout/trtllm_async_rollout.md

@@ -51,11 +51,11 @@ flowchart TB
             space2[" "]
             style space2 fill:none,stroke:none
             
-            subgraph AsyncRollout["<b>TRTLLMAsyncRollout<br/>(per DP rank)</b>"]
-                DPLeader["<b>• DP Leader coordination</b>"]
-                IPCMgmt["<b>• IPC handle management</b>"]
-                HTTPAdapter["<b>• HTTP adapter for server communication</b>"]
-            end
+    subgraph AsyncRollout["<b>ServerAdapter<br/>(per DP rank)</b>"]
+        DPLeader["<b>• DP Leader coordination</b>"]
+        IPCMgmt["<b>• IPC handle management</b>"]
+        HTTPAdapter["<b>• HTTP adapter for server communication</b>"]
+    end
             
             AsyncRollout -->|<b>HTTP/REST API</b>| HTTPServer
             
@@ -223,7 +223,7 @@ flowchart TB
 - Validate placement group configurations
 
 
-### 3.3 `TRTLLMAsyncRollout`
+### 3.3 `ServerAdapter`
 
 **Purpose**: Rollout worker that handles weight updates, memory management, and generation via HTTP adapter.
 
@@ -256,7 +256,7 @@ Each DP rank has one leader (the first TP rank within that DP group), and that l
 ```mermaid
 sequenceDiagram
     participant Client as Client/Actor
-    participant Rollout as TRTLLMAsyncRollout
+    participant Rollout as ServerAdapter
     participant Adapter as AsyncHttpAdapter
     participant Server as TRTLLMHttpServer
     participant AsyncLLM as AsyncLLM Engine

verl/workers/rollout/trtllm_rollout/trtllm_async_server.py

@@ -27,7 +27,7 @@
 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
+from verl.workers.rollout.trtllm_rollout.trtllm_rollout import ServerAdapter
 from verl.workers.rollout.utils import is_valid_ipv6_address, run_unvicorn
 
 logger = logging.getLogger(__file__)
@@ -184,20 +184,20 @@ async def wake_up(self):
             # Call all workers to switch between trainer mode and rollout mode.
             await asyncio.gather(*[worker.wake_up.remote() for worker in self.workers])
         elif self.rollout_mode == RolloutMode.COLOCATED:
-            await self.llm.resume(tags=TRTLLMAsyncRollout.get_full_tags())
+            await self.llm.resume(tags=ServerAdapter.get_full_tags())
         elif self.rollout_mode == RolloutMode.STANDALONE:
             logger.info("skip wake_up in standalone mode")
 
     async def sleep(self):
         if self.rollout_mode == RolloutMode.HYBRID:
             await asyncio.gather(*[worker.sleep.remote() for worker in self.workers])
         elif self.rollout_mode == RolloutMode.COLOCATED:
-            await self.llm.release(tags=TRTLLMAsyncRollout.get_full_tags())
+            await self.llm.release(tags=ServerAdapter.get_full_tags())
         elif self.rollout_mode == RolloutMode.STANDALONE:
             logger.info("skip sleep in standalone mode")
 
 
-_rollout_worker_actor_cls = ray.remote(TRTLLMAsyncRollout)
+_rollout_worker_actor_cls = ray.remote(ServerAdapter)
 
 
 class TRTLLMReplica(RolloutReplica):

verl/workers/rollout/trtllm_rollout/trtllm_rollout.py

@@ -277,7 +277,7 @@ async def update_weights(self, weights: dict[str, str]):
         return await self._make_async_request("update_weights", {"weights": weights})
 
 
-class TRTLLMAsyncRollout(BaseRollout):
+class ServerAdapter(BaseRollout):
     _WEIGHTS_TAGS = [
         "sampler",
         "drafter",
@@ -291,7 +291,7 @@ class TRTLLMAsyncRollout(BaseRollout):
 
     @staticmethod
     def get_full_tags() -> list[str]:
-        return TRTLLMAsyncRollout._WEIGHTS_TAGS + ["kv_cache"]
+        return ServerAdapter._WEIGHTS_TAGS + ["kv_cache"]
 
     def __init__(
         self, config: RolloutConfig, model_config: HFModelConfig, device_mesh: DeviceMesh, replica_rank: int = -1
@@ -322,7 +322,7 @@ def __init__(
             logger.info(f"exclude_dp_size = {self.hybrid_device_mesh['exclude_dp'].size()}")
             self.gpu_id = ray.get_gpu_ids()[0]
             self.replica_rank = self.hybrid_device_mesh["dp"].get_local_rank()
-            assert len(ray.get_gpu_ids()) == 1, "TRTLLMAsyncRollout should run on a single GPU node"
+            assert len(ray.get_gpu_ids()) == 1, "ServerAdapter should run on a single GPU node"
         else:
             rank = int(os.environ["RANK"])
             self.replica_rank = replica_rank
@@ -332,7 +332,7 @@ def __init__(
         assert self.replica_rank >= 0, "replica_rank is not set"
         assert self.is_leader_rank is not None, "is_leader_rank is not set"
 
-        print(f"TRTLLMAsyncRollout, replica_rank: {self.replica_rank}, is_leader_rank: {self.is_leader_rank}")
+        print(f"ServerAdapter, replica_rank: {self.replica_rank}, is_leader_rank: {self.is_leader_rank}")
 
         self.node_ip = ray.util.get_node_ip_address().strip("[]")