diff --git a/langfuse/_client/client.py b/langfuse/_client/client.py index e7be8214b..84bb163fb 100644 --- a/langfuse/_client/client.py +++ b/langfuse/_client/client.py @@ -9,6 +9,7 @@ import re import urllib.parse import uuid +from concurrent.futures import ThreadPoolExecutor from datetime import datetime, timezone from hashlib import sha256 from time import time_ns @@ -2430,6 +2431,7 @@ def run_experiment( composite_evaluator: Optional[CompositeEvaluatorFunction] = None, run_evaluators: List[RunEvaluatorFunction] = [], max_concurrency: int = 50, + run_sync_tasks_in_threads: bool = False, metadata: Optional[Dict[str, str]] = None, _dataset_version: Optional[datetime] = None, ) -> ExperimentResult: @@ -2475,6 +2477,17 @@ def run_experiment( max_concurrency: Maximum number of concurrent task executions (default: 50). Controls the number of items processed simultaneously. Adjust based on API rate limits and system resources. + run_sync_tasks_in_threads: Opt in to running synchronous tasks in a + dedicated thread pool with up to `max_concurrency` workers (default: + False). This enables concurrency for blocking I/O without changing + the default event-loop execution. Async tasks stay on the experiment + loop. When enabled, tasks must be thread-safe: create and close + thread-bound resources, such as SQLite connections, inside the task. + Even `max_concurrency=1` offloads synchronous tasks. Exceptions are + still logged and failed items omitted. Cancellation does not stop + tasks already running in threads, which can continue after the + experiment returns and delay Python process exit. Use finite I/O + timeouts. For cooperative cancellation, prefer async tasks. metadata: Optional metadata dictionary to attach to all experiment traces. This metadata will be included in every trace created during the experiment and is shown as the experiment's metadata in Langfuse. @@ -2617,6 +2630,7 @@ def average_accuracy(*, item_results, **kwargs): composite_evaluator=composite_evaluator, run_evaluators=run_evaluators or [], max_concurrency=max_concurrency, + run_sync_tasks_in_threads=run_sync_tasks_in_threads, metadata=metadata, dataset_version=_dataset_version, ), @@ -2635,6 +2649,7 @@ async def _run_experiment_async( composite_evaluator: Optional[CompositeEvaluatorFunction], run_evaluators: List[RunEvaluatorFunction], max_concurrency: int, + run_sync_tasks_in_threads: bool = False, metadata: Optional[Dict[str, Any]] = None, dataset_version: Optional[datetime] = None, ) -> ExperimentResult: @@ -2673,6 +2688,11 @@ async def _run_experiment_async( # Set up concurrency control semaphore = asyncio.Semaphore(max_concurrency) + task_executor = ( + ThreadPoolExecutor(max_workers=max_concurrency) + if run_sync_tasks_in_threads + else None + ) # Process all items async def process_item(item: ExperimentItem) -> ExperimentItemResult: @@ -2688,11 +2708,18 @@ async def process_item(item: ExperimentItem) -> ExperimentItemResult: description, metadata, experiment_item_version, + task_executor=task_executor, ) # Run all items concurrently - tasks = [process_item(item) for item in data] - item_results = await asyncio.gather(*tasks, return_exceptions=True) + try: + tasks = [process_item(item) for item in data] + item_results = await asyncio.gather(*tasks, return_exceptions=True) + finally: + if task_executor is not None: + # Do not block the event loop on cancellation. Running calls + # cannot be stopped, but queued calls must not start afterward. + task_executor.shutdown(wait=False, cancel_futures=True) # Filter out any exceptions and log errors valid_results: List[ExperimentItemResult] = [] @@ -2829,6 +2856,8 @@ async def _process_experiment_item( experiment_description: Optional[str], experiment_metadata: Optional[Dict[str, Any]] = None, experiment_item_version: Optional[str] = None, + *, + task_executor: Optional[ThreadPoolExecutor] = None, ) -> ExperimentItemResult: with self.start_as_current_observation(name="experiment-item-run") as span: try: @@ -2942,7 +2971,7 @@ async def _process_experiment_item( # metadata instead of leaving no room for the output. task_span.update(metadata=final_observation_metadata) try: - output = await _run_task(task, item) + output = await _run_task(task, item, executor=task_executor) except Exception as e: task_span.update( output=f"Error: {str(e)}", diff --git a/langfuse/_client/datasets.py b/langfuse/_client/datasets.py index 1070c76d3..8a9a25296 100644 --- a/langfuse/_client/datasets.py +++ b/langfuse/_client/datasets.py @@ -89,6 +89,7 @@ def run_experiment( composite_evaluator: Optional[CompositeEvaluatorFunction] = None, run_evaluators: List[RunEvaluatorFunction] = [], max_concurrency: int = 50, + run_sync_tasks_in_threads: bool = False, metadata: Optional[Dict[str, Any]] = None, ) -> ExperimentResult: """Run an experiment on this Langfuse dataset with automatic tracking. @@ -126,6 +127,12 @@ def run_experiment( Useful for computing aggregate statistics across all dataset items. max_concurrency: Maximum number of concurrent task executions (default: 50). Adjust based on API rate limits and system resources. + run_sync_tasks_in_threads: Opt in to concurrent synchronous tasks in a + dedicated pool with up to `max_concurrency` workers (default: False). + Tasks must be thread-safe; create thread-bound resources inside the + task. Async tasks stay on the experiment loop. Cancellation cannot + stop synchronous tasks already running in threads. See + `Langfuse.run_experiment` for resource ownership and timeout guidance. metadata: Optional metadata to attach to the experiment run and all traces. Will be combined with individual item metadata. @@ -300,6 +307,7 @@ def content_diversity(*, item_results, **kwargs): composite_evaluator=composite_evaluator, run_evaluators=run_evaluators, max_concurrency=max_concurrency, + run_sync_tasks_in_threads=run_sync_tasks_in_threads, metadata=metadata, _dataset_version=self.version, ) diff --git a/langfuse/experiment.py b/langfuse/experiment.py index ee29d8611..65986d1b7 100644 --- a/langfuse/experiment.py +++ b/langfuse/experiment.py @@ -6,8 +6,12 @@ """ import asyncio +import inspect import warnings +from concurrent.futures import Future, ThreadPoolExecutor +from contextvars import copy_context from datetime import datetime +from functools import partial from typing import ( TYPE_CHECKING, Any, @@ -642,6 +646,18 @@ class TaskFunction(Protocol): - Return any type of output (will be passed to evaluators) - Can be either synchronous or asynchronous - Should handle their own errors gracefully (exceptions will be logged) + + By default, synchronous task functions run on the experiment's event loop, + so blocking calls execute serially. Set `run_sync_tasks_in_threads=True` on + `run_experiment` to run synchronous tasks concurrently in worker threads. + Asynchronous tasks, including decorated wrappers, stay on the event loop. + + When opting in, synchronous tasks must be thread-safe. Create and close + thread-bound resources such as sqlite3 connections inside the task, rather + than capturing a connection created on the caller's thread. Even + `max_concurrency=1` offloads the task. Task exceptions are logged and failed + items are omitted from the result. Cancellation cannot stop a synchronous + task already running in a thread; use finite timeouts on blocking I/O. """ def __call__( @@ -1056,9 +1072,61 @@ async def _run_evaluator( return [] -async def _run_task(task: TaskFunction, item: ExperimentItem) -> Any: +def _is_async_task(task: Any) -> bool: + """Inspect async task metadata without replacing the callable we invoke.""" + seen: set[int] = set() + check_call = True + for _ in range(100): + if id(task) in seen: + return False + seen.add(id(task)) + if isinstance(task, partial): + task = task.func + continue + try: + unwrapped = inspect.unwrap(task, stop=inspect.iscoroutinefunction) + except ValueError: + return False + if inspect.iscoroutinefunction(unwrapped): + return True + if unwrapped is not task: + task = unwrapped + elif check_call: + task = getattr(task, "__call__", None) + check_call = False + else: + return False + return False + + +async def _run_task( + task: TaskFunction, + item: ExperimentItem, + *, + executor: Optional[ThreadPoolExecutor] = None, +) -> Any: """Run a task function and handle sync/async.""" - result = task(item=item) + result: Any + if executor is None or _is_async_task(task): + result = task(item=item) + else: + # Copy separately for every item so worker spans inherit that item's + # tracing context. Keep the source future to dispose of a late coroutine + # if cancellation prevents the caller from awaiting it. + context = copy_context() + future = executor.submit(context.run, partial(task, item=item)) + try: + result = await asyncio.wrap_future(future) + except asyncio.CancelledError: + + def discard_result(done: Future[Any]) -> None: + if not done.cancelled() and done.exception() is None: + value = done.result() + if asyncio.iscoroutine(value): + value.close() + + future.add_done_callback(discard_result) + raise # Handle async tasks if asyncio.iscoroutine(result): @@ -1163,6 +1231,7 @@ def run_experiment( composite_evaluator: Optional["CompositeEvaluatorFunction"] = None, run_evaluators: List[RunEvaluatorFunction] = [], max_concurrency: int = 50, + run_sync_tasks_in_threads: bool = False, metadata: Optional[Dict[str, str]] = None, _dataset_version: Optional[datetime] = None, ) -> ExperimentResult: @@ -1192,6 +1261,7 @@ def run_experiment( composite_evaluator=composite_evaluator, run_evaluators=run_evaluators, max_concurrency=max_concurrency, + run_sync_tasks_in_threads=run_sync_tasks_in_threads, metadata=merged_metadata, _dataset_version=resolved_dataset_version, ) diff --git a/tests/unit/test_experiment.py b/tests/unit/test_experiment.py index 192577b6b..5f09d30fd 100644 --- a/tests/unit/test_experiment.py +++ b/tests/unit/test_experiment.py @@ -1,21 +1,29 @@ """Tests for ``langfuse.experiment`` — ``RunnerContext`` and ``RegressionError``.""" +import asyncio import inspect import json +import sqlite3 import typing +from concurrent.futures import ThreadPoolExecutor +from contextlib import closing from datetime import datetime, timezone +from functools import partial, wraps +from threading import Barrier, Event, get_ident from typing import get_type_hints from unittest.mock import MagicMock import pytest +from opentelemetry import context as otel_context from opentelemetry import trace as otel_trace_api from langfuse import Evaluation, RegressionError, RunnerContext from langfuse._client.attributes import LangfuseOtelSpanAttributes from langfuse._client.client import Langfuse -from langfuse.api import DatasetItem, DatasetStatus +from langfuse._client.datasets import DatasetClient +from langfuse.api import Dataset, DatasetItem, DatasetStatus from langfuse.batch_evaluation import CompositeEvaluatorFunction -from langfuse.experiment import ExperimentItemResult, ExperimentResult +from langfuse.experiment import ExperimentItemResult, ExperimentResult, _run_task def _noop_task(*, item, **kwargs): # pragma: no cover - never invoked via mock @@ -28,6 +36,159 @@ def _make_ctx(**kwargs) -> RunnerContext: return RunnerContext(client=client, **kwargs) +def _task_form(task, form): + if "callable" in form: + wrapped_task = task + + class CallableTask: + @wraps(wrapped_task) + def __call__(self, **kwargs): + return wrapped_task(**kwargs) + + task = CallableTask() + + if "partial" in form: + task = partial(task, prefix="partial") + + return task + + +def _experiment_runner(client, entrypoint, inputs): + if entrypoint == "client": + return partial(client.run_experiment, data=inputs) + if entrypoint == "context": + return RunnerContext(client=client, data=inputs).run_experiment + + timestamp = datetime(2026, 1, 1, tzinfo=timezone.utc) + dataset = DatasetClient( + dataset=Dataset( + id="task-dataset", + name="task-dataset", + project_id="project", + metadata=None, + created_at=timestamp, + updated_at=timestamp, + ), + items=[ + DatasetItem( + id=str(i), + input=item["input"], + expected_output=None, + metadata=None, + media_references=[], + status=DatasetStatus.ACTIVE, + dataset_id="task-dataset", + dataset_name="task-dataset", + created_at=timestamp, + updated_at=timestamp, + ) + for i, item in enumerate(inputs) + ], + langfuse_client=client, + ) + return dataset.run_experiment + + +@pytest.fixture +def task_executor(): + with ThreadPoolExecutor(max_workers=2) as executor: + yield executor + + +@pytest.mark.asyncio +@pytest.mark.parametrize("offload", [False, True]) +async def test_async_task_stays_on_experiment_event_loop(offload, task_executor): + loop = asyncio.get_running_loop() + + async def task(*, item): + assert asyncio.get_running_loop() is loop + return item["input"] + + assert ( + await _run_task( + task, {"input": "async-result"}, executor=task_executor if offload else None + ) + == "async-result" + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("offload", [False, True]) +async def test_sync_task_returning_coroutine_is_awaited(offload, task_executor): + loop = asyncio.get_running_loop() + + async def async_result(): + assert asyncio.get_running_loop() is loop + return "coroutine-result" + + def task(*, item): + return async_result() + + assert ( + await _run_task( + task, {"input": "question"}, executor=task_executor if offload else None + ) + == "coroutine-result" + ) + + +@pytest.mark.asyncio +async def test_cancelled_sync_factory_closes_its_unawaited_coroutine(task_executor): + loop = asyncio.get_running_loop() + started = asyncio.Event() + release = Event() + results = [] + + async def result(): + return "cancelled-result" + + def task(*, item): + loop.call_soon_threadsafe(started.set) + release.wait(timeout=3) + coroutine = result() + results.append(coroutine) + return coroutine + + running = asyncio.create_task( + _run_task(task, {"input": "A"}, executor=task_executor) + ) + try: + await asyncio.wait_for(started.wait(), timeout=2) + running.cancel() + with pytest.raises(asyncio.CancelledError): + await running + release.set() + await asyncio.to_thread(task_executor.shutdown, wait=True) + assert len(results) == 1 + assert inspect.getcoroutinestate(results[0]) == inspect.CORO_CLOSED + finally: + release.set() + await asyncio.to_thread(task_executor.shutdown, wait=True) + for coroutine in results: + coroutine.close() + + +@pytest.mark.asyncio +@pytest.mark.timeout(3) +@pytest.mark.parametrize("cycle", ["wrapped", "partial"]) +async def test_cyclic_wrapper_metadata_keeps_sync_task_execution(cycle, task_executor): + caller_thread = get_ident() + worker_threads = [] + + def task(*, item): + worker_threads.append(get_ident()) + return item["input"] + + task.__wrapped__ = task if cycle == "wrapped" else partial(task) + + assert ( + await _run_task(task, {"input": "cycle-result"}, executor=task_executor) + == "cycle-result" + ) + assert len(worker_threads) == 1 + assert worker_threads[0] != caller_thread + + class TestRunnerContextDefaults: def test_context_defaults_flow_through(self): ctx_data = [{"input": "a"}] @@ -248,8 +409,325 @@ def _is_optional(annotation) -> bool: class TestExperimentObservationTree: + @pytest.mark.parametrize("entrypoint", ["client", "context", "dataset"]) + @pytest.mark.parametrize("flag", [None, False]) + def test_default_sync_task_preserves_caller_owned_sqlite_connection( + self, entrypoint, flag, langfuse_memory_client + ): + inputs = [{"input": "A"}, {"input": "B"}] + run_experiment = _experiment_runner(langfuse_memory_client, entrypoint, inputs) + kwargs = {} + if flag is not None: + kwargs["run_sync_tasks_in_threads"] = flag + + caller_thread = get_ident() + with closing(sqlite3.connect(":memory:")) as connection: + + def task(*, item): + assert get_ident() == caller_thread + value = item["input"] if isinstance(item, dict) else item.input + return connection.execute("SELECT ?", (value,)).fetchone()[0] + + result = run_experiment( + name="default-sync-task", + task=task, + max_concurrency=2, + **kwargs, + ) + + assert [item.output for item in result.item_results] == ["A", "B"] + + @pytest.mark.asyncio + async def test_opt_in_sync_tasks_use_experiment_concurrency_not_default_pool( + self, langfuse_memory_client + ): + loop = asyncio.get_running_loop() + loop.set_default_executor(ThreadPoolExecutor(max_workers=1)) + barrier = Barrier(4, timeout=3) + + def task(*, item): + barrier.wait() + return item["input"] + + result = await langfuse_memory_client._run_experiment_async( + name="dedicated-sync-task-pool", + run_name="dedicated-sync-task-pool", + description=None, + data=[{"input": i} for i in range(4)], + task=task, + evaluators=[], + composite_evaluator=None, + run_evaluators=[], + max_concurrency=4, + run_sync_tasks_in_threads=True, + ) + + assert [item.output for item in result.item_results] == [0, 1, 2, 3] + + @pytest.mark.asyncio + async def test_cancelling_experiment_does_not_wait_for_running_sync_task( + self, langfuse_memory_client + ): + loop = asyncio.get_running_loop() + started = asyncio.Event() + release = Event() + finished = Event() + + def task(*, item): + loop.call_soon_threadsafe(started.set) + try: + release.wait(timeout=3) + return item["input"] + finally: + finished.set() + + experiment = asyncio.create_task( + langfuse_memory_client._run_experiment_async( + name="cancel-sync-task", + run_name="cancel-sync-task", + description=None, + data=[{"input": "A"}], + task=task, + evaluators=[], + composite_evaluator=None, + run_evaluators=[], + max_concurrency=1, + run_sync_tasks_in_threads=True, + ) + ) + try: + await asyncio.wait_for(started.wait(), timeout=2) + experiment.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(experiment, timeout=1) + assert not finished.is_set() + finally: + release.set() + await asyncio.to_thread(finished.wait, 2) + + def test_thread_affine_sync_task_records_item_failures( + self, langfuse_memory_client, caplog + ): + with closing(sqlite3.connect(":memory:")) as connection: + + def task(*, item, **kwargs): + return connection.execute("SELECT ?", (item["input"],)).fetchone()[0] + + result = langfuse_memory_client.run_experiment( + name="thread-affine-sync-task", + data=[{"input": "A"}, {"input": "B"}], + task=task, + max_concurrency=2, + run_sync_tasks_in_threads=True, + ) + + assert result.item_results == [] + assert "SQLite objects created in a thread" in caplog.text + + def test_sync_task_can_create_its_own_thread_affine_resource( + self, langfuse_memory_client, tmp_path + ): + database = str(tmp_path / "experiment.sqlite") + with closing(sqlite3.connect(database)): + pass + + def task(*, item, **kwargs): + with closing(sqlite3.connect(database)) as connection: + return connection.execute("SELECT ?", (item["input"],)).fetchone()[0] + + result = langfuse_memory_client.run_experiment( + name="worker-owned-sync-task-resource", + data=[{"input": "A"}, {"input": "B"}], + task=task, + max_concurrency=2, + run_sync_tasks_in_threads=True, + ) + + assert [item.output for item in result.item_results] == ["A", "B"] + + @pytest.mark.skipif( + not hasattr(inspect, "markcoroutinefunction"), + reason="explicit coroutine markers require Python 3.12+", + ) + @pytest.mark.parametrize("task_form", ["marked", "partial", "wrapped"]) + @pytest.mark.parametrize("offload", [False, True]) + def test_marked_coroutine_factory_runs_on_experiment_loop( + self, task_form, offload, langfuse_memory_client + ): + def sync_target(*, item, prefix): + return "the wrapped target must not be invoked" + + @wraps(sync_target) + def factory(*, item, prefix="marked"): + asyncio.get_running_loop() + + async def result(): + return f"{prefix}:{item['input']}" + + return result() + + marked = inspect.markcoroutinefunction(factory) + task = marked + expected_prefix = "marked" + if task_form == "partial": + task = partial(marked, prefix="partial") + expected_prefix = "partial" + elif task_form == "wrapped": + + @wraps(marked) + def wrapped(**kwargs): + return marked(**kwargs) + + task = wrapped + + result = langfuse_memory_client.run_experiment( + name="marked-coroutine-factory", + data=[{"input": "question"}], + task=task, + max_concurrency=1, + run_sync_tasks_in_threads=offload, + ) + assert [item.output for item in result.item_results] == [ + f"{expected_prefix}:question" + ] + + @pytest.mark.parametrize( + "task_form,expected_prefix", + [ + ("wrapped", "wrapped"), + ("partial", "partial"), + ("callable", "wrapped"), + ("partial-callable", "partial"), + ], + ) + @pytest.mark.parametrize("offload", [False, True]) + def test_wrapped_async_task_invokes_wrapper_on_event_loop( + self, task_form, expected_prefix, offload, langfuse_memory_client + ): + wrapper_calls = [] + + async def native_async(*, item, prefix): + return f"{prefix}:{item['input']}" + + @wraps(native_async) + def wrapped(*, item, prefix="wrapped"): + wrapper_calls.append(asyncio.get_running_loop().is_running()) + return native_async(item=item, prefix=prefix) + + result = langfuse_memory_client.run_experiment( + name="wrapped-async-loop", + data=[{"input": "question"}], + task=_task_form(wrapped, task_form), + max_concurrency=1, + run_sync_tasks_in_threads=offload, + ) + + assert [item.output for item in result.item_results] == [ + f"{expected_prefix}:question" + ] + assert wrapper_calls == [True] + + @pytest.mark.parametrize( + "task_form,expected_prefix", + [ + ("wrapped", "wrapped"), + ("partial", "partial"), + ("callable", "wrapped"), + ("partial-callable", "partial"), + ], + ) + @pytest.mark.parametrize("offload", [False, True]) + def test_wrapped_async_task_preserves_wrapper_otel_context( + self, task_form, expected_prefix, offload, langfuse_memory_client, get_span + ): + key = otel_context.create_key("wrapped-task-context") + + async def read_context(*, item, prefix): + await asyncio.sleep(0) + with langfuse_memory_client.start_as_current_observation( + name="wrapped-async-child" + ): + return f"{prefix}:{otel_context.get_value(key)}" + + @wraps(read_context) + def contextual_wrapper(*, item, prefix="wrapped"): + token = otel_context.attach( + otel_context.set_value(key, "wrapped-otel-context") + ) + + async def finish(): + try: + return await read_context(item=item, prefix=prefix) + finally: + otel_context.detach(token) + + return finish() + + result = langfuse_memory_client.run_experiment( + name="wrapped-async-context", + data=[{"input": "question"}], + task=_task_form(contextual_wrapper, task_form), + max_concurrency=1, + run_sync_tasks_in_threads=offload, + ) + + assert [item.output for item in result.item_results] == [ + f"{expected_prefix}:wrapped-otel-context" + ] + child = get_span("wrapped-async-child") + parent = get_span("experiment-item-task") + assert child.parent.span_id == parent.context.span_id + assert child.context.trace_id == parent.context.trace_id + + @pytest.mark.parametrize("entrypoint", ["client", "context", "dataset"]) + def test_sync_tasks_run_concurrently_with_trace_context( + self, entrypoint, langfuse_memory_client, memory_exporter + ): + barrier = Barrier(2, timeout=3) + + def task(*, item): + value = item["input"] if isinstance(item, dict) else item.input + with langfuse_memory_client.start_as_current_observation( + name=f"worker-{value}" + ): + barrier.wait() + return value.upper() + + run_experiment = _experiment_runner( + langfuse_memory_client, + entrypoint, + [{"input": "first"}, {"input": "second"}], + ) + result = run_experiment( + name="sync-task-concurrency", + task=task, + max_concurrency=2, + run_sync_tasks_in_threads=True, + ) + + assert [item.output for item in result.item_results] == ["FIRST", "SECOND"] + + spans = memory_exporter.get_finished_spans() + tasks_by_id = { + span.context.span_id: span + for span in spans + if span.name == "experiment-item-task" + } + workers = [span for span in spans if span.name.startswith("worker-")] + assert len(workers) == 2 + assert len({span.parent.span_id for span in workers}) == 2 + for worker in workers: + parent = tasks_by_id[worker.parent.span_id] + assert worker.context.trace_id == parent.context.trace_id + assert ( + worker.attributes[LangfuseOtelSpanAttributes.EXPERIMENT_ID] + == result.experiment_id + ) + + @pytest.mark.parametrize("offload", [False, True]) def test_failed_task_preserves_experiment_attributes_on_item_run( self, + offload, langfuse_memory_client, get_span, ): @@ -262,6 +740,7 @@ def failing_task(**kwargs): task=failing_task, metadata={"run": "metadata"}, max_concurrency=1, + run_sync_tasks_in_threads=offload, ) item_run = get_span("experiment-item-run")