feat(recipe): add ThunderAgent to the Dynamo rollout backend

README.md

@@ -67,6 +67,7 @@ The script requires only `bash`, `git`, `awk`, and `pip`/`pip3` on `PATH`. It do
 | collabllm | [`recipe/collabllm/REQUIRED_VERL.txt`](collabllm/REQUIRED_VERL.txt) |
 | dapo | [`recipe/dapo/REQUIRED_VERL.txt`](dapo/REQUIRED_VERL.txt) |
 | deepeyes | [`recipe/deepeyes/REQUIRED_VERL.txt`](deepeyes/REQUIRED_VERL.txt) |
+| dynamo | [`recipe/dynamo/REQUIRED_VERL.txt`](dynamo/REQUIRED_VERL.txt) |
 | entropy | [`recipe/entropy/REQUIRED_VERL.txt`](entropy/REQUIRED_VERL.txt) |
 | fapo | [`recipe/fapo/REQUIRED_VERL.txt`](fapo/REQUIRED_VERL.txt) |
 | fault_recover | [`recipe/fault_recover/REQUIRED_VERL.txt`](fault_recover/REQUIRED_VERL.txt) |

dynamo/REQUIRED_VERL.txt

@@ -0,0 +1,5 @@
+# Dynamo ThunderAgent — tested core verl commit
+UPSTREAM=https://github.com/verl-project/verl.git
+MODE=pinned_commit
+COMMIT=d82d2777b5dc3e96a8a45168d02660312707ab98
+PIP_INSTALL=pip install verl@git+https://github.com/verl-project/verl.git@d82d2777b5dc3e96a8a45168d02660312707ab98

dynamo/config/dynamo_trainer.yaml

@@ -23,8 +23,13 @@ actor_rollout_ref:
   rollout:
     name: dynamo
     mode: async
+    agent:
+      agent_loop_manager_class: recipe.dynamo.dynamo_agent_loop.DynamoAgentLoopManager
     engine_kwargs:
       dynamo:
+        thunderagent:
+          enabled: true
+          router_block_size: 16
         # All keys optional; sane defaults applied in DynamoHttpServer.
         # namespace: verl_dynamo
         # router_mode: kv               # round-robin | kv | random | least-loaded

dynamo/dynamo_agent_loop.py

@@ -21,6 +21,7 @@
 
 from __future__ import annotations
 
+from contextlib import asynccontextmanager
 from typing import Any, Optional
 from uuid import uuid4
 
@@ -32,6 +33,9 @@
 from verl.workers.rollout.replica import TokenOutput
 from verl.workers.rollout.utils import update_prometheus_config
 
+from .thunderagent import current_program
+from .thunderagent import program_scope as bind_program
+
 
 class DynamoServerManager:
     """Direct manager for the shared Dynamo frontend actor.
@@ -41,10 +45,28 @@ class DynamoServerManager:
     frontend, so verl should only call the single shared Dynamo actor.
     """
 
-    def __init__(self, servers: list[tuple[str, ray.actor.ActorHandle]]):
+    def __init__(
+        self,
+        servers: list[tuple[str, ray.actor.ActorHandle]],
+        *,
+        thunderagent_enabled: bool = False,
+    ):
         if len(servers) != 1:
             raise ValueError(f"DynamoServerManager expects exactly one shared server, got {len(servers)}")
         self.server_address, self.server = servers[0]
+        self.thunderagent_enabled = thunderagent_enabled
+
+    @asynccontextmanager
+    async def program_scope(self):
+        """Bind all turns in one agent-loop run to one Dynamo program."""
+        if not self.thunderagent_enabled:
+            yield
+            return
+        async with bind_program(uuid4().hex, self._finalize_program):
+            yield
+
+    async def _finalize_program(self, session_id: str) -> None:
+        await self.server.finalize_program.remote(session_id=session_id)
 
     async def generate(
         self,
@@ -54,16 +76,32 @@ async def generate(
         sampling_params: dict[str, Any],
         image_data: Optional[list[Any]] = None,
         video_data: Optional[list[Any]] = None,
+        audio_data: Optional[list[Any]] = None,
+        mm_processor_kwargs: Optional[dict[str, Any]] = None,
         **kwargs: Any,
     ) -> TokenOutput:
-        return await self.server.generate.remote(
+        if audio_data is not None or mm_processor_kwargs:
+            raise RuntimeError("Dynamo frontend generate does not support audio inputs or processor options")
+
+        generate_kwargs = dict(
             request_id=request_id or uuid4().hex,
             prompt_ids=prompt_ids,
             sampling_params=sampling_params,
             image_data=image_data,
             video_data=video_data,
             **kwargs,
         )
+        if not self.thunderagent_enabled:
+            return await self.server.generate.remote(**generate_kwargs)
+
+        scope = current_program()
+        if scope is None:
+            raise RuntimeError("Dynamo generation requires an active ThunderAgent program")
+        async with scope.request():
+            return await self.server.generate.remote(
+                **generate_kwargs,
+                thunderagent_session_id=scope.session_id,
+            )
 
 
 class DynamoLLMServerManager(LLMServerManager):
@@ -73,9 +111,9 @@ async def _initialize_llm_servers(self, start_rank: int = 0):
         if self.worker_group is None:
             raise ValueError("Dynamo rollout requires hybrid mode with an actor rollout worker group")
 
-        from recipe.dynamo.dynamo_async_server import DynamoReplica
+        from recipe.dynamo.dynamo_thunderagent import DynamoThunderAgentReplica
 
-        replica = DynamoReplica(
+        replica = DynamoThunderAgentReplica(
             replica_rank=start_rank,
             config=self.rollout_config,
             model_config=self.model_config,
@@ -93,9 +131,25 @@ async def _initialize_llm_servers(self, start_rank: int = 0):
                 raise ValueError("PROMETHEUS needs disable_log_stats==False, but it is currently True.")
             update_prometheus_config(self.rollout_config.prometheus, self.server_addresses, self.rollout_config.name)
 
+    async def _init_global_load_balancer(self) -> None:
+        """Dynamo owns routing behind its single shared frontend."""
+
+    def get_client(self, client_cls=None, **kwargs) -> DynamoServerManager:
+        dynamo_config = (self.rollout_config.engine_kwargs or {}).get("dynamo", {}) or {}
+        thunderagent_config = dynamo_config.get("thunderagent", {}) or {}
+        servers = list(zip(self.server_addresses, self.server_handles, strict=True))
+        return DynamoServerManager(
+            servers,
+            thunderagent_enabled=bool(thunderagent_config.get("enabled", False)),
+        )
+
 
 class DynamoAgentLoopWorker(AgentLoopWorker):
-    """Compatibility wrapper for Dynamo agent loop workers."""
+    """Bind each trajectory to one ThunderAgent program."""
+
+    async def _run_agent_loop(self, *args, **kwargs):
+        async with self.llm_client.program_scope():
+            return await super()._run_agent_loop(*args, **kwargs)
 
 
 class DynamoAgentLoopManager(AgentLoopManager):