From 4a28714a08c19ef50e40116bff5220ca747d86a0 Mon Sep 17 00:00:00 2001
From: sabaimran <narmiabas@gmail.com>
Date: Tue, 28 Jan 2025 18:46:04 -0800
Subject: [PATCH] Add retry logic to code execution and add health check to
 sandbox container

---
 docker-compose.yml                   |  5 +++++
 src/khoj/processor/tools/run_code.py | 26 +++++++++++++++++++++++++-
 2 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/docker-compose.yml b/docker-compose.yml
index 22371182..ea0b603b 100644
--- a/docker-compose.yml
+++ b/docker-compose.yml
@@ -14,6 +14,11 @@ services:
       retries: 5
   sandbox:
     image: ghcr.io/khoj-ai/terrarium:latest
+    healthcheck:
+      test: ["CMD-SHELL", "curl -f http://localhost:8080/health"]
+      interval: 30s
+      timeout: 10s
+      retries: 2
   search:
     image: docker.io/searxng/searxng:latest
     volumes:
diff --git a/src/khoj/processor/tools/run_code.py b/src/khoj/processor/tools/run_code.py
index 6c1eb125..5a38e0f4 100644
--- a/src/khoj/processor/tools/run_code.py
+++ b/src/khoj/processor/tools/run_code.py
@@ -1,8 +1,10 @@
+import asyncio
 import base64
 import datetime
 import logging
 import mimetypes
 import os
+from functools import wraps
 from pathlib import Path
 from typing import Any, Callable, List, NamedTuple, Optional
 
@@ -144,6 +146,28 @@ async def generate_python_code(
     return GeneratedCode(code, input_files, input_links)
 
 
+def async_retry_with_backoff(retries=3, backoff_in_seconds=1):
+    def decorator(func):
+        @wraps(func)
+        async def wrapper(*args, **kwargs):
+            retry_count = 0
+            while retry_count < retries:
+                try:
+                    return await func(*args, **kwargs)
+                except (aiohttp.ClientError, asyncio.TimeoutError) as e:
+                    retry_count += 1
+                    if retry_count == retries:
+                        raise e
+                    wait_time = backoff_in_seconds * (2 ** (retry_count - 1))  # exponential backoff
+                    await asyncio.sleep(wait_time)
+            return await func(*args, **kwargs)
+
+        return wrapper
+
+    return decorator
+
+
+@async_retry_with_backoff(retries=3, backoff_in_seconds=1)
 async def execute_sandboxed_python(code: str, input_data: list[dict], sandbox_url: str = SANDBOX_URL) -> dict[str, Any]:
     """
     Takes code to run as a string and calls the terrarium API to execute it.
@@ -157,7 +181,7 @@ async def execute_sandboxed_python(code: str, input_data: list[dict], sandbox_ur
     data = {"code": cleaned_code, "files": input_data}
 
     async with aiohttp.ClientSession() as session:
-        async with session.post(sandbox_url, json=data, headers=headers) as response:
+        async with session.post(sandbox_url, json=data, headers=headers, timeout=30) as response:
             if response.status == 200:
                 result: dict[str, Any] = await response.json()
                 result["code"] = cleaned_code