diff --git a/langfuse/__init__.py b/langfuse/__init__.py index 23c77c169..e158e6d4a 100644 --- a/langfuse/__init__.py +++ b/langfuse/__init__.py @@ -58,7 +58,30 @@ EvaluatorStats, MapperFunction, ) -from langfuse.experiment import Evaluation, RegressionError, RunnerContext +from langfuse.experiment import ( + AttemptRecord, + CheckpointCorruptedError, + CheckpointMismatchError, + CheckpointStore, + Evaluation, + EvaluatorFailure, + EvaluatorTimeoutError, + ExecutionPolicy, + ExecutionSummary, + ExperimentExecutionError, + FailureBudgetExceededError, + FailurePolicy, + FileCheckpointStore, + InMemoryCheckpointStore, + ItemFailure, + RegressionError, + RetryPolicy, + RunnerContext, + RunTimeoutError, + TaskTimeoutError, + TimeoutPolicy, + is_retryable_error, +) from ._client import client as _client_module from ._client.attributes import LangfuseOtelSpanAttributes @@ -132,6 +155,25 @@ "BatchEvaluationResult", "RunnerContext", "RegressionError", + "ExecutionPolicy", + "RetryPolicy", + "TimeoutPolicy", + "FailurePolicy", + "FileCheckpointStore", + "InMemoryCheckpointStore", + "CheckpointStore", + "ExecutionSummary", + "ItemFailure", + "EvaluatorFailure", + "AttemptRecord", + "ExperimentExecutionError", + "TaskTimeoutError", + "EvaluatorTimeoutError", + "RunTimeoutError", + "FailureBudgetExceededError", + "CheckpointMismatchError", + "CheckpointCorruptedError", + "is_retryable_error", "__version__", "is_default_export_span", "is_langfuse_span", diff --git a/langfuse/_client/client.py b/langfuse/_client/client.py index 5d128567c..fef45d0ad 100644 --- a/langfuse/_client/client.py +++ b/langfuse/_client/client.py @@ -4,10 +4,13 @@ """ import asyncio +import inspect import json import logging import os +import random import re +import time import urllib.parse import uuid from datetime import datetime, timezone @@ -90,10 +93,26 @@ _set_span_attributes_within_limit, ) from langfuse._client.utils import ( - get_sha256_hash_hex, get_string_span_attribute, run_async_safely, ) +from langfuse._experiment.executor import ( + ExperimentExecutor, + ItemProcessingFailed, + _run_with_retry, +) +from langfuse._experiment.item_key import ( + experiment_item_id_for as _experiment_item_id_for, +) +from langfuse._experiment.results import ( + AttemptRecord as _AttemptRecord, +) +from langfuse._experiment.results import ( + EvaluatorFailure as _EvaluatorFailure, +) +from langfuse._experiment.results import ( + ItemFailure as _ItemFailure, +) from langfuse._utils import _get_timestamp, json_path from langfuse._utils.environment import get_common_release_envs from langfuse._utils.parse_error import handle_fern_exception @@ -128,6 +147,7 @@ from langfuse.experiment import ( Evaluation, EvaluatorFunction, + ExecutionPolicy, ExperimentData, ExperimentItem, ExperimentItemResult, @@ -2611,6 +2631,7 @@ def run_experiment( max_concurrency: int = 50, metadata: Optional[Dict[str, str]] = None, experiment_id: Optional[str] = None, + execution: Optional[ExecutionPolicy] = None, _dataset_version: Optional[datetime] = None, ) -> ExperimentResult: """Run an experiment on a dataset with automatic tracing and evaluation. @@ -2813,7 +2834,33 @@ def average_accuracy(*, item_results, **kwargs): - When using Langfuse datasets, results are automatically linked to the dataset for easy comparison - This method works in both sync and async contexts (Jupyter notebooks, web apps, etc.) - Async execution is handled automatically with smart event loop detection + + Raises: + ValueError: If required parameters are missing or invalid, if + `experiment_id` is empty or whitespace-only, or if `execution` + enables resume on local (non-dataset) data without an explicit + `run_name` or `experiment_id` (a new run name and a random + experiment id are generated per call, so the previous run's + checkpoint could never be found). """ + if ( + execution is not None + and execution.resume + and run_name is None + and experiment_id is None + and not any( + not isinstance(item, dict) and getattr(item, "dataset_id", None) + for item in data + ) + ): + message = ( + "resume=True on local (non-dataset) data requires an explicit " + "run_name or experiment_id: without one, each call generates a " + "new run name and a random experiment id, so the checkpoint " + "written by a previous run can never be found." + ) + raise ValueError(message) + return cast( ExperimentResult, run_async_safely( @@ -2831,6 +2878,7 @@ def average_accuracy(*, item_results, **kwargs): max_concurrency=max_concurrency, metadata=metadata, experiment_id=experiment_id, + execution=execution, dataset_version=_dataset_version, ), ), @@ -2850,6 +2898,7 @@ async def _run_experiment_async( max_concurrency: int, metadata: Optional[Dict[str, Any]] = None, experiment_id: Optional[str] = None, + execution: Optional[ExecutionPolicy] = None, dataset_version: Optional[datetime] = None, ) -> ExperimentResult: if experiment_id is not None and ( @@ -2880,56 +2929,80 @@ async def _run_experiment_async( ), None, ) + # The id is deterministic when the caller supplied it or when it is + # derived from project + dataset + run name; otherwise (local data + # without an explicit id) it is random per run, which the executor + # must know to key checkpoints and adopt a resumed id. + experiment_id_is_deterministic = experiment_id is not None if experiment_id is None: experiment_id = self._create_experiment_id( project_id=project_id, dataset_id=experiment_dataset_id, run_name=run_name, ) + experiment_id_is_deterministic = ( + experiment_dataset_id is not None and project_id is not None + ) experiment_item_version = ( self._format_experiment_item_version(dataset_version) if dataset_version is not None else None ) - # Set up concurrency control - semaphore = asyncio.Semaphore(max_concurrency) - - # Process all items - async def process_item(item: ExperimentItem) -> ExperimentItemResult: - async with semaphore: - return await self._process_experiment_item( - item, - task, - evaluators, - composite_evaluator, - experiment_id, - name, - run_name, - description, - metadata, - experiment_item_version, - ) - - # Run all items concurrently - tasks = [process_item(item) for item in data] - item_results = await asyncio.gather(*tasks, return_exceptions=True) - - # Filter out any exceptions and log errors - valid_results: List[ExperimentItemResult] = [] - for i, result in enumerate(item_results): - if isinstance(result, Exception): - langfuse_logger.error("Item %s failed: %s", i, result) - elif isinstance(result, ExperimentItemResult): - valid_results.append(result) # type: ignore + async def processor( + item: ExperimentItem, + item_key: str, + processor_experiment_id: str, + policy: Optional[ExecutionPolicy], + attempt_hook: Optional[Callable[[_AttemptRecord], None]], + evaluator_failures: List[_EvaluatorFailure], + ) -> ExperimentItemResult: + return await self._process_experiment_item( + item, + task, + evaluators, + composite_evaluator, + processor_experiment_id, + name, + run_name, + description, + metadata, + experiment_item_version, + item_key=item_key, + policy=policy, + attempt_hook=attempt_hook, + evaluator_failures=evaluator_failures, + ) - # Run experiment-level evaluators + executor = ExperimentExecutor( + experiment_id=experiment_id, + experiment_id_deterministic=experiment_id_is_deterministic, + experiment_name=name, + run_name=run_name, + data=list(data), + dataset_version=experiment_item_version, + policy=execution, + processor=processor, + max_concurrency=max_concurrency, + ) + exec_result = await executor.run(task=task, evaluators=evaluators) + # A resumed local-data run adopts the experiment id recorded in the + # checkpoint so both runs share one experiment. + experiment_id = exec_result.experiment_id + + # Run experiment-level evaluators. Each run evaluator receives the + # successful item_results (for backwards compatibility) and, only if + # its signature accepts it, the recorded failures so it can compute + # honest rates. run_evaluations: List[Evaluation] = [] for run_evaluator in run_evaluators: try: - evaluations = await _run_evaluator( - run_evaluator, item_results=valid_results + kwargs = self._run_evaluator_kwargs( + run_evaluator, + exec_result.item_results, + exec_result.failures, ) + evaluations = await _run_evaluator(run_evaluator, **kwargs) run_evaluations.extend(evaluations) except Exception as e: langfuse_logger.error("Run evaluator failed: %s", e) @@ -2963,12 +3036,37 @@ async def process_item(item: ExperimentItem) -> ExperimentItemResult: name=name, run_name=run_name, description=description, - item_results=valid_results, + item_results=exec_result.item_results, run_evaluations=run_evaluations, experiment_id=experiment_id, experiment_url=experiment_url, + failures=exec_result.failures, + evaluator_failures=exec_result.evaluator_failures, + execution=exec_result.summary, ) + @staticmethod + def _run_evaluator_kwargs( + run_evaluator: RunEvaluatorFunction, + item_results: List[ExperimentItemResult], + failures: List[_ItemFailure], + ) -> Dict[str, Any]: + """Build kwargs for a run evaluator, passing ``failures`` only if the + evaluator declares an explicit ``failures`` parameter. + + Functions that merely accept ``**kwargs`` are left untouched so + pre-existing run evaluators keep receiving exactly what they received + before.""" + kwargs: Dict[str, Any] = {"item_results": item_results} + try: + sig = inspect.signature(run_evaluator) + except (TypeError, ValueError): + return kwargs + params = sig.parameters + if "failures" in params: + kwargs["failures"] = failures + return kwargs + def _create_experiment_id( self, *, @@ -3013,6 +3111,11 @@ async def _process_experiment_item( experiment_description: Optional[str], experiment_metadata: Optional[Dict[str, Any]] = None, experiment_item_version: Optional[str] = None, + *, + item_key: Optional[str] = None, + policy: Optional[ExecutionPolicy] = None, + attempt_hook: Optional[Callable[[_AttemptRecord], None]] = None, + evaluator_failures: Optional[List[_EvaluatorFailure]] = None, ) -> ExperimentItemResult: with self.start_as_current_observation(name="experiment-item-run") as span: try: @@ -3068,9 +3171,7 @@ async def _process_experiment_item( **experiment_run_metadata, } - experiment_item_id = ( - dataset_item_id or get_sha256_hash_hex(_serialize(input_data))[:16] - ) + experiment_item_id = _experiment_item_id_for(item) experiment_span_attributes = { k: v for k, v in { @@ -3125,15 +3226,69 @@ async def _process_experiment_item( # attributes, so the span attribute limit trims the # metadata instead of leaving no room for the output. task_span.update(metadata=final_observation_metadata) - try: - output = await _run_task(task, item) - except Exception as e: - task_span.update( - output=f"Error: {str(e)}", - level="ERROR", - status_message=str(e), + + task_policy = policy.task_retry if policy is not None else None + task_timeout = ( + policy.timeouts.task_timeout if policy is not None else None + ) + has_retry = ( + task_policy is not None and task_policy.max_attempts > 1 + ) + has_timeout = task_timeout is not None + + if not has_retry and not has_timeout: + # Default path: identical trace shape to today. + task_started_at = time.monotonic() + try: + output = await _run_task(task, item) + except Exception as e: + if attempt_hook is not None: + attempt_hook( + _AttemptRecord( + attempt=1, + started_at=task_started_at, + duration_s=( + time.monotonic() - task_started_at + ), + error_type=type(e).__name__, + error_message=str(e), + retry_after=None, + will_retry=False, + ) + ) + task_span.update( + output=f"Error: {str(e)}", + level="ERROR", + status_message=str(e), + ) + raise + if attempt_hook is not None: + attempt_hook( + _AttemptRecord( + attempt=1, + started_at=task_started_at, + duration_s=(time.monotonic() - task_started_at), + error_type=None, + error_message=None, + retry_after=None, + will_retry=False, + ) + ) + else: + output, task_attempts = await self._run_task_with_policy( + task=task, + item=item, + task_span=task_span, + task_policy=task_policy, + task_timeout=task_timeout, + policy=policy, + attempt_hook=attempt_hook, + item_key=item_key or experiment_item_id, ) - raise + if task_attempts is not None: + self._emit_attempt_observations( + task_span, task_attempts + ) task_span.update(output=output) @@ -3180,12 +3335,41 @@ async def _process_experiment_item( "evaluator_index": evaluator_index, }, ) as evaluator_span: + eval_policy = ( + policy.evaluator_retry + if policy is not None + else None + ) + eval_timeout = ( + policy.timeouts.evaluator_timeout + if policy is not None + else None + ) + has_eval_retry = ( + eval_policy is not None + and eval_policy.max_attempts > 1 + ) + has_eval_timeout = eval_timeout is not None try: - eval_results = await _run_evaluator( - evaluator, - _raise_on_error=True, - **evaluator_input, - ) + if not has_eval_retry and not has_eval_timeout: + eval_results = await _run_evaluator( + evaluator, + _raise_on_error=True, + **evaluator_input, + ) + else: + eval_results = ( + await self._run_evaluator_with_policy( + evaluator=evaluator, + evaluator_input=evaluator_input, + evaluator_span=evaluator_span, + eval_policy=eval_policy, + eval_timeout=eval_timeout, + policy=policy, + item_key=item_key or experiment_item_id, + evaluator_failures=evaluator_failures, + ) + ) evaluator_span.update( output=_serialize_evaluations(eval_results) ) @@ -3305,6 +3489,175 @@ async def _process_experiment_item( experiment_id=experiment_id, ) + async def _run_task_with_policy( + self, + *, + task: Callable, + item: ExperimentItem, + task_span: Any, + task_policy: Any, + task_timeout: Optional[float], + policy: Optional[ExecutionPolicy], + attempt_hook: Optional[Callable[[_AttemptRecord], None]], + item_key: str, + ) -> Tuple[Any, Optional[List[_AttemptRecord]]]: + """Run the task with retry/timeout, emitting attempt observations. + + Returns ``(output, attempts)`` on success. On terminal failure raises + ``ItemProcessingFailed`` carrying an ``ItemFailure`` with the + recorded attempts. + """ + rng = random.Random( + policy.seed if policy is not None and policy.seed is not None else None + ) + clock = time.monotonic + + async def factory() -> Any: + return await _run_task(task, item) + + try: + output, attempts = await _run_with_retry( + factory, + policy=task_policy, + timeout=task_timeout, + clock=clock, + sleep=asyncio.sleep, + rng=rng, + attempt_hook=attempt_hook, + ) + return output, attempts + except asyncio.TimeoutError: + failure = _ItemFailure( + item=item, + item_key=item_key, + stage="timeout", + error_type="TaskTimeoutError", + error_message="task timed out after %s seconds" % task_timeout, + attempts=[], + trace_id=task_span.trace_id, + ) + task_span.update( + output=f"Error: {failure.error_message}", + level="ERROR", + status_message=failure.error_message, + ) + raise ItemProcessingFailed(failure) + except Exception as e: + failure = _ItemFailure( + item=item, + item_key=item_key, + stage="task", + error_type=type(e).__name__, + error_message=str(e), + attempts=[], + trace_id=task_span.trace_id, + ) + task_span.update( + output=f"Error: {str(e)}", + level="ERROR", + status_message=str(e), + ) + raise ItemProcessingFailed(failure) + + async def _run_evaluator_with_policy( + self, + *, + evaluator: Callable, + evaluator_input: Dict[str, Any], + evaluator_span: Any, + eval_policy: Any, + eval_timeout: Optional[float], + policy: Optional[ExecutionPolicy], + item_key: str, + evaluator_failures: Optional[List[_EvaluatorFailure]], + ) -> List[Evaluation]: + """Run one evaluator with retry/timeout, recording an + ``EvaluatorFailure`` on terminal failure instead of raising.""" + rng = random.Random( + policy.seed if policy is not None and policy.seed is not None else None + ) + clock = time.monotonic + + async def factory() -> List[Evaluation]: + return await _run_evaluator( + evaluator, _raise_on_error=True, **evaluator_input + ) + + try: + eval_results, attempts = await _run_with_retry( + factory, + policy=eval_policy, + timeout=eval_timeout, + clock=clock, + sleep=asyncio.sleep, + rng=rng, + attempt_hook=None, + ) + return list(eval_results) + except asyncio.TimeoutError: + failure = _EvaluatorFailure( + item_key=item_key, + evaluator_name=_get_evaluator_name(evaluator), + error_type="EvaluatorTimeoutError", + error_message="evaluator timed out after %s seconds" % eval_timeout, + attempts=[], + ) + if evaluator_failures is not None: + evaluator_failures.append(failure) + evaluator_span.update( + output={"error": failure.error_message}, + level="ERROR", + status_message=failure.error_message, + ) + langfuse_logger.error("Evaluator failed: %s", failure.error_message) + return [] + except Exception as e: + failure = _EvaluatorFailure( + item_key=item_key, + evaluator_name=_get_evaluator_name(evaluator), + error_type=type(e).__name__, + error_message=str(e), + attempts=[], + ) + if evaluator_failures is not None: + evaluator_failures.append(failure) + evaluator_span.update( + output={"error": str(e)}, + level="ERROR", + status_message=str(e), + ) + langfuse_logger.error("Evaluator failed: %s", e) + return [] + + @staticmethod + def _emit_attempt_observations( + task_span: Any, attempts: List[_AttemptRecord] + ) -> None: + """Emit a WARNING child observation for each failed attempt that was + retried, so traces show attempts and backoff separately.""" + for record in attempts: + if record.error_type is None or not record.will_retry: + continue + with task_span.start_as_current_observation( + name="experiment-item-task-attempt", + as_type="span", + level="WARNING", + status_message=record.error_message or "", + metadata={ + "attempt": record.attempt, + "error_type": record.error_type, + "retry_after": record.retry_after, + }, + ) as attempt_obs: + attempt_obs.update( + output={ + "attempt": record.attempt, + "error_type": record.error_type, + "error_message": record.error_message, + "will_retry": record.will_retry, + } + ) + def _create_experiment_run_name( self, *, name: Optional[str] = None, run_name: Optional[str] = None ) -> str: @@ -3331,6 +3684,7 @@ def run_batched_evaluation( _add_observation_scores_to_trace: bool = False, resume_from: Optional[BatchEvaluationResumeToken] = None, verbose: bool = False, + execution: Optional[ExecutionPolicy] = None, ) -> BatchEvaluationResult: """Fetch observations from Langfuse and evaluate each of them. @@ -3396,6 +3750,12 @@ def run_batched_evaluation( (fetch failure or `max_items`). Continues exactly after the last processed page. Pass the same `filter`, or omit it to reuse the token's filter. Default: None. + execution: Optional `ExecutionPolicy` adding retries and timeouts for + mapper/evaluator/composite calls, a run timeout, a failure budget + that stops fetching further pages when exceeded (the result then has + `completed=False`, a resume token, and `execution.aborted=True`), + and detailed `failures` plus an `execution` summary on the result. + Default: None (behavior unchanged). Returns: @@ -3415,6 +3775,8 @@ def run_batched_evaluation( - error_summary: Error types and counts - has_more_items: True if max_items was reached and more items may exist - item_evaluations: Evaluations per observation ID + - failures: Detailed per-item failure records (stage, error, attempts) + - execution: Execution summary with retry/abort statistics Raises: ValueError: If a non-array filter is provided, or the resume token was @@ -3540,6 +3902,7 @@ def composite_evaluator(*, input, output, expected_output, metadata, evaluations max_retries=max_retries, verbose=verbose, resume_from=resume_from, + execution=execution, ) ), ) diff --git a/langfuse/_client/datasets.py b/langfuse/_client/datasets.py index 97b2e497d..dac8e9acd 100644 --- a/langfuse/_client/datasets.py +++ b/langfuse/_client/datasets.py @@ -8,6 +8,7 @@ from langfuse.batch_evaluation import CompositeEvaluatorFunction from langfuse.experiment import ( EvaluatorFunction, + ExecutionPolicy, ExperimentResult, RunEvaluatorFunction, TaskFunction, @@ -91,6 +92,7 @@ def run_experiment( max_concurrency: int = 50, metadata: Optional[Dict[str, Any]] = None, experiment_id: Optional[str] = None, + execution: Optional[ExecutionPolicy] = None, ) -> ExperimentResult: """Run an experiment on this Langfuse dataset with automatic tracking. @@ -311,5 +313,6 @@ def content_diversity(*, item_results, **kwargs): max_concurrency=max_concurrency, metadata=metadata, experiment_id=experiment_id, + execution=execution, _dataset_version=self.version, ) diff --git a/langfuse/_experiment/__init__.py b/langfuse/_experiment/__init__.py new file mode 100644 index 000000000..56c41b724 --- /dev/null +++ b/langfuse/_experiment/__init__.py @@ -0,0 +1,73 @@ +"""@private Experiment execution policies, checkpoints, and orchestration. + +Public names are re-exported from ``langfuse.experiment`` and +``langfuse.__init__``. Import from there instead of this private package. +""" + +from langfuse._experiment.checkpoint import ( + CheckpointItem, + CheckpointStore, + ExperimentCheckpoint, + ExperimentFingerprint, + FileCheckpointStore, + InMemoryCheckpointStore, + compute_data_digest, + experiment_checkpoint_key, +) +from langfuse._experiment.errors import ( + CheckpointCorruptedError, + CheckpointMismatchError, + EvaluatorTimeoutError, + ExperimentExecutionError, + FailureBudgetExceededError, + RunTimeoutError, + TaskTimeoutError, +) +from langfuse._experiment.item_key import ( + experiment_item_id_for, + local_item_input_hash, + make_item_keys, +) +from langfuse._experiment.policy import ( + ExecutionPolicy, + FailurePolicy, + RetryPolicy, + TimeoutPolicy, + is_retryable_error, +) +from langfuse._experiment.results import ( + AttemptRecord, + EvaluatorFailure, + ExecutionSummary, + ItemFailure, +) + +__all__ = [ + "AttemptRecord", + "CheckpointCorruptedError", + "CheckpointItem", + "CheckpointMismatchError", + "CheckpointStore", + "EvaluatorFailure", + "EvaluatorTimeoutError", + "ExecutionPolicy", + "ExecutionSummary", + "ExperimentCheckpoint", + "ExperimentExecutionError", + "ExperimentFingerprint", + "FailureBudgetExceededError", + "FailurePolicy", + "FileCheckpointStore", + "InMemoryCheckpointStore", + "ItemFailure", + "RetryPolicy", + "RunTimeoutError", + "TaskTimeoutError", + "TimeoutPolicy", + "compute_data_digest", + "experiment_checkpoint_key", + "experiment_item_id_for", + "is_retryable_error", + "local_item_input_hash", + "make_item_keys", +] diff --git a/langfuse/_experiment/checkpoint.py b/langfuse/_experiment/checkpoint.py new file mode 100644 index 000000000..219537c8b --- /dev/null +++ b/langfuse/_experiment/checkpoint.py @@ -0,0 +1,336 @@ +"""@private Checkpoint persistence for resumable experiment runs. + +Checkpoints record which items already completed (or permanently failed) so +a crashed or interrupted run can resume without re-paying for completed +task calls. The on-disk format is JSON with a sha256 integrity footer. +""" + +import json +import os +import re +import tempfile +import threading +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional, Protocol, runtime_checkable + +from langfuse._client.utils import get_sha256_hash_hex +from langfuse._experiment.errors import ( + CheckpointCorruptedError, + CheckpointMismatchError, +) + +_SCHEMA_VERSION = 1 +_INTEGRITY_KEY = "__integrity_sha256__" + +# Filenames are built only from [A-Za-z0-9_-]: anything else in the key +# (including dots, which could form ".." traversal sequences) is replaced. +_UNSAFE_KEY_RE = re.compile(r"[^A-Za-z0-9_-]") + + +def experiment_checkpoint_key(experiment_name: str, run_name: str) -> str: + """Deterministic checkpoint key for runs whose ``experiment_id`` is random. + + Dataset runs derive a stable ``experiment_id`` from project, dataset and + run name, so the id itself is the checkpoint key. Local (non-dataset) + runs get a fresh random id on every call, so the key is derived from the + experiment name and run name instead; resuming such a run therefore + requires the caller to pass an explicit ``run_name``. + """ + payload = json.dumps( + ["langfuse-experiment-checkpoint-v1", experiment_name, run_name], + separators=(",", ":"), + ensure_ascii=False, + ) + return get_sha256_hash_hex(payload)[:32] + + +def _sanitize_key(key: str) -> str: + """Return a filesystem-safe filename for ``key``. + + Every character outside ``[A-Za-z0-9_-]`` is replaced with ``_``, so + traversal sequences (``..``) and separators can never appear in the + result. The value is prefixed with a short hash of the raw key so + distinct keys cannot collide after sanitization. Keys containing NUL + are rejected outright. + """ + if "\x00" in key: + message = "checkpoint key must not contain NUL bytes: %r" % key + raise ValueError(message) + safe = _UNSAFE_KEY_RE.sub("_", key) + digest = get_sha256_hash_hex(key)[:16] + return "%s--%s.json" % (digest, safe) + + +@dataclass +class CheckpointItem: + """One item's persisted state.""" + + state: str # "completed" | "failed" + output_json: Optional[str] = None + evaluations: List[Dict[str, Any]] = field(default_factory=list) + trace_id: Optional[str] = None + attempts: int = 0 + error: Optional[str] = None + + def to_dict(self) -> Dict[str, Any]: + return { + "state": self.state, + "output_json": self.output_json, + "evaluations": list(self.evaluations), + "trace_id": self.trace_id, + "attempts": self.attempts, + "error": self.error, + } + + @classmethod + def from_dict(cls, data: Dict[str, Any]) -> "CheckpointItem": + return cls( + state=str(data.get("state", "failed")), + output_json=data.get("output_json"), + evaluations=list(data.get("evaluations") or []), + trace_id=data.get("trace_id"), + attempts=int(data.get("attempts") or 0), + error=data.get("error"), + ) + + +@dataclass +class ExperimentFingerprint: + """Identity of the task+evaluator set a checkpoint belongs to.""" + + task_name: str + evaluator_names: List[str] + data_size: int + data_digest: str # sha256 over sorted item keys + + def to_dict(self) -> Dict[str, Any]: + return { + "task_name": self.task_name, + "evaluator_names": list(self.evaluator_names), + "data_size": self.data_size, + "data_digest": self.data_digest, + } + + @classmethod + def from_dict(cls, data: Dict[str, Any]) -> "ExperimentFingerprint": + return cls( + task_name=str(data.get("task_name") or ""), + evaluator_names=list(data.get("evaluator_names") or []), + data_size=int(data.get("data_size") or 0), + data_digest=str(data.get("data_digest") or ""), + ) + + def matches(self, other: "ExperimentFingerprint") -> bool: + return self.to_dict() == other.to_dict() + + +def compute_data_digest(item_keys: List[str]) -> str: + joined = "\n".join(sorted(item_keys)) + return get_sha256_hash_hex(joined) + + +@dataclass +class ExperimentCheckpoint: + """Full persisted state of one experiment run.""" + + schema_version: int = _SCHEMA_VERSION + experiment_id: str = "" + experiment_name: str = "" + run_name: str = "" + fingerprint: Optional[ExperimentFingerprint] = None + items: Dict[str, CheckpointItem] = field(default_factory=dict) + created_at: float = 0.0 + updated_at: float = 0.0 + status: str = "running" # "running" | "completed" | "aborted" + + def to_payload(self) -> Dict[str, Any]: + return { + "schema_version": self.schema_version, + "experiment_id": self.experiment_id, + "experiment_name": self.experiment_name, + "run_name": self.run_name, + "fingerprint": (self.fingerprint.to_dict() if self.fingerprint else None), + "items": {k: v.to_dict() for k, v in self.items.items()}, + "created_at": self.created_at, + "updated_at": self.updated_at, + "status": self.status, + } + + def serialize(self) -> str: + payload = self.to_payload() + integrity = get_sha256_hash_hex( + json.dumps(payload, sort_keys=True, ensure_ascii=False) + ) + payload[_INTEGRITY_KEY] = integrity + return json.dumps(payload, ensure_ascii=False) + + @classmethod + def deserialize(cls, raw: str) -> "ExperimentCheckpoint": + try: + data = json.loads(raw) + except (json.JSONDecodeError, TypeError) as e: + raise CheckpointCorruptedError("invalid JSON: %s" % e) from e + if not isinstance(data, dict): + raise CheckpointCorruptedError("top-level value is not an object") + stored = data.pop(_INTEGRITY_KEY, None) + if stored is None: + raise CheckpointCorruptedError("missing integrity hash") + recomputed = get_sha256_hash_hex( + json.dumps(data, sort_keys=True, ensure_ascii=False) + ) + if not isinstance(stored, str) or stored != recomputed: + raise CheckpointCorruptedError("integrity hash mismatch") + schema_version = int(data.get("schema_version") or 0) + if schema_version != _SCHEMA_VERSION: + raise CheckpointCorruptedError( + "unsupported schema version %s" % schema_version + ) + fingerprint_data = data.get("fingerprint") + fingerprint = ( + ExperimentFingerprint.from_dict(fingerprint_data) + if fingerprint_data + else None + ) + items_raw = data.get("items") or {} + items = {str(k): CheckpointItem.from_dict(v) for k, v in items_raw.items()} + return cls( + schema_version=schema_version, + experiment_id=str(data.get("experiment_id") or ""), + experiment_name=str(data.get("experiment_name") or ""), + run_name=str(data.get("run_name") or ""), + fingerprint=fingerprint, + items=items, + created_at=float(data.get("created_at") or 0.0), + updated_at=float(data.get("updated_at") or 0.0), + status=str(data.get("status") or "running"), + ) + + def verify_fingerprint(self, expected: ExperimentFingerprint) -> None: + if self.fingerprint is None: + raise CheckpointMismatchError("fingerprint", expected.to_dict(), None) + if not self.fingerprint.matches(expected): + raise CheckpointMismatchError( + "fingerprint", expected.to_dict(), self.fingerprint.to_dict() + ) + + +@runtime_checkable +class CheckpointStore(Protocol): + """Protocol for pluggable checkpoint storage backends. + + ``experiment_key`` is the same key passed to :meth:`load`, so a store + never has to guess the lookup key from the checkpoint payload (for local + data the key differs from ``checkpoint.experiment_id``). + """ + + def load(self, experiment_key: str) -> Optional[ExperimentCheckpoint]: ... + + def save(self, checkpoint: ExperimentCheckpoint, experiment_key: str) -> None: ... + + def delete(self, experiment_key: str) -> None: ... + + +class InMemoryCheckpointStore: + """In-process checkpoint store, useful for tests.""" + + def __init__(self) -> None: + self._store: Dict[str, ExperimentCheckpoint] = {} + self._lock = threading.Lock() + + def load(self, experiment_key: str) -> Optional[ExperimentCheckpoint]: + with self._lock: + return self._store.get(experiment_key) + + def save(self, checkpoint: ExperimentCheckpoint, experiment_key: str) -> None: + with self._lock: + self._store[experiment_key] = checkpoint + + def delete(self, experiment_key: str) -> None: + with self._lock: + self._store.pop(experiment_key, None) + + +class FileCheckpointStore: + """On-disk checkpoint store with atomic writes. + + Writes are debounced to at most one write per ``flush_interval_s`` seconds + (plus a final write on close/abort). Writes run on a worker thread via + ``asyncio.to_thread`` so the event loop is never blocked. + """ + + def __init__(self, directory: str, flush_interval_s: float = 1.0) -> None: + self._directory = directory + self._flush_interval_s = flush_interval_s + self._lock = threading.Lock() + self._last_flush: float = 0.0 + os.makedirs(directory, exist_ok=True) + + def _path_for(self, experiment_key: str) -> str: + return os.path.join(self._directory, _sanitize_key(experiment_key)) + + def load(self, experiment_key: str) -> Optional[ExperimentCheckpoint]: + path = self._path_for(experiment_key) + if not os.path.exists(path): + return None + try: + with open(path, "r", encoding="utf-8") as f: + raw = f.read() + except OSError: + return None + return ExperimentCheckpoint.deserialize(raw) + + def save(self, checkpoint: ExperimentCheckpoint, experiment_key: str) -> None: + """Write ``checkpoint`` atomically to disk under ``experiment_key``. + + ``save`` performs a full write every call; debouncing is the caller's + responsibility (the executor calls it at most once per + ``flush_interval_s``). ``experiment_key`` must be the same key the + store is later loaded with. + """ + path = self._path_for(experiment_key) + with self._lock: + self._write_atomic(path, checkpoint.serialize()) + + @staticmethod + def _write_atomic(path: str, data: str) -> None: + directory = os.path.dirname(path) or "." + fd, tmp_path = tempfile.mkstemp(prefix=".cp-", suffix=".tmp", dir=directory) + try: + with os.fdopen(fd, "w", encoding="utf-8") as f: + f.write(data) + f.flush() + os.fsync(f.fileno()) + os.replace(tmp_path, path) + except Exception: + try: + os.unlink(tmp_path) + except OSError: + pass + raise + + def delete(self, experiment_key: str) -> None: + path = self._path_for(experiment_key) + with self._lock: + if os.path.exists(path): + try: + os.unlink(path) + except OSError: + pass + + def should_flush(self, now: float) -> bool: + return (now - self._last_flush) >= self._flush_interval_s + + def mark_flushed(self, now: float) -> None: + self._last_flush = now + + +__all__ = [ + "CheckpointItem", + "ExperimentFingerprint", + "ExperimentCheckpoint", + "CheckpointStore", + "InMemoryCheckpointStore", + "FileCheckpointStore", + "compute_data_digest", + "experiment_checkpoint_key", +] diff --git a/langfuse/_experiment/errors.py b/langfuse/_experiment/errors.py new file mode 100644 index 000000000..30c6535c5 --- /dev/null +++ b/langfuse/_experiment/errors.py @@ -0,0 +1,98 @@ +"""@private Exception types raised by the experiment execution policies. + +These are intentionally separate from the public ``langfuse.experiment`` types so +that callers can catch execution failures without importing the private package. +""" + +from typing import TYPE_CHECKING, Optional + +if TYPE_CHECKING: + from langfuse._experiment.policy import FailurePolicy + from langfuse.experiment import ExperimentResult + + +class ExperimentExecutionError(Exception): + """Base class for all experiment execution policy errors.""" + + +class TaskTimeoutError(ExperimentExecutionError): + """Raised when a single task call exceeds its timeout budget.""" + + def __init__(self, item_key: str, timeout: float): + self.item_key = item_key + self.timeout = timeout + message = "Task for item %s timed out after %s seconds" % (item_key, timeout) + super().__init__(message) + + +class EvaluatorTimeoutError(ExperimentExecutionError): + """Raised when a single evaluator call exceeds its timeout budget.""" + + def __init__(self, evaluator_name: str, item_key: str, timeout: float): + self.evaluator_name = evaluator_name + self.item_key = item_key + self.timeout = timeout + message = "Evaluator %s for item %s timed out after %s seconds" % ( + evaluator_name, + item_key, + timeout, + ) + super().__init__(message) + + +class RunTimeoutError(ExperimentExecutionError): + """Raised when the whole experiment run exceeds its run deadline.""" + + def __init__(self, timeout: float): + self.timeout = timeout + message = "Experiment run timed out after %s seconds" % timeout + super().__init__(message) + + +class FailureBudgetExceededError(ExperimentExecutionError): + """Raised when the failure budget is exhausted. + + Carries the partial ``ExperimentResult`` so callers can inspect what + completed before the budget tripped. + """ + + def __init__( + self, + failed: int, + attempted: int, + policy: "FailurePolicy", + partial_result: Optional["ExperimentResult"] = None, + ): + self.failed = failed + self.attempted = attempted + self.policy = policy + self.partial_result = partial_result + message = "Failure budget exceeded: %s of %s attempted items failed" % ( + failed, + attempted, + ) + super().__init__(message) + + +class CheckpointMismatchError(ExperimentExecutionError): + """Raised on resume when the checkpoint fingerprint differs.""" + + def __init__(self, field: str, expected: object, actual: object): + self.field = field + self.expected = expected + self.actual = actual + message = "Checkpoint fingerprint mismatch for %r: expected %r, got %r" % ( + field, + expected, + actual, + ) + super().__init__(message) + + +class CheckpointCorruptedError(ExperimentExecutionError): + """Raised when a checkpoint payload fails integrity or schema checks.""" + + def __init__(self, reason: str): + self.reason = reason + message = "Checkpoint corrupted: %s" % reason + super().__init__(message) diff --git a/langfuse/_experiment/executor.py b/langfuse/_experiment/executor.py new file mode 100644 index 000000000..f0cd13e6f --- /dev/null +++ b/langfuse/_experiment/executor.py @@ -0,0 +1,525 @@ +"""@private Orchestrates one experiment run with retry, timeout, and resume. + +The executor owns scheduling, semaphore, failure budget, run deadline, +cancellation, and checkpointing, and delegates per-item processing to a +caller-supplied coroutine. The ``Langfuse`` client wires its existing +``_process_experiment_item`` into that coroutine so the span tree and score +creation stay in the client. +""" + +import asyncio +import json as _json +import random +import time +from dataclasses import dataclass, field +from typing import ( + Any, + Awaitable, + Callable, + Dict, + List, + Optional, + Tuple, +) + +from langfuse._client.attributes import _serialize +from langfuse._experiment.checkpoint import ( + CheckpointItem, + ExperimentCheckpoint, + ExperimentFingerprint, + FileCheckpointStore, + compute_data_digest, + experiment_checkpoint_key, +) +from langfuse._experiment.errors import ( + CheckpointMismatchError, +) +from langfuse._experiment.item_key import make_item_keys +from langfuse._experiment.policy import ( + ExecutionPolicy, +) +from langfuse._experiment.results import ( + AttemptRecord, + EvaluatorFailure, + ExecutionSummary, + ItemFailure, +) +from langfuse._experiment.retrying import ( + _attempts_from_exception, + _run_with_retry, +) +from langfuse.experiment import Evaluation, ExperimentItem, ExperimentItemResult + +# processor(item, item_key, experiment_id, policy, attempt_hook, +# evaluator_failures). ``experiment_id`` is passed per item because a resumed +# local-data run adopts the id stored in the checkpoint. +ItemProcessor = Callable[ + [ + ExperimentItem, + str, + str, + Optional[ExecutionPolicy], + Optional[Callable[[AttemptRecord], None]], + List[EvaluatorFailure], + ], + Awaitable[ExperimentItemResult], +] + +# Grace period (seconds) when waiting for in-flight tasks to finish +# cancelling, so a task that swallows cancellation cannot hang the run. +_CANCEL_GRACE_S = 5.0 + + +class ItemProcessingFailed(Exception): + """Carrier exception raised by ``_process_experiment_item`` on terminal + task failure. Carries the assembled ``ItemFailure``.""" + + def __init__(self, failure: ItemFailure): + self.failure = failure + super().__init__(failure.error_message) + + +@dataclass +class ExecutorResult: + item_results: List[ExperimentItemResult] = field(default_factory=list) + failures: List[ItemFailure] = field(default_factory=list) + evaluator_failures: List[EvaluatorFailure] = field(default_factory=list) + summary: Optional[ExecutionSummary] = None + aborted: bool = False + abort_reason: Optional[str] = None + # The effective experiment id: the constructor id, or the id adopted + # from the checkpoint when a local-data run resumed. + experiment_id: str = "" + + +def _evaluator_name(evaluator: Any) -> str: + return getattr(evaluator, "__name__", "evaluator") + + +def _key_for_item(item: ExperimentItem, dataset_version: Optional[str]) -> str: + keys = make_item_keys([item], dataset_version) + return keys[0] + + +def _evaluation_to_dict(e: Evaluation) -> Dict[str, Any]: + return { + "name": e.name, + "value": e.value, + "comment": e.comment, + "metadata": e.metadata, + "data_type": e.data_type, + "config_id": e.config_id, + } + + +def _evaluation_from_dict(d: Dict[str, Any]) -> Evaluation: + return Evaluation( + name=str(d.get("name")), # type: ignore[arg-type] + value=d.get("value"), # type: ignore[arg-type] + comment=d.get("comment"), + metadata=d.get("metadata"), + data_type=d.get("data_type"), + config_id=d.get("config_id"), + ) + + +class ExperimentExecutor: + """Orchestrates a single experiment run.""" + + def __init__( + self, + *, + experiment_id: str, + experiment_id_deterministic: bool, + experiment_name: str, + run_name: str, + data: List[ExperimentItem], + dataset_version: Optional[str], + policy: Optional[ExecutionPolicy], + processor: ItemProcessor, + max_concurrency: int, + clock: Optional[Callable[[], float]] = None, + sleep: Optional[Callable[[float], Awaitable[None]]] = None, + rng: Optional[random.Random] = None, + ): + if max_concurrency < 1: + message = "max_concurrency must be >= 1, got %s" % max_concurrency + raise ValueError(message) + self.experiment_id = experiment_id + # True when ``experiment_id`` is stable across runs (dataset runs, or + # a caller-supplied experiment_id). Passed in explicitly by the + # client; the executor does not infer it. + self._experiment_id_deterministic = experiment_id_deterministic + self.experiment_name = experiment_name + self.run_name = run_name + self.data = data + self.dataset_version = dataset_version + self.policy = policy + self.processor = processor + self.max_concurrency = max_concurrency + self._clock = clock or time.monotonic + self._sleep = sleep or asyncio.sleep + self._rng = rng or random.Random( + policy.seed if policy is not None and policy.seed is not None else None + ) + self._item_keys = make_item_keys(data, dataset_version) + self._failures: List[ItemFailure] = [] + self._evaluator_failures: List[EvaluatorFailure] = [] + self._aborted = False + self._abort_reason: Optional[str] = None + self._skipped_resumed = 0 + self._retried_items = 0 + self._retried_keys: set = set() + self._total_attempts = 0 + self._item_attempts: Dict[str, int] = {} + self._checkpoint: Optional[ExperimentCheckpoint] = None + + def _fingerprint(self, task: Any, evaluators: List[Any]) -> ExperimentFingerprint: + task_name = getattr(task, "__name__", "task") + evaluator_names = [_evaluator_name(e) for e in evaluators] + return ExperimentFingerprint( + task_name=task_name, + evaluator_names=evaluator_names, + data_size=len(self.data), + data_digest=compute_data_digest(self._item_keys), + ) + + def _checkpoint_key(self) -> str: + """Store/load key for the checkpoint. + + Deterministic ids (dataset runs, caller-supplied ids) key the + checkpoint by the id itself. Local-data runs draw a fresh random id + per call, so their key is derived from the experiment name and run + name instead; see ``experiment_checkpoint_key``. + """ + if self._experiment_id_deterministic: + return self.experiment_id + return experiment_checkpoint_key(self.experiment_name, self.run_name) + + def _load_checkpoint(self, fingerprint: ExperimentFingerprint) -> None: + if self.policy is None or self.policy.checkpoint is None: + return + if not self.policy.resume: + return + cp = self.policy.checkpoint.load(self._checkpoint_key()) + if cp is None: + return + if cp.experiment_id != self.experiment_id: + if self._experiment_id_deterministic: + raise CheckpointMismatchError( + "experiment_id", self.experiment_id, cp.experiment_id + ) + # Local data: the id is random per run, so a resumed run adopts + # the id recorded in the checkpoint. Resumed items and newly + # executed items then share one experiment id. + self.experiment_id = cp.experiment_id + force = self.policy.resume == "force" + if ( + not force + and cp.fingerprint is not None + and not cp.fingerprint.matches(fingerprint) + ): + raise CheckpointMismatchError( + "fingerprint", + fingerprint.to_dict(), + cp.fingerprint.to_dict() if cp.fingerprint else None, + ) + self._checkpoint = cp + + def _init_checkpoint(self, fingerprint: ExperimentFingerprint) -> None: + if self.policy is None or self.policy.checkpoint is None: + return + if self._checkpoint is None: + now = self._clock() + self._checkpoint = ExperimentCheckpoint( + schema_version=1, + experiment_id=self.experiment_id, + experiment_name=self.experiment_name, + run_name=self.run_name, + fingerprint=fingerprint, + items={}, + created_at=now, + updated_at=now, + status="running", + ) + + async def _maybe_flush_checkpoint(self, force: bool = False) -> None: + if ( + self._checkpoint is None + or self.policy is None + or self.policy.checkpoint is None + ): + return + store = self.policy.checkpoint + now = self._clock() + if not force and isinstance(store, FileCheckpointStore): + if not store.should_flush(now): + return + self._checkpoint.updated_at = now + await asyncio.to_thread(store.save, self._checkpoint, self._checkpoint_key()) + if isinstance(store, FileCheckpointStore): + store.mark_flushed(now) + + def _record_checkpoint_item( + self, + item_key: str, + result: Optional[ExperimentItemResult], + failure: Optional[ItemFailure], + ) -> None: + if self._checkpoint is None: + return + if result is not None: + self._checkpoint.items[item_key] = CheckpointItem( + state="completed", + output_json=_serialize(result.output), + evaluations=[_evaluation_to_dict(e) for e in result.evaluations], + trace_id=result.trace_id, + attempts=self._item_attempts.get(item_key, 1), + error=None, + ) + elif failure is not None: + self._checkpoint.items[item_key] = CheckpointItem( + state="failed", + output_json=None, + evaluations=[], + trace_id=failure.trace_id, + attempts=len(failure.attempts), + error=failure.error_message, + ) + + def _make_attempt_hook(self, key: str) -> Callable[[AttemptRecord], None]: + retried_keys = self._retried_keys + + def hook(record: AttemptRecord) -> None: + self._total_attempts += 1 + self._item_attempts[key] = record.attempt + # An item that consumed more than one attempt was retried; count + # it at most once, no matter how many retry attempts it used. + if record.attempt > 1 and key not in retried_keys: + retried_keys.add(key) + self._retried_items += 1 + + return hook + + def _failure_budget_exceeded(self, failed: int, attempted: int) -> bool: + if self.policy is None: + return False + return self.policy.failures.is_exceeded(failed, attempted) + + def _rebuild_result_from_checkpoint( + self, item: ExperimentItem, key: str, cp_item: CheckpointItem + ) -> ExperimentItemResult: + output: Any = None + if cp_item.output_json is not None: + try: + output = _json.loads(cp_item.output_json) + except (ValueError, TypeError): + output = cp_item.output_json + evaluations = [_evaluation_from_dict(d) for d in cp_item.evaluations] + return ExperimentItemResult( + item=item, + output=output, + evaluations=evaluations, + trace_id=cp_item.trace_id, + experiment_id=self.experiment_id, + ) + + async def run( + self, + *, + task: Any, + evaluators: List[Any], + ) -> ExecutorResult: + fingerprint = self._fingerprint(task, evaluators) + self._load_checkpoint(fingerprint) + self._init_checkpoint(fingerprint) + + start = self._clock() + run_timeout = ( + self.policy.timeouts.run_timeout if self.policy is not None else None + ) + run_deadline = start + run_timeout if run_timeout is not None else None + + semaphore = asyncio.Semaphore(self.max_concurrency) + + skip_keys: set = set() + if self._checkpoint is not None: + for key, cp_item in self._checkpoint.items.items(): + if cp_item.state == "completed": + skip_keys.add(key) + + to_run: List[Tuple[int, ExperimentItem, str]] = [] + resumed_results: List[ExperimentItemResult] = [] + for idx, (item, key) in enumerate(zip(self.data, self._item_keys)): + if key in skip_keys: + assert self._checkpoint is not None + cp_item = self._checkpoint.items[key] + resumed_results.append( + self._rebuild_result_from_checkpoint(item, key, cp_item) + ) + self._skipped_resumed += 1 + else: + to_run.append((idx, item, key)) + + in_flight: set = set() + completed_results: List[ExperimentItemResult] = [] + attempted = 0 + failed = 0 + + async def process_one(item: ExperimentItem, key: str) -> None: + nonlocal attempted, failed + async with semaphore: + if self._aborted: + return + attempted += 1 + try: + result = await self.processor( + item, + key, + self.experiment_id, + self.policy, + self._make_attempt_hook(key), + self._evaluator_failures, + ) + completed_results.append(result) + self._record_checkpoint_item(key, result, None) + await self._maybe_flush_checkpoint() + except ItemProcessingFailed as e: + self._failures.append(e.failure) + failed += 1 + self._record_checkpoint_item(key, None, e.failure) + await self._maybe_flush_checkpoint() + except asyncio.CancelledError: + raise + except BaseException as e: # noqa: BLE001 + failure = ItemFailure( + item=item, + item_key=key, + stage="task", + error_type=type(e).__name__, + error_message=str(e), + attempts=_attempts_from_exception(e), + trace_id=None, + ) + self._failures.append(failure) + failed += 1 + self._record_checkpoint_item(key, None, failure) + await self._maybe_flush_checkpoint() + + next_index = 0 + + def _spawn_next() -> bool: + """Schedule the next pending item; return True if one was spawned.""" + nonlocal next_index + if next_index >= len(to_run): + return False + _, item, key = to_run[next_index] + next_index += 1 + t = asyncio.ensure_future(process_one(item, key)) + in_flight.add(t) + t.add_done_callback(in_flight.discard) + return True + + try: + while not self._aborted: + if run_deadline is not None and self._clock() >= run_deadline: + self._aborted = True + self._abort_reason = "run timeout" + break + while len(in_flight) < self.max_concurrency and _spawn_next(): + pass + if not in_flight: + break + wait_timeout = ( + max(0.0, run_deadline - self._clock()) + if run_deadline is not None + else None + ) + await asyncio.wait( + in_flight, + timeout=wait_timeout, + return_when=asyncio.FIRST_COMPLETED, + ) + if ( + run_deadline is not None + and in_flight + and self._clock() >= run_deadline + ): + self._aborted = True + self._abort_reason = "run timeout" + if not self._aborted and self._failure_budget_exceeded( + failed, attempted + ): + self._aborted = True + self._abort_reason = "failure budget exceeded" + except asyncio.CancelledError: + self._aborted = True + self._abort_reason = "cancelled" + for t in in_flight: + t.cancel() + if in_flight: + # Bounded wait: never block the run on tasks that ignore + # cancellation. + await asyncio.wait(in_flight, timeout=_CANCEL_GRACE_S) + await self._maybe_flush_checkpoint(force=True) + raise + finally: + pending = [t for t in in_flight if not t.done()] + for t in pending: + t.cancel() + if pending: + # Bounded wait so a task that swallows cancellation cannot hang + # the whole run. + await asyncio.wait(pending, timeout=_CANCEL_GRACE_S) + + if self._aborted and self._abort_reason is None: + self._abort_reason = "aborted" + + await self._maybe_flush_checkpoint(force=True) + + by_key: Dict[str, ExperimentItemResult] = {} + for r in completed_results: + by_key[_key_for_item(r.item, self.dataset_version)] = r + for r in resumed_results: + by_key[_key_for_item(r.item, self.dataset_version)] = r + + ordered: List[ExperimentItemResult] = [] + for item, key in zip(self.data, self._item_keys): + if key in by_key: + ordered.append(by_key[key]) + + failures_by_stage: Dict[str, int] = {} + for f in self._failures: + failures_by_stage[f.stage] = failures_by_stage.get(f.stage, 0) + 1 + + summary = ExecutionSummary( + total_items=len(self.data), + succeeded=len(completed_results) + len(resumed_results), + failed=len(self._failures), + skipped_resumed=self._skipped_resumed, + retried_items=self._retried_items, + total_attempts=self._total_attempts, + aborted=self._aborted, + abort_reason=self._abort_reason, + duration_s=self._clock() - start, + failures_by_stage=failures_by_stage, + ) + + return ExecutorResult( + item_results=ordered, + failures=list(self._failures), + evaluator_failures=list(self._evaluator_failures), + summary=summary, + aborted=self._aborted, + abort_reason=self._abort_reason, + experiment_id=self.experiment_id, + ) + + +__all__ = [ + "ExperimentExecutor", + "ExecutorResult", + "ItemProcessingFailed", + "ItemProcessor", + "_run_with_retry", +] diff --git a/langfuse/_experiment/item_key.py b/langfuse/_experiment/item_key.py new file mode 100644 index 000000000..b2027741f --- /dev/null +++ b/langfuse/_experiment/item_key.py @@ -0,0 +1,110 @@ +"""@private Stable per-item identity used by experiment checkpoints. + +The key derivation is shared with the ``experiment_item_id`` expression in +``langfuse._client.client`` so the two never diverge. +""" + +from typing import Any, Dict, List, Optional + +from langfuse._client.attributes import _serialize +from langfuse._client.utils import get_sha256_hash_hex +from langfuse.api import DatasetItem + + +def _hash_serialized(value: Any) -> str: + """sha256 hex of the serialized value, treating None as the JSON null.""" + serialized = _serialize(value) + if serialized is None: + serialized = "null" + return get_sha256_hash_hex(serialized)[:16] + + +def local_item_input_hash(input_data: Any) -> str: + """First 16 hex chars of ``sha256(_serialize(input))``. + + This is the same value used as ``experiment_item_id`` for local dict items + in ``Langfuse._process_experiment_item``. Centralizing it here keeps the + checkpoint key and the span attribute in sync. + """ + return _hash_serialized(input_data) + + +def _is_dataset_item(item: object) -> bool: + return isinstance(item, DatasetItem) or ( + not isinstance(item, dict) + and hasattr(item, "dataset_id") + and hasattr(item, "id") + ) + + +def _dataset_item_id(item: object) -> Optional[str]: + return getattr(item, "id", None) + + +def _item_input(item: object) -> Any: + if isinstance(item, dict): + return item.get("input") + return getattr(item, "input", None) + + +def _item_expected_output(item: object) -> Any: + if isinstance(item, dict): + return item.get("expected_output") + return getattr(item, "expected_output", None) + + +def _base_key(item: object, dataset_version: Optional[str]) -> str: + """Compute the base (pre-disambiguation) key for one item.""" + if _is_dataset_item(item): + item_id = _dataset_item_id(item) or "" + if dataset_version is not None: + return "%s@v%s" % (item_id, dataset_version) + return item_id + + input_data = _item_input(item) + base = local_item_input_hash(input_data) + expected_output = _item_expected_output(item) + # Always fold the expected-output hash in so two items with identical + # inputs but different ground truth do not collide in the checkpoint. + expected_hash = _hash_serialized(expected_output) + return "%s:%s" % (base, expected_hash) + + +def make_item_keys(data: List[Any], dataset_version: Optional[str] = None) -> List[str]: + """Return stable keys for every item in ``data``. + + Duplicate keys within the same ``data`` list are disambiguated by + appending a deterministic ``#`` suffix (0-based) so the order + of items fixes the suffix assignment. + """ + base_keys: List[str] = [_base_key(item, dataset_version) for item in data] + + seen: Dict[str, int] = {} + keys: List[str] = [] + for base in base_keys: + count = seen.get(base, 0) + seen[base] = count + 1 + if count == 0 and base not in keys: + keys.append(base) + else: + keys.append("%s#%d" % (base, count)) + return keys + + +def experiment_item_id_for(item: object) -> str: + """Return the ``experiment_item_id`` for an item. + + For dataset items this is the dataset item id; for local dict items this + is the input hash. Mirrors the expression in + ``Langfuse._process_experiment_item``. + """ + if _is_dataset_item(item): + return _dataset_item_id(item) or "" + return local_item_input_hash(_item_input(item)) + + +__all__ = [ + "local_item_input_hash", + "make_item_keys", + "experiment_item_id_for", +] diff --git a/langfuse/_experiment/policy.py b/langfuse/_experiment/policy.py new file mode 100644 index 000000000..f1882b60e --- /dev/null +++ b/langfuse/_experiment/policy.py @@ -0,0 +1,309 @@ +"""@private Execution policies for reliable experiment runs. + +All policies are opt-in. With the defaults (``max_attempts=1``, no timeouts, +no failure budget, no checkpoint) the executor reproduces today's behavior +exactly, except that failures are additionally recorded on the result object. +""" + +import asyncio +import datetime +import email.utils +import random +from dataclasses import dataclass, field +from typing import Callable, Optional, Union + +from langfuse._experiment.checkpoint import CheckpointStore + +# Status codes that should be retried by default. +_RETRYABLE_STATUS_CODES = {408, 409, 425, 429} +_RETRYABLE_STATUS_RANGES = ((500, 599),) + +# Exception class names that are treated as retryable transient errors +# across provider SDKs (openai, anthropic, litellm, ...). +_RETRYABLE_CLASS_NAMES = { + "APIConnectionError", + "APITimeoutError", + "RateLimitError", + "InternalServerError", + "ServiceUnavailableError", +} + +# Exceptions that must never be retried unless an explicit ``retry_on`` +# predicate says so. +_NEVER_RETRY_CLASS_NAMES = {"ValueError", "TypeError"} + + +def _coerce_status_code(exc: BaseException) -> Optional[int]: + """Best-effort extraction of an HTTP status code from an exception.""" + for attr in ("status_code", "status"): + value = getattr(exc, attr, None) + if isinstance(value, int): + return value + response = getattr(exc, "response", None) + if response is not None: + for attr in ("status_code", "status"): + value = getattr(response, attr, None) + if isinstance(value, int): + return value + return None + + +def _coerce_headers(exc: BaseException) -> Optional[object]: + response = getattr(exc, "response", None) + if response is None: + return None + headers = getattr(response, "headers", None) + if headers is None: + return None + return headers # type: ignore[no-any-return] + + +def _parse_retry_after(value: str) -> Optional[float]: + """Parse a ``Retry-After`` header value (seconds or HTTP date).""" + value = value.strip() + if not value: + return None + # Try integer/float seconds first. + try: + return float(value) + except ValueError: + pass + # Fall back to HTTP-date. + try: + parsed = email.utils.parsedate_to_datetime(value) + except (TypeError, ValueError): + return None + + now = datetime.datetime.now(tz=datetime.timezone.utc) + if parsed.tzinfo is None: + parsed = parsed.replace(tzinfo=datetime.timezone.utc) + delta = (parsed - now).total_seconds() + if delta < 0: + return 0.0 + return delta + + +def _retry_after_from_exc(exc: BaseException) -> Optional[float]: + headers = _coerce_headers(exc) + if headers is None: + return None + getter = getattr(headers, "get", None) + if callable(getter): + raw = getter("Retry-After") + else: + try: + raw = headers["Retry-After"] # type: ignore[index] + except (KeyError, TypeError): + raw = None + if raw is None: + return None + if isinstance(raw, (int, float)): + return float(raw) + if isinstance(raw, str): + return _parse_retry_after(raw) + return None + + +def _class_name_mro_retryable(exc: BaseException) -> bool: + """Return True if any class in the exception's MRO has a retryable name.""" + for cls in type(exc).__mro__: + if cls.__name__ in _RETRYABLE_CLASS_NAMES: + return True + return False + + +def _class_name_mro_never_retry(exc: BaseException) -> bool: + for cls in type(exc).__mro__: + if cls.__name__ in _NEVER_RETRY_CLASS_NAMES: + return True + return False + + +def is_retryable_error(exc: BaseException) -> bool: + """Default retry predicate. + + Retries transient HTTP errors (408/409/425/429/5xx), timeout and + connection errors, and provider SDK exception classes whose names match + the openai/anthropic/litellm naming conventions. Never retries + ``asyncio.CancelledError``, ``KeyboardInterrupt``, or user-code + ``ValueError``/``TypeError``. + """ + if isinstance(exc, asyncio.CancelledError): + return False + if isinstance(exc, KeyboardInterrupt): + return False + + if isinstance(exc, (asyncio.TimeoutError, TimeoutError, ConnectionError)): + return True + + if _class_name_mro_retryable(exc): + return True + + status = _coerce_status_code(exc) + if status is not None: + if status in _RETRYABLE_STATUS_CODES: + return True + for low, high in _RETRYABLE_STATUS_RANGES: + if low <= status <= high: + return True + + if _class_name_mro_never_retry(exc): + return False + + return False + + +def _default_retry_on(exc: BaseException) -> bool: + return is_retryable_error(exc) + + +@dataclass(frozen=True) +class RetryPolicy: + """Retry configuration for a task or evaluator call. + + ``max_attempts`` is the total number of attempts including the first one, + so ``max_attempts=1`` disables retries (the default). + """ + + max_attempts: int = 1 + initial_delay: float = 1.0 + backoff_multiplier: float = 2.0 + max_delay: float = 30.0 + jitter: bool = True + retry_on: Optional[Callable[[BaseException], bool]] = None + + def __post_init__(self) -> None: + if self.max_attempts < 1: + message = "max_attempts must be >= 1, got %s" % self.max_attempts + raise ValueError(message) + if self.initial_delay < 0: + message = "initial_delay must be >= 0, got %s" % self.initial_delay + raise ValueError(message) + if self.backoff_multiplier < 0: + message = ( + "backoff_multiplier must be >= 0, got %s" % self.backoff_multiplier + ) + raise ValueError(message) + if self.max_delay < 0: + message = "max_delay must be >= 0, got %s" % self.max_delay + raise ValueError(message) + + def should_retry(self, exc: BaseException) -> bool: + predicate = self.retry_on if self.retry_on is not None else _default_retry_on + try: + return bool(predicate(exc)) + except Exception: + return False + + def delay_for( + self, + attempt: int, + retry_after: Optional[float], + rng: random.Random, + ) -> float: + """Compute the delay (seconds) before the next attempt. + + ``attempt`` is the 1-based index of the attempt that just failed, so + the delay precedes attempt ``attempt + 1``. + """ + if retry_after is not None and retry_after >= 0: + return min(retry_after, self.max_delay) + if attempt < 1: + return 0.0 + # Exponential backoff: initial_delay * backoff^(attempt-1). + delay = self.initial_delay * (self.backoff_multiplier ** (attempt - 1)) + if self.max_delay is not None: + delay = min(delay, self.max_delay) + if self.jitter: + # Full jitter: uniform in [0, delay]. + delay = rng.uniform(0.0, delay) if delay > 0 else 0.0 + return delay + + +@dataclass(frozen=True) +class TimeoutPolicy: + """Per-call and per-run timeout budgets (seconds).""" + + task_timeout: Optional[float] = None + evaluator_timeout: Optional[float] = None + run_timeout: Optional[float] = None + + +@dataclass(frozen=True) +class FailurePolicy: + """Failure budget that can abort a run early. + + When the budget is exceeded the executor stops scheduling new items, + cancels in-flight ones, and marks the result ``aborted``. ``on_threshold`` + controls whether the run is stopped silently (``"stop"``) or whether a + ``FailureBudgetExceededError`` is raised (``"raise"``). + """ + + max_failures: Optional[int] = None + max_failure_rate: Optional[float] = None + min_items_for_rate: int = 10 + on_threshold: str = "stop" + + def __post_init__(self) -> None: + if self.max_failures is not None and self.max_failures < 0: + message = "max_failures must be >= 0, got %s" % self.max_failures + raise ValueError(message) + if self.max_failure_rate is not None and not ( + 0.0 <= self.max_failure_rate <= 1.0 + ): + message = "max_failure_rate must be in [0, 1], got %s" % ( + self.max_failure_rate + ) + raise ValueError(message) + if self.min_items_for_rate < 0: + message = "min_items_for_rate must be >= 0, got %s" % ( + self.min_items_for_rate + ) + raise ValueError(message) + if self.on_threshold not in ("stop", "raise"): + message = "on_threshold must be 'stop' or 'raise', got %r" % ( + self.on_threshold + ) + raise ValueError(message) + + def is_exceeded(self, failed: int, attempted: int) -> bool: + if self.max_failures is not None and failed >= self.max_failures: + return True + if ( + self.max_failure_rate is not None + and attempted >= self.min_items_for_rate + and attempted > 0 + and (failed / attempted) > self.max_failure_rate + ): + return True + return False + + +@dataclass(frozen=True) +class ExecutionPolicy: + """Top-level execution policy for one experiment run.""" + + task_retry: RetryPolicy = field(default_factory=RetryPolicy) + evaluator_retry: RetryPolicy = field(default_factory=RetryPolicy) + timeouts: TimeoutPolicy = field(default_factory=TimeoutPolicy) + failures: FailurePolicy = field(default_factory=FailurePolicy) + checkpoint: Optional[CheckpointStore] = None + resume: Union[bool, str] = False + seed: Optional[int] = None + + def __post_init__(self) -> None: + if isinstance(self.resume, str) and self.resume not in ("force",): + message = "resume must be bool or 'force', got %r" % (self.resume,) + raise ValueError(message) + if self.resume and self.checkpoint is None: + message = "resume=True requires a checkpoint store" + raise ValueError(message) + + +__all__ = [ + "RetryPolicy", + "TimeoutPolicy", + "FailurePolicy", + "ExecutionPolicy", + "is_retryable_error", +] diff --git a/langfuse/_experiment/results.py b/langfuse/_experiment/results.py new file mode 100644 index 000000000..38ccb7362 --- /dev/null +++ b/langfuse/_experiment/results.py @@ -0,0 +1,221 @@ +"""@private Result and summary types produced by the experiment executor.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional + +if TYPE_CHECKING: + pass + +# Stage of the pipeline where a failure happened. ``task``/``validation``/ +# ``timeout``/``cancelled`` are experiment-run stages; ``mapper``/``evaluator``/ +# ``composite``/``score`` are batch-evaluation stages. +FailureStage = Literal[ + "validation", + "task", + "evaluator", + "timeout", + "cancelled", + "mapper", + "composite", + "score", +] + + +class AttemptRecord: + """Record of a single task or evaluator attempt.""" + + __slots__ = ( + "attempt", + "started_at", + "duration_s", + "error_type", + "error_message", + "retry_after", + "will_retry", + ) + + def __init__( + self, + *, + attempt: int, + started_at: float, + duration_s: float, + error_type: Optional[str] = None, + error_message: Optional[str] = None, + retry_after: Optional[float] = None, + will_retry: bool = False, + ): + self.attempt = attempt + self.started_at = started_at + self.duration_s = duration_s + self.error_type = error_type + self.error_message = error_message + self.retry_after = retry_after + self.will_retry = will_retry + + def to_dict(self) -> Dict[str, Any]: + return { + "attempt": self.attempt, + "started_at": self.started_at, + "duration_s": self.duration_s, + "error_type": self.error_type, + "error_message": self.error_message, + "retry_after": self.retry_after, + "will_retry": self.will_retry, + } + + +class ItemFailure: + """Record of an item that could not be processed to completion.""" + + __slots__ = ( + "item", + "item_key", + "stage", + "error_type", + "error_message", + "attempts", + "trace_id", + "evaluator_name", + ) + + def __init__( + self, + *, + item: Any, + item_key: str, + stage: FailureStage, + error_type: str, + error_message: str, + attempts: List[AttemptRecord], + trace_id: Optional[str] = None, + evaluator_name: Optional[str] = None, + ): + self.item = item + self.item_key = item_key + self.stage = stage + self.error_type = error_type + self.error_message = error_message + self.attempts = attempts + self.trace_id = trace_id + self.evaluator_name = evaluator_name + + def to_dict(self) -> Dict[str, Any]: + return { + "item_key": self.item_key, + "stage": self.stage, + "error_type": self.error_type, + "error_message": self.error_message, + "attempts": [a.to_dict() for a in self.attempts], + "trace_id": self.trace_id, + "evaluator_name": self.evaluator_name, + } + + +class EvaluatorFailure: + """Record of an evaluator that failed for an item after all retries.""" + + __slots__ = ( + "item_key", + "evaluator_name", + "error_type", + "error_message", + "attempts", + ) + + def __init__( + self, + *, + item_key: str, + evaluator_name: str, + error_type: str, + error_message: str, + attempts: List[AttemptRecord], + ): + self.item_key = item_key + self.evaluator_name = evaluator_name + self.error_type = error_type + self.error_message = error_message + self.attempts = attempts + + def to_dict(self) -> Dict[str, Any]: + return { + "item_key": self.item_key, + "evaluator_name": self.evaluator_name, + "error_type": self.error_type, + "error_message": self.error_message, + "attempts": [a.to_dict() for a in self.attempts], + } + + +class ExecutionSummary: + """Aggregate statistics for one experiment run.""" + + def __init__( + self, + *, + total_items: int, + succeeded: int, + failed: int, + skipped_resumed: int, + retried_items: int, + total_attempts: int, + aborted: bool, + abort_reason: Optional[str], + duration_s: float, + failures_by_stage: Optional[Dict[str, int]] = None, + ): + self.total_items = total_items + self.succeeded = succeeded + self.failed = failed + self.skipped_resumed = skipped_resumed + self.retried_items = retried_items + self.total_attempts = total_attempts + self.aborted = aborted + self.abort_reason = abort_reason + self.duration_s = duration_s + self.failures_by_stage: Dict[str, int] = failures_by_stage or {} + + def to_dict(self) -> Dict[str, Any]: + return { + "total_items": self.total_items, + "succeeded": self.succeeded, + "failed": self.failed, + "skipped_resumed": self.skipped_resumed, + "retried_items": self.retried_items, + "total_attempts": self.total_attempts, + "aborted": self.aborted, + "abort_reason": self.abort_reason, + "duration_s": self.duration_s, + "failures_by_stage": dict(self.failures_by_stage), + } + + def format(self) -> str: + lines = [ + "Execution Summary:", + " total items: %s" % self.total_items, + " succeeded: %s" % self.succeeded, + " failed: %s" % self.failed, + " skipped (resumed): %s" % self.skipped_resumed, + " retried items: %s" % self.retried_items, + " total attempts: %s" % self.total_attempts, + " duration: %.3fs" % self.duration_s, + ] + if self.failures_by_stage: + lines.append(" failures by stage:") + for stage, count in sorted(self.failures_by_stage.items()): + lines.append(" %s: %s" % (stage, count)) + if self.aborted: + reason = self.abort_reason or "unknown" + lines.append(" aborted: %s" % reason) + return "\n".join(lines) + + +__all__ = [ + "AttemptRecord", + "ItemFailure", + "EvaluatorFailure", + "ExecutionSummary", + "FailureStage", +] diff --git a/langfuse/_experiment/retrying.py b/langfuse/_experiment/retrying.py new file mode 100644 index 000000000..5eff34a9c --- /dev/null +++ b/langfuse/_experiment/retrying.py @@ -0,0 +1,123 @@ +"""@private Shared retry/timeout helper for experiment and batch evaluation. + +Both ``ExperimentExecutor`` and ``BatchEvaluationRunner`` wrap their task, +evaluator, mapper, and composite calls in :func:`_run_with_retry` so retry +and timeout behavior is defined in exactly one place. +""" + +import asyncio +import random +from typing import Any, Awaitable, Callable, List, Optional, Tuple + +from langfuse._experiment.policy import RetryPolicy, _retry_after_from_exc +from langfuse._experiment.results import AttemptRecord + +# Attribute name used to attach the recorded attempts to the terminal +# exception so failure records can carry them. +_ATTEMPTS_ATTR = "_langfuse_attempt_records" + + +async def _run_with_retry( + coro_factory: Callable[[], Awaitable[Any]], + *, + policy: Optional[RetryPolicy], + timeout: Optional[float], + clock: Callable[[], float], + sleep: Callable[[float], Awaitable[None]], + rng: random.Random, + attempt_hook: Optional[Callable[[AttemptRecord], None]], +) -> Tuple[Any, List[AttemptRecord]]: + """Run ``coro_factory()`` with retry and timeout. + + Returns ``(result, attempts)`` on success. Raises the final exception on + terminal failure (after ``max_attempts`` or a non-retryable error); the + recorded attempts are attached to the raised exception under + ``_langfuse_attempt_records``. + """ + retry_policy = policy if policy is not None else RetryPolicy() + attempts: List[AttemptRecord] = [] + max_attempts = retry_policy.max_attempts + last_exc: Optional[BaseException] = None + + for attempt in range(1, max_attempts + 1): + started_at = clock() + attempt_exc: Optional[BaseException] = None + result: Any = None + succeeded = False + try: + if timeout is not None: + result = await asyncio.wait_for(coro_factory(), timeout=timeout) + else: + result = await coro_factory() + succeeded = True + except asyncio.TimeoutError as e: + attempt_exc = e + last_exc = e + except BaseException as e: # noqa: BLE001 + attempt_exc = e + last_exc = e + duration = clock() - started_at + + record = AttemptRecord( + attempt=attempt, + started_at=started_at, + duration_s=duration, + error_type=( + type(attempt_exc).__name__ if attempt_exc is not None else None + ), + error_message=(str(attempt_exc) if attempt_exc is not None else None), + retry_after=_retry_after_from_exc_safe(attempt_exc), + will_retry=False, + ) + if succeeded: + attempts.append(record) + if attempt_hook is not None: + attempt_hook(record) + return result, attempts + + will_retry = attempt < max_attempts and retry_policy.should_retry( + attempt_exc # type: ignore[arg-type] + ) + record.will_retry = will_retry + attempts.append(record) + if attempt_hook is not None: + attempt_hook(record) + + if not will_retry: + break + + delay = retry_policy.delay_for( + attempt=attempt, retry_after=record.retry_after, rng=rng + ) + if delay > 0: + await sleep(delay) + + assert last_exc is not None + try: + setattr(last_exc, _ATTEMPTS_ATTR, attempts) + except Exception: + pass + raise last_exc + + +def _attempts_from_exception(exc: Optional[BaseException]) -> List[AttemptRecord]: + """Return the attempts recorded for a terminal failure, if any.""" + if exc is None: + return [] + return list(getattr(exc, _ATTEMPTS_ATTR, None) or []) + + +def _retry_after_from_exc_safe(exc: Optional[BaseException]) -> Optional[float]: + """Best-effort ``Retry-After`` extraction; never raises.""" + if exc is None: + return None + try: + return _retry_after_from_exc(exc) + except Exception: + return None + + +__all__ = [ + "_run_with_retry", + "_attempts_from_exception", +] diff --git a/langfuse/batch_evaluation.py b/langfuse/batch_evaluation.py index 8ef89c79b..d070d6b6a 100644 --- a/langfuse/batch_evaluation.py +++ b/langfuse/batch_evaluation.py @@ -9,11 +9,13 @@ import asyncio import json +import random import time from typing import ( TYPE_CHECKING, Any, Awaitable, + Callable, Dict, List, Optional, @@ -23,6 +25,13 @@ Union, ) +from langfuse._experiment.policy import ExecutionPolicy +from langfuse._experiment.results import ( + ExecutionSummary, + FailureStage, + ItemFailure, +) +from langfuse._experiment.retrying import _attempts_from_exception, _run_with_retry from langfuse.api import ObservationV2 from langfuse.experiment import Evaluation, EvaluatorFunction from langfuse.logger import langfuse_logger as logger @@ -581,6 +590,8 @@ class BatchEvaluationResult: error_summary: Dictionary mapping error types to occurrence counts. has_more_items: True if the run stopped at max_items and more items may exist. item_evaluations: Dictionary mapping observation IDs to their evaluation results (both regular and composite). + failures: Detailed per-item failure records with stage and attempt history. + execution: Execution summary with retry/timeout/abort statistics (see ExecutionSummary). Examples: Basic result inspection: @@ -676,6 +687,8 @@ def __init__( error_summary: Dict[str, int], has_more_items: bool, item_evaluations: Dict[str, List["Evaluation"]], + failures: Optional[List[ItemFailure]] = None, + execution: Optional[ExecutionSummary] = None, ): """Initialize BatchEvaluationResult with comprehensive statistics. @@ -694,6 +707,10 @@ def __init__( error_summary: Error types and counts. has_more_items: Whether more items may exist beyond max_items. item_evaluations: Dictionary mapping item IDs to their evaluation results. + failures: Detailed per-item failure records (mapper, evaluator, + composite, score, or validation stage). Empty when nothing failed. + execution: Execution summary (retries, timeouts, aborts). Populated + by ``run_async``; None on manually constructed results. Note: All arguments must be provided as keywords. @@ -712,6 +729,8 @@ def __init__( self.error_summary = error_summary self.has_more_items = has_more_items self.item_evaluations = item_evaluations + self.failures: List[ItemFailure] = list(failures or []) + self.execution = execution def __str__(self) -> str: """Return a formatted string representation of the batch evaluation results. @@ -836,6 +855,7 @@ async def run_async( max_retries: int = 3, verbose: bool = False, resume_from: Optional[BatchEvaluationResumeToken] = None, + execution: Optional[ExecutionPolicy] = None, ) -> BatchEvaluationResult: """Run batch evaluation asynchronously. @@ -861,6 +881,10 @@ async def run_async( verbose: If True, log progress to console. resume_from: Resume token from a previous run. If `filter` is omitted, the token's filter is reused. + execution: Optional execution policy adding retries, timeouts, and a + failure budget for mapper/evaluator/composite calls, plus detailed + failure records and an execution summary on the result. With None + (the default), behavior is unchanged from previous versions. Returns: BatchEvaluationResult with comprehensive statistics. @@ -871,6 +895,47 @@ async def run_async( """ start_time = time.time() + run_deadline: Optional[float] = None + if execution is not None and execution.timeouts.run_timeout is not None: + run_deadline = time.monotonic() + execution.timeouts.run_timeout + rng = random.Random( + execution.seed + if execution is not None and execution.seed is not None + else None + ) + execution_failures: List[ItemFailure] = [] + attempt_counters = {"total_attempts": 0} + retried_item_keys: Set[str] = set() + + def _attempt_hook_for(item_key: str) -> Callable[[Any], None]: + def hook(record: Any) -> None: + attempt_counters["total_attempts"] += 1 + if record.attempt > 1: + retried_item_keys.add(item_key) + + return hook + + def _build_execution_summary( + aborted: bool, abort_reason: Optional[str] + ) -> ExecutionSummary: + failures_by_stage: Dict[str, int] = {} + for failure in execution_failures: + failures_by_stage[failure.stage] = ( + failures_by_stage.get(failure.stage, 0) + 1 + ) + return ExecutionSummary( + total_items=total_items_fetched, + succeeded=total_items_processed, + failed=len(execution_failures), + skipped_resumed=0, + retried_items=len(retried_item_keys), + total_attempts=attempt_counters["total_attempts"], + aborted=aborted, + abort_reason=abort_reason, + duration_s=time.time() - start_time, + failures_by_stage=failures_by_stage, + ) + # The token's cursor only continues the query that produced it. if resume_from is not None: if filter is None: @@ -892,6 +957,8 @@ async def run_async( failed_item_ids: List[str] = [] error_summary: Dict[str, int] = {} item_evaluations: Dict[str, List[Evaluation]] = {} + aborted = False + abort_reason: Optional[str] = None previously_processed = resume_from.items_processed if resume_from else 0 evaluator_stats_dict = { @@ -952,6 +1019,13 @@ def build_resume_token() -> BatchEvaluationResumeToken: logger.info("Reached max_items limit (%s)", max_items) break + if run_deadline is not None and time.monotonic() >= run_deadline: + if verbose: + logger.info("Run timeout reached, stopping batch evaluation") + aborted = True + abort_reason = "run timeout" + break + # Clamping the page size keeps the cursor aligned with the last # processed item, so resuming after max_items skips nothing. limit = fetch_batch_size @@ -984,6 +1058,8 @@ def build_resume_token() -> BatchEvaluationResumeToken: error_summary=error_summary, has_more_items=False, item_evaluations=item_evaluations, + failures=execution_failures, + execution=_build_execution_summary(False, None), ) batch_number += 1 @@ -994,7 +1070,12 @@ def build_resume_token() -> BatchEvaluationResumeToken: async def process_item( item: ObservationV2, - ) -> Tuple[str, Union[Tuple[int, int, int, List[Evaluation]], Exception]]: + ) -> Tuple[ + str, + Union[Tuple[int, int, int, List[Evaluation]], Exception], + List[ItemFailure], + ]: + item_failures: List[ItemFailure] = [] async with semaphore: try: result = await self._process_batch_evaluation_item( @@ -1005,10 +1086,32 @@ async def process_item( metadata=metadata, _add_observation_scores_to_trace=_add_observation_scores_to_trace, evaluator_stats_dict=evaluator_stats_dict, + execution=execution, + rng=rng, + attempt_hook_factory=_attempt_hook_for, + item_failures=item_failures, ) - return (item.id, result) + return (item.id, result, item_failures) except Exception as e: - return (item.id, e) + if not item_failures: + # Nothing was recorded yet: the item failed before + # or outside the stages that record their own + # failures. + stage: FailureStage = ( + "validation" if not item.trace_id else "mapper" + ) + item_failures.append( + ItemFailure( + item=item, + item_key=item.id, + stage=stage, + error_type=type(e).__name__, + error_message=str(e), + attempts=_attempts_from_exception(e), + trace_id=item.trace_id, + ) + ) + return (item.id, e, item_failures) items_to_process = self._deduplicate_page( items, skip_ids=previous_page_ids | {resumed_item_id} @@ -1018,7 +1121,10 @@ async def process_item( *[process_item(item) for item in items_to_process] ) - for item, (item_id, result) in zip(items_to_process, results): + for item, (item_id, result, item_failures) in zip( + items_to_process, results + ): + execution_failures.extend(item_failures) if isinstance(result, Exception): total_items_failed += 1 failed_item_ids.append(item_id) @@ -1063,6 +1169,16 @@ async def process_item( has_more = False break + if execution is not None and execution.failures.is_exceeded( + len(execution_failures), + total_items_processed + total_items_failed, + ): + if verbose: + logger.info("Failure budget exceeded, stopping batch evaluation") + aborted = True + abort_reason = "failure budget exceeded" + break + if verbose: logger.info("Flushing scores to Langfuse...") self.client.flush() @@ -1074,6 +1190,29 @@ async def process_item( time.time() - start_time, ) + if aborted: + # A policy abort (run timeout or failure budget) is an incomplete + # run: hand back a token that resumes after the last processed + # page. + return self._build_result( + total_items_fetched=total_items_fetched, + total_items_processed=total_items_processed, + total_items_failed=total_items_failed, + total_scores_created=total_scores_created, + total_composite_scores_created=total_composite_scores_created, + total_evaluations_failed=total_evaluations_failed, + evaluator_stats_dict=evaluator_stats_dict, + resume_token=build_resume_token(), + completed=False, + start_time=start_time, + failed_item_ids=failed_item_ids, + error_summary=error_summary, + has_more_items=False, + item_evaluations=item_evaluations, + failures=execution_failures, + execution=_build_execution_summary(True, abort_reason), + ) + return self._build_result( total_items_fetched=total_items_fetched, total_items_processed=total_items_processed, @@ -1089,6 +1228,8 @@ async def process_item( error_summary=error_summary, has_more_items=has_more, item_evaluations=item_evaluations, + failures=execution_failures, + execution=_build_execution_summary(False, None), ) async def _fetch_batch_with_retry( @@ -1139,6 +1280,10 @@ async def _process_batch_evaluation_item( metadata: Optional[Dict[str, Any]], _add_observation_scores_to_trace: bool, evaluator_stats_dict: Dict[str, EvaluatorStats], + execution: Optional[ExecutionPolicy] = None, + rng: Optional[random.Random] = None, + attempt_hook_factory: Optional[Callable[[str], Callable[[Any], None]]] = None, + item_failures: Optional[List[ItemFailure]] = None, ) -> Tuple[int, int, int, List[Evaluation]]: """Process a single item: map, evaluate, create scores. @@ -1151,6 +1296,10 @@ async def _process_batch_evaluation_item( _add_observation_scores_to_trace: Whether to duplicate observation-level scores at trace level. evaluator_stats_dict: Dictionary tracking evaluator statistics. + execution: Optional execution policy for retries/timeouts. + rng: Random source for retry jitter. + attempt_hook_factory: Returns an attempt hook for an item key. + item_failures: List collecting this item's detailed failures. Returns: Tuple of (scores_created, composite_scores_created, evaluations_failed, all_evaluations). @@ -1160,13 +1309,47 @@ async def _process_batch_evaluation_item( """ if not item.trace_id: message = f"Observation {item.id} has no trace_id" + if item_failures is not None: + item_failures.append( + ItemFailure( + item=item, + item_key=item.id, + stage="validation", + error_type="ValueError", + error_message=message, + attempts=[], + trace_id=None, + ) + ) raise ValueError(message) scores_created = 0 composite_scores_created = 0 evaluations_failed = 0 - evaluator_inputs = await self._run_mapper(mapper, item) + try: + evaluator_inputs = await self._run_stage_with_policy( + lambda: self._run_mapper(mapper, item), + execution=execution, + rng=rng, + attempt_hook=( + attempt_hook_factory(item.id) if attempt_hook_factory else None + ), + ) + except Exception as e: + if item_failures is not None: + item_failures.append( + ItemFailure( + item=item, + item_key=item.id, + stage="mapper", + error_type=type(e).__name__, + error_message=str(e), + attempts=_attempts_from_exception(e), + trace_id=item.trace_id, + ) + ) + raise evaluations: List[Evaluation] = [] for evaluator in evaluators: @@ -1175,12 +1358,19 @@ async def _process_batch_evaluation_item( stats.total_runs += 1 try: - eval_results = await self._run_evaluator_internal( - evaluator, - input=evaluator_inputs.input, - output=evaluator_inputs.output, - expected_output=evaluator_inputs.expected_output, - metadata=evaluator_inputs.metadata, + eval_results = await self._run_stage_with_policy( + lambda: self._run_evaluator_internal( + evaluator, + input=evaluator_inputs.input, + output=evaluator_inputs.output, + expected_output=evaluator_inputs.expected_output, + metadata=evaluator_inputs.metadata, + ), + execution=execution, + rng=rng, + attempt_hook=( + attempt_hook_factory(item.id) if attempt_hook_factory else None + ), ) stats.successful_runs += 1 @@ -1190,6 +1380,19 @@ async def _process_batch_evaluation_item( except Exception as e: stats.failed_runs += 1 evaluations_failed += 1 + if item_failures is not None: + item_failures.append( + ItemFailure( + item=item, + item_key=item.id, + stage="evaluator", + error_type=type(e).__name__, + error_message=str(e), + attempts=_attempts_from_exception(e), + trace_id=item.trace_id, + evaluator_name=evaluator_name, + ) + ) logger.warning( "Evaluator %s failed on item %s: %s", evaluator_name, @@ -1198,36 +1401,94 @@ async def _process_batch_evaluation_item( ) for evaluation in evaluations: - scores_created += self._create_score( - item=item, - evaluation=evaluation, - additional_metadata=metadata, - add_observation_score_to_trace=_add_observation_scores_to_trace, - ) + try: + scores_created += self._create_score( + item=item, + evaluation=evaluation, + additional_metadata=metadata, + add_observation_score_to_trace=_add_observation_scores_to_trace, + ) + except Exception as e: + if item_failures is not None: + item_failures.append( + ItemFailure( + item=item, + item_key=item.id, + stage="score", + error_type=type(e).__name__, + error_message=str(e), + attempts=[], + trace_id=item.trace_id, + ) + ) + raise if composite_evaluator and evaluations: + composite_evals: Optional[List[Evaluation]] = None try: - composite_evals = await self._run_composite_evaluator( - composite_evaluator, - input=evaluator_inputs.input, - output=evaluator_inputs.output, - expected_output=evaluator_inputs.expected_output, - metadata=evaluator_inputs.metadata, - evaluations=evaluations, + composite_evals = await self._run_stage_with_policy( + lambda: self._run_composite_evaluator( + composite_evaluator, + input=evaluator_inputs.input, + output=evaluator_inputs.output, + expected_output=evaluator_inputs.expected_output, + metadata=evaluator_inputs.metadata, + evaluations=evaluations, + ), + execution=execution, + rng=rng, + attempt_hook=( + attempt_hook_factory(item.id) if attempt_hook_factory else None + ), ) + except Exception as e: + if item_failures is not None: + item_failures.append( + ItemFailure( + item=item, + item_key=item.id, + stage="composite", + error_type=type(e).__name__, + error_message=str(e), + attempts=_attempts_from_exception(e), + trace_id=item.trace_id, + ) + ) + logger.warning("Composite evaluator failed on item %s: %s", item.id, e) + if composite_evals is not None: + composite_score_error: Optional[Exception] = None for composite_eval in composite_evals: - composite_scores_created += self._create_score( - item=item, - evaluation=composite_eval, - additional_metadata=metadata, - add_observation_score_to_trace=_add_observation_scores_to_trace, + try: + composite_scores_created += self._create_score( + item=item, + evaluation=composite_eval, + additional_metadata=metadata, + add_observation_score_to_trace=_add_observation_scores_to_trace, + ) + except Exception as e: + composite_score_error = e + if item_failures is not None: + item_failures.append( + ItemFailure( + item=item, + item_key=item.id, + stage="score", + error_type=type(e).__name__, + error_message=str(e), + attempts=[], + trace_id=item.trace_id, + ) + ) + break + if composite_score_error is not None: + logger.warning( + "Composite evaluator failed on item %s: %s", + item.id, + composite_score_error, ) - - evaluations.extend(composite_evals) - - except Exception as e: - logger.warning("Composite evaluator failed on item %s: %s", item.id, e) + else: + evaluations.extend(composite_evals) return ( scores_created, @@ -1236,6 +1497,34 @@ async def _process_batch_evaluation_item( evaluations, ) + async def _run_stage_with_policy( + self, + coro_factory: Callable[[], Awaitable[Any]], + *, + execution: Optional[ExecutionPolicy], + rng: Optional[random.Random], + attempt_hook: Optional[Callable[[Any], None]], + ) -> Any: + """Run one mapper/evaluator/composite call under the execution policy. + + Without a policy the coroutine factory is awaited directly, keeping + the default code path identical to previous versions. With a policy, + the call is retried per ``evaluator_retry`` and bounded by + ``evaluator_timeout`` via the shared ``_run_with_retry`` helper. + """ + if execution is None: + return await coro_factory() + result, _attempts = await _run_with_retry( + coro_factory, + policy=execution.evaluator_retry, + timeout=execution.timeouts.evaluator_timeout, + clock=time.monotonic, + sleep=asyncio.sleep, + rng=rng if rng is not None else random.Random(), + attempt_hook=attempt_hook, + ) + return result + async def _run_evaluator_internal( self, evaluator: EvaluatorFunction, @@ -1467,6 +1756,8 @@ def _build_result( error_summary: Dict[str, int], has_more_items: bool, item_evaluations: Dict[str, List[Evaluation]], + failures: Optional[List[ItemFailure]] = None, + execution: Optional[ExecutionSummary] = None, ) -> BatchEvaluationResult: """Build the final BatchEvaluationResult. @@ -1485,6 +1776,8 @@ def _build_result( error_summary: Error type counts. has_more_items: Whether more items may exist beyond max_items. item_evaluations: Dictionary mapping item IDs to their evaluation results. + failures: Detailed per-item failure records. + execution: Execution summary for policy-driven runs. Returns: BatchEvaluationResult instance. @@ -1506,4 +1799,6 @@ def _build_result( error_summary=error_summary, has_more_items=has_more_items, item_evaluations=item_evaluations, + failures=failures, + execution=execution, ) diff --git a/langfuse/experiment.py b/langfuse/experiment.py index 9782ac1a6..1d1f43205 100644 --- a/langfuse/experiment.py +++ b/langfuse/experiment.py @@ -21,6 +21,33 @@ overload, ) +from langfuse._experiment.checkpoint import ( + CheckpointStore, + FileCheckpointStore, + InMemoryCheckpointStore, +) +from langfuse._experiment.errors import ( + CheckpointCorruptedError, + CheckpointMismatchError, + EvaluatorTimeoutError, + ExperimentExecutionError, + FailureBudgetExceededError, + RunTimeoutError, + TaskTimeoutError, +) +from langfuse._experiment.policy import ( + ExecutionPolicy, + FailurePolicy, + RetryPolicy, + TimeoutPolicy, + is_retryable_error, +) +from langfuse._experiment.results import ( + AttemptRecord, + EvaluatorFailure, + ExecutionSummary, + ItemFailure, +) from langfuse.api import DatasetItem from langfuse.logger import langfuse_logger as logger from langfuse.types import ExperimentScoreType @@ -397,6 +424,9 @@ def __init__( experiment_url: Optional[str] = None, dataset_run_id: Optional[str] = None, dataset_run_url: Optional[str] = None, + failures: Optional[List[ItemFailure]] = None, + evaluator_failures: Optional[List[EvaluatorFailure]] = None, + execution: Optional[ExecutionSummary] = None, ): """Initialize an ExperimentResult with the provided data. @@ -405,6 +435,8 @@ def __init__( run_name: The current experiment run name. description: Optional description of the experiment. item_results: List of results from processing individual dataset items. + Only successfully processed items are included; failed items + are recorded in ``failures``. run_evaluations: List of aggregate evaluation results for the entire run. experiment_id: ID of the experiment run. experiment_url: Optional URL to view results in Langfuse UI. @@ -412,6 +444,14 @@ def __init__( returns `experiment_id`. dataset_run_url: Deprecated alias of `experiment_url`. Used only when `experiment_url` is not provided. + failures: List of items that could not be processed (task failures, + timeouts, cancellations). Always populated, even without an + execution policy; today these failures are only logged. + evaluator_failures: List of evaluator calls that failed after all + retries. The corresponding item still appears in + ``item_results`` with its surviving evaluations. + execution: Aggregate statistics about the run (counts, attempts, + duration, abort status). Populated by the executor. """ self.name = name self.run_name = run_name @@ -420,6 +460,11 @@ def __init__( self.run_evaluations = run_evaluations self.experiment_id = experiment_id self.experiment_url = experiment_url or dataset_run_url + self.failures = failures if failures is not None else [] + self.evaluator_failures = ( + evaluator_failures if evaluator_failures is not None else [] + ) + self.execution = execution @property def dataset_run_id(self) -> str: @@ -627,6 +672,27 @@ def format(self, *, include_item_results: bool = False) -> str: if self.experiment_url: output += f"\n🔗 Experiment:\n {self.experiment_url}" + # Failures section (only when there are failures). + if self.failures or self.evaluator_failures: + output += "\n\n⚠ Failures:" + for failure in self.failures: + output += ( + f"\n • [{failure.stage}] {failure.item_key}: " + f"{failure.error_type}: {failure.error_message}" + ) + if failure.attempts: + output += " (attempts=%s)" % len(failure.attempts) + for eval_failure in self.evaluator_failures: + output += ( + f"\n • [evaluator:{eval_failure.evaluator_name}] " + f"{eval_failure.item_key}: {eval_failure.error_type}: " + f"{eval_failure.error_message}" + ) + output += "\n" + + if self.execution is not None: + output += "\n" + self.execution.format() + return output @@ -1165,6 +1231,7 @@ def run_experiment( max_concurrency: int = 50, metadata: Optional[Dict[str, str]] = None, experiment_id: Optional[str] = None, + execution: Optional[ExecutionPolicy] = None, _dataset_version: Optional[datetime] = None, ) -> ExperimentResult: resolved_data = data if data is not None else self.data @@ -1195,6 +1262,7 @@ def run_experiment( max_concurrency=max_concurrency, metadata=merged_metadata, experiment_id=experiment_id, + execution=execution, _dataset_version=resolved_dataset_version, ) @@ -1251,3 +1319,37 @@ def __init__( else: formatted = "Experiment regression detected" super().__init__(formatted) + + +__all__ = [ + "Evaluation", + "EvaluatorFunction", + "ExperimentData", + "ExperimentItem", + "ExperimentItemResult", + "ExperimentResult", + "RunEvaluatorFunction", + "RunnerContext", + "RegressionError", + "TaskFunction", + # Re-exported execution policy types. + "AttemptRecord", + "CheckpointCorruptedError", + "CheckpointMismatchError", + "CheckpointStore", + "EvaluatorFailure", + "EvaluatorTimeoutError", + "ExecutionPolicy", + "ExecutionSummary", + "ExperimentExecutionError", + "FailureBudgetExceededError", + "FailurePolicy", + "FileCheckpointStore", + "InMemoryCheckpointStore", + "ItemFailure", + "RetryPolicy", + "RunTimeoutError", + "TaskTimeoutError", + "TimeoutPolicy", + "is_retryable_error", +] diff --git a/tests/unit/test_batch_evaluation_execution.py b/tests/unit/test_batch_evaluation_execution.py new file mode 100644 index 000000000..93e58b5e5 --- /dev/null +++ b/tests/unit/test_batch_evaluation_execution.py @@ -0,0 +1,318 @@ +"""Unit tests for execution policies on batch evaluation. + +Uses the local fakes from ``tests/unit/test_batch_evaluation.py`` +(``make_runner`` / ``make_observation``); no Langfuse server required. +""" + +import asyncio +from typing import Any + +import pytest + +from langfuse._experiment.policy import ( + ExecutionPolicy, + FailurePolicy, + RetryPolicy, + TimeoutPolicy, +) +from langfuse.batch_evaluation import BatchEvaluationRunner +from langfuse.experiment import Evaluation + +from .test_batch_evaluation import make_observation, make_runner, mapper + + +async def run(runner: BatchEvaluationRunner, **kwargs: Any): + kwargs.setdefault("mapper", mapper) + kwargs.setdefault("evaluators", [_length_evaluator]) + return await runner.run_async(**kwargs) + + +def _length_evaluator(*, output, **kwargs): + return Evaluation(name="length", value=len(output)) + + +class TestEvaluatorRetry: + @pytest.mark.asyncio + async def test_evaluator_retried_after_transient_errors_then_succeeds(self): + runner, _, client = make_runner([make_observation(1)]) + calls = {"n": 0} + + def flaky_evaluator(*, output, **kwargs): + calls["n"] += 1 + if calls["n"] < 3: + raise ConnectionError("transient") + return Evaluation(name="length", value=len(output)) + + result = await run( + runner, + evaluators=[flaky_evaluator], + execution=ExecutionPolicy( + evaluator_retry=RetryPolicy( + max_attempts=5, initial_delay=0, jitter=False + ) + ), + ) + + assert calls["n"] == 3 + assert result.total_items_processed == 1 + assert result.total_scores_created == 1 + assert result.failures == [] + stats = {s.name: s for s in result.evaluator_stats} + # Stats keep their per-call (not per-attempt) accounting. + assert stats["flaky_evaluator"].successful_runs == 1 + assert stats["flaky_evaluator"].failed_runs == 0 + assert result.execution is not None + # 1 mapper attempt + 3 evaluator attempts. + assert result.execution.total_attempts == 4 + assert result.execution.retried_items == 1 + assert result.execution.aborted is False + client.create_score.assert_called_once() + + @pytest.mark.asyncio + async def test_evaluator_non_retryable_fails_once(self): + runner, _, _ = make_runner([make_observation(1)]) + calls = {"n": 0} + + def broken_evaluator(*, output, **kwargs): + calls["n"] += 1 + raise ValueError("bad input") + + result = await run( + runner, + evaluators=[broken_evaluator], + execution=ExecutionPolicy( + evaluator_retry=RetryPolicy( + max_attempts=5, initial_delay=0, jitter=False + ) + ), + ) + + assert calls["n"] == 1 + assert result.total_evaluations_failed == 1 + assert len(result.failures) == 1 + failure = result.failures[0] + assert failure.stage == "evaluator" + assert failure.error_type == "ValueError" + assert failure.evaluator_name == "broken_evaluator" + assert result.execution is not None + # 1 mapper attempt + 1 evaluator attempt. + assert result.execution.total_attempts == 2 + + +class TestEvaluatorTimeout: + @pytest.mark.asyncio + async def test_evaluator_timeout_recorded_item_still_scored(self): + runner, _, client = make_runner([make_observation(1)]) + + async def slow_evaluator(*, output, **kwargs): + await asyncio.sleep(5) + return Evaluation(name="slow", value=1) + + result = await run( + runner, + evaluators=[slow_evaluator, _length_evaluator], + execution=ExecutionPolicy(timeouts=TimeoutPolicy(evaluator_timeout=0.05)), + ) + + assert result.total_items_processed == 1 + # The slow evaluator failed, the fast one still scored the item. + assert result.total_scores_created == 1 + assert [e.name for e in result.item_evaluations["obs-1"]] == ["length"] + assert len(result.failures) == 1 + failure = result.failures[0] + assert failure.stage == "evaluator" + assert failure.error_type == "TimeoutError" + assert failure.evaluator_name == "slow_evaluator" + assert client.create_score.call_count == 1 + + +class TestMapperFailures: + @pytest.mark.asyncio + async def test_mapper_failure_recorded(self): + runner, _, _ = make_runner([make_observation(1)]) + calls = {"n": 0} + + def failing_mapper(*, item): + calls["n"] += 1 + raise ConnectionError("mapper boom") + + result = await run( + runner, + mapper=failing_mapper, + execution=ExecutionPolicy( + evaluator_retry=RetryPolicy( + max_attempts=3, initial_delay=0, jitter=False + ) + ), + ) + + assert calls["n"] == 3 # mapper retried per the policy + assert result.total_items_failed == 1 + assert result.failed_item_ids == ["obs-1"] + assert len(result.failures) == 1 + failure = result.failures[0] + assert failure.stage == "mapper" + assert failure.error_type == "ConnectionError" + assert failure.item_key == "obs-1" + + @pytest.mark.asyncio + async def test_missing_trace_id_recorded_as_validation(self): + runner, _, _ = make_runner([make_observation(2, trace_id=None)]) + + result = await run(runner, execution=ExecutionPolicy()) + + assert result.total_items_failed == 1 + assert result.error_summary == {"ValueError": 1} + assert len(result.failures) == 1 + assert result.failures[0].stage == "validation" + + @pytest.mark.asyncio + async def test_composite_failure_recorded_item_still_processed(self): + runner, _, client = make_runner([make_observation(1)]) + + def failing_composite(*, evaluations, **kwargs): + raise RuntimeError("composite boom") + + result = await run( + runner, + composite_evaluator=failing_composite, + execution=ExecutionPolicy(), + ) + + assert result.total_items_processed == 1 + assert result.total_scores_created == 1 + assert result.total_composite_scores_created == 0 + assert len(result.failures) == 1 + assert result.failures[0].stage == "composite" + + +class TestFailureBudget: + @pytest.mark.asyncio + async def test_budget_stops_paging_and_token_resumes_from_cursor(self): + runner, api, _ = make_runner([make_observation(i) for i in range(6)]) + + def failing_evaluator(*, output, **kwargs): + raise ConnectionError("judge down") + + first = await run( + runner, + evaluators=[failing_evaluator], + fetch_batch_size=2, + execution=ExecutionPolicy(failures=FailurePolicy(max_failures=2)), + ) + + # Two items on the first page; the evaluator fails on both, the items + # themselves still count as processed (evaluator failures are + # tolerated), then the budget stops further paging. + assert first.completed is False + assert first.total_items_fetched == 2 + assert first.total_items_processed == 2 + assert len(api.calls) == 1 + assert first.execution is not None + assert first.execution.aborted is True + assert "failure budget" in (first.execution.abort_reason or "") + token = first.resume_token + assert token is not None + assert token.cursor == "2" + assert token.last_processed_id == "obs-4" + # Both items were processed (evaluator failures do not fail items). + assert token.items_processed == 2 + + # The token resumes after the aborted page. + resumed = await run( + runner, + fetch_batch_size=2, + resume_from=token, + execution=ExecutionPolicy(failures=FailurePolicy(max_failures=100)), + ) + + assert resumed.completed is True + assert set(resumed.item_evaluations) == {"obs-3", "obs-2", "obs-1", "obs-0"} + + @pytest.mark.asyncio + async def test_budget_counts_evaluator_failures_across_pages(self): + runner, api, _ = make_runner([make_observation(i) for i in range(6)]) + + def half_failing_evaluator(*, output, **kwargs): + if int(output.split()[-1]) % 2 == 0: + raise ConnectionError("flaky judge") + return Evaluation(name="length", value=len(output)) + + result = await run( + runner, + evaluators=[half_failing_evaluator], + fetch_batch_size=2, + execution=ExecutionPolicy( + failures=FailurePolicy(max_failures=2, on_threshold="stop"), + ), + ) + + # Newest first: page 1 = obs-5 (ok) + obs-4 (fail); page 2 = obs-3 + # (ok) + obs-2 (fail) -> 2 evaluator failures -> stop before page 3. + # Evaluator failures do not fail the item, so all 4 attempted items + # count as processed. + assert result.execution is not None + assert result.execution.aborted is True + assert len(result.failures) == 2 + assert result.total_items_processed == 4 + assert len(api.calls) == 2 + + +class TestRunTimeout: + @pytest.mark.asyncio + async def test_run_timeout_stops_between_pages(self): + runner, api, _ = make_runner([make_observation(i) for i in range(6)]) + + async def slow_evaluator(*, output, **kwargs): + await asyncio.sleep(0.4) + return Evaluation(name="length", value=len(output)) + + result = await run( + runner, + evaluators=[slow_evaluator], + fetch_batch_size=2, + execution=ExecutionPolicy(timeouts=TimeoutPolicy(run_timeout=0.2)), + ) + + # The first page takes ~0.4s, which already exceeds the 0.2s budget, + # so the run stops before fetching the second page. + assert result.execution is not None + assert result.execution.aborted is True + assert "run timeout" in (result.execution.abort_reason or "") + assert result.completed is False + assert result.resume_token is not None + assert result.total_items_processed == 2 + assert len(api.calls) == 1 + + +class TestDefaultBehaviour: + @pytest.mark.asyncio + async def test_no_policy_matches_previous_behaviour(self): + runner, api, client = make_runner([make_observation(i) for i in range(5)]) + + result = await run(runner, fetch_batch_size=2) + + assert result.completed is True + assert result.total_items_processed == 5 + assert result.total_scores_created == 5 + assert result.failures == [] + assert result.execution is not None + assert result.execution.aborted is False + assert result.execution.total_attempts == 0 + assert result.execution.succeeded == 5 + assert result.execution.failures_by_stage == {} + assert [call["cursor"] for call in api.calls] == [None, "2", "4"] + client.flush.assert_called_once() + + @pytest.mark.asyncio + async def test_fetch_failure_without_policy_unchanged(self): + runner, api, _ = make_runner([make_observation(i) for i in range(6)]) + api.fail_on_call = 2 + + result = await run(runner, fetch_batch_size=2, max_retries=0) + + assert result.completed is False + assert result.resume_token is not None + assert result.execution is not None + assert result.execution.aborted is False + assert result.failures == [] diff --git a/tests/unit/test_experiment_execution.py b/tests/unit/test_experiment_execution.py new file mode 100644 index 000000000..5d1dd0d5e --- /dev/null +++ b/tests/unit/test_experiment_execution.py @@ -0,0 +1,1440 @@ +"""Unit tests for the experiment execution policies and executor. + +All tests are deterministic and run against the in-memory span exporter; no +Langfuse server or model provider is contacted. +""" + +import asyncio +import json +import os +import random +import re +import time +from datetime import datetime, timezone +from typing import List +from unittest.mock import MagicMock + +import pytest + +from langfuse import ( + AttemptRecord, + Evaluation, + ExecutionPolicy, + FailurePolicy, + RetryPolicy, + RunnerContext, + TimeoutPolicy, + is_retryable_error, +) +from langfuse._client.attributes import LangfuseOtelSpanAttributes +from langfuse._client.client import Langfuse +from langfuse._client.datasets import DatasetClient +from langfuse._experiment.checkpoint import ( + CheckpointItem, + ExperimentCheckpoint, + ExperimentFingerprint, + FileCheckpointStore, + InMemoryCheckpointStore, + compute_data_digest, + experiment_checkpoint_key, +) +from langfuse._experiment.errors import ( + CheckpointCorruptedError, + CheckpointMismatchError, +) +from langfuse._experiment.executor import ( + _run_with_retry, +) +from langfuse._experiment.item_key import ( + experiment_item_id_for, + local_item_input_hash, + make_item_keys, +) +from langfuse._experiment.policy import _parse_retry_after, _retry_after_from_exc +from langfuse.api import DatasetItem, DatasetStatus + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _dataset_item(*, item_id: str, dataset_id: str, input: str) -> DatasetItem: + return DatasetItem( + id=item_id, + status=DatasetStatus.ACTIVE, + input=input, + expected_output="expected %s" % input, + metadata=None, + source_trace_id=None, + source_observation_id=None, + dataset_id=dataset_id, + dataset_name="dataset", + media_references=[], + created_at=datetime(2026, 1, 1, tzinfo=timezone.utc), + updated_at=datetime(2026, 1, 1, tzinfo=timezone.utc), + ) + + +class _FakeStatus: + """Minimal exception carrying a ``status_code`` attribute.""" + + def __init__(self, status_code: int): + self.status_code = status_code + + +class _FakeResponse: + def __init__(self, status_code: int, headers=None): + self.status_code = status_code + self.headers = headers or {} + + +class _FakeHTTPError(Exception): + def __init__(self, status_code: int, headers=None): + self.response = _FakeResponse(status_code, headers) + super().__init__("http %s" % status_code) + + +class _NamedAPIError(Exception): + pass + + +# Register a retryable class name via MRO. +class APIConnectionError(_NamedAPIError): + pass + + +class _NeverRetryValue(ValueError): + pass + + +# --------------------------------------------------------------------------- +# Retry predicate +# --------------------------------------------------------------------------- + + +class TestIsRetryableError: + @pytest.mark.parametrize( + "exc,expected", + [ + # HTTP status codes via status_code attr. + (_FakeStatus(429), True), + (_FakeStatus(408), True), + (_FakeStatus(409), True), + (_FakeStatus(425), True), + (_FakeStatus(500), True), + (_FakeStatus(503), True), + (_FakeStatus(599), True), + (_FakeStatus(400), False), + (_FakeStatus(404), False), + # HTTP status via response.status_code. + (_FakeHTTPError(429), True), + (_FakeHTTPError(500), True), + (_FakeHTTPError(400), False), + # Timeout / connection errors. + (asyncio.TimeoutError(), True), + (TimeoutError(), True), + (ConnectionError(), True), + # Retryable class name via MRO. + (APIConnectionError(), True), + # Never retry cancellation / keyboard / user Value/Type errors. + (asyncio.CancelledError(), False), + (KeyboardInterrupt(), False), + (ValueError("bad"), False), + (TypeError("bad"), False), + (_NeverRetryValue("bad"), False), + # Plain non-retryable. + (RuntimeError("boom"), False), + (KeyError("missing"), False), + ], + ) + def test_predicate_table(self, exc, expected): + assert is_retryable_error(exc) is expected + + def test_retry_after_seconds_header(self): + exc = _FakeHTTPError(429, headers={"Retry-After": "5"}) + assert _retry_after_from_exc(exc) == 5.0 + + def test_retry_after_http_date_header(self): + # A future HTTP date should yield a non-negative delay. + import datetime as _dt + import email.utils + + future = _dt.datetime.now(_dt.timezone.utc) + _dt.timedelta(seconds=10) + exc = _FakeHTTPError( + 429, headers={"Retry-After": email.utils.format_datetime(future)} + ) + delay = _retry_after_from_exc(exc) + assert delay is not None and 0 <= delay <= 10 + + def test_retry_after_missing_header(self): + exc = _FakeHTTPError(429, headers={}) + assert _retry_after_from_exc(exc) is None + + def test_parse_retry_after_int(self): + assert _parse_retry_after("3") == 3.0 + + def test_parse_retry_after_invalid(self): + assert _parse_retry_after("not-a-date") is None + + +# --------------------------------------------------------------------------- +# Backoff math +# --------------------------------------------------------------------------- + + +class TestBackoffMath: + def test_no_jitter_uses_exponential_backoff(self): + policy = RetryPolicy( + max_attempts=4, initial_delay=1.0, backoff_multiplier=2.0, jitter=False + ) + rng = random.Random(0) + # attempt 1 fails -> delay before attempt 2 = 1.0 + assert policy.delay_for(1, None, rng) == 1.0 + assert policy.delay_for(2, None, rng) == 2.0 + assert policy.delay_for(3, None, rng) == 4.0 + + def test_max_delay_caps_backoff(self): + policy = RetryPolicy( + max_attempts=5, + initial_delay=10.0, + backoff_multiplier=10.0, + max_delay=15.0, + jitter=False, + ) + rng = random.Random(0) + assert policy.delay_for(1, None, rng) == 10.0 + assert policy.delay_for(2, None, rng) == 15.0 + + def test_retry_after_overrides_backoff(self): + policy = RetryPolicy(max_attempts=3, jitter=False) + rng = random.Random(0) + assert policy.delay_for(1, 7.0, rng) == 7.0 + # Capped by max_delay. + assert policy.delay_for(1, 100.0, rng) == 30.0 + + def test_jitter_is_within_bounds(self): + policy = RetryPolicy( + max_attempts=3, initial_delay=4.0, backoff_multiplier=1.0, jitter=True + ) + rng = random.Random(123) + for _ in range(20): + delay = policy.delay_for(1, None, rng) + assert 0.0 <= delay <= 4.0 + + def test_invalid_max_attempts_raises(self): + with pytest.raises(ValueError, match="max_attempts"): + RetryPolicy(max_attempts=0) + + def test_invalid_failure_rate_raises(self): + with pytest.raises(ValueError, match="max_failure_rate"): + FailurePolicy(max_failure_rate=1.5) + + def test_resume_requires_checkpoint(self): + with pytest.raises(ValueError, match="checkpoint"): + ExecutionPolicy(resume=True) + + def test_resume_force_accepted(self): + store = InMemoryCheckpointStore() + ExecutionPolicy(checkpoint=store, resume="force") + + +# --------------------------------------------------------------------------- +# _run_with_retry helper +# --------------------------------------------------------------------------- + + +class TestRunWithRetry: + @pytest.mark.asyncio + async def test_retries_then_succeeds(self): + calls = {"n": 0} + + async def factory(): + calls["n"] += 1 + if calls["n"] < 3: + raise _FakeHTTPError(429) + return "ok" + + policy = RetryPolicy(max_attempts=5, initial_delay=0, jitter=False) + rng = random.Random(0) + + sleeps = [] + + async def fake_sleep(d): + sleeps.append(d) + + result, attempts = await _run_with_retry( + factory, + policy=policy, + timeout=None, + clock=time.monotonic, + sleep=fake_sleep, + rng=rng, + attempt_hook=None, + ) + assert result == "ok" + assert calls["n"] == 3 + assert len(attempts) == 3 + assert attempts[0].will_retry is True + assert attempts[1].will_retry is True + assert attempts[2].will_retry is False + assert attempts[2].error_type is None + + @pytest.mark.asyncio + async def test_non_retryable_fails_once(self): + calls = {"n": 0} + + async def factory(): + calls["n"] += 1 + raise ValueError("nope") + + policy = RetryPolicy(max_attempts=5, initial_delay=0, jitter=False) + rng = random.Random(0) + with pytest.raises(ValueError): + await _run_with_retry( + factory, + policy=policy, + timeout=None, + clock=time.monotonic, + sleep=asyncio.sleep, + rng=rng, + attempt_hook=None, + ) + assert calls["n"] == 1 + + @pytest.mark.asyncio + async def test_timeout_raises_timeout_error(self): + async def factory(): + await asyncio.sleep(10) + return "never" + + policy = RetryPolicy(max_attempts=3, initial_delay=0, jitter=False) + rng = random.Random(0) + with pytest.raises(asyncio.TimeoutError): + await _run_with_retry( + factory, + policy=policy, + timeout=0.05, + clock=time.monotonic, + sleep=asyncio.sleep, + rng=rng, + attempt_hook=None, + ) + + @pytest.mark.asyncio + async def test_attempt_hook_invoked(self): + records: List[AttemptRecord] = [] + + def hook(record): + records.append(record) + + async def factory(): + raise _FakeHTTPError(503) + + policy = RetryPolicy(max_attempts=2, initial_delay=0, jitter=False) + rng = random.Random(0) + with pytest.raises(_FakeHTTPError): + await _run_with_retry( + factory, + policy=policy, + timeout=None, + clock=time.monotonic, + sleep=asyncio.sleep, + rng=rng, + attempt_hook=hook, + ) + assert len(records) == 2 + assert records[0].attempt == 1 and records[0].will_retry is True + assert records[1].attempt == 2 and records[1].will_retry is False + + +# --------------------------------------------------------------------------- +# Item keys +# --------------------------------------------------------------------------- + + +class TestItemKeys: + def test_local_item_input_hash_matches_experiment_item_id(self): + item = {"input": "hello"} + assert local_item_input_hash("hello") == experiment_item_id_for(item) + + def test_dataset_item_key_is_dataset_item_id(self): + item = _dataset_item(item_id="i1", dataset_id="d", input="a") + assert experiment_item_id_for(item) == "i1" + + def test_duplicate_input_keys_disambiguated(self): + data = [ + {"input": "same", "expected_output": "x"}, + {"input": "same", "expected_output": "y"}, + ] + keys = make_item_keys(data) + assert len(keys) == 2 + assert keys[0] != keys[1] + + def test_distinct_inputs_keep_distinct_keys(self): + data = [{"input": "a"}, {"input": "b"}] + keys = make_item_keys(data) + assert len(set(keys)) == 2 + + def test_dataset_version_pinned_into_key(self): + item = _dataset_item(item_id="i1", dataset_id="d", input="a") + keys = make_item_keys([item], dataset_version="2026-01-01T00:00:00Z") + assert keys[0] == "i1@v2026-01-01T00:00:00Z" + + +# --------------------------------------------------------------------------- +# Failure policy +# --------------------------------------------------------------------------- + + +class TestFailurePolicy: + def test_max_failures_exceeded(self): + p = FailurePolicy(max_failures=3) + assert not p.is_exceeded(2, 10) + assert p.is_exceeded(3, 10) + + def test_max_failure_rate_respects_min_items(self): + p = FailurePolicy(max_failure_rate=0.3, min_items_for_rate=10) + # 50% failure but below min_items_for_rate -> not exceeded. + assert not p.is_exceeded(2, 4) + # At min_items_for_rate with > 30% -> exceeded. + assert p.is_exceeded(4, 10) + + def test_no_budget_never_exceeded(self): + p = FailurePolicy() + assert not p.is_exceeded(100, 100) + + +# --------------------------------------------------------------------------- +# Checkpoint serialization +# --------------------------------------------------------------------------- + + +class TestCheckpointSerialization: + def test_roundtrip(self): + cp = ExperimentCheckpoint( + experiment_id="exp1", + experiment_name="exp", + run_name="r", + fingerprint=ExperimentFingerprint( + task_name="task", + evaluator_names=["e1"], + data_size=2, + data_digest="abc", + ), + items={ + "k1": CheckpointItem( + state="completed", + output_json='"out"', + evaluations=[{"name": "e1", "value": 1.0}], + trace_id="t1", + attempts=1, + ), + "k2": CheckpointItem(state="failed", error="boom", attempts=2), + }, + created_at=1.0, + updated_at=2.0, + status="running", + ) + raw = cp.serialize() + restored = ExperimentCheckpoint.deserialize(raw) + assert restored.experiment_id == "exp1" + assert restored.items["k1"].state == "completed" + assert restored.items["k1"].output_json == '"out"' + assert restored.items["k2"].state == "failed" + assert restored.items["k2"].error == "boom" + assert restored.fingerprint is not None + assert restored.fingerprint.matches(cp.fingerprint) # type: ignore[arg-type] + + def test_corrupted_integrity_detected(self): + cp = ExperimentCheckpoint(experiment_id="exp1") + raw = cp.serialize() + # Tamper with the payload (remove the integrity key and re-add a bad one). + data = json.loads(raw) + data["experiment_id"] = "tampered" + bad = json.dumps(data) + with pytest.raises(CheckpointCorruptedError): + ExperimentCheckpoint.deserialize(bad) + + def test_fingerprint_mismatch_raises(self): + cp = ExperimentCheckpoint( + experiment_id="exp1", + fingerprint=ExperimentFingerprint( + task_name="task", + evaluator_names=[], + data_size=1, + data_digest="d1", + ), + ) + other = ExperimentFingerprint( + task_name="other", + evaluator_names=[], + data_size=1, + data_digest="d2", + ) + with pytest.raises(CheckpointMismatchError): + cp.verify_fingerprint(other) + + +class TestFileCheckpointStore: + def test_atomic_write_and_load(self, tmp_path): + store = FileCheckpointStore(str(tmp_path)) + cp = ExperimentCheckpoint(experiment_id="exp1", experiment_name="n") + store.save(cp, "exp1") + loaded = store.load("exp1") + assert loaded is not None + assert loaded.experiment_id == "exp1" + + def test_rejects_path_traversal_in_filename(self, tmp_path): + store = FileCheckpointStore(str(tmp_path)) + # The key is hashed + sanitized, so traversal characters cannot escape. + store.save(ExperimentCheckpoint(experiment_id="exp-evil"), "../evil") + files = os.listdir(str(tmp_path)) + # No file should contain ".." or "/". + for f in files: + assert ".." not in f and "/" not in f + + def test_rejects_nul_in_key(self, tmp_path): + store = FileCheckpointStore(str(tmp_path)) + with pytest.raises(ValueError, match="NUL"): + store.save(ExperimentCheckpoint(experiment_id="exp"), "ex\x00p") + + def test_survives_exception_mid_write(self, tmp_path, monkeypatch): + store = FileCheckpointStore(str(tmp_path)) + cp = ExperimentCheckpoint(experiment_id="exp1") + + original_replace = os.replace + call_state = {"called": False} + + def boom(src, dst): + call_state["called"] = True + # Simulate the first os.replace failing; the temp file is cleaned. + raise OSError("disk full") + + monkeypatch.setattr(os, "replace", boom) + with pytest.raises(OSError): + store.save(cp, "exp1") + assert call_state["called"] + # No partial checkpoint should be loadable. + monkeypatch.setattr(os, "replace", original_replace) + assert store.load("exp1") is None + + def test_debounce_flush(self, tmp_path): + store = FileCheckpointStore(str(tmp_path), flush_interval_s=1.0) + assert not store.should_flush(0.0) + assert store.should_flush(1.0) + store.mark_flushed(1.0) + assert not store.should_flush(1.5) + assert store.should_flush(2.5) + + +# --------------------------------------------------------------------------- +# Integration: task retry / timeout / failures recorded +# --------------------------------------------------------------------------- + + +def _run_scores(create_score: MagicMock) -> list: + return [ + call.kwargs + for call in create_score.call_args_list + if call.kwargs.get("dataset_run_id") is not None + ] + + +class TestTaskRetryIntegration: + def test_task_retried_then_succeeds( + self, langfuse_memory_client, monkeypatch, get_span, find_spans + ): + monkeypatch.setattr(langfuse_memory_client, "_get_project_id", lambda: "p") + monkeypatch.setattr(langfuse_memory_client, "create_score", MagicMock()) + calls = {"n": 0} + + def task(*, item): + calls["n"] += 1 + if calls["n"] < 3: + raise _FakeHTTPError(429) + return "ok:%s" % item["input"] + + policy = ExecutionPolicy( + task_retry=RetryPolicy(max_attempts=5, initial_delay=0, jitter=False) + ) + result = langfuse_memory_client.run_experiment( + name="retry", + data=[{"input": "a"}], + task=task, + execution=policy, + max_concurrency=1, + ) + assert calls["n"] == 3 + assert len(result.item_results) == 1 + assert result.item_results[0].output == "ok:a" + assert result.failures == [] + # Attempt observations emitted for the two failed retries. + attempt_spans = find_spans("experiment-item-task-attempt") + assert len(attempt_spans) == 2 + for s in attempt_spans: + assert ( + s.attributes.get(LangfuseOtelSpanAttributes.OBSERVATION_LEVEL) + == "WARNING" + ) + + def test_non_retryable_error_fails_once(self, langfuse_memory_client, monkeypatch): + monkeypatch.setattr(langfuse_memory_client, "_get_project_id", lambda: "p") + monkeypatch.setattr(langfuse_memory_client, "create_score", MagicMock()) + calls = {"n": 0} + + def task(*, item): + calls["n"] += 1 + raise ValueError("user error") + + policy = ExecutionPolicy( + task_retry=RetryPolicy(max_attempts=5, initial_delay=0, jitter=False) + ) + result = langfuse_memory_client.run_experiment( + name="noretry", + data=[{"input": "a"}], + task=task, + execution=policy, + max_concurrency=1, + ) + assert calls["n"] == 1 + assert result.item_results == [] + assert len(result.failures) == 1 + assert result.failures[0].stage == "task" + assert result.failures[0].error_type == "ValueError" + + def test_task_timeout_recorded(self, langfuse_memory_client, monkeypatch): + monkeypatch.setattr(langfuse_memory_client, "_get_project_id", lambda: "p") + monkeypatch.setattr(langfuse_memory_client, "create_score", MagicMock()) + + async def task(*, item): + await asyncio.sleep(5) + return "never" + + policy = ExecutionPolicy( + task_retry=RetryPolicy(max_attempts=1), + timeouts=TimeoutPolicy(task_timeout=0.05), + ) + result = langfuse_memory_client.run_experiment( + name="timeout", + data=[{"input": "a"}], + task=task, + execution=policy, + max_concurrency=1, + ) + assert result.item_results == [] + assert len(result.failures) == 1 + assert result.failures[0].stage == "timeout" + assert result.failures[0].error_type == "TaskTimeoutError" + + +class TestEvaluatorRetryIntegration: + def test_evaluator_retry_and_timeout_recorded_item_succeeds( + self, langfuse_memory_client, monkeypatch + ): + monkeypatch.setattr(langfuse_memory_client, "_get_project_id", lambda: "p") + monkeypatch.setattr(langfuse_memory_client, "create_score", MagicMock()) + calls = {"n": 0} + + def failing_evaluator(*, input, output, expected_output, metadata): + calls["n"] += 1 + if calls["n"] < 3: + raise _FakeHTTPError(503) + raise asyncio.TimeoutError() + + policy = ExecutionPolicy( + evaluator_retry=RetryPolicy(max_attempts=5, initial_delay=0, jitter=False), + timeouts=TimeoutPolicy(evaluator_timeout=0.05), + ) + result = langfuse_memory_client.run_experiment( + name="evalretry", + data=[{"input": "a"}], + task=lambda *, item: item["input"], + evaluators=[failing_evaluator], + execution=policy, + max_concurrency=1, + ) + # Item still succeeds; evaluator failure recorded. + assert len(result.item_results) == 1 + assert result.item_results[0].output == "a" + assert len(result.evaluator_failures) == 1 + assert result.evaluator_failures[0].error_type == "EvaluatorTimeoutError" + assert result.evaluator_failures[0].evaluator_name == "failing_evaluator" + + +class TestFailuresRecordedDefaultPath: + def test_failures_recorded_item_results_unchanged( + self, langfuse_memory_client, monkeypatch + ): + monkeypatch.setattr(langfuse_memory_client, "_get_project_id", lambda: "p") + monkeypatch.setattr(langfuse_memory_client, "create_score", MagicMock()) + + def task(*, item): + if item["input"] == "fail": + raise RuntimeError("boom") + return item["input"] + + result = langfuse_memory_client.run_experiment( + name="mixed", + data=[{"input": "ok"}, {"input": "fail"}], + task=task, + max_concurrency=2, + ) + # Only the successful item appears in item_results. + assert [r.output for r in result.item_results] == ["ok"] + # The failure is recorded. + assert len(result.failures) == 1 + assert result.failures[0].error_type == "RuntimeError" + assert result.execution is not None + assert result.execution.succeeded == 1 + assert result.execution.failed == 1 + + +class TestRunEvaluatorFailuresKwarg: + def test_run_evaluator_receives_failures_kwarg( + self, langfuse_memory_client, monkeypatch + ): + monkeypatch.setattr(langfuse_memory_client, "_get_project_id", lambda: "p") + monkeypatch.setattr(langfuse_memory_client, "create_score", MagicMock()) + + def task(*, item): + if item["input"] == "fail": + raise RuntimeError("boom") + return item["input"] + + seen = {} + + def run_eval(*, item_results, failures=None, **kwargs): + seen["failures"] = failures + return Evaluation(name="count", value=len(item_results)) + + langfuse_memory_client.run_experiment( + name="kwarg", + data=[{"input": "ok"}, {"input": "fail"}], + task=task, + run_evaluators=[run_eval], + max_concurrency=2, + ) + assert seen["failures"] is not None + assert len(seen["failures"]) == 1 + + def test_run_evaluator_without_failures_kwarg_unaffected( + self, langfuse_memory_client, monkeypatch + ): + monkeypatch.setattr(langfuse_memory_client, "_get_project_id", lambda: "p") + monkeypatch.setattr(langfuse_memory_client, "create_score", MagicMock()) + + def task(*, item): + return item["input"] + + def run_eval(*, item_results, **kwargs): + # No `failures` should be passed. + assert "failures" not in kwargs + return Evaluation(name="count", value=len(item_results)) + + result = langfuse_memory_client.run_experiment( + name="nokwarg", + data=[{"input": "a"}], + task=task, + run_evaluators=[run_eval], + max_concurrency=1, + ) + assert len(result.run_evaluations) == 1 + + +# --------------------------------------------------------------------------- +# Failure budget / run timeout / cancellation / resume +# --------------------------------------------------------------------------- + + +class TestFailureBudget: + def test_max_failures_stops_scheduling(self, langfuse_memory_client, monkeypatch): + monkeypatch.setattr(langfuse_memory_client, "_get_project_id", lambda: "p") + monkeypatch.setattr(langfuse_memory_client, "create_score", MagicMock()) + started = {"n": 0} + + def task(*, item): + started["n"] += 1 + raise RuntimeError("boom") + + policy = ExecutionPolicy( + failures=FailurePolicy(max_failures=1, on_threshold="stop"), + ) + result = langfuse_memory_client.run_experiment( + name="budget", + data=[{"input": "a"}, {"input": "b"}, {"input": "c"}], + task=task, + execution=policy, + max_concurrency=1, + ) + assert result.execution is not None + assert result.execution.aborted is True + assert "failure budget" in (result.execution.abort_reason or "") + # With max_concurrency=1 and max_failures=1, only 1 item is attempted. + assert started["n"] == 1 + + def test_max_failure_rate_stops_scheduling( + self, langfuse_memory_client, monkeypatch + ): + monkeypatch.setattr(langfuse_memory_client, "_get_project_id", lambda: "p") + monkeypatch.setattr(langfuse_memory_client, "create_score", MagicMock()) + started = {"n": 0} + + def task(*, item): + started["n"] += 1 + raise RuntimeError("boom") + + policy = ExecutionPolicy( + failures=FailurePolicy(max_failure_rate=0.3, min_items_for_rate=2), + ) + result = langfuse_memory_client.run_experiment( + name="ratebudget", + data=[{"input": "a"}, {"input": "b"}, {"input": "c"}], + task=task, + execution=policy, + max_concurrency=1, + ) + assert result.execution is not None + assert result.execution.aborted is True + + +class TestRunTimeout: + def test_run_timeout_aborts(self, langfuse_memory_client, monkeypatch): + monkeypatch.setattr(langfuse_memory_client, "_get_project_id", lambda: "p") + monkeypatch.setattr(langfuse_memory_client, "create_score", MagicMock()) + + async def task(*, item): + await asyncio.sleep(0.2) + return item["input"] + + policy = ExecutionPolicy(timeouts=TimeoutPolicy(run_timeout=0.05)) + result = langfuse_memory_client.run_experiment( + name="runtimeout", + data=[{"input": "a"}, {"input": "b"}], + task=task, + execution=policy, + max_concurrency=1, + ) + assert result.execution is not None + assert result.execution.aborted is True + assert "run timeout" in (result.execution.abort_reason or "") + + +class TestCancellation: + @pytest.mark.asyncio + async def test_cancellation_writes_checkpoint_and_reraises( + self, langfuse_memory_client, monkeypatch, tmp_path + ): + monkeypatch.setattr(langfuse_memory_client, "_get_project_id", lambda: "p") + monkeypatch.setattr(langfuse_memory_client, "create_score", MagicMock()) + store = FileCheckpointStore(str(tmp_path)) + # asyncio.Event (not threading.Event): a threading.Event.wait() awaited + # in a coroutine blocks the event loop, so the task coroutine would + # never get a chance to set it and the test deadlocks. + started = asyncio.Event() + + async def task(*, item): + started.set() + await asyncio.sleep(5) + return item["input"] + + policy = ExecutionPolicy( + checkpoint=store, + resume=True, + task_retry=RetryPolicy(max_attempts=1), + ) + + async def main(): + coro = langfuse_memory_client._run_experiment_async( + name="cancel", + run_name="r", + description=None, + data=[{"input": "a"}], + task=task, + evaluators=[], + composite_evaluator=None, + run_evaluators=[], + max_concurrency=1, + execution=policy, + ) + task_obj = asyncio.ensure_future(coro) + await started.wait() + task_obj.cancel() + with pytest.raises(asyncio.CancelledError): + await task_obj + + await main() + # The checkpoint is stored under the derived name+run_name key; the + # experiment id is random for local data and cannot serve as the key. + cp = store.load(experiment_checkpoint_key("cancel", "r")) + assert cp is not None + files = [f for f in os.listdir(str(tmp_path)) if f.endswith(".json")] + assert len(files) >= 1 + + +class TestResume: + def test_resume_skips_completed_retries_failed( + self, langfuse_memory_client, monkeypatch + ): + monkeypatch.setattr(langfuse_memory_client, "_get_project_id", lambda: "p") + monkeypatch.setattr(langfuse_memory_client, "create_score", MagicMock()) + store = InMemoryCheckpointStore() + + # Pre-populate a checkpoint with one completed and one failed item. + data = [ + {"input": "done", "expected_output": "done"}, + {"input": "fail", "expected_output": "fail"}, + {"input": "new", "expected_output": "new"}, + ] + keys = make_item_keys(data) + cp = ExperimentCheckpoint( + experiment_id="exp-resume", + experiment_name="resume", + run_name="r", + fingerprint=ExperimentFingerprint( + task_name="task", + evaluator_names=[], + data_size=3, + data_digest=compute_data_digest(keys), + ), + items={ + keys[0]: CheckpointItem( + state="completed", + output_json='"done"', + evaluations=[], + trace_id="t-done", + attempts=1, + ), + keys[1]: CheckpointItem( + state="failed", + error="boom", + attempts=1, + ), + }, + created_at=1.0, + updated_at=2.0, + status="running", + ) + store.save(cp, "exp-resume") + + task_calls = {"n": 0} + + def task(*, item): + task_calls["n"] += 1 + return item["input"] + + policy = ExecutionPolicy(checkpoint=store, resume=True) + result = langfuse_memory_client.run_experiment( + name="resume", + run_name="r", + experiment_id="exp-resume", + data=data, + task=task, + execution=policy, + max_concurrency=1, + ) + # The completed item is NOT re-executed; the failed and new ones are. + assert task_calls["n"] == 2 + assert len(result.item_results) == 3 + # The resumed completed item is rebuilt from the checkpoint. + resumed = [r for r in result.item_results if r.output == "done"] + assert len(resumed) == 1 + assert resumed[0].trace_id == "t-done" + assert result.execution is not None + assert result.execution.skipped_resumed == 1 + + def test_fingerprint_mismatch_raises(self, langfuse_memory_client, monkeypatch): + monkeypatch.setattr(langfuse_memory_client, "_get_project_id", lambda: "p") + monkeypatch.setattr(langfuse_memory_client, "create_score", MagicMock()) + store = InMemoryCheckpointStore() + data = [{"input": "a"}] + keys = make_item_keys(data) + cp = ExperimentCheckpoint( + experiment_id="exp-mismatch", + experiment_name="m", + run_name="r", + fingerprint=ExperimentFingerprint( + task_name="other_task", + evaluator_names=[], + data_size=1, + data_digest=compute_data_digest(keys), + ), + items={}, + created_at=1.0, + updated_at=2.0, + ) + store.save(cp, "exp-mismatch") + + def task(*, item): + return item["input"] + + policy = ExecutionPolicy(checkpoint=store, resume=True) + with pytest.raises(CheckpointMismatchError): + langfuse_memory_client.run_experiment( + name="m", + run_name="r", + experiment_id="exp-mismatch", + data=data, + task=task, + execution=policy, + max_concurrency=1, + ) + + def test_resume_force_ignores_mismatch(self, langfuse_memory_client, monkeypatch): + monkeypatch.setattr(langfuse_memory_client, "_get_project_id", lambda: "p") + monkeypatch.setattr(langfuse_memory_client, "create_score", MagicMock()) + store = InMemoryCheckpointStore() + data = [{"input": "a"}] + keys = make_item_keys(data) + cp = ExperimentCheckpoint( + experiment_id="exp-force", + experiment_name="m", + run_name="r", + fingerprint=ExperimentFingerprint( + task_name="other_task", + evaluator_names=[], + data_size=1, + data_digest=compute_data_digest(keys), + ), + items={}, + created_at=1.0, + updated_at=2.0, + ) + store.save(cp, "exp-force") + + def task(*, item): + return item["input"] + + policy = ExecutionPolicy(checkpoint=store, resume="force") + result = langfuse_memory_client.run_experiment( + name="m", + run_name="r", + experiment_id="exp-force", + data=data, + task=task, + execution=policy, + max_concurrency=1, + ) + assert len(result.item_results) == 1 + + +class TestLocalDataResume: + """Resume must work for local (non-dataset) data too. + + Local data gets a random experiment id per run, so the checkpoint is + keyed by a digest of (experiment name, run name) instead, and a resumed + run adopts the id recorded in the checkpoint. + """ + + @pytest.mark.asyncio + async def test_crash_mid_run_then_resume_skips_completed_and_adopts_id( + self, langfuse_memory_client, monkeypatch + ): + monkeypatch.setattr(langfuse_memory_client, "_get_project_id", lambda: "p") + monkeypatch.setattr(langfuse_memory_client, "create_score", MagicMock()) + store = InMemoryCheckpointStore() + data = [{"input": "a"}, {"input": "b"}, {"input": "c"}] + + started_item3 = asyncio.Event() + block_item3 = True + calls = {"n": 0} + + async def task(*, item): + calls["n"] += 1 + if item["input"] == "c" and block_item3: + started_item3.set() + await asyncio.sleep(30) + return item["input"] + + policy = ExecutionPolicy(checkpoint=store, resume=True) + + async def run_and_cancel(): + # Keyboard-interrupt-style cancellation of the whole run while + # item "c" is in flight; items "a" and "b" already completed. + run_task = asyncio.ensure_future( + langfuse_memory_client._run_experiment_async( + name="local-resume", + run_name="r", + description=None, + data=data, + task=task, + evaluators=[], + composite_evaluator=None, + run_evaluators=[], + max_concurrency=1, + execution=policy, + ) + ) + await started_item3.wait() + run_task.cancel() + with pytest.raises(asyncio.CancelledError): + await run_task + + await run_and_cancel() + calls_after_run1 = calls["n"] + assert calls_after_run1 == 3 + + # The checkpoint lives under the derived name+run_name key, not under + # the random per-run experiment id. + key = experiment_checkpoint_key("local-resume", "r") + cp = store.load(key) + assert cp is not None + run1_experiment_id = cp.experiment_id + assert run1_experiment_id + assert set(cp.items) == set(make_item_keys(data)[:2]) + assert all(cp.items[k].state == "completed" for k in cp.items) + + # Second run with the same run_name: completed items are skipped + # (only item "c" executes) and both runs share one experiment id. + block_item3 = False + result = langfuse_memory_client.run_experiment( + name="local-resume", + run_name="r", + data=data, + task=task, + execution=ExecutionPolicy(checkpoint=store, resume=True), + max_concurrency=1, + ) + assert calls["n"] - calls_after_run1 == 1 + assert [r.output for r in result.item_results] == ["a", "b", "c"] + assert result.experiment_id == run1_experiment_id + assert {r.experiment_id for r in result.item_results} == {run1_experiment_id} + assert result.execution is not None + assert result.execution.skipped_resumed == 2 + assert result.execution.failed == 0 + + def test_local_resume_without_run_name_or_experiment_id_raises( + self, langfuse_memory_client, monkeypatch + ): + monkeypatch.setattr(langfuse_memory_client, "_get_project_id", lambda: "p") + monkeypatch.setattr(langfuse_memory_client, "create_score", MagicMock()) + store = InMemoryCheckpointStore() + policy = ExecutionPolicy(checkpoint=store, resume=True) + + def task(*, item): + return item["input"] if isinstance(item, dict) else item.input + + with pytest.raises(ValueError, match="run_name"): + langfuse_memory_client.run_experiment( + name="local-resume", + data=[{"input": "a"}], + task=task, + execution=policy, + ) + + # Dataset items have a deterministic experiment id, so resuming them + # does not require an explicit run_name. + result = langfuse_memory_client.run_experiment( + name="dataset-resume", + data=[_dataset_item(item_id="i1", dataset_id="d", input="a")], + task=task, + execution=policy, + ) + assert len(result.item_results) == 1 + assert result.failures == [] + + def test_checkpoint_key_is_stable_and_distinct(self): + key = experiment_checkpoint_key("exp", "r") + assert key == experiment_checkpoint_key("exp", "r") + assert key != experiment_checkpoint_key("exp", "other") + assert key != experiment_checkpoint_key("other", "r") + assert len(key) == 32 + + +class TestRetriedItemsCounting: + def test_item_retried_three_times_counts_once( + self, langfuse_memory_client, monkeypatch + ): + monkeypatch.setattr(langfuse_memory_client, "_get_project_id", lambda: "p") + monkeypatch.setattr(langfuse_memory_client, "create_score", MagicMock()) + calls = {"n": 0} + + def task(*, item): + calls["n"] += 1 + if calls["n"] <= 3: + raise _FakeHTTPError(429) + return "ok" + + policy = ExecutionPolicy( + task_retry=RetryPolicy(max_attempts=5, initial_delay=0, jitter=False) + ) + result = langfuse_memory_client.run_experiment( + name="retried-count", + data=[{"input": "a"}], + task=task, + execution=policy, + max_concurrency=1, + ) + assert result.item_results[0].output == "ok" + assert result.execution is not None + # One item was retried, no matter how many retry attempts it used. + assert result.execution.retried_items == 1 + assert result.execution.total_attempts == 4 + + def test_two_items_each_retried_once_count_two( + self, langfuse_memory_client, monkeypatch + ): + monkeypatch.setattr(langfuse_memory_client, "_get_project_id", lambda: "p") + monkeypatch.setattr(langfuse_memory_client, "create_score", MagicMock()) + calls = {"n": 0} + + def task(*, item): + calls["n"] += 1 + # First call per item fails, second succeeds (concurrency 1). + if calls["n"] % 2 == 1: + raise _FakeHTTPError(429) + return item["input"] + + policy = ExecutionPolicy( + task_retry=RetryPolicy(max_attempts=5, initial_delay=0, jitter=False) + ) + result = langfuse_memory_client.run_experiment( + name="retried-count-two", + data=[{"input": "a"}, {"input": "b"}], + task=task, + execution=policy, + max_concurrency=1, + ) + assert result.execution is not None + assert result.execution.retried_items == 2 + assert result.execution.total_attempts == 4 + + def test_terminal_failure_counts_item_once( + self, langfuse_memory_client, monkeypatch + ): + monkeypatch.setattr(langfuse_memory_client, "_get_project_id", lambda: "p") + monkeypatch.setattr(langfuse_memory_client, "create_score", MagicMock()) + + def task(*, item): + raise _FakeHTTPError(503) + + policy = ExecutionPolicy( + task_retry=RetryPolicy(max_attempts=3, initial_delay=0, jitter=False) + ) + result = langfuse_memory_client.run_experiment( + name="retried-count-terminal", + data=[{"input": "a"}], + task=task, + execution=policy, + max_concurrency=1, + ) + assert result.execution is not None + assert result.execution.retried_items == 1 + assert result.execution.total_attempts == 3 + assert len(result.failures) == 1 + + +class TestDefaultPathSpanTree: + def test_no_policy_span_names_and_count_match_pre_change_code( + self, langfuse_memory_client, monkeypatch, memory_exporter + ): + """With no execution policy the run must emit exactly the span tree + the pre-policy code emitted (same names, same count, no attempt + observations) and the same item-level score calls.""" + monkeypatch.setattr(langfuse_memory_client, "_get_project_id", lambda: "p") + create_score = MagicMock() + monkeypatch.setattr(langfuse_memory_client, "create_score", create_score) + + def task(*, item): + return item["input"] + + def eval1(*, input, output, expected_output, metadata): + return Evaluation(name="acc", value=1.0) + + result = langfuse_memory_client.run_experiment( + name="default", + data=[{"input": "a"}], + task=task, + evaluators=[eval1], + max_concurrency=1, + ) + langfuse_memory_client.flush() + + assert len(result.item_results) == 1 + names = [s.name for s in memory_exporter.get_finished_spans()] + # Expected list pinned against the pre-change (main) behavior. + assert sorted(names) == sorted( + [ + "experiment-item-run", + "experiment-item-task", + "experiment-item-evaluation", + "eval1", + ] + ) + assert len(names) == 4 + assert "experiment-item-task-attempt" not in names + item_scores = [ + c.kwargs + for c in create_score.call_args_list + if c.kwargs.get("trace_id") is not None + ] + assert len(item_scores) == 1 + assert item_scores[0]["name"] == "acc" + + +class TestResumeFailureAccounting: + def test_resumed_item_success_not_counted_as_failed( + self, langfuse_memory_client, monkeypatch + ): + """A resumed run that retries a previously-failed checkpoint item and + succeeds must not count it in ``ExecutionSummary.failed``.""" + monkeypatch.setattr(langfuse_memory_client, "_get_project_id", lambda: "p") + monkeypatch.setattr(langfuse_memory_client, "create_score", MagicMock()) + store = InMemoryCheckpointStore() + data = [ + {"input": "done", "expected_output": "done"}, + {"input": "flaky", "expected_output": "flaky"}, + ] + keys = make_item_keys(data) + cp = ExperimentCheckpoint( + experiment_id="exp-accounting", + experiment_name="accounting", + run_name="r", + fingerprint=ExperimentFingerprint( + task_name="task", + evaluator_names=[], + data_size=2, + data_digest=compute_data_digest(keys), + ), + items={ + keys[0]: CheckpointItem( + state="completed", + output_json='"done"', + evaluations=[], + trace_id="t-done", + attempts=1, + ), + keys[1]: CheckpointItem( + state="failed", + error="boom", + attempts=2, + ), + }, + created_at=1.0, + updated_at=2.0, + status="running", + ) + store.save(cp, "exp-accounting") + + def task(*, item): + return item["input"] + + result = langfuse_memory_client.run_experiment( + name="accounting", + run_name="r", + experiment_id="exp-accounting", + data=data, + task=task, + execution=ExecutionPolicy(checkpoint=store, resume=True), + max_concurrency=1, + ) + assert result.execution is not None + # The previously-failed item was retried and succeeded: it counts as + # succeeded, not failed. + assert result.execution.failed == 0 + assert result.execution.succeeded == 2 + assert result.execution.skipped_resumed == 1 + assert result.failures == [] + # The checkpoint no longer marks the item as failed. + reloaded = store.load("exp-accounting") + assert reloaded is not None + assert reloaded.items[keys[1]].state == "completed" + + +class TestNoFunctionLevelImports: + """Function-level imports added by the execution-policy change must live + at module top (repo style; ruff does not currently flag them).""" + + def test_experiment_package_has_no_function_level_imports(self): + import ast + from pathlib import Path + + import langfuse._experiment as experiment_pkg + + package_dir = Path(experiment_pkg.__file__).parent + offenders = [] + for path in sorted(package_dir.glob("*.py")): + tree = ast.parse(path.read_text(encoding="utf-8")) + for node in ast.walk(tree): + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): + for child in ast.walk(node): + if child is node: + continue + if isinstance(child, (ast.Import, ast.ImportFrom)): + offenders.append("%s:%d" % (path.name, child.lineno)) + assert offenders == [] + + def test_client_has_no_local_time_or_random_imports(self): + from pathlib import Path + + import langfuse._client.client as client_module + + source = Path(client_module.__file__).read_text(encoding="utf-8") + assert not re.search(r"^\s+import (time|random) as ", source, re.MULTILINE) + + +class TestDefaultPathIdentical: + def test_no_policy_identical_span_tree_and_scores( + self, langfuse_memory_client, monkeypatch, find_spans + ): + monkeypatch.setattr(langfuse_memory_client, "_get_project_id", lambda: "p") + create_score = MagicMock() + monkeypatch.setattr(langfuse_memory_client, "create_score", create_score) + + def task(*, item): + return item["input"] + + def eval1(*, input, output, expected_output, metadata): + return Evaluation(name="acc", value=1.0) + + result = langfuse_memory_client.run_experiment( + name="default", + data=[{"input": "a"}], + task=task, + evaluators=[eval1], + max_concurrency=1, + ) + # No retry-attempt spans by default. + assert find_spans("experiment-item-task-attempt") == [] + # One item-level score. + item_scores = [ + c.kwargs + for c in create_score.call_args_list + if c.kwargs.get("trace_id") is not None + ] + assert len(item_scores) == 1 + # Failures empty, execution summary present. + assert result.failures == [] + assert result.evaluator_failures == [] + assert result.execution is not None + assert result.execution.succeeded == 1 + assert result.execution.total_attempts == 1 + + +# --------------------------------------------------------------------------- +# Signature drift guard: execution param threaded through +# --------------------------------------------------------------------------- + + +class TestExecutionParamForwarding: + def test_runner_context_forwards_execution(self): + from unittest.mock import MagicMock + + client = MagicMock(spec=Langfuse) + client.run_experiment.return_value = "sentinel" + ctx = RunnerContext(client=client, data=[{"input": "a"}]) + policy = ExecutionPolicy() + ctx.run_experiment(name="exp", task=lambda *, item: None, execution=policy) + kwargs = client.run_experiment.call_args.kwargs + assert kwargs["execution"] is policy + + def test_dataset_client_forwards_execution(self): + langfuse_client = MagicMock() + item = _dataset_item(item_id="i1", dataset_id="d", input="a") + dataset = DatasetClient( + dataset=MagicMock(), + items=[item], + langfuse_client=langfuse_client, + ) + policy = ExecutionPolicy() + dataset.run_experiment(name="exp", task=lambda *, item: None, execution=policy) + assert langfuse_client.run_experiment.call_args.kwargs["execution"] is policy