Don't let the agents fail. Retry on error

README.md

@@ -6,7 +6,7 @@ Todo:
 
 - [x] Use Postgresql as a database
 - [x] Expose a simple read-only website.
-- [ ] Agents should wait and retry on timeout
+- [x] Agents should wait and retry on timeout
 - [ ] Last seen agents
 - [ ] Use background tasks for alerting
 - [ ] Add a command to generate new authentication tokens

argos/agent.py

@@ -1,7 +1,9 @@
 import asyncio
+import logging
 from typing import List
 
 import httpx
+from tenacity import after_log, retry, wait_random
 
 from argos.checks import get_registered_check
 from argos.logging import logger
@@ -63,7 +65,22 @@ async def get_and_complete_tasks(http_client, server, max_tasks):
         return False
 
 
+def log_failure(retry_state):
+    if retry_state.attempt_number < 1:
+        loglevel = logging.INFO
+    else:
+        loglevel = logging.WARNING
+    logger.log(
+        loglevel,
+        "Retrying: attempt %s ended with: %s",
+        retry_state.attempt_number,
+        retry_state.outcome,
+    )
+
+
+@retry(after=log_failure, wait=wait_random(min=1, max=2))
 async def run_agent(server: str, auth: str, max_tasks: int, wait_time: int):
+    logger.info(f"Running agent against {server}")
     headers = {"Authorization": f"Bearer {auth}"}
     async with httpx.AsyncClient(headers=headers) as http_client:
         while True:

argos/commands.py

@@ -35,7 +35,10 @@ def agent(server, auth, max_tasks, wait_time, log_level):
 
     Usage: argos agent https://argos.server "auth-token-here"
     """
-    logging.set_log_level(log_level)
+    click.echo("Starting argos agent. Will retry forever.")
+    from argos.logging import logger
+
+    logger.setLevel(log_level)
     asyncio.run(run_agent(server, auth, max_tasks, wait_time))
 
 

pyproject.toml

@@ -28,6 +28,7 @@ dependencies = [
   "uvicorn>=0.23,<1",
   "Jinja2>=3.0,<4",
   "pydantic-settings>=2.0,<3",
+  "tenacity>=8.2,<9",
 ]
 
 [project.urls]