diff --git a/examples/log_ingestion/README.md b/examples/log_ingestion/README.md new file mode 100644 index 000000000..02ebc7f7b --- /dev/null +++ b/examples/log_ingestion/README.md @@ -0,0 +1,17 @@ +# Policy-aware log ingestion demo + +From `py/`, with `BRAINTRUST_API_KEY` configured through mise: + +```bash +mise exec -- python ../examples/log_ingestion/demo.py --replay +``` + +The demo sends a root/child trace, flushes a score update, relogs in, and sends a final score and comment. `--replay` sends the same memoized feedback records again, preserving their row identities. The JSON output includes the trace link and remaining record count. Inspect the trace with: + +```bash +bt view trace --object-ref project_logs: --trace-id --json +``` + +Expect two spans, final root score `quality=1`, and one comment even after replay. + +`--overflow` lowers the writer's payload limit to exercise signed uploads on servers that advertise `logs3_payload_max_bytes`. The output reports `overflow_supported`; older servers continue using ordinary ingestion. No provider SDK or model calls are needed. diff --git a/examples/log_ingestion/demo.py b/examples/log_ingestion/demo.py new file mode 100644 index 000000000..9a6aff956 --- /dev/null +++ b/examples/log_ingestion/demo.py @@ -0,0 +1,72 @@ +"""Send nested traces and ordered score updates through the ingestion transport. + +From py/: mise exec -- python ../examples/log_ingestion/demo.py +Add --overflow to exercise a signed overflow upload on a supporting server. +""" + +import argparse +import json +import uuid + +import braintrust +from braintrust.logger import _internal_get_global_state + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--project", default="sdk-839-ingestion-demo") + parser.add_argument("--overflow", action="store_true") + parser.add_argument( + "--replay", action="store_true", help="Replay the same final feedback records to check server deduplication" + ) + args = parser.parse_args() + logger = braintrust.init_logger(project=args.project) + run_id = str(uuid.uuid4()) + state = _internal_get_global_state() + writer = state.global_bg_logger() + writer.sync_flush = True + if args.overflow: + writer._max_request_size_override = 1024 + writer._max_request_size_result = None + with logger.start_span( + name="policy-aware-ingestion-demo", input={"run_id": run_id}, metadata={"issue": 839} + ) as root: + with root.start_span(name="child", input="hello") as child: + child.log(output="world", scores={"quality": 1}) + root.log( + output={"status": "delivered", "payload": "ü" * (2000 if args.overflow else 1)}, scores={"quality": 0} + ) + braintrust.flush() + # Two later updates must reach the same row in order, including after relogin. + logger.log_feedback(id=root.id, scores={"quality": 0.5}) + braintrust.flush() + braintrust.login(force_login=True) + logger.log_feedback(id=root.id, scores={"quality": 1}, comment="Verified ordered score update after relogin") + replay_records = writer.queue.drain_all(reserve=True) if args.replay else [] + if replay_records: + writer.queue.restore(replay_records) + braintrust.flush() + if replay_records: + for record in replay_records: + writer.queue.put(record) + braintrust.flush() + print( + json.dumps( + { + "run_id": run_id, + "project_id": logger.id, + "row_id": root.id, + "trace_id": root.root_span_id, + "link": root.link(), + "pending_rows": writer.pending_count, + "overflow_uploads": writer._overflow_upload_count, + "replayed_feedback": bool(replay_records), + "overflow_supported": (writer._max_request_size_result or {}).get("can_use_overflow"), + }, + indent=2, + ) + ) + + +if __name__ == "__main__": + main() diff --git a/py/benchmarks/benches/bench_log_ingestion.py b/py/benchmarks/benches/bench_log_ingestion.py new file mode 100644 index 000000000..7701bbfea --- /dev/null +++ b/py/benchmarks/benches/bench_log_ingestion.py @@ -0,0 +1,51 @@ +"""Preparation and producer cost, independent of HTTP server latency. + +Use benchmarks.log_ingestion for end-to-end healthy/throttled measurements. +""" + +import contextlib +import os +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 _HTTPBackgroundLogger, construct_logs3_data, stringify_with_overflow_meta +from braintrust.util import LazyValue + +from benchmarks._utils import disable_pyperf_psutil +from benchmarks.fixtures import ingestion_rows + + +_ROWS = ingestion_rows() +_RECORDS = [LazyValue(lambda row=row: row, use_mutex=False) for row in _ROWS] + + +def _prepare() -> None: + construct_logs3_data([stringify_with_overflow_meta(row) for row in _ROWS]).encode("utf-8") + + +def main(runner: pyperf.Runner | None = None) -> None: + if runner is None: + disable_pyperf_psutil() + runner = pyperf.Runner() + # Pause the real writer's publisher to isolate producer work; no network calls occur. + os.environ["BRAINTRUST_DISABLE_ATEXIT_FLUSH"] = "1" + writer = _HTTPBackgroundLogger(lambda: contextlib.nullcontext(None)) + writer.sync_flush = True + writer._start() + + def enqueue() -> None: + writer.log(*_RECORDS) + writer.queue.drain_all() + + runner.bench_func("logs3[prepare-100-rows]", _prepare) + runner.bench_func("logs3[enqueue-100-rows]", enqueue) + + +if __name__ == "__main__": + main() diff --git a/py/benchmarks/fixtures.py b/py/benchmarks/fixtures.py index 907983e1b..b95041e31 100644 --- a/py/benchmarks/fixtures.py +++ b/py/benchmarks/fixtures.py @@ -201,3 +201,11 @@ def make_bt_safe_deep_copy_cases() -> list[tuple[str, Any]]: ("circular", make_circular_payload()), ("non-string-keys", make_non_string_key_payload()), ] + + +def ingestion_rows(count: int = 100) -> list[dict[str, Any]]: + """Deterministic trace-shaped rows for preparation and HTTP delivery measurements.""" + return [ + {"id": str(index), "project_id": "benchmark", "input": "hello" * 50, "scores": {"quality": 1}} + for index in range(count) + ] diff --git a/py/benchmarks/log_ingestion.py b/py/benchmarks/log_ingestion.py new file mode 100644 index 000000000..ed010684e --- /dev/null +++ b/py/benchmarks/log_ingestion.py @@ -0,0 +1,119 @@ +"""End-to-end HTTP ingestion measurements (run from py/). + +python -m benchmarks.log_ingestion --output /tmp/ingestion.json +Works on the legacy writer as well, for a same-harness baseline comparison. +""" + +import argparse +import contextlib +import json +import os +import statistics +import threading +import time +import tracemalloc + +from braintrust.api._test_server import scripted_server +from braintrust.api._transport import HTTPConnection +from braintrust.logger import _HTTPBackgroundLogger +from braintrust.util import LazyValue + +from benchmarks.fixtures import ingestion_rows + + +def measure(throttled, rows=2000, *, continuous=False): + rejected = 0 + accepted = 0 + arrivals = [] + connections = set() + lock = threading.Lock() + started = time.monotonic() + + def respond(method, path, body, headers): + nonlocal rejected, accepted + if path == "/version": + return 200, {}, b"{}" + now = time.monotonic() + with lock: + arrivals.append(now - started) + if throttled and now - started < 1: + rejected += 1 + return 429, {"Retry-After": "1"}, b"limited" + accepted += len(json.loads(body)["rows"]) + return "sleep", 0.005, 200, {}, b"ok" + + os.environ["BRAINTRUST_DISABLE_ATEXIT_FLUSH"] = "1" + os.environ["BRAINTRUST_NUM_RETRIES"] = "2" + with scripted_server(respond, persistent=True) as (url, handler): + # Select the branch's ingestion service when available. + try: + from braintrust.api._ingestion import LogIngestionAPI + from braintrust.api._routing import EndpointRouter + + connection = LogIngestionAPI(EndpointRouter(app_url=url, api_url=url), "benchmark", concurrency=4) + source = lambda: contextlib.nullcontext(connection) + except ImportError: + connection = HTTPConnection(url) + source = LazyValue(lambda: connection, use_mutex=False) + writer = _HTTPBackgroundLogger(source) + writer.sync_flush = not continuous + writer._max_request_size_result = {"max_request_size": 6_000_000, "can_use_overflow": False} + latencies = [] + tracemalloc.start() + started = time.monotonic() + peak_outstanding = 0 + for index, row in enumerate(ingestion_rows(rows)): + tick = time.perf_counter() + writer.log(LazyValue(lambda row=row: row, use_mutex=False)) + latencies.append(time.perf_counter() - tick) + with lock: + peak_outstanding = max(peak_outstanding, index + 1 - accepted) + if continuous: + time.sleep(0.0005) + queued = writer.queue.size() + tick = time.monotonic() + writer.flush() + flush_seconds = time.monotonic() - tick + elapsed = time.monotonic() - started + _, peak = tracemalloc.get_traced_memory() + tracemalloc.stop() + if accepted != rows: + raise RuntimeError(f"Benchmark delivered {accepted} of {rows} rows") + writer.sync_flush = False + connection.close() + connections.update(getattr(handler, "connections", [])) + return { + "rows": rows, + "delivered_rows": accepted, + "rejected_requests": rejected, + "requests": len(arrivals), + "queued_rows_at_flush": queued, + "outstanding_rows_peak": peak_outstanding, + "pending_rows_after_flush": getattr(writer, "pending_count", 0), + "rows_per_second": rows / elapsed, + "flush_seconds": flush_seconds, + "producer_p50_us": statistics.median(latencies) * 1e6, + "producer_p99_us": sorted(latencies)[int(len(latencies) * 0.99)] * 1e6, + "peak_allocated_bytes": peak, + "recovery_seconds": max(arrivals) if arrivals else 0, + "connections": len(connections), + } + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output", required=True) + parser.add_argument("--runs", type=int, default=3) + args = parser.parse_args() + results = { + name: [measure(throttled) for _ in range(args.runs)] + for name, throttled in [("healthy", False), ("throttled", True)] + } + results["continuous_throttled"] = [measure(True, continuous=True) for _ in range(args.runs)] + with open(args.output, "w") as output: + json.dump(results, output, indent=2) + print(json.dumps(results, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/py/benchmarks/results/log_ingestion/baseline-http.json b/py/benchmarks/results/log_ingestion/baseline-http.json new file mode 100644 index 000000000..0ae1d118a --- /dev/null +++ b/py/benchmarks/results/log_ingestion/baseline-http.json @@ -0,0 +1,152 @@ +{ + "healthy": [ + { + "rows": 2000, + "delivered_rows": 2000, + "rejected_requests": 0, + "requests": 20, + "queued_rows_at_flush": 2000, + "outstanding_rows_peak": 2000, + "pending_rows_after_flush": 0, + "rows_per_second": 15611.670057581148, + "flush_seconds": 0.1166860421653837, + "producer_p50_us": 3.5830307751893997, + "producer_p99_us": 5.041016265749931, + "peak_allocated_bytes": 7099339, + "recovery_seconds": 0.12039354210719466, + "connections": 12 + }, + { + "rows": 2000, + "delivered_rows": 2000, + "rejected_requests": 0, + "requests": 20, + "queued_rows_at_flush": 2000, + "outstanding_rows_peak": 2000, + "pending_rows_after_flush": 0, + "rows_per_second": 16166.330076030152, + "flush_seconds": 0.1147070829756558, + "producer_p50_us": 2.708984538912773, + "producer_p99_us": 3.2919924706220627, + "peak_allocated_bytes": 6628689, + "recovery_seconds": 0.11542949988506734, + "connections": 12 + }, + { + "rows": 2000, + "delivered_rows": 2000, + "rejected_requests": 0, + "requests": 20, + "queued_rows_at_flush": 2000, + "outstanding_rows_peak": 2000, + "pending_rows_after_flush": 0, + "rows_per_second": 21270.51441778817, + "flush_seconds": 0.08508162503130734, + "producer_p50_us": 2.707820385694504, + "producer_p99_us": 3.4170225262641907, + "peak_allocated_bytes": 6728063, + "recovery_seconds": 0.086067209020257, + "connections": 12 + } + ], + "throttled": [ + { + "rows": 2000, + "delivered_rows": 2000, + "rejected_requests": 12, + "requests": 32, + "queued_rows_at_flush": 2000, + "outstanding_rows_peak": 2000, + "pending_rows_after_flush": 0, + "rows_per_second": 1815.7892402699079, + "flush_seconds": 1.0934249160345644, + "producer_p50_us": 2.4579931050539017, + "producer_p99_us": 3.250082954764366, + "peak_allocated_bytes": 7437227, + "recovery_seconds": 1.0930026669520885, + "connections": 11 + }, + { + "rows": 2000, + "delivered_rows": 2000, + "rejected_requests": 12, + "requests": 32, + "queued_rows_at_flush": 2000, + "outstanding_rows_peak": 2000, + "pending_rows_after_flush": 0, + "rows_per_second": 1765.1538558242505, + "flush_seconds": 1.1217551671434194, + "producer_p50_us": 3.250082954764366, + "producer_p99_us": 8.58306884765625, + "peak_allocated_bytes": 7830173, + "recovery_seconds": 1.1241297498345375, + "connections": 12 + }, + { + "rows": 2000, + "delivered_rows": 2000, + "rejected_requests": 12, + "requests": 32, + "queued_rows_at_flush": 2000, + "outstanding_rows_peak": 2000, + "pending_rows_after_flush": 0, + "rows_per_second": 1746.91229568579, + "flush_seconds": 1.1322082499973476, + "producer_p50_us": 3.8750004023313522, + "producer_p99_us": 4.875008016824722, + "peak_allocated_bytes": 7747542, + "recovery_seconds": 1.1355997079517692, + "connections": 14 + } + ], + "continuous_throttled": [ + { + "rows": 2000, + "delivered_rows": 2000, + "rejected_requests": 1, + "requests": 54, + "queued_rows_at_flush": 5, + "outstanding_rows_peak": 1522, + "pending_rows_after_flush": 0, + "rows_per_second": 1370.3654412536657, + "flush_seconds": 0.014694207813590765, + "producer_p50_us": 20.999927073717117, + "producer_p99_us": 84.25000123679638, + "peak_allocated_bytes": 6003801, + "recovery_seconds": 1.4519599589984864, + "connections": 12 + }, + { + "rows": 2000, + "delivered_rows": 2000, + "rejected_requests": 1, + "requests": 59, + "queued_rows_at_flush": 5, + "outstanding_rows_peak": 1454, + "pending_rows_after_flush": 0, + "rows_per_second": 1309.7469379257598, + "flush_seconds": 0.015948459040373564, + "producer_p50_us": 21.479441784322262, + "producer_p99_us": 85.04185825586319, + "peak_allocated_bytes": 6238428, + "recovery_seconds": 1.5200314579997212, + "connections": 12 + }, + { + "rows": 2000, + "delivered_rows": 2000, + "rejected_requests": 1, + "requests": 57, + "queued_rows_at_flush": 5, + "outstanding_rows_peak": 1514, + "pending_rows_after_flush": 0, + "rows_per_second": 1361.2162604712644, + "flush_seconds": 0.01853725011460483, + "producer_p50_us": 21.416926756501198, + "producer_p99_us": 73.83385673165321, + "peak_allocated_bytes": 6108898, + "recovery_seconds": 1.4629678749479353, + "connections": 12 + } + ] +} diff --git a/py/benchmarks/results/log_ingestion/baseline-pyperf.json b/py/benchmarks/results/log_ingestion/baseline-pyperf.json new file mode 100644 index 000000000..b33b05974 --- /dev/null +++ b/py/benchmarks/results/log_ingestion/baseline-pyperf.json @@ -0,0 +1 @@ +{"benchmarks":[{"metadata":{"loops":512,"name":"logs3[prepare-100-rows]"},"runs":[{"metadata":{"calibrate_loops":512,"date":"2026-10-01 16:23:08.023281","duration":0.4464000838343054,"load_avg_1min":2.1572265625,"mem_max_rss":54296576},"warmups":[[1,0.00026050000451505184],[2,0.000205708434805274],[4,0.00019037502352148294],[8,0.00019450002582743764],[16,0.00020158068218734115],[32,0.00020549868349917233],[64,0.00020666081400122494],[128,0.00020764680448337458],[256,0.00020693131682492094],[512,0.0002054986994153296],[512,0.00020594116222127923],[512,0.00020679785166066722]]},{"metadata":{"date":"2026-10-01 16:23:08.534311","duration":0.3432835419662297,"load_avg_1min":2.46484375,"mem_max_rss":54427648},"values":[0.0002064007167064119,0.00020944173047610093],"warmups":[[512,0.00020664339263021247]]},{"metadata":{"date":"2026-10-01 16:23:09.046880","duration":0.3397464998997748,"load_avg_1min":2.46484375,"mem_max_rss":54296576},"values":[0.0002078627107948705,0.0002041062812168093],"warmups":[[512,0.0002057054025499383]]},{"metadata":{"date":"2026-10-01 16:23:09.565670","duration":0.3462902919854969,"load_avg_1min":2.46484375,"mem_max_rss":54804480},"values":[0.00021096541422593873,0.0002036796872744162],"warmups":[[512,0.00020805078111152397]]},{"metadata":{"date":"2026-10-01 16:23:10.086129","duration":0.34393291594460607,"load_avg_1min":2.46484375,"mem_max_rss":54902784},"values":[0.00020563810130624915,0.00020982934574931278],"warmups":[[512,0.0002058972168015316]]},{"metadata":{"date":"2026-10-01 16:23:10.600437","duration":0.34310254198499024,"load_avg_1min":2.46484375,"mem_max_rss":54083584},"values":[0.00020638159139707568,0.00020571419145198888],"warmups":[[512,0.0002070377599920903]]},{"metadata":{"date":"2026-10-01 16:23:11.130251","duration":0.34579770802520216,"load_avg_1min":2.46484375,"mem_max_rss":54804480},"values":[0.0002061330565084063,0.00020710685384983663],"warmups":[[512,0.00020871435526714777]]},{"metadata":{"date":"2026-10-01 16:23:11.655330","duration":0.3479377501644194,"load_avg_1min":2.46484375,"mem_max_rss":54362112},"values":[0.00020967398995708209,0.0002105125331581803],"warmups":[[512,0.00021220996086412924]]},{"metadata":{"date":"2026-10-01 16:23:12.163039","duration":0.34169133310206234,"load_avg_1min":2.46484375,"mem_max_rss":54525952},"values":[0.0002059946291410597,0.0002012368163377687],"warmups":[[512,0.00020645955237341695]]},{"metadata":{"date":"2026-10-01 16:23:12.681530","duration":0.3423548750579357,"load_avg_1min":2.46484375,"mem_max_rss":54312960},"values":[0.00020773331834789133,0.0002087727048092347],"warmups":[[512,0.0002079940586554585]]},{"metadata":{"date":"2026-10-01 16:23:13.192764","duration":0.347176417009905,"load_avg_1min":2.46484375,"mem_max_rss":54263808},"values":[0.00020784082016689354,0.00020932601000822615],"warmups":[[512,0.00020957210153937922]]}]},{"metadata":{"loops":4096,"name":"logs3[enqueue-100-rows]"},"runs":[{"metadata":{"calibrate_loops":4096,"date":"2026-10-01 16:23:13.881343","duration":0.5201240410096943,"load_avg_1min":2.34716796875,"mem_max_rss":54394880},"warmups":[[1,4.3582869693636894e-05],[2,2.954143565148115e-05],[4,2.9166520107537508e-05],[8,2.798961941152811e-05],[16,2.915626100730151e-05],[32,2.946615859400481e-05],[64,2.9821607313351706e-05],[128,3.007714985869825e-05],[256,3.033561006304808e-05],[512,2.9690594146813964e-05],[1024,3.028812602678954e-05],[2048,3.0445516131294426e-05],[4096,3.0189788844836585e-05],[4096,2.9865183364563563e-05],[4096,3.0375223843748245e-05]]},{"metadata":{"date":"2026-10-01 16:23:14.467580","duration":0.41324012516997755,"load_avg_1min":2.34716796875,"mem_max_rss":54935552},"values":[3.07065837432674e-05,3.116065477115626e-05],"warmups":[[4096,3.0452148450876848e-05]]},{"metadata":{"date":"2026-10-01 16:23:15.100002","duration":0.3997991669457406,"load_avg_1min":2.34716796875,"mem_max_rss":54870016},"values":[3.0256276374984736e-05,3.0404103767978086e-05],"warmups":[[4096,3.0126078115699784e-05]]},{"metadata":{"date":"2026-10-01 16:23:15.683226","duration":0.404148125089705,"load_avg_1min":2.34716796875,"mem_max_rss":54444032},"values":[3.09909565316957e-05,3.09626972807564e-05],"warmups":[[4096,3.074750782161573e-05]]},{"metadata":{"date":"2026-10-01 16:23:16.256792","duration":0.4002597499638796,"load_avg_1min":2.34716796875,"mem_max_rss":54919168},"values":[3.0372802768852125e-05,3.079607154177211e-05],"warmups":[[4096,3.0208740213311103e-05]]},{"metadata":{"date":"2026-10-01 16:23:16.842717","duration":0.4193739171605557,"load_avg_1min":2.34716796875,"mem_max_rss":54198272},"values":[3.12156372501704e-05,3.113588479664031e-05],"warmups":[[4096,3.212866215562826e-05]]},{"metadata":{"date":"2026-10-01 16:23:17.422617","duration":0.40613137488253415,"load_avg_1min":2.34716796875,"mem_max_rss":54345728},"values":[3.0597473141824594e-05,3.072897339961855e-05],"warmups":[[4096,3.0474772245270287e-05]]},{"metadata":{"date":"2026-10-01 16:23:18.001822","duration":0.3993266250472516,"load_avg_1min":2.34716796875,"mem_max_rss":54067200},"values":[3.0252726048729528e-05,3.0347086692472658e-05],"warmups":[[4096,3.027502441454999e-05]]},{"metadata":{"date":"2026-10-01 16:23:18.571800","duration":0.39767983299680054,"load_avg_1min":2.3193359375,"mem_max_rss":54050816},"values":[3.034726972828139e-05,3.0369537341812247e-05],"warmups":[[4096,3.0260274172633217e-05]]},{"metadata":{"date":"2026-10-01 16:23:19.166414","duration":0.4040500419214368,"load_avg_1min":2.3193359375,"mem_max_rss":54575104},"values":[3.059329225152396e-05,3.066180420319142e-05],"warmups":[[4096,3.0377960172245366e-05]]},{"metadata":{"date":"2026-10-01 16:23:19.746842","duration":0.40492012491449714,"load_avg_1min":2.3193359375,"mem_max_rss":54542336},"values":[3.033798219576056e-05,3.065036014504585e-05],"warmups":[[4096,3.039963769424503e-05]]}]}],"metadata":{"cpu_count":12,"hostname":"A-Prasad","perf_version":"2.10.0","platform":"macOS-26.6.2-arm64-arm-64bit-Mach-O","python_cflags":"-fno-strict-overflow -Wsign-compare -Wunreachable-code -DNDEBUG -g -O3 -Wall -O3 -arch arm64 -mmacosx-version-min=11.0 -Wno-nullability-completeness -Wno-expansion-to-defined -Wno-undef-prefix -fPIC","python_compiler":"Clang 22.1.3","python_config_args":"'--build=aarch64-apple-darwin' '--host=aarch64-apple-darwin' '--prefix=/install' '--with-openssl=/tmp/tmps_46j10_/tools/deps' '--with-system-expat' '--with-system-libmpdec' '--without-ensurepip' '--with-tail-call-interp' 'MODULE_BUILDTYPE=static' '--enable-loadable-sqlite-extensions' '--enable-static-libpython-for-interpreter' '--enable-shared' '--with-mimalloc' '--enable-optimizations' '--enable-experimental-jit=yes-off' '--with-lto' '--with-build-python=/tmp/tmps_46j10_/tools/host/bin/python3.14' 'ac_cv_lib_intl_textdomain=no' 'ac_cv_func_ptsname_r=no' 'ac_cv_func_explicit_bzero=no' '--with-dbmliborder=ndbm' 'build_alias=aarch64-apple-darwin' 'host_alias=aarch64-apple-darwin' 'PKG_CONFIG=pkg-config --static --define-prefix' 'PKG_CONFIG_PATH=/tmp/tmps_46j10_/tools/deps/share/pkgconfig:/tmp/tmps_46j10_/tools/deps/lib/pkgconfig' 'CC=clang' 'CFLAGS=-O3 -arch arm64 -mmacosx-version-min=11.0 -Wno-nullability-completeness -Wno-expansion-to-defined -Wno-undef-prefix -fPIC -Werror=unguarded-availability-new' 'LDFLAGS=-arch arm64 -mmacosx-version-min=11.0 -Wl,-headerpad,40' 'CPPFLAGS=-O3 -arch arm64 -mmacosx-version-min=11.0 -Wno-nullability-completeness -Wno-expansion-to-defined -Wno-undef-prefix -fPIC -Werror=unguarded-availability-new' 'PROFILE_TASK=-m test --pgo -j 12' 'BOLT_COMMON_FLAGS=' 'BOLT_APPLY_FLAGS='","python_executable":"/Users/abhijeetprasad/.superset/worktrees/braintrust-sdk-python/abhi/recondite-flier/py/.venv/bin/python","python_implementation":"cpython","python_version":"3.14.7 (64-bit)","timer":"mach_absolute_time(), resolution: 41.7 ns","unit":"second"},"version":"1.0"} diff --git a/py/benchmarks/results/log_ingestion/branch-http.json b/py/benchmarks/results/log_ingestion/branch-http.json new file mode 100644 index 000000000..862aa7289 --- /dev/null +++ b/py/benchmarks/results/log_ingestion/branch-http.json @@ -0,0 +1,152 @@ +{ + "healthy": [ + { + "rows": 2000, + "delivered_rows": 2000, + "rejected_requests": 0, + "requests": 20, + "queued_rows_at_flush": 2000, + "outstanding_rows_peak": 2000, + "pending_rows_after_flush": 0, + "rows_per_second": 18458.052673980543, + "flush_seconds": 0.09588183392770588, + "producer_p50_us": 3.6659184843301773, + "producer_p99_us": 6.166053935885429, + "peak_allocated_bytes": 4282213, + "recovery_seconds": 0.10145404213108122, + "connections": 4 + }, + { + "rows": 2000, + "delivered_rows": 2000, + "rejected_requests": 0, + "requests": 20, + "queued_rows_at_flush": 2000, + "outstanding_rows_peak": 2000, + "pending_rows_after_flush": 0, + "rows_per_second": 20416.288128412518, + "flush_seconds": 0.08968308288604021, + "producer_p50_us": 2.4170149117708206, + "producer_p99_us": 3.5830307751893997, + "peak_allocated_bytes": 3935177, + "recovery_seconds": 0.09193216706626117, + "connections": 4 + }, + { + "rows": 2000, + "delivered_rows": 2000, + "rejected_requests": 0, + "requests": 20, + "queued_rows_at_flush": 2000, + "outstanding_rows_peak": 2000, + "pending_rows_after_flush": 0, + "rows_per_second": 19361.131074439974, + "flush_seconds": 0.09469783282838762, + "producer_p50_us": 2.541113644838333, + "producer_p99_us": 3.6659184843301773, + "peak_allocated_bytes": 3935001, + "recovery_seconds": 0.096680958988145, + "connections": 4 + } + ], + "throttled": [ + { + "rows": 2000, + "delivered_rows": 2000, + "rejected_requests": 4, + "requests": 24, + "queued_rows_at_flush": 2000, + "outstanding_rows_peak": 2000, + "pending_rows_after_flush": 0, + "rows_per_second": 1502.2712463648581, + "flush_seconds": 1.3229867918416858, + "producer_p50_us": 2.4582259356975555, + "producer_p99_us": 3.5830307751893997, + "peak_allocated_bytes": 4448453, + "recovery_seconds": 1.3232313750777394, + "connections": 4 + }, + { + "rows": 2000, + "delivered_rows": 2000, + "rejected_requests": 4, + "requests": 24, + "queued_rows_at_flush": 2000, + "outstanding_rows_peak": 2000, + "pending_rows_after_flush": 0, + "rows_per_second": 1508.9977766705422, + "flush_seconds": 1.3129435000009835, + "producer_p50_us": 2.4170149117708206, + "producer_p99_us": 6.37490302324295, + "peak_allocated_bytes": 4425713, + "recovery_seconds": 1.3177713339682668, + "connections": 4 + }, + { + "rows": 2000, + "delivered_rows": 2000, + "rejected_requests": 4, + "requests": 24, + "queued_rows_at_flush": 2000, + "outstanding_rows_peak": 2000, + "pending_rows_after_flush": 0, + "rows_per_second": 1516.4792480534306, + "flush_seconds": 1.3100886249449104, + "producer_p50_us": 2.416083589196205, + "producer_p99_us": 3.00002284348011, + "peak_allocated_bytes": 4425189, + "recovery_seconds": 1.3112428339663893, + "connections": 4 + } + ], + "continuous_throttled": [ + { + "rows": 2000, + "delivered_rows": 2000, + "rejected_requests": 1, + "requests": 33, + "queued_rows_at_flush": 0, + "outstanding_rows_peak": 1580, + "pending_rows_after_flush": 0, + "rows_per_second": 1473.396808121645, + "flush_seconds": 0.007899792166426778, + "producer_p50_us": 5.458015948534012, + "producer_p99_us": 37.20889799296856, + "peak_allocated_bytes": 4176479, + "recovery_seconds": 1.35044787498191, + "connections": 4 + }, + { + "rows": 2000, + "delivered_rows": 2000, + "rejected_requests": 1, + "requests": 33, + "queued_rows_at_flush": 4, + "outstanding_rows_peak": 1588, + "pending_rows_after_flush": 0, + "rows_per_second": 1478.3942842725528, + "flush_seconds": 0.0124506251886487, + "producer_p50_us": 6.041489541530609, + "producer_p99_us": 41.74979403614998, + "peak_allocated_bytes": 4219856, + "recovery_seconds": 1.3460560832172632, + "connections": 4 + }, + { + "rows": 2000, + "delivered_rows": 2000, + "rejected_requests": 1, + "requests": 41, + "queued_rows_at_flush": 12, + "outstanding_rows_peak": 1554, + "pending_rows_after_flush": 0, + "rows_per_second": 1340.3736459956085, + "flush_seconds": 0.009492750046774745, + "producer_p50_us": 8.167000487446785, + "producer_p99_us": 90.87519720196724, + "peak_allocated_bytes": 4283861, + "recovery_seconds": 1.4851050418801606, + "connections": 4 + } + ] +} diff --git a/py/benchmarks/results/log_ingestion/branch-pyperf.json b/py/benchmarks/results/log_ingestion/branch-pyperf.json new file mode 100644 index 000000000..bbd67aa69 --- /dev/null +++ b/py/benchmarks/results/log_ingestion/branch-pyperf.json @@ -0,0 +1 @@ +{"benchmarks":[{"metadata":{"loops":512,"name":"logs3[prepare-100-rows]"},"runs":[{"metadata":{"calibrate_loops":512,"date":"2026-10-01 16:33:39.858575","duration":0.5096242078579962,"load_avg_1min":3.6337890625,"mem_max_rss":55083008},"warmups":[[1,0.00023504206910729408],[2,0.00020041700918227434],[4,0.00023213523672893643],[8,0.00021514063701033592],[16,0.00021784637647215277],[32,0.00021619530889438465],[64,0.00021252864098642021],[128,0.00020838216369156726],[256,0.00021260986341076205],[512,0.0002117078456649324],[512,0.00021363492942327866],[512,0.00032818782619870035]]},{"metadata":{"date":"2026-10-01 16:33:40.382346","duration":0.3470744590740651,"load_avg_1min":3.6337890625,"mem_max_rss":55033856},"values":[0.00021299788477335824,0.00021817407196067506],"warmups":[[512,0.00021567928115473478]]},{"metadata":{"date":"2026-10-01 16:33:40.988314","duration":0.35921195801347494,"load_avg_1min":3.6337890625,"mem_max_rss":55443456},"values":[0.000215645996377134,0.00021471842592291068],"warmups":[[512,0.0002129268382304872]]},{"metadata":{"date":"2026-10-01 16:33:41.528156","duration":0.3584591660182923,"load_avg_1min":3.6337890625,"mem_max_rss":54624256},"values":[0.00021505704717128538,0.00021147933011889108],"warmups":[[512,0.0002148891599063063]]},{"metadata":{"date":"2026-10-01 16:33:42.069498","duration":0.3519462919794023,"load_avg_1min":3.6337890625,"mem_max_rss":54771712},"values":[0.00021254174816931481,0.00021372688661358552],"warmups":[[512,0.00021433365827761008]]},{"metadata":{"date":"2026-10-01 16:33:42.592596","duration":0.34996945899911225,"load_avg_1min":3.6337890625,"mem_max_rss":54722560},"values":[0.00020816430651393603,0.00020818977964154328],"warmups":[[512,0.00020915592585879494]]},{"metadata":{"date":"2026-10-01 16:33:43.120506","duration":0.34918650006875396,"load_avg_1min":3.6337890625,"mem_max_rss":54509568},"values":[0.00020958821687599993,0.00020886287529719993],"warmups":[[512,0.0002109452325385064]]},{"metadata":{"date":"2026-10-01 16:33:43.640808","duration":0.34697845787741244,"load_avg_1min":3.6337890625,"mem_max_rss":54558720},"values":[0.00020574251357174944,0.00020792431678273715],"warmups":[[512,0.00021029467779953848]]},{"metadata":{"date":"2026-10-01 16:33:44.166680","duration":0.3428519170265645,"load_avg_1min":3.6337890625,"mem_max_rss":54509568},"values":[0.00021043652350272168,0.0002082728669847711],"warmups":[[512,0.00021033512484791572]]},{"metadata":{"date":"2026-10-01 16:33:44.680133","duration":0.3466305839829147,"load_avg_1min":3.4228515625,"mem_max_rss":54444032},"values":[0.00020954296860509203,0.00020949576946804882],"warmups":[[512,0.00020963387123629218]]},{"metadata":{"date":"2026-10-01 16:33:45.201021","duration":0.34849387500435114,"load_avg_1min":3.4228515625,"mem_max_rss":54214656},"values":[0.00020953955072400277,0.00021248421080599655],"warmups":[[512,0.00020954272486051195]]}]},{"metadata":{"loops":4096,"name":"logs3[enqueue-100-rows]"},"runs":[{"metadata":{"calibrate_loops":4096,"date":"2026-10-01 16:33:45.913310","duration":0.5390088330022991,"load_avg_1min":3.4228515625,"mem_max_rss":54591488},"warmups":[[1,4.895799793303013e-05],[2,3.418745473027229e-05],[4,3.0135503038764e-05],[8,2.955738455057144e-05],[16,3.0455688829533756e-05],[32,3.09270981233567e-05],[64,3.1284500437323004e-05],[128,3.111588375759311e-05],[256,3.087483582930872e-05],[512,3.1231283173838165e-05],[1024,3.130285654151521e-05],[2048,3.132472750166926e-05],[4096,3.149635813315399e-05],[4096,3.134935496973412e-05],[4096,3.1386586670123506e-05]]},{"metadata":{"date":"2026-10-01 16:33:46.490783","duration":0.3982332921586931,"load_avg_1min":3.4228515625,"mem_max_rss":54444032},"values":[3.073358158189876e-05,3.091226193419061e-05],"warmups":[[4096,3.073799661024168e-05]]},{"metadata":{"date":"2026-10-01 16:33:47.066999","duration":0.411022667074576,"load_avg_1min":3.4228515625,"mem_max_rss":54738944},"values":[3.096644087463574e-05,3.0887512252775196e-05],"warmups":[[4096,3.107479858499573e-05]]},{"metadata":{"date":"2026-10-01 16:33:47.658963","duration":0.4100068327970803,"load_avg_1min":3.4228515625,"mem_max_rss":54558720},"values":[3.116254686119646e-05,3.103071094301413e-05],"warmups":[[4096,3.102487181649849e-05]]},{"metadata":{"date":"2026-10-01 16:33:48.241725","duration":0.4116921660024673,"load_avg_1min":3.4228515625,"mem_max_rss":54460416},"values":[3.111916112175095e-05,3.1081736324267695e-05],"warmups":[[4096,3.116294357141669e-05]]},{"metadata":{"date":"2026-10-01 16:33:48.823889","duration":0.40472033293917775,"load_avg_1min":3.4228515625,"mem_max_rss":54476800},"values":[3.10143737465296e-05,3.0762257836158824e-05],"warmups":[[4096,3.075356028148235e-05]]},{"metadata":{"date":"2026-10-01 16:33:49.408508","duration":0.41863066679798067,"load_avg_1min":3.4228515625,"mem_max_rss":54509568},"values":[3.1513041278685705e-05,3.171038821392358e-05],"warmups":[[4096,3.102163697121796e-05]]},{"metadata":{"date":"2026-10-01 16:33:50.001855","duration":0.4131154171191156,"load_avg_1min":3.228515625,"mem_max_rss":54214656},"values":[3.102931714238366e-05,3.184017941748607e-05],"warmups":[[4096,3.1015757087970997e-05]]},{"metadata":{"date":"2026-10-01 16:33:50.584215","duration":0.4149068749975413,"load_avg_1min":3.228515625,"mem_max_rss":54165504},"values":[3.149046824546531e-05,3.1332784203641495e-05],"warmups":[[4096,3.1187510046493117e-05]]},{"metadata":{"date":"2026-10-01 16:33:51.182815","duration":0.42037391709163785,"load_avg_1min":3.228515625,"mem_max_rss":55033856},"values":[3.12997944433846e-05,3.172418212216144e-05],"warmups":[[4096,3.16866558023321e-05]]},{"metadata":{"date":"2026-10-01 16:33:51.776360","duration":0.4158359579741955,"load_avg_1min":3.228515625,"mem_max_rss":54771712},"values":[3.1254933844593324e-05,3.2562520289047825e-05],"warmups":[[4096,3.1172882074770314e-05]]}]}],"metadata":{"cpu_count":12,"hostname":"A-Prasad","perf_version":"2.10.0","platform":"macOS-26.6.2-arm64-arm-64bit-Mach-O","python_cflags":"-fno-strict-overflow -Wsign-compare -Wunreachable-code -DNDEBUG -g -O3 -Wall -O3 -arch arm64 -mmacosx-version-min=11.0 -Wno-nullability-completeness -Wno-expansion-to-defined -Wno-undef-prefix -fPIC","python_compiler":"Clang 22.1.3","python_config_args":"'--build=aarch64-apple-darwin' '--host=aarch64-apple-darwin' '--prefix=/install' '--with-openssl=/tmp/tmps_46j10_/tools/deps' '--with-system-expat' '--with-system-libmpdec' '--without-ensurepip' '--with-tail-call-interp' 'MODULE_BUILDTYPE=static' '--enable-loadable-sqlite-extensions' '--enable-static-libpython-for-interpreter' '--enable-shared' '--with-mimalloc' '--enable-optimizations' '--enable-experimental-jit=yes-off' '--with-lto' '--with-build-python=/tmp/tmps_46j10_/tools/host/bin/python3.14' 'ac_cv_lib_intl_textdomain=no' 'ac_cv_func_ptsname_r=no' 'ac_cv_func_explicit_bzero=no' '--with-dbmliborder=ndbm' 'build_alias=aarch64-apple-darwin' 'host_alias=aarch64-apple-darwin' 'PKG_CONFIG=pkg-config --static --define-prefix' 'PKG_CONFIG_PATH=/tmp/tmps_46j10_/tools/deps/share/pkgconfig:/tmp/tmps_46j10_/tools/deps/lib/pkgconfig' 'CC=clang' 'CFLAGS=-O3 -arch arm64 -mmacosx-version-min=11.0 -Wno-nullability-completeness -Wno-expansion-to-defined -Wno-undef-prefix -fPIC -Werror=unguarded-availability-new' 'LDFLAGS=-arch arm64 -mmacosx-version-min=11.0 -Wl,-headerpad,40' 'CPPFLAGS=-O3 -arch arm64 -mmacosx-version-min=11.0 -Wno-nullability-completeness -Wno-expansion-to-defined -Wno-undef-prefix -fPIC -Werror=unguarded-availability-new' 'PROFILE_TASK=-m test --pgo -j 12' 'BOLT_COMMON_FLAGS=' 'BOLT_APPLY_FLAGS='","python_executable":"/Users/abhijeetprasad/.superset/worktrees/braintrust-sdk-python/abhi/recondite-flier/py/.venv/bin/python","python_implementation":"cpython","python_version":"3.14.7 (64-bit)","timer":"mach_absolute_time(), resolution: 41.7 ns","unit":"second"},"version":"1.0"} diff --git a/py/benchmarks/results/log_ingestion/environment.json b/py/benchmarks/results/log_ingestion/environment.json new file mode 100644 index 000000000..116bbe5d9 --- /dev/null +++ b/py/benchmarks/results/log_ingestion/environment.json @@ -0,0 +1,19 @@ +{ + "baseline_commit": "6537ec61d7809e9c06e8ba9d1955ca6b3eb53cac", + "measured_at": "2026-10-01", + "python": "3.14.7", + "platform": "macOS-26.6.2-arm64-arm-64bit-Mach-O", + "cpu_count": 12, + "baseline_concurrency": 12, + "branch_concurrency": 4, + "http_runs_per_scenario": 3, + "rows_per_run": 2000, + "batch_size": 100, + "service_latency_seconds": 0.005, + "throttle_duration_seconds": 1, + "retry_after_seconds": 1, + "continuous_producer_interval_seconds": 0.0005, + "tracemalloc_enabled": true, + "persistent_http": true, + "pyperf_mode": "--fast" +} diff --git a/py/src/braintrust/api/_ingestion.py b/py/src/braintrust/api/_ingestion.py new file mode 100644 index 000000000..c1e5d09ab --- /dev/null +++ b/py/src/braintrust/api/_ingestion.py @@ -0,0 +1,130 @@ +"""Internal single-attempt ingestion service and destination admission control.""" + +import contextlib +import hashlib +import threading +import time +import weakref +from typing import Any + +from requests.adapters import HTTPAdapter + +from ..env import BraintrustEnv +from ._routing import EndpointRouter, RequestTarget +from ._service import ResourceAPI +from ._transport import RetryRequestExceptionsAdapter, Transport +from .errors import BraintrustHTTPError +from .policies import DEFAULT_RETRYABLE_STATUSES, RetryMode, RetryPolicy + + +class IngestionDeferred(Exception): + """Admission changed after the writer scheduled an attempt.""" + + +class IngestionDestination: + def __init__(self, concurrency: int): + self.lock = threading.Lock() + self.concurrency = concurrency + self.active = 0 + self.cooldown_until = 0.0 + self.next_start = 0.0 + self.recovery_starts = 0 + + def delay(self) -> float: + with self.lock: + return max(0.0, self.cooldown_until - time.monotonic(), self.next_start - time.monotonic()) + + def acquire(self) -> None: + with self.lock: + now = time.monotonic() + if self.active >= self.concurrency or now < max(self.cooldown_until, self.next_start): + raise IngestionDeferred() + self.active += 1 + if self.recovery_starts: + self.next_start = now + 0.05 + self.recovery_starts -= 1 + + def release(self, retry_after: float | None = None) -> None: + with self.lock: + if retry_after is not None: + self.cooldown_until = max(self.cooldown_until, time.monotonic() + retry_after) + self.recovery_starts = self.concurrency + self.active -= 1 + + +_destinations: weakref.WeakValueDictionary[tuple[str, str], IngestionDestination] = weakref.WeakValueDictionary() +_destinations_lock = threading.Lock() + + +class LogIngestionAPI(ResourceAPI): + """Dedicated pools, existing router/auth, and no HTTP retries underneath the writer.""" + + def __init__(self, router: EndpointRouter, api_key: str, *, concurrency: int, adapter: HTTPAdapter | None = None): + if concurrency < 1: + raise ValueError("Log concurrency must be positive") + self._validate_adapter(adapter) + timeout = BraintrustEnv.HTTP_TIMEOUT.get(60.0) + super().__init__( + Transport(adapter=adapter, request_timeout=timeout, persist_cookies=False, pool_maxsize=concurrency), + router, + api_key, + ) + self.storage = Transport( + adapter=adapter, request_timeout=timeout, persist_cookies=False, pool_maxsize=concurrency + ) + key = (router.resolve(RequestTarget.API, "/logs3"), hashlib.sha256(api_key.encode()).hexdigest()) + with _destinations_lock: + destination = _destinations.get(key) + if destination is None: + destination = IngestionDestination(concurrency) + _destinations[key] = destination + self.destination = destination + + @property + def router(self) -> EndpointRouter: + return self._router + + @property + def transport(self) -> Transport: + return self._transport + + @staticmethod + def _validate_adapter(adapter: HTTPAdapter | None) -> None: + if adapter is not None and ( + adapter.max_retries.total not in (0, False) + or isinstance(adapter, RetryRequestExceptionsAdapter) + and adapter.base_num_retries > 0 + ): + raise ValueError("Log ingestion requires a single-attempt HTTP adapter; the writer owns retries") + + def request(self, method: str, path: str, **kwargs: Any): + return self._request(RequestTarget.API, method, path, retry_mode=RetryMode.LOG_INGESTION, **kwargs) + + def request_json(self, method: str, path: str, **kwargs: Any): + return self._request_json(RequestTarget.API, method, path, retry_mode=RetryMode.LOG_INGESTION, **kwargs) + + @contextlib.contextmanager + def attempt(self, *, retry_after_on_429: float = 1.0): + """Admit one writer attempt and publish any server cooldown before releasing it.""" + self.destination.acquire() + retry_after = None + try: + yield + except BraintrustHTTPError as error: + if error.status_code in DEFAULT_RETRYABLE_STATUSES: + retry_after = error.retry_after + if retry_after is None and error.status_code == 429: + retry_after = retry_after_on_429 + raise + finally: + self.destination.release(retry_after) + + def version(self, policy: RetryPolicy | None = None): + with self.attempt(): + return self._request_json( + RequestTarget.API, "GET", "/version", retry_mode=RetryMode.LOG_INGESTION, retry_policy=policy + ) + + def close(self) -> None: + self._transport.close() + self.storage.close() diff --git a/py/src/braintrust/api/_test_server.py b/py/src/braintrust/api/_test_server.py index 2d577b543..b6da7fa7f 100644 --- a/py/src/braintrust/api/_test_server.py +++ b/py/src/braintrust/api/_test_server.py @@ -2,16 +2,19 @@ import contextlib import http.server +import socket import socketserver import threading import time @contextlib.contextmanager -def scripted_server(script): +def scripted_server(script, *, persistent=False): """Run a local server driven by sequential actions or a request callback.""" class ScriptedHandler(http.server.BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" if persistent else "HTTP/1.0" + connections = set() request_count = 0 requests = [] @@ -24,6 +27,9 @@ def do_GET(self): def do_POST(self): self._handle() + def do_PUT(self): + self._handle() + def do_PATCH(self): self._handle() @@ -31,6 +37,7 @@ def do_DELETE(self): self._handle() def _handle(self): + type(self).connections.add(self.client_address) request_number = type(self).request_count type(self).request_count += 1 content_length = int(self.headers.get("Content-Length", "0")) @@ -43,6 +50,8 @@ def _handle(self): ) if action == "close": + self.close_connection = True + self.connection.shutdown(socket.SHUT_RDWR) self.connection.close() return diff --git a/py/src/braintrust/api/_transport.py b/py/src/braintrust/api/_transport.py index 509dabd38..44cf503d7 100644 --- a/py/src/braintrust/api/_transport.py +++ b/py/src/braintrust/api/_transport.py @@ -228,6 +228,7 @@ def __init__( enable_sdk_retries: bool | None = None, request_timeout: float | None = None, persist_cookies: bool = True, + pool_maxsize: int | None = None, sleep: Callable[[float], None] = time.sleep, monotonic: Callable[[], float] = time.monotonic, wall_clock: Callable[[], float] = time.time, @@ -242,9 +243,26 @@ def __init__( if request_timeout is not None and request_timeout <= 0: raise ValueError("request_timeout must be positive") self._request_timeout = request_timeout + if pool_maxsize is not None and pool_maxsize < 1: + raise ValueError("pool_maxsize must be positive") if adapter is not None: self.session.mount("http://", adapter) self.session.mount("https://", adapter) + elif pool_maxsize is not None and self._owns_session: + pooled_adapter = HTTPAdapter( + # requests does not expose urllib3's pool_timeout. With + # pool_block=True a request can therefore wait forever for a + # slot after urllib3 clears pools during interpreter shutdown. + pool_connections=pool_maxsize, + pool_maxsize=pool_maxsize, + pool_block=False, + max_retries=0, + ) + replaced_adapters = set(self.session.adapters.values()) + self.session.mount("http://", pooled_adapter) + self.session.mount("https://", pooled_adapter) + for replaced_adapter in replaced_adapters: + replaced_adapter.close() self._sleep = sleep self._monotonic = monotonic self._wall_clock = wall_clock diff --git a/py/src/braintrust/env.py b/py/src/braintrust/env.py index f08afb668..ab4555541 100644 --- a/py/src/braintrust/env.py +++ b/py/src/braintrust/env.py @@ -217,6 +217,8 @@ class BraintrustEnv: API_URL = EnvVar("BRAINTRUST_API_URL", EnvParser.STRING) PROXY_URL = EnvVar("BRAINTRUST_PROXY_URL", EnvParser.STRING) HTTP_TIMEOUT = EnvVar("BRAINTRUST_HTTP_TIMEOUT", EnvParser.FLOAT) + LOG_MAX_CONCURRENCY = EnvVar("BRAINTRUST_LOG_MAX_CONCURRENCY", EnvParser.INT) + LOG_FLUSH_TIMEOUT = EnvVar("BRAINTRUST_LOG_FLUSH_TIMEOUT", EnvParser.FLOAT) SYNC_FLUSH = EnvVar("BRAINTRUST_SYNC_FLUSH", EnvParser.BOOL) MAX_REQUEST_SIZE = EnvVar("BRAINTRUST_MAX_REQUEST_SIZE", EnvParser.INT) DEFAULT_BATCH_SIZE = EnvVar("BRAINTRUST_DEFAULT_BATCH_SIZE", EnvParser.INT) diff --git a/py/src/braintrust/framework.py b/py/src/braintrust/framework.py index b352252e9..19bcad16a 100644 --- a/py/src/braintrust/framework.py +++ b/py/src/braintrust/framework.py @@ -808,9 +808,9 @@ async def run_to_completion(): return ret finally: if experiment: - experiment.flush() + experiment.state.flush_best_effort() elif state is not None: - state.flush() + state.flush_best_effort() return run_to_completion @@ -1709,13 +1709,15 @@ def report_progress(event: TaskProgressEvent): async def ensure_spans_flushed(): # Flush native Braintrust spans if experiment: - await asyncio.get_event_loop().run_in_executor(None, lambda: experiment.state.flush()) + await asyncio.get_event_loop().run_in_executor( + None, lambda: experiment.state.flush_best_effort() + ) elif state: - await asyncio.get_event_loop().run_in_executor(None, lambda: state.flush()) + await asyncio.get_event_loop().run_in_executor(None, lambda: state.flush_best_effort()) else: - from braintrust.logger import flush as flush_logger + from braintrust.logger import _state - await asyncio.get_event_loop().run_in_executor(None, flush_logger) + await asyncio.get_event_loop().run_in_executor(None, lambda: _state.flush_best_effort()) # Also flush OTEL spans if registered if state: diff --git a/py/src/braintrust/logger.py b/py/src/braintrust/logger.py index 2e7a7adf4..4394a9e3b 100644 --- a/py/src/braintrust/logger.py +++ b/py/src/braintrust/logger.py @@ -39,12 +39,13 @@ from urllib.parse import quote, urlencode import chevron -import exceptiongroup from braintrust.functions.stream import BraintrustStream from requests import exceptions as requests_exceptions from requests.adapters import HTTPAdapter from . import context, id_gen +from .api._ingestion import IngestionDeferred as _IngestionDeferred +from .api._ingestion import LogIngestionAPI as _LogIngestionAPI from .api._routing import normalize_proxy_url from .api._transport import HTTPConnection from .api._transport import RetryRequestExceptionsAdapter as RetryRequestExceptionsAdapter @@ -57,6 +58,8 @@ BraintrustTransportError, ) from .api.policies import DEFAULT_RETRYABLE_STATUSES, is_retryable_request_exception +from .api.policies import RetryMode as _RetryMode +from .api.policies import RetryPolicy as _RetryPolicy from .bt_json import bt_dumps, bt_safe_deep_copy from .db_fields import ( AUDIT_METADATA_FIELD, @@ -111,6 +114,7 @@ from .util import ( GLOBAL_PROJECT, LazyValue, + _reset_lazy_value_locks_after_fork, add_azure_blob_headers, bt_iscoroutinefunction, coalesce, @@ -522,7 +526,7 @@ class _LoaderLoginOptions: cache_namespace: str -_LoaderResource = TypeVar("_LoaderResource", bound=BraintrustClient) +_LoaderResource = TypeVar("_LoaderResource", bound=BraintrustClient | _LogIngestionAPI) class _LoaderLoginEntry(Generic[_LoaderResource]): @@ -557,6 +561,14 @@ def release(self) -> None: if should_close: self._close() + @contextlib.contextmanager + def lease(self) -> Iterator[_LoaderResource]: + resource = self.acquire() + try: + yield resource + finally: + self.release() + def evict(self) -> None: with self._lock: self._evicted = True @@ -600,20 +612,18 @@ def __init__(self): self._context_manager = None self._context_manager_lock = threading.Lock() self._client_lock = threading.RLock() + self._login_operation_lock = threading.RLock() + self._ingestion_cache: LRUCache[str, _LoaderLoginEntry[_LogIngestionAPI]] = LRUCache( + max_size=1, + on_remove=self._evict_loader_login_entry, + ) - def default_get_api_conn(): - self.login() - return self.api_conn() - - # Any time we re-log in, we directly update the api_conn inside the - # logger. This is preferable to replacing the whole logger, which would - # create the possibility of multiple loggers floating around. + # Keep one writer across logins. Prepared batches lease their original + # transport; later batches resolve current authentication and routing. # # We lazily-initialize the logger so that it does any initialization # (including reading env variables) upon the first actual usage. - self._global_bg_logger = LazyValue( - lambda: _HTTPBackgroundLogger(LazyValue(default_get_api_conn, use_mutex=True)), use_mutex=True - ) + self._global_bg_logger = LazyValue(lambda: _HTTPBackgroundLogger(self.ingestion_service), use_mutex=True) self._id_generator = None @@ -654,6 +664,7 @@ def default_get_api_conn(): self._otel_flush_callback: Any | None = None def reset_login_info(self): + self._ingestion_cache.clear() self._loader_api_client_cache.clear() self.app_url: str | None = None @@ -736,6 +747,7 @@ async def flush_otel(self) -> None: def copy_state(self, other: "BraintrustState"): """Copy login information from another BraintrustState instance.""" self._loader_api_client_cache.clear() + self._ingestion_cache.clear() self.__dict__.update( { k: v @@ -753,7 +765,9 @@ def copy_state(self, other: "BraintrustState"): "_last_otel_setting", "_context_manager_lock", "_client_lock", + "_login_operation_lock", "_loader_api_client_cache", + "_ingestion_cache", ) } ) @@ -765,29 +779,55 @@ def login( org_name: str | None = None, force_login: bool = False, ) -> None: - with self._client_lock: - if not force_login and self.logged_in: - # We have already logged in. If any provided login inputs disagree - # with our existing settings, raise an Exception warning the user to - # try again with `force_login=True`. - def check_updated_param(varname, arg, orig): - if arg is not None and orig is not None and arg != orig: - raise Exception( - f"Re-logging in with different {varname} ({arg}) than original ({orig}). To force re-login, pass `force_login=True`" - ) - - sanitized_api_key = HTTPConnection.sanitize_token(api_key) if api_key else None - check_updated_param("app_url", app_url, self.app_url) - check_updated_param("api_key", sanitized_api_key, self.login_token) - check_updated_param("org_name", org_name, self.org_name) - return - + with self._login_operation_lock: + with self._client_lock: + if not force_login and self.logged_in: + # We have already logged in. If any provided login inputs disagree + # with our existing settings, raise an Exception warning the user to + # try again with `force_login=True`. + def check_updated_param(varname, arg, orig): + if arg is not None and orig is not None and arg != orig: + raise Exception( + f"Re-logging in with different {varname} ({arg}) than original ({orig}). To force re-login, pass `force_login=True`" + ) + + sanitized_api_key = HTTPConnection.sanitize_token(api_key) if api_key else None + check_updated_param("app_url", app_url, self.app_url) + check_updated_param("api_key", sanitized_api_key, self.login_token) + check_updated_param("org_name", org_name, self.org_name) + return + fork_login = getattr(self, "_fork_login", None) or (None, None, None) state = login_to_state( - app_url=app_url, - api_key=api_key, - org_name=org_name, + # Do not let records queued under the old identity be prepared + # after the state swap and sent with the new user's key. + # Seal them into a prepared batch before replacing credentials. + app_url=app_url or fork_login[0], + api_key=api_key or fork_login[1], + org_name=org_name or fork_login[2], ) - self.copy_state(state) + self._copy_login_state_after_sealing(state) + + def _copy_login_state_after_sealing(self, new_state: "BraintrustState") -> None: + if not self.logged_in or not self._global_bg_logger.has_succeeded: + with self._client_lock: + self.copy_state(new_state) + self._fork_login = None + return + writer = self._global_bg_logger.get() + while True: + while not writer.flush_lock.acquire(timeout=1.0): + pass + try: + with writer.queue._mutex: + if not writer.queue._queue: + with self._client_lock: + self.copy_state(new_state) + self._fork_login = None + return + deadline = time.monotonic() + writer.flush_timeout + writer._prepare_batches(writer.default_batch_size, deadline, append=bool(writer._retained_count)) + finally: + writer.flush_lock.release() def api_client(self) -> BraintrustOpenApiClient: """Return the lazily bootstrapped OpenAPI client.""" @@ -800,6 +840,32 @@ def api_client(self) -> BraintrustOpenApiClient: raise RuntimeError("Braintrust API client was not initialized during login") return self._client.openapi + @contextlib.contextmanager + def ingestion_service(self) -> "Iterator[_LogIngestionAPI]": + """Lease a pool for the current credentials; retained batches keep their lease.""" + self.api_client() + with contextlib.ExitStack() as stack: + with self._client_lock: + client = self._client + if client is None: + raise RuntimeError("Braintrust API client was not initialized during login") + router = dataclasses.replace(client.router) + api_key = client.api_key + key = bt_dumps([router.api_url, hashlib.sha256(api_key.encode()).hexdigest()]) + service = stack.enter_context( + self._cached_loader_resource( + self._ingestion_cache, + key, + lambda: _LogIngestionAPI( + router, + api_key, + concurrency=BraintrustEnv.LOG_MAX_CONCURRENCY.get(4), + adapter=_http_adapter, + ), + ) + ) + yield service + def app_conn(self): if not self._app_conn: if not self.app_url: @@ -851,11 +917,8 @@ def _cached_loader_resource( entry = _LoaderLoginEntry(factory) cache.set(cache_key, entry) - resource = entry.acquire() - try: + with entry.lease() as resource: yield resource - finally: - entry.release() @contextlib.contextmanager def loader_api_client(self, options: "_LoaderLoginOptions") -> "Iterator[BraintrustOpenApiClient]": @@ -882,6 +945,13 @@ def login_replace_api_conn(self, api_conn: "HTTPConnection"): def flush(self): self._global_bg_logger.get().flush() + def flush_best_effort(self) -> None: + """Flush during a user operation without turning telemetry failure into app failure.""" + try: + self.flush() + except BraintrustLogFlushError as error: + _logger.warning("Failed to flush Braintrust logs; records remain queued: %s", error) + def enforce_queue_size_limit(self, enforce: bool) -> None: """ Set queue size limit enforcement for the global background logger. @@ -906,6 +976,9 @@ def set_http_adapter(adapter: HTTPAdapter) -> None: Braintrust uses the `requests` library, so the adapter should be an instance of `requests.adapters.HTTPAdapter`. Alternatively, consider sub-classing our `RetryRequestExceptionsAdapter` to get automatic retries on network-related exceptions. + Log ingestion requires adapters with retries disabled because the writer owns + its retry budget. Other legacy callers retain their configured adapter policies. + :param adapter: The adapter to use. """ @@ -926,6 +999,8 @@ def set_http_adapter(adapter: HTTPAdapter) -> None: # adapter. Eviction closes them once any active requests release their lease; # subsequent loads recreate them with the new global adapter. _state._loader_api_client_cache.clear() + with _state._client_lock: + _state._ingestion_cache.clear() # Sometimes we'd like to launch network requests concurrently. We provide a @@ -1068,24 +1143,71 @@ def pop(self): BACKGROUND_LOGGER_BASE_SLEEP_TIME_S = 1.0 +class BraintrustLogFlushError(RuntimeError): + """Explicit delivery failure. Undelivered rows remain in the writer.""" + + def __init__(self, pending_count: int, errors: Sequence[Exception]): + self.pending_count = pending_count + self.errors = tuple(errors) + super().__init__( + f"Log flush failed with {pending_count} records still pending: " + "; ".join(str(e) for e in errors) + ) + + +class _ExpiredOverflowUpload(RuntimeError): + """A signed URL expired; request a fresh URL before retrying the batch.""" + + +@dataclasses.dataclass(eq=False) +class _PreparedLogBatch: + item_count: int + payload: bytes + overflow_rows: list[dict[str, Any]] | None + service: _LogIngestionAPI | None = None + wave: int = 0 + overflow_upload: dict[str, Any] | None = None + overflow_reference: bytes | None = None + uploaded: bool = False + attempts: int = 0 + ready_at: float = 0.0 + error: Exception | None = None + permanent: bool = False + dumped: bool = False + + def exhausted(self, num_tries: int) -> bool: + return self.permanent or self.attempts >= num_tries + + # 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 # the backend in a deterministic order. class _HTTPBackgroundLogger: - def __init__(self, api_conn: LazyValue[HTTPConnection]): - self.api_conn = api_conn + def __init__(self, service_source: "Callable[[], contextlib.AbstractContextManager[_LogIngestionAPI]]"): + self._service_source = service_source + self._replaced_entry: _LoaderLoginEntry[_LogIngestionAPI] | None = None + self._pending: list[_PreparedLogBatch] = [] + self._delivery_errors: list[Exception] = [] + self._wave_counter = 0 + self._retained_count = 0 + self._pending_attachments: list[BaseAttachment] = [] + self._pending_service: _LogIngestionAPI | None = None + self._pending_context = contextlib.ExitStack() self._export_customizers: tuple[SpanCustomizer, ...] = () self.outfile = sys.stderr self.flush_lock = threading.RLock() self._max_request_size_override: int | None = None self._max_request_size_result: dict[str, Any] | None = None - self._max_request_size_lock = threading.Lock() + self._limit_service: _LogIngestionAPI | None = None self.sync_flush = BraintrustEnv.SYNC_FLUSH.get(False) 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 + self.num_tries = max(1, BraintrustEnv.NUM_RETRIES.get(2) + 1) + self.max_concurrency = BraintrustEnv.LOG_MAX_CONCURRENCY.get(4) + self.flush_timeout = BraintrustEnv.LOG_FLUSH_TIMEOUT.get(60.0) + if self.max_concurrency < 1 or self.flush_timeout <= 0: + raise ValueError("Log concurrency and flush timeout must be positive") queue_maxsize = BraintrustEnv.QUEUE_SIZE.get(None) self.queue_maxsize = DEFAULT_QUEUE_SIZE if queue_maxsize is None else queue_maxsize self.queue_drop_logging_period = BraintrustEnv.QUEUE_DROP_LOGGING_PERIOD.get(60.0) @@ -1139,20 +1261,57 @@ def _start(self): self.thread.start() self.started = True + def reset_after_fork(self) -> None: + """Start a clean child writer without replaying the parent's in-flight rows.""" + self.flush_lock = threading.RLock() + self.start_thread_lock = threading.RLock() + self.queue.reset_after_fork() + self.thread = threading.Thread(target=self._publisher, daemon=True) + self.started = False + self._pending = [] + self._retained_count = 0 + self._pending_attachments = [] + self._pending_service = None + self._pending_context = contextlib.ExitStack() + self._delivery_errors = [] + self._wave_counter = 0 + self._replaced_entry = None + def _finalize(self): self.logger.debug("Flushing final log events...") - self.flush() + try: + self.flush() + except Exception as exc: + try: + print(str(exc), file=self.outfile) + except ValueError: + pass # stderr may already be closed during interpreter shutdown. + finally: + if self._replaced_entry is not None: + self._replaced_entry.evict() def _publisher(self): while True: # Wait for some data on the queue before trying to flush. - self.queue.wait_for_items() + self.queue.wait_for_items(timeout=1.0 if self._retained_count else None) while self.sync_flush: time.sleep(0.1) + service = self._pending_service + if service is not None and (delay := service.destination.delay()) > 0: + time.sleep(min(delay, 1.0)) + continue + try: - self.flush() + # Transient failures remain eligible for background retries after + # the per-flush attempt budget expires. A permanent response is + # reported and discarded by _deliver_batches so it cannot wedge + # later rows indefinitely. + for batch in self._pending: + if not batch.permanent and batch.exhausted(self.num_tries): + batch.attempts = 0 + self.flush(_background=True) except: # Print exception but don't worry if stderr is closed because the process is shutting down. try: @@ -1162,99 +1321,275 @@ def _publisher(self): pass else: raise + # Restored, unprepared records leave the queue event set. Back off + # when version admission or construction fails rather than spinning. + time.sleep(1.0) - def _get_max_request_size(self) -> dict[str, Any]: + def _get_max_request_size(self, service: _LogIngestionAPI, deadline: float) -> dict[str, Any]: if self._max_request_size_result is not None: return self._max_request_size_result - with self._max_request_size_lock: - if self._max_request_size_result is not None: - return self._max_request_size_result - server_limit: int | None = None - try: - conn = self.api_conn.get() - info = conn.get_json("version") - limit = info.get("logs3_payload_max_bytes") - if isinstance(limit, (int, float)) and int(limit) > 0: - server_limit = int(limit) - except Exception as e: - print(f"Failed to fetch version info for payload limit: {e}", file=self.outfile) - valid_server_limit = server_limit if server_limit is not None and server_limit > 0 else None - can_use_overflow = valid_server_limit is not None - max_request_size = DEFAULT_MAX_REQUEST_SIZE - if self._max_request_size_override is not None: - max_request_size = ( - min(self._max_request_size_override, valid_server_limit) - if valid_server_limit is not None - else self._max_request_size_override - ) - elif valid_server_limit is not None: - max_request_size = valid_server_limit - self._max_request_size_result = { - "max_request_size": max_request_size, - "can_use_overflow": can_use_overflow, - } - return self._max_request_size_result + server_limit: int | None = None + try: + info = service.version(self._ingestion_policy(deadline)) + limit = info.get("logs3_payload_max_bytes") + if isinstance(limit, (int, float)) and int(limit) > 0: + server_limit = int(limit) + except _IngestionDeferred: + raise + except Exception as e: + print(f"Failed to fetch version info for payload limit: {e}", file=self.outfile) + max_request_size = DEFAULT_MAX_REQUEST_SIZE + if self._max_request_size_override is not None: + max_request_size = ( + min(self._max_request_size_override, server_limit) + if server_limit is not None + else self._max_request_size_override + ) + elif server_limit is not None: + max_request_size = server_limit + self._max_request_size_result = { + "max_request_size": max_request_size, + "can_use_overflow": server_limit is not None, + } + return self._max_request_size_result - def flush(self, batch_size: int | None = None): - if batch_size is None: - batch_size = self.default_batch_size + @staticmethod + def _ingestion_policy(deadline: float) -> _RetryPolicy: + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError("Log flush deadline exceeded") + return _RetryPolicy( + max_attempts=1, max_elapsed_time=None, timeout=min(BraintrustEnv.HTTP_TIMEOUT.get(60.0), remaining) + ) - # We cannot have multiple threads flushing in parallel, because the - # order of published elements would be undefined. - with self.flush_lock: - # Drain the queue. - wrapped_items = self.queue.drain_all() + @property + def pending_count(self) -> int: + return self._retained_count + self.queue.size() + + def flush(self, batch_size: int | None = None, *, timeout: float | None = None, _background: bool = False): + """Deliver retained and queued rows or raise; cooldowns never consume worker threads.""" + timeout = self.flush_timeout if timeout is None else timeout + if timeout <= 0: + raise ValueError("flush timeout must be positive") + deadline = time.monotonic() + timeout + if not self.flush_lock.acquire(timeout=timeout): + raise BraintrustLogFlushError( + self.pending_count, [TimeoutError("Timed out waiting for another log delivery")] + ) + try: + if not _background: + for batch in self._pending: + if not batch.permanent: + batch.attempts = 0 + # Flush one queued wave after any retained wave. Concurrent producers + # belong to subsequent publisher cycles, so a healthy flush cannot starve. + prepare_queued = True + while self._retained_count or prepare_queued: + if not self._retained_count and not self.queue.size(): + break + if time.monotonic() >= deadline: + raise BraintrustLogFlushError(self.pending_count, [TimeoutError("Log flush deadline exceeded")]) + if not self._retained_count: + prepare_queued = False + self._prepare_batches(self.default_batch_size if batch_size is None else batch_size, deadline) + self._deliver_batches(deadline) + if self._pending: + if _background: + return # Keep retained work parked; only explicit flush promises delivery. + errors = [batch.error for batch in self._pending if batch.error is not None] + raise BraintrustLogFlushError( + self.pending_count, errors or [TimeoutError("Log destination is cooling down")] + ) + # Attempt every attachment, retaining only failures for the next flush. + attachment_errors: list[Exception] = [] + for attachment in list(self._pending_attachments): + try: + result = attachment.upload() + if result["upload_status"] == "error": + raise RuntimeError(result.get("error_message")) + except Exception as error: + attachment_errors.append(error) + else: + self._pending_attachments.remove(attachment) + if attachment_errors: + raise BraintrustLogFlushError(self.pending_count, attachment_errors) + self.queue.release(self._retained_count) + self._retained_count = 0 + self._pending_service = None + self._pending_context.close() + if self._delivery_errors and not _background: + errors = self._delivery_errors[:] + self._delivery_errors.clear() + raise BraintrustLogFlushError(self.pending_count, errors) + if self._delivery_errors and not _background: + errors = self._delivery_errors[:] + self._delivery_errors.clear() + raise BraintrustLogFlushError(self.pending_count, errors) + except _IngestionDeferred as error: + raise BraintrustLogFlushError( + self.pending_count, [TimeoutError("Log destination is cooling down")] + ) from error + finally: + self.flush_lock.release() + def _prepare_batches(self, batch_size: int, deadline: float, *, append: bool = False) -> None: + wrapped_items = self.queue.drain_all(reserve=True) + service_context = contextlib.ExitStack() + try: all_items, attachments = self._unwrap_lazy_values(wrapped_items) - if len(all_items) == 0: + if not all_items: + self.queue.release(len(wrapped_items)) + service_context.close() return - - # Construct batches of records to flush in parallel. - all_items_with_meta = [stringify_with_overflow_meta(item) for item in all_items] - max_request_size_result = self._get_max_request_size() + service = service_context.enter_context(self._service_source()) + if self._limit_service is not None and service is not self._limit_service: + self._max_request_size_result = None + self._limit_service = service + try: + limit = self._get_max_request_size(service, deadline) + except _IngestionDeferred: + limit = { + "max_request_size": self._max_request_size_override or DEFAULT_MAX_REQUEST_SIZE, + "can_use_overflow": False, + } + items = [stringify_with_overflow_meta(item) for item in all_items] batches = batch_items( - items=all_items_with_meta, + items, 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=limit["max_request_size"] // 2, + get_byte_size=lambda item: item.overflow_meta.byte_size, ) + prepared = [self._prepare_batch(batch, limit) for batch in batches] + for batch in prepared: + batch.service = service + batch.wave = self._wave_counter + self._wave_counter += 1 + if append: + self._pending.extend(prepared) + self._pending_attachments.extend(attachments) + else: + self._pending = prepared + self._pending_attachments = attachments + self._pending_service = self._pending[0].service if self._pending else service + self._pending_context.callback(service_context.close) + self.queue.release(len(wrapped_items) - len(items)) + self._retained_count += len(items) + except BaseException: + service_context.close() + self.queue.restore(wrapped_items) + raise - post_promises = [] - try: - post_promises = [ - HTTP_REQUEST_THREAD_POOL.submit(self._submit_logs_request, batch, max_request_size_result) - for batch in batches - ] - except RuntimeError: - # If the thread pool has shut down, e.g. because the process - # is terminating, run the requests the old fashioned way. - for batch in batches: - self._submit_logs_request(batch, max_request_size_result) - - concurrent.futures.wait(post_promises) - # Raise any exceptions from the promises as one group. - post_promise_exceptions = [e for e in (f.exception() for f in post_promises) if e is not None] - if post_promise_exceptions: - raise exceptiongroup.BaseExceptionGroup( - f"Encountered the following errors while logging:", post_promise_exceptions - ) + def _prepare_batch(self, items: Sequence[LogItemWithMeta], limit: dict[str, Any]) -> _PreparedLogBatch: + payload = construct_logs3_data(items).encode("utf-8") + overflow_rows = None + if limit["can_use_overflow"] and len(payload) > limit["max_request_size"]: + overflow_rows = [ + { + "object_ids": item.overflow_meta.object_ids, + "has_comment": item.overflow_meta.has_comment, + "is_delete": item.overflow_meta.is_delete, + "input_row": {"byte_size": item.overflow_meta.byte_size}, + } + for item in items + ] + if self.all_publish_payloads_dir: + self._write_payload_to_dir(payload_dir=self.all_publish_payloads_dir, payload=payload.decode("utf-8")) + return _PreparedLogBatch(len(items), payload, overflow_rows) - attachment_errors: list[Exception] = [] - for attachment in attachments: - try: - result = attachment.upload() - if result["upload_status"] == "error": - raise RuntimeError(result.get("error_message")) - except Exception as e: - attachment_errors.append(e) - - if len(attachment_errors) == 1: - raise attachment_errors[0] - elif len(attachment_errors) > 1: - raise exceptiongroup.ExceptionGroup( - "Encountered errors while uploading attachments", - attachment_errors, + def _deliver_batches(self, deadline: float) -> None: + if not self._pending: + return + active: dict[concurrent.futures.Future, _PreparedLogBatch] = {} + # A dedicated executor bounds submitted work and isolates ingestion from attachments/API calls. + with concurrent.futures.ThreadPoolExecutor(max_workers=self.max_concurrency) as executor: + while self._pending or active: + now = time.monotonic() + wave = min((batch.wave for batch in self._pending), default=None) + wave_batches = [batch for batch in self._pending if batch.wave == wave] + service = wave_batches[0].service if wave_batches else self._pending_service + if service is None: + return + if now < deadline: + for batch in wave_batches: + if len(active) >= self.max_concurrency or service.destination.delay() > 0: + break + if batch in active.values() or batch.exhausted(self.num_tries) or batch.ready_at > now: + continue + try: + future = executor.submit(self._submit_logs_request, batch, batch.service, deadline) + except RuntimeError: + # Python shuts down executors before atexit callbacks. + # Preserve serial delivery through the same scheduler. + future = concurrent.futures.Future() + try: + self._submit_logs_request(batch, batch.service, deadline) + except Exception as error: + future.set_exception(error) + else: + future.set_result(None) + active[future] = batch + if not active: + eligible = [batch for batch in wave_batches if not batch.exhausted(self.num_tries)] + if not eligible or now >= deadline: + return + delay = max(service.destination.delay(), min(max(0.0, batch.ready_at - now) for batch in eligible)) + if delay >= deadline - now: + return + time.sleep(min(max(delay, 0.001), deadline - now)) + continue + done, _ = concurrent.futures.wait( + active, timeout=max(0.0, deadline - now), return_when=concurrent.futures.FIRST_COMPLETED ) + if not done: + # Requests have finite timeouts. Account for their results before handing off retention. + done, _ = concurrent.futures.wait(active) + for future in done: + batch = active.pop(future) + error = future.exception() + if isinstance(error, _IngestionDeferred): + continue + batch.attempts += 1 + if error is None: + self._pending.remove(batch) + self._pending_service = self._pending[0].service if self._pending else None + if not self._pending_attachments: + self.queue.release(batch.item_count) + self._retained_count -= batch.item_count + continue + batch.error = error + batch.permanent = not self._retryable_ingestion_error(error) + batch.ready_at = time.monotonic() + BACKGROUND_LOGGER_BASE_SLEEP_TIME_S * ( + 2 ** (batch.attempts - 1) + ) + if batch.exhausted(self.num_tries): + if self.failed_publish_payloads_dir and not batch.dumped: + self._write_payload_to_dir( + payload_dir=self.failed_publish_payloads_dir, payload=batch.payload.decode("utf-8") + ) + batch.dumped = True + print(f"Log batch failed after {batch.attempts} attempts: {error}", file=self.outfile) + if batch.permanent: + self._pending.remove(batch) + self._pending_service = self._pending[0].service if self._pending else None + self._delivery_errors.append(error) + if not self._pending_attachments: + self.queue.release(batch.item_count) + self._retained_count -= batch.item_count + elif batch.exhausted(self.num_tries): + # Let the background publisher start a fresh bounded + # retry cycle while keeping the original prepared bytes. + batch.ready_at = max(batch.ready_at, time.monotonic() + 1.0) + + @staticmethod + def _retryable_ingestion_error(error: Exception) -> bool: + if isinstance(error, TimeoutError): + return True + if isinstance(error, BraintrustHTTPError): + return error.status_code in DEFAULT_RETRYABLE_STATUSES + if isinstance(error, _ExpiredOverflowUpload): + return True + cause = error.__cause__ if isinstance(error, BraintrustTransportError) else error + return isinstance(cause, requests_exceptions.RequestException) and is_retryable_request_exception(cause) def _unwrap_lazy_values( self, wrapped_items: Sequence[_LazyLogRecord] @@ -1278,159 +1613,98 @@ def _unwrap_lazy_values( return merged_items, attachments except Exception as e: - errmsg = "Encountered error when constructing records to flush" - is_retrying = i + 1 < self.num_tries - if is_retrying: - errmsg += ". Retrying" + if i + 1 == self.num_tries: + raise Exception("Encountered error when constructing records to flush") from e + print("Encountered error when constructing records to flush. Retrying", file=self.outfile) + traceback.print_exc(file=self.outfile) + sleep_time_s = BACKGROUND_LOGGER_BASE_SLEEP_TIME_S * (2**i) + print(f"Sleeping for {sleep_time_s}s", file=self.outfile) + time.sleep(sleep_time_s) - if not is_retrying and self.sync_flush: - raise Exception(errmsg) from e - else: - print(errmsg, file=self.outfile) - traceback.print_exc(file=self.outfile) - if is_retrying: - sleep_time_s = BACKGROUND_LOGGER_BASE_SLEEP_TIME_S * (2**i) - print(f"Sleeping for {sleep_time_s}s", file=self.outfile) - time.sleep(sleep_time_s) - - print( - f"Failed to construct log records to flush after {self.num_tries} attempts. Dropping batch", - file=self.outfile, - ) - return [], [] + raise AssertionError("record construction retry loop exited unexpectedly") def _request_logs3_overflow_upload( - self, conn: HTTPConnection, payload_size_bytes: int, rows: list[dict[str, Any]] + self, + service: _LogIngestionAPI, + payload_size_bytes: int, + rows: list[dict[str, Any]], + policy: _RetryPolicy, ) -> dict[str, Any]: - try: - resp = conn.post( - "/logs3/overflow", - json={"content_type": "application/json", "size_bytes": payload_size_bytes, "rows": rows}, - ) - resp.raise_for_status() - payload = resp.json() - except Exception as e: - raise RuntimeError(f"Failed to request logs3 overflow upload URL: {e}") from e - + payload = service.request_json( + "POST", + "/logs3/overflow", + json={"content_type": "application/json", "size_bytes": payload_size_bytes, "rows": rows}, + retry_policy=policy, + ) method = payload.get("method") - if method not in ("PUT", "POST"): - raise RuntimeError(f"Invalid response from API server (method must be PUT or POST): {payload}") signed_url = payload.get("signedUrl") headers = payload.get("headers") fields = payload.get("fields") key = payload.get("key") - if not isinstance(signed_url, str) or not isinstance(key, str): - raise RuntimeError(f"Invalid response from API server: {payload}") - if method == "PUT" and not isinstance(headers, dict): - raise RuntimeError(f"Invalid response from API server: {payload}") - if method == "POST" and not isinstance(fields, dict): - raise RuntimeError(f"Invalid response from API server: {payload}") - + if method not in ("PUT", "POST") or not isinstance(signed_url, str) or not isinstance(key, str): + raise RuntimeError("Invalid logs3 overflow upload response") + if method == "PUT" and not isinstance(headers, dict) or method == "POST" and not isinstance(fields, dict): + raise RuntimeError("Missing logs3 overflow upload headers or fields") if method == "PUT": add_azure_blob_headers(headers, signed_url) - return { "method": method, "signed_url": signed_url, - "headers": headers if isinstance(headers, dict) else {}, - "fields": fields if isinstance(fields, dict) else {}, + "headers": headers or {}, + "fields": fields or {}, "key": key, } - def _upload_logs3_overflow_payload(self, upload: dict[str, Any], payload: str) -> None: - obj_conn = HTTPConnection(base_url="", adapter=_http_adapter) - method = upload["method"] - if method == "POST": - fields = upload.get("fields") - if not isinstance(fields, dict): - raise RuntimeError("Missing logs3 overflow upload fields") - content_type = fields.get("Content-Type", "application/json") - headers = {k: v for k, v in upload.get("headers", {}).items() if k.lower() != "content-type"} - obj_response = obj_conn.post( - upload["signed_url"], - headers=headers, - data=fields, - files={"file": ("logs3.json", payload.encode("utf-8"), content_type)}, - ) + def _upload_logs3_overflow_payload( + self, service: _LogIngestionAPI, upload: dict[str, Any], payload: bytes, policy: _RetryPolicy + ) -> None: + if upload["method"] == "POST": + fields = upload["fields"] + kwargs: dict[str, Any] = { + "headers": {k: v for k, v in upload["headers"].items() if k.lower() != "content-type"}, + "data": fields, + "files": {"file": ("logs3.json", payload, fields.get("Content-Type", "application/json"))}, + } else: - obj_response = obj_conn.put( - upload["signed_url"], - headers=upload["headers"], - data=payload.encode("utf-8"), - ) - 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) - 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 - if self.all_publish_payloads_dir: - _HTTPBackgroundLogger._write_payload_to_dir(payload_dir=self.all_publish_payloads_dir, payload=dataStr) - overflow_upload: dict[str, Any] | None = None - overflow_rows = ( - [ - { - "object_ids": item.overflow_meta.object_ids, - "has_comment": item.overflow_meta.has_comment, - "is_delete": item.overflow_meta.is_delete, - "input_row": {"byte_size": item.overflow_meta.byte_size}, - } - for item in items - ] - if use_overflow - else None + kwargs = {"headers": upload["headers"], "data": payload} + service.storage.request( + upload["method"], upload["signed_url"], retry_mode=_RetryMode.LOG_INGESTION, retry_policy=policy, **kwargs ) - for i in range(self.num_tries): - start_time = time.time() - resp = None - error = None - try: - 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) - 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")) - except Exception as e: - error = e - if error is None and resp is not None and resp.ok: - if overflow_rows: - self._overflow_upload_count += 1 - return - has_response = error is None and resp is not None - is_413 = has_response and resp.status_code == 413 - resp_errmsg = f"{resp.status_code}: {resp.text}" if has_response else str(error) - - should_retry = i + 1 < self.num_tries and not is_413 - if not should_retry and self.failed_publish_payloads_dir: - _HTTPBackgroundLogger._write_payload_to_dir( - payload_dir=self.failed_publish_payloads_dir, payload=dataStr + def _submit_logs_request(self, batch: _PreparedLogBatch, service: _LogIngestionAPI, deadline: float): + with service.attempt(retry_after_on_429=BACKGROUND_LOGGER_BASE_SLEEP_TIME_S * (2**batch.attempts)): + if batch.overflow_rows: + if batch.overflow_upload is None: + batch.overflow_upload = self._request_logs3_overflow_upload( + service, len(batch.payload), batch.overflow_rows, self._ingestion_policy(deadline) + ) + if not batch.uploaded: + try: + self._upload_logs3_overflow_payload( + service, batch.overflow_upload, batch.payload, self._ingestion_policy(deadline) + ) + except BraintrustHTTPError as error: + if error.status_code == 403: + batch.overflow_upload = None + raise _ExpiredOverflowUpload( + "Signed overflow upload URL was rejected; refreshing it" + ) from error + raise + batch.uploaded = True + if batch.overflow_reference is None: + batch.overflow_reference = bt_dumps( + construct_logs3_overflow_request(batch.overflow_upload["key"], len(batch.payload)) + ).encode("utf-8") + service.request( + "POST", + "/logs3", + data=batch.overflow_reference, + headers={"Content-Type": "application/json"}, + retry_policy=self._ingestion_policy(deadline), ) - self._log_failed_payloads_dir() - - retrying_text = " Retrying" if should_retry else "" - errmsg = f"log request failed. Elapsed time: {time.time() - start_time} seconds. Payload size: {payload_bytes}.{retrying_text} Error: {resp_errmsg}" - if not should_retry and self.sync_flush: - raise Exception(errmsg) - print(errmsg, file=self.outfile) - - if is_413: - return - if should_retry: - sleep_time_s = BACKGROUND_LOGGER_BASE_SLEEP_TIME_S * (2**i) - print(f"Sleeping for {sleep_time_s}s", file=self.outfile) - time.sleep(sleep_time_s) - - print(f"log request failed after {self.num_tries} retries. Dropping batch", file=self.outfile) + self._overflow_upload_count += 1 + else: + service.request("POST", "/logs3", data=batch.payload, retry_policy=self._ingestion_policy(deadline)) def _dump_dropped_events(self, wrapped_items): publish_payloads_dir = [x for x in [self.all_publish_payloads_dir, self.failed_publish_payloads_dir] if x] @@ -1480,7 +1754,18 @@ def _log_failed_payloads_dir(self): # Should only be called by BraintrustState. def internal_replace_api_conn(self, api_conn: HTTPConnection): - self.api_conn = LazyValue(lambda: api_conn, use_mutex=False) + # Compatibility hook for extensions; SDK ingestion uses the policy-aware service. + from .api._routing import EndpointRouter + + router = EndpointRouter(app_url=api_conn.base_url, api_url=api_conn.base_url) + token, adapter = api_conn.token or "", _http_adapter + with self.flush_lock: + if self._replaced_entry is not None: + self._replaced_entry.evict() + self._replaced_entry = _LoaderLoginEntry( + lambda: _LogIngestionAPI(router, token, concurrency=self.max_concurrency, adapter=adapter) + ) + self._service_source = self._replaced_entry.lease def set_masking_function(self, masking_function: Callable[[Any], Any] | None): """Set or update the masking function.""" @@ -1497,12 +1782,37 @@ def _internal_get_global_state() -> BraintrustState: _internal_reset_global_state() + + +def _reset_global_state_after_fork() -> None: + """Replace inherited synchronization and transport state in a forked child.""" + global _state + _reset_lazy_value_locks_after_fork() + _state._fork_login = (_state.app_url, _state.login_token, _state.org_name) + if _state._global_bg_logger.has_succeeded: + _state._global_bg_logger.value.reset_after_fork() + _state._client_lock = threading.RLock() + _state._login_operation_lock = threading.RLock() + _state._loader_api_client_cache = LRUCache(max_size=16, on_remove=_state._evict_loader_login_entry) + _state._ingestion_cache = LRUCache(max_size=1, on_remove=_state._evict_loader_login_entry) + _state._client = None + _state._app_conn = None + _state._api_conn = None + _state._proxy_conn = None + _state.logged_in = False + _state._context_manager_lock = threading.Lock() + + +if hasattr(os, "register_at_fork"): + os.register_at_fork(after_in_child=_reset_global_state_after_fork) + + _logger = logging.getLogger("braintrust") @contextlib.contextmanager def _internal_with_custom_background_logger(): - custom_logger = _HTTPBackgroundLogger(LazyValue(lambda: _state.api_conn(), use_mutex=True)) + custom_logger = _HTTPBackgroundLogger(_state.ingestion_service) _state._override_bg_logger.logger = custom_logger try: yield custom_logger @@ -2486,9 +2796,6 @@ def login_to_state( proxy_connection.set_token(api_key) proxy_connection.make_long_lived() - # Replace the global logger's api_conn with this one. - state.login_replace_api_conn(conn) - return state @@ -3116,7 +3423,11 @@ def start_span( def flush(): - """Flush any pending rows to the server.""" + """Deliver pending rows or raise BraintrustLogFlushError, retaining undelivered rows. + + BRAINTRUST_LOG_FLUSH_TIMEOUT bounds cooldown waits (60 seconds by default). + A later explicit flush retries exhausted transient batches without bypassing cooldowns. + """ _state.global_bg_logger().flush() @@ -4569,7 +4880,7 @@ def summarize( """ # Flush our events to the API, and to the data warehouse, to ensure that the link we print # includes the new experiment. - self.flush() + self.state.flush_best_effort() state = self._get_state() project_url = f"{state.app_public_url}/app/{encode_uri_component(state.org_name)}/p/{encode_uri_component(self.project.name)}" @@ -5532,7 +5843,7 @@ def summarize(self, summarize_data: bool = True) -> "DatasetSummary": """ # Flush our events to the API, and to the data warehouse, to ensure that the link we print # includes the new experiment. - self.flush() + self.state.flush_best_effort() state = self._get_state() response = state.api_client().datasets.get_dataset_id_summarize(self.id, summarize_data=summarize_data) raw_data_summary = response.get("data_summary") @@ -6003,7 +6314,7 @@ def log( self.last_start_time = span.end() if not self.async_flush: - self.flush() + self.state.flush_best_effort() return span.id @@ -6102,7 +6413,7 @@ def _emit_log_record( ) if not self.async_flush: - self.flush() + self.state.flush_best_effort() return span.id diff --git a/py/src/braintrust/queue.py b/py/src/braintrust/queue.py index ff6fc6cfc..e2e9bdbf7 100644 --- a/py/src/braintrust/queue.py +++ b/py/src/braintrust/queue.py @@ -33,6 +33,7 @@ def __init__(self, maxsize: int = 0): self._queue: deque[T] = deque(maxlen=maxsize) self._has_items_event = threading.Event() self._total_dropped = 0 + self._reserved = 0 self._enforce_size_limit = False def enforce_queue_size_limit(self, enforce: bool) -> None: @@ -77,7 +78,7 @@ def put(self, item: T) -> list[T]: return dropped - def drain_all(self) -> list[T]: + def drain_all(self, *, reserve: bool = False) -> list[T]: """ Drain all items from the queue. @@ -89,6 +90,8 @@ def drain_all(self) -> list[T]: if len(self._queue) == 0: return [] + if reserve: + self._reserved += len(self._queue) old_queue = self._queue self._queue = deque(maxlen=self.maxsize) @@ -97,6 +100,26 @@ def drain_all(self) -> list[T]: return list(old_queue) if old_queue else [] + def release(self, count: int) -> None: + """Release capacity after retained rows have been delivered.""" + with self._mutex: + self._reserved -= count + + def restore(self, items: list[T]) -> None: + """Return unprepared records ahead of newer records without losing their capacity.""" + with self._mutex: + self._reserved -= len(items) + self._queue.extendleft(reversed(items)) + if items: + self._has_items_event.set() + + def reset_after_fork(self) -> None: + """Discard inherited pending rows and replace locks owned by vanished threads.""" + self._mutex = threading.Lock() + self._queue = deque(maxlen=self.maxsize) + self._has_items_event = threading.Event() + self._reserved = 0 + def size(self) -> int: """ Get the current size of the queue. diff --git a/py/src/braintrust/test_git_metadata_vcr.py b/py/src/braintrust/test_git_metadata_vcr.py index 5b5a3c3a1..660de013a 100644 --- a/py/src/braintrust/test_git_metadata_vcr.py +++ b/py/src/braintrust/test_git_metadata_vcr.py @@ -98,6 +98,8 @@ def _normalize_vcr_request(request): row.pop("created", None) row.pop("root_span_id", None) row.pop("span_id", None) + # This counter is process-wide and depends on which tests ran first. + row.get("span_attributes", {}).pop("exec_counter", None) metrics = row.get("metrics", {}) metrics.pop("start", None) metrics.pop("end", None) diff --git a/py/src/braintrust/test_log_ingestion.py b/py/src/braintrust/test_log_ingestion.py new file mode 100644 index 000000000..e03d636ce --- /dev/null +++ b/py/src/braintrust/test_log_ingestion.py @@ -0,0 +1,669 @@ +"""Writer regressions exercised through real HTTP requests.""" + +import contextlib +import json + +import pytest +from braintrust.api._ingestion import LogIngestionAPI +from braintrust.api._routing import EndpointRouter +from braintrust.api._test_server import scripted_server +from braintrust.logger import _HTTPBackgroundLogger +from braintrust.util import LazyValue + + +def test_permanent_failure_is_reported_and_released(monkeypatch): + monkeypatch.setenv("BRAINTRUST_DISABLE_ATEXIT_FLUSH", "1") + with scripted_server([(413, {}, b"Payload Too Large")]) as (url, handler): + connection = LogIngestionAPI(EndpointRouter(app_url=url, api_url=url), "test", concurrency=4) + writer = _HTTPBackgroundLogger(lambda: contextlib.nullcontext(connection)) + writer._max_request_size_result = {"max_request_size": 10**9, "can_use_overflow": False} + writer.queue.put(LazyValue(lambda: {"id": "score-row", "scores": {"quality": 1}}, use_mutex=False)) + from braintrust.logger import BraintrustLogFlushError + + with pytest.raises(BraintrustLogFlushError, match="413"): + writer.flush() + writer.flush() + assert handler.request_count == 1 + assert json.loads(handler.requests[0][2])["rows"][0]["id"] == "score-row" + assert writer.pending_count == 0 + connection.close() + + +@pytest.fixture +def ingestion_writer(monkeypatch): + monkeypatch.setenv("BRAINTRUST_DISABLE_ATEXIT_FLUSH", "1") + with contextlib.ExitStack() as stack: + + def make(script, *, concurrency=4, capacity=100, persistent=True): + url, handler = stack.enter_context(scripted_server(script, persistent=persistent)) + service = stack.enter_context( + contextlib.closing( + LogIngestionAPI( + EndpointRouter(app_url=url, api_url=url), + "test", + concurrency=concurrency, + ) + ) + ) + writer = _HTTPBackgroundLogger(lambda: contextlib.nullcontext(service)) + writer.max_concurrency = concurrency + writer.queue.maxsize = capacity + writer._max_request_size_result = {"max_request_size": 10**9, "can_use_overflow": False} + return writer, service, handler + + yield make + + +def enqueue(writer, *rows): + for row in rows: + writer.queue.put(LazyValue(lambda row=row: row, use_mutex=False)) + + +@pytest.mark.parametrize("status,retry_after", [(429, "120"), (503, "120"), (429, "date"), (429, None)]) +def test_cooldown_blocks_new_batches_and_is_scoped_to_destination(ingestion_writer, status, retry_after): + import time + + from braintrust.logger import BraintrustLogFlushError + + if retry_after == "date": + import datetime + from email.utils import format_datetime + + retry_after = format_datetime( + datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(seconds=120), usegmt=True + ) + headers = {} if retry_after is None else {"Retry-After": retry_after} + writer, service, handler = ingestion_writer([(status, headers, b"later")], concurrency=1) + enqueue(writer, {"id": "old"}) + started = time.monotonic() + with pytest.raises(BraintrustLogFlushError): + writer.flush(timeout=0.1) + assert time.monotonic() - started < 0.5 + enqueue(writer, {"id": "new"}) + with pytest.raises(BraintrustLogFlushError): + writer.flush(timeout=0.1) + assert writer.pending_count == 2 + assert handler.request_count == 1 + assert service.destination.delay() > 0.5 + # A separate writer with the same URL/credentials must also honor the cooldown. + other = _HTTPBackgroundLogger(lambda: contextlib.nullcontext(service)) + other._max_request_size_result = writer._max_request_size_result + enqueue(other, {"id": "other"}) + with pytest.raises(BraintrustLogFlushError): + other.flush(timeout=0.1) + assert handler.request_count == 1 + independent, _, independent_handler = ingestion_writer([(200, {}, b"ok")]) + enqueue(independent, {"id": "independent"}) + independent.flush() + assert independent_handler.request_count == 1 + + +def test_retry_exhaustion_retains_payload_and_orders_later_updates(ingestion_writer, tmp_path): + from braintrust.logger import BraintrustLogFlushError + + writer, _, handler = ingestion_writer([(503, {}, b"later"), (503, {}, b"later"), (200, {}, b"ok")], capacity=2) + writer.num_tries = 2 + writer.failed_publish_payloads_dir = str(tmp_path) + enqueue(writer, {"id": "row", "scores": {"quality": 0}}) + with pytest.raises(BraintrustLogFlushError) as failure: + writer.flush() + assert failure.value.pending_count == 1 + assert handler.request_count == 2 + assert len(list(tmp_path.glob("*.json"))) == 1 + prepared = writer._pending[0].payload + enqueue(writer, {"id": "row", "_is_merge": True, "scores": {"quality": 1}}) + # Producers remain non-throwing while an earlier wave is retained. + writer.queue.put(LazyValue(lambda: {"id": "later"}, use_mutex=False)) + writer.flush() + assert writer.pending_count == 0 + assert handler.requests[0][2] == handler.requests[1][2] == handler.requests[2][2] == prepared + assert json.loads(handler.requests[3][2])["rows"][0]["scores"] == {"quality": 1} + + +def test_bounded_concurrent_recovery_is_staggered_and_reuses_connections(ingestion_writer): + import threading + import time + from concurrent.futures import ThreadPoolExecutor + + arrivals = [] + accepted = [] + active = 0 + maximum = 0 + recover_at = time.monotonic() + 0.4 + lock = threading.Lock() + + def respond(method, path, body, headers): + nonlocal active, maximum + with lock: + active += 1 + maximum = max(maximum, active) + now = time.monotonic() + arrivals.append(now) + time.sleep(0.015) + with lock: + active -= 1 + if now < recover_at: + return 429, {"Retry-After": "1"}, b"limited" + accepted.append(now) + return 200, {}, b"ok" + + writer, service, handler = ingestion_writer(respond, concurrency=3) + other = _HTTPBackgroundLogger(lambda: contextlib.nullcontext(service)) + other.max_concurrency = 3 + other._max_request_size_result = writer._max_request_size_result + enqueue(writer, *({"id": str(i)} for i in range(6))) + enqueue(other, *({"id": str(i)} for i in range(6, 9))) + with ThreadPoolExecutor(max_workers=2) as pool: + futures = [pool.submit(target.flush, batch_size=1) for target in (writer, other)] + for future in futures: + future.result(timeout=5) + assert maximum <= 3 + assert len(accepted) == 9 + assert handler.request_count <= 12 + assert len(handler.connections) <= 3 + assert min(accepted) - max(t for t in arrivals if t < recover_at) >= 1 + assert all(b - a >= 0.04 for a, b in zip(accepted[:3], accepted[1:3])) + assert writer.pending_count == other.pending_count == 0 + + +@pytest.mark.parametrize("method", ["PUT", "POST"]) +@pytest.mark.parametrize("failure_stage", ["ingest", "upload", "url"]) +def test_overflow_retries_reuse_signed_upload_and_isolate_credentials(monkeypatch, method, failure_stage): + from braintrust.logger import BraintrustLogFlushError + + monkeypatch.setenv("BRAINTRUST_DISABLE_ATEXIT_FLUSH", "1") + logs = 0 + urls = 0 + uploads = [] + with contextlib.ExitStack() as stack: + + def storage_response(verb, path, body, headers): + uploads.append((verb, body, headers.get("Authorization"))) + if failure_stage == "upload" and len(uploads) == 1: + return 503, {}, b"retry" + return 200, {}, b"ok" + + storage_url, _ = stack.enter_context(scripted_server(storage_response, persistent=True)) + + def response(verb, path, body, headers): + nonlocal logs, urls + if path == "/version": + return 200, {}, b'{"logs3_payload_max_bytes": 80}' + if path == "/logs3/overflow": + urls += 1 + if failure_stage == "url" and urls == 1: + return 503, {}, b"retry" + return ( + 200, + {}, + json.dumps( + { + "method": method, + "signedUrl": storage_url, + "headers": {"Content-Type": "application/json"}, + "fields": {"key": "object"}, + "key": "object", + } + ).encode(), + ) + logs += 1 + return (503, {}, b"retry") if failure_stage == "ingest" and logs == 1 else (200, {}, b"ok") + + writer, service, handler = stack.enter_context(_writer_context(response)) + writer.num_tries = 1 + enqueue(writer, {"id": "large", "input": "ü" * 300}) + with pytest.raises(BraintrustLogFlushError): + writer.flush() + payload = writer._pending[0].payload + writer.flush() + assert len(uploads) == (2 if failure_stage == "upload" else 1) + assert all(upload[0] == method and upload[2] is None and payload in upload[1] for upload in uploads) + assert writer._overflow_upload_count == 1 + assert [r[1] for r in handler.requests].count("/logs3/overflow") == (2 if failure_stage == "url" else 1) + log_requests = [r for r in handler.requests if r[1] == "/logs3"] + assert len(log_requests) == (2 if failure_stage == "ingest" else 1) + assert all(request[2] == log_requests[0][2] for request in log_requests) + assert all(r[3] == "Bearer test" for r in handler.requests) + + +# This context also makes version negotiation part of the HTTP regression coverage. +@contextlib.contextmanager +def _writer_context(script): + with scripted_server(script, persistent=True) as (url, handler): + with contextlib.closing( + LogIngestionAPI(EndpointRouter(app_url=url, api_url=url), "test", concurrency=4) + ) as service: + writer = _HTTPBackgroundLogger(lambda: contextlib.nullcontext(service)) + yield writer, service, handler + + +def test_relogin_and_adapter_replacement_keep_retained_batches_on_original_credentials(monkeypatch): + from braintrust import logger + from braintrust.logger import BraintrustLogFlushError, BraintrustState + from requests.adapters import HTTPAdapter + + monkeypatch.setenv("BRAINTRUST_DISABLE_ATEXIT_FLUSH", "1") + monkeypatch.delenv("BRAINTRUST_API_URL", raising=False) + monkeypatch.delenv("BRAINTRUST_PROXY_URL", raising=False) + monkeypatch.delenv("BRAINTRUST_ORG_NAME", raising=False) + calls = 0 + outage = True + + def respond(method, path, body, headers): + nonlocal calls + if path == "/api/apikey/login": + return 200, {}, json.dumps({"org_info": [{"id": "org", "name": "org", "api_url": url}]}).encode() + if path == "/version": + return 200, {}, b"{}" + if path == "/logs3": + calls += 1 + return (503, {}, b"later") if outage else (200, {}, b"ok") + return 200, {}, b'{"ok": true}' + + with scripted_server(respond, persistent=True) as (url, handler): + state = BraintrustState() + monkeypatch.setattr(logger, "_state", state) + monkeypatch.setattr(logger, "_http_adapter", None) + state.login(app_url=url, api_key="original") + writer = state.global_bg_logger() + writer.num_tries = 1 + enqueue(writer, {"id": "retained"}) + with pytest.raises(BraintrustLogFlushError): + writer.flush() + original = writer._pending_service + owned_adapter = original.transport.session.get_adapter(url) + assert len(owned_adapter.poolmanager.pools) == 1 + assert owned_adapter._pool_block is False + enqueue(writer, {"id": "queued-before-relogin"}) + state.login(app_url=url, api_key="replacement", force_login=True) + caller_adapter = HTTPAdapter() + logger.set_http_adapter(caller_adapter) + outage = False + # Evicted while leased: old pools remain usable until the retained wave is delivered. + assert len(owned_adapter.poolmanager.pools) == 1 + enqueue(writer, {"id": "new"}) + writer.flush() + assert len(owned_adapter.poolmanager.pools) == 0 + requests = [request for request in handler.requests if request[1] == "/logs3"] + assert [request[3] for request in requests] == [ + "Bearer original", + "Bearer original", + "Bearer original", + "Bearer replacement", + ] + assert requests[0][2] == requests[1][2] + assert state.api_conn().get_json("ping") == {"ok": True} + # Clearing owned sessions leaves the shared, caller-owned adapter open. + state._ingestion_cache.clear() + assert len(caller_adapter.poolmanager.pools) > 0 + state._client.close() + state.api_conn().close() + state.app_conn().close() + assert len(caller_adapter.poolmanager.pools) > 0 + caller_adapter.close() + + +def test_ingestion_rejects_adapter_retries_without_changing_legacy_policy(ingestion_writer): + from braintrust.api._transport import HTTPConnection, RetryRequestExceptionsAdapter + from requests.adapters import HTTPAdapter + + writer, service, handler = ingestion_writer([(200, {}, b"ok")]) + for adapter in (HTTPAdapter(max_retries=2), RetryRequestExceptionsAdapter(base_num_retries=2)): + with pytest.raises(ValueError, match="single-attempt"): + LogIngestionAPI(service.router, "test", concurrency=4, adapter=adapter) + connection = HTTPConnection(service.router.api_url, adapter=adapter) + assert connection.get("ping").status_code == 200 + connection.close() + assert len(adapter.poolmanager.pools) > 0 + adapter.close() + assert handler.request_count == 2 + + +def test_legacy_connection_replacement_leases_retained_batches_and_closes_owned_pools(ingestion_writer): + from braintrust.api._transport import HTTPConnection + from braintrust.logger import BraintrustLogFlushError + + writer, service, handler = ingestion_writer([(503, {}, b"later"), (200, {}, b"ok")]) + connection = HTTPConnection(service.router.api_url) + try: + connection.set_token("original") + writer.internal_replace_api_conn(connection) + writer.num_tries = 1 + enqueue(writer, {"id": "retained"}) + with pytest.raises(BraintrustLogFlushError): + writer.flush() + original = writer._pending_service + original_adapter = original.transport.session.get_adapter(connection.base_url) + assert len(original_adapter.poolmanager.pools) == 1 + + connection.set_token("replacement") + writer.internal_replace_api_conn(connection) + assert len(original_adapter.poolmanager.pools) == 1 + enqueue(writer, {"id": "new"}) + writer.flush() + assert len(original_adapter.poolmanager.pools) == 0 + requests = [request for request in handler.requests if request[1] == "/logs3"] + assert [request[3] for request in requests] == ["Bearer original", "Bearer original", "Bearer replacement"] + assert requests[0][2] == requests[1][2] + + replacement_adapter = writer._limit_service.transport.session.get_adapter(connection.base_url) + assert len(replacement_adapter.poolmanager.pools) == 1 + writer.internal_replace_api_conn(connection) + assert len(replacement_adapter.poolmanager.pools) == 0 + enqueue(writer, {"id": "last"}) + writer.flush() + final_adapter = writer._limit_service.transport.session.get_adapter(connection.base_url) + assert len(final_adapter.poolmanager.pools) == 1 + writer._finalize() + assert len(final_adapter.poolmanager.pools) == 0 + finally: + connection.close() + + +def test_concurrent_explicit_flush_timeout_leaves_accounting_intact(ingestion_writer): + import threading + from concurrent.futures import ThreadPoolExecutor + + from braintrust.logger import BraintrustLogFlushError + + entered = threading.Event() + released = threading.Event() + + def respond(method, path, body, headers): + entered.set() + assert released.wait(3) + return 200, {}, b"ok" + + writer, _, handler = ingestion_writer(respond) + enqueue(writer, {"id": "row"}) + with ThreadPoolExecutor(max_workers=1) as pool: + future = pool.submit(writer.flush) + assert entered.wait(3) + try: + with pytest.raises(BraintrustLogFlushError, match="waiting for another log delivery"): + writer.flush(timeout=0.05) + assert writer.pending_count == 1 + finally: + released.set() + future.result(timeout=3) + assert writer.pending_count == 0 + assert handler.request_count == 1 + + +def test_explicit_flush_waits_for_active_delivery_then_sends_queued_rows(ingestion_writer): + import threading + import time + from concurrent.futures import ThreadPoolExecutor + + entered = threading.Event() + released = threading.Event() + + def respond(method, path, body, headers): + entered.set() + assert released.wait(3) + return 200, {}, b"ok" + + writer, _, handler = ingestion_writer(respond) + enqueue(writer, {"id": "first"}) + with ThreadPoolExecutor(max_workers=2) as pool: + first = pool.submit(writer.flush) + assert entered.wait(3) + enqueue(writer, {"id": "second"}) + second = pool.submit(writer.flush) + time.sleep(0.05) + released.set() + first.result(timeout=3) + second.result(timeout=3) + assert writer.pending_count == 0 + assert handler.request_count == 2 + + +def test_background_permanent_failure_is_dropped_and_later_rows_continue(ingestion_writer): + import time + + from braintrust.queue import LogQueue + + writer, _, handler = ingestion_writer([(413, {}, b"too large"), (200, {}, b"ok")]) + writer.queue = LogQueue(maxsize=1) + writer.log(LazyValue(lambda: {"id": "failed"}, use_mutex=False)) + deadline = time.monotonic() + 3 + while handler.request_count < 1 or writer.pending_count: + assert time.monotonic() < deadline + time.sleep(0.01) + writer.log(LazyValue(lambda: {"id": "new"}, use_mutex=False)) + deadline = time.monotonic() + 3 + while handler.request_count < 2 or writer.pending_count: + assert time.monotonic() < deadline + time.sleep(0.01) + assert writer.pending_count == 0 + assert handler.request_count == 2 + + +def test_fork_reset_replaces_pending_lazy_value_lock(): + from braintrust.util import _reset_lazy_value_locks_after_fork + + lazy_value = LazyValue(lambda: {"id": "child"}, use_mutex=True) + lazy_value.mutex.acquire() + _reset_lazy_value_locks_after_fork() + assert lazy_value.get() == {"id": "child"} + + +def test_background_permanent_failure_is_reported_by_next_explicit_flush(ingestion_writer): + import time + + from braintrust.logger import BraintrustLogFlushError + + writer, _, handler = ingestion_writer([(401, {}, b"unauthorized")]) + writer.log(LazyValue(lambda: {"id": "failed"}, use_mutex=False)) + deadline = time.monotonic() + 3 + while handler.request_count < 1 or writer.pending_count: + assert time.monotonic() < deadline + time.sleep(0.01) + with pytest.raises(BraintrustLogFlushError, match="401"): + writer.flush() + + +def test_background_writer_resumes_after_retry_budget_during_outage(ingestion_writer): + import time + + writer, _, handler = ingestion_writer([(503, {}, b"outage")] * 3 + [(200, {}, b"ok")]) + writer.num_tries = 2 + writer.log(LazyValue(lambda: {"id": "survives-outage"}, use_mutex=False)) + deadline = time.monotonic() + 8 + while writer.pending_count or handler.request_count < 4: + assert time.monotonic() < deadline + time.sleep(0.02) + assert writer.pending_count == 0 + assert handler.request_count == 4 + + +@pytest.mark.parametrize("status", [400, 401, 403, 413]) +def test_permanent_log_rejections_do_not_block_later_batches(ingestion_writer, status): + from braintrust.logger import BraintrustLogFlushError + + def respond(method, path, body, headers): + row = json.loads(body)["rows"][0] + return (status, {}, b"rejected") if row["id"] == "bad" else (200, {}, b"ok") + + writer, _, handler = ingestion_writer(respond, concurrency=1) + enqueue(writer, {"id": "bad"}, {"id": "good"}) + with pytest.raises(BraintrustLogFlushError): + writer.flush(batch_size=1) + assert writer.pending_count == 0 + assert handler.request_count == 2 + assert [json.loads(request[2])["rows"][0]["id"] for request in handler.requests] == ["bad", "good"] + + +def test_expired_signed_url_is_refreshed_after_403(monkeypatch): + monkeypatch.setenv("BRAINTRUST_DISABLE_ATEXIT_FLUSH", "1") + uploads = 0 + urls = 0 + with contextlib.ExitStack() as stack: + + def storage_response(method, path, body, headers): + nonlocal uploads + uploads += 1 + return (403, {}, b"expired") if uploads == 1 else (200, {}, b"ok") + + storage_url, _ = stack.enter_context(scripted_server(storage_response, persistent=True)) + + def api_response(method, path, body, headers): + nonlocal urls + if path == "/version": + return 200, {}, b'{"logs3_payload_max_bytes": 80}' + if path == "/logs3/overflow": + urls += 1 + return ( + 200, + {}, + json.dumps( + {"method": "PUT", "signedUrl": storage_url, "headers": {}, "key": f"key-{urls}"} + ).encode(), + ) + return 200, {}, b"ok" + + writer, _, handler = stack.enter_context(_writer_context(api_response)) + enqueue(writer, {"id": "large", "input": "x" * 300}) + writer.flush() + assert urls == 2 + assert uploads == 2 + assert handler.request_count == 4 + + +def test_unprepared_records_are_retained_after_local_resolution_failure(ingestion_writer): + writer, _, handler = ingestion_writer([(200, {}, b"ok")]) + writer.num_tries = 1 + calls = 0 + + def resolve(): + nonlocal calls + calls += 1 + if calls == 1: + raise ValueError("cannot resolve yet") + return {"id": "row"} + + writer.queue.put(LazyValue(resolve, use_mutex=False)) + with pytest.raises(Exception, match="constructing records"): + writer.flush() + assert writer.pending_count == 1 + assert handler.request_count == 0 + writer.flush() + assert writer.pending_count == 0 + assert handler.request_count == 1 + + +def test_successful_batches_release_capacity_while_failed_batches_remain(ingestion_writer): + from braintrust.logger import BraintrustLogFlushError + + def respond(method, path, body, headers): + row = json.loads(body)["rows"][0] + return (413, {}, b"too large") if row["id"] == "failed" else (200, {}, b"ok") + + writer, _, handler = ingestion_writer(respond, concurrency=2, capacity=2) + enqueue(writer, {"id": "failed"}, {"id": "delivered"}) + with pytest.raises(BraintrustLogFlushError) as failure: + writer.flush(batch_size=1) + assert failure.value.pending_count == 0 + enqueue(writer, {"id": "later"}) + assert writer.pending_count == 1 + writer.flush() + assert handler.request_count == 3 + + +def test_explicit_flush_finishes_its_wave_while_producers_keep_logging(ingestion_writer): + def respond(method, path, body, headers): + row = json.loads(body)["rows"][0] + if row["id"] == "initial": + enqueue(writer, {"id": "next-wave"}) + return 200, {}, b"ok" + + writer, _, handler = ingestion_writer(respond) + enqueue(writer, {"id": "initial"}) + writer.flush() + assert handler.request_count == 1 + assert writer.pending_count == 1 + writer.flush() + assert handler.request_count == 2 + assert writer.pending_count == 0 + + +def test_interpreter_shutdown_flushes_without_an_available_thread_pool(ingestion_writer): + import os + import subprocess + import sys + + _, service, handler = ingestion_writer([(200, {}, b"ok")]) + program = f""" +import contextlib +from braintrust.api._ingestion import LogIngestionAPI +from braintrust.api._routing import EndpointRouter +from braintrust.logger import _HTTPBackgroundLogger +from braintrust.util import LazyValue +service = LogIngestionAPI(EndpointRouter(app_url={service.router.api_url!r}, api_url={service.router.api_url!r}), "test", concurrency=4) +writer = _HTTPBackgroundLogger(lambda: contextlib.nullcontext(service)) +writer.sync_flush = True +writer._max_request_size_result = {{"max_request_size": 1000000, "can_use_overflow": False}} +writer.log(LazyValue(lambda: {{"id": "shutdown-row"}}, use_mutex=False)) +service.transport.session.get_adapter({service.router.api_url!r}).poolmanager.clear() +""" + process = subprocess.run( + [sys.executable, "-c", program], + env={**os.environ, "BRAINTRUST_DISABLE_ATEXIT_FLUSH": "false"}, + capture_output=True, + text=True, + timeout=10, + ) + assert process.returncode == 0, process.stderr + assert handler.request_count == 1, process.stderr + assert json.loads(handler.requests[0][2])["rows"][0]["id"] == "shutdown-row" + + +def test_forked_child_does_not_inherit_parent_pending_rows(monkeypatch): + import os + + if not hasattr(os, "fork"): + pytest.skip("fork is not available") + from braintrust import logger + from braintrust.logger import BraintrustState + + monkeypatch.setenv("BRAINTRUST_DISABLE_ATEXIT_FLUSH", "1") + state = BraintrustState() + monkeypatch.setattr(logger, "_state", state) + writer = state.global_bg_logger() + enqueue(writer, {"id": "parent-only"}) + child = os.fork() + if child == 0: + os._exit(0 if state.global_bg_logger().pending_count == 0 else 3) + _, status = os.waitpid(child, 0) + assert os.waitstatus_to_exitcode(status) == 0 + assert writer.pending_count == 1 + + +def test_new_background_writer_does_not_spin_during_shared_long_cooldown(ingestion_writer): + import time + + from braintrust.logger import BraintrustLogFlushError + + writer, service, handler = ingestion_writer([(429, {"Retry-After": "120"}, b"later")]) + enqueue(writer, {"id": "throttled"}) + with pytest.raises(BraintrustLogFlushError): + writer.flush(timeout=0.1) + # This writer must negotiate /version after the same destination recovers. + other = _HTTPBackgroundLogger(lambda: contextlib.nullcontext(service)) + observed = [] + + def mask(value): + if value == "observe": + observed.append(value) + return value + + other.set_masking_function(mask) + other.log(LazyValue(lambda: {"id": "waiting", "input": "observe"}, use_mutex=False)) + deadline = time.monotonic() + 3 + while not observed: + assert time.monotonic() < deadline + time.sleep(0.01) + time.sleep(0.1) + assert len(observed) == 1 + assert other.pending_count == 1 + assert handler.request_count == 1 diff --git a/py/src/braintrust/test_logger.py b/py/src/braintrust/test_logger.py index 8e0bef213..a25a44f54 100644 --- a/py/src/braintrust/test_logger.py +++ b/py/src/braintrust/test_logger.py @@ -254,65 +254,11 @@ def test_http_background_logger_atexit_flush(monkeypatch, disable_atexit_flush, monkeypatch.setenv("BRAINTRUST_DISABLE_ATEXIT_FLUSH", disable_atexit_flush) with patch("atexit.register") as mock_register: - _HTTPBackgroundLogger(LazyValue(MagicMock, use_mutex=False)) + _HTTPBackgroundLogger(logger._state.ingestion_service) assert mock_register.called is registers -class TestHTTPBackgroundLoggerLogs3(TestCase): - def test_submit_logs_request_413_skips_retries(self) -> None: - """Any 413 while publishing ``/logs3`` cannot succeed on retry with the same payload. - - ``sync_flush`` controls whether the terminal failure raises instead of printing. - """ - from braintrust.logger import ( - LogItemWithMeta, - Logs3OverflowInputRow, - _HTTPBackgroundLogger, - ) - - item = LogItemWithMeta( - str_value="{}", - overflow_meta=Logs3OverflowInputRow( - object_ids={}, - has_comment=False, - is_delete=False, - byte_size=2, - ), - ) - max_result = {"max_request_size": 10**9, "can_use_overflow": True} - - for response_text in ("Request Too Long", "", "Payload Too Large"): - for sync_flush in (False, True): - with self.subTest(response_text=response_text, sync_flush=sync_flush): - mock_resp = MagicMock() - mock_resp.ok = False - mock_resp.status_code = 413 - mock_resp.text = response_text - - mock_conn = MagicMock() - mock_conn.post.return_value = mock_resp - - bg = _HTTPBackgroundLogger(LazyValue(lambda: mock_conn, use_mutex=False)) - bg.num_tries = 5 - bg.sync_flush = sync_flush - bg.failed_publish_payloads_dir = "/tmp/failed-payloads" - - with patch.object(_HTTPBackgroundLogger, "_write_payload_to_dir") as mock_write_payload: - if sync_flush: - with self.assertRaises(Exception) as cm: - bg._submit_logs_request([item], max_result) - self.assertIn("413", str(cm.exception)) - else: - bg._submit_logs_request([item], max_result) - - self.assertEqual(mock_conn.post.call_count, 1) - mock_write_payload.assert_called_once() - self.assertEqual( - mock_write_payload.call_args.kwargs["payload_dir"], bg.failed_publish_payloads_dir - ) - - def test_load_prompt_async_signature_matches_load_prompt(): assert ( inspect.signature(braintrust.load_prompt_async).parameters diff --git a/py/src/braintrust/test_span_customizer.py b/py/src/braintrust/test_span_customizer.py index e61fc07db..d4493b102 100644 --- a/py/src/braintrust/test_span_customizer.py +++ b/py/src/braintrust/test_span_customizer.py @@ -1,4 +1,5 @@ import asyncio +import contextlib import inspect import json import logging @@ -11,6 +12,9 @@ import pytest from braintrust import Attachment, SpanCustomizer, SpanExportData, auto_instrument, logger, set_span_customizers +from braintrust.api._ingestion import LogIngestionAPI +from braintrust.api._routing import EndpointRouter +from braintrust.api._test_server import scripted_server from braintrust.functions.stream import BraintrustJsonChunk, BraintrustStream from braintrust.integrations.utils import _resolved_attachment_from_bytes from braintrust.span_customizer import _customize_span_export @@ -19,6 +23,37 @@ from braintrust.util import LazyValue +@pytest.fixture +def http_ingestion(monkeypatch): + monkeypatch.setenv("BRAINTRUST_DISABLE_ATEXIT_FLUSH", "1") + with contextlib.ExitStack() as stack: + + def make(retry=False): + payloads = [] + + def respond(method, path, body, headers): + if path == "/version": + return 200, {}, b"{}" + payloads.append(body) + if retry and len(payloads) == 1: + return 503, {}, b"retry transport" + return 200, {}, b"ok" + + url, _ = stack.enter_context(scripted_server(respond)) + service = stack.enter_context( + contextlib.closing( + LogIngestionAPI( + EndpointRouter(app_url=url, api_url=url), + "test", + concurrency=4, + ) + ) + ) + return service, payloads + + yield make + + @pytest.fixture(autouse=True) def reset_customizers(): set_span_customizers(None) @@ -273,7 +308,7 @@ def on_span_export(self, data): @pytest.mark.parametrize("backend", ["memory", "http"]) @pytest.mark.parametrize("include_healthy", [False, True], ids=["all-dropped", "mixed-batch"]) def test_dropped_records_skip_attachments_masking_and_upload( - monkeypatch, with_memory_logger, test_logger, caplog, backend, include_healthy + monkeypatch, with_memory_logger, test_logger, caplog, backend, include_healthy, http_ingestion ): attachment = Attachment(data=b"private", filename="private.txt", content_type="text/plain") upload = MagicMock() @@ -299,12 +334,11 @@ def on_span_export(self, data): pending = list(with_memory_logger.logs) with_memory_logger.logs.clear() monkeypatch.setenv("BRAINTRUST_DISABLE_ATEXIT_FLUSH", "1") - connection = MagicMock() - connection.post.return_value = SimpleNamespace(ok=True) + connection, payloads = http_ingestion() background = ( logger._MemoryBackgroundLogger() if backend == "memory" - else logger._HTTPBackgroundLogger(LazyValue(lambda: connection, use_mutex=False)) + else logger._HTTPBackgroundLogger(lambda: contextlib.nullcontext(connection)) ) masked = [] @@ -330,10 +364,10 @@ def mask(value): rows, attachments = background._unwrap_lazy_values(pending) assert attachments == [] if include_healthy: - sent = json.loads(connection.post.call_args.kwargs["data"])["rows"] + sent = json.loads(payloads[-1])["rows"] assert [row["id"] for row in sent] == [healthy.id] else: - connection.post.assert_not_called() + assert payloads == [] expected_ids = [healthy.id] if healthy else [] assert [row["id"] for row in rows] == expected_ids @@ -458,7 +492,7 @@ def on_span_export(self, data): def test_customization_precedes_attachments_masking_and_reuses_records_on_retry( - monkeypatch, with_memory_logger, test_logger, caplog + monkeypatch, with_memory_logger, test_logger, caplog, http_ingestion ): attachment = Attachment(data=b"private", filename="private.txt", content_type="text/plain") rejected_attachment = Attachment(data=b"secret", filename="secret.txt", content_type="text/plain") @@ -499,18 +533,8 @@ def resolve_later_record(): pending.append(LazyValue(resolve_later_record, use_mutex=False)) monkeypatch.setenv("BRAINTRUST_DISABLE_ATEXIT_FLUSH", "1") - monkeypatch.setattr(logger.time, "sleep", lambda _: None) - connection = MagicMock() - payloads = [] - - def send(_path, *, data): - payloads.append(data) - if len(payloads) == 1: - raise ConnectionError("retry transport") - return SimpleNamespace(ok=True) - - connection.post.side_effect = send - background = logger._HTTPBackgroundLogger(LazyValue(lambda: connection, use_mutex=False)) + connection, payloads = http_ingestion(retry=True) + background = logger._HTTPBackgroundLogger(lambda: contextlib.nullcontext(connection)) background.num_tries = 2 background.sync_flush = True masked = [] @@ -542,20 +566,10 @@ def mask(value): assert "private" not in caplog.text -def _retrying_http_logger(monkeypatch, pending): +def _retrying_http_logger(monkeypatch, pending, http_ingestion): monkeypatch.setenv("BRAINTRUST_DISABLE_ATEXIT_FLUSH", "1") - monkeypatch.setattr(logger.time, "sleep", lambda _: None) - connection = MagicMock() - payloads = [] - - def send(_path, *, data): - payloads.append(data) - if len(payloads) == 1: - raise ConnectionError("retry transport") - return SimpleNamespace(ok=True) - - connection.post.side_effect = send - background = logger._HTTPBackgroundLogger(LazyValue(lambda: connection, use_mutex=False)) + connection, payloads = http_ingestion(retry=True) + background = logger._HTTPBackgroundLogger(lambda: contextlib.nullcontext(connection)) background.num_tries = 2 background.sync_flush = True background._max_request_size_result = {"max_request_size": 6_000_000, "can_use_overflow": False} @@ -566,7 +580,7 @@ def send(_path, *, data): @pytest.mark.parametrize("error", [GeneratorExit, asyncio.CancelledError]) def test_base_exception_hook_failures_drop_only_that_record( - monkeypatch, with_memory_logger, test_logger, caplog, error + monkeypatch, with_memory_logger, test_logger, caplog, error, http_ingestion ): invocations = [] @@ -582,7 +596,7 @@ def on_span_export(self, data): healthy = test_logger.start_span(input="safe") pending = list(with_memory_logger.logs) with_memory_logger.logs.clear() - background, payloads = _retrying_http_logger(monkeypatch, pending) + background, payloads = _retrying_http_logger(monkeypatch, pending, http_ingestion) with caplog.at_level(logging.ERROR, logger="braintrust"): background.flush() @@ -594,7 +608,9 @@ def on_span_export(self, data): assert "private" not in caplog.text -def test_redacting_integration_attachment_prevents_upload(monkeypatch, with_memory_logger, test_logger): +def test_redacting_integration_attachment_prevents_upload( + monkeypatch, with_memory_logger, test_logger, http_ingestion +): # Integrations convert inline media into Attachment objects at capture time. attachment = _resolved_attachment_from_bytes(b"private image", "image/png", prefix="input").attachment upload = MagicMock() @@ -616,7 +632,7 @@ def on_span_export(self, data): ) pending = list(with_memory_logger.logs) with_memory_logger.logs.clear() - background, payloads = _retrying_http_logger(monkeypatch, pending) + background, payloads = _retrying_http_logger(monkeypatch, pending, http_ingestion) background.flush() assert seen == [attachment] @@ -627,7 +643,7 @@ def on_span_export(self, data): @pytest.mark.parametrize("backend", ["memory", "http"]) def test_masking_remains_logger_local_and_runs_on_merged_manual_records( - monkeypatch, with_memory_logger, test_logger, backend + monkeypatch, with_memory_logger, test_logger, backend, http_ingestion ): class Customize(SpanCustomizer): def on_span_export(self, data): @@ -641,10 +657,11 @@ def on_span_export(self, data): set_span_customizers([Customize()]) monkeypatch.setenv("BRAINTRUST_DISABLE_ATEXIT_FLUSH", "1") + connection, _ = http_ingestion() background = ( logger._MemoryBackgroundLogger() if backend == "memory" - else logger._HTTPBackgroundLogger(LazyValue(lambda: MagicMock(), use_mutex=False)) + else logger._HTTPBackgroundLogger(lambda: contextlib.nullcontext(connection)) ) other_background = logger._MemoryBackgroundLogger() diff --git a/py/src/braintrust/util.py b/py/src/braintrust/util.py index 0f4d09270..6d6f5404f 100644 --- a/py/src/braintrust/util.py +++ b/py/src/braintrust/util.py @@ -3,6 +3,7 @@ import sys import threading import urllib.parse +import weakref from collections.abc import Callable, Mapping from dataclasses import dataclass from typing import Any, Generic, Literal, TypedDict, TypeVar @@ -175,6 +176,14 @@ class _LazyValuePendingState: _LazyValueState = _LazyValueResolvedState[T] | _LazyValuePendingState +_lazy_values: weakref.WeakSet = weakref.WeakSet() + + +def _reset_lazy_value_locks_after_fork() -> None: + """Replace locks inherited from threads that do not exist in a forked child.""" + for lazy_value in list(_lazy_values): + lazy_value.reset_after_fork() + class LazyValue(Generic[T]): """A simple wrapper around a callable object which computes the value @@ -185,6 +194,11 @@ def __init__(self, callable: Callable[[], T], use_mutex: bool): self.callable = callable self.mutex = threading.Lock() if use_mutex else None self._state: _LazyValueState[T] = _LazyValuePendingState() + _lazy_values.add(self) + + def reset_after_fork(self) -> None: + if self.mutex is not None: + self.mutex = threading.Lock() @property def has_succeeded(self) -> bool: