refactor: update is_episode_claimed function to include unregister_if_not_claimed argument and adjust related logic

Author: binary-husky

ajet/task_runner/base_runner.py

@@ -59,7 +59,7 @@ def should_interrupt_hard_fn() -> bool:
                 return True
             if (observation_window["stop"] is not None) and observation_window["stop"][task_thread_index]: # check soft condition
                 # if soft condition met, check if episode is claimed
-                has_claimed = is_episode_claimed(self.config, workflow_task.episode_uuid)
+                has_claimed = is_episode_claimed(self.config, workflow_task.episode_uuid, unregister_if_not_claimed=True)
                 if not has_claimed:
                     # if not claimed by now (ENGINE.ROLLING_POST), this episode will never be claimed again, so we can hard stop
                     return True

ajet/tuner_lib/experimental/as_swarm_client.py

@@ -11,6 +11,7 @@
 from ajet.schema.task import WorkflowOutput, Task
 from ajet.copilot.job import AgentJetJob
 from ajet.utils.thread_executors import BoundedThreadPoolExecutor
+from ajet.utils.cache import cache_with_ttl
 from ajet.tuner_lib.as_oai_baseurl_apikey import OpenaiBaseUrlAndApiKey
 from ajet.tuner_lib.experimental.swarm_overwatch_utils import CurrentBatchRolloutPoolInformation
 from ajet.tuner_lib.experimental.interchange_utils import (
@@ -31,6 +32,8 @@
 # To prevent stale records from accumulating, do not need to be changed
 CLEAN_RECORD_TIMEOUT = 10
 START_EPISODE_RETRY_DELAY = 15
+TROTTLE_EPISODE_RETRY_DELAY = 2
+WAIT_MORE_AVAIL_EPISODE_RETRY_DELAY = 2
 
 def raise_for_status_with_detail(resp):
     try:
@@ -237,7 +240,7 @@ def _begin_episode_auto_retry(self, discard_episode_timeout=600, episode_type="t
                     should_throttle = self._should_throttle(throttle_policy, pool_info)
                     if should_throttle:
                         self.logger_info(f"Throttle policy is active, delaying episode ...")
-                        retry_delay = START_EPISODE_RETRY_DELAY
+                        retry_delay = TROTTLE_EPISODE_RETRY_DELAY
                         continue
 
                 # connect remote server to claim an episode
@@ -273,18 +276,22 @@ def _begin_episode_auto_retry(self, discard_episode_timeout=600, episode_type="t
                         episode_uuid=episode_uuid
                     )
                 else:
-                    need_wait_scenarios =[
+                    need_snap_scenarios =[
                         "Engine is syncing weights",
                         "Engine is in post-rolling phase",
+                    ]
+                    need_wait_scenarios =[
                         "No available episodes to claim.",
-                        "SwarmThrottlePolicy",
                     ]
-                    if any(scenario in data.fail_cause for scenario in need_wait_scenarios):
+                    if any(scenario in data.fail_cause for scenario in need_snap_scenarios):
                         if time.time() - self.previous_warning_time > 60:
                             self.logger_info(f"{data.fail_cause}. Retrying ...")
                             self.previous_warning_time = time.time()
                         retry_delay = START_EPISODE_RETRY_DELAY
                         continue
+                    elif any(scenario in data.fail_cause for scenario in need_wait_scenarios):
+                        retry_delay = WAIT_MORE_AVAIL_EPISODE_RETRY_DELAY
+                        continue
                     else:
                         logger.warning(f"Failed to claim episode: {data.fail_cause}. Retrying ...")
                         retry_delay = START_EPISODE_RETRY_DELAY
@@ -464,6 +471,7 @@ def _wait_until_status_change_to(self, desired_status="ENGINE.ROLLING", verbose=
                 logger.error(f"Error polling engine status: {e}")
                 time.sleep(5)
 
+    @cache_with_ttl(ttl=0.5)
     def get_engine_status(self) -> Tuple[str, dict]:
         try:
             resp = httpx.get(

ajet/tuner_lib/experimental/as_swarm_server.py

@@ -127,12 +127,14 @@ async def _revert_episode_to_unclaimed(episode_uuid: str, shared_mem_dict, share
         if ep_key(episode_uuid) not in shared_mem_dict:
             logger.warning(f"Episode record for {episode_uuid} not found in shared memory. It may have been already processed by another thread. Skipping unclaim.")
             return
-        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
+
+        with shared_mem_dict_lock:
+            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
         # _context_tracker_reset_blocking(episode_uuid, shared_mem_dict)   # must async
@@ -533,6 +535,15 @@ async def claim_episode(req: ClaimEpisodeRequest):
                 shared_mem_dict["unclaimed_episodes"] = shared_mem_dict["unclaimed_episodes"][1:]
 
                 # get episode
+                if ep_key(episode_uuid) not in shared_mem_dict:
+                    return ClaimEpisodeResponse(
+                        success=False,
+                        client_uuid=req.client_uuid,
+                        episode_uuid="",
+                        openai_base_url="",
+                        openai_api_key="",
+                        fail_cause="No available episodes to claim. Try again (maybe 2 minutes) later.",
+                    )
                 es: EpisodeStatus = shared_mem_dict[ep_key(episode_uuid)]
                 es.episode_status = "claimed"
                 es.episode_type = req.episode_type
@@ -684,12 +695,7 @@ async def is_episode_claimed(req: CheckWhetherEpisodeClaimedRequest):
             return BoolResponse(success=True)
         else:
             if req.unregister_if_not_claimed:
-                # remove from shared memory to avoid stale records
-                with shared_mem_dict_lock:
-                    if ep_key(req.episode_uuid) in shared_mem_dict:
-                        del shared_mem_dict[ep_key(req.episode_uuid)]
-                    if req.episode_uuid in shared_mem_dict["unclaimed_episodes"]:
-                        shared_mem_dict["unclaimed_episodes"].remove(req.episode_uuid)
+                _delete_episode_record(req.episode_uuid, shared_mem_dict, shared_mem_dict_lock)
             return BoolResponse(success=False)
 
     @app.post("/get_episode_buffer", response_model=EpisodeBufferResponse)

ajet/tuner_lib/experimental/interchange_utils.py

@@ -132,10 +132,10 @@ def http_change_engine_status(config, new_status: str, new_status_detail: str|No
     logger.success(f"Changed engine status to {new_status}")
 
 
-def is_episode_claimed(config, episode_uuid: str) -> bool:
+def is_episode_claimed(config, episode_uuid: str, unregister_if_not_claimed: bool) -> bool:
     resp = httpx.post(
         f"{get_interchange_server_url(config)}/is_episode_claimed",
-        json={"episode_uuid": episode_uuid, "unregister_if_not_claimed": True},
+        json={"episode_uuid": episode_uuid, "unregister_if_not_claimed": unregister_if_not_claimed},
         timeout=5
     )
     resp.raise_for_status()