update rollout result when stop_condition_callback

Author: binary-husky

ajet/task_rollout/native_parallel_worker.py

@@ -187,6 +187,7 @@ def rollout_swarm(  # noqa: C901
         completed_task_id_map_ct: Dict[str, List[BaseContextTracker]] = {}
         executor_lock = threading.Lock()
 
+        # count tasks to see whether we have reach the finish line for next weight update
         def count_tasks(completed_task_id_map_ct):
             total_completed_episodes = 0
             total_completed_tasks = 0
@@ -290,7 +291,9 @@ def stop_condition_callback(completed_task_id_map_ct):
             if stop_condition(completed_task_id_map_ct):
                 if not is_already_soft_stopped():
                     stop_all_threads_soft()
+                    update_rollout_result_array_preview(observation_window, completed_task_id_map_ct)
                 return True
+            update_rollout_result_array_preview(observation_window, completed_task_id_map_ct)
             return False
 
         # submit initial tasks
@@ -388,7 +391,6 @@ def update_rollout_result_array_preview(observation_window, completed_task_id_ma
             tracker.current_batch_success_rate = float(task_success_rate)
             tracker.current_batch_reward = float(task_scalar_reward)
 
-        # for debugging
         update_rollout_result_array_preview(observation_window, completed_task_id_map_ct)
         self._write_swarm_rollout_dynamic_log(observation_window)