pre commit fix

Author: 0russwest0

claw_r1/agent_flow/agent_flow.py

@@ -231,9 +231,18 @@ async def gateway_compute_reward(
         resp.raise_for_status()
         return resp.json()
 
-    async def gateway_submit_steps(self, steps: list[Step]) -> int:
+    async def gateway_submit_steps(
+        self,
+        steps: list[Step],
+        channel: str | None = None,
+    ) -> int:
         """Submit Steps to DataPool via ``POST /submit_steps``.
 
+        Args:
+            steps: Steps to submit.
+            channel: DataPool channel name.  When *None* the Gateway uses
+                its default (``"train"``).
+
         Returns the number of accepted steps.
         """
         client = _get_http_client()
@@ -253,10 +262,9 @@ async def gateway_submit_steps(self, steps: list[Step]) -> int:
                     "metadata": _json_safe(s.metadata),
                 }
             )
-        resp = await client.post(
-            f"{self.gateway_url}/submit_steps",
-            json={"steps": payloads},
-        )
+        url = f"{self.gateway_url}/submit_steps"
+        params = {"channel": channel} if channel else None
+        resp = await client.post(url, json={"steps": payloads}, params=params)
         resp.raise_for_status()
         return resp.json()["accepted"]
 

claw_r1/agent_flow/multi_step_agent_flow.py

@@ -18,7 +18,9 @@
 logger = logging.getLogger(__file__)
 logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN"))
 
-FEEDBACK_PROMPT = "Please verify your answer step by step and provide the final answer again in the format #### <number>."
+FEEDBACK_PROMPT = (
+    "Please verify your answer step by step and provide the final answer again in the format #### <number>."
+)
 
 
 @register("multi_step_agent")
@@ -35,10 +37,9 @@ async def run(self, sampling_params: dict[str, Any], **kwargs) -> int:
         trajectory_uid = uuid4().hex
         prompt_uid = str(kwargs.get("uid", uuid4().hex))
         messages = list(kwargs["raw_prompt"])
+        channel = kwargs.pop("channel", None)
         metadata = {k: v for k, v in kwargs.items() if k != "agent_name"}
 
-        steps: list[Step] = []
-
         for turn in range(self.max_turns):
             prompt_ids = await self.apply_chat_template(messages)
             prompt_ids = prompt_ids[-self.prompt_length :]
@@ -66,9 +67,10 @@ async def run(self, sampling_params: dict[str, Any], **kwargs) -> int:
             )
 
             if not is_last:
-                await self.gateway_submit_steps([step])
+                await self.gateway_submit_steps([step], channel=channel)
                 response_text = self.tokenizer.decode(
-                    response_ids, skip_special_tokens=True,
+                    response_ids,
+                    skip_special_tokens=True,
                 )
                 messages.append({"role": "assistant", "content": response_text})
                 messages.append({"role": "user", "content": FEEDBACK_PROMPT})
@@ -81,6 +83,6 @@ async def run(self, sampling_params: dict[str, Any], **kwargs) -> int:
                     dataset_fields=dataset_fields,
                 )
                 step.reward = reward_result["reward_score"]
-                await self.gateway_submit_steps([step])
+                await self.gateway_submit_steps([step], channel=channel)
 
         return self.max_turns

claw_r1/agent_flow/single_step_single_turn_agent_flow.py

@@ -63,6 +63,7 @@ async def run(self, sampling_params: dict[str, Any], **kwargs) -> int:
         trajectory_uid = uuid4().hex
         prompt_uid = str(kwargs.get("uid", uuid4().hex))
 
+        channel = kwargs.pop("channel", None)
         metadata = {k: v for k, v in kwargs.items() if k != "agent_name"}
 
         step = Step(
@@ -75,5 +76,5 @@ async def run(self, sampling_params: dict[str, Any], **kwargs) -> int:
             is_last=True,
             metadata=metadata,
         )
-        await self.gateway_submit_steps([step])
+        await self.gateway_submit_steps([step], channel=channel)
         return 1

claw_r1/async_main.py

