From ce536fce615fd55cec9042235e8d8754af7f3bcd Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Mon, 5 Oct 2026 18:17:31 -0400 Subject: [PATCH 1/4] feat(workflow-eval): add batch tasks and scorers --- py/src/braintrust/__init__.py | 1 + py/src/braintrust/test_workflow_eval.py | 106 +++- .../type_tests/test_workflow_eval.py | 55 +- py/src/braintrust/workflow_eval.py | 484 +++++++++++++++++- 4 files changed, 618 insertions(+), 28 deletions(-) diff --git a/py/src/braintrust/__init__.py b/py/src/braintrust/__init__.py index 93a885fb6..850d8e4bb 100644 --- a/py/src/braintrust/__init__.py +++ b/py/src/braintrust/__init__.py @@ -85,6 +85,7 @@ def is_equal(expected, output): from .span_customizer import set_span_customizers as set_span_customizers from .util import BT_IS_ASYNC_ATTRIBUTE as BT_IS_ASYNC_ATTRIBUTE from .util import MarkAsyncWrapper as MarkAsyncWrapper +from .workflow_eval import * def wrap_anthropic(client: _Any) -> _Any: diff --git a/py/src/braintrust/test_workflow_eval.py b/py/src/braintrust/test_workflow_eval.py index c7f6fd8b5..5b126629b 100644 --- a/py/src/braintrust/test_workflow_eval.py +++ b/py/src/braintrust/test_workflow_eval.py @@ -10,6 +10,10 @@ from .test_helpers import init_test_exp, with_memory_logger, with_simulate_login # noqa: F401 from .util import LazyValue from .workflow_eval import ( + WorkflowBatchingOptions, + WorkflowBatchItemResult, + WorkflowBatchScorer, + WorkflowBatchTask, WorkflowEvalCompletedResult, WorkflowEvalMemoryStore, WorkflowEvalRedisStore, @@ -289,7 +293,7 @@ async def collect(_submission, context): started = await workflow_eval.start(no_send_logs=True) submission_id = next(iter(submitted)) - completed = await workflow_eval.process_submission_result(started.run_id, external_id=f"external-{submission_id}") + completed = await workflow_eval.process_submission_result(external_id=f"external-{submission_id}") assert isinstance(completed, WorkflowEvalCompletedResult) assert ( @@ -297,6 +301,106 @@ async def collect(_submission, context): ).status == "completed" +@pytest.mark.asyncio +async def test_batch_tasks_and_scorers_submit_and_collect_multiple_items(): + task_batch_sizes = [] + score_batch_sizes = [] + + async def submit_task(items, _context): + task_batch_sizes.append(len(items)) + return {"inputs": [item.input for item in items]} + + async def collect_task(submission, _context): + return [ + WorkflowBatchItemResult(custom_id=f"item-{index}", result=WorkflowTaskResult(output=value * 2)) + for index, value in enumerate(submission["inputs"]) + ] + + async def submit_score(items, _context): + score_batch_sizes.append(len(items)) + return {"values": [(item.output, item.expected) for item in items]} + + async def collect_score(submission, _context): + for index, (output, expected) in enumerate(submission["values"]): + yield WorkflowBatchItemResult( + custom_id=f"item-{index}", result=WorkflowScorerResult(score=output == expected) + ) + + completion = WorkflowSubmissionCompletionPoll(lambda _submission, _context: WorkflowSubmissionPoll("complete")) + workflow = define_workflow_eval( + "project", + store=WorkflowEvalMemoryStore(), + data=[{"id": str(value), "input": value, "expected": value * 2} for value in range(3)], + task=WorkflowBatchTask( + batching=WorkflowBatchingOptions(max_size=2), + submit=submit_task, + completion=completion, + collect=collect_task, + ), + scores=[ + WorkflowBatchScorer( + name="batch-score", + batching=WorkflowBatchingOptions(max_size=2), + submit=submit_score, + completion=completion, + collect=collect_score, + ) + ], + ) + + started = await workflow.start(no_send_logs=True) + after_tasks = await workflow.poll(started.run_id) + completed = await workflow.poll(started.run_id) + + assert isinstance(after_tasks, WorkflowEvalWaitingResult) + assert after_tasks.pending.poll == 2 + assert isinstance(completed, WorkflowEvalCompletedResult) + assert completed.summary.scores["batch-score"].score == 1 + assert task_batch_sizes == [2, 1] + assert score_batch_sizes == [2, 1] + + +@pytest.mark.asyncio +async def test_batch_item_failures_skip_scoring(): + score_inputs = [] + + async def submit(items, _context): + return {"inputs": [item.input for item in items]} + + async def collect(submission, _context): + return [ + WorkflowBatchItemResult(custom_id="item-0", result=WorkflowTaskResult(output="good")), + WorkflowBatchItemResult(custom_id="item-1", error=ValueError("provider rejected item")), + # item-2 is intentionally missing and should also fail. + ] + + def score(output): + score_inputs.append(output) + return 1 + + workflow = define_workflow_eval( + "project", + store=WorkflowEvalMemoryStore(), + data=[{"id": str(value), "input": value} for value in range(3)], + task=WorkflowBatchTask( + batching=WorkflowBatchingOptions(max_size=3), + submit=submit, + completion=WorkflowSubmissionCompletionPoll( + lambda _submission, _context: WorkflowSubmissionPoll("complete") + ), + collect=collect, + ), + scores=[score], + ) + + started = await workflow.start(no_send_logs=True) + completed = await workflow.poll(started.run_id) + + assert isinstance(started, WorkflowEvalWaitingResult) + assert isinstance(completed, WorkflowEvalCompletedResult) + assert score_inputs == ["good"] + + @pytest.mark.asyncio async def test_collect_rejects_array_results(): submission_ids = [] diff --git a/py/src/braintrust/type_tests/test_workflow_eval.py b/py/src/braintrust/type_tests/test_workflow_eval.py index 484ddd220..589b4ce4f 100644 --- a/py/src/braintrust/type_tests/test_workflow_eval.py +++ b/py/src/braintrust/type_tests/test_workflow_eval.py @@ -1,10 +1,15 @@ """Static and runtime type coverage for the experimental workflow eval API.""" -from typing import TypedDict +from typing import Sequence, TypedDict import pytest from braintrust import EvalCase from braintrust.workflow_eval import ( + WorkflowBatchingOptions, + WorkflowBatchItem, + WorkflowBatchItemResult, + WorkflowBatchScorer, + WorkflowBatchTask, WorkflowEval, WorkflowEvalMemoryStore, WorkflowScorer, @@ -68,6 +73,54 @@ async def poll_submission(submission: Submission, context: WorkflowSubmissionCon ) +async def submit_batch_task( + items: Sequence[WorkflowBatchItem[WorkflowTaskItem[str, str]]], context: WorkflowSubmissionContext +) -> Submission: + assert items[0].item.input + return {"id": context.submission_id} + + +async def collect_batch_task( + submission: Submission, _context: WorkflowSubmissionContext +) -> list[WorkflowBatchItemResult[WorkflowTaskResult[int]]]: + return [WorkflowBatchItemResult(custom_id="item-0", result=WorkflowTaskResult(output=len(submission["id"])))] + + +async def submit_batch_score( + items: Sequence[WorkflowBatchItem[WorkflowScorerItem[str, int, str]]], context: WorkflowSubmissionContext +) -> Submission: + assert items[0].item.output > 0 + return {"id": context.submission_id} + + +async def collect_batch_score( + _submission: Submission, _context: WorkflowSubmissionContext +) -> list[WorkflowBatchItemResult[WorkflowScorerResult]]: + return [WorkflowBatchItemResult(custom_id="item-0", result=WorkflowScorerResult(score=1))] + + +batch_task: WorkflowBatchTask[str, int, str, Submission] = WorkflowBatchTask( + batching=WorkflowBatchingOptions(max_size=10), + submit=submit_batch_task, + completion=WorkflowSubmissionCompletionPoll(poll_submission), + collect=collect_batch_task, +) +batch_score: WorkflowBatchScorer[str, int, str, Submission] = WorkflowBatchScorer( + name="batch-score", + batching=WorkflowBatchingOptions(max_size=10, max_wait_ms=100), + submit=submit_batch_score, + completion=WorkflowSubmissionCompletionPoll(poll_submission), + collect=collect_batch_score, +) +batch_workflow_eval: WorkflowEval[str, int, str] = define_workflow_eval( + "project", + store=WorkflowEvalMemoryStore(), + data=[EvalCase(id="case", input="input")], + task=batch_task, + scores=[batch_score], +) + + @pytest.mark.asyncio async def test_workflow_eval_types_at_runtime(): result = await workflow_eval.start(no_send_logs=True) diff --git a/py/src/braintrust/workflow_eval.py b/py/src/braintrust/workflow_eval.py index 01d5ae55a..32ae9f185 100644 --- a/py/src/braintrust/workflow_eval.py +++ b/py/src/braintrust/workflow_eval.py @@ -1,4 +1,4 @@ -"""Experimental workflow evaluations with one asynchronous provider submission per case/trial.""" +"""Experimental workflow evaluations with asynchronous provider submissions.""" import asyncio import base64 @@ -133,6 +133,45 @@ class WorkflowTaskResult(Generic[Output]): tags: list[str] | None = None +@dataclasses.dataclass(frozen=True) +class WorkflowBatchingOptions: + """Controls how workflow items are grouped into provider batches.""" + + max_size: int + max_wait_ms: int | float | None = None + + def __post_init__(self) -> None: + if not isinstance(self.max_size, int) or isinstance(self.max_size, bool) or self.max_size < 1: + raise ValueError("batching.max_size must be a positive integer") + if self.max_wait_ms is not None and ( + isinstance(self.max_wait_ms, bool) + or not isinstance(self.max_wait_ms, (int, float)) + or not self.max_wait_ms >= 0 + ): + raise ValueError("batching.max_wait_ms must be a non-negative number") + + +@dataclasses.dataclass(frozen=True) +class WorkflowBatchItem(Generic[Input]): + """One provider-safe item in a batch. Return its custom_id from collect.""" + + custom_id: str + item: Input + + def __getattr__(self, name: str) -> Any: + # Expose the task/scorer item fields directly alongside custom_id. + return getattr(self.item, name) + + +@dataclasses.dataclass(frozen=True) +class WorkflowBatchItemResult(Generic[Output]): + """A result or failure for one item returned by a batch collector.""" + + custom_id: str + result: Output | None = None + error: Any = None + + @dataclasses.dataclass(frozen=True) class WorkflowScorerItem(Generic[Input, Output, Expected]): """One completed case/trial passed to a workflow scorer submission callback.""" @@ -153,6 +192,53 @@ class WorkflowScorerResult: score: OneOrMoreScores +@dataclasses.dataclass(frozen=True) +class WorkflowBatchTask(Generic[Input, Output, Expected, SubmissionData]): + """Submits multiple cases/trials in one asynchronous provider operation.""" + + batching: WorkflowBatchingOptions + submit: Callable[ + [Sequence[WorkflowBatchItem[WorkflowTaskItem[Input, Expected]]], WorkflowSubmissionContext], + SubmissionData | Awaitable[SubmissionData], + ] + completion: WorkflowSubmissionCompletion[SubmissionData] + collect: Callable[ + [SubmissionData, WorkflowSubmissionContext], + Iterable[WorkflowBatchItemResult[WorkflowTaskResult[Output]]] + | AsyncIterable[WorkflowBatchItemResult[WorkflowTaskResult[Output]]] + | Awaitable[ + Iterable[WorkflowBatchItemResult[WorkflowTaskResult[Output]]] + | AsyncIterable[WorkflowBatchItemResult[WorkflowTaskResult[Output]]] + ], + ] + + +@dataclasses.dataclass(frozen=True) +class WorkflowBatchScorer(Generic[Input, Output, Expected, SubmissionData]): + """Submits multiple completed cases/trials in one asynchronous scorer operation.""" + + name: str + batching: WorkflowBatchingOptions + submit: Callable[ + [Sequence[WorkflowBatchItem[WorkflowScorerItem[Input, Output, Expected]]], WorkflowSubmissionContext], + SubmissionData | Awaitable[SubmissionData], + ] + completion: WorkflowSubmissionCompletion[SubmissionData] + collect: Callable[ + [SubmissionData, WorkflowSubmissionContext], + Iterable[WorkflowBatchItemResult[WorkflowScorerResult]] + | AsyncIterable[WorkflowBatchItemResult[WorkflowScorerResult]] + | Awaitable[ + Iterable[WorkflowBatchItemResult[WorkflowScorerResult]] + | AsyncIterable[WorkflowBatchItemResult[WorkflowScorerResult]] + ], + ] + + def __post_init__(self) -> None: + if not self.name: + raise ValueError("WorkflowBatchScorer name must be a non-empty string") + + @dataclasses.dataclass(frozen=True) class WorkflowTask(Generic[Input, Output, Expected, SubmissionData]): """Submits one asynchronous provider operation per case/trial and collects its task result.""" @@ -338,8 +424,16 @@ class _WorkflowEvalConfig(Generic[Input, Output, Expected]): project_name: str store: WorkflowEvalStore data: EvalData[Input, Expected] - task: EvalTask[Input, Output, Expected] | WorkflowTask[Input, Output, Expected, Any] - scores: Sequence[EvalScorer[Input, Output, Expected] | WorkflowScorer[Input, Output, Expected, Any]] + task: ( + EvalTask[Input, Output, Expected] + | WorkflowTask[Input, Output, Expected, Any] + | WorkflowBatchTask[Input, Output, Expected, Any] + ) + scores: Sequence[ + EvalScorer[Input, Output, Expected] + | WorkflowScorer[Input, Output, Expected, Any] + | WorkflowBatchScorer[Input, Output, Expected, Any] + ] classifiers: Sequence[EvalClassifier[Input, Output, Expected]] case_id: Callable[[EvalCase[Input, Expected]], str | Awaitable[str]] | None experiment_name: str | None @@ -377,7 +471,7 @@ async def poll(self, run_id: str) -> WorkflowEvalResult: return await _WorkflowEvalRunner(self._config).poll(run_id) async def process_submission_result( - self, run_id: str, *, submission_id: str | None = None, external_id: str | None = None + self, run_id: str | None = None, *, submission_id: str | None = None, external_id: str | None = None ) -> WorkflowEvalResult: return await _WorkflowEvalRunner(self._config).process_submission_result( run_id, submission_id=submission_id, external_id=external_id @@ -389,8 +483,15 @@ def define_workflow_eval( *, store: WorkflowEvalStore, data: EvalData[Input, Expected], - task: EvalTask[Input, Output, Expected] | WorkflowTask[Input, Output, Expected, Any], - scores: Sequence[EvalScorer[Input, Output, Expected] | WorkflowScorer[Input, Output, Expected, Any]] | None = None, + task: EvalTask[Input, Output, Expected] + | WorkflowTask[Input, Output, Expected, Any] + | WorkflowBatchTask[Input, Output, Expected, Any], + scores: Sequence[ + EvalScorer[Input, Output, Expected] + | WorkflowScorer[Input, Output, Expected, Any] + | WorkflowBatchScorer[Input, Output, Expected, Any] + ] + | None = None, classifiers: Sequence[EvalClassifier[Input, Output, Expected]] | None = None, case_id: Callable[[EvalCase[Input, Expected]], str | Awaitable[str]] | None = None, experiment_name: str | None = None, @@ -411,7 +512,8 @@ def define_workflow_eval( ) -> WorkflowEval[Input, Output, Expected]: """Define an experimental evaluation that can pause across provider submissions. - WorkflowTask and WorkflowScorer submit one case/trial and collect one result. + WorkflowTask and WorkflowScorer submit one case/trial; WorkflowBatchTask and + WorkflowBatchScorer submit multiple cases/trials in one provider operation. Submission data and collected values must be JSON serializable. Scorers and classifiers start once their case's task is persisted and logged. @@ -515,6 +617,9 @@ def _key(self, run_id: str, kind: str, identifier: str | None = None) -> str: suffix = f"/{identifier}" if identifier is not None else "" return f"{_SCHEMA_PREFIX}/{self.definition_key}/{run_id}/{kind}{suffix}" + def _external_key(self, external_id: str) -> str: + return f"{_SCHEMA_PREFIX}/{self.definition_key}/external/{_stable_hex(external_id, length=32)}" + async def _read(self, run_id: str, kind: str, identifier: str | None = None) -> Any | None: value = await self.store.read(self._key(run_id, kind, identifier)) return _decode(value) if value is not None else None @@ -529,6 +634,17 @@ async def _required(self, run_id: str, kind: str, identifier: str | None = None) async def _write(self, run_id: str, kind: str, value: Any, identifier: str | None = None) -> None: await self.store.write(self._key(run_id, kind, identifier), _json_bytes(value)) + async def _record_external_locator(self, run_id: str, submission_id: str, external_id: str | None) -> None: + if external_id is not None: + key = self._external_key(external_id) + locator = {"run_id": run_id, "submission_id": submission_id} + entry = await self.store.get_or_set(key, _json_bytes(locator)) + if not entry.created: + existing = _decode(entry.value) + if existing is not None and existing != locator: + # Ambiguous provider IDs require the caller to supply run_id. + await self.store.write(key, _json_bytes(None)) + async def _claim(self, run_id: str, action: str) -> bool: result = await self.store.get_or_set(self._key(run_id, "claim", action), b"1") return result.created @@ -610,8 +726,8 @@ async def _resolve_data(self, experiment: Experiment | None) -> list[EvalCase[In return [value if isinstance(value, EvalCase) else EvalCase.from_dict(value) for value in values] def _uses_submission_processor(self) -> bool: - return isinstance(self.config.task, WorkflowTask) or any( - isinstance(score, WorkflowScorer) for score in self.config.scores + return isinstance(self.config.task, (WorkflowTask, WorkflowBatchTask)) or any( + isinstance(score, (WorkflowScorer, WorkflowBatchScorer)) for score in self.config.scores ) async def _save_completed(self, run: Mapping[str, Any], summary: ExperimentSummary) -> WorkflowEvalCompletedResult: @@ -739,6 +855,31 @@ async def poll(self, run_id: str) -> WorkflowEvalResult: submissions = await self._submission_records(run) async def poll_submission(submission: dict[str, Any], processor: Any) -> None: + if submission["status"] == "failed": + if not submission.get("batch"): + return + context = WorkflowSubmissionContext(run_id=run_id, submission_id=submission["id"]) + if submission["kind"] == "task": + values = [ + await self._task_item(run, item_id, self._parameters(run)) + for item_id in submission["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)] + submission["submission_data"] = _json_value(await self._call(processor.submit, items, context)) + completion = processor.completion + if isinstance(completion, WorkflowSubmissionCompletionWebhook): + external_id = await self._call(completion.get_external_id, submission["submission_data"], context) + if not isinstance(external_id, str) or not external_id.strip(): + raise ValueError("Submission webhook get_external_id must return a non-empty string") + submission["external_id"] = external_id + await self._record_external_locator(run_id, submission["id"], external_id) + submission["mode"] = completion.mode + submission["status"] = "submitted" + await self._write(run_id, "submission", submission, submission["id"]) + await self._persist_submission_progress(run_id, submission) + return completion = processor.completion assert isinstance(completion, WorkflowSubmissionCompletionPoll) context = WorkflowSubmissionContext(run_id=run_id, submission_id=submission["id"]) @@ -760,7 +901,8 @@ async def poll_submission(submission: dict[str, Any], processor: Any) -> None: *( poll_submission(submission, processor) for submission, processor in submissions - if submission["status"] == "submitted" and submission["mode"] == "poll" + if submission["status"] == "failed" + or (submission["status"] == "submitted" and submission["mode"] == "poll") ), return_exceptions=True, ) @@ -775,10 +917,19 @@ async def poll_submission(submission: dict[str, Any], processor: Any) -> None: return status async def process_submission_result( - self, run_id: str, *, submission_id: str | None, external_id: str | None + self, run_id: str | None, *, submission_id: str | None, external_id: str | None ) -> WorkflowEvalResult: if not submission_id and not external_id: raise ValueError("process_submission_result requires submission_id or external_id") + if run_id is None: + if external_id is None: + raise ValueError("run_id is required when submission_id is used") + locator = await self.store.read(self._external_key(external_id)) + resolved = _decode(locator) if locator is not None else None + if not isinstance(resolved, Mapping): + raise ValueError("No submission matches this external_id") + run_id = resolved["run_id"] + submission_id = submission_id or resolved["submission_id"] run = await self._required(run_id, "run") external_submission_id = await self._read(run_id, "external", external_id) if external_id is not None else None if submission_id and external_submission_id and submission_id != external_submission_id: @@ -797,12 +948,12 @@ async def process_submission_result( else next( score for score in self.config.scores - if isinstance(score, WorkflowScorer) and score.name == submission["scorer_name"] + if isinstance(score, (WorkflowScorer, WorkflowBatchScorer)) and score.name == submission["scorer_name"] ) ) await self._collect_submission(run, submission, processor) # Only the completed case needs advancing; other cases can still be pending. - return await self._advance(run, item_ids=[submission["item_id"]]) + return await self._advance(run, item_ids=submission.get("item_ids", [submission.get("item_id")])) def _submission_specs(self, run: Mapping[str, Any]) -> list[tuple[dict[str, Any], Any]]: specs: list[tuple[dict[str, Any], Any]] = [] @@ -835,11 +986,29 @@ def _submission_specs(self, run: Mapping[str, Any]) -> list[tuple[dict[str, Any] async def _submission_records(self, run: Mapping[str, Any]) -> list[tuple[dict[str, Any], Any]]: records: list[tuple[dict[str, Any], Any]] = [] + seen: set[str] = set() for spec, processor in self._submission_specs(run): record = await self._read(run["run_id"], "submission", spec["id"]) if record is not None: await self._persist_submission_progress(run["run_id"], record) records.append((record, processor)) + seen.add(record["id"]) + batch_processors: dict[str, Any] = {} + if isinstance(self.config.task, WorkflowBatchTask): + batch_processors["task"] = self.config.task + for scorer in self.config.scores: + if isinstance(scorer, WorkflowBatchScorer): + batch_processors[f"score:{scorer.name}"] = scorer + for item_id in run["case_ids"]: + for stage, processor in batch_processors.items(): + submission_id = await self._read(run["run_id"], "batch-link", f"{stage}/{item_id}") + if not submission_id or submission_id in seen: + continue + record = await self._read(run["run_id"], "submission", submission_id) + if record is not None: + await self._persist_submission_progress(run["run_id"], record) + records.append((record, processor)) + seen.add(submission_id) return records async def _case(self, run_id: str, item_id: str) -> dict[str, Any]: @@ -904,8 +1073,111 @@ async def _submit_submission(self, run: Mapping[str, Any], spec: dict[str, Any], "status": "submitted", } await self._write(run_id, "submission", submission, spec["id"]) + await self._record_external_locator(run_id, spec["id"], external_id) await self._persist_submission_progress(run_id, submission) + async def _schedule_batch_stage( + self, + run: Mapping[str, Any], + processor: WorkflowBatchTask[Any, Any, Any, Any] | WorkflowBatchScorer[Any, Any, Any, Any], + *, + kind: Literal["task", "score"], + ) -> None: + run_id = run["run_id"] + scorer_name = processor.name if kind == "score" else None + stage = f"score:{scorer_name}" if scorer_name else "task" + candidates: list[str] = [] + for item_id in run["case_ids"]: + if await self._read(run_id, "batch-link", f"{stage}/{item_id}") is not None: + continue + if kind == "task": + if await self._read(run_id, "task-result", item_id) is None: + candidates.append(item_id) + else: + task_result = await self._read(run_id, "task-result", item_id) + score_result = await self._read(run_id, _stage_kind("score-result", cast(str, scorer_name)), item_id) + if task_result is not None and task_result.get("error") is None and score_result is None: + candidates.append(item_id) + if not candidates: + return + + max_size = processor.batching.max_size + submit_partial = kind == "task" + if kind == "score": + completed_tasks = sum( + [await self._read(run_id, "task-result", item_id) is not None for item_id in run["case_ids"]] + ) + if completed_tasks == len(run["case_ids"]): + submit_partial = True + elif processor.batching.max_wait_ms is not None and len(candidates) < max_size: + window_key = f"{stage}/{candidates[0]}" + first_wait = await self.store.get_or_set( + self._key(run_id, "batch-wait", window_key), str(asyncio.get_running_loop().time()).encode() + ) + started = float(first_wait.value.decode()) + submit_partial = asyncio.get_running_loop().time() - started >= processor.batching.max_wait_ms / 1000 + if len(candidates) < max_size and not submit_partial: + return + batches = [candidates[index : index + max_size] for index in range(0, len(candidates), max_size)] + if not submit_partial and batches and len(batches[-1]) < max_size: + batches.pop() + + async def submit_batch(item_ids: list[str]) -> None: + submission_id = f"submission-{_stable_hex(run_id, stage, *item_ids, length=32)}" + spec: dict[str, Any] = { + "id": submission_id, + "kind": kind, + "item_ids": item_ids, + "batch": True, + } + if scorer_name is not None: + spec["scorer_name"] = scorer_name + if not await self._claim(run_id, f"submit:{submission_id}"): + return + for item_id in item_ids: + await self._write(run_id, "batch-link", submission_id, f"{stage}/{item_id}") + context = WorkflowSubmissionContext(run_id=run_id, submission_id=submission_id) + if kind == "task": + 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 item_ids] + submitted_items = [WorkflowBatchItem(f"item-{index}", value) for index, value in enumerate(values)] + completion = processor.completion + try: + submission_data = _json_value(await self._call(processor.submit, submitted_items, context)) + except Exception: + submission = { + **spec, + "submission_data": None, + "external_id": None, + "mode": completion.mode, + "status": "failed", + "attempts": 1, + } + await self._write(run_id, "submission", submission, submission_id) + await self._persist_submission_progress(run_id, submission) + raise + external_id = None + if isinstance(completion, WorkflowSubmissionCompletionWebhook): + external_id = await self._call(completion.get_external_id, submission_data, context) + if not isinstance(external_id, str) or not external_id.strip(): + raise ValueError("Submission webhook get_external_id must return a non-empty string") + submission = { + **spec, + "submission_data": submission_data, + "external_id": external_id, + "mode": completion.mode, + "status": "submitted", + } + await self._write(run_id, "submission", submission, submission_id) + await self._record_external_locator(run_id, submission_id, external_id) + await self._persist_submission_progress(run_id, submission) + + results = await asyncio.gather(*(submit_batch(item_ids) for item_ids in batches), return_exceptions=True) + errors = [result for result in results if isinstance(result, BaseException)] + if errors: + _raise_callback_errors(errors) + async def _persist_submission_progress(self, run_id: str, submission: Mapping[str, Any]) -> None: if submission.get("external_id") is not None: await self._write(run_id, "external", submission["id"], submission["external_id"]) @@ -919,6 +1191,12 @@ async def _collect_submission(self, run: Mapping[str, Any], submission: dict[str run_id = run["run_id"] if submission["status"] != "complete": context = WorkflowSubmissionContext(run_id=run_id, submission_id=submission["id"]) + if submission.get("batch"): + await self._collect_batch_submission(run, submission, processor, context) + submission["status"] = "complete" + await self._write(run_id, "submission", submission, submission["id"]) + await self._persist_submission_progress(run_id, submission) + return result = await self._call(processor.collect, submission["submission_data"], context) result_type = WorkflowTaskResult if submission["kind"] == "task" else WorkflowScorerResult if isinstance(result, Mapping): @@ -949,6 +1227,87 @@ async def _collect_submission(self, run: Mapping[str, Any], submission: dict[str # Replay repairs an interrupted progress update without collecting again. await self._persist_submission_progress(run_id, submission) + async def _collect_batch_submission( + self, + run: Mapping[str, Any], + submission: dict[str, Any], + processor: WorkflowBatchTask[Any, Any, Any, Any] | WorkflowBatchScorer[Any, Any, Any, Any], + context: WorkflowSubmissionContext, + ) -> None: + run_id = run["run_id"] + async with self._callback_semaphore: + if inspect.isasyncgenfunction(processor.collect): + collected = processor.collect(submission["submission_data"], context) + else: + collected = await await_or_run( + asyncio.get_running_loop(), processor.collect, submission["submission_data"], context + ) + if inspect.isawaitable(collected): + collected = await collected + if isinstance(collected, AsyncIterable): + entries = [entry async for entry in collected] + else: + entries = list(collected) + item_ids = submission["item_ids"] + expected = {f"item-{index}": item_id for index, item_id in enumerate(item_ids)} + outcomes: dict[str, WorkflowBatchItemResult[Any]] = {} + for entry in entries: + if isinstance(entry, Mapping): + entry = WorkflowBatchItemResult(**entry) + if not isinstance(entry, WorkflowBatchItemResult): + raise TypeError("Batch collect callback entries must be WorkflowBatchItemResult values") + if entry.custom_id not in expected: + raise ValueError(f"Batch collect returned unknown custom_id {entry.custom_id!r}") + if entry.custom_id in outcomes: + raise ValueError(f"Batch collect returned custom_id {entry.custom_id!r} more than once") + outcomes[entry.custom_id] = entry + for custom_id, item_id in expected.items(): + entry = outcomes.get(custom_id) + if entry is None: + error = f"Batch submission {submission['id']} returned no result for {custom_id}" + result = None + elif entry.error is not None: + error = str(entry.error) + result = None + else: + error = None + result = entry.result + if submission["kind"] == "task": + if error is None: + if isinstance(result, Mapping): + result = WorkflowTaskResult(**result) + if not isinstance(result, WorkflowTaskResult): + raise TypeError("Batch task result must be a WorkflowTaskResult") + case = await self._case(run_id, item_id) + metadata = {**case["metadata"], **(result.metadata or {})} + tags = result.tags if result.tags is not None else case.get("tags") + await self._write( + run_id, + "task-result", + {"output": _json_value(result.output), "metadata": metadata, "tags": tags}, + item_id, + ) + else: + await self._write( + run_id, + "task-result", + {"output": None, "metadata": {}, "tags": None, "error": error}, + item_id, + ) + elif error is None: + if isinstance(result, Mapping): + result = WorkflowScorerResult(**result) + if not isinstance(result, WorkflowScorerResult): + raise TypeError("Batch scorer result must be a WorkflowScorerResult") + await self._write( + run_id, + _stage_kind("score-result", submission["scorer_name"]), + _json_value(result.score), + item_id, + ) + else: + await self._write(run_id, _stage_kind("score-error", submission["scorer_name"]), error, item_id) + async def _waiting_status(self, run: Mapping[str, Any]) -> WorkflowEvalWaitingResult: pending = {} for mode in ("poll", "webhook"): @@ -1113,8 +1472,14 @@ async def _log_task(self, run: Mapping[str, Any], case: dict[str, Any], experime SpanTypeAttribute.TASK, input=case["datum"]["input"], ) as span: - span.log(output=task_result["output"]) - root.log(output=task_result["output"], metadata=task_result["metadata"], tags=task_result.get("tags")) + if task_result.get("error") is not None: + span.log(error=task_result["error"]) + else: + span.log(output=task_result["output"]) + if task_result.get("error") is not None: + root.log(error=task_result["error"]) + else: + root.log(output=task_result["output"], metadata=task_result["metadata"], tags=task_result.get("tags")) if root is not NOOP_SPAN: await self._flush_logs() await self._write( @@ -1216,6 +1581,31 @@ async def _log_score(self, run: Mapping[str, Any], case: dict[str, Any], name: s await self._flush_logs() await self._write(run_id, log_kind, True, item_id) + async def _log_score_error(self, run: Mapping[str, Any], case: dict[str, Any], name: str, error: str) -> None: + run_id = run["run_id"] + item_id = case["id"] + log_kind = _stage_kind("score-log", name) + if await self._read(run_id, log_kind, item_id) is not None: + return + task_result = await self._required(run_id, "task-result", item_id) + root = await self._root_for_case(run, item_id) + propagated = merge_dicts({**(root.propagated_event or {})}, {"span_attributes": {"purpose": "scorer"}}) + with root: + with self._start_child( + root, + run_id, + item_id, + f"score:{name}", + name, + SpanTypeAttribute.SCORE, + input=_scorer_args(case, task_result), + propagated_event=propagated, + ) as span: + span.log(error=error) + if root is not NOOP_SPAN: + await self._flush_logs() + await self._write(run_id, log_kind, True, item_id) + async def _run_classifier(self, run: Mapping[str, Any], case: dict[str, Any], classifier: Any, name: str) -> None: run_id = run["run_id"] item_id = case["id"] @@ -1296,9 +1686,18 @@ async def _advance_case( return False await self._log_task(run, case, experiment) else: - await self._run_ordinary_task(run, case, experiment, self._parameters(run)) + if not isinstance(self.config.task, WorkflowBatchTask): + await self._run_ordinary_task(run, case, experiment, self._parameters(run)) + elif await self._read(run_id, "task-result", item_id) is None: + return False + else: + await self._log_task(run, case, experiment) if await self._read(run_id, "task-log", item_id) is None: return False + task_result = await self._required(run_id, "task-result", item_id) + if task_result.get("error") is not None: + await self.store.add_to_set(self._key(run_id, "progress", "cases"), item_id) + return True for spec, processor in self._submission_specs(case_run): if spec["kind"] == "score": @@ -1310,6 +1709,15 @@ async def _advance_case( complete = False continue await self._log_score(run, case, name) + elif isinstance(scorer, WorkflowBatchScorer): + if await self._read(run_id, _stage_kind("score-result", name), item_id) is None: + error = await self._read(run_id, _stage_kind("score-error", name), item_id) + if error is None: + complete = False + continue + await self._log_score_error(run, case, name, error) + else: + await self._log_score(run, case, name) else: await self._run_ordinary_score(run, case, scorer, name) if await self._read(run_id, _stage_kind("score-log", name), item_id) is None: @@ -1335,7 +1743,7 @@ async def _advance( score() if inspect.isclass(score) and is_scorer(score) else score for score in self.config.scores ] scorer_names = [ - score.name if isinstance(score, WorkflowScorer) else _scorer_name(score, index) + score.name if isinstance(score, (WorkflowScorer, WorkflowBatchScorer)) else _scorer_name(score, index) for index, score in enumerate(resolved_scores) ] if len(scorer_names) != len(set(scorer_names)): @@ -1346,14 +1754,28 @@ async def _advance( if len(classifier_names) != len(set(classifier_names)): raise ValueError("Workflow evaluation classifier names must be unique") - results = await asyncio.gather( - *( - self._advance_case(run, item_id, experiment, resolved_scores, scorer_names, classifier_names) - for item_id in (item_ids if item_ids is not None else run["case_ids"]) - ), - return_exceptions=True, - ) - errors = [result for result in results if isinstance(result, BaseException)] + if isinstance(self.config.task, WorkflowBatchTask): + await self._schedule_batch_stage(run, self.config.task, kind="task") + + async def advance_cases() -> list[BaseException]: + results = await asyncio.gather( + *( + self._advance_case(run, item_id, experiment, resolved_scores, scorer_names, classifier_names) + for item_id in (item_ids if item_ids is not None else run["case_ids"]) + ), + return_exceptions=True, + ) + return [result for result in results if isinstance(result, BaseException)] + + errors = await advance_cases() + if errors: + _raise_callback_errors(errors) + for scorer in self.config.scores: + if isinstance(scorer, WorkflowBatchScorer): + await self._schedule_batch_stage(run, scorer, kind="score") + # Newly submitted batch scorer work is now represented in the waiting + # count; this pass also logs results collected by webhook/poll calls. + errors = await advance_cases() if errors: _raise_callback_errors(errors) if await self.store.get_set_size(self._key(run_id, "progress", "cases")) != len(run["case_ids"]): @@ -1385,11 +1807,15 @@ async def _summary( task_result = await self._required(run["run_id"], "task-result", item_id) scores: dict[str, float | None] = {} for name in scorer_names: - raw = await self._required(run["run_id"], _stage_kind("score-result", name), item_id) + raw = await self._read(run["run_id"], _stage_kind("score-result", name), item_id) + if raw is None: + continue for score in self._prepare_scores(raw, name): scores[score.name] = score.score classifications: dict[str, list[Any]] = {} for name in classifier_names: + if task_result.get("error") is not None: + continue raw = await self._required(run["run_id"], _stage_kind("classification-result", name), item_id) for value in raw: classification = Classification.from_dict(value) @@ -1404,6 +1830,7 @@ async def _summary( expected=datum.get("expected"), metadata=task_result["metadata"], tags=task_result.get("tags"), + error=RuntimeError(task_result["error"]) if task_result.get("error") is not None else None, ) ) evaluator = Evaluator( @@ -1419,6 +1846,11 @@ async def _summary( __all__ = [ + "WorkflowBatchItem", + "WorkflowBatchItemResult", + "WorkflowBatchScorer", + "WorkflowBatchTask", + "WorkflowBatchingOptions", "WorkflowSubmissionCompletionPoll", "WorkflowSubmissionCompletionWebhook", "WorkflowSubmissionContext", From 8ec2da9d31e3f89a0514cf95dd57c96570698f7d Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Mon, 5 Oct 2026 18:33:01 -0400 Subject: [PATCH 2/4] fix(workflow-eval): claim batch items atomically --- py/src/braintrust/test_workflow_eval.py | 104 ++++++++++++++++++++++-- py/src/braintrust/workflow_eval.py | 100 +++++++++++++++++------ 2 files changed, 172 insertions(+), 32 deletions(-) diff --git a/py/src/braintrust/test_workflow_eval.py b/py/src/braintrust/test_workflow_eval.py index 5b126629b..b993daea8 100644 --- a/py/src/braintrust/test_workflow_eval.py +++ b/py/src/braintrust/test_workflow_eval.py @@ -49,6 +49,26 @@ async def test_memory_store_is_atomic_and_copies_values(): assert await store.read("key") == b"first" +@pytest.mark.asyncio +async def test_memory_store_claim_many_is_atomic_for_overlapping_batches(): + store = WorkflowEvalMemoryStore() + + claims = await asyncio.gather( + store.claim_many(["batch/a", "batch/b"], b"batch-1"), + store.claim_many(["batch/b", "batch/c"], b"batch-2"), + ) + + assert sum(claims) == 1 + if claims[0]: + assert await store.read("batch/a") == b"batch-1" + assert await store.read("batch/b") == b"batch-1" + assert await store.read("batch/c") is None + else: + assert await store.read("batch/a") is None + assert await store.read("batch/b") == b"batch-2" + assert await store.read("batch/c") == b"batch-2" + + class _SyncRedis: def __init__(self): self.values = {} @@ -56,11 +76,20 @@ def __init__(self): self.sets = {} self.expirations = {} - def eval(self, script, numkeys, key, member, ttl_ms): - assert "SADD" in script and "PEXPIRE" in script - assert numkeys == 1 - self.sets.setdefault(key, set()).add(member) - self.expirations[key] = ttl_ms + def eval(self, script, numkeys, *args): + if "SADD" in script and "PEXPIRE" in script: + assert numkeys == 1 + key, member, ttl_ms = args + self.sets.setdefault(key, set()).add(member) + self.expirations[key] = ttl_ms + return 1 + assert "EXISTS" in script and numkeys == len(args) - 2 + keys, member, ttl_ms = args[:numkeys], args[-2], args[-1] + if any(key in self.values for key in keys): + return 0 + for key in keys: + self.values[key] = member + self.expirations[key] = ttl_ms return 1 def scard(self, key): @@ -107,13 +136,18 @@ async def test_redis_store_supports_sync_and_async_redis_py(client_type): assert existing.value == b"first" assert all(key.startswith("test:") for key, _ in client.calls) assert all(options["px"] == 1234 for _, options in client.calls) + assert await store.claim_many(["batch/a", "batch/b"], b"batch-1") is True + assert await store.claim_many(["batch/b", "batch/c"], b"batch-2") is False + assert await store.read("batch/a") == b"batch-1" + assert await store.read("batch/b") == b"batch-1" + assert await store.read("batch/c") is None assert await store.get_set_size("progress") == 0 await store.add_to_set("progress", "a") await store.add_to_set("progress", "a") await store.add_to_set("progress", "b") assert await store.get_set_size("progress") == 2 assert client.sets == {"test:progress": {"a", "b"}} - assert client.expirations == {"test:progress": 1234} + assert client.expirations == {"test:batch/a": 1234, "test:batch/b": 1234, "test:progress": 1234} @pytest.mark.asyncio @@ -401,6 +435,64 @@ def score(output): assert score_inputs == ["good"] +@pytest.mark.asyncio +async def test_batch_scorer_wait_window_uses_persisted_wall_clock(monkeypatch): + now = [1_000.0] + monkeypatch.setattr(workflow_eval_module.time, "time", lambda: now[0]) + ready = {"a"} + score_batches = [] + + async def task_submit(item, _context): + return {"input": item.input} + + async def task_collect(submission, _context): + return WorkflowTaskResult(output=submission["input"]) + + def task_poll(submission, _context): + return WorkflowSubmissionPoll("complete" if submission["input"] in ready else "pending") + + async def score_submit(items, _context): + score_batches.append([item.output for item in items]) + return {"size": len(items)} + + async def score_collect(submission, _context): + return [ + WorkflowBatchItemResult(custom_id=f"item-{index}", result=WorkflowScorerResult(score=1)) + for index in range(submission["size"]) + ] + + workflow = define_workflow_eval( + "project", + store=WorkflowEvalMemoryStore(), + data=[{"id": value, "input": value} for value in ("a", "b")], + task=WorkflowTask( + submit=task_submit, + completion=WorkflowSubmissionCompletionPoll(task_poll), + collect=task_collect, + ), + scores=[ + WorkflowBatchScorer( + name="batch-score", + batching=WorkflowBatchingOptions(max_size=2, max_wait_ms=100), + submit=score_submit, + completion=WorkflowSubmissionCompletionPoll( + lambda _submission, _context: WorkflowSubmissionPoll("pending") + ), + collect=score_collect, + ) + ], + ) + + started = await workflow.start(no_send_logs=True) + await workflow.poll(started.run_id) + assert score_batches == [] + + now[0] += 0.101 + await workflow.poll(started.run_id) + + assert score_batches == [["a"]] + + @pytest.mark.asyncio async def test_collect_rejects_array_results(): submission_ids = [] diff --git a/py/src/braintrust/workflow_eval.py b/py/src/braintrust/workflow_eval.py index 32ae9f185..a81767683 100644 --- a/py/src/braintrust/workflow_eval.py +++ b/py/src/braintrust/workflow_eval.py @@ -7,6 +7,7 @@ import hashlib import inspect import threading +import time import uuid from collections.abc import AsyncIterable, Awaitable, Callable, Iterable, Mapping, Sequence from typing import Any, Generic, Literal, Protocol, TypeVar, cast @@ -288,6 +289,10 @@ async def write(self, key: str, value: bytes) -> None: ... async def get_or_set(self, key: str, value: bytes) -> WorkflowEvalStoreEntry: ... + async def claim_many(self, keys: Sequence[str], value: bytes) -> bool: + """Atomically store value for all keys iff none of them exists.""" + ... + async def add_to_set(self, key: str, member: str) -> None: """Atomically add a unique member; retain sets as long as run records.""" ... @@ -320,6 +325,14 @@ async def get_or_set(self, key: str, value: bytes) -> WorkflowEvalStoreEntry: self._values[key] = bytes(value) return WorkflowEvalStoreEntry(value=bytes(value), created=True) + async def claim_many(self, keys: Sequence[str], value: bytes) -> bool: + with self._lock: + if any(key in self._values for key in keys): + return False + for key in keys: + self._values[key] = bytes(value) + return True + async def add_to_set(self, key: str, member: str) -> None: with self._lock: self._sets.setdefault(key, set()).add(member) @@ -379,6 +392,25 @@ async def get_or_set(self, key: str, value: bytes) -> WorkflowEvalStoreEntry: raise TypeError("WorkflowEvalRedisStore expected atomic SET to return str, bytes, or None") return WorkflowEvalStoreEntry(value=base64.b64decode(existing), created=False) + async def claim_many(self, keys: Sequence[str], value: bytes) -> bool: + if not keys: + return True + redis_keys = [f"{self.key_prefix}{key}" for key in keys] + encoded = base64.b64encode(value).decode("ascii") + script = """ + for i = 1, #KEYS do + if redis.call('EXISTS', KEYS[i]) == 1 then return 0 end + end + for i = 1, #KEYS do + redis.call('SET', KEYS[i], ARGV[1], 'PX', ARGV[2]) + end + return 1 + """ + claimed = await self._call(self.client.eval, script, len(redis_keys), *redis_keys, encoded, self.ttl_ms) + if claimed not in (0, 1, b"0", b"1"): + raise TypeError("WorkflowEvalRedisStore expected atomic batch claim to return 0 or 1") + return claimed in (1, b"1") + async def add_to_set(self, key: str, member: str) -> None: await self._call( self.client.eval, @@ -1086,18 +1118,25 @@ async def _schedule_batch_stage( run_id = run["run_id"] scorer_name = processor.name if kind == "score" else None stage = f"score:{scorer_name}" if scorer_name else "task" - candidates: list[str] = [] - for item_id in run["case_ids"]: - if await self._read(run_id, "batch-link", f"{stage}/{item_id}") is not None: - continue - if kind == "task": - if await self._read(run_id, "task-result", item_id) is None: - candidates.append(item_id) - else: - task_result = await self._read(run_id, "task-result", item_id) - score_result = await self._read(run_id, _stage_kind("score-result", cast(str, scorer_name)), item_id) - if task_result is not None and task_result.get("error") is None and score_result is None: - candidates.append(item_id) + + async def find_candidates() -> list[str]: + result: list[str] = [] + for item_id in run["case_ids"]: + if await self._read(run_id, "batch-link", f"{stage}/{item_id}") is not None: + continue + if kind == "task": + if await self._read(run_id, "task-result", item_id) is None: + result.append(item_id) + else: + task_result = await self._read(run_id, "task-result", item_id) + score_result = await self._read( + run_id, _stage_kind("score-result", cast(str, scorer_name)), item_id + ) + if task_result is not None and task_result.get("error") is None and score_result is None: + result.append(item_id) + return result + + candidates = await find_candidates() if not candidates: return @@ -1112,17 +1151,20 @@ async def _schedule_batch_stage( elif processor.batching.max_wait_ms is not None and len(candidates) < max_size: window_key = f"{stage}/{candidates[0]}" first_wait = await self.store.get_or_set( - self._key(run_id, "batch-wait", window_key), str(asyncio.get_running_loop().time()).encode() + self._key(run_id, "batch-wait", window_key), str(time.time()).encode() ) started = float(first_wait.value.decode()) - submit_partial = asyncio.get_running_loop().time() - started >= processor.batching.max_wait_ms / 1000 + submit_partial = time.time() - started >= processor.batching.max_wait_ms / 1000 if len(candidates) < max_size and not submit_partial: return - batches = [candidates[index : index + max_size] for index in range(0, len(candidates), max_size)] - if not submit_partial and batches and len(batches[-1]) < max_size: - batches.pop() - async def submit_batch(item_ids: list[str]) -> None: + def make_batches(item_ids: list[str]) -> list[list[str]]: + batches = [item_ids[index : index + max_size] for index in range(0, len(item_ids), max_size)] + if not submit_partial and batches and len(batches[-1]) < max_size: + batches.pop() + return batches + + async def submit_batch(item_ids: list[str]) -> bool: submission_id = f"submission-{_stable_hex(run_id, stage, *item_ids, length=32)}" spec: dict[str, Any] = { "id": submission_id, @@ -1132,10 +1174,9 @@ async def submit_batch(item_ids: list[str]) -> None: } if scorer_name is not None: spec["scorer_name"] = scorer_name - if not await self._claim(run_id, f"submit:{submission_id}"): - return - for item_id in item_ids: - await self._write(run_id, "batch-link", submission_id, f"{stage}/{item_id}") + claim_keys = [self._key(run_id, "batch-link", f"{stage}/{item_id}") for item_id in item_ids] + if not await self.store.claim_many(claim_keys, _json_bytes(submission_id)): + return False context = WorkflowSubmissionContext(run_id=run_id, submission_id=submission_id) if kind == "task": values = [await self._task_item(run, item_id, self._parameters(run)) for item_id in item_ids] @@ -1172,11 +1213,18 @@ async def submit_batch(item_ids: list[str]) -> None: await self._write(run_id, "submission", submission, submission_id) await self._record_external_locator(run_id, submission_id, external_id) await self._persist_submission_progress(run_id, submission) + return True - results = await asyncio.gather(*(submit_batch(item_ids) for item_ids in batches), return_exceptions=True) - errors = [result for result in results if isinstance(result, BaseException)] - if errors: - _raise_callback_errors(errors) + batches = make_batches(candidates) + while batches: + results = await asyncio.gather(*(submit_batch(item_ids) for item_ids in batches), return_exceptions=True) + errors = [result for result in results if isinstance(result, BaseException)] + if errors: + _raise_callback_errors(errors) + if all(result is True for result in results): + break + candidates = await find_candidates() + batches = make_batches(candidates) async def _persist_submission_progress(self, run_id: str, submission: Mapping[str, Any]) -> None: if submission.get("external_id") is not None: From d81dea7f684ed600d328325c1461f83066a3d149 Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Mon, 5 Oct 2026 18:46:13 -0400 Subject: [PATCH 3/4] fix(workflow-eval): recover batch webhook lookup failures --- py/src/braintrust/test_workflow_eval.py | 121 ++++++++++++++++++++++++ py/src/braintrust/workflow_eval.py | 51 +++++++--- 2 files changed, 158 insertions(+), 14 deletions(-) diff --git a/py/src/braintrust/test_workflow_eval.py b/py/src/braintrust/test_workflow_eval.py index b993daea8..f39d4fa9d 100644 --- a/py/src/braintrust/test_workflow_eval.py +++ b/py/src/braintrust/test_workflow_eval.py @@ -493,6 +493,127 @@ async def score_collect(submission, _context): assert score_batches == [["a"]] +@pytest.mark.asyncio +async def test_batch_scorer_wait_window_starts_for_remainder_while_full_batches_submit(monkeypatch): + now = [1_000.0] + monkeypatch.setattr(workflow_eval_module.time, "time", lambda: now[0]) + ready = {"a", "b", "c"} + score_batches = [] + + async def task_submit(item, _context): + return {"input": item.input} + + async def task_collect(submission, _context): + return WorkflowTaskResult(output=submission["input"]) + + def task_poll(submission, _context): + return WorkflowSubmissionPoll("complete" if submission["input"] in ready else "pending") + + async def score_submit(items, _context): + score_batches.append([item.output for item in items]) + return {"size": len(items)} + + async def score_collect(submission, _context): + return [ + WorkflowBatchItemResult(custom_id=f"item-{index}", result=WorkflowScorerResult(score=1)) + for index in range(submission["size"]) + ] + + workflow = define_workflow_eval( + "project", + store=WorkflowEvalMemoryStore(), + data=[{"id": value, "input": value} for value in ("a", "b", "c", "d")], + task=WorkflowTask( + submit=task_submit, + completion=WorkflowSubmissionCompletionPoll(task_poll), + collect=task_collect, + ), + scores=[ + WorkflowBatchScorer( + name="batch-score", + batching=WorkflowBatchingOptions(max_size=2, max_wait_ms=100), + submit=score_submit, + completion=WorkflowSubmissionCompletionPoll( + lambda _submission, _context: WorkflowSubmissionPoll("pending") + ), + collect=score_collect, + ) + ], + ) + + started = await workflow.start(no_send_logs=True) + await workflow.poll(started.run_id) + assert score_batches == [["a", "b"]] + + now[0] += 0.101 + await workflow.poll(started.run_id) + + assert score_batches == [["a", "b"], ["c"]] + + +@pytest.mark.asyncio +async def test_batch_webhook_id_failure_keeps_submission_recoverable(): + get_id_calls = 0 + score_submit_calls = 0 + + async def task_submit(items, _context): + return {"inputs": [item.input for item in items]} + + async def task_collect(submission, _context): + return [WorkflowBatchItemResult(custom_id="item-0", result=WorkflowTaskResult(output=submission["inputs"][0]))] + + async def score_submit(items, _context): + nonlocal score_submit_calls + score_submit_calls += 1 + return {"external_id": "score-batch"} + + def get_score_external_id(submission, _context): + nonlocal get_id_calls + get_id_calls += 1 + if get_id_calls == 1: + raise RuntimeError("temporary id lookup failure") + return submission["external_id"] + + async def score_collect(_submission, _context): + return [WorkflowBatchItemResult(custom_id="item-0", result=WorkflowScorerResult(score=1))] + + workflow = define_workflow_eval( + "project", + store=WorkflowEvalMemoryStore(), + data=[{"id": "case", "input": "value"}], + task=WorkflowBatchTask( + batching=WorkflowBatchingOptions(max_size=1), + submit=task_submit, + completion=WorkflowSubmissionCompletionPoll( + lambda _submission, _context: WorkflowSubmissionPoll("complete") + ), + collect=task_collect, + ), + scores=[ + WorkflowBatchScorer( + name="batch-score", + batching=WorkflowBatchingOptions(max_size=1), + submit=score_submit, + completion=WorkflowSubmissionCompletionWebhook(get_score_external_id), + collect=score_collect, + ) + ], + ) + + started = await workflow.start(no_send_logs=True) + with pytest.raises(RuntimeError, match="temporary id lookup failure"): + await workflow.poll(started.run_id) + + waiting = await workflow.poll(started.run_id) + assert isinstance(waiting, WorkflowEvalWaitingResult) + assert waiting.pending.webhook == 1 + assert score_submit_calls == 1 + assert get_id_calls == 2 + + completed = await workflow.process_submission_result(external_id="score-batch") + assert isinstance(completed, WorkflowEvalCompletedResult) + + @pytest.mark.asyncio async def test_collect_rejects_array_results(): submission_ids = [] diff --git a/py/src/braintrust/workflow_eval.py b/py/src/braintrust/workflow_eval.py index a81767683..42705719d 100644 --- a/py/src/braintrust/workflow_eval.py +++ b/py/src/braintrust/workflow_eval.py @@ -887,6 +887,19 @@ async def poll(self, run_id: str) -> WorkflowEvalResult: submissions = await self._submission_records(run) async def poll_submission(submission: dict[str, Any], processor: Any) -> None: + if submission["status"] == "external_id_pending": + completion = processor.completion + assert isinstance(completion, WorkflowSubmissionCompletionWebhook) + context = WorkflowSubmissionContext(run_id=run_id, submission_id=submission["id"]) + external_id = await self._call(completion.get_external_id, submission["submission_data"], context) + if not isinstance(external_id, str) or not external_id.strip(): + raise ValueError("Submission webhook get_external_id must return a non-empty string") + submission["external_id"] = external_id + submission["status"] = "submitted" + await self._write(run_id, "submission", submission, submission["id"]) + await self._record_external_locator(run_id, submission["id"], external_id) + await self._persist_submission_progress(run_id, submission) + return if submission["status"] == "failed": if not submission.get("batch"): return @@ -934,6 +947,7 @@ async def poll_submission(submission: dict[str, Any], processor: Any) -> None: poll_submission(submission, processor) for submission, processor in submissions if submission["status"] == "failed" + or submission["status"] == "external_id_pending" or (submission["status"] == "submitted" and submission["mode"] == "poll") ), return_exceptions=True, @@ -1148,13 +1162,16 @@ async def find_candidates() -> list[str]: ) if completed_tasks == len(run["case_ids"]): submit_partial = True - elif processor.batching.max_wait_ms is not None and len(candidates) < max_size: - window_key = f"{stage}/{candidates[0]}" - first_wait = await self.store.get_or_set( - self._key(run_id, "batch-wait", window_key), str(time.time()).encode() - ) - started = float(first_wait.value.decode()) - submit_partial = time.time() - started >= processor.batching.max_wait_ms / 1000 + elif processor.batching.max_wait_ms is not None: + full_item_count = (len(candidates) // max_size) * max_size + remainder = candidates[full_item_count:] + if remainder: + window_key = f"{stage}/{remainder[0]}" + first_wait = await self.store.get_or_set( + self._key(run_id, "batch-wait", window_key), str(time.time()).encode() + ) + started = float(first_wait.value.decode()) + submit_partial = time.time() - started >= processor.batching.max_wait_ms / 1000 if len(candidates) < max_size and not submit_partial: return @@ -1198,19 +1215,25 @@ async def submit_batch(item_ids: list[str]) -> bool: await self._write(run_id, "submission", submission, submission_id) await self._persist_submission_progress(run_id, submission) raise - external_id = None - if isinstance(completion, WorkflowSubmissionCompletionWebhook): - external_id = await self._call(completion.get_external_id, submission_data, context) - if not isinstance(external_id, str) or not external_id.strip(): - raise ValueError("Submission webhook get_external_id must return a non-empty string") submission = { **spec, "submission_data": submission_data, - "external_id": external_id, + "external_id": None, "mode": completion.mode, - "status": "submitted", + "status": "external_id_pending" + if isinstance(completion, WorkflowSubmissionCompletionWebhook) + else "submitted", } await self._write(run_id, "submission", submission, submission_id) + await self._persist_submission_progress(run_id, submission) + external_id = None + if isinstance(completion, WorkflowSubmissionCompletionWebhook): + external_id = await self._call(completion.get_external_id, submission_data, context) + if not isinstance(external_id, str) or not external_id.strip(): + raise ValueError("Submission webhook get_external_id must return a non-empty string") + submission["external_id"] = external_id + submission["status"] = "submitted" + await self._write(run_id, "submission", submission, submission_id) await self._record_external_locator(run_id, submission_id, external_id) await self._persist_submission_progress(run_id, submission) return True From 49e76af819ab77cfa0681eff736b5edaa0e68661 Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Mon, 5 Oct 2026 19:21:47 -0400 Subject: [PATCH 4/4] fix(workflow-eval): make batch retries recoverable --- py/src/braintrust/test_workflow_eval.py | 222 +++++++++++++++-- py/src/braintrust/workflow_eval.py | 317 +++++++++++++++++------- 2 files changed, 439 insertions(+), 100 deletions(-) diff --git a/py/src/braintrust/test_workflow_eval.py b/py/src/braintrust/test_workflow_eval.py index f39d4fa9d..1a03d8c13 100644 --- a/py/src/braintrust/test_workflow_eval.py +++ b/py/src/braintrust/test_workflow_eval.py @@ -50,23 +50,41 @@ async def test_memory_store_is_atomic_and_copies_values(): @pytest.mark.asyncio -async def test_memory_store_claim_many_is_atomic_for_overlapping_batches(): +async def test_memory_store_reserve_batch_is_atomic_for_overlapping_batches(): store = WorkflowEvalMemoryStore() claims = await asyncio.gather( - store.claim_many(["batch/a", "batch/b"], b"batch-1"), - store.claim_many(["batch/b", "batch/c"], b"batch-2"), + store.reserve_batch( + ["workflow/{run-1}/links/a", "workflow/{run-1}/links/b"], + b"batch-1", + "workflow/{run-1}/submissions/batch-1", + b"record-1", + "workflow/{run-1}/leases/batch-1", + b"lease-1", + 1000, + ), + store.reserve_batch( + ["workflow/{run-1}/links/b", "workflow/{run-1}/links/c"], + b"batch-2", + "workflow/{run-1}/submissions/batch-2", + b"record-2", + "workflow/{run-1}/leases/batch-2", + b"lease-2", + 1000, + ), ) assert sum(claims) == 1 if claims[0]: - assert await store.read("batch/a") == b"batch-1" - assert await store.read("batch/b") == b"batch-1" - assert await store.read("batch/c") is None + assert await store.read("workflow/{run-1}/links/a") == b"batch-1" + assert await store.read("workflow/{run-1}/links/b") == b"batch-1" + assert await store.read("workflow/{run-1}/links/c") is None + assert await store.read("workflow/{run-1}/submissions/batch-1") == b"record-1" else: - assert await store.read("batch/a") is None - assert await store.read("batch/b") == b"batch-2" - assert await store.read("batch/c") == b"batch-2" + assert await store.read("workflow/{run-1}/links/a") is None + assert await store.read("workflow/{run-1}/links/b") == b"batch-2" + assert await store.read("workflow/{run-1}/links/c") == b"batch-2" + assert await store.read("workflow/{run-1}/submissions/batch-2") == b"record-2" class _SyncRedis: @@ -83,6 +101,28 @@ def eval(self, script, numkeys, *args): self.sets.setdefault(key, set()).add(member) self.expirations[key] = ttl_ms return 1 + if "KEYS[#KEYS - 1]" in script: + keys = args[:numkeys] + hash_tags = {key.split("{", 1)[1].split("}", 1)[0] for key in keys} + assert len(hash_tags) == 1 + claim_value, submission_value, lease_value, ttl_ms, lease_ttl_ms = args[numkeys:] + if any(key in self.values for key in keys): + return 0 + for key in keys[:-2]: + self.values[key] = claim_value + self.expirations[key] = ttl_ms + self.values[keys[-2]] = submission_value + self.expirations[keys[-2]] = ttl_ms + self.values[keys[-1]] = lease_value + self.expirations[keys[-1]] = lease_ttl_ms + return 1 + if "redis.call('GET', KEYS[1]) == ARGV[1]" in script: + key, expected = args[:2] + if self.values.get(key) == expected: + self.values.pop(key, None) + self.expirations.pop(key, None) + return 1 + return 0 assert "EXISTS" in script and numkeys == len(args) - 2 keys, member, ttl_ms = args[:numkeys], args[-2], args[-1] if any(key in self.values for key in keys): @@ -136,18 +176,36 @@ async def test_redis_store_supports_sync_and_async_redis_py(client_type): assert existing.value == b"first" assert all(key.startswith("test:") for key, _ in client.calls) assert all(options["px"] == 1234 for _, options in client.calls) - assert await store.claim_many(["batch/a", "batch/b"], b"batch-1") is True - assert await store.claim_many(["batch/b", "batch/c"], b"batch-2") is False - assert await store.read("batch/a") == b"batch-1" - assert await store.read("batch/b") == b"batch-1" - assert await store.read("batch/c") is None assert await store.get_set_size("progress") == 0 await store.add_to_set("progress", "a") await store.add_to_set("progress", "a") await store.add_to_set("progress", "b") assert await store.get_set_size("progress") == 2 assert client.sets == {"test:progress": {"a", "b"}} - assert client.expirations == {"test:batch/a": 1234, "test:batch/b": 1234, "test:progress": 1234} + assert client.expirations == {"test:progress": 1234} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("client_type", [_SyncRedis, _AsyncRedis]) +async def test_redis_store_atomically_reserves_batch_in_one_cluster_slot(client_type): + client = client_type() + store = WorkflowEvalRedisStore(client, key_prefix="test:", ttl_ms=1234) + + reserved = await store.reserve_batch( + ["workflow/{run-1}/links/a", "workflow/{run-1}/links/b"], + b"submission-1", + "workflow/{run-1}/submissions/submission-1", + b'{"status":"failed"}', + "workflow/{run-1}/leases/submission-1", + b"lease-1", + 5000, + ) + + assert reserved is True + assert await store.read("workflow/{run-1}/submissions/submission-1") == b'{"status":"failed"}' + assert await store.acquire_lease("workflow/{run-1}/leases/submission-1", b"lease-2", 5000) is False + await store.release_lease("workflow/{run-1}/leases/submission-1", b"lease-1") + assert await store.acquire_lease("workflow/{run-1}/leases/submission-1", b"lease-2", 5000) is True @pytest.mark.asyncio @@ -614,6 +672,140 @@ async def score_collect(_submission, _context): assert isinstance(completed, WorkflowEvalCompletedResult) +@pytest.mark.asyncio +async def test_failed_batch_retry_is_leased_and_webhook_lookup_is_recoverable(): + score_submit_calls = 0 + external_id_calls = 0 + retry_started = asyncio.Event() + allow_retry_to_finish = asyncio.Event() + + async def task_submit(items, _context): + return {"inputs": [item.input for item in items]} + + async def task_collect(submission, _context): + return [WorkflowBatchItemResult(custom_id="item-0", result=WorkflowTaskResult(output=submission["inputs"][0]))] + + async def score_submit(_items, _context): + nonlocal score_submit_calls + score_submit_calls += 1 + if score_submit_calls == 1: + raise RuntimeError("initial provider submit failed") + retry_started.set() + await allow_retry_to_finish.wait() + return {"external_id": "retried-score-job"} + + def get_external_id(submission, _context): + nonlocal external_id_calls + external_id_calls += 1 + if external_id_calls == 1: + raise RuntimeError("temporary webhook ID lookup failure") + return submission["external_id"] + + async def score_collect(_submission, _context): + return [WorkflowBatchItemResult(custom_id="item-0", result=WorkflowScorerResult(score=1))] + + workflow = define_workflow_eval( + "project", + store=WorkflowEvalMemoryStore(), + data=[{"id": "case", "input": "value"}], + task=WorkflowBatchTask( + batching=WorkflowBatchingOptions(max_size=1), + submit=task_submit, + completion=WorkflowSubmissionCompletionPoll( + lambda _submission, _context: WorkflowSubmissionPoll("complete") + ), + collect=task_collect, + ), + scores=[ + WorkflowBatchScorer( + name="batch-score", + batching=WorkflowBatchingOptions(max_size=1), + submit=score_submit, + completion=WorkflowSubmissionCompletionWebhook(get_external_id), + collect=score_collect, + ) + ], + ) + + started = await workflow.start(no_send_logs=True) + with pytest.raises(RuntimeError, match="initial provider submit failed"): + await workflow.poll(started.run_id) + + retry = asyncio.create_task(workflow.poll(started.run_id)) + await retry_started.wait() + concurrent = await workflow.poll(started.run_id) + assert isinstance(concurrent, WorkflowEvalWaitingResult) + allow_retry_to_finish.set() + with pytest.raises(RuntimeError, match="temporary webhook ID lookup failure"): + await retry + + waiting = await workflow.poll(started.run_id) + assert isinstance(waiting, WorkflowEvalWaitingResult) + assert waiting.pending.webhook == 1 + assert score_submit_calls == 2 + assert external_id_calls == 2 + completed = await workflow.process_submission_result(external_id="retried-score-job") + assert isinstance(completed, WorkflowEvalCompletedResult) + + +@pytest.mark.asyncio +async def test_batch_reservation_placeholder_recovers_after_pre_submit_store_failure(): + class _FailOnceStore(WorkflowEvalMemoryStore): + def __init__(self): + super().__init__() + self.fail_case_read = False + self.run_id = None + + async def reserve_batch(self, keys, value, submission_key, submission_value, lease_key, lease_value, ttl_ms): + self.run_id = keys[0].split("{")[1].split("}")[0] + self.fail_case_read = True + return await super().reserve_batch( + keys, value, submission_key, submission_value, lease_key, lease_value, ttl_ms + ) + + async def read(self, key): + if self.fail_case_read and "/case/" in key: + self.fail_case_read = False + raise RuntimeError("transient case read failure") + return await super().read(key) + + store = _FailOnceStore() + submit_calls = 0 + + async def submit(items, _context): + nonlocal submit_calls + submit_calls += 1 + return {"inputs": [item.input for item in items]} + + async def collect(submission, _context): + return [WorkflowBatchItemResult(custom_id="item-0", result=WorkflowTaskResult(output=submission["inputs"][0]))] + + workflow = define_workflow_eval( + "project", + store=store, + data=[{"id": "case", "input": "value"}], + task=WorkflowBatchTask( + batching=WorkflowBatchingOptions(max_size=1), + submit=submit, + completion=WorkflowSubmissionCompletionPoll( + lambda _submission, _context: WorkflowSubmissionPoll("complete") + ), + collect=collect, + ), + ) + + with pytest.raises(RuntimeError, match="transient case read failure"): + await workflow.start(no_send_logs=True) + assert store.run_id is not None + stored_submissions = [value for key, value in store._values.items() if "/submission/" in key] + assert len(stored_submissions) == 1 + + await workflow.poll(store.run_id) + completed = await workflow.poll(store.run_id) + assert isinstance(completed, WorkflowEvalCompletedResult) + assert submit_calls == 1 + + @pytest.mark.asyncio async def test_collect_rejects_array_results(): submission_ids = [] diff --git a/py/src/braintrust/workflow_eval.py b/py/src/braintrust/workflow_eval.py index 42705719d..b9ee968e4 100644 --- a/py/src/braintrust/workflow_eval.py +++ b/py/src/braintrust/workflow_eval.py @@ -70,6 +70,7 @@ DEFAULT_MAX_CONCURRENCY = 10 DEFAULT_REDIS_TTL_MS = 1_000 * 60 * 60 * 24 * 7 +_BATCH_LEASE_TTL_MS = 5 * 60 * 1000 _SCHEMA_PREFIX = "workflow-eval/python/v1" @@ -289,9 +290,20 @@ async def write(self, key: str, value: bytes) -> None: ... async def get_or_set(self, key: str, value: bytes) -> WorkflowEvalStoreEntry: ... - async def claim_many(self, keys: Sequence[str], value: bytes) -> bool: - """Atomically store value for all keys iff none of them exists.""" - ... + async def reserve_batch( + self, + keys: Sequence[str], + value: bytes, + submission_key: str, + submission_value: bytes, + lease_key: str, + lease_value: bytes, + lease_ttl_ms: int, + ) -> bool: ... + + async def acquire_lease(self, key: str, value: bytes, ttl_ms: int) -> bool: ... + + async def release_lease(self, key: str, value: bytes) -> None: ... async def add_to_set(self, key: str, member: str) -> None: """Atomically add a unique member; retain sets as long as run records.""" @@ -306,33 +318,74 @@ class WorkflowEvalMemoryStore: def __init__(self) -> None: self._values: dict[str, bytes] = {} self._sets: dict[str, set[str]] = {} + self._expirations: dict[str, float] = {} self._lock = threading.Lock() + def _expire(self, key: str) -> None: + expires_at = self._expirations.get(key) + if expires_at is not None and expires_at <= time.monotonic(): + self._values.pop(key, None) + self._expirations.pop(key, None) + async def read(self, key: str) -> bytes | None: with self._lock: + self._expire(key) value = self._values.get(key) return bytes(value) if value is not None else None async def write(self, key: str, value: bytes) -> None: with self._lock: self._values[key] = bytes(value) + self._expirations.pop(key, None) async def get_or_set(self, key: str, value: bytes) -> WorkflowEvalStoreEntry: with self._lock: + self._expire(key) existing = self._values.get(key) if existing is not None: return WorkflowEvalStoreEntry(value=bytes(existing), created=False) self._values[key] = bytes(value) return WorkflowEvalStoreEntry(value=bytes(value), created=True) - async def claim_many(self, keys: Sequence[str], value: bytes) -> bool: + async def reserve_batch( + self, + keys: Sequence[str], + value: bytes, + submission_key: str, + submission_value: bytes, + lease_key: str, + lease_value: bytes, + lease_ttl_ms: int, + ) -> bool: with self._lock: - if any(key in self._values for key in keys): + check_keys = [*keys, submission_key, lease_key] + for key in check_keys: + self._expire(key) + if any(key in self._values for key in check_keys): return False for key in keys: self._values[key] = bytes(value) + self._values[submission_key] = bytes(submission_value) + self._values[lease_key] = bytes(lease_value) + self._expirations[lease_key] = time.monotonic() + lease_ttl_ms / 1000 return True + async def acquire_lease(self, key: str, value: bytes, ttl_ms: int) -> bool: + with self._lock: + self._expire(key) + if key in self._values: + return False + self._values[key] = bytes(value) + self._expirations[key] = time.monotonic() + ttl_ms / 1000 + return True + + async def release_lease(self, key: str, value: bytes) -> None: + with self._lock: + self._expire(key) + if self._values.get(key) == value: + self._values.pop(key, None) + self._expirations.pop(key, None) + async def add_to_set(self, key: str, member: str) -> None: with self._lock: self._sets.setdefault(key, set()).add(member) @@ -392,25 +445,60 @@ async def get_or_set(self, key: str, value: bytes) -> WorkflowEvalStoreEntry: raise TypeError("WorkflowEvalRedisStore expected atomic SET to return str, bytes, or None") return WorkflowEvalStoreEntry(value=base64.b64decode(existing), created=False) - async def claim_many(self, keys: Sequence[str], value: bytes) -> bool: - if not keys: - return True - redis_keys = [f"{self.key_prefix}{key}" for key in keys] - encoded = base64.b64encode(value).decode("ascii") + async def reserve_batch( + self, + keys: Sequence[str], + value: bytes, + submission_key: str, + submission_value: bytes, + lease_key: str, + lease_value: bytes, + lease_ttl_ms: int, + ) -> bool: + redis_keys = [f"{self.key_prefix}{key}" for key in [*keys, submission_key, lease_key]] + encoded_claim = base64.b64encode(value).decode("ascii") + encoded_submission = base64.b64encode(submission_value).decode("ascii") script = """ for i = 1, #KEYS do if redis.call('EXISTS', KEYS[i]) == 1 then return 0 end end - for i = 1, #KEYS do - redis.call('SET', KEYS[i], ARGV[1], 'PX', ARGV[2]) + for i = 1, #KEYS - 2 do + redis.call('SET', KEYS[i], ARGV[1], 'PX', ARGV[4]) end + redis.call('SET', KEYS[#KEYS - 1], ARGV[2], 'PX', ARGV[4]) + redis.call('SET', KEYS[#KEYS], ARGV[3], 'PX', ARGV[5]) return 1 """ - claimed = await self._call(self.client.eval, script, len(redis_keys), *redis_keys, encoded, self.ttl_ms) + claimed = await self._call( + self.client.eval, + script, + len(redis_keys), + *redis_keys, + encoded_claim, + encoded_submission, + base64.b64encode(lease_value).decode("ascii"), + self.ttl_ms, + lease_ttl_ms, + ) if claimed not in (0, 1, b"0", b"1"): - raise TypeError("WorkflowEvalRedisStore expected atomic batch claim to return 0 or 1") + raise TypeError("WorkflowEvalRedisStore expected atomic batch reservation to return 0 or 1") return claimed in (1, b"1") + async def acquire_lease(self, key: str, value: bytes, ttl_ms: int) -> bool: + result = await self._call( + self.client.set, f"{self.key_prefix}{key}", base64.b64encode(value).decode("ascii"), px=ttl_ms, nx=True + ) + return result is True or result == b"OK" or result == "OK" + + async def release_lease(self, key: str, value: bytes) -> None: + await self._call( + self.client.eval, + "if redis.call('GET', KEYS[1]) == ARGV[1] then return redis.call('DEL', KEYS[1]) else return 0 end", + 1, + f"{self.key_prefix}{key}", + base64.b64encode(value).decode("ascii"), + ) + async def add_to_set(self, key: str, member: str) -> None: await self._call( self.client.eval, @@ -649,6 +737,28 @@ def _key(self, run_id: str, kind: str, identifier: str | None = None) -> str: suffix = f"/{identifier}" if identifier is not None else "" return f"{_SCHEMA_PREFIX}/{self.definition_key}/{run_id}/{kind}{suffix}" + def _batch_key(self, run_id: str, kind: str, identifier: str) -> str: + # Batch reservation scripts touch several keys, which must share a Redis Cluster slot. + return f"{_SCHEMA_PREFIX}/{self.definition_key}/{{{run_id}}}/{kind}/{identifier}" + + async def _read_submission(self, run_id: str, submission_id: str) -> Any | None: + batch_record = await self.store.read(self._batch_key(run_id, "submission", submission_id)) + if batch_record is not None: + return _decode(batch_record) + return await self._read(run_id, "submission", submission_id) + + async def _read_batch_link(self, run_id: str, stage: str, item_id: str) -> str | None: + value = await self.store.read(self._batch_key(run_id, "batch-link", f"{stage}/{item_id}")) + if value is None: + value = await self.store.read(self._key(run_id, "batch-link", f"{stage}/{item_id}")) + return _decode(value) if value is not None else None + + async def _write_submission(self, run_id: str, submission: Mapping[str, Any]) -> None: + if submission.get("batch"): + await self.store.write(self._batch_key(run_id, "submission", submission["id"]), _json_bytes(submission)) + else: + await self._write(run_id, "submission", submission, submission["id"]) + def _external_key(self, external_id: str) -> str: return f"{_SCHEMA_PREFIX}/{self.definition_key}/external/{_stable_hex(external_id, length=32)}" @@ -896,35 +1006,60 @@ async def poll_submission(submission: dict[str, Any], processor: Any) -> None: raise ValueError("Submission webhook get_external_id must return a non-empty string") submission["external_id"] = external_id submission["status"] = "submitted" - await self._write(run_id, "submission", submission, submission["id"]) + await self._write_submission(run_id, submission) await self._record_external_locator(run_id, submission["id"], external_id) await self._persist_submission_progress(run_id, submission) return if submission["status"] == "failed": if not submission.get("batch"): return - context = WorkflowSubmissionContext(run_id=run_id, submission_id=submission["id"]) - if submission["kind"] == "task": - values = [ - await self._task_item(run, item_id, self._parameters(run)) - for item_id in submission["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)] - submission["submission_data"] = _json_value(await self._call(processor.submit, items, context)) - completion = processor.completion - if isinstance(completion, WorkflowSubmissionCompletionWebhook): - external_id = await self._call(completion.get_external_id, submission["submission_data"], context) - if not isinstance(external_id, str) or not external_id.strip(): - raise ValueError("Submission webhook get_external_id must return a non-empty string") - submission["external_id"] = external_id - await self._record_external_locator(run_id, submission["id"], external_id) - submission["mode"] = completion.mode - submission["status"] = "submitted" - await self._write(run_id, "submission", submission, submission["id"]) - await self._persist_submission_progress(run_id, submission) - 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): + return + try: + current = await self._read_submission(run_id, submission["id"]) + if current is None or current["status"] != "failed": + return + submission = current + context = WorkflowSubmissionContext(run_id=run_id, submission_id=submission["id"]) + if submission["kind"] == "task": + values = [ + await self._task_item(run, item_id, self._parameters(run)) + for item_id in submission["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)] + completion = processor.completion + try: + submission_data = _json_value(await self._call(processor.submit, items, context)) + except Exception: + submission["attempts"] = submission.get("attempts", 0) + 1 + await self._write_submission(run_id, submission) + raise + submission.update( + submission_data=submission_data, + external_id=None, + mode=completion.mode, + status="external_id_pending" + if isinstance(completion, WorkflowSubmissionCompletionWebhook) + else "submitted", + ) + await self._write_submission(run_id, submission) + await self._persist_submission_progress(run_id, submission) + if isinstance(completion, WorkflowSubmissionCompletionWebhook): + external_id = await self._call(completion.get_external_id, submission_data, context) + if not isinstance(external_id, str) or not external_id.strip(): + raise ValueError("Submission webhook get_external_id must return a non-empty string") + submission["external_id"] = external_id + submission["status"] = "submitted" + await self._write_submission(run_id, submission) + await self._record_external_locator(run_id, submission["id"], external_id) + await self._persist_submission_progress(run_id, submission) + return + finally: + await self.store.release_lease(lease_key, lease_value) completion = processor.completion assert isinstance(completion, WorkflowSubmissionCompletionPoll) context = WorkflowSubmissionContext(run_id=run_id, submission_id=submission["id"]) @@ -981,7 +1116,7 @@ async def process_submission_result( if submission_id and external_submission_id and submission_id != external_submission_id: raise ValueError("submission_id and external_id identify different submissions") submission_id = submission_id or external_submission_id - submission = await self._read(run_id, "submission", submission_id) if submission_id else None + submission = await self._read_submission(run_id, submission_id) if submission_id else None if submission is None: raise ValueError("No submission matches this result") if external_id is not None and submission.get("external_id") != external_id: @@ -1034,7 +1169,7 @@ async def _submission_records(self, run: Mapping[str, Any]) -> list[tuple[dict[s records: list[tuple[dict[str, Any], Any]] = [] seen: set[str] = set() for spec, processor in self._submission_specs(run): - record = await self._read(run["run_id"], "submission", spec["id"]) + record = await self._read_submission(run["run_id"], spec["id"]) if record is not None: await self._persist_submission_progress(run["run_id"], record) records.append((record, processor)) @@ -1047,10 +1182,10 @@ async def _submission_records(self, run: Mapping[str, Any]) -> list[tuple[dict[s batch_processors[f"score:{scorer.name}"] = scorer for item_id in run["case_ids"]: for stage, processor in batch_processors.items(): - submission_id = await self._read(run["run_id"], "batch-link", f"{stage}/{item_id}") + submission_id = await self._read_batch_link(run["run_id"], stage, item_id) if not submission_id or submission_id in seen: continue - record = await self._read(run["run_id"], "submission", submission_id) + record = await self._read_submission(run["run_id"], submission_id) if record is not None: await self._persist_submission_progress(run["run_id"], record) records.append((record, processor)) @@ -1136,7 +1271,7 @@ async def _schedule_batch_stage( async def find_candidates() -> list[str]: result: list[str] = [] for item_id in run["case_ids"]: - if await self._read(run_id, "batch-link", f"{stage}/{item_id}") is not None: + if await self._read_batch_link(run_id, stage, item_id) is not None: continue if kind == "task": if await self._read(run_id, "task-result", item_id) is None: @@ -1191,52 +1326,64 @@ async def submit_batch(item_ids: list[str]) -> bool: } if scorer_name is not None: spec["scorer_name"] = scorer_name - claim_keys = [self._key(run_id, "batch-link", f"{stage}/{item_id}") for item_id in item_ids] - if not await self.store.claim_many(claim_keys, _json_bytes(submission_id)): - return False - context = WorkflowSubmissionContext(run_id=run_id, submission_id=submission_id) - if kind == "task": - 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 item_ids] - submitted_items = [WorkflowBatchItem(f"item-{index}", value) for index, value in enumerate(values)] + claim_keys = [self._batch_key(run_id, "batch-link", f"{stage}/{item_id}") for item_id in item_ids] + submission_key = self._batch_key(run_id, "submission", submission_id) + lease_key = self._batch_key(run_id, "retry-lease", submission_id) + lease_value = uuid.uuid4().hex.encode() completion = processor.completion - try: - submission_data = _json_value(await self._call(processor.submit, submitted_items, context)) - except Exception: - submission = { - **spec, - "submission_data": None, - "external_id": None, - "mode": completion.mode, - "status": "failed", - "attempts": 1, - } - await self._write(run_id, "submission", submission, submission_id) - await self._persist_submission_progress(run_id, submission) - raise - submission = { + reservation = { **spec, - "submission_data": submission_data, + "submission_data": None, "external_id": None, "mode": completion.mode, - "status": "external_id_pending" - if isinstance(completion, WorkflowSubmissionCompletionWebhook) - else "submitted", + "status": "failed", + "attempts": 0, } - await self._write(run_id, "submission", submission, submission_id) - await self._persist_submission_progress(run_id, submission) - external_id = None - if isinstance(completion, WorkflowSubmissionCompletionWebhook): - external_id = await self._call(completion.get_external_id, submission_data, context) - if not isinstance(external_id, str) or not external_id.strip(): - raise ValueError("Submission webhook get_external_id must return a non-empty string") - submission["external_id"] = external_id - submission["status"] = "submitted" - await self._write(run_id, "submission", submission, submission_id) - await self._record_external_locator(run_id, submission_id, external_id) - await self._persist_submission_progress(run_id, submission) - return True + if not await self.store.reserve_batch( + claim_keys, + _json_bytes(submission_id), + submission_key, + _json_bytes(reservation), + lease_key, + lease_value, + _BATCH_LEASE_TTL_MS, + ): + return False + try: + context = WorkflowSubmissionContext(run_id=run_id, submission_id=submission_id) + if kind == "task": + 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 item_ids] + submitted_items = [WorkflowBatchItem(f"item-{index}", value) for index, value in enumerate(values)] + try: + submission_data = _json_value(await self._call(processor.submit, submitted_items, context)) + except Exception: + reservation["attempts"] = 1 + await self._write_submission(run_id, reservation) + await self._persist_submission_progress(run_id, reservation) + raise + reservation.update( + submission_data=submission_data, + status="external_id_pending" + if isinstance(completion, WorkflowSubmissionCompletionWebhook) + else "submitted", + ) + await self._write_submission(run_id, reservation) + await self._persist_submission_progress(run_id, reservation) + external_id = None + if isinstance(completion, WorkflowSubmissionCompletionWebhook): + external_id = await self._call(completion.get_external_id, submission_data, context) + if not isinstance(external_id, str) or not external_id.strip(): + raise ValueError("Submission webhook get_external_id must return a non-empty string") + reservation["external_id"] = external_id + reservation["status"] = "submitted" + await self._write_submission(run_id, reservation) + await self._record_external_locator(run_id, submission_id, external_id) + await self._persist_submission_progress(run_id, reservation) + return True + finally: + await self.store.release_lease(lease_key, lease_value) batches = make_batches(candidates) while batches: @@ -1265,7 +1412,7 @@ async def _collect_submission(self, run: Mapping[str, Any], submission: dict[str if submission.get("batch"): await self._collect_batch_submission(run, submission, processor, context) submission["status"] = "complete" - await self._write(run_id, "submission", submission, submission["id"]) + await self._write_submission(run_id, submission) await self._persist_submission_progress(run_id, submission) return result = await self._call(processor.collect, submission["submission_data"], context) @@ -1294,7 +1441,7 @@ async def _collect_submission(self, run: Mapping[str, Any], submission: dict[str run_id, _stage_kind("score-result", submission["scorer_name"]), _json_value(result.score), item_id ) submission["status"] = "complete" - await self._write(run_id, "submission", submission, submission["id"]) + await self._write_submission(run_id, submission) # Replay repairs an interrupted progress update without collecting again. await self._persist_submission_progress(run_id, submission)