make rollout more robust

Author: binary-husky

ajet/task_rollout/native_parallel_worker.py

@@ -61,9 +61,9 @@ def step_status_printer(self, observation_window):
             if start == -1:
                 print_buf += [f"[finished]:{count} threads"]
         print(f"Rollout progress ({token_gen_per_sec_str}): " + "  //  ".join(print_buf))
-        if "info" in observation_window:
-            print_buf2 = "\t".join(observation_window["info"])
-            print(print_buf2)
+        # if "info" in observation_window:
+        #     print_buf2 = "\t".join(observation_window["info"])
+        #     print(print_buf2)
 
     def rollout_static(
         self,

ajet/task_runner/swarm_runner.py

@@ -88,6 +88,7 @@ def register_episode_and_wait_output(
                     logger.warning(f"Received reset command for episode {episode_uuid}.")
                     context_tracker.reset()
                     zmq_socket.send_string("ack")
+                    continue
                 elif message == "RUNNER.SPECIAL.ABORT":
                     logger.warning(f"Received abort command for episode {episode_uuid}.")
                     context_tracker.reset()
@@ -104,8 +105,8 @@ def register_episode_and_wait_output(
             raise exc
 
         finally:
+            tuner.terminate_episode()   # this is very important to avoid resource leak
             zmq_socket.close()
-            tuner.terminate_episode()
             if ipc_path and os.path.exists(ipc_path): os.remove(ipc_path)
 
         return final_output

ajet/tuner_lib/weight_tuner/experimental/as_swarm_client.py

@@ -26,6 +26,7 @@ def __init__(self, server_url: str):
         self.server_url = server_url
         self.client_uuid = str(uuid.uuid4())
         self.previous_warning_time = 0
+        self.record_episode_expire_time = {}
 
 
     def begin_episode(self, allow_discard_timeout=60, episode_type="train") -> Tuple[str, OpenaiBaseUrlAndApiKey]:
@@ -48,6 +49,7 @@ def begin_episode(self, allow_discard_timeout=60, episode_type="train") -> Tuple
                 resp.raise_for_status()
                 data = ClaimEpisodeResponse.model_validate(resp.json())
                 episode_uuid = data.episode_uuid
+                self.record_episode_expire_time[episode_uuid] = time.time() + allow_discard_timeout
 
                 if data.success:
                     episode_uuid = data.episode_uuid
@@ -82,6 +84,11 @@ def end_episode(self, task:Task, episode_uuid: str, workflow_output: WorkflowOut
             logger.error("No episode to end.")
             return
 
+        remain_time = self.record_episode_expire_time.get(episode_uuid, 0) - time.time()
+        if remain_time < 0:
+            logger.warning(f"Episode {episode_uuid} has expired (expired {remain_time} seconds ago). Please use a larger `allow_discard_timeout` when `begin_episode`. Skipping end_episode.")
+            return
+
         try:
             task_id = task.task_id
             workflow_output.metadata["task_id"] = task_id
@@ -131,7 +138,7 @@ def abort_episode(self, episode_uuid: str):
             data = EndEpisodeResponse.model_validate(resp.json())
 
             if data.success:
-                logger.info(f"Ended episode {episode_uuid}")
+                logger.info(f"Aborted episode {episode_uuid}")
             else:
                 logger.error(f"Failed to end episode {episode_uuid}")
 

ajet/tuner_lib/weight_tuner/experimental/as_swarm_server.py

@@ -61,20 +61,20 @@ def register_enable_swarm_mode_routes(
     # ------------------------------------------------------------------------------------------------
 
     async def find_claimed_episodes_that_need_to_be_unclaimed() -> List[str]:
-        result = []
+        to_unclaim_episodes = []
         current_time = time.time()
 
         for k, v in shared_mem_dict.items():
             if is_key_epsisode_status(k):
                 es:EpisodeStatus = v
                 if es.episode_status == "claimed":
                     if (current_time - es.latest_activity_timestamp) > es.allow_discard_timeout:
-                        result.append(es.episode_uuid)
+                        to_unclaim_episodes.append(es.episode_uuid)
 
-        for episode_uuid in result:
+        for episode_uuid in to_unclaim_episodes:
             await _revert_episode_to_unclaimed(episode_uuid, shared_mem_dict, shared_mem_dict_lock)
 
-        return result
+        return to_unclaim_episodes
 
     def _context_tracker_reset_blocking(episode_uuid, shared_mem_dict):  # must async
         # send message to context tracker
@@ -110,6 +110,8 @@ def _context_tracker_reset_blocking(episode_uuid, shared_mem_dict):  # must asyn
     async def _revert_episode_to_unclaimed(episode_uuid: str, shared_mem_dict, shared_mem_dict_lock):
         # check status again, because other thread may have changed it
         if shared_mem_dict[ep_key(episode_uuid)].episode_status != "claimed":
+            if episode_uuid in shared_mem_dict['unclaimed_episodes']: pass
+            else: shared_mem_dict['unclaimed_episodes'] += [episode_uuid]
             return
 
         # reset context tracker
@@ -126,17 +128,15 @@ async def _revert_episode_to_unclaimed(episode_uuid: str, shared_mem_dict, share
             es.allow_discard_timeout = -1
             with shared_mem_dict_lock:
                 shared_mem_dict[ep_key(episode_uuid)] = es
-                if episode_uuid in shared_mem_dict['unclaimed_episodes']:
-                    pass
-                else:
-                    shared_mem_dict['unclaimed_episodes'] += [episode_uuid]
+                if episode_uuid in shared_mem_dict['unclaimed_episodes']: pass
+                else: shared_mem_dict['unclaimed_episodes'] += [episode_uuid]
 
     def _delete_episode_record(episode_uuid: str, shared_mem_dict, shared_mem_dict_lock):
 
         with shared_mem_dict_lock:
             # remove episode record
             if ep_key(episode_uuid) in shared_mem_dict:
-                del shared_mem_dict[ep_key(episode_uuid)]
+                del shared_mem_dict[ep_key(episode_uuid)]   # RM--
                 logger.info(f"Deleted episode record for {episode_uuid}.")
             # remove from unclaimed list if present
             if episode_uuid in shared_mem_dict['unclaimed_episodes']:
@@ -499,7 +499,17 @@ async def end_episode(req: EndEpisodeRequest):
 
         # send workflow_output to zmq
         assert 'episodes' in shared_mem_dict
-        episode_type = shared_mem_dict[ep_key(episode_uuid)].episode_type
+        ep_stat = shared_mem_dict[ep_key(episode_uuid)]
+        episode_type = ep_stat.episode_type
+        episode_status = ep_stat.episode_status
+        client_uuid_recorded = ep_stat.client_uuid
+        if client_uuid_recorded != client_uuid:
+            logger.error(f"[server] Episode {episode_uuid} is claimed by different client: {client_uuid_recorded}, but got {client_uuid}.")
+            raise HTTPException(status_code=404, detail=f"Episode {episode_uuid} is claimed by different client: {client_uuid_recorded}, but got {client_uuid}.")
+
+        if episode_status != "claimed":
+            logger.error(f"[server] Episode {episode_uuid} is not in claimed status.")
+            raise HTTPException(status_code=400, detail=f"Episode {episode_uuid} is not in claimed status, maybe you take too long to submit.")
 
         if episode_type == "train":
             # _register_final_episode_output_blocking(episode_uuid, workflow_output, shared_mem_dict, shared_mem_dict_lock)    # must async

ajet/utils/thread_executors.py

@@ -1,22 +1,42 @@
+from concurrent.futures import ThreadPoolExecutor
 from ajet.utils.sington import singleton
-import concurrent.futures
-
+import threading
 
 
 @singleton
 class SharedInterchangeThreadExecutor:
     def __init__(self, max_workers=64):
-        self.executor = concurrent.futures.ThreadPoolExecutor(max_workers=max_workers)
+        self.executor = ThreadPoolExecutor(max_workers=max_workers)
 
-    def get_shared_executor(self) -> concurrent.futures.ThreadPoolExecutor:
+    def get_shared_executor(self) -> ThreadPoolExecutor:
         return self.executor
 
 
 
 @singleton
 class SharedInferenceTrackerThreadExecutor:
     def __init__(self, max_workers=64):
-        self.executor = concurrent.futures.ThreadPoolExecutor(max_workers=max_workers)
+        self.executor = ThreadPoolExecutor(max_workers=max_workers)
 
-    def get_shared_executor(self) -> concurrent.futures.ThreadPoolExecutor:
+    def get_shared_executor(self) -> ThreadPoolExecutor:
         return self.executor
+
+
+class BoundedThreadPoolExecutor:
+    def __init__(self, max_workers, max_queue_size=100):
+        self.executor = ThreadPoolExecutor(max_workers=max_workers)
+        self.semaphore = threading.Semaphore(max_queue_size)
+
+    def submit(self, fn, *args, **kwargs):
+        self.semaphore.acquire()
+
+        def wrapped_fn(*args, **kwargs):
+            try:
+                return fn(*args, **kwargs)
+            finally:
+                self.semaphore.release()
+
+        return self.executor.submit(wrapped_fn, *args, **kwargs)
+
+    def shutdown(self, wait=True):
+        self.executor.shutdown(wait=wait)

tutorial/example_academic_trans/trans.py

@@ -49,7 +49,7 @@ def execute_agent(task: Task, api_baseurl_key: OpenaiBaseUrlAndApiKey):
         grader = TranslationQualityGrader(
             model=OpenAIChatModel(base_url=grader_base_url, api_key=grader_api_key, model="qwen3-max-2026-01-23")
         )
-        grader_score = asyncio.run(grader.aevaluate(original_text=abstract, translation=final_translation))
+        grader_score = asyncio.run(asyncio.wait_for(grader.aevaluate(original_text=abstract, translation=final_translation), timeout=120))
         raw_reward = grader_score.score
         print(f"Grader Score: {grader_score.score}, Reason: {grader_score.reason}, Metadata: {grader_score.metadata}")
     return WorkflowOutput(reward=raw_reward, metadata={
@@ -111,7 +111,8 @@ def detect_hard_proper_nouns(messages, base_url, api_key, abstract, rough_transl
     response = client.chat.completions.create(
         model="qwen3-max-2026-01-23",
         messages=messages,
-        extra_body={"enable_thinking":True}
+        timeout=60,
+        # extra_body={"enable_thinking":True}
     )
     fix_nouns = response.choices[0].message.content
     messages += [

tutorial/example_academic_trans/trans_roll.py

@@ -8,7 +8,7 @@
 from ajet.tuner_lib.weight_tuner.experimental.as_swarm_client import SwarmClient
 from ajet.default_config.ajet_default import AjetTaskReader, HuggingfaceDatRepo
 from ajet.tuner_lib.weight_tuner.as_oai_baseurl_apikey import OpenaiBaseUrlAndApiKey
-from ajet import WorkflowOutput
+from ajet.utils.thread_executors import BoundedThreadPoolExecutor
 from ajet.schema.task import Task
 from ajet.task_reader import RouterTaskReader
 from ajet.utils.retry import retry_with_backoff
@@ -56,7 +56,7 @@ def main():
             model=REMOTE_TRAIN_MODEL_01,
             batch_size=REMOTE_BATCH_SIZE,
             grpo_n=LOCAL_GRPO_N,
-        )
+        ),
     )
 
     def rollout(task):
@@ -80,20 +80,17 @@ def rollout(task):
             logger.exception("Exception during rollout group", e)
 
     task_batch = []
+    executor = BoundedThreadPoolExecutor(max_workers=LOCAL_MAX_PARALLEL, max_queue_size=LOCAL_MAX_PARALLEL*2)
     for i, task in enumerate(dataset.generate_training_tasks()):
         task_batch += [task]
 
         if len(task_batch) == REMOTE_BATCH_SIZE:
             print('*********** beginning a new batch of tasks... ***********')
-            with ThreadPoolExecutor(max_workers=LOCAL_MAX_PARALLEL) as executor:
-                for task in task_batch:
-                    executor.submit(rollout, task)
-            executor.shutdown(wait=True)
+            for task in task_batch:
+                executor.submit(rollout, task)
             task_batch = []
-            print('*********** tasks completed, wait a minute... ***********')
-            time.sleep(60)
-
 
+    executor.shutdown(wait=True)
     return None