add client side rollout status monitor

Author: binary-husky

ajet/task_rollout/native_parallel_worker.py

@@ -20,9 +20,12 @@
 from ajet.schema.trajectory import Sample
 from ajet.task_rollout.single_worker import BaseRolloutManager
 from ajet.context_tracker.basic_tracker import BaseContextTracker
-from ajet.tuner_lib.weight_tuner.experimental.interchange_utils import http_change_engine_status
-
-DEBUG = True
+from ajet.tuner_lib.weight_tuner.experimental.interchange_utils import (
+    http_change_engine_status,
+    http_update_rollout_pool_information,
+    CurrentBatchRolloutPoolInformation,
+    DEBUG,
+)
 
 
 def spawn_thread_shared_observation_window(n_threads) -> Dict[str, List[int | bool | str]]:
@@ -302,10 +305,14 @@ def stop_all_threads_hard():
 
         def update_rollout_result_array_preview(observation_window, completed_task_id_map_ct: Dict[str, List[BaseContextTracker]]):
             buffer = ""
+            completed_tasks_details = {}
             for task_id, tracker_arr in completed_task_id_map_ct.items():
                 buffer += f"Task {task_id} (completed {len(tracker_arr)} episodes):\n"
+                episode_uuids = []
                 for ct in tracker_arr:
                     buffer += f"\tEpisode: {ct.episode_uuid}\tTimelines: {len(ct.saved_timelines)}\tLLM_Calls: {ct.llm_call_cnt}\tReward: {ct.reward_structure.performance_reward}\n"
+                    episode_uuids.append(ct.episode_uuid)
+                completed_tasks_details[task_id] = episode_uuids
             buffer += f"\n"
             buffer += f"\n"
             counts = count_tasks(completed_task_id_map_ct)
