Use fresh worker caches for async batch predictions

great_ai/deploy/great_ai.py

@@ -205,7 +205,11 @@ class GreatAI(Generic[T, V]):
                 evaluations run part of the CI.
         """
 
-        wrapped_function = self._wrapped_func
+        # Futures in a warmed async cache cannot be serialized into worker
+        # processes. Build a fresh cache around the original decorated function.
+        wrapped_function = freeze_arguments(
+            self._get_cached_traced_function(self._wrapped_func.__wrapped__)
+        )
 
         def inner(value: Any) -> T:
             return (

tests/test_prediction_regressions.py

@@ -58,6 +58,20 @@ def test_http_preserves_nested_models(isolated_context: None, is_async: bool) ->
         assert response.json()["output"] == 5
 
 
+@pytest.mark.asyncio
+async def test_batch_after_async_prediction(isolated_context: None) -> None:
+    @GreatAI.create
+    async def predict(value: int) -> int:
+        return value + 1
+
+    assert (await predict(1)).output == 2
+    assert [trace.output for trace in predict.process_batch([1, 2], concurrency=1)] == [
+        2,
+        3,
+    ]
+    assert (await predict(1)).output == 2
+
+
 def test_flattening_preserves_reserved_trace_fields() -> None:
     trace = Trace(
         trace_id="actual-id",