From 0001b69a1317f44717147feae5de4523b13fba06 Mon Sep 17 00:00:00 2001 From: Ankur Goyal Date: Sun, 4 Oct 2026 23:15:56 -0700 Subject: [PATCH 1/3] optimize ingest pathway --- py/benchmarks/benches/bench_log_queue.py | 42 ++++++++ py/src/braintrust/logger.py | 78 +++++++++------ py/src/braintrust/merge_row_batch.py | 10 +- py/src/braintrust/queue.py | 3 +- py/src/braintrust/test_logger.py | 111 ++++++++++++++++++++++ py/src/braintrust/test_merge_row_batch.py | 3 + py/src/braintrust/test_queue.py | 18 ++++ 7 files changed, 226 insertions(+), 39 deletions(-) create mode 100644 py/benchmarks/benches/bench_log_queue.py diff --git a/py/benchmarks/benches/bench_log_queue.py b/py/benchmarks/benches/bench_log_queue.py new file mode 100644 index 000000000..76ca55951 --- /dev/null +++ b/py/benchmarks/benches/bench_log_queue.py @@ -0,0 +1,42 @@ +import pathlib +import sys + +import pyperf + + +if __package__ in (None, ""): + sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[2])) + +from braintrust.logger import _enrich_attachments, _extract_attachments +from braintrust.queue import LogQueue + +from benchmarks._utils import disable_pyperf_psutil +from benchmarks.fixtures import make_large_payload, make_medium_payload, make_small_payload + + +_queue = LogQueue(maxsize=1000) + + +def _bench_queue() -> None: + _queue.drain_all() + for i in range(1000): + _queue.put(i) + + +def main(runner: pyperf.Runner | None = None) -> None: + if runner is None: + disable_pyperf_psutil() + runner = pyperf.Runner() + + runner.bench_func("log_queue.put[1000]", _bench_queue) + for name, payload in ( + ("small", make_small_payload()), + ("medium", make_medium_payload()), + ("large", make_large_payload()), + ): + runner.bench_func(f"logger.extract_attachments[{name}]", _extract_attachments, payload, []) + runner.bench_func(f"logger.enrich_attachments[{name}]", _enrich_attachments, payload) + + +if __name__ == "__main__": + main() diff --git a/py/src/braintrust/logger.py b/py/src/braintrust/logger.py index 21155a2dd..eacbffaea 100644 --- a/py/src/braintrust/logger.py +++ b/py/src/braintrust/logger.py @@ -993,7 +993,7 @@ def stringify_with_overflow_meta(item: dict[str, Any]) -> LogItemWithMeta: def utf8_byte_length(value: str) -> int: - return len(value.encode("utf-8")) + return len(value) if value.isascii() else len(value.encode("utf-8")) _LazyLogRecord = LazyValue[dict[str, Any]] | LazyValue[dict[str, Any] | None] @@ -1216,8 +1216,10 @@ def flush(self, batch_size: int | None = None): batches = batch_items( items=all_items_with_meta, batch_max_num_items=batch_size, - batch_max_num_bytes=max_request_size_result["max_request_size"] // 2, - get_byte_size=lambda item: len(item.str_value), + batch_max_num_bytes=max( + 0, max_request_size_result["max_request_size"] - utf8_byte_length(construct_logs3_data([])) + 1 + ), + get_byte_size=lambda item: item.overflow_meta.byte_size + 1, ) post_promises = [] @@ -1338,7 +1340,7 @@ def _request_logs3_overflow_upload( "key": key, } - def _upload_logs3_overflow_payload(self, upload: dict[str, Any], payload: str) -> None: + def _upload_logs3_overflow_payload(self, upload: dict[str, Any], payload: bytes) -> None: obj_conn = HTTPConnection(base_url="", adapter=_http_adapter) method = upload["method"] if method == "POST": @@ -1351,20 +1353,21 @@ def _upload_logs3_overflow_payload(self, upload: dict[str, Any], payload: str) - upload["signed_url"], headers=headers, data=fields, - files={"file": ("logs3.json", payload.encode("utf-8"), content_type)}, + files={"file": ("logs3.json", payload, content_type)}, ) else: obj_response = obj_conn.put( upload["signed_url"], headers=upload["headers"], - data=payload.encode("utf-8"), + data=payload, ) obj_response.raise_for_status() def _submit_logs_request(self, items: Sequence[LogItemWithMeta], max_request_size_result: dict[str, Any]): conn = self.api_conn.get() dataStr = construct_logs3_data(items) - payload_bytes = utf8_byte_length(dataStr) + data_bytes = dataStr.encode("utf-8") + payload_bytes = len(data_bytes) max_request_size = max_request_size_result["max_request_size"] can_use_overflow = max_request_size_result["can_use_overflow"] use_overflow = can_use_overflow and payload_bytes > max_request_size @@ -1392,14 +1395,14 @@ def _submit_logs_request(self, items: Sequence[LogItemWithMeta], max_request_siz if overflow_rows: if overflow_upload is None: current_upload = self._request_logs3_overflow_upload(conn, payload_bytes, overflow_rows) - self._upload_logs3_overflow_payload(current_upload, dataStr) + self._upload_logs3_overflow_payload(current_upload, data_bytes) overflow_upload = current_upload resp = conn.post( "/logs3", json=construct_logs3_overflow_request(overflow_upload["key"], payload_bytes), ) else: - resp = conn.post("/logs3", data=dataStr.encode("utf-8")) + resp = conn.post("/logs3", data=data_bytes) except Exception as e: error = e if error is None and resp is not None and resp.ok: @@ -3229,29 +3232,36 @@ def _extract_attachments(event: dict[str, Any], attachments: list["BaseAttachmen :param attachments: Flat array of extracted attachments (output parameter). """ - def _helper(v: Any) -> Any: - # Base case: Attachment or ExternalAttachment. - if isinstance(v, BaseAttachment): - attachments.append(v) - return v.reference # Attachment cannot be nested. + def _helper(v: Any, items: Any) -> None: + for k, v2 in items: + tv = type(v2) + if tv is dict: + _helper(v2, v2.items()) + continue + if tv is list: + _helper(v2, ((i, v2[i]) for i in range(len(v2)))) + continue + if tv is str or tv is int or tv is float or v2 is None or v2 is True or v2 is False: + continue - # Recursive case: object. - if isinstance(v, dict): - for k, v2 in v.items(): - v[k] = _helper(v2) - return v + # Base case: Attachment or ExternalAttachment. + if isinstance(v2, BaseAttachment): + attachments.append(v2) + v[k] = v2.reference # Attachment cannot be nested. - # Recursive case: array. - if isinstance(v, list): - for i in range(len(v)): - v[i] = _helper(v[i]) - return v + # Recursive case: object. + elif isinstance(v2, dict): + _helper(v2, v2.items()) - # Base case: non object. - return v # Nothing to explore recursively. + # Recursive case: array. + elif isinstance(v2, list): + _helper(v2, ((i, v2[i]) for i in range(len(v2)))) - for k, v in event.items(): - event[k] = _helper(v) + # Base case: non object. + else: + continue # Nothing to explore recursively. + + _helper(event, event.items()) def _enrich_attachments(event: TMutableMapping) -> TMutableMapping: @@ -3261,6 +3271,8 @@ def _enrich_attachments(event: TMutableMapping) -> TMutableMapping: :returns: The same event instance as the input. """ + traversable = (dict, list) + def _helper(v: Any) -> Any: if isinstance(v, dict): # Base case: AttachmentReference. @@ -3269,20 +3281,24 @@ def _helper(v: Any) -> Any: else: # Recursive case: object. for k, v2 in v.items(): - v[k] = _helper(v2) + if isinstance(v2, traversable): + v[k] = _helper(v2) return v # Recursive case: array. if isinstance(v, list): for i in range(len(v)): - v[i] = _helper(v[i]) + v2 = v[i] + if isinstance(v2, traversable): + v[i] = _helper(v2) return v # Base case: non object. return v # Nothing to explore recursively. for k, v in event.items(): - event[k] = _helper(v) + if isinstance(v, traversable): + event[k] = _helper(v) return event diff --git a/py/src/braintrust/merge_row_batch.py b/py/src/braintrust/merge_row_batch.py index 066744cec..bd3a1ead8 100644 --- a/py/src/braintrust/merge_row_batch.py +++ b/py/src/braintrust/merge_row_batch.py @@ -159,11 +159,6 @@ def get_byte_size(item: T) -> int: batch: list[T] = [] batch_len = 0 - def add_to_batch(item: T) -> None: - nonlocal batch_len - batch.append(item) - batch_len += get_byte_size(item) - def flush_batch() -> None: nonlocal batch, batch_len output.append(batch) @@ -173,11 +168,12 @@ def flush_batch() -> None: for item in items: item_size = get_byte_size(item) if len(batch) > 0 and not ( - (batch_max_num_bytes is None or item_size + batch_len < batch_max_num_bytes) + (batch_max_num_bytes is None or item_size + batch_len <= batch_max_num_bytes) and (batch_max_num_items is None or len(batch) < batch_max_num_items) ): flush_batch() - add_to_batch(item) + batch.append(item) + batch_len += item_size if len(batch) > 0: flush_batch() diff --git a/py/src/braintrust/queue.py b/py/src/braintrust/queue.py index ff6fc6cfc..cdb0ba364 100644 --- a/py/src/braintrust/queue.py +++ b/py/src/braintrust/queue.py @@ -59,6 +59,7 @@ def put(self, item: T) -> list[T]: """ with self._mutex: dropped = [] + was_empty = not self._queue if not self._enforce_size_limit: # For queues with enforcement disabled, deque auto-drops silently @@ -72,7 +73,7 @@ def put(self, item: T) -> list[T]: self._queue.append(item) # Signal that items are available if queue was not empty before or item was added - if len(self._queue) > 0: + if was_empty: self._has_items_event.set() return dropped diff --git a/py/src/braintrust/test_logger.py b/py/src/braintrust/test_logger.py index 597ab4f77..d927c38d2 100644 --- a/py/src/braintrust/test_logger.py +++ b/py/src/braintrust/test_logger.py @@ -261,6 +261,47 @@ def test_http_background_logger_atexit_flush(monkeypatch, disable_atexit_flush, class TestHTTPBackgroundLoggerLogs3(TestCase): + def test_flush_uses_request_byte_budget(self) -> None: + from braintrust.logger import _HTTPBackgroundLogger, construct_logs3_data, stringify_with_overflow_meta + + rows = [{"id": str(i), "dataset_id": "dataset", "input": "δΈ­ζ–‡πŸ™‚" * 20} for i in range(6)] + limit = len(construct_logs3_data([stringify_with_overflow_meta(row) for row in rows[:3]]).encode()) + conn = MagicMock() + conn.post.return_value.ok = True + + with patch("atexit.register"): + bg = _HTTPBackgroundLogger(LazyValue(lambda: conn, use_mutex=False)) + bg._max_request_size_result = {"max_request_size": limit, "can_use_overflow": True} + for row in rows: + bg.queue.put(LazyValue(lambda row=row: row, use_mutex=False)) + + bg.flush(batch_size=100) + + self.assertEqual(conn.post.call_count, 2) + posted_rows = [] + for call in conn.post.call_args_list: + self.assertEqual(call.args, ("/logs3",)) + self.assertLessEqual(len(call.kwargs["data"]), limit) + posted_rows.extend(json.loads(call.kwargs["data"])["rows"]) + self.assertCountEqual(posted_rows, rows) + + def test_flush_keeps_oversized_row_separate(self) -> None: + from braintrust.logger import _HTTPBackgroundLogger + + rows = [{"id": str(i), "dataset_id": "dataset", "input": "x" * size} for i, size in enumerate((10, 2000, 10))] + conn = MagicMock() + conn.post.return_value.ok = True + with patch("atexit.register"): + bg = _HTTPBackgroundLogger(LazyValue(lambda: conn, use_mutex=False)) + bg._max_request_size_result = {"max_request_size": 500, "can_use_overflow": False} + for row in rows: + bg.queue.put(LazyValue(lambda row=row: row, use_mutex=False)) + + bg.flush(batch_size=100) + + posted = [json.loads(call.kwargs["data"])["rows"] for call in conn.post.call_args_list] + self.assertCountEqual(posted, [[row] for row in rows]) + def test_submit_logs_request_413_skips_retries(self) -> None: """Any 413 while publishing ``/logs3`` cannot succeed on retry with the same payload. @@ -313,6 +354,40 @@ def test_submit_logs_request_413_skips_retries(self) -> None: mock_write_payload.call_args.kwargs["payload_dir"], bg.failed_publish_payloads_dir ) + def test_submit_logs_request_preserves_overflow_payload(self) -> None: + from braintrust.logger import _HTTPBackgroundLogger, construct_logs3_data, stringify_with_overflow_meta + + item = stringify_with_overflow_meta({"id": "row", "dataset_id": "dataset", "input": "δΈ­ζ–‡πŸ™‚" * 100}) + expected = construct_logs3_data([item]).encode("utf-8") + for method in ("PUT", "POST"): + with self.subTest(method=method): + conn, object_conn = MagicMock(), MagicMock() + conn.post.return_value.ok = True + upload = { + "method": method, + "signed_url": "https://example.com/upload", + "headers": {}, + "fields": {"Content-Type": "application/json"}, + "key": "overflow-key", + } + with patch("atexit.register"): + bg = _HTTPBackgroundLogger(LazyValue(lambda: conn, use_mutex=False)) + with ( + patch.object(bg, "_request_logs3_overflow_upload", return_value=upload) as request_upload, + patch.object(logger, "HTTPConnection", return_value=object_conn), + ): + bg._submit_logs_request([item], {"max_request_size": 100, "can_use_overflow": True}) + + assert request_upload.call_args.args[1] == len(expected) + assert request_upload.call_args.args[2][0]["input_row"]["byte_size"] == item.overflow_meta.byte_size + if method == "PUT": + assert object_conn.put.call_args.kwargs["data"] == expected + else: + assert object_conn.post.call_args.kwargs["files"]["file"][1] == expected + assert conn.post.call_args.kwargs["json"]["rows"]["key"] == "overflow-key" + assert conn.post.call_args.kwargs["json"]["rows"]["size_bytes"] == len(expected) + assert bg._overflow_upload_count == 1 + def test_load_prompt_async_signature_matches_load_prompt(): assert ( @@ -835,6 +910,42 @@ def test_extract_attachments_with_attachments(): } +@pytest.mark.parametrize("subclasses", [False, True]) +def test_attachment_roundtrip_nested_containers(subclasses): + from collections import UserDict + + class Dict(dict): + pass + + class List(list): + def __iter__(self): + return iter(()) + + dict_type, list_type = (Dict, List) if subclasses else (dict, list) + attachment = Attachment(data=b"data", filename="file", content_type="text/plain") + external = ExternalAttachment(url="s3://bucket/file", filename="file", content_type="text/plain") + values = list_type([None, False, 42, 1.5, "text", attachment, dict_type(file=external)]) + event = dict_type(input=dict_type(values=values)) + if subclasses: + event = UserDict(event) + attachments = [] + + _extract_attachments(event, attachments) + + assert attachments == [attachment, external] + assert values[:5] == [None, False, 42, 1.5, "text"] + assert values[5] is attachment.reference + assert values[6]["file"] is external.reference + + assert logger._enrich_attachments(event) is event + assert event["input"]["values"] is values + assert isinstance(values[5], logger.ReadonlyAttachment) + assert values[5].reference is attachment.reference + assert isinstance(values[6]["file"], logger.ReadonlyAttachment) + assert values[6]["file"].reference is external.reference + assert values[:5] == [None, False, 42, 1.5, "text"] + + def _test_prompt(content: str, options: dict | None = None) -> Prompt: """Create a chat prompt with a single user message.""" return Prompt.from_prompt_data( diff --git a/py/src/braintrust/test_merge_row_batch.py b/py/src/braintrust/test_merge_row_batch.py index 7e373a3c3..745334d73 100644 --- a/py/src/braintrust/test_merge_row_batch.py +++ b/py/src/braintrust/test_merge_row_batch.py @@ -133,6 +133,9 @@ def test_skip_fields(self): class BatchItemsTest(unittest.TestCase): + def test_exact_byte_limit(self): + self.assertEqual(batch_items(["aa", "bb", "c"], batch_max_num_bytes=4), [["aa", "bb"], ["c"]]) + def test_basic(self): a = "x" * 1 b = "x" * 2 diff --git a/py/src/braintrust/test_queue.py b/py/src/braintrust/test_queue.py index ab26e46ab..3dde12df3 100644 --- a/py/src/braintrust/test_queue.py +++ b/py/src/braintrust/test_queue.py @@ -30,6 +30,24 @@ def test_log_queue_basic_operations(): assert queue.size() == 0 +@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: + for item in range(5): + queue.put(item) + assert signal.call_count == 1 + assert queue.wait_for_items(timeout=0) + assert queue.drain_all() == [3, 4] + assert not queue.wait_for_items(timeout=0) + queue.put(5) + assert signal.call_count == 2 + assert queue.wait_for_items(timeout=0) + + 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 From 92898ee30a91a42d61a45f6f1c14714fae4f1b96 Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Mon, 5 Oct 2026 10:47:11 -0400 Subject: [PATCH 2/3] better assertion on test --- py/src/braintrust/test_logger.py | 18 +++++++++++------- 1 file changed, 11 insertions(+), 7 deletions(-) diff --git a/py/src/braintrust/test_logger.py b/py/src/braintrust/test_logger.py index d927c38d2..38ba9d85e 100644 --- a/py/src/braintrust/test_logger.py +++ b/py/src/braintrust/test_logger.py @@ -265,25 +265,30 @@ def test_flush_uses_request_byte_budget(self) -> None: from braintrust.logger import _HTTPBackgroundLogger, construct_logs3_data, stringify_with_overflow_meta rows = [{"id": str(i), "dataset_id": "dataset", "input": "δΈ­ζ–‡πŸ™‚" * 20} for i in range(6)] - limit = len(construct_logs3_data([stringify_with_overflow_meta(row) for row in rows[:3]]).encode()) + serialized_rows = [stringify_with_overflow_meta(row) for row in rows] + three_row_size = len(construct_logs3_data(serialized_rows[:3]).encode("utf-8")) + four_row_size = len(construct_logs3_data(serialized_rows[:4]).encode("utf-8")) + self.assertGreater(four_row_size, three_row_size) conn = MagicMock() conn.post.return_value.ok = True with patch("atexit.register"): bg = _HTTPBackgroundLogger(LazyValue(lambda: conn, use_mutex=False)) - bg._max_request_size_result = {"max_request_size": limit, "can_use_overflow": True} + bg._max_request_size_result = {"max_request_size": three_row_size, "can_use_overflow": True} for row in rows: bg.queue.put(LazyValue(lambda row=row: row, use_mutex=False)) bg.flush(batch_size=100) self.assertEqual(conn.post.call_count, 2) - posted_rows = [] + posted_batches = [] for call in conn.post.call_args_list: self.assertEqual(call.args, ("/logs3",)) - self.assertLessEqual(len(call.kwargs["data"]), limit) - posted_rows.extend(json.loads(call.kwargs["data"])["rows"]) - self.assertCountEqual(posted_rows, rows) + request_bytes = call.kwargs["data"] + self.assertLessEqual(len(request_bytes), three_row_size) + posted_batches.append(json.loads(request_bytes)["rows"]) + self.assertEqual(posted_batches, [rows[:3], rows[3:]]) + self.assertEqual(len(conn.post.call_args_list[0].kwargs["data"]), three_row_size) def test_flush_keeps_oversized_row_separate(self) -> None: from braintrust.logger import _HTTPBackgroundLogger @@ -943,7 +948,6 @@ def __iter__(self): assert values[5].reference is attachment.reference assert isinstance(values[6]["file"], logger.ReadonlyAttachment) assert values[6]["file"].reference is external.reference - assert values[:5] == [None, False, 42, 1.5, "text"] def _test_prompt(content: str, options: dict | None = None) -> Prompt: From 3aa38ec0725f940bc3dde9303212640f3e2e8b28 Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Mon, 5 Oct 2026 10:55:39 -0400 Subject: [PATCH 3/3] test(logger): allow concurrent request ordering --- py/src/braintrust/test_logger.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/py/src/braintrust/test_logger.py b/py/src/braintrust/test_logger.py index 38ba9d85e..19dc3add3 100644 --- a/py/src/braintrust/test_logger.py +++ b/py/src/braintrust/test_logger.py @@ -282,13 +282,16 @@ def test_flush_uses_request_byte_budget(self) -> None: self.assertEqual(conn.post.call_count, 2) posted_batches = [] + request_sizes = [] for call in conn.post.call_args_list: self.assertEqual(call.args, ("/logs3",)) request_bytes = call.kwargs["data"] self.assertLessEqual(len(request_bytes), three_row_size) + request_sizes.append(len(request_bytes)) posted_batches.append(json.loads(request_bytes)["rows"]) + posted_batches.sort(key=lambda batch: batch[0]["id"]) self.assertEqual(posted_batches, [rows[:3], rows[3:]]) - self.assertEqual(len(conn.post.call_args_list[0].kwargs["data"]), three_row_size) + self.assertEqual(request_sizes, [three_row_size, three_row_size]) def test_flush_keeps_oversized_row_separate(self) -> None: from braintrust.logger import _HTTPBackgroundLogger