add llm_call_count to EpisodeStatus and update related logic in swarm server and monitoring

Author: binary-husky

ajet/tuner_lib/experimental/as_oai_model_server.py

@@ -169,10 +169,11 @@ async def chat_completions(request: Request, authorization: str = Header(None)):
             if ep_key(episode_uuid) not in shared_mem_dict:
                 raise HTTPException(status_code=404, detail=f"Episode {episode_uuid} not found.")
 
-            # update activate timestamp
+            # update activate timestamp and increment llm call counter
             with shared_mem_dict_lock:
                 es:EpisodeStatus = shared_mem_dict[ep_key(episode_uuid)]
                 es.latest_activity_timestamp = time.time()
+                es.llm_call_count += 1
                 shared_mem_dict[ep_key(episode_uuid)] = es
 
         # Add to received queue

ajet/tuner_lib/experimental/as_swarm_server.py

@@ -132,6 +132,7 @@ async def _revert_episode_to_unclaimed(episode_uuid: str, shared_mem_dict, share
             es.episode_status = "registered"
             es.client_uuid = ""
             es.latest_activity_timestamp = time.time()
+            es.llm_call_count = 0
             es.discard_episode_timeout = -1
             with shared_mem_dict_lock:
                 shared_mem_dict[ep_key(episode_uuid)] = es
@@ -461,6 +462,7 @@ async def register_episode(req: RegisterEpisodeRequest):
             discard_episode_timeout=-1,
         )
         es.latest_activity_timestamp = time.time()
+        es.llm_call_count = 0
 
         with shared_mem_dict_lock:
             shared_mem_dict[ep_key(episode_uuid)] = es
@@ -519,6 +521,7 @@ async def claim_episode(req: ClaimEpisodeRequest):
                 es.episode_type = req.episode_type
                 es.client_uuid = req.client_uuid
                 es.latest_activity_timestamp = time.time()
+                es.llm_call_count = 0
                 es.discard_episode_timeout = req.discard_episode_timeout
 
                 shared_mem_dict[ep_key(episode_uuid)] = es
@@ -716,6 +719,7 @@ async def get_current_batch_rollout_pool_information():
                             "episode_status": es.episode_status,
                             "time_since_last_activity": f"{time_since_last_activity:.1f}s",
                             "discard_episode_timeout": f"{es.discard_episode_timeout:.1f}s",
+                            "llm_call_count": str(es.llm_call_count),
                         }
             pool_info.running_episode_details = running_episode_details if running_episode_details else None
 

ajet/tuner_lib/experimental/interchange_utils.py

@@ -62,6 +62,7 @@ class EpisodeStatus(BaseModel):
     zmq_listen_result_addr: str = ""
     latest_activity_timestamp: float = time.time()
     discard_episode_timeout: float
+    llm_call_count: int = 0
     debug_log: List[str] = []
 
 class EpisodeBufferResponse(BaseModel):

ajet/tuner_lib/experimental/swarm_overwatch_utils.py

@@ -12,7 +12,7 @@ class CurrentBatchRolloutPoolInformation(BaseModel):
     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
-    running_episode_details: Dict[str, Dict[str, str]] | None = None # episode_uuid -> { "episode_status": ..., "time_since_last_activity": ..., "discard_episode_timeout": ...}
+    running_episode_details: Dict[str, Dict[str, str]] | None = None # episode_uuid -> { "episode_status": ..., "time_since_last_activity": ..., "discard_episode_timeout": ..., "llm_call_count": ...}
     engine_status: str | None = None
     global_step: int | None = None
     booting_start_time: float | None = None  # timestamp when ENGINE.BOOTING started

ajet/utils/swarm_overwatch.py

@@ -244,11 +244,12 @@ def create_running_episodes_table(
         )
 
         table.add_column("Episode UUID", style="cyan", no_wrap=True, width=20, overflow="ellipsis")
-        table.add_column("Status", style="green", width=15)
-        table.add_column("Last Req / Patience", style="yellow", width=30)
+        table.add_column("Status", style="green")
+        table.add_column("LLM Calls", style="magenta", justify="right")
+        table.add_column("Last Req / Patience", style="yellow")
 
         if not info.running_episode_details:
-            table.add_row("[dim]No running episodes[/dim]", "", "")
+            table.add_row("[dim]No running episodes[/dim]", "", "", "")
             return table
 
         # Sort by time since last activity (descending)
@@ -261,15 +262,17 @@ def create_running_episodes_table(
         for episode_uuid, details in sorted_episodes[:30]:
             last_req = details["time_since_last_activity"]
             patience = details.get("discard_episode_timeout", "N/A")
+            llm_calls = details.get("llm_call_count", "0")
             table.add_row(
                 episode_uuid[:40] if len(episode_uuid) > 40 else episode_uuid,
                 details["episode_status"],
+                llm_calls,
                 f"{last_req} / {patience}",
             )
 
         if len(sorted_episodes) > 30:
             table.add_row(
-                f"[dim]... and {len(sorted_episodes) - 30} more episodes[/dim]", "", ""
+                f"[dim]... and {len(sorted_episodes) - 30} more episodes[/dim]", "", "", ""
             )
 
         return table

tutorial/example_academic_trans/trans_roll.py

@@ -13,7 +13,7 @@
 LOCAL_NUM_EPOCH = 1
 LOCAL_MAX_PARALLEL = 32
 LOCAL_DATASET_PATH = "/mnt/data_cpfs/qingxu.fu/agentjet/agentjet/tmp/arxiv_papers/train.parquet"
-REMOTE_SWARM_URL = "http://localhost:10099" # Change to your swarm remote url
+REMOTE_SWARM_URL = "http://localhost:10086" # Change to your swarm remote url
 
 # --------- configurations that take effect remotely -------------
 REMOTE_BATCH_SIZE = 8