impl multi rollout stop condition

Author: binary-husky

ajet/backbone/trainer_verl.py

@@ -610,10 +610,10 @@ def fit(self):  # noqa: C901
                             traj.reward_structure.madness for traj in context_tracker_arr
                         ]
                         # reward = [traj.reward_structure.raw_reward for traj in context_tracker_arr]
-                        round_cnt = [traj.round_cnt for traj in context_tracker_arr]
+                        llm_call_cnt = [traj.llm_call_cnt for traj in context_tracker_arr]
                         metrics.update(
                             {
-                                "critic/round_cnt": np.mean(round_cnt),
+                                "critic/llm_call_cnt": np.mean(llm_call_cnt),
                                 "critic/madness_rate": np.mean(madness_rate),
                                 "critic/success_rate": np.mean(success_rate),
                                 "critic/real_success_rate": np.mean(

ajet/context_tracker/base_tracker.py

@@ -144,7 +144,7 @@ def __init__(self, config, tokenizer, workflow_task: WorkflowTask, **kwargs):
 
         # meta data attributes
         self.tag = ""
-        self.round_cnt = 0
+        self.llm_call_cnt = 0
         self.generated_token_cnt = 0
         self.current_batch_success_rate: float = float("-inf")
         self.current_batch_reward: float = float("-inf")
@@ -171,7 +171,7 @@ def reset(self):
         self.current_batch_success_rate: float = float("-inf")
         self.current_batch_reward: float = float("-inf")
         self.already_mad_flag: bool = False
-        self.round_cnt = 0
+        self.llm_call_cnt = 0
         self.log_metrics: Optional[Dict[str, Union[float, List[float], Dict[str, Any]]]] = None
 
     def group_tokenize(self):

ajet/context_tracker/multiagent_tracking.py

@@ -201,6 +201,8 @@ def step_prepare(self, messages: List[dict], tools: List = [], timeline_uuid: st
         custom_sampling_params = {}
         if not context_safe:
             self.context_overflow = True
+            logger.warning(f"[{self.workflow_task.episode_uuid}] Stop tracking timelines because {info}.")
+
 
         self.timeline_cache[timeline_uuid] = timeline
         return context_safe, token_overflow, info, converted_message, custom_sampling_params, tools
@@ -218,7 +220,7 @@ def step_track(
         assert timeline_uuid in self.timeline_cache, "Timeline UUID not found in cache. Please ensure `step_prepare` is called before `step_track`."
 
         # round ++
-        self.round_cnt += 1
+        self.llm_call_cnt += 1
 
         # get timeline from cache
         timeline = self.timeline_cache.pop(timeline_uuid, [])

ajet/default_config/ajet_default.yaml

@@ -293,7 +293,17 @@ ajet:
     max_fastapi_threads: 512  # 64 or 128 is fine
     max_inference_tracker_threads: 64 # recommend to be equal to `ajet.rollout.max_env_worker`
     already_started: False # do not edit, used by `swarm`
-
+  # what is the stop condition for swarm mode sample collection
+  # "rollout_until_finish_enough_episodes":
+  #   AgentJet simply counts the number of completed episodes, and stop when it has collected [>= (ajet.data.train_batch_size * ajet.rollout.num_repeat)] samples
+  # "rollout_until_finish_enough_tasks":
+  #   AgentJet will identify the **task_id** of each episode, and stop when it has collected [>= ajet.data.train_batch_size] unique & FINISHED **task_id**.
+  #     (Hint: a **task_id** is considered "FINISHED" when [>= ajet.rollout.num_repeat] episodes of this **task_id** have been completed.)
+  # "rollout_until_finish_enough_non_dummy_tasks":
+  #   AgentJet will identify the **task_id** of each episode, and stop when it has collected [>= ajet.data.train_batch_size] unique & FINISHED & NON-DUMMY **task_id**.
+  #     (Hint: a **task_id** is considered "NON-DUMMY" at least one of **episodes** of **task_id** has **different** reward value.)
+  swarm_mode_sample_collection_method: "rollout_until_finish_enough_tasks"
+  swarm_mode_sample_collection_max_cached_episodes: 9999
 
   task_runner:
     # submit llm infer submit method
@@ -303,7 +313,6 @@ ajet:
     wrapper_type: "asyncio-with-gc"
     # - wrapper_type: "asyncio-with-gc":  safe, with periodic garbage collection to prevent event loop leaks (recommended)
     # - wrapper_type: "asyncio":          fast, but may cause event loop leak in long run
-    # - wrapper_type: "multi-processing": safe, but resource consuming
 
     # when `wrapper_type` is `multi-processing`, the timeout for each task
     wrapper_multiprocessing_timeout: 3600  # in seconds

ajet/default_config/ajet_ts_default.yaml

@@ -33,6 +33,17 @@ ajet:
     max_inference_tracker_threads: 64 # recommend to be equal to `ajet.rollout.max_env_worker`
     already_started: False # do not edit, used by `swarm`
 
+  # the method to determine when to stop rollout in swarm mode. Options:
+  # "rollout_until_finish_enough_episodes":
+  #   AgentJet simply counts the number of completed episodes, and stop when it has collected [>= (ajet.data.train_batch_size * ajet.rollout.num_repeat)] samples
+  # "rollout_until_finish_enough_tasks":
+  #   AgentJet will identify the **task_id** of each episode, and stop when it has collected [>= ajet.data.train_batch_size] unique & FINISHED **task_id**.
+  #     (Hint: a **task_id** is considered "FINISHED" when [>= ajet.rollout.num_repeat] episodes of this **task_id** have been completed.)
+  # "rollout_until_finish_enough_non_dummy_tasks":
+  #   AgentJet will identify the **task_id** of each episode, and stop when it has collected [>= ajet.data.train_batch_size] unique & FINISHED & NON-DUMMY **task_id**.
+  #     (Hint: a **task_id** is considered "NON-DUMMY" at least one of **episodes** of **task_id** has **different** reward value.)
+  swarm_mode_sample_collection_method: "rollout_until_finish_enough_tasks"
+
   rollout:
     # maximum number of parallel environments / simulate workers
     max_env_worker: 128

ajet/default_config/verl/verl_default.yaml

@@ -6,7 +6,6 @@ ajet:
   rollout:
     step_skip_action: 0
     submit_oversample_multiplier: 1.5
-    enable_oversample: False
 
 actor_rollout_ref:
   actor:

ajet/task_rollout/async_llm_bridge.py

@@ -5,28 +5,23 @@
 import uuid
 from typing import Any, Callable, Dict, List, Literal, Union
 
-
-
 from loguru import logger
 from omegaconf import DictConfig
 from pydantic import BaseModel
 from vllm.entrypoints.openai.tool_parsers.hermes_tool_parser import Hermes2ProToolParser
 from vllm.outputs import RequestOutput as VerlVllmRequestOutput
-
 from agentscope.model import ChatResponse as AgentScopeChatResponse
 from openai.types.chat.chat_completion import ChatCompletion as OpenAIChatCompletion
 
-ChatResponse = Union[OpenAIChatCompletion, AgentScopeChatResponse]
-
-from ajet.context_tracker.multiagent_tracking import (
-    MultiAgentContextTracker,
-)
-from ajet.schema.convertion import convert_llm_proxy_response_to_oai_response
-from ajet.schema.convertion import convert_llm_proxy_response_to_agentscope_response
 from ajet.schema.logprob import TokenAndProb
+from ajet.utils.tokenizer import ajet_apply_chat_template
 from ajet.utils.async_utils import run_async_coroutine_with_timeout
 from ajet.utils.testing_utils import _mock_if_test_mode, _test_if_test_mode
-from ajet.utils.tokenizer import ajet_apply_chat_template
+from ajet.schema.convertion import convert_llm_proxy_response_to_oai_response
+from ajet.schema.convertion import convert_llm_proxy_response_to_agentscope_response
+from ajet.context_tracker.multiagent_tracking import MultiAgentContextTracker
+
+ChatResponse = Union[OpenAIChatCompletion, AgentScopeChatResponse]
 
 
 class AjetStandardLlmBridgeRequest(BaseModel):