perf: Optimize ingest delivery for logs

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()

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
 

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()

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

py/src/braintrust/test_logger.py

@@ -261,6 +261,55 @@ 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)]
+        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": 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_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(request_sizes, [three_row_size, three_row_size])
+
+    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 +362,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 +918,41 @@ 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
+
+
 def _test_prompt(content: str, options: dict | None = None) -> Prompt:
     """Create a chat prompt with a single user message."""
     return Prompt.from_prompt_data(