@@ -26,9 +26,9 @@
 from verl.trainer.ppo.utils import Role
 from verl.utils.device import auto_set_device
 
-
 # -- Resource pool helpers ------------------------------------------------
 
+
 def _create_resource_pool_manager(config, roles: list[Role]) -> ResourcePoolManager:
     """Build separate resource pools for trainer vs. rollout roles."""
     resource_pool_spec = {}
@@ -74,6 +74,7 @@ def _create_role_worker_mapping(config):
 
 # -- AsyncTaskRunner ------------------------------------------------------
 
+
 @ray.remote(num_cpus=1)
 class AsyncTaskRunner:
     """Orchestrates async training: creates components, wires them, starts loops."""
@@ -93,8 +94,8 @@ def _initialize(self, config):
         pprint(OmegaConf.to_container(config, resolve=True))
         OmegaConf.resolve(config)
 
-        from verl.utils.fs import copy_to_local
         from verl.utils import hf_processor, hf_tokenizer
+        from verl.utils.fs import copy_to_local
 
         local_path = copy_to_local(
             config.actor_rollout_ref.model.path,
@@ -103,7 +104,9 @@ def _initialize(self, config):
         trust_remote_code = config.data.get("trust_remote_code", False)
         tokenizer = hf_tokenizer(local_path, trust_remote_code=trust_remote_code)
         processor = hf_processor(
-            local_path, trust_remote_code=trust_remote_code, use_fast=True,
+            local_path,
+            trust_remote_code=trust_remote_code,
+            use_fast=True,
         )
 
         role_worker_mapping, ray_worker_group_cls = _create_role_worker_mapping(config)
@@ -139,7 +142,9 @@ def _initialize(self, config):
         max_queue_size = config.async_training.get("max_queue_size", None)
         data_pool_name = "data_pool"
         data_pool = DataPool.options(name=data_pool_name).remote(
-            data_pool_config, verl_backend, max_queue_size=max_queue_size,
+            data_pool_config,
+            verl_backend,
+            max_queue_size=max_queue_size,
         )
         self.components["data_pool"] = data_pool
         self.components["data_pool_name"] = data_pool_name
@@ -213,17 +218,12 @@ def _create_rollouter(self, config):
     def _create_trainer(self, config):
         from claw_r1.async_trainer import AsyncTrainer
 
-        trainer_roles = [
-            role for role in self.components["role_worker_mapping"]
-            if role != Role.Rollout
-        ]
+        trainer_roles = [role for role in self.components["role_worker_mapping"] if role != Role.Rollout]
         trainer = AsyncTrainer.remote(
             config=config,
             tokenizer=self.components["tokenizer"],
             role_worker_mapping={
-                role: cls
-                for role, cls in self.components["role_worker_mapping"].items()
-                if role != Role.Rollout
+                role: cls for role, cls in self.components["role_worker_mapping"].items() if role != Role.Rollout
             },
             resource_pool_manager=_create_resource_pool_manager(config, trainer_roles),
             ray_worker_group_cls=self.components["ray_worker_group_cls"],
@@ -266,6 +266,7 @@ def _run(self):
 
 # -- Hydra entry point ----------------------------------------------------
 
+
 @hydra.main(config_path="config", config_name="async_ppo_trainer", version_base=None)
 def main(config):
     auto_set_device(config)
@@ -275,9 +276,7 @@ def main(config):
         ray_init_kwargs = config.ray_kwargs.get("ray_init", {})
         runtime_env_kwargs = ray_init_kwargs.get("runtime_env", {})
         runtime_env = OmegaConf.merge(default_runtime_env, runtime_env_kwargs)
-        ray_init_kwargs = OmegaConf.create(
-            {**ray_init_kwargs, "runtime_env": runtime_env}
-        )
+        ray_init_kwargs = OmegaConf.create({**ray_init_kwargs, "runtime_env": runtime_env})
         ray.init(namespace="claw_r1_async", **OmegaConf.to_container(ray_init_kwargs))
 
     if not hasattr(config, "async_training"):