From c982a5e26e7c9c7506e97f6b3f51b09db673f20d Mon Sep 17 00:00:00 2001 From: barryj0915 Date: Tue, 28 Jul 2026 14:01:21 -0400 Subject: [PATCH 1/2] changing spiller size limit and adding config option --- CHANGELOG.md | 1 + README.md | 1 + braintrust_migrate/attachments.py | 22 ++++++++++------------ braintrust_migrate/cli.py | 1 + braintrust_migrate/config.py | 14 ++++++++++++++ braintrust_migrate/streaming_utils.py | 14 +++++++++----- tests/unit/test_config.py | 10 ++++++++++ tests/unit/test_logs_oversize_event_e2e.py | 17 +++++++---------- tests/unit/test_streaming_helpers.py | 20 +++++++++++++++++++- 9 files changed, 72 insertions(+), 28 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 0218f15..b47b9fe 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. +- Lower the default streaming event spill threshold to 3MB and make it configurable through `MIGRATION_EVENTS_MAX_EVENT_BYTES`, preventing individual spans below the 20MB per-span limit from exceeding intermediary gateway request limits. ### Fixed diff --git a/README.md b/README.md index d2a2fef..bcfeb3c 100644 --- a/README.md +++ b/README.md @@ -168,6 +168,7 @@ These settings control BTQL-based streaming for high-volume resources. | `MIGRATION_EVENTS_FETCH_LIMIT` | — | `1000` | BTQL fetch page size (rows per query) | | `MIGRATION_EVENTS_FETCH_GROUP_SIZE` | — | `25` | Number of experiment or dataset ids to group into one BTQL event stream | | `MIGRATION_EVENTS_FLUSH_MAX_ROWS` | — | `5000` | Buffered flush threshold shared by logs, experiment events, and dataset events | +| `MIGRATION_EVENTS_MAX_EVENT_BYTES` | — | `3145728` (3MB) | Spill large event fields to Braintrust-managed attachments above this serialized size | | `MIGRATION_EVENTS_USE_SEEN_DB` | — | `true` | Use SQLite store for deduplication | | `MIGRATION_LOGS_FETCH_LIMIT` | `--logs-fetch-limit` | *(inherits)* | Override fetch limit for logs only | | `MIGRATION_LOGS_INSERT_BATCH_SIZE` | `--logs-insert-batch-size` | `5000` | Deprecated alias for `MIGRATION_EVENTS_FLUSH_MAX_ROWS` | diff --git a/braintrust_migrate/attachments.py b/braintrust_migrate/attachments.py index ca97c41..5b7fbd4 100644 --- a/braintrust_migrate/attachments.py +++ b/braintrust_migrate/attachments.py @@ -21,11 +21,9 @@ AttachmentRef = dict[str, Any] -# Braintrust enforces a ~20MB per-span limit on individual logging upload -# requests (this also bounds the logs3 overflow path). We spill below it with -# headroom so the small inline attachment reference plus the rest of the row -# stays comfortably under the cap. -DEFAULT_MAX_EVENT_BYTES = 18 * 1024 * 1024 +# Keep individual events below intermediary gateway request limits, including +# envelope overhead around the serialized event. +DEFAULT_MAX_EVENT_BYTES = 3 * 1024 * 1024 # Fields eligible to be spilled into a JSON attachment when a row is too large, # ordered loosely by how likely they are to hold the bulk of the payload. We @@ -269,13 +267,13 @@ async def upload_bytes_as_attachment( class OversizeFieldSpiller: """Spill oversized event fields into JSON attachments on the destination. - Braintrust enforces a ~20MB per-span limit on individual logging upload - requests (this also bounds the logs3 overflow path). Some migrated spans - exceed it because a single field — typically ``input``/``output``/ - ``metadata`` — holds a very large JSON blob. This rewrites such fields into - ``braintrust_attachment`` references uploaded to the destination object - store, shrinking the inline row below the cap while keeping the data - accessible via the attachment viewer in the UI. + Logging requests can pass through gateways with limits below Braintrust's + per-span limit. Some migrated spans exceed those request limits because a + single field — typically ``input``/``output``/``metadata`` — holds a very + large JSON blob. This rewrites such fields into ``braintrust_attachment`` + references uploaded to the destination object store, shrinking the inline + row below the cap while keeping the data accessible via the attachment + viewer in the UI. """ dest_client: BraintrustClient diff --git a/braintrust_migrate/cli.py b/braintrust_migrate/cli.py index d15db58..97a94aa 100644 --- a/braintrust_migrate/cli.py +++ b/braintrust_migrate/cli.py @@ -414,6 +414,7 @@ async def _migrate_main( dry_run=dry_run, logs_fetch_limit=config.migration.logs_fetch_limit, events_flush_max_rows=config.migration.events_flush_max_rows, + events_max_event_bytes=config.migration.events_max_event_bytes, logs_insert_batch_size=config.migration.logs_insert_batch_size, created_after=config.migration.created_after, created_before=config.migration.created_before, diff --git a/braintrust_migrate/config.py b/braintrust_migrate/config.py index bc9d854..6e7b2c7 100644 --- a/braintrust_migrate/config.py +++ b/braintrust_migrate/config.py @@ -214,6 +214,15 @@ class MigrationConfig(BaseModel): "one flush is committed." ), ) + events_max_event_bytes: int = Field( + default=3 * 1024 * 1024, + ge=1, + le=500 * 1024 * 1024, + description=( + "Maximum serialized size of a streaming event before large fields " + "are spilled to Braintrust-managed attachments." + ), + ) logs_insert_batch_size: int = Field( default=5_000, ge=1, @@ -449,6 +458,7 @@ def from_env(cls) -> "Config": # MIGRATION_EVENTS_USE_SEEN_DB=true # MIGRATION_EVENTS_FETCH_GROUP_SIZE=25 # MIGRATION_EVENTS_FLUSH_MAX_ROWS=5000 + # MIGRATION_EVENTS_MAX_EVENT_BYTES=3145728 # # Resource-specific overrides (optional): # MIGRATION_LOGS_FETCH_LIMIT, MIGRATION_EXPERIMENT_EVENTS_FETCH_LIMIT, etc. @@ -467,6 +477,9 @@ def _get_bool(specific_key: str, unified_key: str, default: str) -> bool: os.getenv("MIGRATION_EVENTS_FLUSH_MAX_ROWS") or os.getenv("MIGRATION_LOGS_INSERT_BATCH_SIZE", "5000") ) + events_max_event_bytes = int( + os.getenv("MIGRATION_EVENTS_MAX_EVENT_BYTES", str(3 * 1024 * 1024)) + ) # Logs logs_fetch_limit = _get_int( @@ -598,6 +611,7 @@ def _get_bool(specific_key: str, unified_key: str, default: str) -> bool: insert_request_headroom_ratio=insert_request_headroom_ratio, logs_fetch_limit=logs_fetch_limit, events_flush_max_rows=events_flush_max_rows, + events_max_event_bytes=events_max_event_bytes, logs_insert_batch_size=logs_insert_batch_size, logs_use_version_snapshot=logs_use_version_snapshot, logs_use_seen_db=logs_use_seen_db, diff --git a/braintrust_migrate/streaming_utils.py b/braintrust_migrate/streaming_utils.py index e175c45..592a5bb 100644 --- a/braintrust_migrate/streaming_utils.py +++ b/braintrust_migrate/streaming_utils.py @@ -232,7 +232,7 @@ def coerce_int_config( # expose these as ClassVars for back-compat and tests. STREAMING_FLUSH_MAX_ROWS = 5_000 STREAMING_FLUSH_MAX_BYTES = 25 * 1024 * 1024 -STREAMING_MAX_EVENT_BYTES = 18 * 1024 * 1024 +STREAMING_MAX_EVENT_BYTES = 3 * 1024 * 1024 STREAMING_EVENT_FETCH_GROUP_SIZE = 25 @@ -253,9 +253,8 @@ class StreamingConfig: def resolve(cls, source_client: Any, dest_client: Any) -> StreamingConfig: """Resolve config from the dest client's migration_config (or source's). - ``events_flush_max_rows`` / ``events_fetch_group_size`` are the only - env-overridable knobs; the byte caps are fixed. Falls back to the module - defaults and is tolerant of lightweight test doubles / missing config. + Falls back to the module defaults and is tolerant of lightweight test + doubles / missing config. """ cfg = getattr(dest_client, "migration_config", None) or getattr( source_client, "migration_config", None @@ -265,7 +264,12 @@ def resolve(cls, source_client: Any, dest_client: Any) -> StreamingConfig: cfg, "events_flush_max_rows", STREAMING_FLUSH_MAX_ROWS, minimum=1 ), sdk_flush_max_bytes=STREAMING_FLUSH_MAX_BYTES, - max_event_bytes=STREAMING_MAX_EVENT_BYTES, + max_event_bytes=coerce_int_config( + cfg, + "events_max_event_bytes", + STREAMING_MAX_EVENT_BYTES, + minimum=1, + ), event_fetch_group_size=coerce_int_config( cfg, "events_fetch_group_size", diff --git a/tests/unit/test_config.py b/tests/unit/test_config.py index 7333725..b2c32ec 100644 --- a/tests/unit/test_config.py +++ b/tests/unit/test_config.py @@ -218,6 +218,16 @@ def test_unified_events_flush_max_rows_from_env(self, monkeypatch): assert config.migration.events_flush_max_rows == TEST_EVENTS_FLUSH_MAX_ROWS assert config.migration.logs_insert_batch_size == TEST_EVENTS_FLUSH_MAX_ROWS + def test_events_max_event_bytes_from_env(self, monkeypatch): + """Test the shared oversized-event spill threshold env var.""" + monkeypatch.setenv("BT_SOURCE_API_KEY", "source-test-key") + monkeypatch.setenv("BT_DEST_API_KEY", "dest-test-key") + monkeypatch.setenv("MIGRATION_EVENTS_MAX_EVENT_BYTES", "2097152") + + config = Config.from_env() + + assert config.migration.events_max_event_bytes == 2 * 1024 * 1024 + def test_project_map_file_from_env(self, monkeypatch, tmp_path: Path): """Test loading project name mapping from env-provided file.""" path = tmp_path / "project-map.json" diff --git a/tests/unit/test_logs_oversize_event_e2e.py b/tests/unit/test_logs_oversize_event_e2e.py index 5158c63..651033e 100644 --- a/tests/unit/test_logs_oversize_event_e2e.py +++ b/tests/unit/test_logs_oversize_event_e2e.py @@ -1,7 +1,7 @@ """End-to-end-ish test of the logs migrator spilling an oversized event. Drives the real LogsMigrator.migrate_all loop over a page of 5 events, one of -which is larger than the 20MB per-span logging limit. Source BTQL and the +which is larger than the streaming event threshold. Source BTQL and the destination attachment-upload handshake are mocked; the SDK logs writer is stubbed to capture exactly the rows that would be sent to /logs3. We then assert the oversized row was spilled to an attachment and every row is under the cap. @@ -19,9 +19,7 @@ import braintrust_migrate.resources.logs as logs_module from braintrust_migrate.resources.logs import LogsMigrator -# The SaaS cap (LOGS_MAX_SPAN_BYTES) that produced the original 400 on -# /logs3/overflow. Confirmed in braintrust api-ts (LOGS_MAX_SPAN_MB=20). -LOGS_MAX_SPAN_BYTES = 20 * 1024 * 1024 +MAX_EVENT_BYTES = 3 * 1024 * 1024 class _SourceStub: @@ -108,8 +106,8 @@ async def write_rows(self, rows: list[dict[str, Any]]) -> None: @pytest.mark.asyncio async def test_oversized_event_among_normal_events_is_spilled(tmp_path: Path) -> None: - # One ~21MB event (over the 20MB cap) mixed with four small ones. - big_input = {"transcript": "x" * (21 * 1024 * 1024)} + # One ~4MB event (over the default spill threshold) mixed with four small ones. + big_input = {"transcript": "x" * (4 * 1024 * 1024)} events: list[dict[str, Any]] = [ { "id": f"e{i}", @@ -123,8 +121,7 @@ async def test_oversized_event_among_normal_events_is_spilled(tmp_path: Path) -> events[2]["input"] = big_input # Sanity: the oversized event really would exceed the span cap as-is. assert ( - len(json.dumps(events[2], ensure_ascii=False).encode("utf-8")) - > LOGS_MAX_SPAN_BYTES + len(json.dumps(events[2], ensure_ascii=False).encode("utf-8")) > MAX_EVENT_BYTES ) source = _SourceStub(btql_pages=[events]) @@ -168,10 +165,10 @@ async def test_oversized_event_among_normal_events_is_spilled(tmp_path: Path) -> assert spilled_input["filename"] == "input.json" assert spilled_input["content_type"] == "application/json" - # ...and every row actually handed to /logs3 is under the real 20MB cap. + # ...and every row handed to /logs3 is under the configured threshold. for row in dest.inserted_rows: row_bytes = len(json.dumps(row, ensure_ascii=False).encode("utf-8")) - assert row_bytes < LOGS_MAX_SPAN_BYTES + assert row_bytes < MAX_EVENT_BYTES # The four normal events are untouched (no needless spilling/uploads). for i in (0, 1, 3, 4): diff --git a/tests/unit/test_streaming_helpers.py b/tests/unit/test_streaming_helpers.py index 3c8b247..1357be7 100644 --- a/tests/unit/test_streaming_helpers.py +++ b/tests/unit/test_streaming_helpers.py @@ -9,12 +9,14 @@ import json from pathlib import Path +from types import SimpleNamespace from typing import Any import httpx from braintrust_migrate.streaming_utils import ( EventsStreamState, + StreamingConfig, approx_event_size_bytes, build_stream_progress, count_attachment_refs, @@ -24,6 +26,17 @@ ) +def test_streaming_config_uses_configured_max_event_bytes(): + source = SimpleNamespace(migration_config=None) + destination = SimpleNamespace( + migration_config=SimpleNamespace(events_max_event_bytes=2 * 1024 * 1024) + ) + + config = StreamingConfig.resolve(source, destination) + + assert config.max_event_bytes == 2 * 1024 * 1024 + + def _http_status_error(status: int) -> httpx.HTTPStatusError: request = httpx.Request("POST", "https://api.example/logs3") response = httpx.Response(status, request=request) @@ -45,7 +58,12 @@ def test_approx_event_size_bytes(): def test_count_attachment_refs_walks_nested_structures(): - ref = {"type": "braintrust_attachment", "key": "k", "filename": "f", "content_type": "application/json"} + ref = { + "type": "braintrust_attachment", + "key": "k", + "filename": "f", + "content_type": "application/json", + } event = { "input": {"a": ref, "b": [ref, {"c": ref}]}, # 3 refs "metadata": {"type": "braintrust_attachment"}, # missing key -> not counted From fce31df7cc6e657c80309cc54142368893561691 Mon Sep 17 00:00:00 2001 From: barryj0915 Date: Wed, 5 Aug 2026 13:56:59 -0700 Subject: [PATCH 2/2] SDK request payload splitting --- CHANGELOG.md | 2 +- README.md | 1 - braintrust_migrate/attachments.py | 22 +++++----- braintrust_migrate/cli.py | 1 - braintrust_migrate/config.py | 14 ------- braintrust_migrate/sdk_logs.py | 40 +++++++++++++----- braintrust_migrate/streaming_utils.py | 14 +++---- tests/unit/test_config.py | 10 ----- tests/unit/test_logs_oversize_event_e2e.py | 17 ++++---- tests/unit/test_sdk_logs.py | 47 +++++++++++++++++++++- tests/unit/test_streaming_helpers.py | 20 +-------- 11 files changed, 105 insertions(+), 83 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index b47b9fe..523194e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,7 +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. -- Lower the default streaming event spill threshold to 3MB and make it configurable through `MIGRATION_EVENTS_MAX_EVENT_BYTES`, preventing individual spans below the 20MB per-span limit from exceeding intermediary gateway request limits. +- 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/README.md b/README.md index bcfeb3c..d2a2fef 100644 --- a/README.md +++ b/README.md @@ -168,7 +168,6 @@ These settings control BTQL-based streaming for high-volume resources. | `MIGRATION_EVENTS_FETCH_LIMIT` | — | `1000` | BTQL fetch page size (rows per query) | | `MIGRATION_EVENTS_FETCH_GROUP_SIZE` | — | `25` | Number of experiment or dataset ids to group into one BTQL event stream | | `MIGRATION_EVENTS_FLUSH_MAX_ROWS` | — | `5000` | Buffered flush threshold shared by logs, experiment events, and dataset events | -| `MIGRATION_EVENTS_MAX_EVENT_BYTES` | — | `3145728` (3MB) | Spill large event fields to Braintrust-managed attachments above this serialized size | | `MIGRATION_EVENTS_USE_SEEN_DB` | — | `true` | Use SQLite store for deduplication | | `MIGRATION_LOGS_FETCH_LIMIT` | `--logs-fetch-limit` | *(inherits)* | Override fetch limit for logs only | | `MIGRATION_LOGS_INSERT_BATCH_SIZE` | `--logs-insert-batch-size` | `5000` | Deprecated alias for `MIGRATION_EVENTS_FLUSH_MAX_ROWS` | diff --git a/braintrust_migrate/attachments.py b/braintrust_migrate/attachments.py index 5b7fbd4..ca97c41 100644 --- a/braintrust_migrate/attachments.py +++ b/braintrust_migrate/attachments.py @@ -21,9 +21,11 @@ AttachmentRef = dict[str, Any] -# Keep individual events below intermediary gateway request limits, including -# envelope overhead around the serialized event. -DEFAULT_MAX_EVENT_BYTES = 3 * 1024 * 1024 +# Braintrust enforces a ~20MB per-span limit on individual logging upload +# requests (this also bounds the logs3 overflow path). We spill below it with +# headroom so the small inline attachment reference plus the rest of the row +# stays comfortably under the cap. +DEFAULT_MAX_EVENT_BYTES = 18 * 1024 * 1024 # Fields eligible to be spilled into a JSON attachment when a row is too large, # ordered loosely by how likely they are to hold the bulk of the payload. We @@ -267,13 +269,13 @@ async def upload_bytes_as_attachment( class OversizeFieldSpiller: """Spill oversized event fields into JSON attachments on the destination. - Logging requests can pass through gateways with limits below Braintrust's - per-span limit. Some migrated spans exceed those request limits because a - single field — typically ``input``/``output``/``metadata`` — holds a very - large JSON blob. This rewrites such fields into ``braintrust_attachment`` - references uploaded to the destination object store, shrinking the inline - row below the cap while keeping the data accessible via the attachment - viewer in the UI. + Braintrust enforces a ~20MB per-span limit on individual logging upload + requests (this also bounds the logs3 overflow path). Some migrated spans + exceed it because a single field — typically ``input``/``output``/ + ``metadata`` — holds a very large JSON blob. This rewrites such fields into + ``braintrust_attachment`` references uploaded to the destination object + store, shrinking the inline row below the cap while keeping the data + accessible via the attachment viewer in the UI. """ dest_client: BraintrustClient diff --git a/braintrust_migrate/cli.py b/braintrust_migrate/cli.py index 97a94aa..d15db58 100644 --- a/braintrust_migrate/cli.py +++ b/braintrust_migrate/cli.py @@ -414,7 +414,6 @@ async def _migrate_main( dry_run=dry_run, logs_fetch_limit=config.migration.logs_fetch_limit, events_flush_max_rows=config.migration.events_flush_max_rows, - events_max_event_bytes=config.migration.events_max_event_bytes, logs_insert_batch_size=config.migration.logs_insert_batch_size, created_after=config.migration.created_after, created_before=config.migration.created_before, diff --git a/braintrust_migrate/config.py b/braintrust_migrate/config.py index 6e7b2c7..bc9d854 100644 --- a/braintrust_migrate/config.py +++ b/braintrust_migrate/config.py @@ -214,15 +214,6 @@ class MigrationConfig(BaseModel): "one flush is committed." ), ) - events_max_event_bytes: int = Field( - default=3 * 1024 * 1024, - ge=1, - le=500 * 1024 * 1024, - description=( - "Maximum serialized size of a streaming event before large fields " - "are spilled to Braintrust-managed attachments." - ), - ) logs_insert_batch_size: int = Field( default=5_000, ge=1, @@ -458,7 +449,6 @@ def from_env(cls) -> "Config": # MIGRATION_EVENTS_USE_SEEN_DB=true # MIGRATION_EVENTS_FETCH_GROUP_SIZE=25 # MIGRATION_EVENTS_FLUSH_MAX_ROWS=5000 - # MIGRATION_EVENTS_MAX_EVENT_BYTES=3145728 # # Resource-specific overrides (optional): # MIGRATION_LOGS_FETCH_LIMIT, MIGRATION_EXPERIMENT_EVENTS_FETCH_LIMIT, etc. @@ -477,9 +467,6 @@ def _get_bool(specific_key: str, unified_key: str, default: str) -> bool: os.getenv("MIGRATION_EVENTS_FLUSH_MAX_ROWS") or os.getenv("MIGRATION_LOGS_INSERT_BATCH_SIZE", "5000") ) - events_max_event_bytes = int( - os.getenv("MIGRATION_EVENTS_MAX_EVENT_BYTES", str(3 * 1024 * 1024)) - ) # Logs logs_fetch_limit = _get_int( @@ -611,7 +598,6 @@ def _get_bool(specific_key: str, unified_key: str, default: str) -> bool: insert_request_headroom_ratio=insert_request_headroom_ratio, logs_fetch_limit=logs_fetch_limit, events_flush_max_rows=events_flush_max_rows, - events_max_event_bytes=events_max_event_bytes, logs_insert_batch_size=logs_insert_batch_size, logs_use_version_snapshot=logs_use_version_snapshot, logs_use_seen_db=logs_use_seen_db, 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/braintrust_migrate/streaming_utils.py b/braintrust_migrate/streaming_utils.py index 592a5bb..e175c45 100644 --- a/braintrust_migrate/streaming_utils.py +++ b/braintrust_migrate/streaming_utils.py @@ -232,7 +232,7 @@ def coerce_int_config( # expose these as ClassVars for back-compat and tests. STREAMING_FLUSH_MAX_ROWS = 5_000 STREAMING_FLUSH_MAX_BYTES = 25 * 1024 * 1024 -STREAMING_MAX_EVENT_BYTES = 3 * 1024 * 1024 +STREAMING_MAX_EVENT_BYTES = 18 * 1024 * 1024 STREAMING_EVENT_FETCH_GROUP_SIZE = 25 @@ -253,8 +253,9 @@ class StreamingConfig: def resolve(cls, source_client: Any, dest_client: Any) -> StreamingConfig: """Resolve config from the dest client's migration_config (or source's). - Falls back to the module defaults and is tolerant of lightweight test - doubles / missing config. + ``events_flush_max_rows`` / ``events_fetch_group_size`` are the only + env-overridable knobs; the byte caps are fixed. Falls back to the module + defaults and is tolerant of lightweight test doubles / missing config. """ cfg = getattr(dest_client, "migration_config", None) or getattr( source_client, "migration_config", None @@ -264,12 +265,7 @@ def resolve(cls, source_client: Any, dest_client: Any) -> StreamingConfig: cfg, "events_flush_max_rows", STREAMING_FLUSH_MAX_ROWS, minimum=1 ), sdk_flush_max_bytes=STREAMING_FLUSH_MAX_BYTES, - max_event_bytes=coerce_int_config( - cfg, - "events_max_event_bytes", - STREAMING_MAX_EVENT_BYTES, - minimum=1, - ), + max_event_bytes=STREAMING_MAX_EVENT_BYTES, event_fetch_group_size=coerce_int_config( cfg, "events_fetch_group_size", diff --git a/tests/unit/test_config.py b/tests/unit/test_config.py index b2c32ec..7333725 100644 --- a/tests/unit/test_config.py +++ b/tests/unit/test_config.py @@ -218,16 +218,6 @@ def test_unified_events_flush_max_rows_from_env(self, monkeypatch): assert config.migration.events_flush_max_rows == TEST_EVENTS_FLUSH_MAX_ROWS assert config.migration.logs_insert_batch_size == TEST_EVENTS_FLUSH_MAX_ROWS - def test_events_max_event_bytes_from_env(self, monkeypatch): - """Test the shared oversized-event spill threshold env var.""" - monkeypatch.setenv("BT_SOURCE_API_KEY", "source-test-key") - monkeypatch.setenv("BT_DEST_API_KEY", "dest-test-key") - monkeypatch.setenv("MIGRATION_EVENTS_MAX_EVENT_BYTES", "2097152") - - config = Config.from_env() - - assert config.migration.events_max_event_bytes == 2 * 1024 * 1024 - def test_project_map_file_from_env(self, monkeypatch, tmp_path: Path): """Test loading project name mapping from env-provided file.""" path = tmp_path / "project-map.json" diff --git a/tests/unit/test_logs_oversize_event_e2e.py b/tests/unit/test_logs_oversize_event_e2e.py index 651033e..5158c63 100644 --- a/tests/unit/test_logs_oversize_event_e2e.py +++ b/tests/unit/test_logs_oversize_event_e2e.py @@ -1,7 +1,7 @@ """End-to-end-ish test of the logs migrator spilling an oversized event. Drives the real LogsMigrator.migrate_all loop over a page of 5 events, one of -which is larger than the streaming event threshold. Source BTQL and the +which is larger than the 20MB per-span logging limit. Source BTQL and the destination attachment-upload handshake are mocked; the SDK logs writer is stubbed to capture exactly the rows that would be sent to /logs3. We then assert the oversized row was spilled to an attachment and every row is under the cap. @@ -19,7 +19,9 @@ import braintrust_migrate.resources.logs as logs_module from braintrust_migrate.resources.logs import LogsMigrator -MAX_EVENT_BYTES = 3 * 1024 * 1024 +# The SaaS cap (LOGS_MAX_SPAN_BYTES) that produced the original 400 on +# /logs3/overflow. Confirmed in braintrust api-ts (LOGS_MAX_SPAN_MB=20). +LOGS_MAX_SPAN_BYTES = 20 * 1024 * 1024 class _SourceStub: @@ -106,8 +108,8 @@ async def write_rows(self, rows: list[dict[str, Any]]) -> None: @pytest.mark.asyncio async def test_oversized_event_among_normal_events_is_spilled(tmp_path: Path) -> None: - # One ~4MB event (over the default spill threshold) mixed with four small ones. - big_input = {"transcript": "x" * (4 * 1024 * 1024)} + # One ~21MB event (over the 20MB cap) mixed with four small ones. + big_input = {"transcript": "x" * (21 * 1024 * 1024)} events: list[dict[str, Any]] = [ { "id": f"e{i}", @@ -121,7 +123,8 @@ async def test_oversized_event_among_normal_events_is_spilled(tmp_path: Path) -> events[2]["input"] = big_input # Sanity: the oversized event really would exceed the span cap as-is. assert ( - len(json.dumps(events[2], ensure_ascii=False).encode("utf-8")) > MAX_EVENT_BYTES + len(json.dumps(events[2], ensure_ascii=False).encode("utf-8")) + > LOGS_MAX_SPAN_BYTES ) source = _SourceStub(btql_pages=[events]) @@ -165,10 +168,10 @@ async def test_oversized_event_among_normal_events_is_spilled(tmp_path: Path) -> assert spilled_input["filename"] == "input.json" assert spilled_input["content_type"] == "application/json" - # ...and every row handed to /logs3 is under the configured threshold. + # ...and every row actually handed to /logs3 is under the real 20MB cap. for row in dest.inserted_rows: row_bytes = len(json.dumps(row, ensure_ascii=False).encode("utf-8")) - assert row_bytes < MAX_EVENT_BYTES + assert row_bytes < LOGS_MAX_SPAN_BYTES # The four normal events are untouched (no needless spilling/uploads). for i in (0, 1, 3, 4): 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 diff --git a/tests/unit/test_streaming_helpers.py b/tests/unit/test_streaming_helpers.py index 1357be7..3c8b247 100644 --- a/tests/unit/test_streaming_helpers.py +++ b/tests/unit/test_streaming_helpers.py @@ -9,14 +9,12 @@ import json from pathlib import Path -from types import SimpleNamespace from typing import Any import httpx from braintrust_migrate.streaming_utils import ( EventsStreamState, - StreamingConfig, approx_event_size_bytes, build_stream_progress, count_attachment_refs, @@ -26,17 +24,6 @@ ) -def test_streaming_config_uses_configured_max_event_bytes(): - source = SimpleNamespace(migration_config=None) - destination = SimpleNamespace( - migration_config=SimpleNamespace(events_max_event_bytes=2 * 1024 * 1024) - ) - - config = StreamingConfig.resolve(source, destination) - - assert config.max_event_bytes == 2 * 1024 * 1024 - - def _http_status_error(status: int) -> httpx.HTTPStatusError: request = httpx.Request("POST", "https://api.example/logs3") response = httpx.Response(status, request=request) @@ -58,12 +45,7 @@ def test_approx_event_size_bytes(): def test_count_attachment_refs_walks_nested_structures(): - ref = { - "type": "braintrust_attachment", - "key": "k", - "filename": "f", - "content_type": "application/json", - } + ref = {"type": "braintrust_attachment", "key": "k", "filename": "f", "content_type": "application/json"} event = { "input": {"a": ref, "b": [ref, {"c": ref}]}, # 3 refs "metadata": {"type": "braintrust_attachment"}, # missing key -> not counted