fix(workflow-eval): retry failed per-case submissions on poll

py/src/braintrust/test_workflow_eval.py

@@ -361,6 +361,38 @@ async def poll(_submission, _context):
     assert (await workflow.poll(started.run_id)).status == "completed"
 
 
+@pytest.mark.asyncio
+async def test_failed_submit_is_retried_on_poll():
+    submit_calls = []
+
+    async def submit(item, context):
+        submit_calls.append(context.run_id)
+        if len(submit_calls) == 1:
+            raise RuntimeError("transient submit failure")
+        return {"input": item.input}
+
+    workflow = define_workflow_eval(
+        "project",
+        store=WorkflowEvalMemoryStore(),
+        data=[{"id": "case", "input": 1}],
+        task=WorkflowTask(
+            submit=submit,
+            completion=WorkflowSubmissionCompletionPoll(
+                lambda _submission, _context: WorkflowSubmissionPoll("complete")
+            ),
+            collect=lambda submission, _context: WorkflowTaskResult(output=submission["input"]),
+        ),
+    )
+    with pytest.raises(RuntimeError, match="transient submit failure"):
+        await workflow.start(no_send_logs=True)
+    run_id = submit_calls[0]
+    assert (await workflow.status(run_id)).pending.poll == 1
+
+    assert (await workflow.poll(run_id)).status == "waiting"
+    assert (await workflow.poll(run_id)).status == "completed"
+    assert len(submit_calls) == 2
+
+
 @pytest.mark.asyncio
 async def test_webhook_result_can_be_matched_by_external_id():
     submitted = {}
@@ -1324,7 +1356,7 @@ async def submit(item, context):
     with pytest.raises(RuntimeError, match="submission failed"):
         await workflow.start(no_send_logs=True)
     assert calls == [1, 2]
-    assert (await workflow.status(run_ids[0])).pending.poll == 2
+    assert (await workflow.status(run_ids[0])).pending.poll == 3
 
 
 @pytest.mark.asyncio

py/src/braintrust/workflow_eval.py

@@ -1014,8 +1014,6 @@ async def poll_submission(submission: dict[str, Any], processor: Any) -> None:
                 await self._persist_submission_progress(run_id, submission)
                 return
             if submission["status"] == "failed":
-                if not submission.get("batch"):
-                    return
                 lease_key = self._batch_key(run_id, "retry-lease", submission["id"])
                 lease_value = uuid.uuid4().hex.encode()
                 if not await self.store.acquire_lease(lease_key, lease_value, _BATCH_LEASE_TTL_MS):
@@ -1026,14 +1024,16 @@ async def poll_submission(submission: dict[str, Any], processor: Any) -> None:
                         return
                     submission = current
                     context = WorkflowSubmissionContext(run_id=run_id, submission_id=submission["id"])
+                    item_ids = submission["item_ids"] if submission.get("batch") else [submission["item_id"]]
                     if submission["kind"] == "task":
-                        values = [
-                            await self._task_item(run, item_id, self._parameters(run))
-                            for item_id in submission["item_ids"]
-                        ]
+                        values = [await self._task_item(run, item_id, self._parameters(run)) for item_id in item_ids]
                     else:
-                        values = [await self._scorer_item(run, item_id) for item_id in submission["item_ids"]]
-                    items = [WorkflowBatchItem(f"item-{index}", value) for index, value in enumerate(values)]
+                        values = [await self._scorer_item(run, item_id) for item_id in item_ids]
+                    items = (
+                        [WorkflowBatchItem(f"item-{index}", value) for index, value in enumerate(values)]
+                        if submission.get("batch")
+                        else values[0]
+                    )
                     completion = processor.completion
                     try:
                         submission_data = _json_value(await self._call(processor.submit, items, context))
@@ -1241,9 +1241,21 @@ async def _submit_submission(self, run: Mapping[str, Any], spec: dict[str, Any],
         else:
             item = await self._scorer_item(run, spec["item_id"])
         context = WorkflowSubmissionContext(run_id=run_id, submission_id=spec["id"])
-        submission_data = await self._call(processor.submit, item, context)
-        submission_data = _json_value(submission_data)
         completion = processor.completion
+        try:
+            submission_data = _json_value(await self._call(processor.submit, item, context))
+        except Exception:
+            failed = {
+                **spec,
+                "submission_data": None,
+                "external_id": None,
+                "mode": completion.mode,
+                "status": "failed",
+                "attempts": 1,
+            }
+            await self._write_submission(run_id, failed)
+            await self._persist_submission_progress(run_id, failed)
+            raise
         external_id = None
         if isinstance(completion, WorkflowSubmissionCompletionWebhook):
             external_id = await self._call(completion.get_external_id, submission_data, context)