fix: hook up startup_recovery with new functions

Author: omar-ridges

api/src/backend/queries/agents.py

@@ -183,3 +183,7 @@ async def agent_startup_recovery(conn: asyncpg.Connection):
     # Legacy status recovery for backward compatibility
     await conn.execute("UPDATE miner_agents SET status = 'awaiting_screening_1' WHERE status = 'screening'")
     await conn.execute("UPDATE miner_agents SET status = 'waiting' WHERE status = 'evaluation'")  # Legacy alias
+
+@db_operation
+async def set_agent_status_by_version_id(conn: asyncpg.Connection, version_id: str, status: str):
+    await conn.execute("UPDATE miner_agents SET status = $1 WHERE version_id = $2", status, version_id)

api/src/backend/queries/agents.pyi

@@ -15,4 +15,5 @@ async def set_approved_agents_to_awaiting_screening() -> List[MinerAgent]: ...
 async def get_all_approved_version_ids() -> List[str]: ...
 async def set_agent_status(version_id: str, status: str): ...
 async def upload_miner_agent(version_id: str, miner_hotkey: str, agent_name: str, version_num: int, ip_address: str): ...
-async def agent_startup_recovery() -> None: ...
+async def agent_startup_recovery() -> None: ...
+async def set_agent_status_by_version_id(version_id: str, status: str): ...

api/src/endpoints/model_replacers.py

@@ -3,10 +3,10 @@
 """
 
 from api.src.backend.entities import AgentStatus, MinerAgent
-from api.src.backend.queries.agents import set_agent_status
-from api.src.backend.queries.evaluations import get_evaluation_by_evaluation_id, get_running_evaluations, get_stuck_evaluations, get_waiting_evaluations, cancel_dangling_evaluation_runs, reset_evaluation_to_waiting
+from api.src.backend.queries.agents import get_top_agent, set_agent_status
+from api.src.backend.queries.evaluations import get_running_evaluations, get_stuck_evaluations, get_waiting_evaluations, cancel_dangling_evaluation_runs, reset_evaluation_to_waiting, update_evaluation_to_error
 from api.src.backend.queries.agents import agent_startup_recovery
-from api.src.utils.config import SCREENING_1_THRESHOLD, SCREENING_2_THRESHOLD
+from api.src.endpoints.screener import atomically_update_agent_status, finish_evaluation, prune_queue
 from loggers.logging_utils import get_logger
 
 logger = get_logger(__name__)
@@ -38,44 +38,51 @@ async def startup_recovery():
     # Reset running evaluations
     running_evals = await get_running_evaluations()
     for eval_row in running_evals:
-        evaluation_id = eval_row["evaluation_id"]
-        evaluation = await get_evaluation_by_evaluation_id(evaluation_id)
-        if evaluation:
-            if evaluation.is_screening:
-                await evaluation.error("Disconnected from screener (error code 2)")
-            else:
-                # await evaluation.reset_to_waiting()
-                # set evaluation to waiting, and its runs to cancelled
-                await reset_evaluation_to_waiting(evaluation_id)
-                # set agent status to waiting
-                agent_version_id = evaluation.version_id
-                await set_agent_status(
-                    version_id=agent_version_id,
-                    status=AgentStatus.waiting.value
-                )
+        evaluation_id = eval_row.evaluation_id
+        agent_version_id = eval_row.version_id
+        from api.src.models.screener import Screener
+        is_screening = Screener.get_stage(eval_row.validator_hotkey) is not None
+        if is_screening:
+            await update_evaluation_to_error(evaluation_id, "Disconnected from screener (error code 2)")
+            await atomically_update_agent_status(version_id=agent_version_id)
+        else:
+            # set evaluation to waiting, and its runs to cancelled
+            await reset_evaluation_to_waiting(evaluation_id)
+            # set agent status to waiting
+            await set_agent_status(
+                version_id=agent_version_id,
+                status=AgentStatus.waiting.value
+            )
 
     # Check for running evaluations that should be auto-completed
     stuck_evaluations = await get_stuck_evaluations()
 
     for stuck_eval in stuck_evaluations:
-        evaluation = await get_evaluation_by_evaluation_id(stuck_eval.evaluation_id)
-        if evaluation:
-            logger.info(f"Auto-completing stuck evaluation {evaluation.evaluation_id} during startup recovery")
-            # During startup recovery, don't trigger notifications
-            _ = await evaluation.finish()
+        evaluation_id = stuck_eval.evaluation_id
+        # evaluation = await get_evaluation_by_evaluation_id(evaluation_id)
+        validator_hotkey = stuck_eval.validator_hotkey
+
+        logger.info(f"Auto-completing stuck evaluation {evaluation_id} during startup recovery")
+        # During startup recovery, don't trigger notifications
+        _ = await finish_evaluation(evaluation_id, validator_hotkey, errored=True, reason="Platform restarted")
 
     # Cancel waiting screenings for all screener types
     waiting_screenings = await get_waiting_evaluations()
     for screening_row in waiting_screenings:
-        evaluation = await get_evaluation_by_evaluation_id(screening_row.evaluation_id)
-        if evaluation:
-            await evaluation.error("Disconnected from screener (error code 3)")
+        evaluation_id = screening_row.evaluation_id
+        evaluation_version_id = screening_row.version_id
+
+        # await evaluation.error("Disconnected from screener (error code 3)")
+        await update_evaluation_to_error(evaluation_id, "Disconnected from screener (error code 3)")
+        await atomically_update_agent_status(version_id=evaluation_version_id)
 
     # Cancel dangling evaluation runs
     await cancel_dangling_evaluation_runs()
 
     # Prune low-scoring evaluations that should not continue waiting
-    # await Evaluation.prune_low_waiting(conn)
+    top_agent = await get_top_agent()
+    if top_agent:
+        await prune_queue(top_agent)
 
     logger.info("Application startup recovery completed with multi-stage screening support")