Merge pull request #1200 from Kiln-AI/dchiang/KIL-478/retry-eval-fail

Add retry to AsyncJobRunner

Author: chiang-daniel

libs/core/kiln_ai/adapters/eval/eval_runner.py

@@ -3,6 +3,8 @@
 from dataclasses import dataclass
 from typing import AsyncGenerator, Dict, List, Literal, Set
 
+import litellm
+
 from kiln_ai.adapters.adapter_registry import load_skills_for_task
 from kiln_ai.adapters.eval.base_eval import BaseEval
 from kiln_ai.adapters.eval.registry import eval_adapter_from_type
@@ -12,7 +14,7 @@
 from kiln_ai.datamodel.eval import EvalConfig, EvalDataType, EvalRun, EvalScores
 from kiln_ai.datamodel.task import TaskRunConfig
 from kiln_ai.datamodel.task_run import TaskRun, Usage
-from kiln_ai.utils.async_job_runner import AsyncJobRunner, Progress
+from kiln_ai.utils.async_job_runner import AsyncJobRunner, Progress, RetryableError
 
 logger = logging.getLogger(__name__)
 
@@ -188,6 +190,7 @@ async def run(self, concurrency: int = 25) -> AsyncGenerator[Progress, None]:
             concurrency=concurrency,
             jobs=jobs,
             run_job_fn=self.run_job,
+            max_retries=2,
         )
         async for progress in runner.run():
             yield progress
@@ -261,8 +264,37 @@ async def run_job(self, job: EvalJob) -> bool:
 
             return True
         except Exception as e:
+            if _is_retryable_error(e):
+                logger.error(
+                    f"Transient error running eval job for dataset item {job.item.id}: {e}",
+                    exc_info=True,
+                )
+                raise RetryableError(str(e)) from e
             logger.error(
                 f"Error running eval job for dataset item {job.item.id}: {e}",
                 exc_info=True,
             )
-            return False
+            raise
+
+
+def _is_retryable_error(e: BaseException) -> bool:
+    if isinstance(
+        e,
+        (
+            litellm.RateLimitError,
+            litellm.APIConnectionError,
+            litellm.InternalServerError,
+            litellm.ServiceUnavailableError,
+            litellm.BadGatewayError,
+            litellm.JSONSchemaValidationError,
+        ),
+    ):
+        return True
+
+    # ValueError thrown by Kiln's adapter when structured output doesn't match schema
+    if isinstance(
+        e, ValueError
+    ) and "This task requires a specific output schema" in str(e):
+        return True
+
+    return False

libs/core/kiln_ai/adapters/eval/test_eval_runner.py

@@ -1,10 +1,11 @@
 from typing import Dict
 from unittest.mock import AsyncMock, patch
 
+import litellm
 import pytest
 
 from kiln_ai.adapters.eval.base_eval import BaseEval
-from kiln_ai.adapters.eval.eval_runner import EvalJob, EvalRunner
+from kiln_ai.adapters.eval.eval_runner import EvalJob, EvalRunner, _is_retryable_error
 from kiln_ai.adapters.ml_model_list import ModelProviderName
 from kiln_ai.datamodel import (
     DataSource,
@@ -608,9 +609,9 @@ async def test_run_job_invalid_evaluator(
         "kiln_ai.adapters.eval.eval_runner.eval_adapter_from_type",
         return_value=lambda *args, **kwargs: object(),
     ):
-        success = await mock_eval_runner.run_job(job)
+        with pytest.raises(ValueError):
+            await mock_eval_runner.run_job(job)
 
-    assert success is False
     assert len(mock_eval_config.runs()) == 0
 
 
@@ -640,9 +641,9 @@ async def run_task_and_eval(self, eval_job_item: TaskRun):
         "kiln_ai.adapters.eval.eval_runner.eval_adapter_from_type",
         return_value=lambda *args, **kwargs: ErrorEvaluator(*args, **kwargs),
     ):
-        success = await mock_eval_runner.run_job(job)
+        with pytest.raises(ValueError):
+            await mock_eval_runner.run_job(job)
 
-    assert success is False
     assert len(mock_eval_config.runs()) == 0
 
 
