diff --git a/CHANGELOG.md b/CHANGELOG.md index 0218f15..523194e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,6 +13,7 @@ The format is based on Keep a Changelog and this project follows Semantic Versio ### Changed - Skip bundle-backed code functions during migration. A code function's compiled bundle is produced by the push/eval build pipeline and is not exposed by the API, so it can't be recreated in the destination — migrating it produces a broken function. These are now skipped (recorded with `skip_reason="code_bundle_not_migratable"` and logged with name/slug so they can be re-pushed manually). Inline code functions (which carry their source) and all other function types continue to migrate normally. +- Split aggregate SDK logging submissions by prepared payload size so batches of individually small spans also remain below the configured request threshold. ### Fixed diff --git a/braintrust_migrate/sdk_logs.py b/braintrust_migrate/sdk_logs.py index 10f336f..fac01b8 100644 --- a/braintrust_migrate/sdk_logs.py +++ b/braintrust_migrate/sdk_logs.py @@ -6,7 +6,9 @@ from collections.abc import Mapping, Sequence from typing import Any +from braintrust_migrate.batching import iter_ordered_batches_by_count_and_bytes from braintrust_migrate.client import BraintrustClient +from braintrust_migrate.streaming_utils import StreamingConfig PROJECT_LOGS_LOG_ID = "g" @@ -49,21 +51,41 @@ def _prepare_row(self, row: dict[str, Any]) -> dict[str, Any]: **self._object_id_fields, } + def _max_request_bytes(self) -> int: + stream_config = StreamingConfig.resolve(self._dest_client, self._dest_client) + migration_config = getattr(self._dest_client, "migration_config", None) + insert_max_request_bytes = int( + getattr(migration_config, "insert_max_request_bytes", 6 * 1024 * 1024) + ) + insert_request_headroom_ratio = float( + getattr(migration_config, "insert_request_headroom_ratio", 0.75) + ) + return min( + stream_config.max_event_bytes, + int(insert_max_request_bytes * insert_request_headroom_ratio), + ) + def write_rows_sync(self, rows: Sequence[dict[str, Any]]) -> None: self._ensure_logger() assert self._background_logger is not None assert self._lazy_value_cls is not None - events = [ - self._lazy_value_cls( - lambda prepared=self._prepare_row(row): prepared, - use_mutex=False, - ) - for row in rows - ] - if events: + prepared_rows = [self._prepare_row(row) for row in rows] + for batch in iter_ordered_batches_by_count_and_bytes( + prepared_rows, + max_items=max(1, len(prepared_rows)), + max_bytes=self._max_request_bytes(), + exact_wrapper_bytes=True, + ): + events = [ + self._lazy_value_cls( + lambda prepared=prepared: prepared, + use_mutex=False, + ) + for prepared in batch + ] self._background_logger.log(*events) - self._background_logger.flush() + self._background_logger.flush() async def write_rows(self, rows: Sequence[dict[str, Any]]) -> None: await asyncio.to_thread(self.write_rows_sync, rows) diff --git a/tests/unit/test_sdk_logs.py b/tests/unit/test_sdk_logs.py index 562175d..5b11166 100644 --- a/tests/unit/test_sdk_logs.py +++ b/tests/unit/test_sdk_logs.py @@ -5,12 +5,14 @@ from typing import Any from braintrust_migrate.sdk_logs import ( - SDKDatasetWriter, PROJECT_LOGS_LOG_ID, + SDKDatasetWriter, SDKExperimentWriter, SDKProjectLogsWriter, ) +EXPECTED_SPLIT_FLUSHES = 2 + class _FakeLazyValue: def __init__(self, fn, use_mutex: bool = False) -> None: @@ -39,10 +41,13 @@ def __init__(self, api_conn: _FakeLazyValue) -> None: self.api_conn = api_conn self.sync_flush = False self.logged_rows: list[dict[str, Any]] = [] + self.log_calls: list[list[dict[str, Any]]] = [] self.flush_count = 0 def log(self, *args: _FakeLazyValue) -> None: - self.logged_rows.extend(arg.get() for arg in args) + rows = [arg.get() for arg in args] + self.log_calls.append(rows) + self.logged_rows.extend(rows) def flush(self) -> None: self.flush_count += 1 @@ -85,6 +90,44 @@ def test_sdk_project_logs_writer_adds_logs3_object_ids(monkeypatch) -> None: assert conn.long_lived is True +def test_sdk_project_logs_writer_splits_aggregate_payloads(monkeypatch) -> None: + fake_logger_module = ModuleType("braintrust.logger") + fake_logger_module.HTTPConnection = _FakeHTTPConnection + fake_logger_module._HTTPBackgroundLogger = _FakeBackgroundLogger + fake_util_module = ModuleType("braintrust.util") + fake_util_module.LazyValue = _FakeLazyValue + + monkeypatch.setitem(sys.modules, "braintrust.logger", fake_logger_module) + monkeypatch.setitem(sys.modules, "braintrust.util", fake_util_module) + + dest_client = SimpleNamespace( + org_config=SimpleNamespace( + url="https://api.example.com", + api_key="secret-token", + ), + migration_config=SimpleNamespace( + insert_max_request_bytes=1_000, + insert_request_headroom_ratio=1.0, + ), + ) + + writer = SDKProjectLogsWriter(dest_client, "dest-project-id") + writer.write_rows_sync( + [ + {"id": "row1", "input": "x" * 700}, + {"id": "row2", "input": "y" * 700}, + ] + ) + + logger = writer._background_logger + assert logger is not None + assert logger.flush_count == EXPECTED_SPLIT_FLUSHES + assert [[row["id"] for row in call] for call in logger.log_calls] == [ + ["row1"], + ["row2"], + ] + + def test_sdk_experiment_writer_adds_experiment_id(monkeypatch) -> None: fake_logger_module = ModuleType("braintrust.logger") fake_logger_module.HTTPConnection = _FakeHTTPConnection