From dd3881f85c7b1e410de6b86d5b10b8f7de0424ee Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Mon, 26 Oct 2026 19:49:51 -0400 Subject: [PATCH] feat(logger): batch log rows with a short flush delay MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit > 🤖 **This text was AI-generated.** The HTTP logger sent rows as soon as they were queued, so nearby events often used separate requests and span start and end rows could not be merged. Add a short batching window for long-lived Python processes, while keeping detected serverless environments immediate by default. ```text span.log() / span.end() flush() / atexit │ │ ▼ │ ┌─────────────────────────┐ │ │ queue [r1][r2] … [rN] │ │ └────────────┬────────────┘ │ ▼ │ ╔═════════════════════╗ yes │ ║ N ≥ batch threshold ║ ───────────────┐ │ ╚══════════╤══════════╝ │ │ │ no ▼ ▼ ▼ ┌─────────────────────────────┐ ◷ wait until oldest row ──────▶ │ drain ─▶ merge ─▶ POST │ + BRAINTRUST_FLUSH_DELAY_MS │ same-id rows /logs3 │ └──────────────┬──────────────┘ 1000ms long-lived Python │ done 0ms detected serverless ▼ rows queued meanwhile get their own deadline ──▶ ◷ ``` ```text before r1 r2 r3 r4 after r1 r2 r3 r4 │ │ │ │ │ │ │ │ ▼ ▼ ▼ ▼ └─────┴──┬──┴─────┘ [P] [P] [P] [P] ◷ ≤ 1s ▼ 4 requests, start + end [P] rows sent separately 1 request, start + end rows merged into one ``` An explicit flush and normal process exit still send immediately. BRAINTRUST_FLUSH_DELAY_MS adjusts the delay; BRAINTRUST_DEFAULT_BATCH_SIZE controls the batch threshold and maximum merged rows per request. Call flush() before terminating with a signal or os._exit(), which bypass Python's exit handlers. Refs: https://github.com/braintrustdata/braintrust-sdk-javascript/pull/2653 --- py/src/braintrust/env.py | 1 + py/src/braintrust/logger.py | 29 +++++- py/src/braintrust/queue.py | 41 ++++++-- py/src/braintrust/test_logger.py | 155 +++++++++++++++++++++++++++++++ py/src/braintrust/test_queue.py | 42 ++++++++- 5 files changed, 257 insertions(+), 11 deletions(-) diff --git a/py/src/braintrust/env.py b/py/src/braintrust/env.py index 3214b290c..98babd6db 100644 --- a/py/src/braintrust/env.py +++ b/py/src/braintrust/env.py @@ -218,6 +218,7 @@ class BraintrustEnv: PROXY_URL = EnvVar("BRAINTRUST_PROXY_URL", EnvParser.STRING) HTTP_TIMEOUT = EnvVar("BRAINTRUST_HTTP_TIMEOUT", EnvParser.FLOAT) SYNC_FLUSH = EnvVar("BRAINTRUST_SYNC_FLUSH", EnvParser.BOOL) + FLUSH_DELAY_MS = EnvVar("BRAINTRUST_FLUSH_DELAY_MS", EnvParser.FLOAT) MAX_REQUEST_SIZE = EnvVar("BRAINTRUST_MAX_REQUEST_SIZE", EnvParser.INT) DEFAULT_BATCH_SIZE = EnvVar("BRAINTRUST_DEFAULT_BATCH_SIZE", EnvParser.INT) NUM_RETRIES = EnvVar("BRAINTRUST_NUM_RETRIES", EnvParser.INT) diff --git a/py/src/braintrust/logger.py b/py/src/braintrust/logger.py index 4c6dae50b..41dbb0277 100644 --- a/py/src/braintrust/logger.py +++ b/py/src/braintrust/logger.py @@ -1217,6 +1217,27 @@ def pop(self) -> list[dict[str, Any]]: BACKGROUND_LOGGER_BASE_SLEEP_TIME_S = 1.0 +def _default_flush_delay_ms() -> float: + # Serverless and scale-to-zero platforms can freeze background threads + # between invocations, before a delayed flush has a chance to run. + serverless = ( + any( + os.environ.get(name) + for name in ( + "VERCEL", + "NETLIFY", + "AWS_LAMBDA_FUNCTION_NAME", + "K_SERVICE", + "FUNCTION_TARGET", + "FUNCTIONS_WORKER_RUNTIME", + ) + ) + or os.environ.get("AWS_EXECUTION_ENV", "").startswith("AWS_Lambda_") + or os.environ.get("GAE_ENV") == "standard" + ) + return 0.0 if serverless else 1000.0 + + # We should only have one instance of this object in # 'BraintrustState._global_bg_logger'. Be careful about spawning multiple # instances of this class, because concurrent _BackgroundLoggers will not log to @@ -1232,6 +1253,9 @@ def __init__(self, api_conn: LazyValue[HTTPConnection]): self._max_request_size_lock = threading.Lock() self.sync_flush = BraintrustEnv.SYNC_FLUSH.get(False) + default_flush_delay_ms = _default_flush_delay_ms() + flush_delay_ms = BraintrustEnv.FLUSH_DELAY_MS.get(default_flush_delay_ms) + self.flush_delay_ms = flush_delay_ms if flush_delay_ms >= 0 else default_flush_delay_ms self._max_request_size_override = BraintrustEnv.MAX_REQUEST_SIZE.get(None) self.default_batch_size = BraintrustEnv.DEFAULT_BATCH_SIZE.get(100) self.num_tries = BraintrustEnv.NUM_RETRIES.get(2) + 1 @@ -1294,8 +1318,9 @@ def _finalize(self): def _publisher(self): while True: - # Wait for some data on the queue before trying to flush. - self.queue.wait_for_items() + # Batch until the oldest queued row's deadline, or until a full + # batch is available. Explicit and atexit flushes bypass this wait. + self.queue.wait_for_batch(self.default_batch_size, self.flush_delay_ms / 1000) while self.sync_flush: time.sleep(0.1) diff --git a/py/src/braintrust/queue.py b/py/src/braintrust/queue.py index cdb0ba364..582714df3 100644 --- a/py/src/braintrust/queue.py +++ b/py/src/braintrust/queue.py @@ -1,5 +1,6 @@ import threading from collections import deque +from time import monotonic from typing import TypeVar from .util import eprint @@ -31,7 +32,9 @@ def __init__(self, maxsize: int = 0): self.maxsize = maxsize self._mutex = threading.Lock() self._queue: deque[T] = deque(maxlen=maxsize) - self._has_items_event = threading.Event() + self._condition = threading.Condition(self._mutex) + self._batch_min_size: int | None = None + self._pending_since: float | None = None self._total_dropped = 0 self._enforce_size_limit = False @@ -72,9 +75,11 @@ def put(self, item: T) -> list[T]: self._total_dropped += 1 self._queue.append(item) - # Signal that items are available if queue was not empty before or item was added if was_empty: - self._has_items_event.set() + self._pending_since = monotonic() + self._condition.notify_all() + elif self._batch_min_size is not None and len(self._queue) >= self._batch_min_size: + self._condition.notify_all() return dropped @@ -92,9 +97,11 @@ def drain_all(self) -> list[T]: old_queue = self._queue self._queue = deque(maxlen=self.maxsize) + self._pending_since = None - # Clear the event since queue is now empty - self._has_items_event.clear() + # An explicit flush may drain the queue while the publisher waits. + if self._batch_min_size is not None: + self._condition.notify_all() return list(old_queue) if old_queue else [] @@ -118,4 +125,26 @@ def wait_for_items(self, timeout: float | None = None) -> bool: Returns: True if items became available, False if timeout occurred. """ - return self._has_items_event.wait(timeout=timeout) + with self._condition: + return self._condition.wait_for(lambda: bool(self._queue), timeout=timeout) + + def wait_for_batch(self, min_size: int, max_delay: float) -> None: + """Wait for a full batch or the oldest pending row's deadline. + + Used by the single background publisher. New arrivals do not extend + the deadline; draining the queue starts a new window for the next row. + Like wait_for_items(), the queue may be empty by the time this returns. + """ + with self._condition: + self._batch_min_size = min_size + try: + while True: + if self._pending_since is None: + self._condition.wait() + continue + remaining = self._pending_since + max_delay - monotonic() + if len(self._queue) >= min_size or remaining <= 0: + return + self._condition.wait(timeout=remaining) + finally: + self._batch_min_size = None diff --git a/py/src/braintrust/test_logger.py b/py/src/braintrust/test_logger.py index b2e7ea2a1..9bb585fac 100644 --- a/py/src/braintrust/test_logger.py +++ b/py/src/braintrust/test_logger.py @@ -261,6 +261,161 @@ def test_http_background_logger_atexit_flush(monkeypatch, disable_atexit_flush, assert mock_register.called is registers +_SERVERLESS_ENV_VARS = ( + "VERCEL", + "NETLIFY", + "AWS_EXECUTION_ENV", + "AWS_LAMBDA_FUNCTION_NAME", + "K_SERVICE", + "FUNCTION_TARGET", + "GAE_ENV", + "FUNCTIONS_WORKER_RUNTIME", +) + + +@pytest.fixture +def make_batching_logger(monkeypatch): + for name in (*_SERVERLESS_ENV_VARS, "BRAINTRUST_FLUSH_DELAY_MS", "BRAINTRUST_DEFAULT_BATCH_SIZE"): + monkeypatch.delenv(name, raising=False) + monkeypatch.setenv("BRAINTRUST_DISABLE_ATEXIT_FLUSH", "true") + monkeypatch.setenv("BRAINTRUST_SYNC_FLUSH", "false") + loggers = [] + + def create(): + conn = MagicMock() + conn.post.return_value.ok = True + bg = logger._HTTPBackgroundLogger(LazyValue(lambda: conn, use_mutex=False)) + bg._max_request_size_result = {"max_request_size": 6 * 1024 * 1024, "can_use_overflow": False} + loggers.append(bg) + return bg, conn + + yield create + for bg in loggers: + bg.sync_flush = False + bg.flush() + + +@pytest.mark.parametrize( + ("name", "value", "expected"), + [ + (None, None, 1000), + ("VERCEL", "1", 0), + ("NETLIFY", "true", 0), + ("AWS_EXECUTION_ENV", "AWS_Lambda_python3.14", 0), + ("AWS_LAMBDA_FUNCTION_NAME", "fn", 0), + ("K_SERVICE", "service", 0), + ("FUNCTION_TARGET", "handler", 0), + ("GAE_ENV", "standard", 0), + ("FUNCTIONS_WORKER_RUNTIME", "python", 0), + ("AWS_EXECUTION_ENV", "AWS_ECS_FARGATE", 1000), + ("GAE_ENV", "flex", 1000), + ], +) +def test_http_background_logger_flush_delay_default(make_batching_logger, monkeypatch, name, value, expected): + if name: + monkeypatch.setenv(name, value) + bg, _ = make_batching_logger() + assert bg.flush_delay_ms == expected + + +@pytest.mark.parametrize( + ("serverless", "value", "expected"), + [(False, "0", 0), (False, "0.5", 0.5), (False, "-1", 1000), (False, "bad", 1000), (True, "250", 250)], +) +def test_http_background_logger_flush_delay_override(make_batching_logger, monkeypatch, serverless, value, expected): + if serverless: + monkeypatch.setenv("VERCEL", "1") + monkeypatch.setenv("BRAINTRUST_FLUSH_DELAY_MS", value) + bg, _ = make_batching_logger() + assert bg.flush_delay_ms == expected + + +def test_http_background_logger_batches_and_merges_nearby_rows(make_batching_logger): + bg, conn = make_batching_logger() + bg.flush_delay_ms = 150 + sent = threading.Event() + conn.post.side_effect = lambda *args, **kwargs: sent.set() or MagicMock(ok=True) + + bg.log(LazyValue(lambda: {"id": "span", "dataset_id": "dataset", "metrics": {"start": 1}}, use_mutex=False)) + assert not sent.wait(0.03) + bg.log( + LazyValue( + lambda: {"id": "span", "dataset_id": "dataset", "_is_merge": True, "metrics": {"end": 2}}, use_mutex=False + ) + ) + assert sent.wait(2) + assert conn.post.call_count == 1 + rows = json.loads(conn.post.call_args.kwargs["data"])["rows"] + assert len(rows) == 1 + assert rows[0]["metrics"] == {"start": 1, "end": 2} + + +def test_http_background_logger_full_batch_skips_delay(make_batching_logger): + bg, conn = make_batching_logger() + bg.flush_delay_ms = 60_000 + bg.default_batch_size = 3 + sent = threading.Event() + conn.post.side_effect = lambda *args, **kwargs: sent.set() or MagicMock(ok=True) + for i in range(2): + bg.log(LazyValue(lambda i=i: {"id": str(i), "dataset_id": "dataset"}, use_mutex=False)) + assert not sent.wait(0.03) + bg.log(LazyValue(lambda: {"id": "2", "dataset_id": "dataset"}, use_mutex=False)) + assert sent.wait(2) + assert len(json.loads(conn.post.call_args.kwargs["data"])["rows"]) == 3 + + +def test_http_background_logger_exit_flush_skips_delay(make_batching_logger): + bg, conn = make_batching_logger() + bg.flush_delay_ms = 60_000 + bg.log(LazyValue(lambda: {"id": "span", "dataset_id": "dataset"}, use_mutex=False)) + bg._finalize() + conn.post.assert_called_once() + assert bg.queue.size() == 0 + + +def test_http_background_logger_rows_during_flush_get_their_own_delay(make_batching_logger): + bg, conn = make_batching_logger() + bg.flush_delay_ms = 150 + uploading = threading.Event() + release_upload = threading.Event() + second_sent = threading.Event() + + def post(*args, **kwargs): + rows = json.loads(kwargs["data"])["rows"] + if rows[0]["id"] == "first": + uploading.set() + assert release_upload.wait(2) + else: + second_sent.set() + return MagicMock(ok=True) + + conn.post.side_effect = post + try: + bg.log(LazyValue(lambda: {"id": "first", "dataset_id": "dataset"}, use_mutex=False)) + assert uploading.wait(2) + bg.log(LazyValue(lambda: {"id": "second", "dataset_id": "dataset"}, use_mutex=False)) + finally: + release_upload.set() + assert not second_sent.wait(0.03) + assert second_sent.wait(2) + assert conn.post.call_count == 2 + + +@pytest.mark.parametrize("sync_flush", [False, True]) +def test_http_background_logger_zero_delay_respects_sync_flush(make_batching_logger, sync_flush): + bg, conn = make_batching_logger() + bg.sync_flush = sync_flush + bg.flush_delay_ms = 0 + sent = threading.Event() + conn.post.side_effect = lambda *args, **kwargs: sent.set() or MagicMock(ok=True) + bg.log(LazyValue(lambda: {"id": "span", "dataset_id": "dataset"}, use_mutex=False)) + if sync_flush: + assert not sent.wait(0.03) + bg.flush() + assert sent.wait(2) + conn.post.assert_called_once() + + class TestHTTPBackgroundLoggerLogs3(TestCase): def test_flush_uses_request_byte_budget(self) -> None: from braintrust.logger import _HTTPBackgroundLogger, construct_logs3_data, stringify_with_overflow_meta diff --git a/py/src/braintrust/test_queue.py b/py/src/braintrust/test_queue.py index 3dde12df3..d78b50ce0 100644 --- a/py/src/braintrust/test_queue.py +++ b/py/src/braintrust/test_queue.py @@ -1,6 +1,7 @@ import asyncio import threading import time +from unittest.mock import patch import pytest from braintrust.queue import DEFAULT_QUEUE_SIZE, LogQueue @@ -32,11 +33,9 @@ def test_log_queue_basic_operations(): @pytest.mark.parametrize("enforce", [False, True]) def test_log_queue_signals_once_until_drained(enforce): - from unittest.mock import patch - queue = LogQueue(maxsize=2) queue.enforce_queue_size_limit(enforce) - with patch.object(queue._has_items_event, "set", wraps=queue._has_items_event.set) as signal: + with patch.object(queue._condition, "notify_all", wraps=queue._condition.notify_all) as signal: for item in range(5): queue.put(item) assert signal.call_count == 1 @@ -48,6 +47,43 @@ def test_log_queue_signals_once_until_drained(enforce): assert queue.wait_for_items(timeout=0) +def test_log_queue_batch_deadline_is_not_extended_by_arrivals_or_drops(monkeypatch): + now = 0.0 + monkeypatch.setattr("braintrust.queue.monotonic", lambda: now) + queue = LogQueue(maxsize=2) + queue.enforce_queue_size_limit(True) + queue.put("first") + arrivals = iter([(0.4, "second"), (0.8, "third"), (1.0, None)]) + + def wake(timeout=None): + nonlocal now + now, item = next(arrivals) + # Simulate a producer running while Condition.wait releases its lock. + queue._mutex.release() + try: + if item is not None: + queue.put(item) + finally: + queue._mutex.acquire() + + with patch.object(queue._condition, "wait", side_effect=wake) as wait: + queue.wait_for_batch(min_size=10, max_delay=1) + + assert [c.kwargs["timeout"] for c in wait.call_args_list] == pytest.approx([1, 0.6, 0.2]) + assert queue.drain_all() == ["second", "third"] + + +def test_log_queue_batch_already_overdue_does_not_wait(monkeypatch): + now = 0.0 + monkeypatch.setattr("braintrust.queue.monotonic", lambda: now) + queue = LogQueue(maxsize=10) + queue.put("queued during upload") + now = 2 + with patch.object(queue._condition, "wait") as wait: + queue.wait_for_batch(min_size=10, max_delay=1) + wait.assert_not_called() + + def test_log_queue_drop_behavior(): """Test queue drops oldest items when full, including single and multiple drops""" # Test basic drop behavior with size 2