@@ -825,9 +826,40 @@ async def run_task_and_eval(self, eval_job_item: TaskRun):
         "kiln_ai.adapters.eval.eval_runner.eval_adapter_from_type",
         return_value=lambda *args, **kwargs: MockEvaluator(*args, **kwargs),
     ):
-        success = await mock_eval_runner.run_job(job)
+        with pytest.raises(ValueError):
+            await mock_eval_runner.run_job(job)
 
     # For full_trace evals, None trace should fail and not save a run
-    assert success is False
     eval_runs = mock_eval_config.runs()
     assert len(eval_runs) == 0
+
+
+@pytest.mark.parametrize(
+    "error",
+    [
+        litellm.RateLimitError("rate limited", "provider", "model", None),
+        litellm.APIConnectionError("connection failed", "provider", "model", None),
+        litellm.InternalServerError("server error", "provider", "model", None),
+        litellm.ServiceUnavailableError("unavailable", "provider", "model", None),
+        litellm.BadGatewayError("bad gateway", "provider", "model", None),
+        litellm.JSONSchemaValidationError("schema error", "provider", "model", None),
+        ValueError(
+            "This task requires a specific output schema. While the model produced JSON, that JSON didn't meet the schema."
+        ),
+    ],
+)
+def test_is_retryable_error_returns_true(error):
+    assert _is_retryable_error(error) is True
+
+
+@pytest.mark.parametrize(
+    "error",
+    [
+        ValueError("some other value error"),
+        RuntimeError("runtime error"),
+        KeyError("missing key"),
+        TypeError("type error"),
+    ],
+)
+def test_is_retryable_error_returns_false(error):
+    assert _is_retryable_error(error) is False

libs/core/kiln_ai/utils/async_job_runner.py

@@ -15,6 +15,12 @@ class Progress:
     errors: int
 
 
+class RetryableError(Exception):
+    """Raise from run_job_fn to signal a transient failure that should be retried."""
+
+    pass
+
+
 class AsyncJobRunnerObserver(Generic[T]):
     async def on_error(self, job: T, error: Exception):
         """
@@ -42,10 +48,18 @@ def __init__(
         run_job_fn: Callable[[T], Awaitable[bool]],
         concurrency: int = 1,
         observers: List[AsyncJobRunnerObserver[T]] | None = None,
+        max_retries: int = 0,
+        retry_delay: float = 1.0,  # in seconds
     ):
         if concurrency < 1:
             raise ValueError("concurrency must be ≥ 1")
+        if max_retries < 0:
+            raise ValueError("max_retries must be >= 0")
+        if retry_delay < 0:
+            raise ValueError("retry_delay must be >= 0")
         self.concurrency = concurrency
+        self.max_retries = max_retries
+        self.retry_delay = retry_delay
         self.jobs = jobs
         self.run_job_fn = run_job_fn
         self.observers = observers or []
@@ -132,15 +146,32 @@ async def _run_worker(
                 # worker can end when the queue is empty
                 break
 
-            try:
-                await self.notify_job_start(job)
-                result = await run_job_fn(job)
-                if result:
-                    await self.notify_success(job)
-            except Exception as e:
-                logger.error("Job failed to complete", exc_info=True)
-                await self.notify_error(job, e)
-                result = False
+            await self.notify_job_start(job)
+            result = False
+            last_error: Exception | None = None
+            for attempt in range(1 + self.max_retries):
+                is_last_attempt = attempt == self.max_retries
+                try:
+                    result = await run_job_fn(job)
+                    last_error = None
+                    break
+                except RetryableError as e:
+                    result = False
+                    last_error = e
+                    if is_last_attempt:
+                        logger.error("Job failed to complete", exc_info=e)
+                        break
+                    await asyncio.sleep(self.retry_delay)
+                except Exception as e:
+                    result = False
+                    last_error = e
+                    logger.error("Job failed to complete", exc_info=e)
+                    break
+
+            if result:
+                await self.notify_success(job)
+            elif last_error is not None:
+                await self.notify_error(job, last_error)
 
             try:
                 await status_queue.put(result)