Retry failed and cancelled async predictions instead of caching errors

great_ai/deploy/great_ai.py

@@ -265,7 +265,7 @@ class GreatAI(Generic[T, V]):
                 result = func(*args, **kwargs)
                 return cast(T, t.finalise(output=result))
 
-        @alru_cache(maxsize=get_context().prediction_cache_size)
+        @alru_cache(maxsize=get_context().prediction_cache_size, cache_exceptions=False)
         async def func_in_tracing_context_async(
             *args: Any,
             do_not_persist_traces: bool = False,

great_ai/external/async_lru.py

@@ -148,9 +148,9 @@ def alru_cache(
                     _cache_hit(wrapped, key)
                     return await asyncio.shield(fut)
 
-                exc = fut._exception
+                exc = None if fut.cancelled() else fut.exception()
 
-                if exc is None or cache_exceptions:
+                if not fut.cancelled() and (exc is None or cache_exceptions):
                     _cache_hit(wrapped, key)
                     return fut.result()
 

tests/test_cache_regressions.py

@@ -1,5 +1,8 @@
+import asyncio
 from typing import Dict, Tuple
 
+import pytest
+
 from great_ai import GreatAI
 
 
@@ -15,3 +18,27 @@ def test_nested_tuple_inputs_can_be_cached(isolated_context: None) -> None:
     assert predict(({"x": 7},)).output == 7
     assert predict(({"x": 7},)).output == 7
     assert calls == 1
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("cancelled", [False, True])
+async def test_async_predictions_retry_after_failure(
+    isolated_context: None, cancelled: bool
+) -> None:
+    calls = 0
+
+    @GreatAI.create
+    async def predict(value: int) -> int:
+        nonlocal calls
+        calls += 1
+        if calls == 1:
+            if cancelled:
+                raise asyncio.CancelledError()
+            raise ValueError("temporary failure")
+        return value
+
+    with pytest.raises(asyncio.CancelledError if cancelled else ValueError):
+        await predict(7)
+    assert (await predict(7)).output == 7
+    assert (await predict(7)).output == 7
+    assert calls == 2