From b9317eb88656e29a11c5aa088b743da2e95f67ba Mon Sep 17 00:00:00 2001 From: paultancre-bt Date: Thu, 22 Oct 2026 11:56:01 -0700 Subject: [PATCH] fix(workflow-eval): retry failed per-case submissions on poll A WorkflowTask or WorkflowScorer submit that raised left its case claimed with no submission record, so the run stayed waiting forever. Record the failure like batch submissions do, so the next poll() resubmits it. This matches the JS SDK. Co-Authored-By: Claude Opus 5.5 --- py/src/braintrust/test_workflow_eval.py | 34 ++++++++++++++++++++++++- py/src/braintrust/workflow_eval.py | 32 +++++++++++++++-------- 2 files changed, 55 insertions(+), 11 deletions(-) diff --git a/py/src/braintrust/test_workflow_eval.py b/py/src/braintrust/test_workflow_eval.py index 1a03d8c13..1b613b3cf 100644 --- a/py/src/braintrust/test_workflow_eval.py +++ b/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 diff --git a/py/src/braintrust/workflow_eval.py b/py/src/braintrust/workflow_eval.py index 6619a91e9..686cac0cc 100644 --- a/py/src/braintrust/workflow_eval.py +++ b/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)