@@ -314,6 +321,20 @@ def update_rollout_result_array_preview(observation_window, completed_task_id_ma
             buffer += f"Total completed non-dummy tasks: {counts['total_completed_non_dummy_tasks']} (target {n_batch_task})\n"
             buffer += f"Current stop condition: {self.config.ajet.swarm_mode_sample_collection_method}\n"
             observation_window["info"][-1] = buffer
+
+            # Update rollout pool information via API
+            pool_info = CurrentBatchRolloutPoolInformation(
+                completed_episodes=counts['total_completed_episodes'],
+                completed_episode_target=n_batch_task * rollout_n,
+                completed_tasks=counts['total_completed_tasks'],
+                completed_task_target=n_batch_task,
+                completed_non_dummy_tasks=counts['total_completed_non_dummy_tasks'],
+                completed_non_dummy_task_target=n_batch_task,
+                task_expected_num_repeat=rollout_n,
+                completed_tasks_details=completed_tasks_details,
+            )
+            http_update_rollout_pool_information(self.config, pool_info)
+
             return
 
         # loop and wait until stop condition is met, then stop threads and collect results
@@ -360,9 +381,6 @@ def update_rollout_result_array_preview(observation_window, completed_task_id_ma
             update_rollout_result_array_preview(observation_window, completed_task_id_map_ct)
             self._write_swarm_rollout_dynamic_log(observation_window)
 
-        time.sleep(10)
-        raise RuntimeError("DEBUG")
-        # return all trackers
         return tracker_array
 
 

ajet/task_runner/swarm_runner.py

@@ -12,12 +12,11 @@
 from ajet.task_runner.base_runner import BaseAgentRunner
 from ajet.utils.retry import SwarmReceiveAbortException
 from ajet.tuner_lib.weight_tuner.experimental.interchange_utils import http_register_episode, get_zmq_socket, is_episode_claimed
+from ajet.tuner_lib.weight_tuner.experimental.interchange_utils import DEBUG
 from loguru import logger
 from ajet import Workflow
 from typing import Callable
 
-DEBUG = True
-# DEBUG = False
 
 context = zmq.Context()
 atexit.register(context.term)

ajet/tuner_lib/weight_tuner/experimental/as_oai_model_client.py

@@ -15,15 +15,14 @@
 from ajet.tuner_lib.weight_tuner.experimental.as_oai_model_server import InterchangeCompletionRequest, API_KEY_PREFIX
 from ajet.utils.thread_executors import SharedInferenceTrackerThreadExecutor, SharedInterchangeThreadExecutor
 from ajet.tuner_lib.weight_tuner.experimental.interchange_utils import get_zmq_socket, is_episode_claimed
+from ajet.tuner_lib.weight_tuner.experimental.interchange_utils import DEBUG
 
 context = zmq.Context()
 atexit.register(context.term)
 
 if TYPE_CHECKING:
     from ajet.context_tracker.multiagent_tracking import MultiAgentContextTracker
 
-# DEBUG = False
-DEBUG = True
 
 def generate_auth_token(agent_name, target_tag, episode_uuid, episode_address):
     """

ajet/tuner_lib/weight_tuner/experimental/as_oai_model_server.py

@@ -35,6 +35,7 @@
 
 from ajet.utils.networking import find_free_port, get_host_ip
 from ajet.tuner_lib.weight_tuner.experimental.interchange_utils import EpisodeStatus
+from ajet.tuner_lib.weight_tuner.experimental.interchange_utils import DEBUG
 
 API_KEY_PREFIX = "sk-ajet-"
 
@@ -54,8 +55,6 @@ class HealthCheckRequest(BaseModel):
 
 # Create FastAPI app
 SERVER_SHUTDOWN_EVENT = threading.Event()
-# DEBUG = False
-DEBUG = True
 
 context = zmq.Context()
 atexit.register(context.term)

ajet/tuner_lib/weight_tuner/experimental/as_swarm_client.py

@@ -17,6 +17,7 @@
     EndEpisodeResponse,
     EpisodeStatus,
     EpisodeBufferResponse,
+    CurrentBatchRolloutPoolInformation,
 )
 
 
@@ -379,3 +380,20 @@ def stop_engine(self):
             self._wait_until_status_change_to(desired_status="ENGINE.OFFLINE")
         except Exception as e:
             logger.error(f"Error stopping engine: {e}")
+
+    def get_rollout_stat(self) -> CurrentBatchRolloutPoolInformation:
+        """
+        Get the current batch rollout pool information from the Swarm server.
+        Returns statistics about completed episodes, tasks, and progress.
+        """
+        try:
+            resp = httpx.get(
+                f"{self.server_url}/get_current_batch_rollout_pool_information",
+                timeout=10
+            )
+            resp.raise_for_status()
+            data = CurrentBatchRolloutPoolInformation.model_validate(resp.json())
+            return data
+        except Exception as e:
+            logger.error(f"Error getting rollout statistics: {e}")
+            return CurrentBatchRolloutPoolInformation()

ajet/tuner_lib/weight_tuner/experimental/as_swarm_server.py

@@ -10,6 +10,7 @@
 from fastapi import FastAPI, HTTPException
 from multiprocessing.managers import DictProxy
 from typing import Coroutine, Optional, Tuple, List
+from ajet.tuner_lib.weight_tuner.experimental.interchange_utils import DEBUG
 from ajet.tuner_lib.weight_tuner.experimental.interchange_utils import (
     SyncTrainConfigRequest,
     ClaimEpisodeRequest,
@@ -23,11 +24,10 @@
     BoolResponse,
     RegisterEpisodeRequest,
     UpdateEngineStatusRequest,
+    CurrentBatchRolloutPoolInformation,
     VALID_STATUSES,
 )
 
-DEBUG = True
-# DEBUG = False
 RCVTIMEO = 2 * 1000
 RCVTIMEO_OUT = 300 * 1000
 RCVTIMEO_WAIT_N = RCVTIMEO_OUT // RCVTIMEO
@@ -55,6 +55,9 @@ def register_enable_swarm_mode_routes(
     if 'unclaimed_episodes' not in shared_mem_dict:
         shared_mem_dict['unclaimed_episodes'] = []
 
+    if 'current_batch_rollout_pool_information' not in shared_mem_dict:
+        shared_mem_dict['current_batch_rollout_pool_information'] = CurrentBatchRolloutPoolInformation()
+
     # ------------------------------------------------------------------------------------------------
     # ------ Recycle claimed episodes that client failed to complete in (promised) time --------------
     # ---------------------------------  claimed -> unclaimed ----------------------------------------
@@ -613,6 +616,28 @@ async def get_episode_buffer():
         return EpisodeBufferResponse(buffer=result)
 
 
+    @app.post("/update_current_batch_rollout_pool_information", response_model=BoolResponse)
+    async def update_current_batch_rollout_pool_information(req: CurrentBatchRolloutPoolInformation):
+        """Update the current batch rollout pool information."""
+        try:
+            with shared_mem_dict_lock:
+                shared_mem_dict['current_batch_rollout_pool_information'] = req
+            return BoolResponse(success=True)
+        except Exception as e:
+            logger.error(f"Error updating current batch rollout pool information: {e}")
+            return BoolResponse(success=False, failure_reason=str(e))
+
+
+    @app.get("/get_current_batch_rollout_pool_information", response_model=CurrentBatchRolloutPoolInformation)
+    async def get_current_batch_rollout_pool_information():
+        """Get the current batch rollout pool information."""
+        try:
+            return shared_mem_dict.get('current_batch_rollout_pool_information', CurrentBatchRolloutPoolInformation())
+        except Exception as e:
+            logger.error(f"Error getting current batch rollout pool information: {e}")
+            return CurrentBatchRolloutPoolInformation()
+
+
 
 
     # --------------------------------------------------------------------
@@ -751,4 +776,4 @@ def kill_process_tree(shared_mem_dict_lock=None, shared_mem_dict=None):
         except:
             pass
 
-        return {"success": False, "error": str(e)}
+        return {"success": False, "error": str(e)}

ajet/tuner_lib/weight_tuner/experimental/interchange_utils.py

@@ -1,7 +1,7 @@
 import os
 import time
 import httpx
-from typing import List
+from typing import List, Dict
 from pydantic import BaseModel
 from loguru import logger
 from ajet.schema.task import WorkflowOutput
@@ -81,8 +81,19 @@ class UpdateEngineStatusRequest(BaseModel):
     engine_status: str = ""
 
 
-# DEBUG = False
-DEBUG = True
+class CurrentBatchRolloutPoolInformation(BaseModel):
+    completed_episodes: int = 0
+    completed_episode_target: int = 0
+    completed_tasks: int = 0
+    completed_task_target: int = 0
+    completed_non_dummy_tasks: int = 0
+    completed_non_dummy_task_target: int = 0
+    task_expected_num_repeat: int = 0
+    completed_tasks_details: Dict[str, List[str]] = {}  # task_id -> list of episode_uuids
+
+
+DEBUG = False
+# DEBUG = True
 
 def get_interchange_server_url(config):
     port = os.getenv("AJET_DAT_INTERCHANGE_PORT")
@@ -155,6 +166,26 @@ def http_register_episode(config,
     return True
 
 
+def http_update_rollout_pool_information(config, pool_info: CurrentBatchRolloutPoolInformation):
+    """
+    Update the rollout pool information on the interchange server.
+
+    Args:
+        config: The configuration object
+        pool_info: CurrentBatchRolloutPoolInformation object with rollout statistics
+    """
+    try:
+        resp = httpx.post(
+            f"{get_interchange_server_url(config)}/update_current_batch_rollout_pool_information",
+            json=pool_info.model_dump(),
+            timeout=5
+        )
+        resp.raise_for_status()
+    except Exception as e:
+        if DEBUG:
+            logger.warning(f"Failed to update rollout pool information: {e}")
+
+
 def get_zmq_socket(config, episode_uuid: str, tag: str = ""):
     interchange_method = config.ajet.interchange_server.interchange_method
     if interchange_method == 'tcp':