improve readability

Author: binary-husky

ajet/backbone/trainer_verl.py

@@ -563,8 +563,7 @@ def fit(self):  # noqa: C901
                 # pass global_steps to trace
                 gen_batch.meta_info["global_steps"] = self.global_steps
                 is_last_step = self.global_steps >= self.total_training_steps
-                from ajet import bp
-                bp("BATCH")
+
                 with marked_timer("step", timing_raw):
                     # generate a batch
                     logger.info("rollout step begin")
@@ -597,6 +596,8 @@ def fit(self):  # noqa: C901
                         context_tracker_arr: List[BaseContextTracker] = self.parallel_env.rollout(
                             tasks, mode="sample", epoch=f"train.{epoch}"
                         )
+                        from ajet import bp
+                        bp("BATCH")
                         logger.info("end fit rollout")
                         gen_batch_output = self.parallel_env.to_dataproto(context_tracker_arr)
                         logger.info("end dataproto convertion")

ajet/context_tracker/base_tracker.py

@@ -113,34 +113,45 @@ def replace_token_ids(
 class BaseTracker(object):
     def __init__(self, config, tokenizer, workflow_task: WorkflowTask, **kwargs):
 
+        # disable read only mode
+        self._read_only = False
+        self._discarded = False
+
+        # task related info
         self.workflow_task = workflow_task
         self.task_batch_index = self.workflow_task.task_batch_index
         self.task_tag: str = self.workflow_task.task_tag
         self.task_id: str = self.workflow_task.task_id
         self.episode_uuid = self.workflow_task.episode_uuid
 
-        self.config = config
+        # tokenizer
         self.tokenizer = tokenizer
+        self.blackout_token_combo = tokenizer.encode("<|im_start|>assistant\n")
+        self._im_start_token_id = tokenizer.encode("<|im_start|>")[0]
+
+        # config
+        self.config = config
         self.saved_timelines: List[List[ExtendedMessage]] = []
         self.current_context_status = ""
+
+        # length control
         max_response_length = self.config.ajet.rollout.max_response_length_in_one_turn
         max_model_len: int = self.config.ajet.rollout.max_model_len
         self.max_seq_length: int = max_model_len - max_response_length
-        self.blackout_token_combo = tokenizer.encode("<|im_start|>assistant\n")
-        self._im_start_token_id = tokenizer.encode("<|im_start|>")[0]
-        self.generated_token_cnt = 0
-        self.terminal_rewards_dict = {}
-        self.discarded = False
-        self.is_terminated = False
-        self.reward_structure: Union[Reward, None] = None
-        self.context_time_cost = 0
+
+        self.generation_prompt_token = None
+        self.log_metrics: Optional[Dict[str, Union[float, List[float], Dict[str, Any]]]] = None  # Initialize workflow_metadata to store tool statistics
+
+        # meta data attributes
         self.tag = ""
+        self.round_cnt = 0
+        self.generated_token_cnt = 0
         self.current_batch_success_rate: float = float("-inf")
         self.current_batch_reward: float = float("-inf")
+
+        # reward and madness detection
+        self.reward_structure: Union[Reward, None] = None
         self.already_mad_flag: bool = False
-        self.round_cnt = 0
-        self.generation_prompt_token = None
-        self.log_metrics: Optional[Dict[str, Union[float, List[float], Dict[str, Any]]]] = None  # Initialize workflow_metadata to store tool statistics
 
         assert (
             self.config.ajet.data.max_prompt_length
@@ -149,13 +160,13 @@ def __init__(self, config, tokenizer, workflow_task: WorkflowTask, **kwargs):
         )
 
     def reset(self):
+        # disable read only mode
+        self._read_only = False
+        self._discarded = False
+
         self.saved_timelines: List[List[ExtendedMessage]] = []
         self.current_context_status = ""
-        self.terminal_rewards_dict = {}
-        self.discarded = False
-        self.is_terminated = False
         self.reward_structure: Union[Reward, None] = None
-        self.context_time_cost = 0
         self.tag = ""
         self.current_batch_success_rate: float = float("-inf")
         self.current_batch_reward: float = float("-inf")

ajet/context_tracker/basic_tracker.py

@@ -24,7 +24,6 @@ class BaseContextTracker(BaseTracker):
         full_context (List[ExtendedMessage]): List of all messages in the conversation
         current_context_status (str): Current status of the context
         max_seq_length (int): Maximum sequence length for the context window
-        terminal_rewards_dict (dict): Dictionary storing terminal rewards
     """
 
     def __init__(self, config, tokenizer, **kwargs):

ajet/context_tracker/multiagent_tracking.py

@@ -216,7 +216,12 @@ def step_track(
         timeline_uuid: str = "",
     ):
         assert timeline_uuid in self.timeline_cache, "Timeline UUID not found in cache. Please ensure `step_prepare` is called before `step_track`."
-        timeline = self.timeline_cache.get(timeline_uuid, [])
+
+        # round ++
+        self.round_cnt += 1
+
+        # get timeline from cache
+        timeline = self.timeline_cache.pop(timeline_uuid, [])
         if not self.already_mad_flag:
             if (
                 compute_string_madness(
@@ -291,6 +296,11 @@ def save_llm_interaction_timeline(self, tools, llm_ext_msg, timeline):
         for i in range(1, len(timeline)):
             assert not timeline[i].first_message
 
+        # no longer write anything
+        if self._read_only:
+            logger.exception("Timeline is in read-only mode, should not save new timeline. Please report a github issue if you see this error.")
+            return
+
         # save to self.saved_timelines
         self.saved_timelines += [copy.deepcopy(timeline)]
 
@@ -556,6 +566,8 @@ def generate_log(self, task_id=None, global_step="NA"):
     def group_merge(self) -> List[List[ExtendedMessage]]:
         timeline_merging_policy: TimelineMergingPolicyConfig = self.config.ajet.context_tracker.timeline_merging_policy
         self.saved_timelines = merge_tracker_timelines(self.saved_timelines, timeline_merging_policy)
+        self._read_only = True
+
         return self.saved_timelines
 
 

ajet/task_rollout/native_parallel_worker.py

@@ -255,7 +255,7 @@ def rollout_dynamic(  # noqa: C901
                     completed_task_futures = [f for f in task_future_array if f.done()]
                     completed_results = [f.result() for f in completed_task_futures]
                     completed_results = [
-                        tracker for tracker in completed_results if not tracker.discarded
+                        tracker for tracker in completed_results if not tracker._discarded
                     ]
                     reward = [
                         tracker.reward_structure.performance_reward for tracker in completed_results
@@ -306,7 +306,7 @@ def rollout_dynamic(  # noqa: C901
                         )
                     time.sleep(5)
 
-            # We have enough number of samples, but we need to wait for all threads to finish, including discarded threads
+            # We have enough number of samples, but we need to wait for all threads to finish, including ._discarded threads
             tic = -1
             while any(f.running() for task_future_array in futures for f in task_future_array):
                 tic += 1
@@ -325,7 +325,7 @@ def rollout_dynamic(  # noqa: C901
                 completed_task_futures = [f for f in task_future_array if f.done()]
                 completed_results = [f.result() for f in completed_task_futures]
                 completed_results = [
-                    tracker for tracker in completed_results if not tracker.discarded
+                    tracker for tracker in completed_results if not tracker._discarded
                 ]
                 task_cmd_reward_array = [
                     tracker.reward_structure.performance_reward for tracker in completed_results
@@ -409,7 +409,7 @@ def rollout_dynamic(  # noqa: C901
                 completed_task_futures = [f for f in task_future_array if f.done()]
                 completed_results = [f.result() for f in completed_task_futures]
                 completed_results = [
-                    tracker for tracker in completed_results if not tracker.discarded
+                    tracker for tracker in completed_results if not tracker._discarded
                 ]
                 # in-group success rate and reward
                 task_cmd_reward_array = [
@@ -583,6 +583,19 @@ def stop_all_threads_hard():
         for ct_list in completed_task_id_map_ct.values():
             tracker_array.extend(ct_list)
 
+
+        # TODO: support multi-step reward
+        task_success_rate = np.mean(
+            [tracker.reward_structure.success_rate for tracker in tracker_array]
+        )
+        task_scalar_reward = np.mean(
+            [tracker.reward_structure.final_scalar_reward for tracker in tracker_array]
+        )
+
+        for tracker in tracker_array:
+            tracker.current_batch_success_rate = float(task_success_rate)
+            tracker.current_batch_reward = float(task_scalar_reward)
+
         # return all trackers
         return tracker_array
 

ajet/task_rollout/single_worker.py

@@ -160,14 +160,18 @@ def rollout_env_worker_loop(
         **kwargs,
     ):
         try:
+
             cnt = 1
+
             while True:
 
-                if observation_window["stop"][task_thread_index]:
-                    print('rollout_env_worker_loop received stop signal, exiting...')
+                if observation_window["stop"][task_thread_index]:           # since we use multi-threading, the best way to communicate with main thread is through shared memory.
                     return
 
-                observation_window["info"][task_thread_index] = str(cnt)
+                observation_window["info"][task_thread_index] = str(cnt)    # observe how many iterations have been done in the loop
+
+                # Let's begin working on the task, the result `tracker` will contain everything: reward, llm calls, conversation history, etc.
+                # Later we will gather all trackers and do post-processing, generating samples for VeRL.
                 tracker = self.rollout_env_worker(
                     task=task,
                     task_batch_index=task_batch_index,
@@ -185,7 +189,9 @@ def rollout_env_worker_loop(
                             completed_task_id_map_ct[tracker.task_id] = [tracker]
                         else:
                             completed_task_id_map_ct[tracker.task_id] += [tracker]
+
                 cnt += 1
+
                 if observation_window["stop"][task_thread_index]:
                     return
                 else:

ajet/task_runner/swarm_runner.py

@@ -130,13 +130,15 @@ def execute(self, workflow_task: WorkflowTask) -> BaseContextTracker:
             print(f'Exiting workflow worker due to interrupt signal for episode {workflow_task.episode_uuid}.')
             raise SwarmReceiveAbortException(f"Episode {workflow_task.episode_uuid} aborted due to interrupt signal.")
 
+        # context tracker will trace and gather everything we need for training
         context_tracker = MultiAgentContextTracker(
             llm_inference_fn=self.llm_inference_fn,
             tokenizer=self.tokenizer,
             config=self.config,
             workflow_task = workflow_task,
             **hooks,
         )
+        # tuner will handle the communication and provide `baseurl_apikey`
         tuner = AjetTuner(
             context_tracker=context_tracker,
             llm_inference_fn=self.llm_inference_fn,