From 87a15e5e035115b4c2ca72fe7764f6ef7aa3cfa5 Mon Sep 17 00:00:00 2001 From: lforst <8118419+lforst@users.noreply.github.com> Date: Tue, 1 Sep 2026 15:26:32 +0000 Subject: [PATCH 1/3] feat: Add batch evals api --- py/src/braintrust/__init__.py | 1 + py/src/braintrust/durable_eval.py | 1450 +++++++++++++++++ py/src/braintrust/logger.py | 71 +- py/src/braintrust/test_durable_eval.py | 600 +++++++ .../type_tests/test_durable_eval.py | 81 + 5 files changed, 2198 insertions(+), 5 deletions(-) create mode 100644 py/src/braintrust/durable_eval.py create mode 100644 py/src/braintrust/test_durable_eval.py create mode 100644 py/src/braintrust/type_tests/test_durable_eval.py diff --git a/py/src/braintrust/__init__.py b/py/src/braintrust/__init__.py index 1e66fced..1047b63d 100644 --- a/py/src/braintrust/__init__.py +++ b/py/src/braintrust/__init__.py @@ -63,6 +63,7 @@ def is_equal(expected, output): from .audit import * from .auto import auto_instrument as auto_instrument from .dataset_pipeline import * +from .durable_eval import * from .framework import * from .framework2 import * from .functions.invoke import * diff --git a/py/src/braintrust/durable_eval.py b/py/src/braintrust/durable_eval.py new file mode 100644 index 00000000..95da3975 --- /dev/null +++ b/py/src/braintrust/durable_eval.py @@ -0,0 +1,1450 @@ +"""Experimental durable evaluation support for asynchronous batch providers.""" + +import asyncio +import base64 +import dataclasses +import functools +import hashlib +import inspect +import threading +import uuid +from collections.abc import AsyncIterable, Awaitable, Callable, Iterable, Mapping, Sequence +from typing import Any, Generic, Literal, Protocol, TypeVar, cast + +from .bt_json import bt_dumps, bt_loads +from .env import BraintrustEnv +from .framework import ( + BaseExperiment, + DictEvalHooks, + EvalCase, + EvalClassifier, + EvalData, + EvalResult, + EvalScorer, + EvalTask, + Evaluator, + OneOrMoreScores, + _build_classification_span_output, + _build_span_metadata, + _classifier_name, + _get_persisted_base_experiment_id, + _normalize_score, + _scorer_name, + _validate_classification_result, + _validated_object_reference, + await_or_run, + build_local_summary, + call_user_fn, + init_experiment, + run_evaluator, +) +from .git_fields import GitMetadataSettings, RepoInfo +from .logger import ( + NOOP_SPAN, + BraintrustState, + Dataset, + Experiment, + ExperimentSummary, + Metadata, + Span, + _internal_get_global_state, + _internal_resume_span, + _internal_start_span_with_initial_merge, + span_components_to_object_id, +) +from .logger import init as _init_experiment +from .parameters import EvalParameters, RemoteEvalParameters, ValidatedParameters, validate_parameters +from .score import Classification, Score, ScoreLike, is_score, is_scorer +from .span_identifier_v3 import span_object_type_v3_to_typed_string +from .span_identifier_v4 import SpanComponentsV4 +from .span_types import SpanTypeAttribute +from .trace import LocalTrace +from .util import get_signature, merge_dicts + + +Input = TypeVar("Input") +Output = TypeVar("Output") +Expected = TypeVar("Expected") +SubmissionData = TypeVar("SubmissionData") + +DEFAULT_BATCH_SIZE = 1_000 +DEFAULT_REDIS_TTL_MS = 1_000 * 60 * 60 * 24 * 7 +_SCHEMA_PREFIX = "durable-eval/python/v1" + + +@dataclasses.dataclass(frozen=True) +class BatchContext: + """Identifiers supplied to a durable batch processor callback.""" + + run_id: str + batch_id: str + + +@dataclasses.dataclass(frozen=True) +class BatchPollResult: + """The current state returned by a polling completion callback.""" + + status: Literal["pending", "complete", "failed"] + error: Any = None + + +@dataclasses.dataclass(frozen=True) +class BatchCompletionPoll(Generic[SubmissionData]): + """Configures a batch processor whose provider is checked by polling.""" + + poll: Callable[[SubmissionData, BatchContext], BatchPollResult | Awaitable[BatchPollResult]] + mode: Literal["poll"] = dataclasses.field(default="poll", init=False) + + +@dataclasses.dataclass(frozen=True) +class BatchCompletionWebhook(Generic[SubmissionData]): + """Configures a batch processor completed by an incoming webhook.""" + + get_external_id: Callable[[SubmissionData, BatchContext], str | Awaitable[str]] + mode: Literal["webhook"] = dataclasses.field(default="webhook", init=False) + + +BatchCompletion = BatchCompletionPoll[SubmissionData] | BatchCompletionWebhook[SubmissionData] + + +@dataclasses.dataclass(frozen=True) +class BatchTaskItem(Generic[Input, Expected]): + """A stable task item submitted to a provider batch.""" + + id: str + input: Input + expected: Expected | None + metadata: Metadata + tags: list[str] | None + parameters: ValidatedParameters | None + trial_index: int + + +@dataclasses.dataclass(frozen=True) +class BatchTaskResult(Generic[Output]): + """A collected result for one task item.""" + + id: str + output: Output + metadata: Metadata | None = None + tags: list[str] | None = None + + +@dataclasses.dataclass(frozen=True) +class BatchScorerItem(Generic[Input, Output, Expected]): + """A stable scorer item submitted to a provider batch.""" + + id: str + input: Input + output: Output + expected: Expected | None + metadata: Metadata + tags: list[str] | None + trial_index: int + + +@dataclasses.dataclass(frozen=True) +class BatchScorerResult: + """A collected score for one scorer item.""" + + id: str + score: OneOrMoreScores + + +@dataclasses.dataclass(frozen=True) +class BatchTask(Generic[Input, Output, Expected, SubmissionData]): + """Runs an evaluation task through asynchronous provider batch operations.""" + + submit: Callable[[list[BatchTaskItem[Input, Expected]], BatchContext], SubmissionData | Awaitable[SubmissionData]] + completion: BatchCompletion[SubmissionData] + collect: Callable[ + [SubmissionData, BatchContext], list[BatchTaskResult[Output]] | Awaitable[list[BatchTaskResult[Output]]] + ] + batch_size: int = DEFAULT_BATCH_SIZE + + def __post_init__(self) -> None: + if not isinstance(self.batch_size, int) or isinstance(self.batch_size, bool) or self.batch_size < 1: + raise ValueError("BatchTask batch_size must be a positive integer") + + +@dataclasses.dataclass(frozen=True) +class BatchScorer(Generic[Input, Output, Expected, SubmissionData]): + """Runs an evaluation scorer through asynchronous provider batch operations.""" + + name: str + submit: Callable[ + [list[BatchScorerItem[Input, Output, Expected]], BatchContext], SubmissionData | Awaitable[SubmissionData] + ] + completion: BatchCompletion[SubmissionData] + collect: Callable[[SubmissionData, BatchContext], list[BatchScorerResult] | Awaitable[list[BatchScorerResult]]] + batch_size: int = DEFAULT_BATCH_SIZE + + def __post_init__(self) -> None: + if not self.name: + raise ValueError("BatchScorer name must be a non-empty string") + if not isinstance(self.batch_size, int) or isinstance(self.batch_size, bool) or self.batch_size < 1: + raise ValueError("BatchScorer batch_size must be a positive integer") + + +@dataclasses.dataclass(frozen=True) +class DurableEvalStoreEntry: + """Result of an atomic durable-store get-or-set operation.""" + + value: bytes + created: bool + + +class DurableEvalStore(Protocol): + """Minimal persistence interface used by durable evaluations.""" + + async def read(self, key: str) -> bytes | None: ... + + async def write(self, key: str, value: bytes) -> None: ... + + async def get_or_set(self, key: str, value: bytes) -> DurableEvalStoreEntry: ... + + +class DurableEvalMemoryStore: + """Process-local durable evaluation state, intended for tests and local runs.""" + + def __init__(self) -> None: + self._values: dict[str, bytes] = {} + self._lock = threading.Lock() + + async def read(self, key: str) -> bytes | None: + with self._lock: + 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) + + async def get_or_set(self, key: str, value: bytes) -> DurableEvalStoreEntry: + with self._lock: + existing = self._values.get(key) + if existing is not None: + return DurableEvalStoreEntry(value=bytes(existing), created=False) + self._values[key] = bytes(value) + return DurableEvalStoreEntry(value=bytes(value), created=True) + + +class DurableEvalRedisStore: + """Durable state backed by an existing sync or async redis-py client.""" + + def __init__( + self, + client: Any, + *, + key_prefix: str = "braintrust-eval:", + ttl_ms: int = DEFAULT_REDIS_TTL_MS, + ) -> None: + if not isinstance(ttl_ms, int) or isinstance(ttl_ms, bool) or ttl_ms < 1: + raise ValueError("DurableEvalRedisStore ttl_ms must be a positive integer") + self.client = client + self.key_prefix = key_prefix + self.ttl_ms = ttl_ms + + async def _call(self, method: Callable[..., Any], *args: Any, **kwargs: Any) -> Any: + if inspect.iscoroutinefunction(method): + return await method(*args, **kwargs) + value = await asyncio.get_running_loop().run_in_executor(None, functools.partial(method, *args, **kwargs)) + if inspect.isawaitable(value): + return await value + return value + + async def read(self, key: str) -> bytes | None: + value = await self._call(self.client.get, f"{self.key_prefix}{key}") + if value is None: + return None + if isinstance(value, bytes): + value = value.decode("ascii") + if not isinstance(value, str): + raise TypeError("DurableEvalRedisStore expected GET to return str, bytes, or None") + return base64.b64decode(value) + + async def write(self, key: str, value: bytes) -> None: + encoded = base64.b64encode(value).decode("ascii") + await self._call(self.client.set, f"{self.key_prefix}{key}", encoded, px=self.ttl_ms) + + async def get_or_set(self, key: str, value: bytes) -> DurableEvalStoreEntry: + redis_key = f"{self.key_prefix}{key}" + encoded = base64.b64encode(value).decode("ascii") + existing = await self._call(self.client.set, redis_key, encoded, px=self.ttl_ms, nx=True, get=True) + if existing is None: + return DurableEvalStoreEntry(value=bytes(value), created=True) + if isinstance(existing, bytes): + existing = existing.decode("ascii") + if not isinstance(existing, str): + raise TypeError("DurableEvalRedisStore expected atomic SET to return str, bytes, or None") + return DurableEvalStoreEntry(value=base64.b64decode(existing), created=False) + + +@dataclasses.dataclass(frozen=True) +class DurableEvalPending: + """Counts of submitted provider batches awaiting completion.""" + + poll: int + webhook: int + + +@dataclasses.dataclass(frozen=True) +class DurableEvalWaitingResult: + run_id: str + pending: DurableEvalPending + status: Literal["waiting"] = dataclasses.field(default="waiting", init=False) + + +@dataclasses.dataclass(frozen=True) +class DurableEvalCompletedResult: + run_id: str + pending: DurableEvalPending + summary: ExperimentSummary + status: Literal["completed"] = dataclasses.field(default="completed", init=False) + + +@dataclasses.dataclass(frozen=True) +class DurableEvalFailedResult: + run_id: str + batch_id: str + error: Any + pending: DurableEvalPending + status: Literal["failed"] = dataclasses.field(default="failed", init=False) + + +DurableEvalResult = DurableEvalWaitingResult | DurableEvalCompletedResult | DurableEvalFailedResult + + +@dataclasses.dataclass(frozen=True) +class _DurableEvalConfig(Generic[Input, Output, Expected]): + project_name: str + store: DurableEvalStore + data: EvalData[Input, Expected] + task: EvalTask[Input, Output, Expected] | BatchTask[Input, Output, Expected, Any] + scores: Sequence[EvalScorer[Input, Output, Expected] | BatchScorer[Input, Output, Expected, Any]] + classifiers: Sequence[EvalClassifier[Input, Output, Expected]] + case_id: Callable[[EvalCase[Input, Expected]], str | Awaitable[str]] | None + experiment_name: str | None + trial_count: int + metadata: Metadata | None + tags: Sequence[str] | None + is_public: bool + project_id: str | None + base_experiment_name: str | None + base_experiment_id: str | None + git_metadata_settings: GitMetadataSettings | None + repo_info: RepoInfo | None + description: str | None + summarize_scores: bool + parameters: EvalParameters | RemoteEvalParameters | None + state: BraintrustState | None + + +class DurableEval(Generic[Input, Output, Expected]): + """A durable evaluation definition that can be started and resumed.""" + + def __init__(self, config: _DurableEvalConfig[Input, Output, Expected]) -> None: + self._config = config + + async def start( + self, parameters: Mapping[str, Any] | None = None, *, no_send_logs: bool = False + ) -> DurableEvalResult: + return await _DurableEvalRunner(self._config).start(parameters, no_send_logs=no_send_logs) + + async def status(self, run_id: str) -> DurableEvalResult: + return await _DurableEvalRunner(self._config).status(run_id) + + async def poll(self, run_id: str) -> DurableEvalResult: + return await _DurableEvalRunner(self._config).poll(run_id) + + async def process_batch_result( + self, run_id: str, *, batch_id: str | None = None, external_id: str | None = None + ) -> DurableEvalResult: + return await _DurableEvalRunner(self._config).process_batch_result( + run_id, batch_id=batch_id, external_id=external_id + ) + + +def define_durable_eval( + project_name: str, + *, + store: DurableEvalStore, + data: EvalData[Input, Expected], + task: EvalTask[Input, Output, Expected] | BatchTask[Input, Output, Expected, Any], + scores: Sequence[EvalScorer[Input, Output, Expected] | BatchScorer[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, + trial_count: int = 1, + metadata: Metadata | None = None, + tags: Sequence[str] | None = None, + is_public: bool = False, + project_id: str | None = None, + base_experiment_name: str | None = None, + base_experiment_id: str | None = None, + git_metadata_settings: GitMetadataSettings | None = None, + repo_info: RepoInfo | None = None, + description: str | None = None, + summarize_scores: bool = True, + parameters: EvalParameters | RemoteEvalParameters | None = None, + state: BraintrustState | None = None, +) -> DurableEval[Input, Output, Expected]: + """Define an experimental evaluation that can pause across provider batches.""" + if not isinstance(trial_count, int) or isinstance(trial_count, bool) or trial_count < 1: + raise ValueError("trial_count must be a positive integer") + return DurableEval( + _DurableEvalConfig( + project_name=project_name, + store=store, + data=data, + task=task, + scores=list(scores or []), + classifiers=list(classifiers or []), + case_id=case_id, + experiment_name=experiment_name, + trial_count=trial_count, + metadata=metadata, + tags=tags, + is_public=is_public, + project_id=project_id, + base_experiment_name=base_experiment_name, + base_experiment_id=base_experiment_id, + git_metadata_settings=git_metadata_settings, + repo_info=repo_info, + description=description, + summarize_scores=summarize_scores, + parameters=parameters, + state=state, + ) + ) + + +def _json_bytes(value: Any) -> bytes: + return bt_dumps(value).encode("utf-8") + + +def _decode(value: bytes) -> Any: + return bt_loads(value.decode("utf-8")) + + +def _json_value(value: Any) -> Any: + """Validate and normalize a value at the durable JSON boundary.""" + return _decode(_json_bytes(value)) + + +def _stable_hex(*parts: str, length: int) -> str: + return hashlib.sha256("\0".join(parts).encode("utf-8")).hexdigest()[:length] + + +def _stable_uuid(*parts: str) -> str: + digest = hashlib.sha256("\0".join(parts).encode("utf-8")).digest() + return str(uuid.UUID(bytes=digest[:16], version=5)) + + +def _stage_kind(prefix: str, name: str) -> str: + return f"{prefix}-{_stable_hex(name, length=32)}" + + +def _scorer_args(case: Mapping[str, Any], task_result: Mapping[str, Any]) -> dict[str, Any]: + datum = case["datum"] + return { + "input": datum["input"], + "output": task_result["output"], + "expected": datum.get("expected"), + "metadata": task_result["metadata"], + "id": case["case_id"], + "tags": task_result.get("tags"), + } + + +def _score_fields(result: ScoreLike) -> dict[str, Any]: + return {key: value for key, value in result.as_dict().items() if key not in ("metadata", "name")} + + +class _DurableEvalRunner(Generic[Input, Output, Expected]): + def __init__(self, config: _DurableEvalConfig[Input, Output, Expected]) -> None: + self.config = config + self.store = config.store + self.eval_name = config.experiment_name or config.project_name + self.definition_key = _stable_hex(config.project_name, self.eval_name, length=32) + + 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}" + + 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 + + async def _required(self, run_id: str, kind: str, identifier: str | None = None) -> Any: + value = await self._read(run_id, kind, identifier) + if value is None: + target = f" {identifier!r}" if identifier is not None else "" + raise ValueError(f"Unknown durable evaluation {kind}{target} for run {run_id!r}") + return value + + 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 _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 + + async def _call(self, fn: Callable[..., Any], *args: Any, **kwargs: Any) -> Any: + return await await_or_run(asyncio.get_running_loop(), fn, *args, **kwargs) + + async def _flush_logs(self) -> None: + state = self.config.state or _internal_get_global_state() + background_logger = state.global_bg_logger() + await asyncio.get_running_loop().run_in_executor(None, background_logger.flush) + + def _parameters(self, run: Mapping[str, Any]) -> ValidatedParameters | None: + raw = run.get("parameters") + if self.config.parameters is None: + return cast(ValidatedParameters | None, raw) + return validate_parameters(raw or {}, self.config.parameters) + + def _experiment(self, run: Mapping[str, Any]) -> Experiment | None: + if run["no_send_logs"]: + return None + experiment_parameters = None + if isinstance(self.config.parameters, RemoteEvalParameters) and self.config.parameters.id is not None: + experiment_parameters = {"id": self.config.parameters.id} + if self.config.parameters.version is not None: + experiment_parameters["version"] = self.config.parameters.version + dataset = self.config.data if isinstance(self.config.data, Dataset) else None + return init_experiment( + project_name=self.config.project_name if self.config.project_id is None else None, + project_id=self.config.project_id, + experiment_name=run["experiment_name"], + description=self.config.description, + metadata=self.config.metadata, + tags=self.config.tags, + is_public=self.config.is_public, + update=True, + base_experiment=self.config.base_experiment_name, + base_experiment_id=self.config.base_experiment_id, + git_metadata_settings=self.config.git_metadata_settings, + repo_info=self.config.repo_info, + dataset=dataset, + parameters=experiment_parameters, + state=self.config.state, + ) + + async def _resolve_data(self, experiment: Experiment | None) -> list[EvalCase[Input, Expected]]: + data: Any = self.config.data + if inspect.isclass(data): + data = data() + if isinstance(data, BaseExperiment): + if experiment is None: + raise ValueError("Cannot use BaseExperiment without sending logs") + base_name = data.name + if base_name is None: + base = experiment.fetch_base_experiment() + if base is None: + raise ValueError("BaseExperiment failed to resolve a base experiment") + base_name = base.name + data = _init_experiment( + project=self.config.project_name if self.config.project_id is None else None, + project_id=self.config.project_id, + experiment=base_name, + open=True, + set_current=False, + state=self.config.state, + ).as_dataset() + elif callable(data) and not isinstance(data, Dataset): + data = await self._call(data) + if inspect.isawaitable(data): + data = await data + + values: list[Any] = [] + if isinstance(data, AsyncIterable): + async for value in data: + values.append(value) + else: + values.extend(data) + return [value if isinstance(value, EvalCase) else EvalCase.from_dict(value) for value in values] + + def _uses_batch_processor(self) -> bool: + return isinstance(self.config.task, BatchTask) or any( + isinstance(score, BatchScorer) for score in self.config.scores + ) + + async def _save_completed(self, run: Mapping[str, Any], summary: ExperimentSummary) -> DurableEvalCompletedResult: + completed = {**run, "status": "completed", "summary": _json_value(summary.as_dict())} + await self._write(run["run_id"], "run", completed) + return DurableEvalCompletedResult( + run_id=run["run_id"], pending=DurableEvalPending(poll=0, webhook=0), summary=summary + ) + + async def _save_failed( + self, run: Mapping[str, Any], batch: Mapping[str, Any], error: Any + ) -> DurableEvalFailedResult: + normalized_error = _json_value(error) + failure = {"batch_id": batch["id"], "error": normalized_error} + failed_run = {**run, "status": "failed", "failure": failure} + await self._write(run["run_id"], "run", failed_run) + await self._write( + run["run_id"], + "batch", + {**batch, "status": "failed", "error": normalized_error}, + batch["id"], + ) + return DurableEvalFailedResult( + run_id=run["run_id"], + batch_id=batch["id"], + error=normalized_error, + pending=DurableEvalPending(poll=0, webhook=0), + ) + + async def _run_ordinary_evaluator( + self, + run: Mapping[str, Any], + data: list[EvalCase[Input, Expected]], + experiment: Experiment | None, + parameters: ValidatedParameters, + ) -> DurableEvalCompletedResult: + evaluator = Evaluator( + project_name=self.config.project_name, + eval_name=self.eval_name, + data=data, + task=cast(EvalTask[Input, Output, Expected], self.config.task), + scores=cast(Sequence[EvalScorer[Input, Output, Expected]], self.config.scores), + classifiers=list(self.config.classifiers), + experiment_name=run["experiment_name"], + metadata=self.config.metadata, + tags=self.config.tags, + trial_count=self.config.trial_count, + is_public=self.config.is_public, + update=True, + project_id=self.config.project_id, + base_experiment_name=self.config.base_experiment_name, + base_experiment_id=self.config.base_experiment_id, + git_metadata_settings=self.config.git_metadata_settings, + repo_info=self.config.repo_info, + description=self.config.description, + summarize_scores=self.config.summarize_scores, + parameters=self.config.parameters, + parameter_values=cast(dict[str, Any], parameters), + ) + result = await run_evaluator( + experiment, + evaluator, + position=None, + filters=[], + state=self.config.state, + ) + return await self._save_completed(run, result.summary) + + async def _persist_cases(self, run: dict[str, Any], data: Sequence[EvalCase[Input, Expected]]) -> None: + seen_case_ids: set[str] = set() + item_ids: list[str] = [] + for datum in data: + case_id = datum.id + if not case_id and self.config.case_id is not None: + case_id = await self._call(self.config.case_id, datum) + if not isinstance(case_id, str) or not case_id: + raise ValueError("Every durable evaluation case must have a non-empty id or be assigned by case_id") + if case_id in seen_case_ids: + raise ValueError(f"Durable evaluation case IDs must be unique; found duplicate {case_id!r}") + seen_case_ids.add(case_id) + + trial_count = datum.trial_count if datum.trial_count is not None else self.config.trial_count + if not isinstance(trial_count, int) or isinstance(trial_count, bool) or trial_count < 1: + raise ValueError(f"trial_count for case {case_id!r} must be a positive integer") + for trial_index in range(trial_count): + item_id = f"{case_id}:trial:{trial_index}" + item_ids.append(item_id) + await self._write( + run["run_id"], + "case", + { + "id": item_id, + "case_id": case_id, + "trial_index": trial_index, + "datum": _json_value( + {field.name: getattr(datum, field.name) for field in dataclasses.fields(datum)} + ), + "metadata": _json_value(dict(datum.metadata or {})), + "tags": list(datum.tags) if datum.tags is not None else None, + }, + item_id, + ) + run["case_ids"] = item_ids + await self._write(run["run_id"], "run", run) + + async def start(self, parameters: Mapping[str, Any] | None, *, no_send_logs: bool) -> DurableEvalResult: + validated_parameters = validate_parameters(parameters or {}, self.config.parameters) + run_id = str(uuid.uuid4()) + run = { + "run_id": run_id, + "experiment_name": self.config.experiment_name or f"{self.eval_name}-{run_id}", + "no_send_logs": no_send_logs, + "parameters": _json_value(validated_parameters), + "status": "running", + "case_ids": [], + } + created = await self.store.get_or_set(self._key(run_id, "run"), _json_bytes(run)) + if not created.created: + raise RuntimeError(f"Durable evaluation run ID collision: {run_id}") + + experiment = self._experiment(run) + data = await self._resolve_data(experiment) + if not self._uses_batch_processor(): + return await self._run_ordinary_evaluator(run, data, experiment, validated_parameters) + + await self._persist_cases(run, data) + return await self._advance(run, experiment=experiment) + + async def status(self, run_id: str) -> DurableEvalResult: + run = await self._required(run_id, "run") + if run["status"] == "completed": + return DurableEvalCompletedResult( + run_id=run_id, + pending=DurableEvalPending(poll=0, webhook=0), + summary=ExperimentSummary.from_dict_deep(run["summary"]), + ) + if run["status"] == "failed": + failure = run["failure"] + return DurableEvalFailedResult( + run_id=run_id, + batch_id=failure["batch_id"], + error=failure.get("error"), + pending=DurableEvalPending(poll=0, webhook=0), + ) + return await self._waiting_status(run) + + async def poll(self, run_id: str) -> DurableEvalResult: + run = await self._required(run_id, "run") + if run["status"] in ("completed", "failed"): + return await self.status(run_id) + # Snapshot first: newly submitted downstream batches wait until the next poll call. + batches = await self._batch_records(run) + for batch, processor in batches: + if batch["status"] != "submitted" or batch["mode"] != "poll": + continue + completion = processor.completion + assert isinstance(completion, BatchCompletionPoll) + context = BatchContext(run_id=run_id, batch_id=batch["id"]) + result = await self._call(completion.poll, batch["submission_data"], context) + if not isinstance(result, BatchPollResult): + if isinstance(result, Mapping): + result = BatchPollResult(**result) + else: + raise TypeError("Batch poll callback must return BatchPollResult") + if result.status not in ("pending", "complete", "failed"): + raise ValueError(f"Batch poll callback returned unsupported status {result.status!r}") + if result.status == "failed": + return await self._save_failed(run, batch, result.error) + if result.status == "complete": + await self._collect_batch(run, batch, processor) + return await self._advance(run) + + async def process_batch_result( + self, run_id: str, *, batch_id: str | None, external_id: str | None + ) -> DurableEvalResult: + if batch_id is None and external_id is None: + raise ValueError("process_batch_result requires batch_id or external_id") + run = await self._required(run_id, "run") + if run["status"] in ("completed", "failed"): + return await self.status(run_id) + batches = await self._batch_records(run) + by_batch = None + if batch_id is not None: + by_batch = next(((batch, processor) for batch, processor in batches if batch["id"] == batch_id), None) + + by_external = None + if external_id is not None: + by_external = next( + ((batch, processor) for batch, processor in batches if batch.get("external_id") == external_id), + None, + ) + + if by_batch is not None and by_external is not None and by_batch[0]["id"] != by_external[0]["id"]: + raise ValueError("batch_id and external_id identify different batches") + match = by_batch or by_external + if match is None: + raise ValueError("No submitted durable batch matches this result") + batch, processor = match + if batch["status"] != "complete": + await self._collect_batch(run, batch, processor) + return await self._advance(run) + + def _batch_id(self, run_id: str, kind: str, index: int, item_ids: Sequence[str]) -> str: + return f"batch-{_stable_hex(run_id, kind, str(index), *item_ids, length=32)}" + + async def _task_complete(self, run: Mapping[str, Any]) -> bool: + for item_id in run["case_ids"]: + if await self._read(run["run_id"], "task-result", item_id) is None: + return False + return True + + async def _batch_specs(self, run: Mapping[str, Any]) -> list[tuple[dict[str, Any], Any]]: + run_id = run["run_id"] + item_ids = run["case_ids"] + specs: list[tuple[dict[str, Any], Any]] = [] + if isinstance(self.config.task, BatchTask): + for index, offset in enumerate(range(0, len(item_ids), self.config.task.batch_size)): + ids = item_ids[offset : offset + self.config.task.batch_size] + specs.append( + ( + { + "id": self._batch_id(run_id, "task", index, ids), + "kind": "task", + "item_ids": ids, + }, + self.config.task, + ) + ) + if await self._task_complete(run): + for scorer in self.config.scores: + if not isinstance(scorer, BatchScorer): + continue + for index, offset in enumerate(range(0, len(item_ids), scorer.batch_size)): + ids = item_ids[offset : offset + scorer.batch_size] + specs.append( + ( + { + "id": self._batch_id(run_id, f"score:{scorer.name}", index, ids), + "kind": "score", + "scorer_name": scorer.name, + "item_ids": ids, + }, + scorer, + ) + ) + return specs + + async def _batch_records(self, run: Mapping[str, Any]) -> list[tuple[dict[str, Any], Any]]: + records: list[tuple[dict[str, Any], Any]] = [] + for spec, processor in await self._batch_specs(run): + record = await self._read(run["run_id"], "batch", spec["id"]) + if record is not None: + records.append((record, processor)) + return records + + async def _case(self, run_id: str, item_id: str) -> dict[str, Any]: + return cast(dict[str, Any], await self._required(run_id, "case", item_id)) + + async def _task_item( + self, run: Mapping[str, Any], item_id: str, parameters: ValidatedParameters | None + ) -> BatchTaskItem[Any, Any]: + case = await self._case(run["run_id"], item_id) + datum = case["datum"] + return BatchTaskItem( + id=item_id, + input=datum["input"], + expected=datum.get("expected"), + metadata=dict(case["metadata"]), + tags=case.get("tags"), + parameters=parameters, + trial_index=case["trial_index"], + ) + + async def _scorer_item(self, run: Mapping[str, Any], item_id: str) -> BatchScorerItem[Any, Any, Any]: + case = await self._case(run["run_id"], item_id) + task = await self._required(run["run_id"], "task-result", item_id) + datum = case["datum"] + return BatchScorerItem( + id=item_id, + input=datum["input"], + output=task["output"], + expected=datum.get("expected"), + metadata=dict(task["metadata"]), + tags=task.get("tags"), + trial_index=case["trial_index"], + ) + + async def _submit_batch(self, run: Mapping[str, Any], spec: dict[str, Any], processor: Any) -> None: + run_id = run["run_id"] + if await self._read(run_id, "batch", spec["id"]) is not None: + return + if not await self._claim(run_id, f"submit:{spec['id']}"): + return + parameters = self._parameters(run) + if spec["kind"] == "task": + items = [await self._task_item(run, item_id, parameters) for item_id in spec["item_ids"]] + else: + items = [await self._scorer_item(run, item_id) for item_id in spec["item_ids"]] + context = BatchContext(run_id=run_id, batch_id=spec["id"]) + submission_data = await self._call(processor.submit, items, context) + submission_data = _json_value(submission_data) + completion = processor.completion + external_id = None + if isinstance(completion, BatchCompletionWebhook): + external_id = await self._call(completion.get_external_id, submission_data, context) + if not isinstance(external_id, str) or not external_id: + raise ValueError("Batch webhook get_external_id must return a non-empty string") + await self._write( + run_id, + "batch", + { + **spec, + "submission_data": submission_data, + "external_id": external_id, + "mode": completion.mode, + "status": "submitted", + }, + spec["id"], + ) + + async def _collect_batch(self, run: Mapping[str, Any], batch: dict[str, Any], processor: Any) -> None: + run_id = run["run_id"] + if batch["status"] == "complete": + return + context = BatchContext(run_id=run_id, batch_id=batch["id"]) + raw_results = await self._call(processor.collect, batch["submission_data"], context) + if not isinstance(raw_results, list): + raise TypeError("Batch collect callback must return a list") + result_type = BatchTaskResult if batch["kind"] == "task" else BatchScorerResult + results = [value if isinstance(value, result_type) else result_type(**value) for value in raw_results] + result_ids = [value.id for value in results] + expected_ids = batch["item_ids"] + if len(result_ids) != len(set(result_ids)): + raise ValueError(f"Batch {batch['id']} returned duplicate item IDs") + unknown = sorted(set(result_ids) - set(expected_ids)) + missing = sorted(set(expected_ids) - set(result_ids)) + if unknown or missing: + raise ValueError( + f"Batch {batch['id']} result IDs did not match submitted items; unknown={unknown}, missing={missing}" + ) + + if batch["kind"] == "task": + for result in results: + case = await self._case(run_id, result.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}, + result.id, + ) + else: + for result in results: + await self._write( + run_id, + _stage_kind("score-result", batch["scorer_name"]), + _json_value(result.score), + result.id, + ) + batch["status"] = "complete" + await self._write(run_id, "batch", batch, batch["id"]) + + async def _waiting_status(self, run: Mapping[str, Any]) -> DurableEvalWaitingResult: + pending = {"poll": 0, "webhook": 0} + for batch, _ in await self._batch_records(run): + if batch["status"] == "submitted": + pending[batch["mode"]] += 1 + return DurableEvalWaitingResult( + run_id=run["run_id"], pending=DurableEvalPending(poll=pending["poll"], webhook=pending["webhook"]) + ) + + def _span_ids(self, run_id: str, item_id: str, stage: str) -> tuple[str, str, str]: + if BraintrustEnv.LEGACY_IDS: + root_span_id = _stable_uuid(run_id, item_id, "root") + span_id = root_span_id if stage == "root" else _stable_uuid(run_id, item_id, stage) + return _stable_uuid(run_id, item_id, f"row:{stage}"), span_id, root_span_id + root_span_id = _stable_hex(run_id, item_id, "root", length=32) + span_id = root_span_id[:16] if stage == "root" else _stable_hex(run_id, item_id, stage, length=16) + row_id = _stable_uuid(run_id, item_id, stage) + return row_id, span_id, root_span_id + + def _start_root(self, experiment: Experiment | None, run: Mapping[str, Any], case: Mapping[str, Any]) -> Span: + if experiment is None: + return NOOP_SPAN + datum = case["datum"] + event_dataset = experiment.dataset or (self.config.data if isinstance(self.config.data, Dataset) else None) + if ( + event_dataset is not None + and isinstance(datum.get("id"), str) + and datum["id"] + and isinstance(datum.get("_xact_id"), str) + and datum["_xact_id"] + ): + origin = { + "object_type": "dataset", + "object_id": event_dataset.id, + "id": datum["id"], + "_xact_id": datum["_xact_id"], + **({"created": datum["created"]} if isinstance(datum.get("created"), str) else {}), + } + else: + origin = _validated_object_reference(datum.get("origin")) + row_id, span_id, root_span_id = self._span_ids(run["run_id"], case["id"], "root") + return _internal_start_span_with_initial_merge( + "eval", + parent=experiment.export(), + span_id=span_id, + root_span_id=root_span_id, + state=self.config.state, + type=SpanTypeAttribute.EVAL, + id=row_id, + input=datum["input"], + expected=datum.get("expected"), + metadata={ + **case["metadata"], + "durable_eval": { + "run_id": run["run_id"], + "case_id": case["case_id"], + "trial_index": case["trial_index"], + }, + }, + tags=case.get("tags"), + **({"origin": origin} if origin is not None else {}), + ) + + def _start_child( + self, + root: Span, + run_id: str, + item_id: str, + stage: str, + name: str, + span_type: SpanTypeAttribute, + **event: Any, + ) -> Span: + row_id, span_id, _ = self._span_ids(run_id, item_id, stage) + span_attributes: dict[str, Any] = {"type": span_type} + if span_type != SpanTypeAttribute.TASK: + span_attributes["purpose"] = "scorer" + return root.start_span( + name, + span_attributes=span_attributes, + internal={"initial_span_write_as_merge": True, "span_id": span_id}, + id=row_id, + **event, + ) + + async def _run_ordinary_task( + self, + run: Mapping[str, Any], + case: dict[str, Any], + experiment: Experiment | None, + parameters: ValidatedParameters | None, + ) -> None: + run_id = run["run_id"] + item_id = case["id"] + if await self._read(run_id, "task-result", item_id) is not None: + await self._log_task(run, case, experiment) + return + if not await self._claim(run_id, f"task:{item_id}"): + return + datum = case["datum"] + metadata = dict(case["metadata"]) + hooks = DictEvalHooks( + metadata, + expected=datum.get("expected"), + trial_index=case["trial_index"], + tags=case.get("tags"), + parameters=parameters, + ) + root = self._start_root(experiment, run, case) + task = cast(Callable[..., Any], self.config.task) + task_args: list[Any] = [datum["input"]] + try: + if len(get_signature(task).parameters) == 2: + task_args.append(hooks) + except Exception: + pass + with root: + with self._start_child( + root, + run_id, + item_id, + "task", + "task", + SpanTypeAttribute.TASK, + input=datum["input"], + ) as span: + hooks.set_span(span) + output = await self._call(task, *task_args) + span.log(output=output) + tags = list(hooks.tags) if hooks.tags else None + task_result = { + "output": _json_value(output), + "metadata": _json_value(metadata), + "tags": tags, + } + await self._write(run_id, "task-result", task_result, item_id) + root.log(output=output, metadata=metadata, tags=tags) + if root is not NOOP_SPAN: + await self._flush_logs() + await self._write( + run_id, + "task-log", + {"root_span": root.export() if experiment is not None else None}, + item_id, + ) + + async def _log_task(self, run: Mapping[str, Any], case: dict[str, Any], experiment: Experiment | None) -> None: + run_id = run["run_id"] + item_id = case["id"] + if await self._read(run_id, "task-log", item_id) is not None: + return + task_result = await self._required(run_id, "task-result", item_id) + root = self._start_root(experiment, run, case) + with root: + with self._start_child( + root, + run_id, + item_id, + "task", + "task", + 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 root is not NOOP_SPAN: + await self._flush_logs() + await self._write( + run_id, + "task-log", + {"root_span": root.export() if experiment is not None else None}, + item_id, + ) + + async def _root_for_case(self, run: Mapping[str, Any], item_id: str) -> Span: + task_log = await self._required(run["run_id"], "task-log", item_id) + exported = task_log.get("root_span") + return _internal_resume_span(exported, self.config.state) if exported else NOOP_SPAN + + async def _trace_for_case(self, run: Mapping[str, Any], item_id: str) -> LocalTrace | None: + task_log = await self._required(run["run_id"], "task-log", item_id) + exported = task_log.get("root_span") + if not exported: + return None + components = SpanComponentsV4.from_str(exported) + if not components.root_span_id: + raise ValueError("Persisted durable evaluation root span is missing its root span ID") + trace_state = self.config.state or _internal_get_global_state() + + async def ensure_spans_flushed() -> None: + await asyncio.get_running_loop().run_in_executor(None, trace_state.flush) + await trace_state.flush_otel() + + return LocalTrace( + object_type=span_object_type_v3_to_typed_string(components.object_type), + object_id=span_components_to_object_id(components), + root_span_id=components.root_span_id, + ensure_spans_flushed=ensure_spans_flushed, + state=trace_state, + ) + + def _prepare_scores(self, raw: Any, name: str) -> list[ScoreLike]: + if isinstance(raw, dict): + raw = _normalize_score(raw, "When returning a dict, it must be a valid Score object.") + if isinstance(raw, Iterable) and not isinstance(raw, (str, bytes, Mapping)): + return [ + _normalize_score(value, "When returning an array of scores, each score must be a valid Score object.") + for value in raw + ] + if is_score(raw): + return [raw] + return [Score(name=name, score=raw)] + + async def _run_ordinary_score(self, run: Mapping[str, Any], case: dict[str, Any], scorer: Any, name: str) -> None: + run_id = run["run_id"] + item_id = case["id"] + result_kind = _stage_kind("score-result", name) + if await self._read(run_id, result_kind, item_id) is None: + if not await self._claim(run_id, f"score:{name}:{item_id}"): + return + task_result = await self._required(run_id, "task-result", item_id) + fn = scorer.eval_async if hasattr(scorer, "eval_async") else scorer + trace = await self._trace_for_case(run, item_id) + raw = await call_user_fn( + asyncio.get_running_loop(), + fn, + **_scorer_args(case, task_result), + trace=trace, + ) + await self._write(run_id, result_kind, _json_value(raw), item_id) + await self._log_score(run, case, name) + + async def _log_score(self, run: Mapping[str, Any], case: dict[str, Any], name: 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 + raw = await self._required(run_id, _stage_kind("score-result", name), item_id) + results = self._prepare_scores(raw, name) + 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: + output = ( + {result.name: _score_fields(result) for result in results} + if len(results) != 1 + else _score_fields(results[0]) + ) + scores = {result.name: result.score for result in results} + span.log(output=output, metadata=_build_span_metadata(results), scores=scores) + root.log(scores=scores) + 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"] + result_kind = _stage_kind("classification-result", name) + if await self._read(run_id, result_kind, item_id) is None: + if not await self._claim(run_id, f"classification:{name}:{item_id}"): + return + task_result = await self._required(run_id, "task-result", item_id) + raw = await call_user_fn( + asyncio.get_running_loop(), + classifier, + **_scorer_args(case, task_result), + trace=None, + ) + if raw is None: + values: list[Any] = [] + elif isinstance(raw, Iterable) and not isinstance(raw, (str, bytes, Mapping)): + values = list(raw) + else: + values = [raw] + classifications = [_validate_classification_result(value, name) for value in values] + await self._write(run_id, result_kind, [value.as_dict() for value in classifications], item_id) + await self._log_classifier(run, case, name) + + async def _log_classifier(self, run: Mapping[str, Any], case: dict[str, Any], name: str) -> None: + run_id = run["run_id"] + item_id = case["id"] + log_kind = _stage_kind("classification-log", name) + if await self._read(run_id, log_kind, item_id) is not None: + return + raw = await self._required(run_id, _stage_kind("classification-result", name), item_id) + classifications = [Classification.from_dict(value) for value in raw] + task_result = await self._required(run_id, "task-result", item_id) + root = await self._root_for_case(run, item_id) + with root: + with self._start_child( + root, + run_id, + item_id, + f"classification:{name}", + name, + SpanTypeAttribute.CLASSIFIER, + input=_scorer_args(case, task_result), + ) as span: + if classifications: + span.log( + output=_build_classification_span_output(classifications), + metadata=_build_span_metadata(classifications), + ) + grouped: dict[str, list[Any]] = {} + for result in classifications: + grouped.setdefault(cast(str, result.name), []).append(result.as_item()) + root.log(classifications=grouped) + else: + span.log(output={}, metadata=None) + if root is not NOOP_SPAN: + await self._flush_logs() + await self._write(run_id, log_kind, True, item_id) + + async def _advance_tasks( + self, + run: Mapping[str, Any], + experiment: Experiment | None, + parameters: ValidatedParameters | None, + ) -> tuple[Experiment | None, bool]: + run_id = run["run_id"] + if isinstance(self.config.task, BatchTask): + for spec, processor in await self._batch_specs(run): + if spec["kind"] == "task": + await self._submit_batch(run, spec, processor) + else: + if experiment is None and not run["no_send_logs"]: + experiment = self._experiment(run) + for item_id in run["case_ids"]: + case = await self._case(run_id, item_id) + await self._run_ordinary_task(run, case, experiment, parameters) + + if not await self._task_complete(run): + return experiment, False + + if experiment is None and not run["no_send_logs"]: + experiment = self._experiment(run) + if isinstance(self.config.task, BatchTask): + for item_id in run["case_ids"]: + await self._log_task(run, await self._case(run_id, item_id), experiment) + for item_id in run["case_ids"]: + if await self._read(run_id, "task-log", item_id) is None: + return experiment, False + return experiment, True + + async def _advance_scorers( + self, + run: Mapping[str, Any], + scorers: Sequence[Any], + scorer_names: Sequence[str], + ) -> None: + run_id = run["run_id"] + for spec, processor in await self._batch_specs(run): + if spec["kind"] == "score": + await self._submit_batch(run, spec, processor) + + for scorer, name in zip(scorers, scorer_names): + for item_id in run["case_ids"]: + if isinstance(scorer, BatchScorer): + if await self._read(run_id, _stage_kind("score-result", name), item_id) is not None: + await self._log_score(run, await self._case(run_id, item_id), name) + else: + await self._run_ordinary_score(run, await self._case(run_id, item_id), scorer, name) + + async def _advance_classifiers( + self, + run: Mapping[str, Any], + classifiers: Sequence[EvalClassifier[Input, Output, Expected]], + classifier_names: Sequence[str], + ) -> None: + run_id = run["run_id"] + for classifier, name in zip(classifiers, classifier_names): + for item_id in run["case_ids"]: + await self._run_classifier(run, await self._case(run_id, item_id), classifier, name) + + async def _stage_records_complete(self, run: Mapping[str, Any], prefix: str, names: Sequence[str]) -> bool: + complete = True + for name in names: + result_kind = _stage_kind(f"{prefix}-result", name) + log_kind = _stage_kind(f"{prefix}-log", name) + for item_id in run["case_ids"]: + if ( + await self._read(run["run_id"], result_kind, item_id) is None + or await self._read(run["run_id"], log_kind, item_id) is None + ): + complete = False + return complete + + async def _advance(self, run: Mapping[str, Any], *, experiment: Experiment | None = None) -> DurableEvalResult: + run = await self._required(run["run_id"], "run") + run_id = run["run_id"] + if run["status"] in ("completed", "failed"): + return await self.status(run_id) + experiment, tasks_complete = await self._advance_tasks(run, experiment, self._parameters(run)) + if not tasks_complete: + return await self._waiting_status(run) + + resolved_scores = [ + score() if inspect.isclass(score) and is_scorer(score) else score for score in self.config.scores + ] + scorer_names = [ + score.name if isinstance(score, BatchScorer) else _scorer_name(score, index) + for index, score in enumerate(resolved_scores) + ] + if len(scorer_names) != len(set(scorer_names)): + raise ValueError("Durable evaluation scorer names must be unique") + + await self._advance_scorers(run, resolved_scores, scorer_names) + + classifiers = list(self.config.classifiers) + classifier_names = [_classifier_name(classifier, index) for index, classifier in enumerate(classifiers)] + if len(classifier_names) != len(set(classifier_names)): + raise ValueError("Durable evaluation classifier names must be unique") + await self._advance_classifiers(run, classifiers, classifier_names) + + scores_complete = await self._stage_records_complete(run, "score", scorer_names) + classifiers_complete = await self._stage_records_complete(run, "classification", classifier_names) + if not scores_complete or not classifiers_complete: + return await self._waiting_status(run) + + if not await self._claim(run_id, "finalize"): + current = await self._required(run_id, "run") + if current["status"] == "completed": + return await self.status(run_id) + return await self._waiting_status(current) + summary = await self._summary(run, scorer_names, classifier_names, experiment) + return await self._save_completed(run, summary) + + async def _summary( + self, + run: Mapping[str, Any], + scorer_names: Sequence[str], + classifier_names: Sequence[str], + experiment: Experiment | None, + ) -> ExperimentSummary: + if experiment is not None: + comparison_experiment_id = self.config.base_experiment_id + if comparison_experiment_id is None: + comparison_experiment_id = _get_persisted_base_experiment_id(experiment) + return experiment.summarize( + summarize_scores=self.config.summarize_scores, + comparison_experiment_id=comparison_experiment_id, + ) + + results: list[EvalResult[Any, Any, Any]] = [] + for item_id in run["case_ids"]: + case = await self._case(run["run_id"], item_id) + 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) + for score in self._prepare_scores(raw, name): + scores[score.name] = score.score + classifications: dict[str, list[Any]] = {} + for name in classifier_names: + raw = await self._required(run["run_id"], _stage_kind("classification-result", name), item_id) + for value in raw: + classification = Classification.from_dict(value) + classifications.setdefault(cast(str, classification.name), []).append(classification.as_item()) + datum = case["datum"] + results.append( + EvalResult( + input=datum["input"], + output=task_result["output"], + scores=scores, + classifications=classifications or None, + expected=datum.get("expected"), + metadata=task_result["metadata"], + tags=task_result.get("tags"), + ) + ) + evaluator = Evaluator( + project_name=self.config.project_name, + eval_name=self.eval_name, + data=[], + task=cast(Any, lambda value: value), + scores=[], + experiment_name=run["experiment_name"], + metadata=self.config.metadata, + ) + return build_local_summary(evaluator, cast(Any, results)) + + +__all__ = [ + "BatchCompletionPoll", + "BatchCompletionWebhook", + "BatchContext", + "BatchPollResult", + "BatchScorer", + "BatchScorerItem", + "BatchScorerResult", + "BatchTask", + "BatchTaskItem", + "BatchTaskResult", + "DurableEval", + "DurableEvalCompletedResult", + "DurableEvalFailedResult", + "DurableEvalMemoryStore", + "DurableEvalPending", + "DurableEvalRedisStore", + "DurableEvalResult", + "DurableEvalStore", + "DurableEvalStoreEntry", + "DurableEvalWaitingResult", + "define_durable_eval", +] diff --git a/py/src/braintrust/logger.py b/py/src/braintrust/logger.py index 70cf41b6..bd0c300a 100644 --- a/py/src/braintrust/logger.py +++ b/py/src/braintrust/logger.py @@ -173,6 +173,18 @@ class SpanInternalOptions(TypedDict, total=False): `context.span_origin.instrumentation.name`. Set by SDK integrations (`openai-auto`, `anthropic-auto`, etc.).""" + initial_span_write_as_merge: bool + """Emit the initial row as a merge for an SDK-owned deterministic span.""" + + skip_initial_span_write: bool + """Rehydrate an existing span without emitting an initial row.""" + + span_id: str + """SDK-controlled span ID.""" + + root_span_id: str + """SDK-controlled root span ID.""" + T = TypeVar("T") TMapping = TypeVar("TMapping", bound=Mapping[str, Any]) @@ -3091,6 +3103,33 @@ def flush(): _state.global_bg_logger().flush() +def _internal_start_span_with_initial_merge( + name: str, + *, + parent: str, + span_id: str, + root_span_id: str, + state: BraintrustState | None = None, + type: SpanTypeAttribute | None = None, + span_attributes: SpanAttributes | Mapping[str, Any] | None = None, + **event: Any, +) -> Span: + """Start a deterministic SDK-owned span whose first row is merge-safe.""" + return start_span( + name=name, + parent=parent, + state=state, + type=type, + span_attributes=span_attributes, + internal={ + "initial_span_write_as_merge": True, + "span_id": span_id, + "root_span_id": root_span_id, + }, + **event, + ) + + def _check_org_info(state, org_info, org_name): if len(org_info) == 0: raise ValueError("This user is not part of any organizations.") @@ -3986,6 +4025,26 @@ def update_span(exported: str, **event: Any) -> None: ) +def _internal_resume_span(exported: str, state: BraintrustState | None = None) -> Span: + """Rehydrate an exported root span so SDK work can continue later.""" + components = SpanComponentsV4.from_str(exported) + if not components.row_id or not components.span_id or not components.root_span_id: + raise ValueError("Only exported root spans can be resumed") + return SpanImpl( + parent_object_type=components.object_type, + parent_object_id=LazyValue(_span_components_to_object_id_lambda(components), use_mutex=False), + parent_compute_object_metadata_args=components.compute_object_metadata_args, + parent_span_ids=None, + event={"id": components.row_id}, + propagated_event=components.propagated_event, + span_id=components.span_id, + root_span_id=components.root_span_id, + state=state, + lookup_span_parent=False, + internal={"skip_initial_span_write": True}, + ) + + @dataclasses.dataclass class ParentSpanIds: span_id: str @@ -4744,7 +4803,8 @@ def __init__( span_attributes=dict(**{"type": type, "name": name, **span_attributes}, exec_counter=exec_counter), created=datetime.datetime.now(datetime.timezone.utc).isoformat(), ) - self._instrumentation = (internal or {}).get("instrumentation") or "braintrust-python-logger" + internal = internal or {} + self._instrumentation = internal.get("instrumentation") or "braintrust-python-logger" internal_data["context"] = merge_span_origin_context( caller_location or {}, self._instrumentation, @@ -4759,8 +4819,8 @@ def __init__( # Resolve all span IDs (span_id, root_span_id, span_parents) span_ids = _resolve_span_ids( - span_id=span_id, - root_span_id=root_span_id, + span_id=internal.get("span_id", span_id), + root_span_id=internal.get("root_span_id", root_span_id), parent_span_ids=parent_span_ids, lookup_span_parent=lookup_span_parent, id_generator=self.state.id_generator, @@ -4772,8 +4832,9 @@ def __init__( # The first log is a replacement, but subsequent logs to the same span # object will be merges. - self._is_merge = False - self.log_internal(event=event, internal_data=internal_data) + self._is_merge = internal.get("initial_span_write_as_merge", False) + if not internal.get("skip_initial_span_write", False): + self.log_internal(event=event, internal_data=internal_data) self._is_merge = True @property diff --git a/py/src/braintrust/test_durable_eval.py b/py/src/braintrust/test_durable_eval.py new file mode 100644 index 00000000..fd54326f --- /dev/null +++ b/py/src/braintrust/test_durable_eval.py @@ -0,0 +1,600 @@ +"""Tests for the experimental durable eval API.""" + +import asyncio +from unittest.mock import patch + +import pytest + +from . import durable_eval as durable_eval_module +from .durable_eval import ( + BatchCompletionPoll, + BatchCompletionWebhook, + BatchPollResult, + BatchScorer, + BatchScorerResult, + BatchTask, + BatchTaskResult, + DurableEvalCompletedResult, + DurableEvalFailedResult, + DurableEvalMemoryStore, + DurableEvalRedisStore, + DurableEvalWaitingResult, + define_durable_eval, +) +from .logger import BraintrustState, Dataset, ObjectMetadata, ProjectDatasetMetadata +from .test_helpers import init_test_exp, with_memory_logger, with_simulate_login # noqa: F401 +from .util import LazyValue + + +@pytest.mark.asyncio +async def test_memory_store_is_atomic_and_copies_values(): + store = DurableEvalMemoryStore() + value = bytearray(b"first") + + first = await store.get_or_set("key", value) + value[:] = b"other" + second = await store.get_or_set("key", b"second") + + assert first.created is True + assert first.value == b"first" + assert second.created is False + assert second.value == b"first" + assert await store.read("key") == b"first" + + +class _SyncRedis: + def __init__(self): + self.values = {} + self.calls = [] + + def get(self, key): + return self.values.get(key) + + def set(self, key, value, **kwargs): + self.calls.append((key, kwargs)) + old = self.values.get(key) + if kwargs.get("nx") and old is not None: + return old if kwargs.get("get") else None + self.values[key] = value + return old if kwargs.get("get") else True + + +class _AsyncRedis(_SyncRedis): + async def get(self, key): + return super().get(key) + + async def set(self, key, value, **kwargs): + return super().set(key, value, **kwargs) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("client_type", [_SyncRedis, _AsyncRedis]) +async def test_redis_store_supports_sync_and_async_redis_py(client_type): + client = client_type() + store = DurableEvalRedisStore(client, key_prefix="test:", ttl_ms=1234) + + await store.write("one", b"value") + assert await store.read("one") == b"value" + assert (await store.get_or_set("two", b"first")).created is True + existing = await store.get_or_set("two", b"second") + + assert existing.created is False + 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) + + +@pytest.mark.asyncio +async def test_ordinary_durable_eval_completes_locally_and_is_idempotent(): + task_calls = [] + score_calls = [] + + def task(input, hooks): + task_calls.append((input, hooks.trial_index)) + hooks.metadata["task"] = True + hooks.tags = ["updated"] + return input * 2 + + def scorer(output, expected, metadata, tags): + score_calls.append((output, metadata, tags)) + return output == expected + + durable_eval = define_durable_eval( + "project", + store=DurableEvalMemoryStore(), + data=[{"id": "case", "input": 2, "expected": 4, "metadata": {"source": "test"}}], + task=task, + scores=[scorer], + trial_count=2, + ) + + result = await durable_eval.start(no_send_logs=True) + repeated = await durable_eval.status(result.run_id) + + assert isinstance(result, DurableEvalCompletedResult) + assert isinstance(repeated, DurableEvalCompletedResult) + assert result.summary.scores["scorer"].score == 1 + assert task_calls == [(2, 0), (2, 1)] + assert score_calls == [ + (4, {"source": "test", "task": True}, ["updated"]), + (4, {"source": "test", "task": True}, ["updated"]), + ] + + +@pytest.mark.asyncio +async def test_ordinary_only_runs_do_not_require_case_ids_and_get_new_run_ids(): + durable_eval = define_durable_eval( + "project", + store=DurableEvalMemoryStore(), + data=[{"input": 1}], + task=lambda value: value, + scores=[lambda: 1], + ) + + first = await durable_eval.start(no_send_logs=True) + second = await durable_eval.start(no_send_logs=True) + + assert first.status == "completed" + assert second.status == "completed" + assert first.run_id != second.run_id + + +@pytest.mark.asyncio +async def test_poll_advances_task_then_mixed_scorers_without_polling_new_batches(): + submitted = {} + ready = set() + calls = [] + + async def task_submit(items, context): + calls.append(("submit-task", context.batch_id)) + submitted[context.batch_id] = items + return {"provider_id": context.batch_id} + + async def task_collect(_submission, context): + calls.append(("collect-task", context.batch_id)) + return [BatchTaskResult(id=item.id, output=item.input * 2) for item in submitted[context.batch_id]] + + async def score_submit(items, context): + calls.append(("submit-score", context.batch_id)) + submitted[context.batch_id] = items + return {"provider_id": context.batch_id} + + async def score_collect(_submission, context): + calls.append(("collect-score", context.batch_id)) + return [ + BatchScorerResult(id=item.id, score=item.output == item.expected) for item in submitted[context.batch_id] + ] + + async def poll(_submission, context): + calls.append(("poll", context.batch_id)) + return BatchPollResult("complete" if context.batch_id in ready else "pending") + + task = BatchTask( + submit=task_submit, + completion=BatchCompletionPoll(poll), + collect=task_collect, + batch_size=2, + ) + batch_score = BatchScorer( + name="batch", + submit=score_submit, + completion=BatchCompletionPoll(poll), + collect=score_collect, + batch_size=2, + ) + durable_eval = define_durable_eval( + "project", + store=DurableEvalMemoryStore(), + data=[{"id": str(i), "input": i, "expected": i * 2} for i in range(3)], + task=task, + scores=[lambda output, expected: output == expected, batch_score], + ) + + started = await durable_eval.start(no_send_logs=True) + assert isinstance(started, DurableEvalWaitingResult) + assert started.pending.poll == 2 + task_batch_ids = set(submitted) + ready.update(task_batch_ids) + + after_tasks = await durable_eval.poll(started.run_id) + assert isinstance(after_tasks, DurableEvalWaitingResult) + assert after_tasks.pending.poll == 2 + score_batch_ids = set(submitted) - task_batch_ids + assert score_batch_ids + assert not any(call == ("poll", batch_id) for batch_id in score_batch_ids for call in calls) + + ready.update(score_batch_ids) + completed = await durable_eval.poll(started.run_id) + assert isinstance(completed, DurableEvalCompletedResult) + assert completed.summary.scores["scorer_0"].score == 1 + assert completed.summary.scores["batch"].score == 1 + calls_before_status = list(calls) + assert (await durable_eval.status(started.run_id)).status == "completed" + assert calls == calls_before_status + + +@pytest.mark.asyncio +async def test_failed_poll_is_persisted_as_a_terminal_result(): + poll_calls = 0 + + async def poll(_submission, _context): + nonlocal poll_calls + poll_calls += 1 + return BatchPollResult("failed", error={"code": "provider_failed"}) + + durable_eval = define_durable_eval( + "project", + store=DurableEvalMemoryStore(), + data=[{"id": "case", "input": 1}], + task=BatchTask( + submit=lambda _items, context: {"id": context.batch_id}, + completion=BatchCompletionPoll(poll), + collect=lambda _submission, _context: pytest.fail("failed batches must not be collected"), + ), + ) + + started = await durable_eval.start(no_send_logs=True) + failed = await durable_eval.poll(started.run_id) + repeated = await durable_eval.status(started.run_id) + + assert isinstance(failed, DurableEvalFailedResult) + assert repeated == failed + assert failed.error == {"code": "provider_failed"} + assert failed.pending.poll == 0 + assert failed.pending.webhook == 0 + assert poll_calls == 1 + + +@pytest.mark.asyncio +async def test_webhook_result_can_be_matched_by_external_id(): + submitted = {} + + async def submit(items, context): + submitted[context.batch_id] = items + return {"id": f"external-{context.batch_id}"} + + async def collect(_submission, context): + return [BatchTaskResult(id=item.id, output=item.input) for item in submitted[context.batch_id]] + + durable_eval = define_durable_eval( + "project", + store=DurableEvalMemoryStore(), + data=[{"id": "case", "input": "ok"}], + task=BatchTask( + submit=submit, + completion=BatchCompletionWebhook(lambda submission, _context: submission["id"]), + collect=collect, + ), + ) + started = await durable_eval.start(no_send_logs=True) + batch_id = next(iter(submitted)) + + completed = await durable_eval.process_batch_result(started.run_id, external_id=f"external-{batch_id}") + + assert isinstance(completed, DurableEvalCompletedResult) + assert (await durable_eval.process_batch_result(started.run_id, batch_id=batch_id)).status == "completed" + + +@pytest.mark.asyncio +async def test_collect_requires_exactly_one_result_per_item(): + batch_ids = [] + + async def submit(_items, context): + batch_ids.append(context.batch_id) + return {"id": context.batch_id} + + durable_eval = define_durable_eval( + "project", + store=DurableEvalMemoryStore(), + data=[{"id": "a", "input": 1}, {"id": "b", "input": 2}], + task=BatchTask( + submit=submit, + completion=BatchCompletionWebhook(lambda submission, _context: submission["id"]), + collect=lambda _submission, _context: [BatchTaskResult(id="a:trial:0", output=1)], + ), + ) + started = await durable_eval.start(no_send_logs=True) + + with pytest.raises(ValueError, match="missing=.*b:trial:0"): + await durable_eval.process_batch_result(started.run_id, batch_id=batch_ids[0]) + + +@pytest.mark.asyncio +async def test_case_ids_must_be_stable_and_unique(): + batch_task = BatchTask( + submit=lambda _items, _context: {"id": "unused"}, + completion=BatchCompletionPoll(lambda _submission, _context: BatchPollResult("pending")), + collect=lambda _submission, _context: [], + ) + missing = define_durable_eval("project", store=DurableEvalMemoryStore(), data=[{"input": 1}], task=batch_task) + with pytest.raises(ValueError, match="non-empty id"): + await missing.start(no_send_logs=True) + + duplicate = define_durable_eval( + "project", + store=DurableEvalMemoryStore(), + data=[{"id": "same", "input": 1}, {"id": "same", "input": 2}], + task=batch_task, + ) + with pytest.raises(ValueError, match="duplicate"): + await duplicate.start(no_send_logs=True) + + +@pytest.mark.asyncio +async def test_case_persistence_does_not_deepcopy_inputs(): + class SerializableWithoutDeepcopy: + def __deepcopy__(self, _memo): + raise AssertionError("input was deep-copied") + + def model_dump(self, **_kwargs): + return {"value": "serialized"} + + submitted = [] + + async def submit(items, _context): + submitted.extend(items) + return {"id": "pending"} + + durable_eval = define_durable_eval( + "project", + store=DurableEvalMemoryStore(), + data=[{"id": "case", "input": SerializableWithoutDeepcopy()}], + task=BatchTask( + submit=submit, + completion=BatchCompletionPoll(lambda _submission, _context: BatchPollResult("pending")), + collect=lambda _submission, _context: [], + ), + ) + + result = await durable_eval.start(no_send_logs=True) + + assert result.status == "waiting" + assert submitted[0].input == {"value": "serialized"} + + +@pytest.mark.asyncio +async def test_concurrent_webhooks_claim_downstream_submission_once(): + task_items = [] + collects_started = 0 + both_collecting = asyncio.Event() + score_submissions = 0 + + async def submit_task(items, _context): + task_items.extend(items) + return {"id": "task-provider"} + + async def collect_task(_submission, _context): + nonlocal collects_started + collects_started += 1 + if collects_started == 2: + both_collecting.set() + await both_collecting.wait() + return [BatchTaskResult(id=item.id, output=item.input * 2) for item in task_items] + + async def submit_score(_items, _context): + nonlocal score_submissions + score_submissions += 1 + return {"id": "score-provider"} + + durable_eval = define_durable_eval( + "project", + store=DurableEvalMemoryStore(), + data=[{"id": "case", "input": 2, "expected": 4}], + task=BatchTask( + submit=submit_task, + completion=BatchCompletionWebhook(lambda submission, _context: submission["id"]), + collect=collect_task, + ), + scores=[ + BatchScorer( + name="exact", + submit=submit_score, + completion=BatchCompletionWebhook(lambda submission, _context: submission["id"]), + collect=lambda _submission, _context: [], + ) + ], + ) + started = await durable_eval.start(no_send_logs=True) + + await asyncio.gather( + durable_eval.process_batch_result(started.run_id, external_id="task-provider"), + durable_eval.process_batch_result(started.run_id, external_id="task-provider"), + ) + + assert score_submissions == 1 + current = await durable_eval.status(started.run_id) + assert current.status == "waiting" + assert current.pending.webhook == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("legacy_ids", [False, True]) +async def test_durable_logging_uses_stable_spans_and_resume_metadata( + monkeypatch, with_memory_logger, with_simulate_login, legacy_ids +): + if legacy_ids: + monkeypatch.setenv("BRAINTRUST_LEGACY_IDS", "true") + else: + monkeypatch.delenv("BRAINTRUST_LEGACY_IDS", raising=False) + local_eval = define_durable_eval( + "project", + store=DurableEvalMemoryStore(), + data=[{"id": "case", "input": 1, "expected": 2}], + task=lambda value: value + 1, + scores=[lambda output, expected: output == expected], + ) + local_summary = (await local_eval.start(no_send_logs=True)).summary + experiment = init_test_exp("durable", "project") + monkeypatch.setattr(durable_eval_module, "init_experiment", lambda **_kwargs: experiment) + monkeypatch.setattr(experiment, "summarize", lambda **_kwargs: local_summary) + submitted = [] + trace_configurations = [] + + async def submit(items, _context): + submitted.extend(items) + return {"id": "task"} + + def scorer(output, expected, trace): + trace_configurations.append(trace.get_configuration()) + return output == expected + + durable_eval = define_durable_eval( + "project", + store=DurableEvalMemoryStore(), + data=[{"id": "case", "input": 1, "expected": 2}], + task=BatchTask( + submit=submit, + completion=BatchCompletionPoll(lambda _submission, _context: BatchPollResult("complete")), + collect=lambda _submission, _context: [ + BatchTaskResult(id=item.id, output=item.input + 1) for item in submitted + ], + ), + scores=[scorer], + experiment_name="durable", + ) + + waiting = await durable_eval.start() + result = await durable_eval.poll(waiting.run_id) + logs = with_memory_logger.pop() + + assert result.status == "completed" + assert len(logs) == 3 + roots = [row for row in logs if not row["span_parents"]] + assert len(roots) == 1 + assert roots[0]["metadata"]["durable_eval"] == { + "run_id": result.run_id, + "case_id": "case", + "trial_index": 0, + } + assert trace_configurations == [ + { + "object_type": "experiment", + "object_id": experiment.id, + "root_span_id": roots[0]["root_span_id"], + } + ] + assert len({row["span_id"] for row in logs}) == 3 + await durable_eval.status(result.run_id) + assert with_memory_logger.pop() == [] + + +@pytest.mark.asyncio +async def test_durable_logging_flushes_before_persisting_log_markers( + monkeypatch, with_memory_logger, with_simulate_login +): + local_eval = define_durable_eval( + "project", + store=DurableEvalMemoryStore(), + data=[{"id": "case", "input": 1}], + task=lambda value: value, + ) + local_summary = (await local_eval.start(no_send_logs=True)).summary + experiment = init_test_exp("durable", "project") + monkeypatch.setattr(durable_eval_module, "init_experiment", lambda **_kwargs: experiment) + monkeypatch.setattr(experiment, "summarize", lambda **_kwargs: local_summary) + + flush_count = 0 + marker_flush_counts = [] + original_flush = with_memory_logger.flush + + def flush(*args, **kwargs): + nonlocal flush_count + flush_count += 1 + return original_flush(*args, **kwargs) + + monkeypatch.setattr(with_memory_logger, "flush", flush) + + class RecordingStore(DurableEvalMemoryStore): + async def write(self, key, value): + if "/task-log/" in key or "/score-log-" in key or "/classification-log-" in key: + marker_flush_counts.append(flush_count) + await super().write(key, value) + + submitted = [] + + async def submit(items, _context): + submitted.extend(items) + return {"id": "task"} + + durable_eval = define_durable_eval( + "project", + store=RecordingStore(), + data=[{"id": "case", "input": 1, "expected": 1}], + task=BatchTask( + submit=submit, + completion=BatchCompletionPoll(lambda _submission, _context: BatchPollResult("complete")), + collect=lambda _submission, _context: [ + BatchTaskResult(id=item.id, output=item.input) for item in submitted + ], + ), + scores=[lambda output, expected: output == expected], + classifiers=[lambda output: {"id": "positive"}], + experiment_name="durable", + ) + + waiting = await durable_eval.start() + completed = await durable_eval.poll(waiting.run_id) + + assert completed.status == "completed" + assert len(marker_flush_counts) == 3 + assert all(count > 0 for count in marker_flush_counts) + + +@pytest.mark.asyncio +async def test_durable_dataset_rows_preserve_dataset_origin(monkeypatch, with_memory_logger, with_simulate_login): + project_metadata = ObjectMetadata(id="test-project", name="test-project", full_info={}) + dataset_metadata = ObjectMetadata(id="active-dataset", name="test-dataset", full_info={}) + dataset = Dataset( + lazy_metadata=LazyValue( + lambda: ProjectDatasetMetadata(project=project_metadata, dataset=dataset_metadata), + use_mutex=False, + ), + state=BraintrustState(), + ) + row = { + "id": "dataset-row", + "_xact_id": "dataset-xact", + "created": "2026-06-02T00:00:00.000Z", + "input": 1, + } + local_summary = ( + await define_durable_eval( + "project", store=DurableEvalMemoryStore(), data=[{"input": 1}], task=lambda value: value + ).start(no_send_logs=True) + ).summary + experiment = init_test_exp("durable", "project") + monkeypatch.setattr(durable_eval_module, "init_experiment", lambda **_kwargs: experiment) + monkeypatch.setattr(experiment, "summarize", lambda **_kwargs: local_summary) + + submitted = [] + + async def submit(items, _context): + submitted.extend(items) + return {"id": "task"} + + durable_eval = define_durable_eval( + "project", + store=DurableEvalMemoryStore(), + data=dataset, + task=BatchTask( + submit=submit, + completion=BatchCompletionPoll(lambda _submission, _context: BatchPollResult("complete")), + collect=lambda _submission, _context: [ + BatchTaskResult(id=item.id, output=item.input) for item in submitted + ], + ), + experiment_name="durable", + ) + + with patch.object(dataset, "_refetch", return_value=[row]): + waiting = await durable_eval.start() + await durable_eval.poll(waiting.run_id) + + root = next(log for log in with_memory_logger.pop() if not log["span_parents"]) + assert root["origin"] == { + "object_type": "dataset", + "object_id": "active-dataset", + "id": "dataset-row", + "_xact_id": "dataset-xact", + "created": "2026-06-02T00:00:00.000Z", + } diff --git a/py/src/braintrust/type_tests/test_durable_eval.py b/py/src/braintrust/type_tests/test_durable_eval.py new file mode 100644 index 00000000..6afe409b --- /dev/null +++ b/py/src/braintrust/type_tests/test_durable_eval.py @@ -0,0 +1,81 @@ +"""Static and runtime type coverage for the experimental durable eval API.""" + +from typing import TypedDict + +import pytest +from braintrust import ( + BatchCompletionPoll, + BatchContext, + BatchPollResult, + BatchScorer, + BatchScorerItem, + BatchScorerResult, + BatchTask, + BatchTaskItem, + BatchTaskResult, + DurableEval, + DurableEvalFailedResult, + DurableEvalMemoryStore, + EvalCase, + define_durable_eval, +) + + +class Submission(TypedDict): + id: str + + +async def submit_task(items: list[BatchTaskItem[str, str]], context: BatchContext) -> Submission: + assert items + return {"id": context.batch_id} + + +async def collect_task(submission: Submission, context: BatchContext) -> list[BatchTaskResult[int]]: + return [BatchTaskResult(id="case:trial:0", output=len(submission["id"] + context.run_id))] + + +async def submit_score(items: list[BatchScorerItem[str, int, str]], context: BatchContext) -> Submission: + assert items + return {"id": context.batch_id} + + +async def collect_score(submission: Submission, context: BatchContext) -> list[BatchScorerResult]: + assert submission["id"] == context.batch_id + return [BatchScorerResult(id="case:trial:0", score=1)] + + +async def poll_batch(submission: Submission, context: BatchContext) -> BatchPollResult: + assert submission["id"] == context.batch_id + return BatchPollResult(status="pending") + + +task: BatchTask[str, int, str, Submission] = BatchTask( + submit=submit_task, + completion=BatchCompletionPoll(poll_batch), + collect=collect_task, +) +score: BatchScorer[str, int, str, Submission] = BatchScorer( + name="score", + submit=submit_score, + completion=BatchCompletionPoll(poll_batch), + collect=collect_score, +) +durable_eval: DurableEval[str, int, str] = define_durable_eval( + "project", + store=DurableEvalMemoryStore(), + data=[EvalCase(id="case", input="input", expected="expected")], + task=task, + scores=[score], +) + + +@pytest.mark.asyncio +async def test_durable_eval_types_at_runtime(): + result = await durable_eval.start(no_send_logs=True) + assert result.status == "waiting" + assert result.pending.poll == 1 + + +def consume_failed_result(result: DurableEvalFailedResult) -> object: + assert result.status == "failed" + return result.error From bb1283b8ad4fa8a49595091bf6b48fb87dc8b32c Mon Sep 17 00:00:00 2001 From: lforst <8118419+lforst@users.noreply.github.com> Date: Tue, 1 Sep 2026 15:49:18 +0000 Subject: [PATCH 2/3] Update PR #725 --- py/src/braintrust/durable_eval.py | 3 ++- py/src/braintrust/test_durable_eval.py | 11 +++++++++-- 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/py/src/braintrust/durable_eval.py b/py/src/braintrust/durable_eval.py index 95da3975..2b7de3e5 100644 --- a/py/src/braintrust/durable_eval.py +++ b/py/src/braintrust/durable_eval.py @@ -1205,11 +1205,12 @@ async def _run_classifier(self, run: Mapping[str, Any], case: dict[str, Any], cl if not await self._claim(run_id, f"classification:{name}:{item_id}"): return task_result = await self._required(run_id, "task-result", item_id) + trace = await self._trace_for_case(run, item_id) raw = await call_user_fn( asyncio.get_running_loop(), classifier, **_scorer_args(case, task_result), - trace=None, + trace=trace, ) if raw is None: values: list[Any] = [] diff --git a/py/src/braintrust/test_durable_eval.py b/py/src/braintrust/test_durable_eval.py index fd54326f..1cd38a4d 100644 --- a/py/src/braintrust/test_durable_eval.py +++ b/py/src/braintrust/test_durable_eval.py @@ -429,6 +429,7 @@ async def test_durable_logging_uses_stable_spans_and_resume_metadata( monkeypatch.setattr(experiment, "summarize", lambda **_kwargs: local_summary) submitted = [] trace_configurations = [] + classifier_trace_configurations = [] async def submit(items, _context): submitted.extend(items) @@ -438,6 +439,10 @@ def scorer(output, expected, trace): trace_configurations.append(trace.get_configuration()) return output == expected + def classifier(output, trace): + classifier_trace_configurations.append(trace.get_configuration()) + return {"id": "positive"} + durable_eval = define_durable_eval( "project", store=DurableEvalMemoryStore(), @@ -450,6 +455,7 @@ def scorer(output, expected, trace): ], ), scores=[scorer], + classifiers=[classifier], experiment_name="durable", ) @@ -458,7 +464,7 @@ def scorer(output, expected, trace): logs = with_memory_logger.pop() assert result.status == "completed" - assert len(logs) == 3 + assert len(logs) == 4 roots = [row for row in logs if not row["span_parents"]] assert len(roots) == 1 assert roots[0]["metadata"]["durable_eval"] == { @@ -473,7 +479,8 @@ def scorer(output, expected, trace): "root_span_id": roots[0]["root_span_id"], } ] - assert len({row["span_id"] for row in logs}) == 3 + assert classifier_trace_configurations == trace_configurations + assert len({row["span_id"] for row in logs}) == 4 await durable_eval.status(result.run_id) assert with_memory_logger.pop() == [] From 98505a04cffaf2ba0db1f7a4f4428679f52d860a Mon Sep 17 00:00:00 2001 From: lforst <8118419+lforst@users.noreply.github.com> Date: Tue, 8 Sep 2026 17:18:32 +0000 Subject: [PATCH 3/3] Update PR #725 --- py/src/braintrust/__init__.py | 2 +- py/src/braintrust/test_durable_eval.py | 607 ----------- py/src/braintrust/test_workflow_eval.py | 991 ++++++++++++++++++ .../type_tests/test_durable_eval.py | 81 -- .../type_tests/test_workflow_eval.py | 75 ++ .../{durable_eval.py => workflow_eval.py} | 785 +++++++------- 6 files changed, 1455 insertions(+), 1086 deletions(-) delete mode 100644 py/src/braintrust/test_durable_eval.py create mode 100644 py/src/braintrust/test_workflow_eval.py delete mode 100644 py/src/braintrust/type_tests/test_durable_eval.py create mode 100644 py/src/braintrust/type_tests/test_workflow_eval.py rename py/src/braintrust/{durable_eval.py => workflow_eval.py} (65%) diff --git a/py/src/braintrust/__init__.py b/py/src/braintrust/__init__.py index 1047b63d..6952ebe2 100644 --- a/py/src/braintrust/__init__.py +++ b/py/src/braintrust/__init__.py @@ -63,7 +63,6 @@ def is_equal(expected, output): from .audit import * from .auto import auto_instrument as auto_instrument from .dataset_pipeline import * -from .durable_eval import * from .framework import * from .framework2 import * from .functions.invoke import * @@ -88,3 +87,4 @@ def is_equal(expected, output): from .sandbox import register_sandbox as register_sandbox from .util import BT_IS_ASYNC_ATTRIBUTE as BT_IS_ASYNC_ATTRIBUTE from .util import MarkAsyncWrapper as MarkAsyncWrapper +from .workflow_eval import * diff --git a/py/src/braintrust/test_durable_eval.py b/py/src/braintrust/test_durable_eval.py deleted file mode 100644 index 1cd38a4d..00000000 --- a/py/src/braintrust/test_durable_eval.py +++ /dev/null @@ -1,607 +0,0 @@ -"""Tests for the experimental durable eval API.""" - -import asyncio -from unittest.mock import patch - -import pytest - -from . import durable_eval as durable_eval_module -from .durable_eval import ( - BatchCompletionPoll, - BatchCompletionWebhook, - BatchPollResult, - BatchScorer, - BatchScorerResult, - BatchTask, - BatchTaskResult, - DurableEvalCompletedResult, - DurableEvalFailedResult, - DurableEvalMemoryStore, - DurableEvalRedisStore, - DurableEvalWaitingResult, - define_durable_eval, -) -from .logger import BraintrustState, Dataset, ObjectMetadata, ProjectDatasetMetadata -from .test_helpers import init_test_exp, with_memory_logger, with_simulate_login # noqa: F401 -from .util import LazyValue - - -@pytest.mark.asyncio -async def test_memory_store_is_atomic_and_copies_values(): - store = DurableEvalMemoryStore() - value = bytearray(b"first") - - first = await store.get_or_set("key", value) - value[:] = b"other" - second = await store.get_or_set("key", b"second") - - assert first.created is True - assert first.value == b"first" - assert second.created is False - assert second.value == b"first" - assert await store.read("key") == b"first" - - -class _SyncRedis: - def __init__(self): - self.values = {} - self.calls = [] - - def get(self, key): - return self.values.get(key) - - def set(self, key, value, **kwargs): - self.calls.append((key, kwargs)) - old = self.values.get(key) - if kwargs.get("nx") and old is not None: - return old if kwargs.get("get") else None - self.values[key] = value - return old if kwargs.get("get") else True - - -class _AsyncRedis(_SyncRedis): - async def get(self, key): - return super().get(key) - - async def set(self, key, value, **kwargs): - return super().set(key, value, **kwargs) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("client_type", [_SyncRedis, _AsyncRedis]) -async def test_redis_store_supports_sync_and_async_redis_py(client_type): - client = client_type() - store = DurableEvalRedisStore(client, key_prefix="test:", ttl_ms=1234) - - await store.write("one", b"value") - assert await store.read("one") == b"value" - assert (await store.get_or_set("two", b"first")).created is True - existing = await store.get_or_set("two", b"second") - - assert existing.created is False - 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) - - -@pytest.mark.asyncio -async def test_ordinary_durable_eval_completes_locally_and_is_idempotent(): - task_calls = [] - score_calls = [] - - def task(input, hooks): - task_calls.append((input, hooks.trial_index)) - hooks.metadata["task"] = True - hooks.tags = ["updated"] - return input * 2 - - def scorer(output, expected, metadata, tags): - score_calls.append((output, metadata, tags)) - return output == expected - - durable_eval = define_durable_eval( - "project", - store=DurableEvalMemoryStore(), - data=[{"id": "case", "input": 2, "expected": 4, "metadata": {"source": "test"}}], - task=task, - scores=[scorer], - trial_count=2, - ) - - result = await durable_eval.start(no_send_logs=True) - repeated = await durable_eval.status(result.run_id) - - assert isinstance(result, DurableEvalCompletedResult) - assert isinstance(repeated, DurableEvalCompletedResult) - assert result.summary.scores["scorer"].score == 1 - assert task_calls == [(2, 0), (2, 1)] - assert score_calls == [ - (4, {"source": "test", "task": True}, ["updated"]), - (4, {"source": "test", "task": True}, ["updated"]), - ] - - -@pytest.mark.asyncio -async def test_ordinary_only_runs_do_not_require_case_ids_and_get_new_run_ids(): - durable_eval = define_durable_eval( - "project", - store=DurableEvalMemoryStore(), - data=[{"input": 1}], - task=lambda value: value, - scores=[lambda: 1], - ) - - first = await durable_eval.start(no_send_logs=True) - second = await durable_eval.start(no_send_logs=True) - - assert first.status == "completed" - assert second.status == "completed" - assert first.run_id != second.run_id - - -@pytest.mark.asyncio -async def test_poll_advances_task_then_mixed_scorers_without_polling_new_batches(): - submitted = {} - ready = set() - calls = [] - - async def task_submit(items, context): - calls.append(("submit-task", context.batch_id)) - submitted[context.batch_id] = items - return {"provider_id": context.batch_id} - - async def task_collect(_submission, context): - calls.append(("collect-task", context.batch_id)) - return [BatchTaskResult(id=item.id, output=item.input * 2) for item in submitted[context.batch_id]] - - async def score_submit(items, context): - calls.append(("submit-score", context.batch_id)) - submitted[context.batch_id] = items - return {"provider_id": context.batch_id} - - async def score_collect(_submission, context): - calls.append(("collect-score", context.batch_id)) - return [ - BatchScorerResult(id=item.id, score=item.output == item.expected) for item in submitted[context.batch_id] - ] - - async def poll(_submission, context): - calls.append(("poll", context.batch_id)) - return BatchPollResult("complete" if context.batch_id in ready else "pending") - - task = BatchTask( - submit=task_submit, - completion=BatchCompletionPoll(poll), - collect=task_collect, - batch_size=2, - ) - batch_score = BatchScorer( - name="batch", - submit=score_submit, - completion=BatchCompletionPoll(poll), - collect=score_collect, - batch_size=2, - ) - durable_eval = define_durable_eval( - "project", - store=DurableEvalMemoryStore(), - data=[{"id": str(i), "input": i, "expected": i * 2} for i in range(3)], - task=task, - scores=[lambda output, expected: output == expected, batch_score], - ) - - started = await durable_eval.start(no_send_logs=True) - assert isinstance(started, DurableEvalWaitingResult) - assert started.pending.poll == 2 - task_batch_ids = set(submitted) - ready.update(task_batch_ids) - - after_tasks = await durable_eval.poll(started.run_id) - assert isinstance(after_tasks, DurableEvalWaitingResult) - assert after_tasks.pending.poll == 2 - score_batch_ids = set(submitted) - task_batch_ids - assert score_batch_ids - assert not any(call == ("poll", batch_id) for batch_id in score_batch_ids for call in calls) - - ready.update(score_batch_ids) - completed = await durable_eval.poll(started.run_id) - assert isinstance(completed, DurableEvalCompletedResult) - assert completed.summary.scores["scorer_0"].score == 1 - assert completed.summary.scores["batch"].score == 1 - calls_before_status = list(calls) - assert (await durable_eval.status(started.run_id)).status == "completed" - assert calls == calls_before_status - - -@pytest.mark.asyncio -async def test_failed_poll_is_persisted_as_a_terminal_result(): - poll_calls = 0 - - async def poll(_submission, _context): - nonlocal poll_calls - poll_calls += 1 - return BatchPollResult("failed", error={"code": "provider_failed"}) - - durable_eval = define_durable_eval( - "project", - store=DurableEvalMemoryStore(), - data=[{"id": "case", "input": 1}], - task=BatchTask( - submit=lambda _items, context: {"id": context.batch_id}, - completion=BatchCompletionPoll(poll), - collect=lambda _submission, _context: pytest.fail("failed batches must not be collected"), - ), - ) - - started = await durable_eval.start(no_send_logs=True) - failed = await durable_eval.poll(started.run_id) - repeated = await durable_eval.status(started.run_id) - - assert isinstance(failed, DurableEvalFailedResult) - assert repeated == failed - assert failed.error == {"code": "provider_failed"} - assert failed.pending.poll == 0 - assert failed.pending.webhook == 0 - assert poll_calls == 1 - - -@pytest.mark.asyncio -async def test_webhook_result_can_be_matched_by_external_id(): - submitted = {} - - async def submit(items, context): - submitted[context.batch_id] = items - return {"id": f"external-{context.batch_id}"} - - async def collect(_submission, context): - return [BatchTaskResult(id=item.id, output=item.input) for item in submitted[context.batch_id]] - - durable_eval = define_durable_eval( - "project", - store=DurableEvalMemoryStore(), - data=[{"id": "case", "input": "ok"}], - task=BatchTask( - submit=submit, - completion=BatchCompletionWebhook(lambda submission, _context: submission["id"]), - collect=collect, - ), - ) - started = await durable_eval.start(no_send_logs=True) - batch_id = next(iter(submitted)) - - completed = await durable_eval.process_batch_result(started.run_id, external_id=f"external-{batch_id}") - - assert isinstance(completed, DurableEvalCompletedResult) - assert (await durable_eval.process_batch_result(started.run_id, batch_id=batch_id)).status == "completed" - - -@pytest.mark.asyncio -async def test_collect_requires_exactly_one_result_per_item(): - batch_ids = [] - - async def submit(_items, context): - batch_ids.append(context.batch_id) - return {"id": context.batch_id} - - durable_eval = define_durable_eval( - "project", - store=DurableEvalMemoryStore(), - data=[{"id": "a", "input": 1}, {"id": "b", "input": 2}], - task=BatchTask( - submit=submit, - completion=BatchCompletionWebhook(lambda submission, _context: submission["id"]), - collect=lambda _submission, _context: [BatchTaskResult(id="a:trial:0", output=1)], - ), - ) - started = await durable_eval.start(no_send_logs=True) - - with pytest.raises(ValueError, match="missing=.*b:trial:0"): - await durable_eval.process_batch_result(started.run_id, batch_id=batch_ids[0]) - - -@pytest.mark.asyncio -async def test_case_ids_must_be_stable_and_unique(): - batch_task = BatchTask( - submit=lambda _items, _context: {"id": "unused"}, - completion=BatchCompletionPoll(lambda _submission, _context: BatchPollResult("pending")), - collect=lambda _submission, _context: [], - ) - missing = define_durable_eval("project", store=DurableEvalMemoryStore(), data=[{"input": 1}], task=batch_task) - with pytest.raises(ValueError, match="non-empty id"): - await missing.start(no_send_logs=True) - - duplicate = define_durable_eval( - "project", - store=DurableEvalMemoryStore(), - data=[{"id": "same", "input": 1}, {"id": "same", "input": 2}], - task=batch_task, - ) - with pytest.raises(ValueError, match="duplicate"): - await duplicate.start(no_send_logs=True) - - -@pytest.mark.asyncio -async def test_case_persistence_does_not_deepcopy_inputs(): - class SerializableWithoutDeepcopy: - def __deepcopy__(self, _memo): - raise AssertionError("input was deep-copied") - - def model_dump(self, **_kwargs): - return {"value": "serialized"} - - submitted = [] - - async def submit(items, _context): - submitted.extend(items) - return {"id": "pending"} - - durable_eval = define_durable_eval( - "project", - store=DurableEvalMemoryStore(), - data=[{"id": "case", "input": SerializableWithoutDeepcopy()}], - task=BatchTask( - submit=submit, - completion=BatchCompletionPoll(lambda _submission, _context: BatchPollResult("pending")), - collect=lambda _submission, _context: [], - ), - ) - - result = await durable_eval.start(no_send_logs=True) - - assert result.status == "waiting" - assert submitted[0].input == {"value": "serialized"} - - -@pytest.mark.asyncio -async def test_concurrent_webhooks_claim_downstream_submission_once(): - task_items = [] - collects_started = 0 - both_collecting = asyncio.Event() - score_submissions = 0 - - async def submit_task(items, _context): - task_items.extend(items) - return {"id": "task-provider"} - - async def collect_task(_submission, _context): - nonlocal collects_started - collects_started += 1 - if collects_started == 2: - both_collecting.set() - await both_collecting.wait() - return [BatchTaskResult(id=item.id, output=item.input * 2) for item in task_items] - - async def submit_score(_items, _context): - nonlocal score_submissions - score_submissions += 1 - return {"id": "score-provider"} - - durable_eval = define_durable_eval( - "project", - store=DurableEvalMemoryStore(), - data=[{"id": "case", "input": 2, "expected": 4}], - task=BatchTask( - submit=submit_task, - completion=BatchCompletionWebhook(lambda submission, _context: submission["id"]), - collect=collect_task, - ), - scores=[ - BatchScorer( - name="exact", - submit=submit_score, - completion=BatchCompletionWebhook(lambda submission, _context: submission["id"]), - collect=lambda _submission, _context: [], - ) - ], - ) - started = await durable_eval.start(no_send_logs=True) - - await asyncio.gather( - durable_eval.process_batch_result(started.run_id, external_id="task-provider"), - durable_eval.process_batch_result(started.run_id, external_id="task-provider"), - ) - - assert score_submissions == 1 - current = await durable_eval.status(started.run_id) - assert current.status == "waiting" - assert current.pending.webhook == 1 - - -@pytest.mark.asyncio -@pytest.mark.parametrize("legacy_ids", [False, True]) -async def test_durable_logging_uses_stable_spans_and_resume_metadata( - monkeypatch, with_memory_logger, with_simulate_login, legacy_ids -): - if legacy_ids: - monkeypatch.setenv("BRAINTRUST_LEGACY_IDS", "true") - else: - monkeypatch.delenv("BRAINTRUST_LEGACY_IDS", raising=False) - local_eval = define_durable_eval( - "project", - store=DurableEvalMemoryStore(), - data=[{"id": "case", "input": 1, "expected": 2}], - task=lambda value: value + 1, - scores=[lambda output, expected: output == expected], - ) - local_summary = (await local_eval.start(no_send_logs=True)).summary - experiment = init_test_exp("durable", "project") - monkeypatch.setattr(durable_eval_module, "init_experiment", lambda **_kwargs: experiment) - monkeypatch.setattr(experiment, "summarize", lambda **_kwargs: local_summary) - submitted = [] - trace_configurations = [] - classifier_trace_configurations = [] - - async def submit(items, _context): - submitted.extend(items) - return {"id": "task"} - - def scorer(output, expected, trace): - trace_configurations.append(trace.get_configuration()) - return output == expected - - def classifier(output, trace): - classifier_trace_configurations.append(trace.get_configuration()) - return {"id": "positive"} - - durable_eval = define_durable_eval( - "project", - store=DurableEvalMemoryStore(), - data=[{"id": "case", "input": 1, "expected": 2}], - task=BatchTask( - submit=submit, - completion=BatchCompletionPoll(lambda _submission, _context: BatchPollResult("complete")), - collect=lambda _submission, _context: [ - BatchTaskResult(id=item.id, output=item.input + 1) for item in submitted - ], - ), - scores=[scorer], - classifiers=[classifier], - experiment_name="durable", - ) - - waiting = await durable_eval.start() - result = await durable_eval.poll(waiting.run_id) - logs = with_memory_logger.pop() - - assert result.status == "completed" - assert len(logs) == 4 - roots = [row for row in logs if not row["span_parents"]] - assert len(roots) == 1 - assert roots[0]["metadata"]["durable_eval"] == { - "run_id": result.run_id, - "case_id": "case", - "trial_index": 0, - } - assert trace_configurations == [ - { - "object_type": "experiment", - "object_id": experiment.id, - "root_span_id": roots[0]["root_span_id"], - } - ] - assert classifier_trace_configurations == trace_configurations - assert len({row["span_id"] for row in logs}) == 4 - await durable_eval.status(result.run_id) - assert with_memory_logger.pop() == [] - - -@pytest.mark.asyncio -async def test_durable_logging_flushes_before_persisting_log_markers( - monkeypatch, with_memory_logger, with_simulate_login -): - local_eval = define_durable_eval( - "project", - store=DurableEvalMemoryStore(), - data=[{"id": "case", "input": 1}], - task=lambda value: value, - ) - local_summary = (await local_eval.start(no_send_logs=True)).summary - experiment = init_test_exp("durable", "project") - monkeypatch.setattr(durable_eval_module, "init_experiment", lambda **_kwargs: experiment) - monkeypatch.setattr(experiment, "summarize", lambda **_kwargs: local_summary) - - flush_count = 0 - marker_flush_counts = [] - original_flush = with_memory_logger.flush - - def flush(*args, **kwargs): - nonlocal flush_count - flush_count += 1 - return original_flush(*args, **kwargs) - - monkeypatch.setattr(with_memory_logger, "flush", flush) - - class RecordingStore(DurableEvalMemoryStore): - async def write(self, key, value): - if "/task-log/" in key or "/score-log-" in key or "/classification-log-" in key: - marker_flush_counts.append(flush_count) - await super().write(key, value) - - submitted = [] - - async def submit(items, _context): - submitted.extend(items) - return {"id": "task"} - - durable_eval = define_durable_eval( - "project", - store=RecordingStore(), - data=[{"id": "case", "input": 1, "expected": 1}], - task=BatchTask( - submit=submit, - completion=BatchCompletionPoll(lambda _submission, _context: BatchPollResult("complete")), - collect=lambda _submission, _context: [ - BatchTaskResult(id=item.id, output=item.input) for item in submitted - ], - ), - scores=[lambda output, expected: output == expected], - classifiers=[lambda output: {"id": "positive"}], - experiment_name="durable", - ) - - waiting = await durable_eval.start() - completed = await durable_eval.poll(waiting.run_id) - - assert completed.status == "completed" - assert len(marker_flush_counts) == 3 - assert all(count > 0 for count in marker_flush_counts) - - -@pytest.mark.asyncio -async def test_durable_dataset_rows_preserve_dataset_origin(monkeypatch, with_memory_logger, with_simulate_login): - project_metadata = ObjectMetadata(id="test-project", name="test-project", full_info={}) - dataset_metadata = ObjectMetadata(id="active-dataset", name="test-dataset", full_info={}) - dataset = Dataset( - lazy_metadata=LazyValue( - lambda: ProjectDatasetMetadata(project=project_metadata, dataset=dataset_metadata), - use_mutex=False, - ), - state=BraintrustState(), - ) - row = { - "id": "dataset-row", - "_xact_id": "dataset-xact", - "created": "2026-06-02T00:00:00.000Z", - "input": 1, - } - local_summary = ( - await define_durable_eval( - "project", store=DurableEvalMemoryStore(), data=[{"input": 1}], task=lambda value: value - ).start(no_send_logs=True) - ).summary - experiment = init_test_exp("durable", "project") - monkeypatch.setattr(durable_eval_module, "init_experiment", lambda **_kwargs: experiment) - monkeypatch.setattr(experiment, "summarize", lambda **_kwargs: local_summary) - - submitted = [] - - async def submit(items, _context): - submitted.extend(items) - return {"id": "task"} - - durable_eval = define_durable_eval( - "project", - store=DurableEvalMemoryStore(), - data=dataset, - task=BatchTask( - submit=submit, - completion=BatchCompletionPoll(lambda _submission, _context: BatchPollResult("complete")), - collect=lambda _submission, _context: [ - BatchTaskResult(id=item.id, output=item.input) for item in submitted - ], - ), - experiment_name="durable", - ) - - with patch.object(dataset, "_refetch", return_value=[row]): - waiting = await durable_eval.start() - await durable_eval.poll(waiting.run_id) - - root = next(log for log in with_memory_logger.pop() if not log["span_parents"]) - assert root["origin"] == { - "object_type": "dataset", - "object_id": "active-dataset", - "id": "dataset-row", - "_xact_id": "dataset-xact", - "created": "2026-06-02T00:00:00.000Z", - } diff --git a/py/src/braintrust/test_workflow_eval.py b/py/src/braintrust/test_workflow_eval.py new file mode 100644 index 00000000..c7f6fd8b --- /dev/null +++ b/py/src/braintrust/test_workflow_eval.py @@ -0,0 +1,991 @@ +"""Tests for the experimental workflow eval API.""" + +import asyncio +from unittest.mock import patch + +import pytest + +from . import workflow_eval as workflow_eval_module +from .logger import BraintrustState, Dataset, ObjectMetadata, ProjectDatasetMetadata +from .test_helpers import init_test_exp, with_memory_logger, with_simulate_login # noqa: F401 +from .util import LazyValue +from .workflow_eval import ( + WorkflowEvalCompletedResult, + WorkflowEvalMemoryStore, + WorkflowEvalRedisStore, + WorkflowEvalWaitingResult, + WorkflowScorer, + WorkflowScorerResult, + WorkflowSubmissionCompletionPoll, + WorkflowSubmissionCompletionWebhook, + WorkflowSubmissionPoll, + WorkflowTask, + WorkflowTaskResult, + define_workflow_eval, +) + + +@pytest.mark.asyncio +async def test_memory_store_is_atomic_and_copies_values(): + store = WorkflowEvalMemoryStore() + value = bytearray(b"first") + + first = await store.get_or_set("key", value) + value[:] = b"other" + second = await store.get_or_set("key", b"second") + + assert first.created is True + assert first.value == b"first" + assert second.created is False + assert second.value == b"first" + assert await store.read("key") == b"first" + assert await store.get_set_size("key") == 0 + await asyncio.gather(*(store.add_to_set("key", str(i % 3)) for i in range(30))) + assert await store.get_set_size("key") == 3 + assert await store.read("key") == b"first" + + +class _SyncRedis: + def __init__(self): + self.values = {} + self.calls = [] + 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 + return 1 + + def scard(self, key): + return len(self.sets.get(key, set())) + + def get(self, key): + return self.values.get(key) + + def set(self, key, value, **kwargs): + self.calls.append((key, kwargs)) + old = self.values.get(key) + if kwargs.get("nx") and old is not None: + return old if kwargs.get("get") else None + self.values[key] = value + return old if kwargs.get("get") else True + + +class _AsyncRedis(_SyncRedis): + async def eval(self, *args): + return super().eval(*args) + + async def scard(self, key): + return super().scard(key) + + async def get(self, key): + return super().get(key) + + async def set(self, key, value, **kwargs): + return super().set(key, value, **kwargs) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("client_type", [_SyncRedis, _AsyncRedis]) +async def test_redis_store_supports_sync_and_async_redis_py(client_type): + client = client_type() + store = WorkflowEvalRedisStore(client, key_prefix="test:", ttl_ms=1234) + + await store.write("one", b"value") + assert await store.read("one") == b"value" + assert (await store.get_or_set("two", b"first")).created is True + existing = await store.get_or_set("two", b"second") + + assert existing.created is False + 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.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} + + +@pytest.mark.asyncio +async def test_ordinary_workflow_eval_completes_locally_and_is_idempotent(): + task_calls = [] + score_calls = [] + + def task(input, hooks): + task_calls.append((input, hooks.trial_index)) + hooks.metadata["task"] = True + hooks.tags = ["updated"] + return input * 2 + + def scorer(output, expected, metadata, tags): + score_calls.append((output, metadata, tags)) + return output == expected + + workflow_eval = define_workflow_eval( + "project", + store=WorkflowEvalMemoryStore(), + data=[{"id": "case", "input": 2, "expected": 4, "metadata": {"source": "test"}}], + task=task, + scores=[scorer], + trial_count=2, + ) + + result = await workflow_eval.start(no_send_logs=True) + repeated = await workflow_eval.status(result.run_id) + + assert isinstance(result, WorkflowEvalCompletedResult) + assert isinstance(repeated, WorkflowEvalCompletedResult) + assert result.summary.scores["scorer"].score == 1 + assert sorted(task_calls) == [(2, 0), (2, 1)] + assert score_calls == [ + (4, {"source": "test", "task": True}, ["updated"]), + (4, {"source": "test", "task": True}, ["updated"]), + ] + + +@pytest.mark.asyncio +async def test_ordinary_only_runs_do_not_require_case_ids_and_get_new_run_ids(): + workflow_eval = define_workflow_eval( + "project", + store=WorkflowEvalMemoryStore(), + data=[{"input": 1}], + task=lambda value: value, + scores=[lambda: 1], + ) + + first = await workflow_eval.start(no_send_logs=True) + second = await workflow_eval.start(no_send_logs=True) + + assert first.status == "completed" + assert second.status == "completed" + assert first.run_id != second.run_id + + +@pytest.mark.asyncio +async def test_poll_advances_task_then_mixed_scorers_without_polling_new_submissions(): + submitted = {} + ready = set() + calls = [] + + async def task_submit(item, context): + calls.append(("submit-task", context.submission_id)) + submitted[context.submission_id] = item + return {"provider_id": context.submission_id} + + async def task_collect(_submission, context): + calls.append(("collect-task", context.submission_id)) + return WorkflowTaskResult(output=submitted[context.submission_id].input * 2) + + async def score_submit(item, context): + calls.append(("submit-score", context.submission_id)) + submitted[context.submission_id] = item + return {"provider_id": context.submission_id} + + async def score_collect(_submission, context): + calls.append(("collect-score", context.submission_id)) + item = submitted[context.submission_id] + return WorkflowScorerResult(score=item.output == item.expected) + + async def poll(_submission, context): + calls.append(("poll", context.submission_id)) + return WorkflowSubmissionPoll("complete" if context.submission_id in ready else "pending") + + task = WorkflowTask( + submit=task_submit, + completion=WorkflowSubmissionCompletionPoll(poll), + collect=task_collect, + ) + submission_score = WorkflowScorer( + name="submission", + submit=score_submit, + completion=WorkflowSubmissionCompletionPoll(poll), + collect=score_collect, + ) + workflow_eval = define_workflow_eval( + "project", + store=WorkflowEvalMemoryStore(), + data=[{"id": str(i), "input": i, "expected": i * 2} for i in range(3)], + task=task, + scores=[lambda output, expected: output == expected, submission_score], + ) + + started = await workflow_eval.start(no_send_logs=True) + assert isinstance(started, WorkflowEvalWaitingResult) + assert started.pending.poll == 3 + task_submission_ids = set(submitted) + ready.update(task_submission_ids) + + after_tasks = await workflow_eval.poll(started.run_id) + assert isinstance(after_tasks, WorkflowEvalWaitingResult) + assert after_tasks.pending.poll == 3 + score_submission_ids = set(submitted) - task_submission_ids + assert score_submission_ids + assert not any(call == ("poll", submission_id) for submission_id in score_submission_ids for call in calls) + + ready.update(score_submission_ids) + completed = await workflow_eval.poll(started.run_id) + assert isinstance(completed, WorkflowEvalCompletedResult) + assert completed.summary.scores["scorer_0"].score == 1 + assert completed.summary.scores["submission"].score == 1 + calls_before_status = list(calls) + assert (await workflow_eval.status(started.run_id)).status == "completed" + assert calls == calls_before_status + + +@pytest.mark.asyncio +async def test_failed_poll_raises_and_can_be_retried(): + failure = RuntimeError("provider failed") + ready = False + + async def poll(_submission, _context): + return WorkflowSubmissionPoll("complete") if ready else WorkflowSubmissionPoll("failed", error=failure) + + workflow = define_workflow_eval( + "project", + store=WorkflowEvalMemoryStore(), + data=[{"id": "case", "input": 1}], + task=WorkflowTask( + submit=lambda item, context: {"id": context.submission_id}, + completion=WorkflowSubmissionCompletionPoll(poll), + collect=lambda _submission, _context: WorkflowTaskResult(output=1), + ), + ) + started = await workflow.start(no_send_logs=True) + with pytest.raises(RuntimeError, match="provider failed") as raised: + await workflow.poll(started.run_id) + assert raised.value is failure + assert (await workflow.status(started.run_id)).pending.poll == 1 + ready = True + assert (await workflow.poll(started.run_id)).status == "completed" + + +@pytest.mark.asyncio +async def test_webhook_result_can_be_matched_by_external_id(): + submitted = {} + + async def submit(item, context): + submitted[context.submission_id] = item + return {"id": f"external-{context.submission_id}"} + + async def collect(_submission, context): + return WorkflowTaskResult(output=submitted[context.submission_id].input) + + workflow_eval = define_workflow_eval( + "project", + store=WorkflowEvalMemoryStore(), + data=[{"id": "case", "input": "ok"}], + task=WorkflowTask( + submit=submit, + completion=WorkflowSubmissionCompletionWebhook(lambda submission, _context: submission["id"]), + collect=collect, + ), + ) + 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}") + + assert isinstance(completed, WorkflowEvalCompletedResult) + assert ( + await workflow_eval.process_submission_result(started.run_id, submission_id=submission_id) + ).status == "completed" + + +@pytest.mark.asyncio +async def test_collect_rejects_array_results(): + submission_ids = [] + + async def submit(_item, context): + submission_ids.append(context.submission_id) + return {"id": context.submission_id} + + workflow_eval = define_workflow_eval( + "project", + store=WorkflowEvalMemoryStore(), + data=[{"id": "a", "input": 1}, {"id": "b", "input": 2}], + task=WorkflowTask( + submit=submit, + completion=WorkflowSubmissionCompletionWebhook(lambda submission, _context: submission["id"]), + collect=lambda _submission, _context: [WorkflowTaskResult(output=1)], + ), + ) + started = await workflow_eval.start(no_send_logs=True) + + with pytest.raises(TypeError, match="single WorkflowTaskResult"): + await workflow_eval.process_submission_result(started.run_id, submission_id=submission_ids[0]) + + +@pytest.mark.asyncio +async def test_case_ids_must_be_stable_and_unique(): + submission_task = WorkflowTask( + submit=lambda _item, _context: {"id": "unused"}, + completion=WorkflowSubmissionCompletionPoll(lambda _submission, _context: WorkflowSubmissionPoll("pending")), + collect=lambda _submission, _context: [], + ) + missing = define_workflow_eval( + "project", store=WorkflowEvalMemoryStore(), data=[{"input": 1}], task=submission_task + ) + with pytest.raises(ValueError, match="non-empty id"): + await missing.start(no_send_logs=True) + + duplicate = define_workflow_eval( + "project", + store=WorkflowEvalMemoryStore(), + data=[{"id": "same", "input": 1}, {"id": "same", "input": 2}], + task=submission_task, + ) + with pytest.raises(ValueError, match="duplicate"): + await duplicate.start(no_send_logs=True) + + +@pytest.mark.asyncio +async def test_case_persistence_does_not_deepcopy_inputs(): + class SerializableWithoutDeepcopy: + def __deepcopy__(self, _memo): + raise AssertionError("input was deep-copied") + + def model_dump(self, **_kwargs): + return {"value": "serialized"} + + submitted = [] + + async def submit(item, _context): + submitted.append(item) + return {"id": "pending"} + + workflow_eval = define_workflow_eval( + "project", + store=WorkflowEvalMemoryStore(), + data=[{"id": "case", "input": SerializableWithoutDeepcopy()}], + task=WorkflowTask( + submit=submit, + completion=WorkflowSubmissionCompletionPoll( + lambda _submission, _context: WorkflowSubmissionPoll("pending") + ), + collect=lambda _submission, _context: [], + ), + ) + + result = await workflow_eval.start(no_send_logs=True) + + assert result.status == "waiting" + assert submitted[0].input == {"value": "serialized"} + + +@pytest.mark.asyncio +async def test_concurrent_webhooks_claim_downstream_submission_once(): + task_items = [] + collects_started = 0 + both_collecting = asyncio.Event() + score_submissions = 0 + + async def submit_task(item, _context): + task_items.append(item) + return {"id": "task-provider"} + + async def collect_task(_submission, _context): + nonlocal collects_started + collects_started += 1 + if collects_started == 2: + both_collecting.set() + await both_collecting.wait() + return WorkflowTaskResult(output=task_items[0].input * 2) + + async def submit_score(_item, _context): + nonlocal score_submissions + score_submissions += 1 + return {"id": "score-provider"} + + workflow_eval = define_workflow_eval( + "project", + store=WorkflowEvalMemoryStore(), + data=[{"id": "case", "input": 2, "expected": 4}], + task=WorkflowTask( + submit=submit_task, + completion=WorkflowSubmissionCompletionWebhook(lambda submission, _context: submission["id"]), + collect=collect_task, + ), + scores=[ + WorkflowScorer( + name="exact", + submit=submit_score, + completion=WorkflowSubmissionCompletionWebhook(lambda submission, _context: submission["id"]), + collect=lambda _submission, _context: [], + ) + ], + ) + started = await workflow_eval.start(no_send_logs=True) + + await asyncio.gather( + workflow_eval.process_submission_result(started.run_id, external_id="task-provider"), + workflow_eval.process_submission_result(started.run_id, external_id="task-provider"), + ) + + assert score_submissions == 1 + current = await workflow_eval.status(started.run_id) + assert current.status == "waiting" + assert current.pending.webhook == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("legacy_ids", [False, True]) +async def test_workflow_logging_uses_stable_spans_and_resume_metadata( + monkeypatch, with_memory_logger, with_simulate_login, legacy_ids +): + if legacy_ids: + monkeypatch.setenv("BRAINTRUST_LEGACY_IDS", "true") + else: + monkeypatch.delenv("BRAINTRUST_LEGACY_IDS", raising=False) + local_eval = define_workflow_eval( + "project", + store=WorkflowEvalMemoryStore(), + data=[{"id": "case", "input": 1, "expected": 2}], + task=lambda value: value + 1, + scores=[lambda output, expected: output == expected], + ) + local_summary = (await local_eval.start(no_send_logs=True)).summary + experiment = init_test_exp("workflow", "project") + monkeypatch.setattr(workflow_eval_module, "init_experiment", lambda **_kwargs: experiment) + monkeypatch.setattr(experiment, "summarize", lambda **_kwargs: local_summary) + submitted = [] + trace_configurations = [] + classifier_trace_configurations = [] + + async def submit(item, _context): + submitted.append(item) + return {"id": "task"} + + def scorer(output, expected, trace): + trace_configurations.append(trace.get_configuration()) + return output == expected + + def classifier(output, trace): + classifier_trace_configurations.append(trace.get_configuration()) + return {"id": "positive"} + + workflow_eval = define_workflow_eval( + "project", + store=WorkflowEvalMemoryStore(), + data=[{"id": "case", "input": 1, "expected": 2}], + task=WorkflowTask( + submit=submit, + completion=WorkflowSubmissionCompletionPoll( + lambda _submission, _context: WorkflowSubmissionPoll("complete") + ), + collect=lambda _submission, _context: WorkflowTaskResult(output=submitted[0].input + 1), + ), + scores=[scorer], + classifiers=[classifier], + experiment_name="workflow", + ) + + waiting = await workflow_eval.start() + result = await workflow_eval.poll(waiting.run_id) + logs = with_memory_logger.pop() + + assert result.status == "completed" + assert len(logs) == 4 + roots = [row for row in logs if not row["span_parents"]] + assert len(roots) == 1 + assert roots[0]["metadata"]["workflow_eval"] == { + "run_id": result.run_id, + "case_id": "case", + "trial_index": 0, + } + assert trace_configurations == [ + { + "object_type": "experiment", + "object_id": experiment.id, + "root_span_id": roots[0]["root_span_id"], + } + ] + assert classifier_trace_configurations == trace_configurations + assert len({row["span_id"] for row in logs}) == 4 + await workflow_eval.status(result.run_id) + assert with_memory_logger.pop() == [] + + +@pytest.mark.asyncio +async def test_workflow_logging_flushes_before_persisting_log_markers( + monkeypatch, with_memory_logger, with_simulate_login +): + local_eval = define_workflow_eval( + "project", + store=WorkflowEvalMemoryStore(), + data=[{"id": "case", "input": 1}], + task=lambda value: value, + ) + local_summary = (await local_eval.start(no_send_logs=True)).summary + experiment = init_test_exp("workflow", "project") + monkeypatch.setattr(workflow_eval_module, "init_experiment", lambda **_kwargs: experiment) + monkeypatch.setattr(experiment, "summarize", lambda **_kwargs: local_summary) + + flush_count = 0 + marker_flush_counts = [] + original_flush = with_memory_logger.flush + + def flush(*args, **kwargs): + nonlocal flush_count + flush_count += 1 + return original_flush(*args, **kwargs) + + monkeypatch.setattr(with_memory_logger, "flush", flush) + + class RecordingStore(WorkflowEvalMemoryStore): + async def write(self, key, value): + if "/task-log/" in key or "/score-log-" in key or "/classification-log-" in key: + marker_flush_counts.append(flush_count) + await super().write(key, value) + + submitted = [] + + async def submit(item, _context): + submitted.append(item) + return {"id": "task"} + + workflow_eval = define_workflow_eval( + "project", + store=RecordingStore(), + data=[{"id": "case", "input": 1, "expected": 1}], + task=WorkflowTask( + submit=submit, + completion=WorkflowSubmissionCompletionPoll( + lambda _submission, _context: WorkflowSubmissionPoll("complete") + ), + collect=lambda _submission, _context: WorkflowTaskResult(output=submitted[0].input), + ), + scores=[lambda output, expected: output == expected], + classifiers=[lambda output: {"id": "positive"}], + experiment_name="workflow", + ) + + waiting = await workflow_eval.start() + completed = await workflow_eval.poll(waiting.run_id) + + assert completed.status == "completed" + assert len(marker_flush_counts) == 3 + assert all(count > 0 for count in marker_flush_counts) + + +@pytest.mark.asyncio +async def test_workflow_dataset_rows_preserve_dataset_origin(monkeypatch, with_memory_logger, with_simulate_login): + project_metadata = ObjectMetadata(id="test-project", name="test-project", full_info={}) + dataset_metadata = ObjectMetadata(id="active-dataset", name="test-dataset", full_info={}) + dataset = Dataset( + lazy_metadata=LazyValue( + lambda: ProjectDatasetMetadata(project=project_metadata, dataset=dataset_metadata), + use_mutex=False, + ), + state=BraintrustState(), + ) + row = { + "id": "dataset-row", + "_xact_id": "dataset-xact", + "created": "2026-06-02T00:00:00.000Z", + "input": 1, + } + local_summary = ( + await define_workflow_eval( + "project", store=WorkflowEvalMemoryStore(), data=[{"input": 1}], task=lambda value: value + ).start(no_send_logs=True) + ).summary + experiment = init_test_exp("workflow", "project") + monkeypatch.setattr(workflow_eval_module, "init_experiment", lambda **_kwargs: experiment) + monkeypatch.setattr(experiment, "summarize", lambda **_kwargs: local_summary) + + submitted = [] + + async def submit(item, _context): + submitted.append(item) + return {"id": "task"} + + workflow_eval = define_workflow_eval( + "project", + store=WorkflowEvalMemoryStore(), + data=dataset, + task=WorkflowTask( + submit=submit, + completion=WorkflowSubmissionCompletionPoll( + lambda _submission, _context: WorkflowSubmissionPoll("complete") + ), + collect=lambda _submission, _context: WorkflowTaskResult(output=submitted[0].input), + ), + experiment_name="workflow", + ) + + with patch.object(dataset, "_refetch", return_value=[row]): + waiting = await workflow_eval.start() + await workflow_eval.poll(waiting.run_id) + + root = next(log for log in with_memory_logger.pop() if not log["span_parents"]) + assert root["origin"] == { + "object_type": "dataset", + "object_id": "active-dataset", + "id": "dataset-row", + "_xact_id": "dataset-xact", + "created": "2026-06-02T00:00:00.000Z", + } + + +@pytest.mark.asyncio +async def test_only_completed_cases_start_scorers_and_classifiers(): + ready = set() + submitted = {} + local_scores = [] + classifications = [] + + async def submit(item, context): + submitted[context.submission_id] = item + return {"id": context.submission_id} + + async def poll(submission, _context): + return WorkflowSubmissionPoll("complete" if submission["id"] in ready else "pending") + + async def score(output): + local_scores.append(output) + return 1 + + async def classify(output): + classifications.append(output) + return {"id": "positive"} + + workflow = define_workflow_eval( + "project", + store=WorkflowEvalMemoryStore(), + data=[{"id": str(i), "input": i} for i in range(3)], + task=WorkflowTask( + submit=submit, + completion=WorkflowSubmissionCompletionPoll(poll), + collect=lambda submission, _context: WorkflowTaskResult(output=submitted[submission["id"]].input), + ), + scores=[ + score, + WorkflowScorer( + name="remote", + submit=submit, + completion=WorkflowSubmissionCompletionPoll(poll), + collect=lambda _submission, _context: WorkflowScorerResult(score=1), + ), + ], + classifiers=[classify], + ) + started = await workflow.start(no_send_logs=True) + task_ids = set(submitted) + ready.add(next(key for key, item in submitted.items() if item.input == 1)) + partial = await workflow.poll(started.run_id) + assert partial.pending.poll == 3 # two tasks and the ready case's scorer + assert local_scores == [1] + assert classifications == [1] + score_ids = set(submitted) - task_ids + assert len(score_ids) == 1 + assert submitted[next(iter(score_ids))].output == 1 + ready.update(score_ids) + partial = await workflow.poll(started.run_id) + assert partial.status == "waiting" + assert partial.pending.poll == 2 + assert local_scores == [1] + ready.update(task_ids) + await workflow.poll(started.run_id) + ready.update(submitted) + completed = await workflow.poll(started.run_id) + assert completed.status == "completed" + assert sorted(local_scores) == [0, 1, 2] + assert sorted(classifications) == [0, 1, 2] + assert completed.summary.scores["remote"].score == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("max_concurrency", [None, 1, 3]) +async def test_provider_callback_concurrency_is_bounded(max_concurrency): + active = {"submit": 0, "poll": 0, "collect": 0} + peak = active.copy() + + async def callback(stage): + active[stage] += 1 + peak[stage] = max(peak[stage], active[stage]) + await asyncio.sleep(0.001) + active[stage] -= 1 + + async def submit(item, context): + await callback("submit") + return {"id": context.submission_id, "input": item.input} + + async def poll(_submission, _context): + await callback("poll") + return WorkflowSubmissionPoll("complete") + + async def collect(submission, _context): + await callback("collect") + return WorkflowTaskResult(output=submission["input"]) + + options = {} if max_concurrency is None else {"max_concurrency": max_concurrency} + workflow = define_workflow_eval( + "project", + store=WorkflowEvalMemoryStore(), + data=[{"id": str(i), "input": i} for i in range(25)], + task=WorkflowTask(submit=submit, completion=WorkflowSubmissionCompletionPoll(poll), collect=collect), + **options, + ) + started = await workflow.start(no_send_logs=True) + assert started.pending.poll == 25 + assert (await workflow.poll(started.run_id)).status == "completed" + limit = max_concurrency or 10 + assert peak["submit"] == limit + assert peak["poll"] == limit + assert 0 < peak["collect"] <= limit + assert all(value == 0 for value in active.values()) + + +@pytest.mark.parametrize("max_concurrency", [0, -1, 1.5, True, "2", None]) +def test_invalid_concurrency_is_rejected(max_concurrency): + with pytest.raises(ValueError, match="max_concurrency must be a positive integer"): + define_workflow_eval( + "project", + store=WorkflowEvalMemoryStore(), + data=[], + task=lambda value: value, + max_concurrency=max_concurrency, + ) + + +@pytest.mark.asyncio +async def test_poll_error_does_not_prevent_independent_case_progress(): + scores = [] + + async def poll(submission, _context): + if submission["input"] == 0: + raise RuntimeError("provider unavailable") + return WorkflowSubmissionPoll("complete") + + async def score(output): + scores.append(output) + return 1 + + workflow = define_workflow_eval( + "project", + store=WorkflowEvalMemoryStore(), + data=[{"id": str(i), "input": i} for i in range(3)], + task=WorkflowTask( + submit=lambda item, _context: {"input": item.input}, + completion=WorkflowSubmissionCompletionPoll(poll), + collect=lambda submission, _context: WorkflowTaskResult(output=submission["input"]), + ), + scores=[score], + max_concurrency=1, + ) + started = await workflow.start(no_send_logs=True) + with pytest.raises(RuntimeError, match="provider unavailable"): + await workflow.poll(started.run_id) + assert scores == [1, 2] + assert (await workflow.status(started.run_id)).pending.poll == 1 + + +@pytest.mark.asyncio +async def test_submission_error_waits_for_independent_submissions(): + calls = [] + run_ids = [] + + async def submit(item, context): + run_ids.append(context.run_id) + if item.input == 0: + raise RuntimeError("submission failed") + await asyncio.sleep(0) + calls.append(item.input) + return {"input": item.input} + + workflow = define_workflow_eval( + "project", + store=WorkflowEvalMemoryStore(), + data=[{"id": str(i), "input": i} for i in range(3)], + task=WorkflowTask( + submit=submit, + completion=WorkflowSubmissionCompletionPoll( + lambda _submission, _context: WorkflowSubmissionPoll("pending") + ), + collect=lambda submission, _context: WorkflowTaskResult(output=submission["input"]), + ), + max_concurrency=1, + ) + 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 + + +@pytest.mark.asyncio +async def test_webhooks_resume_fresh_definitions_and_validate_both_locators(): + store = WorkflowEvalMemoryStore() + submitted = {} + collected = [] + + async def submit(item, context): + submitted[item.input] = context.submission_id + return {"id": item.input} + + async def collect(submission, _context): + collected.append(submission["id"]) + await asyncio.sleep(0) + return WorkflowTaskResult(output=submission["id"]) + + def definition(): + return define_workflow_eval( + "project", + store=store, + data=[{"id": value, "input": value} for value in ("a", "b")], + task=WorkflowTask( + submit=submit, + collect=collect, + completion=WorkflowSubmissionCompletionWebhook(lambda submission, _context: submission["id"]), + ), + ) + + started = await definition().start(no_send_logs=True) + workflow = definition() + for locators in [ + {"submission_id": submitted["a"], "external_id": "b"}, + {"submission_id": submitted["a"], "external_id": "unknown"}, + {"submission_id": "unknown", "external_id": "a"}, + ]: + with pytest.raises(ValueError, match="different submissions"): + await workflow.process_submission_result(started.run_id, **locators) + with pytest.raises(ValueError, match="requires submission_id or external_id"): + await workflow.process_submission_result(started.run_id) + with pytest.raises(ValueError, match="No submission matches"): + await workflow.process_submission_result(started.run_id, external_id="unknown") + await asyncio.gather( + workflow.process_submission_result(started.run_id, submission_id=submitted["a"]), + workflow.process_submission_result(started.run_id, external_id="b"), + ) + assert (await workflow.status(started.run_id)).status == "completed" + assert sorted(collected) == ["a", "b"] + await workflow.process_submission_result(started.run_id, external_id="a") + assert sorted(collected) == ["a", "b"] + + +@pytest.mark.asyncio +async def test_workflow_trials_preserve_parameters_metadata_and_tags(): + task_items = [] + score_items = [] + + async def submit_task(item, context): + task_items.append(item) + return {"id": context.submission_id, "trial": item.trial_index} + + async def submit_score(item, _context): + score_items.append(item) + return {"id": item.id} + + workflow = define_workflow_eval( + "project", + store=WorkflowEvalMemoryStore(), + case_id=lambda datum: "case", + data=[{"input": 1, "expected": 2, "metadata": {"original": True}, "tags": ["old"], "trial_count": 2}], + task=WorkflowTask( + submit=submit_task, + completion=WorkflowSubmissionCompletionPoll( + lambda _submission, _context: WorkflowSubmissionPoll("complete") + ), + collect=lambda submission, _context: WorkflowTaskResult( + output=submission["trial"], metadata={"new": True}, tags=["updated"] + ), + ), + scores=[ + WorkflowScorer( + name="remote", + submit=submit_score, + completion=WorkflowSubmissionCompletionWebhook(lambda submission, _context: submission["id"]), + collect=lambda _submission, _context: WorkflowScorerResult(score=1), + ) + ], + ) + started = await workflow.start({"temperature": 0.5}, no_send_logs=True) + assert {item.id for item in task_items} == {"case:trial:0", "case:trial:1"} + assert all(item.parameters == {"temperature": 0.5} for item in task_items) + assert sorted(item.trial_index for item in task_items) == [0, 1] + assert (await workflow.poll(started.run_id)).pending.webhook == 2 + assert all(item.metadata == {"original": True, "new": True} for item in score_items) + assert all(item.tags == ["updated"] and item.expected == 2 for item in score_items) + for item in score_items: + result = await workflow.process_submission_result(started.run_id, external_id=item.id) + assert result.status == "completed" + + +@pytest.mark.asyncio +async def test_ordinary_tasks_with_workflow_scorers_and_empty_data(): + submitted = [] + + async def submit(item, _context): + submitted.append(item) + return {"id": item.id} + + scorer = WorkflowScorer( + name="remote", + submit=submit, + completion=WorkflowSubmissionCompletionPoll(lambda _submission, _context: WorkflowSubmissionPoll("complete")), + collect=lambda _submission, _context: WorkflowScorerResult(score=1), + ) + for data in [[], [{"id": "case", "input": 1}]]: + workflow = define_workflow_eval( + "project", store=WorkflowEvalMemoryStore(), data=data, task=lambda value: value + 1, scores=[scorer] + ) + started = await workflow.start(no_send_logs=True) + if data: + assert started.pending.poll == 1 + assert submitted[0].output == 2 + else: + assert started.status == "completed" + assert (await workflow.poll(started.run_id)).status == "completed" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ["poll", "webhook"]) +async def test_replay_repairs_interrupted_completion_progress(mode): + class InterruptedStore(WorkflowEvalMemoryStore): + interrupted = False + + async def add_to_set(self, key, member): + if key.endswith("/complete") and not self.interrupted: + self.interrupted = True + raise RuntimeError("interrupted progress") + await super().add_to_set(key, member) + + collects = [] + + async def collect(submission, _context): + collects.append(submission["id"]) + return WorkflowTaskResult(output=1) + + workflow = define_workflow_eval( + "project", + store=InterruptedStore(), + data=[{"id": value, "input": value} for value in ("a", "b")], + task=WorkflowTask( + submit=lambda item, _context: {"id": item.input}, + collect=collect, + completion=WorkflowSubmissionCompletionWebhook(lambda submission, _context: submission["id"]) + if mode == "webhook" + else WorkflowSubmissionCompletionPoll( + lambda submission, _context: WorkflowSubmissionPoll( + "complete" if submission["id"] == "a" else "pending" + ) + ), + ), + ) + started = await workflow.start(no_send_logs=True) + with pytest.raises(RuntimeError, match="interrupted progress"): + if mode == "webhook": + await workflow.process_submission_result(started.run_id, external_id="a") + else: + await workflow.poll(started.run_id) + if mode == "webhook": + result = await workflow.process_submission_result(started.run_id, external_id="a") + else: + result = await workflow.poll(started.run_id) + assert result.pending.poll + result.pending.webhook == 1 + assert collects == ["a"] diff --git a/py/src/braintrust/type_tests/test_durable_eval.py b/py/src/braintrust/type_tests/test_durable_eval.py deleted file mode 100644 index 6afe409b..00000000 --- a/py/src/braintrust/type_tests/test_durable_eval.py +++ /dev/null @@ -1,81 +0,0 @@ -"""Static and runtime type coverage for the experimental durable eval API.""" - -from typing import TypedDict - -import pytest -from braintrust import ( - BatchCompletionPoll, - BatchContext, - BatchPollResult, - BatchScorer, - BatchScorerItem, - BatchScorerResult, - BatchTask, - BatchTaskItem, - BatchTaskResult, - DurableEval, - DurableEvalFailedResult, - DurableEvalMemoryStore, - EvalCase, - define_durable_eval, -) - - -class Submission(TypedDict): - id: str - - -async def submit_task(items: list[BatchTaskItem[str, str]], context: BatchContext) -> Submission: - assert items - return {"id": context.batch_id} - - -async def collect_task(submission: Submission, context: BatchContext) -> list[BatchTaskResult[int]]: - return [BatchTaskResult(id="case:trial:0", output=len(submission["id"] + context.run_id))] - - -async def submit_score(items: list[BatchScorerItem[str, int, str]], context: BatchContext) -> Submission: - assert items - return {"id": context.batch_id} - - -async def collect_score(submission: Submission, context: BatchContext) -> list[BatchScorerResult]: - assert submission["id"] == context.batch_id - return [BatchScorerResult(id="case:trial:0", score=1)] - - -async def poll_batch(submission: Submission, context: BatchContext) -> BatchPollResult: - assert submission["id"] == context.batch_id - return BatchPollResult(status="pending") - - -task: BatchTask[str, int, str, Submission] = BatchTask( - submit=submit_task, - completion=BatchCompletionPoll(poll_batch), - collect=collect_task, -) -score: BatchScorer[str, int, str, Submission] = BatchScorer( - name="score", - submit=submit_score, - completion=BatchCompletionPoll(poll_batch), - collect=collect_score, -) -durable_eval: DurableEval[str, int, str] = define_durable_eval( - "project", - store=DurableEvalMemoryStore(), - data=[EvalCase(id="case", input="input", expected="expected")], - task=task, - scores=[score], -) - - -@pytest.mark.asyncio -async def test_durable_eval_types_at_runtime(): - result = await durable_eval.start(no_send_logs=True) - assert result.status == "waiting" - assert result.pending.poll == 1 - - -def consume_failed_result(result: DurableEvalFailedResult) -> object: - assert result.status == "failed" - return result.error diff --git a/py/src/braintrust/type_tests/test_workflow_eval.py b/py/src/braintrust/type_tests/test_workflow_eval.py new file mode 100644 index 00000000..a43b7683 --- /dev/null +++ b/py/src/braintrust/type_tests/test_workflow_eval.py @@ -0,0 +1,75 @@ +"""Static and runtime type coverage for the experimental workflow eval API.""" + +from typing import TypedDict + +import pytest +from braintrust import ( + EvalCase, + WorkflowEval, + WorkflowEvalMemoryStore, + WorkflowScorer, + WorkflowScorerItem, + WorkflowScorerResult, + WorkflowSubmissionCompletionPoll, + WorkflowSubmissionContext, + WorkflowSubmissionPoll, + WorkflowTask, + WorkflowTaskItem, + WorkflowTaskResult, + define_workflow_eval, +) + + +class Submission(TypedDict): + id: str + + +async def submit_task(item: WorkflowTaskItem[str, str], context: WorkflowSubmissionContext) -> Submission: + assert item.id + return {"id": context.submission_id} + + +async def collect_task(submission: Submission, context: WorkflowSubmissionContext) -> WorkflowTaskResult[int]: + return WorkflowTaskResult(output=len(submission["id"] + context.run_id)) + + +async def submit_score(item: WorkflowScorerItem[str, int, str], context: WorkflowSubmissionContext) -> Submission: + assert item.id + return {"id": context.submission_id} + + +async def collect_score(submission: Submission, context: WorkflowSubmissionContext) -> WorkflowScorerResult: + assert submission["id"] == context.submission_id + return WorkflowScorerResult(score=1) + + +async def poll_submission(submission: Submission, context: WorkflowSubmissionContext) -> WorkflowSubmissionPoll: + assert submission["id"] == context.submission_id + return WorkflowSubmissionPoll(status="pending") + + +task: WorkflowTask[str, int, str, Submission] = WorkflowTask( + submit=submit_task, + completion=WorkflowSubmissionCompletionPoll(poll_submission), + collect=collect_task, +) +score: WorkflowScorer[str, int, str, Submission] = WorkflowScorer( + name="score", + submit=submit_score, + completion=WorkflowSubmissionCompletionPoll(poll_submission), + collect=collect_score, +) +workflow_eval: WorkflowEval[str, int, str] = define_workflow_eval( + "project", + store=WorkflowEvalMemoryStore(), + data=[EvalCase(id="case", input="input", expected="expected")], + task=task, + scores=[score], +) + + +@pytest.mark.asyncio +async def test_workflow_eval_types_at_runtime(): + result = await workflow_eval.start(no_send_logs=True) + assert result.status == "waiting" + assert result.pending.poll == 1 diff --git a/py/src/braintrust/durable_eval.py b/py/src/braintrust/workflow_eval.py similarity index 65% rename from py/src/braintrust/durable_eval.py rename to py/src/braintrust/workflow_eval.py index 2b7de3e5..01d5ae55 100644 --- a/py/src/braintrust/durable_eval.py +++ b/py/src/braintrust/workflow_eval.py @@ -1,4 +1,4 @@ -"""Experimental durable evaluation support for asynchronous batch providers.""" +"""Experimental workflow evaluations with one asynchronous provider submission per case/trial.""" import asyncio import base64 @@ -67,21 +67,21 @@ Expected = TypeVar("Expected") SubmissionData = TypeVar("SubmissionData") -DEFAULT_BATCH_SIZE = 1_000 +DEFAULT_MAX_CONCURRENCY = 10 DEFAULT_REDIS_TTL_MS = 1_000 * 60 * 60 * 24 * 7 -_SCHEMA_PREFIX = "durable-eval/python/v1" +_SCHEMA_PREFIX = "workflow-eval/python/v1" @dataclasses.dataclass(frozen=True) -class BatchContext: - """Identifiers supplied to a durable batch processor callback.""" +class WorkflowSubmissionContext: + """Identifiers supplied to a workflow submission processor callback.""" run_id: str - batch_id: str + submission_id: str @dataclasses.dataclass(frozen=True) -class BatchPollResult: +class WorkflowSubmissionPoll: """The current state returned by a polling completion callback.""" status: Literal["pending", "complete", "failed"] @@ -89,27 +89,31 @@ class BatchPollResult: @dataclasses.dataclass(frozen=True) -class BatchCompletionPoll(Generic[SubmissionData]): - """Configures a batch processor whose provider is checked by polling.""" +class WorkflowSubmissionCompletionPoll(Generic[SubmissionData]): + """Configures a submission processor whose provider is checked by polling.""" - poll: Callable[[SubmissionData, BatchContext], BatchPollResult | Awaitable[BatchPollResult]] + poll: Callable[ + [SubmissionData, WorkflowSubmissionContext], WorkflowSubmissionPoll | Awaitable[WorkflowSubmissionPoll] + ] mode: Literal["poll"] = dataclasses.field(default="poll", init=False) @dataclasses.dataclass(frozen=True) -class BatchCompletionWebhook(Generic[SubmissionData]): - """Configures a batch processor completed by an incoming webhook.""" +class WorkflowSubmissionCompletionWebhook(Generic[SubmissionData]): + """Configures a submission processor completed by an incoming webhook.""" - get_external_id: Callable[[SubmissionData, BatchContext], str | Awaitable[str]] + get_external_id: Callable[[SubmissionData, WorkflowSubmissionContext], str | Awaitable[str]] mode: Literal["webhook"] = dataclasses.field(default="webhook", init=False) -BatchCompletion = BatchCompletionPoll[SubmissionData] | BatchCompletionWebhook[SubmissionData] +WorkflowSubmissionCompletion = ( + WorkflowSubmissionCompletionPoll[SubmissionData] | WorkflowSubmissionCompletionWebhook[SubmissionData] +) @dataclasses.dataclass(frozen=True) -class BatchTaskItem(Generic[Input, Expected]): - """A stable task item submitted to a provider batch.""" +class WorkflowTaskItem(Generic[Input, Expected]): + """One case/trial passed to a workflow task submission callback.""" id: str input: Input @@ -121,18 +125,17 @@ class BatchTaskItem(Generic[Input, Expected]): @dataclasses.dataclass(frozen=True) -class BatchTaskResult(Generic[Output]): +class WorkflowTaskResult(Generic[Output]): """A collected result for one task item.""" - id: str output: Output metadata: Metadata | None = None tags: list[str] | None = None @dataclasses.dataclass(frozen=True) -class BatchScorerItem(Generic[Input, Output, Expected]): - """A stable scorer item submitted to a provider batch.""" +class WorkflowScorerItem(Generic[Input, Output, Expected]): + """One completed case/trial passed to a workflow scorer submission callback.""" id: str input: Input @@ -144,71 +147,74 @@ class BatchScorerItem(Generic[Input, Output, Expected]): @dataclasses.dataclass(frozen=True) -class BatchScorerResult: +class WorkflowScorerResult: """A collected score for one scorer item.""" - id: str score: OneOrMoreScores @dataclasses.dataclass(frozen=True) -class BatchTask(Generic[Input, Output, Expected, SubmissionData]): - """Runs an evaluation task through asynchronous provider batch operations.""" +class WorkflowTask(Generic[Input, Output, Expected, SubmissionData]): + """Submits one asynchronous provider operation per case/trial and collects its task result.""" - submit: Callable[[list[BatchTaskItem[Input, Expected]], BatchContext], SubmissionData | Awaitable[SubmissionData]] - completion: BatchCompletion[SubmissionData] + submit: Callable[ + [WorkflowTaskItem[Input, Expected], WorkflowSubmissionContext], SubmissionData | Awaitable[SubmissionData] + ] + completion: WorkflowSubmissionCompletion[SubmissionData] collect: Callable[ - [SubmissionData, BatchContext], list[BatchTaskResult[Output]] | Awaitable[list[BatchTaskResult[Output]]] + [SubmissionData, WorkflowSubmissionContext], WorkflowTaskResult[Output] | Awaitable[WorkflowTaskResult[Output]] ] - batch_size: int = DEFAULT_BATCH_SIZE - - def __post_init__(self) -> None: - if not isinstance(self.batch_size, int) or isinstance(self.batch_size, bool) or self.batch_size < 1: - raise ValueError("BatchTask batch_size must be a positive integer") @dataclasses.dataclass(frozen=True) -class BatchScorer(Generic[Input, Output, Expected, SubmissionData]): - """Runs an evaluation scorer through asynchronous provider batch operations.""" +class WorkflowScorer(Generic[Input, Output, Expected, SubmissionData]): + """Submits one asynchronous provider operation per case/trial and collects its score.""" name: str submit: Callable[ - [list[BatchScorerItem[Input, Output, Expected]], BatchContext], SubmissionData | Awaitable[SubmissionData] + [WorkflowScorerItem[Input, Output, Expected], WorkflowSubmissionContext], + SubmissionData | Awaitable[SubmissionData], + ] + completion: WorkflowSubmissionCompletion[SubmissionData] + collect: Callable[ + [SubmissionData, WorkflowSubmissionContext], WorkflowScorerResult | Awaitable[WorkflowScorerResult] ] - completion: BatchCompletion[SubmissionData] - collect: Callable[[SubmissionData, BatchContext], list[BatchScorerResult] | Awaitable[list[BatchScorerResult]]] - batch_size: int = DEFAULT_BATCH_SIZE def __post_init__(self) -> None: if not self.name: - raise ValueError("BatchScorer name must be a non-empty string") - if not isinstance(self.batch_size, int) or isinstance(self.batch_size, bool) or self.batch_size < 1: - raise ValueError("BatchScorer batch_size must be a positive integer") + raise ValueError("WorkflowScorer name must be a non-empty string") @dataclasses.dataclass(frozen=True) -class DurableEvalStoreEntry: - """Result of an atomic durable-store get-or-set operation.""" +class WorkflowEvalStoreEntry: + """Result of an atomic workflow-store get-or-set operation.""" value: bytes created: bool -class DurableEvalStore(Protocol): - """Minimal persistence interface used by durable evaluations.""" +class WorkflowEvalStore(Protocol): + """Minimal persistence interface used by workflow evaluations.""" async def read(self, key: str) -> bytes | None: ... async def write(self, key: str, value: bytes) -> None: ... - async def get_or_set(self, key: str, value: bytes) -> DurableEvalStoreEntry: ... + async def get_or_set(self, key: str, value: bytes) -> WorkflowEvalStoreEntry: ... + + async def add_to_set(self, key: str, member: str) -> None: + """Atomically add a unique member; retain sets as long as run records.""" + ... + async def get_set_size(self, key: str) -> int: ... -class DurableEvalMemoryStore: - """Process-local durable evaluation state, intended for tests and local runs.""" + +class WorkflowEvalMemoryStore: + """Process-local workflow evaluation state, intended for tests and local runs.""" def __init__(self) -> None: self._values: dict[str, bytes] = {} + self._sets: dict[str, set[str]] = {} self._lock = threading.Lock() async def read(self, key: str) -> bytes | None: @@ -220,17 +226,25 @@ async def write(self, key: str, value: bytes) -> None: with self._lock: self._values[key] = bytes(value) - async def get_or_set(self, key: str, value: bytes) -> DurableEvalStoreEntry: + async def get_or_set(self, key: str, value: bytes) -> WorkflowEvalStoreEntry: with self._lock: existing = self._values.get(key) if existing is not None: - return DurableEvalStoreEntry(value=bytes(existing), created=False) + return WorkflowEvalStoreEntry(value=bytes(existing), created=False) self._values[key] = bytes(value) - return DurableEvalStoreEntry(value=bytes(value), created=True) + return WorkflowEvalStoreEntry(value=bytes(value), created=True) + async def add_to_set(self, key: str, member: str) -> None: + with self._lock: + self._sets.setdefault(key, set()).add(member) -class DurableEvalRedisStore: - """Durable state backed by an existing sync or async redis-py client.""" + async def get_set_size(self, key: str) -> int: + with self._lock: + return len(self._sets.get(key, set())) + + +class WorkflowEvalRedisStore: + """Workflow state backed by an existing sync or async redis-py client.""" def __init__( self, @@ -240,7 +254,7 @@ def __init__( ttl_ms: int = DEFAULT_REDIS_TTL_MS, ) -> None: if not isinstance(ttl_ms, int) or isinstance(ttl_ms, bool) or ttl_ms < 1: - raise ValueError("DurableEvalRedisStore ttl_ms must be a positive integer") + raise ValueError("WorkflowEvalRedisStore ttl_ms must be a positive integer") self.client = client self.key_prefix = key_prefix self.ttl_ms = ttl_ms @@ -260,72 +274,77 @@ async def read(self, key: str) -> bytes | None: if isinstance(value, bytes): value = value.decode("ascii") if not isinstance(value, str): - raise TypeError("DurableEvalRedisStore expected GET to return str, bytes, or None") + raise TypeError("WorkflowEvalRedisStore expected GET to return str, bytes, or None") return base64.b64decode(value) async def write(self, key: str, value: bytes) -> None: encoded = base64.b64encode(value).decode("ascii") await self._call(self.client.set, f"{self.key_prefix}{key}", encoded, px=self.ttl_ms) - async def get_or_set(self, key: str, value: bytes) -> DurableEvalStoreEntry: + async def get_or_set(self, key: str, value: bytes) -> WorkflowEvalStoreEntry: redis_key = f"{self.key_prefix}{key}" encoded = base64.b64encode(value).decode("ascii") existing = await self._call(self.client.set, redis_key, encoded, px=self.ttl_ms, nx=True, get=True) if existing is None: - return DurableEvalStoreEntry(value=bytes(value), created=True) + return WorkflowEvalStoreEntry(value=bytes(value), created=True) if isinstance(existing, bytes): existing = existing.decode("ascii") if not isinstance(existing, str): - raise TypeError("DurableEvalRedisStore expected atomic SET to return str, bytes, or None") - return DurableEvalStoreEntry(value=base64.b64decode(existing), created=False) + raise TypeError("WorkflowEvalRedisStore expected atomic SET to return str, bytes, or None") + return WorkflowEvalStoreEntry(value=base64.b64decode(existing), created=False) + + async def add_to_set(self, key: str, member: str) -> None: + await self._call( + self.client.eval, + "redis.call('SADD', KEYS[1], ARGV[1]); redis.call('PEXPIRE', KEYS[1], ARGV[2]); return 1", + 1, + f"{self.key_prefix}{key}", + member, + self.ttl_ms, + ) + + async def get_set_size(self, key: str) -> int: + return int(await self._call(self.client.scard, f"{self.key_prefix}{key}")) @dataclasses.dataclass(frozen=True) -class DurableEvalPending: - """Counts of submitted provider batches awaiting completion.""" +class WorkflowEvalPending: + """Counts of submitted provider submissions awaiting completion.""" poll: int webhook: int @dataclasses.dataclass(frozen=True) -class DurableEvalWaitingResult: +class WorkflowEvalWaitingResult: run_id: str - pending: DurableEvalPending + pending: WorkflowEvalPending status: Literal["waiting"] = dataclasses.field(default="waiting", init=False) @dataclasses.dataclass(frozen=True) -class DurableEvalCompletedResult: +class WorkflowEvalCompletedResult: run_id: str - pending: DurableEvalPending + pending: WorkflowEvalPending summary: ExperimentSummary status: Literal["completed"] = dataclasses.field(default="completed", init=False) -@dataclasses.dataclass(frozen=True) -class DurableEvalFailedResult: - run_id: str - batch_id: str - error: Any - pending: DurableEvalPending - status: Literal["failed"] = dataclasses.field(default="failed", init=False) - - -DurableEvalResult = DurableEvalWaitingResult | DurableEvalCompletedResult | DurableEvalFailedResult +WorkflowEvalResult = WorkflowEvalWaitingResult | WorkflowEvalCompletedResult @dataclasses.dataclass(frozen=True) -class _DurableEvalConfig(Generic[Input, Output, Expected]): +class _WorkflowEvalConfig(Generic[Input, Output, Expected]): project_name: str - store: DurableEvalStore + store: WorkflowEvalStore data: EvalData[Input, Expected] - task: EvalTask[Input, Output, Expected] | BatchTask[Input, Output, Expected, Any] - scores: Sequence[EvalScorer[Input, Output, Expected] | BatchScorer[Input, Output, Expected, Any]] + task: EvalTask[Input, Output, Expected] | WorkflowTask[Input, Output, Expected, Any] + scores: Sequence[EvalScorer[Input, Output, Expected] | WorkflowScorer[Input, Output, Expected, Any]] classifiers: Sequence[EvalClassifier[Input, Output, Expected]] case_id: Callable[[EvalCase[Input, Expected]], str | Awaitable[str]] | None experiment_name: str | None trial_count: int + max_concurrency: int metadata: Metadata | None tags: Sequence[str] | None is_public: bool @@ -340,42 +359,43 @@ class _DurableEvalConfig(Generic[Input, Output, Expected]): state: BraintrustState | None -class DurableEval(Generic[Input, Output, Expected]): - """A durable evaluation definition that can be started and resumed.""" +class WorkflowEval(Generic[Input, Output, Expected]): + """A workflow evaluation definition that can be started and resumed.""" - def __init__(self, config: _DurableEvalConfig[Input, Output, Expected]) -> None: + def __init__(self, config: _WorkflowEvalConfig[Input, Output, Expected]) -> None: self._config = config async def start( self, parameters: Mapping[str, Any] | None = None, *, no_send_logs: bool = False - ) -> DurableEvalResult: - return await _DurableEvalRunner(self._config).start(parameters, no_send_logs=no_send_logs) + ) -> WorkflowEvalResult: + return await _WorkflowEvalRunner(self._config).start(parameters, no_send_logs=no_send_logs) - async def status(self, run_id: str) -> DurableEvalResult: - return await _DurableEvalRunner(self._config).status(run_id) + async def status(self, run_id: str) -> WorkflowEvalResult: + return await _WorkflowEvalRunner(self._config).status(run_id) - async def poll(self, run_id: str) -> DurableEvalResult: - return await _DurableEvalRunner(self._config).poll(run_id) + async def poll(self, run_id: str) -> WorkflowEvalResult: + return await _WorkflowEvalRunner(self._config).poll(run_id) - async def process_batch_result( - self, run_id: str, *, batch_id: str | None = None, external_id: str | None = None - ) -> DurableEvalResult: - return await _DurableEvalRunner(self._config).process_batch_result( - run_id, batch_id=batch_id, external_id=external_id + async def process_submission_result( + self, run_id: str, *, 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 ) -def define_durable_eval( +def define_workflow_eval( project_name: str, *, - store: DurableEvalStore, + store: WorkflowEvalStore, data: EvalData[Input, Expected], - task: EvalTask[Input, Output, Expected] | BatchTask[Input, Output, Expected, Any], - scores: Sequence[EvalScorer[Input, Output, Expected] | BatchScorer[Input, Output, Expected, Any]] | None = None, + task: EvalTask[Input, Output, Expected] | WorkflowTask[Input, Output, Expected, Any], + scores: Sequence[EvalScorer[Input, Output, Expected] | WorkflowScorer[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, trial_count: int = 1, + max_concurrency: int = DEFAULT_MAX_CONCURRENCY, metadata: Metadata | None = None, tags: Sequence[str] | None = None, is_public: bool = False, @@ -388,12 +408,25 @@ def define_durable_eval( summarize_scores: bool = True, parameters: EvalParameters | RemoteEvalParameters | None = None, state: BraintrustState | None = None, -) -> DurableEval[Input, Output, Expected]: - """Define an experimental evaluation that can pause across provider batches.""" +) -> 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. + Submission data and collected values must be JSON serializable. + Scorers and classifiers start once their case's task is persisted and logged. + + Each poll checks existing submissions once; newly submitted work waits for a + later invocation. Provider callbacks use max_concurrency (default 10). + Callback failures are raised after independent work advances. + Webhooks resume via process_submission_result with submission_id or external_id. + Collection callbacks must tolerate replay, including concurrent delivery. + """ if not isinstance(trial_count, int) or isinstance(trial_count, bool) or trial_count < 1: raise ValueError("trial_count must be a positive integer") - return DurableEval( - _DurableEvalConfig( + if not isinstance(max_concurrency, int) or isinstance(max_concurrency, bool) or max_concurrency < 1: + raise ValueError("max_concurrency must be a positive integer") + return WorkflowEval( + _WorkflowEvalConfig( project_name=project_name, store=store, data=data, @@ -403,6 +436,7 @@ def define_durable_eval( case_id=case_id, experiment_name=experiment_name, trial_count=trial_count, + max_concurrency=max_concurrency, metadata=metadata, tags=tags, is_public=is_public, @@ -428,7 +462,7 @@ def _decode(value: bytes) -> Any: def _json_value(value: Any) -> Any: - """Validate and normalize a value at the durable JSON boundary.""" + """Validate and normalize a value at the workflow JSON boundary.""" return _decode(_json_bytes(value)) @@ -461,10 +495,19 @@ def _score_fields(result: ScoreLike) -> dict[str, Any]: return {key: value for key, value in result.as_dict().items() if key not in ("metadata", "name")} -class _DurableEvalRunner(Generic[Input, Output, Expected]): - def __init__(self, config: _DurableEvalConfig[Input, Output, Expected]) -> None: +def _raise_callback_errors(errors: Sequence[BaseException]) -> None: + if len(errors) == 1: + raise errors[0] + raise RuntimeError( + "Workflow submission callbacks failed: " + "; ".join(str(error) for error in errors) + ) from errors[0] + + +class _WorkflowEvalRunner(Generic[Input, Output, Expected]): + def __init__(self, config: _WorkflowEvalConfig[Input, Output, Expected]) -> None: self.config = config self.store = config.store + self._callback_semaphore = asyncio.Semaphore(config.max_concurrency) self.eval_name = config.experiment_name or config.project_name self.definition_key = _stable_hex(config.project_name, self.eval_name, length=32) @@ -480,7 +523,7 @@ async def _required(self, run_id: str, kind: str, identifier: str | None = None) value = await self._read(run_id, kind, identifier) if value is None: target = f" {identifier!r}" if identifier is not None else "" - raise ValueError(f"Unknown durable evaluation {kind}{target} for run {run_id!r}") + raise ValueError(f"Unknown workflow evaluation {kind}{target} for run {run_id!r}") return value async def _write(self, run_id: str, kind: str, value: Any, identifier: str | None = None) -> None: @@ -491,7 +534,8 @@ async def _claim(self, run_id: str, action: str) -> bool: return result.created async def _call(self, fn: Callable[..., Any], *args: Any, **kwargs: Any) -> Any: - return await await_or_run(asyncio.get_running_loop(), fn, *args, **kwargs) + async with self._callback_semaphore: + return await await_or_run(asyncio.get_running_loop(), fn, *args, **kwargs) async def _flush_logs(self) -> None: state = self.config.state or _internal_get_global_state() @@ -565,36 +609,16 @@ async def _resolve_data(self, experiment: Experiment | None) -> list[EvalCase[In values.extend(data) return [value if isinstance(value, EvalCase) else EvalCase.from_dict(value) for value in values] - def _uses_batch_processor(self) -> bool: - return isinstance(self.config.task, BatchTask) or any( - isinstance(score, BatchScorer) for score in self.config.scores + def _uses_submission_processor(self) -> bool: + return isinstance(self.config.task, WorkflowTask) or any( + isinstance(score, WorkflowScorer) for score in self.config.scores ) - async def _save_completed(self, run: Mapping[str, Any], summary: ExperimentSummary) -> DurableEvalCompletedResult: + async def _save_completed(self, run: Mapping[str, Any], summary: ExperimentSummary) -> WorkflowEvalCompletedResult: completed = {**run, "status": "completed", "summary": _json_value(summary.as_dict())} await self._write(run["run_id"], "run", completed) - return DurableEvalCompletedResult( - run_id=run["run_id"], pending=DurableEvalPending(poll=0, webhook=0), summary=summary - ) - - async def _save_failed( - self, run: Mapping[str, Any], batch: Mapping[str, Any], error: Any - ) -> DurableEvalFailedResult: - normalized_error = _json_value(error) - failure = {"batch_id": batch["id"], "error": normalized_error} - failed_run = {**run, "status": "failed", "failure": failure} - await self._write(run["run_id"], "run", failed_run) - await self._write( - run["run_id"], - "batch", - {**batch, "status": "failed", "error": normalized_error}, - batch["id"], - ) - return DurableEvalFailedResult( - run_id=run["run_id"], - batch_id=batch["id"], - error=normalized_error, - pending=DurableEvalPending(poll=0, webhook=0), + return WorkflowEvalCompletedResult( + run_id=run["run_id"], pending=WorkflowEvalPending(poll=0, webhook=0), summary=summary ) async def _run_ordinary_evaluator( @@ -603,7 +627,7 @@ async def _run_ordinary_evaluator( data: list[EvalCase[Input, Expected]], experiment: Experiment | None, parameters: ValidatedParameters, - ) -> DurableEvalCompletedResult: + ) -> WorkflowEvalCompletedResult: evaluator = Evaluator( project_name=self.config.project_name, eval_name=self.eval_name, @@ -615,6 +639,7 @@ async def _run_ordinary_evaluator( metadata=self.config.metadata, tags=self.config.tags, trial_count=self.config.trial_count, + max_concurrency=self.config.max_concurrency, is_public=self.config.is_public, update=True, project_id=self.config.project_id, @@ -644,9 +669,9 @@ async def _persist_cases(self, run: dict[str, Any], data: Sequence[EvalCase[Inpu if not case_id and self.config.case_id is not None: case_id = await self._call(self.config.case_id, datum) if not isinstance(case_id, str) or not case_id: - raise ValueError("Every durable evaluation case must have a non-empty id or be assigned by case_id") + raise ValueError("Every workflow evaluation case must have a non-empty id or be assigned by case_id") if case_id in seen_case_ids: - raise ValueError(f"Durable evaluation case IDs must be unique; found duplicate {case_id!r}") + raise ValueError(f"Workflow evaluation case IDs must be unique; found duplicate {case_id!r}") seen_case_ids.add(case_id) trial_count = datum.trial_count if datum.trial_count is not None else self.config.trial_count @@ -673,7 +698,7 @@ async def _persist_cases(self, run: dict[str, Any], data: Sequence[EvalCase[Inpu run["case_ids"] = item_ids await self._write(run["run_id"], "run", run) - async def start(self, parameters: Mapping[str, Any] | None, *, no_send_logs: bool) -> DurableEvalResult: + async def start(self, parameters: Mapping[str, Any] | None, *, no_send_logs: bool) -> WorkflowEvalResult: validated_parameters = validate_parameters(parameters or {}, self.config.parameters) run_id = str(uuid.uuid4()) run = { @@ -686,140 +711,134 @@ async def start(self, parameters: Mapping[str, Any] | None, *, no_send_logs: boo } created = await self.store.get_or_set(self._key(run_id, "run"), _json_bytes(run)) if not created.created: - raise RuntimeError(f"Durable evaluation run ID collision: {run_id}") + raise RuntimeError(f"Workflow evaluation run ID collision: {run_id}") experiment = self._experiment(run) data = await self._resolve_data(experiment) - if not self._uses_batch_processor(): + if not self._uses_submission_processor(): return await self._run_ordinary_evaluator(run, data, experiment, validated_parameters) await self._persist_cases(run, data) return await self._advance(run, experiment=experiment) - async def status(self, run_id: str) -> DurableEvalResult: + async def status(self, run_id: str) -> WorkflowEvalResult: run = await self._required(run_id, "run") if run["status"] == "completed": - return DurableEvalCompletedResult( + return WorkflowEvalCompletedResult( run_id=run_id, - pending=DurableEvalPending(poll=0, webhook=0), + pending=WorkflowEvalPending(poll=0, webhook=0), summary=ExperimentSummary.from_dict_deep(run["summary"]), ) - if run["status"] == "failed": - failure = run["failure"] - return DurableEvalFailedResult( - run_id=run_id, - batch_id=failure["batch_id"], - error=failure.get("error"), - pending=DurableEvalPending(poll=0, webhook=0), - ) return await self._waiting_status(run) - async def poll(self, run_id: str) -> DurableEvalResult: + async def poll(self, run_id: str) -> WorkflowEvalResult: run = await self._required(run_id, "run") - if run["status"] in ("completed", "failed"): + if run["status"] == "completed": return await self.status(run_id) - # Snapshot first: newly submitted downstream batches wait until the next poll call. - batches = await self._batch_records(run) - for batch, processor in batches: - if batch["status"] != "submitted" or batch["mode"] != "poll": - continue + # Snapshot first: newly submitted downstream work waits for the next poll call. + submissions = await self._submission_records(run) + + async def poll_submission(submission: dict[str, Any], processor: Any) -> None: completion = processor.completion - assert isinstance(completion, BatchCompletionPoll) - context = BatchContext(run_id=run_id, batch_id=batch["id"]) - result = await self._call(completion.poll, batch["submission_data"], context) - if not isinstance(result, BatchPollResult): - if isinstance(result, Mapping): - result = BatchPollResult(**result) - else: - raise TypeError("Batch poll callback must return BatchPollResult") + assert isinstance(completion, WorkflowSubmissionCompletionPoll) + context = WorkflowSubmissionContext(run_id=run_id, submission_id=submission["id"]) + result = await self._call(completion.poll, submission["submission_data"], context) + if isinstance(result, Mapping): + result = WorkflowSubmissionPoll(**result) + if not isinstance(result, WorkflowSubmissionPoll): + raise TypeError("Submission poll callback must return WorkflowSubmissionPoll") if result.status not in ("pending", "complete", "failed"): - raise ValueError(f"Batch poll callback returned unsupported status {result.status!r}") + raise ValueError(f"Submission poll callback returned unsupported status {result.status!r}") if result.status == "failed": - return await self._save_failed(run, batch, result.error) + if isinstance(result.error, BaseException): + raise result.error + raise RuntimeError(str(result.error)) if result.status == "complete": - await self._collect_batch(run, batch, processor) - return await self._advance(run) - - async def process_batch_result( - self, run_id: str, *, batch_id: str | None, external_id: str | None - ) -> DurableEvalResult: - if batch_id is None and external_id is None: - raise ValueError("process_batch_result requires batch_id or external_id") + await self._collect_submission(run, submission, processor) + + results = await asyncio.gather( + *( + poll_submission(submission, processor) + for submission, processor in submissions + if submission["status"] == "submitted" and submission["mode"] == "poll" + ), + return_exceptions=True, + ) + # Persist independent progress before reporting provider callback failures. + errors = [result for result in results if isinstance(result, BaseException)] + try: + status = await self._advance(run) + except Exception as error: + errors.append(error) + if errors: + _raise_callback_errors(errors) + return status + + async def process_submission_result( + self, run_id: str, *, 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") run = await self._required(run_id, "run") - if run["status"] in ("completed", "failed"): + 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: + 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 + if submission is None: + raise ValueError("No submission matches this result") + if external_id is not None and submission.get("external_id") != external_id: + raise ValueError("submission_id and external_id identify different submissions") + if run["status"] == "completed": return await self.status(run_id) - batches = await self._batch_records(run) - by_batch = None - if batch_id is not None: - by_batch = next(((batch, processor) for batch, processor in batches if batch["id"] == batch_id), None) - - by_external = None - if external_id is not None: - by_external = next( - ((batch, processor) for batch, processor in batches if batch.get("external_id") == external_id), - None, + processor = ( + self.config.task + if submission["kind"] == "task" + else next( + score + for score in self.config.scores + if isinstance(score, WorkflowScorer) 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"]]) - if by_batch is not None and by_external is not None and by_batch[0]["id"] != by_external[0]["id"]: - raise ValueError("batch_id and external_id identify different batches") - match = by_batch or by_external - if match is None: - raise ValueError("No submitted durable batch matches this result") - batch, processor = match - if batch["status"] != "complete": - await self._collect_batch(run, batch, processor) - return await self._advance(run) - - def _batch_id(self, run_id: str, kind: str, index: int, item_ids: Sequence[str]) -> str: - return f"batch-{_stable_hex(run_id, kind, str(index), *item_ids, length=32)}" - - async def _task_complete(self, run: Mapping[str, Any]) -> bool: - for item_id in run["case_ids"]: - if await self._read(run["run_id"], "task-result", item_id) is None: - return False - return True - - async def _batch_specs(self, run: Mapping[str, Any]) -> list[tuple[dict[str, Any], Any]]: - run_id = run["run_id"] - item_ids = run["case_ids"] + def _submission_specs(self, run: Mapping[str, Any]) -> list[tuple[dict[str, Any], Any]]: specs: list[tuple[dict[str, Any], Any]] = [] - if isinstance(self.config.task, BatchTask): - for index, offset in enumerate(range(0, len(item_ids), self.config.task.batch_size)): - ids = item_ids[offset : offset + self.config.task.batch_size] + for item_id in run["case_ids"]: + if isinstance(self.config.task, WorkflowTask): specs.append( ( { - "id": self._batch_id(run_id, "task", index, ids), + "id": f"submission-{_stable_hex(run['run_id'], 'task', item_id, length=32)}", "kind": "task", - "item_ids": ids, + "item_id": item_id, }, self.config.task, ) ) - if await self._task_complete(run): for scorer in self.config.scores: - if not isinstance(scorer, BatchScorer): - continue - for index, offset in enumerate(range(0, len(item_ids), scorer.batch_size)): - ids = item_ids[offset : offset + scorer.batch_size] + if isinstance(scorer, WorkflowScorer): specs.append( ( { - "id": self._batch_id(run_id, f"score:{scorer.name}", index, ids), + "id": f"submission-{_stable_hex(run['run_id'], 'score', scorer.name, item_id, length=32)}", "kind": "score", "scorer_name": scorer.name, - "item_ids": ids, + "item_id": item_id, }, scorer, ) ) return specs - async def _batch_records(self, run: Mapping[str, Any]) -> list[tuple[dict[str, Any], Any]]: + async def _submission_records(self, run: Mapping[str, Any]) -> list[tuple[dict[str, Any], Any]]: records: list[tuple[dict[str, Any], Any]] = [] - for spec, processor in await self._batch_specs(run): - record = await self._read(run["run_id"], "batch", spec["id"]) + 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)) return records @@ -828,10 +847,10 @@ async def _case(self, run_id: str, item_id: str) -> dict[str, Any]: async def _task_item( self, run: Mapping[str, Any], item_id: str, parameters: ValidatedParameters | None - ) -> BatchTaskItem[Any, Any]: + ) -> WorkflowTaskItem[Any, Any]: case = await self._case(run["run_id"], item_id) datum = case["datum"] - return BatchTaskItem( + return WorkflowTaskItem( id=item_id, input=datum["input"], expected=datum.get("expected"), @@ -841,11 +860,11 @@ async def _task_item( trial_index=case["trial_index"], ) - async def _scorer_item(self, run: Mapping[str, Any], item_id: str) -> BatchScorerItem[Any, Any, Any]: + async def _scorer_item(self, run: Mapping[str, Any], item_id: str) -> WorkflowScorerItem[Any, Any, Any]: case = await self._case(run["run_id"], item_id) task = await self._required(run["run_id"], "task-result", item_id) datum = case["datum"] - return BatchScorerItem( + return WorkflowScorerItem( id=item_id, input=datum["input"], output=task["output"], @@ -855,89 +874,89 @@ async def _scorer_item(self, run: Mapping[str, Any], item_id: str) -> BatchScore trial_index=case["trial_index"], ) - async def _submit_batch(self, run: Mapping[str, Any], spec: dict[str, Any], processor: Any) -> None: + async def _submit_submission(self, run: Mapping[str, Any], spec: dict[str, Any], processor: Any) -> None: run_id = run["run_id"] - if await self._read(run_id, "batch", spec["id"]) is not None: + existing = await self._read(run_id, "submission", spec["id"]) + if existing is not None: + await self._persist_submission_progress(run_id, existing) return if not await self._claim(run_id, f"submit:{spec['id']}"): return parameters = self._parameters(run) if spec["kind"] == "task": - items = [await self._task_item(run, item_id, parameters) for item_id in spec["item_ids"]] + item = await self._task_item(run, spec["item_id"], parameters) else: - items = [await self._scorer_item(run, item_id) for item_id in spec["item_ids"]] - context = BatchContext(run_id=run_id, batch_id=spec["id"]) - submission_data = await self._call(processor.submit, items, context) + 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 external_id = None - if isinstance(completion, BatchCompletionWebhook): + 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: - raise ValueError("Batch webhook get_external_id must return a non-empty string") - await self._write( - run_id, - "batch", - { - **spec, - "submission_data": submission_data, - "external_id": external_id, - "mode": completion.mode, - "status": "submitted", - }, - spec["id"], - ) - - async def _collect_batch(self, run: Mapping[str, Any], batch: dict[str, Any], processor: Any) -> None: - run_id = run["run_id"] - if batch["status"] == "complete": - return - context = BatchContext(run_id=run_id, batch_id=batch["id"]) - raw_results = await self._call(processor.collect, batch["submission_data"], context) - if not isinstance(raw_results, list): - raise TypeError("Batch collect callback must return a list") - result_type = BatchTaskResult if batch["kind"] == "task" else BatchScorerResult - results = [value if isinstance(value, result_type) else result_type(**value) for value in raw_results] - result_ids = [value.id for value in results] - expected_ids = batch["item_ids"] - if len(result_ids) != len(set(result_ids)): - raise ValueError(f"Batch {batch['id']} returned duplicate item IDs") - unknown = sorted(set(result_ids) - set(expected_ids)) - missing = sorted(set(expected_ids) - set(result_ids)) - if unknown or missing: - raise ValueError( - f"Batch {batch['id']} result IDs did not match submitted items; unknown={unknown}, missing={missing}" + 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, spec["id"]) + await self._persist_submission_progress(run_id, submission) + + 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"]) + await self.store.add_to_set(self._key(run_id, "progress", f"{submission['mode']}/submitted"), submission["id"]) + if submission["status"] == "complete": + await self.store.add_to_set( + self._key(run_id, "progress", f"{submission['mode']}/complete"), submission["id"] ) - if batch["kind"] == "task": - for result in results: - case = await self._case(run_id, result.id) + async def _collect_submission(self, run: Mapping[str, Any], submission: dict[str, Any], processor: Any) -> None: + run_id = run["run_id"] + if submission["status"] != "complete": + context = WorkflowSubmissionContext(run_id=run_id, submission_id=submission["id"]) + result = await self._call(processor.collect, submission["submission_data"], context) + result_type = WorkflowTaskResult if submission["kind"] == "task" else WorkflowScorerResult + if isinstance(result, Mapping): + result = result_type(**result) + if not isinstance(result, result_type): + raise TypeError(f"Submission collect callback must return a single {result_type.__name__}") + item_id = submission["item_id"] + if submission["kind"] == "task": + 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}, - result.id, + { + "output": _json_value(result.output), + "metadata": metadata, + "tags": tags, + }, + item_id, ) - else: - for result in results: + else: await self._write( - run_id, - _stage_kind("score-result", batch["scorer_name"]), - _json_value(result.score), - result.id, + run_id, _stage_kind("score-result", submission["scorer_name"]), _json_value(result.score), item_id ) - batch["status"] = "complete" - await self._write(run_id, "batch", batch, batch["id"]) - - async def _waiting_status(self, run: Mapping[str, Any]) -> DurableEvalWaitingResult: - pending = {"poll": 0, "webhook": 0} - for batch, _ in await self._batch_records(run): - if batch["status"] == "submitted": - pending[batch["mode"]] += 1 - return DurableEvalWaitingResult( - run_id=run["run_id"], pending=DurableEvalPending(poll=pending["poll"], webhook=pending["webhook"]) + submission["status"] = "complete" + await self._write(run_id, "submission", submission, submission["id"]) + # Replay repairs an interrupted progress update without collecting again. + await self._persist_submission_progress(run_id, submission) + + async def _waiting_status(self, run: Mapping[str, Any]) -> WorkflowEvalWaitingResult: + pending = {} + for mode in ("poll", "webhook"): + completed = await self.store.get_set_size(self._key(run["run_id"], "progress", f"{mode}/complete")) + submitted = await self.store.get_set_size(self._key(run["run_id"], "progress", f"{mode}/submitted")) + pending[mode] = submitted - completed + return WorkflowEvalWaitingResult( + run_id=run["run_id"], pending=WorkflowEvalPending(poll=pending["poll"], webhook=pending["webhook"]) ) def _span_ids(self, run_id: str, item_id: str, stage: str) -> tuple[str, str, str]: @@ -984,7 +1003,7 @@ def _start_root(self, experiment: Experiment | None, run: Mapping[str, Any], cas expected=datum.get("expected"), metadata={ **case["metadata"], - "durable_eval": { + "workflow_eval": { "run_id": run["run_id"], "case_id": case["case_id"], "trial_index": case["trial_index"], @@ -1117,7 +1136,7 @@ async def _trace_for_case(self, run: Mapping[str, Any], item_id: str) -> LocalTr return None components = SpanComponentsV4.from_str(exported) if not components.root_span_id: - raise ValueError("Persisted durable evaluation root span is missing its root span ID") + raise ValueError("Persisted workflow evaluation root span is missing its root span ID") trace_state = self.config.state or _internal_get_global_state() async def ensure_spans_flushed() -> None: @@ -1257,117 +1276,90 @@ async def _log_classifier(self, run: Mapping[str, Any], case: dict[str, Any], na await self._flush_logs() await self._write(run_id, log_kind, True, item_id) - async def _advance_tasks( + async def _advance_case( self, run: Mapping[str, Any], + item_id: str, experiment: Experiment | None, - parameters: ValidatedParameters | None, - ) -> tuple[Experiment | None, bool]: - run_id = run["run_id"] - if isinstance(self.config.task, BatchTask): - for spec, processor in await self._batch_specs(run): - if spec["kind"] == "task": - await self._submit_batch(run, spec, processor) - else: - if experiment is None and not run["no_send_logs"]: - experiment = self._experiment(run) - for item_id in run["case_ids"]: - case = await self._case(run_id, item_id) - await self._run_ordinary_task(run, case, experiment, parameters) - - if not await self._task_complete(run): - return experiment, False - - if experiment is None and not run["no_send_logs"]: - experiment = self._experiment(run) - if isinstance(self.config.task, BatchTask): - for item_id in run["case_ids"]: - await self._log_task(run, await self._case(run_id, item_id), experiment) - for item_id in run["case_ids"]: - if await self._read(run_id, "task-log", item_id) is None: - return experiment, False - return experiment, True - - async def _advance_scorers( - self, - run: Mapping[str, Any], scorers: Sequence[Any], scorer_names: Sequence[str], - ) -> None: - run_id = run["run_id"] - for spec, processor in await self._batch_specs(run): - if spec["kind"] == "score": - await self._submit_batch(run, spec, processor) - - for scorer, name in zip(scorers, scorer_names): - for item_id in run["case_ids"]: - if isinstance(scorer, BatchScorer): - if await self._read(run_id, _stage_kind("score-result", name), item_id) is not None: - await self._log_score(run, await self._case(run_id, item_id), name) - else: - await self._run_ordinary_score(run, await self._case(run_id, item_id), scorer, name) - - async def _advance_classifiers( - self, - run: Mapping[str, Any], - classifiers: Sequence[EvalClassifier[Input, Output, Expected]], classifier_names: Sequence[str], - ) -> None: + ) -> bool: run_id = run["run_id"] - for classifier, name in zip(classifiers, classifier_names): - for item_id in run["case_ids"]: - await self._run_classifier(run, await self._case(run_id, item_id), classifier, name) + case = await self._case(run_id, item_id) + case_run = {**run, "case_ids": [item_id]} + if isinstance(self.config.task, WorkflowTask): + for spec, processor in self._submission_specs(case_run): + if spec["kind"] == "task": + await self._submit_submission(run, spec, processor) + if await self._read(run_id, "task-result", item_id) is None: + return False + await self._log_task(run, case, experiment) + else: + await self._run_ordinary_task(run, case, experiment, self._parameters(run)) + if await self._read(run_id, "task-log", item_id) is None: + return False - async def _stage_records_complete(self, run: Mapping[str, Any], prefix: str, names: Sequence[str]) -> bool: + for spec, processor in self._submission_specs(case_run): + if spec["kind"] == "score": + await self._submit_submission(run, spec, processor) complete = True - for name in names: - result_kind = _stage_kind(f"{prefix}-result", name) - log_kind = _stage_kind(f"{prefix}-log", name) - for item_id in run["case_ids"]: - if ( - await self._read(run["run_id"], result_kind, item_id) is None - or await self._read(run["run_id"], log_kind, item_id) is None - ): + for scorer, name in zip(scorers, scorer_names): + if isinstance(scorer, WorkflowScorer): + if await self._read(run_id, _stage_kind("score-result", name), item_id) is None: complete = False + continue + 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: + complete = False + for classifier, name in zip(self.config.classifiers, classifier_names): + await self._run_classifier(run, case, classifier, name) + if await self._read(run_id, _stage_kind("classification-log", name), item_id) is None: + complete = False + if complete: + await self.store.add_to_set(self._key(run_id, "progress", "cases"), item_id) return complete - async def _advance(self, run: Mapping[str, Any], *, experiment: Experiment | None = None) -> DurableEvalResult: + async def _advance( + self, run: Mapping[str, Any], *, experiment: Experiment | None = None, item_ids: Sequence[str] | None = None + ) -> WorkflowEvalResult: run = await self._required(run["run_id"], "run") run_id = run["run_id"] - if run["status"] in ("completed", "failed"): + if run["status"] == "completed": return await self.status(run_id) - experiment, tasks_complete = await self._advance_tasks(run, experiment, self._parameters(run)) - if not tasks_complete: - return await self._waiting_status(run) - + if experiment is None and not run["no_send_logs"]: + experiment = self._experiment(run) resolved_scores = [ score() if inspect.isclass(score) and is_scorer(score) else score for score in self.config.scores ] scorer_names = [ - score.name if isinstance(score, BatchScorer) else _scorer_name(score, index) + score.name if isinstance(score, WorkflowScorer) else _scorer_name(score, index) for index, score in enumerate(resolved_scores) ] if len(scorer_names) != len(set(scorer_names)): - raise ValueError("Durable evaluation scorer names must be unique") - - await self._advance_scorers(run, resolved_scores, scorer_names) - - classifiers = list(self.config.classifiers) - classifier_names = [_classifier_name(classifier, index) for index, classifier in enumerate(classifiers)] + raise ValueError("Workflow evaluation scorer names must be unique") + classifier_names = [ + _classifier_name(classifier, index) for index, classifier in enumerate(self.config.classifiers) + ] if len(classifier_names) != len(set(classifier_names)): - raise ValueError("Durable evaluation classifier names must be unique") - await self._advance_classifiers(run, classifiers, classifier_names) - - scores_complete = await self._stage_records_complete(run, "score", scorer_names) - classifiers_complete = await self._stage_records_complete(run, "classification", classifier_names) - if not scores_complete or not classifiers_complete: + 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 errors: + _raise_callback_errors(errors) + if await self.store.get_set_size(self._key(run_id, "progress", "cases")) != len(run["case_ids"]): return await self._waiting_status(run) - if not await self._claim(run_id, "finalize"): - current = await self._required(run_id, "run") - if current["status"] == "completed": - return await self.status(run_id) - return await self._waiting_status(current) + return await self.status(run_id) summary = await self._summary(run, scorer_names, classifier_names, experiment) return await self._save_completed(run, summary) @@ -1427,25 +1419,24 @@ async def _summary( __all__ = [ - "BatchCompletionPoll", - "BatchCompletionWebhook", - "BatchContext", - "BatchPollResult", - "BatchScorer", - "BatchScorerItem", - "BatchScorerResult", - "BatchTask", - "BatchTaskItem", - "BatchTaskResult", - "DurableEval", - "DurableEvalCompletedResult", - "DurableEvalFailedResult", - "DurableEvalMemoryStore", - "DurableEvalPending", - "DurableEvalRedisStore", - "DurableEvalResult", - "DurableEvalStore", - "DurableEvalStoreEntry", - "DurableEvalWaitingResult", - "define_durable_eval", + "WorkflowSubmissionCompletionPoll", + "WorkflowSubmissionCompletionWebhook", + "WorkflowSubmissionContext", + "WorkflowSubmissionPoll", + "WorkflowScorer", + "WorkflowScorerItem", + "WorkflowScorerResult", + "WorkflowTask", + "WorkflowTaskItem", + "WorkflowTaskResult", + "WorkflowEval", + "WorkflowEvalCompletedResult", + "WorkflowEvalMemoryStore", + "WorkflowEvalPending", + "WorkflowEvalRedisStore", + "WorkflowEvalResult", + "WorkflowEvalStore", + "WorkflowEvalStoreEntry", + "WorkflowEvalWaitingResult", + "define_workflow_eval", ]