PullRequest: 945 rollout succuess

Merge branch test/tmp1102 of [email protected]:inclusionAI/AReaL.git into asystem/gh
https://code.alipay.com/inclusionAI/AReaL/pull_requests/945

Reviewed-by: 楚财 <[email protected]>


* .

Author: 峯回

areal/controller/rollout_controller.py

@@ -301,7 +301,7 @@ async def _wait_callback(self, worker: Worker):
         tik = time.time()
         while result == "NO_RESULT" and time.time() - tik < self.config.request_timeout:
             result = await self.scheduler.async_call_engine(
-                worker.id, "wait_quiet", count=1, timeout=1, max_retries=1
+                worker.id, "wait_quiet", count=1, timeout=3600, max_retries=1
             )
 
         # The RPCServer will return None if the
@@ -402,6 +402,8 @@ def wait(self, count: int, timeout: float | None = None) -> DistributedBatch:
             # Check capacity before submitting
             capacity = self.get_capacity()
             # Submit pending tasks
+            self.logger.info(f"Capacity: {capacity}, pending inputs: {len(self._pending_inputs)}")
+
             for _ in range(capacity):
                 if len(self._pending_inputs) == 0:
                     break
@@ -544,6 +546,7 @@ def prepare_batch(
                         # Capacity exhausted during batch submission, stop and wait
                         break
             try:
+                self.logger.info("Wait for batch...")
                 return self.wait(dataloader.batch_size, timeout=1)
             except TimeoutError:
                 pass

areal/core/staleness_manager.py

@@ -97,6 +97,7 @@ def get_capacity(self, current_version: int) -> int:
 
             # Return the minimum of both constraints
             capacity = min(concurrency_capacity, staleness_capacity)
+            print(f"Capacity: {capacity}, max_concurrent_rollouts: {max_concurrent_rollouts}, rollout_stat: {self.rollout_stat}, consumer_bs: {consumer_bs}")
             return capacity
 
     def on_rollout_submitted(self) -> None:

areal/examples/configs/my001/on_policy.yaml

@@ -54,6 +54,10 @@ rollout:
   model_path: ${tokenizer_path}
   storage_path: "${storage_prefix}/checkpoints"
   seed: ${seed}
+  max_concurrent_rollouts: 64
+  queue_size: null
+  consumer_batch_size: ${train_dataset.batch_size}
+  max_head_offpolicyness: 2
   engine_config:
     attention_backend: "triton"
     disable_custom_all_reduce: true

areal/extension/asystem/ascheduler/__init__.py

@@ -227,8 +227,9 @@ async def async_call_engine(
             EngineCallError: If the method call fails.
         """
         logger.info(f"Async calling '{method}' on worker {worker_id}")
+        max_retries = kwargs.pop("max_retries", 3)
         return await self.rpc_client.async_call_engine(
-            worker_id, method, 3, *args, **kwargs
+            worker_id, method, max_retries, *args, **kwargs
         )
 
     async def create_engine(self, worker_id: str, engine: str, *args, **kwargs) -> Any: