diff --git a/langfuse/_client/client.py b/langfuse/_client/client.py index fe80d7e85..00214db20 100644 --- a/langfuse/_client/client.py +++ b/langfuse/_client/client.py @@ -2735,8 +2735,8 @@ async def process_item(item: ExperimentItem) -> ExperimentItemResult: except Exception as e: langfuse_logger.error("Failed to store run evaluation: %s", e) - # Flush scores and traces - self.flush() + # Flush scores and traces off the event loop; flush() blocks on network I/O. + await asyncio.to_thread(self.flush) return ExperimentResult( name=name, diff --git a/langfuse/_client/resource_manager.py b/langfuse/_client/resource_manager.py index b5e104812..e633cdfaa 100644 --- a/langfuse/_client/resource_manager.py +++ b/langfuse/_client/resource_manager.py @@ -593,12 +593,36 @@ def flush(self) -> None: self.tracer_provider.force_flush() langfuse_logger.debug("Successfully flushed OTEL tracer provider") - self._score_queue.join() + for score_consumer in self._score_consumers: + if score_consumer.is_alive(): + score_consumer.wake() + + self._join_score_queue() langfuse_logger.debug("Successfully flushed score queue") self._media_upload_queue.join() langfuse_logger.debug("Successfully flushed media upload queue") + def _join_score_queue(self) -> None: + """Wait until the score queue is drained or no consumer is left to drain it. + + Queue.join() would block forever when the consumers stop first, e.g. + when shutdown() runs concurrently and a wake signal is never taken. + """ + score_queue = self._score_queue + + with score_queue.all_tasks_done: + while score_queue.unfinished_tasks: + if not any(c.is_alive() for c in self._score_consumers): + langfuse_logger.debug( + "Score consumers stopped with %s queue item(s) unprocessed", + score_queue.unfinished_tasks, + ) + + return + + score_queue.all_tasks_done.wait(timeout=0.1) + def shutdown(self) -> None: self._shutdown = True diff --git a/langfuse/_task_manager/score_consumer.py b/langfuse/_task_manager/score_consumer.py index cdeb81404..21d9f5052 100644 --- a/langfuse/_task_manager/score_consumer.py +++ b/langfuse/_task_manager/score_consumer.py @@ -8,6 +8,10 @@ import backoff import httpx +from langfuse._client.environment_variables import ( + LANGFUSE_FLUSH_AT, + LANGFUSE_FLUSH_INTERVAL, +) from langfuse._utils.parse_error import handle_fern_exception from langfuse._utils.serializer import EventSerializer from langfuse.api import LangfuseAPI @@ -32,7 +36,7 @@ _SCORES_BODY_LIMIT_BYTES, ) MAX_BATCH_SIZE_BYTES = min(_configured_batch_size_bytes, _SCORES_BODY_LIMIT_BYTES) -_SHUTDOWN_SENTINEL = object() +_WAKE_SIGNAL = object() def _is_retryable(error: Exception) -> bool: @@ -71,6 +75,15 @@ def __init__( self._score_queue = score_queue self._identifier = identifier self._api_client = api_client + + env_flush_at = os.environ.get(LANGFUSE_FLUSH_AT) + if flush_at is None and env_flush_at is not None: + flush_at = int(env_flush_at) + + env_flush_interval = os.environ.get(LANGFUSE_FLUSH_INTERVAL) + if flush_interval is None and env_flush_interval is not None: + flush_interval = float(env_flush_interval) + self._flush_at = flush_at or 15 self._flush_interval = flush_interval or 1 self._max_retries = max_retries or 3 @@ -99,7 +112,7 @@ def _next(self) -> List[Dict[str, Any]]: block=True, timeout=self._flush_interval - elapsed ) - if score is _SHUTDOWN_SENTINEL: + if score is _WAKE_SIGNAL: self._score_queue.task_done() break @@ -199,16 +212,17 @@ def upload(self) -> None: for _ in batch: self._score_queue.task_done() - def pause(self) -> None: - """Pause the consumer.""" - self.running = False + def wake(self) -> None: try: - self._score_queue.put(_SHUTDOWN_SENTINEL, block=False) + self._score_queue.put(_WAKE_SIGNAL, block=False) except Full: - # If the queue is full, the consumer will wake up naturally while - # draining items, so a dedicated shutdown signal is not required. pass + def pause(self) -> None: + """Pause the consumer.""" + self.running = False + self.wake() + def _upload_batch(self, batch: List[Dict[str, Any]]) -> None: logger.debug("API: Sending batch of %s scores to Langfuse API", len(batch)) diff --git a/tests/unit/test_experiment.py b/tests/unit/test_experiment.py index fea3587b2..deca750c0 100644 --- a/tests/unit/test_experiment.py +++ b/tests/unit/test_experiment.py @@ -1,7 +1,9 @@ """Tests for ``langfuse.experiment`` — ``RunnerContext`` and ``RegressionError``.""" +import asyncio import inspect import json +import time import typing from datetime import datetime, timezone from typing import get_type_hints @@ -779,3 +781,39 @@ def test_tracing_disabled_skips_project_id_lookup(self, monkeypatch): get_project_id.assert_not_called() assert result.experiment_url is None + + +class TestExperimentEventLoop: + @pytest.mark.asyncio + async def test_final_flush_does_not_block_event_loop( + self, langfuse_memory_client, monkeypatch + ): + monkeypatch.setattr(langfuse_memory_client, "_get_project_id", lambda: "p") + monkeypatch.setattr(langfuse_memory_client, "flush", lambda: time.sleep(0.5)) + stop = asyncio.Event() + gaps: list = [] + + async def heartbeat(): + last = time.monotonic() + while not stop.is_set(): + await asyncio.sleep(0.01) + now = time.monotonic() + gaps.append(now - last) + last = now + + beat = asyncio.create_task(heartbeat()) + await langfuse_memory_client._run_experiment_async( + name="exp", + run_name="run", + description=None, + data=[{"input": "a"}], + task=lambda *, item, **kwargs: item["input"], + evaluators=[], + composite_evaluator=None, + run_evaluators=[], + max_concurrency=1, + ) + stop.set() + await beat + + assert max(gaps) < 0.25 diff --git a/tests/unit/test_score_consumer.py b/tests/unit/test_score_consumer.py index a2818d495..cc6611d29 100644 --- a/tests/unit/test_score_consumer.py +++ b/tests/unit/test_score_consumer.py @@ -243,6 +243,72 @@ def test_does_not_retry_client_errors(status, caplog): assert caplog.records +def test_reads_flush_at_from_env(monkeypatch): + monkeypatch.setenv("LANGFUSE_FLUSH_AT", "2") + recorder = Recorder(_responder(ACCEPTED)) + queue: Queue = Queue() + for name in ("a", "b", "c"): + queue.put(_score(name)) + + _drain(_consumer(recorder, queue), queue) + + assert [[s["name"] for s in body] for body in recorder.bodies] == [ + ["a", "b"], + ["c"], + ] + + +def test_reads_flush_interval_from_env(monkeypatch): + monkeypatch.setenv("LANGFUSE_FLUSH_INTERVAL", "0.05") + recorder = Recorder(_responder(ACCEPTED)) + consumer = ScoreConsumer( + score_queue=Queue(), + identifier=0, + api_client=LangfuseAPI( + base_url="http://langfuse.test", + username="pk", + password="sk", + httpx_client=httpx.Client(transport=httpx.MockTransport(recorder)), + ), + ) + + started = time.monotonic() + consumer.upload() + + # The default window is 1 second. + assert time.monotonic() - started < 0.5 + + +def test_constructor_flush_settings_win_over_env(monkeypatch): + monkeypatch.setenv("LANGFUSE_FLUSH_AT", "1") + recorder = Recorder(_responder(ACCEPTED)) + queue: Queue = Queue() + for name in ("a", "b", "c"): + queue.put(_score(name)) + + _drain(_consumer(recorder, queue, flush_at=3), queue) + + assert [[s["name"] for s in body] for body in recorder.bodies] == [["a", "b", "c"]] + + +def test_wake_sends_pending_batch_before_flush_interval(): + recorder = Recorder(_responder(ACCEPTED)) + queue: Queue = Queue() + consumer = _consumer(recorder, queue, flush_interval=30) + consumer.start() + queue.put(_score("a")) + + consumer.wake() + finished = threading.Event() + threading.Thread(target=lambda: (queue.join(), finished.set()), daemon=True).start() + + assert finished.wait(timeout=5) + assert [[s["name"] for s in body] for body in recorder.bodies] == [["a"]] + assert consumer.is_alive() + consumer.pause() + consumer.join(timeout=5) + + def test_pause_sends_carried_over_score(monkeypatch): first_request_started = threading.Event() release_first_request = threading.Event() @@ -351,6 +417,61 @@ def test_create_score_posts_to_scores_endpoint(client_factory): assert sent[1]["timestamp"].endswith("Z") +def test_flush_sends_pending_score_without_waiting_for_flush_interval( + client_factory, +): + recorder = Recorder(_responder(ACCEPTED)) + client = client_factory(recorder, flush_interval=5) + + client.create_score(name="a", value=1.0, trace_id="a" * 32) + started = time.monotonic() + client.flush() + elapsed = time.monotonic() - started + + assert [s["name"] for body in recorder.bodies for s in body] == ["a"] + assert elapsed < 1 + + +def test_flush_after_shutdown_returns(client_factory): + client = client_factory(Recorder(_responder(ACCEPTED))) + client.shutdown() + + flush = threading.Thread(target=client.flush, daemon=True) + flush.start() + flush.join(timeout=2) + + assert not flush.is_alive() + + +def test_flush_returns_when_consumer_stops_after_liveness_check(client_factory): + client = client_factory(Recorder(_responder(ACCEPTED))) + consumer = client._resources._score_consumers[0] + wake = consumer.wake + + def stop_then_wake() -> None: + # Shutdown stops the consumer between flush()'s liveness check and its + # wake signal, so nothing acknowledges the signal. + consumer.running = False + wake() + consumer.join(timeout=5) + wake() + + consumer.wake = stop_then_wake # type: ignore[method-assign] + + flush = threading.Thread(target=client.flush, daemon=True) + flush.start() + flush.join(timeout=2) + flush_hung = flush.is_alive() + # Acknowledge the orphaned signal so a hung flush() cannot block teardown. + score_queue = client._resources._score_queue + while not score_queue.empty(): + score_queue.get_nowait() + score_queue.task_done() + + assert not consumer.is_alive() + assert not flush_hung + + def test_create_score_payload_matches_request_model_serialization(client_factory): recorder = Recorder(_responder(ACCEPTED)) client = client_factory(recorder, environment="staging")