diff --git a/py/src/braintrust/test_workflow_eval.py b/py/src/braintrust/test_workflow_eval.py index 1a03d8c1..1b613b3c 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 6619a91e..686cac0c 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)