diff --git a/py/benchmarks/benches/bench_import_time.py b/py/benchmarks/benches/bench_import_time.py new file mode 100644 index 000000000..25b1387ab --- /dev/null +++ b/py/benchmarks/benches/bench_import_time.py @@ -0,0 +1,23 @@ +import pathlib +import sys + +import pyperf + + +if __package__ in (None, ""): + sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[2])) + +from benchmarks._utils import disable_pyperf_psutil + + +def main(runner: pyperf.Runner | None = None) -> None: + if runner is None: + disable_pyperf_psutil() + runner = pyperf.Runner() + + runner.bench_command("import_time[python_startup]", [sys.executable, "-c", "pass"]) + runner.bench_command("import_time[braintrust]", [sys.executable, "-c", "import braintrust"]) + + +if __name__ == "__main__": + main() diff --git a/py/src/braintrust/audit.py b/py/src/braintrust/audit.py index 55344db29..01396887e 100644 --- a/py/src/braintrust/audit.py +++ b/py/src/braintrust/audit.py @@ -2,8 +2,11 @@ Utilities for working with audit headers. """ +# PEP 810: defer these imports until first use on Python 3.15+. Ignored on older versions. +__lazy_modules__ = {"gzip"} + import base64 -import gzip +import gzip as _gzip import json from typing import TypedDict @@ -17,6 +20,6 @@ class AuditResource(TypedDict): def parse_audit_resources(hdr: str) -> list[AuditResource]: j = json.loads(hdr) if j["v"] == 1: - return json.loads(gzip.decompress(base64.b64decode(j["p"]))) + return json.loads(_gzip.decompress(base64.b64decode(j["p"]))) else: raise ValueError(f"Unsupported audit resources protocol version: {j['v']}") diff --git a/py/src/braintrust/framework.py b/py/src/braintrust/framework.py index 58a8c08ce..3b4c2e52f 100644 --- a/py/src/braintrust/framework.py +++ b/py/src/braintrust/framework.py @@ -1,9 +1,13 @@ +# PEP 810: defer these imports until first use on Python 3.15+. Ignored on older versions. +__lazy_modules__ = {"asyncio", "tqdm.asyncio", "tqdm.auto"} + import abc import asyncio import contextvars import dataclasses import inspect import json +import os import re import sys import threading @@ -14,7 +18,6 @@ from collections.abc import Awaitable, Callable, Coroutine, Iterable, Iterator, Mapping, Sequence from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager -from multiprocessing import cpu_count from typing import ( Any, Generic, @@ -1304,7 +1307,7 @@ def init_experiment( class EvalThreadPoolSingleton: def __init__(self): self._thread_pool = None - self._max_workers = cpu_count() + self._max_workers = os.cpu_count() def set_max_workers(self, max_workers): assert self._thread_pool is None, "Cannot set max_workers. Thread pool has already been initialized" diff --git a/py/src/braintrust/framework2.py b/py/src/braintrust/framework2.py index 7540659c0..914ac5c9b 100644 --- a/py/src/braintrust/framework2.py +++ b/py/src/braintrust/framework2.py @@ -1,9 +1,12 @@ +# PEP 810: defer these imports until first use on Python 3.15+. Ignored on older versions. +__lazy_modules__ = {"slugify"} + import dataclasses import json from collections.abc import Callable, Mapping, Sequence from typing import TYPE_CHECKING, Any, overload -import slugify +import slugify as _slugify from braintrust.logger import _internal_get_global_state, api_conn, login from .framework import _is_lazy_load, bcolors # type: ignore @@ -201,7 +204,7 @@ def create( name = f"Tool {self._task_counter}" assert name is not None if not slug: - slug = slugify.slugify(name) + slug = _slugify.slugify(name) f = CodeFunction( project=self.project, handler=handler, @@ -296,7 +299,7 @@ def create( if not name: name = f"Prompt {self._task_counter}" if not slug: - slug = slugify.slugify(name) + slug = _slugify.slugify(name) tool_functions: list[CodeFunction | SavedFunctionId] = [] raw_tools: list[ToolFunctionDefinition] = [] @@ -363,7 +366,7 @@ def create( metadata: Mapping[str, Any] | None = None, ) -> EvalParameters: if slug is None or len(slug) == 0: - slug = slugify.slugify(name) + slug = _slugify.slugify(name) parameters = CodeParameters( project=self.project, @@ -491,7 +494,7 @@ def create( else: name = f"Scorer {self._task_counter}" if slug is None or len(slug) == 0: - slug = slugify.slugify(name) + slug = _slugify.slugify(name) if handler is not None: # code scorer assert parameters is not None @@ -592,7 +595,7 @@ def create( else: name = f"Classifier {self._task_counter}" if slug is None or len(slug) == 0: - slug = slugify.slugify(name) + slug = _slugify.slugify(name) f = CodeFunction( project=self.project, diff --git a/py/src/braintrust/functions/invoke.py b/py/src/braintrust/functions/invoke.py index 08bf65594..c023d365c 100644 --- a/py/src/braintrust/functions/invoke.py +++ b/py/src/braintrust/functions/invoke.py @@ -1,4 +1,7 @@ -import asyncio +# PEP 810: defer these imports until first use on Python 3.15+. Ignored on older versions. +__lazy_modules__ = {"asyncio"} + +import asyncio as _asyncio from collections.abc import Sequence from typing import Any, Literal, TypedDict, TypeVar, overload @@ -320,7 +323,7 @@ async def invoke_async( When ``stream=True``, the returned ``BraintrustStream`` supports ``async for`` and ``final_value_async()``. """ - return await asyncio.to_thread( + return await _asyncio.to_thread( invoke, function_id=function_id, version=version, diff --git a/py/src/braintrust/functions/stream.py b/py/src/braintrust/functions/stream.py index 36ca86540..919c490bb 100644 --- a/py/src/braintrust/functions/stream.py +++ b/py/src/braintrust/functions/stream.py @@ -5,7 +5,10 @@ with utility methods to make them easy to log and convert into various formats. """ -import asyncio +# PEP 810: defer these imports until first use on Python 3.15+. Ignored on older versions. +__lazy_modules__ = {"asyncio"} + +import asyncio as _asyncio import dataclasses import json from collections.abc import AsyncIterator, Generator, Iterable, Iterator @@ -112,7 +115,7 @@ def __init__(self, base_stream: SSEClient | list[BraintrustStreamChunk]): else: self.stream = base_stream self._async_iterator: Iterator[BraintrustStreamChunk] | None = None - self._pending_async_read: asyncio.Task[BraintrustStreamChunk | object] | None = None + self._pending_async_read: _asyncio.Task[BraintrustStreamChunk | object] | None = None self._memoized_final_value = _UNSET def _parse_sse_stream(self, sse_client: SSEClient) -> Generator[BraintrustStreamChunk, None, None]: @@ -183,7 +186,7 @@ def final_value(self): async def final_value_async(self): """Consume the stream asynchronously and return its combined final value.""" - return await asyncio.to_thread(self.final_value) + return await _asyncio.to_thread(self.final_value) def __iter__(self): """ @@ -202,12 +205,12 @@ async def __anext__(self) -> BraintrustStreamChunk: self._async_iterator = iter(self.stream) self.stream = self._async_iterator if self._pending_async_read is None: - self._pending_async_read = asyncio.create_task(asyncio.to_thread(_next_or_end, self._async_iterator)) + self._pending_async_read = _asyncio.create_task(_asyncio.to_thread(_next_or_end, self._async_iterator)) pending_read = self._pending_async_read try: - chunk = await asyncio.shield(pending_read) - except asyncio.CancelledError: + chunk = await _asyncio.shield(pending_read) + except _asyncio.CancelledError: if pending_read.cancelled(): self._pending_async_read = None raise diff --git a/py/src/braintrust/integrations/openrouter/test_openrouter.py b/py/src/braintrust/integrations/openrouter/test_openrouter.py index 576913d5c..7f4473cc9 100644 --- a/py/src/braintrust/integrations/openrouter/test_openrouter.py +++ b/py/src/braintrust/integrations/openrouter/test_openrouter.py @@ -1,3 +1,4 @@ +import importlib import inspect import os import time @@ -12,9 +13,6 @@ openrouter = pytest.importorskip("openrouter") from openrouter import OpenRouter -from openrouter.chat import Chat -from openrouter.embeddings import Embeddings -from openrouter.responses import Responses PROJECT_NAME = "test-openrouter-sdk" @@ -249,6 +247,14 @@ async def test_wrap_openrouter_responses_send_async_stream(memory_logger): def test_openrouter_integration_setup_is_idempotent(monkeypatch): + # Resolve the classes at call time rather than at module import: openrouter's + # lazy sub-SDK loader drops and re-imports a module when CPython < 3.12's import + # lock raises a spurious KeyError, which would leave module-level names pointing + # at a stale class that setup() never patches. + Chat = importlib.import_module("openrouter.chat").Chat + Embeddings = importlib.import_module("openrouter.embeddings").Embeddings + Responses = importlib.import_module("openrouter.responses").Responses + first_send = inspect.getattr_static(Chat, "send") first_generate = inspect.getattr_static(Embeddings, "generate") first_responses_send = inspect.getattr_static(Responses, "send") diff --git a/py/src/braintrust/logger.py b/py/src/braintrust/logger.py index b12079c7c..45e8355fd 100644 --- a/py/src/braintrust/logger.py +++ b/py/src/braintrust/logger.py @@ -1,4 +1,7 @@ -import asyncio +# PEP 810: defer these imports until first use on Python 3.15+. Ignored on older versions. +__lazy_modules__ = {"asyncio"} + +import asyncio as _asyncio import atexit import base64 import builtins @@ -25,7 +28,6 @@ from collections import Counter from collections.abc import Callable, Iterator, Mapping, MutableMapping, Sequence from functools import partial, wraps -from multiprocessing import cpu_count from types import TracebackType from typing import ( TYPE_CHECKING, @@ -937,7 +939,7 @@ def set_http_adapter(adapter: HTTPAdapter) -> None: # Sometimes we'd like to launch network requests concurrently. We provide a # thread pool to accomplish this. Use a multiple of number of CPU cores to limit # concurrency. -HTTP_REQUEST_THREAD_POOL = concurrent.futures.ThreadPoolExecutor(max_workers=cpu_count()) +HTTP_REQUEST_THREAD_POOL = concurrent.futures.ThreadPoolExecutor(max_workers=os.cpu_count()) def api_conn(): @@ -2202,7 +2204,7 @@ async def load_prompt_async( api_key=api_key, org_name=org_name, ) - await asyncio.to_thread(lambda: prompt.name) + await _asyncio.to_thread(lambda: prompt.name) return prompt diff --git a/py/src/braintrust/parameters.py b/py/src/braintrust/parameters.py index c0a97f993..180cf34af 100644 --- a/py/src/braintrust/parameters.py +++ b/py/src/braintrust/parameters.py @@ -1,5 +1,8 @@ """Evaluation parameters support for Python SDK.""" +# PEP 810: defer these imports until first use on Python 3.15+. Ignored on older versions. +__lazy_modules__ = {"jsonschema", "jsonschema.exceptions"} + from collections.abc import Mapping from copy import deepcopy from dataclasses import dataclass diff --git a/py/src/braintrust/prompt_cache/disk_cache.py b/py/src/braintrust/prompt_cache/disk_cache.py index b74080249..ab20e0be4 100644 --- a/py/src/braintrust/prompt_cache/disk_cache.py +++ b/py/src/braintrust/prompt_cache/disk_cache.py @@ -7,6 +7,9 @@ and handles file system errors gracefully. """ +# PEP 810: defer these imports until first use on Python 3.15+. Ignored on older versions. +__lazy_modules__ = {"gzip"} + import gzip import hashlib import json diff --git a/py/src/braintrust/sandbox.py b/py/src/braintrust/sandbox.py index c2433c5a8..2739899b1 100644 --- a/py/src/braintrust/sandbox.py +++ b/py/src/braintrust/sandbox.py @@ -1,3 +1,6 @@ +# PEP 810: defer these imports until first use on Python 3.15+. Ignored on older versions. +__lazy_modules__ = {"slugify"} + import dataclasses import sys from collections.abc import Sequence diff --git a/py/src/braintrust/test_lazy_integration_imports.py b/py/src/braintrust/test_lazy_integration_imports.py index 7672fe306..9f8f3afe0 100644 --- a/py/src/braintrust/test_lazy_integration_imports.py +++ b/py/src/braintrust/test_lazy_integration_imports.py @@ -2,10 +2,17 @@ import subprocess import sys +import pytest + + +requires_lazy_imports = pytest.mark.skipif( + sys.version_info < (3, 15), reason="__lazy_modules__ (PEP 810) is a no-op before Python 3.15" +) + def _run_import_script(script: str) -> dict[str, object]: result = subprocess.run([sys.executable, "-c", script], check=True, capture_output=True, text=True) - return json.loads(result.stdout) + return json.loads(result.stdout.splitlines()[-1]) def test_import_braintrust_does_not_load_integrations_or_providers(): @@ -82,3 +89,74 @@ def test_integrations_package_lazily_resolves_public_classes(): "loaded_openai": True, "loaded_anthropic": False, } + + +@requires_lazy_imports +def test_import_braintrust_defers_lazy_modules(): + loaded = _run_import_script( + """ +import json +import sys + +import braintrust + +declared = set() +for name, module in list(sys.modules.items()): + if name.startswith("braintrust"): + declared.update(getattr(module, "__lazy_modules__", ())) + +print(json.dumps({ + "declared": sorted(declared), + "loaded": sorted(name for name in declared | {"multiprocessing"} if name in sys.modules), +})) +""" + ) + + assert {"asyncio", "jsonschema", "tqdm.auto"} <= set(loaded["declared"]) + assert loaded["loaded"] == [] + + +@requires_lazy_imports +def test_eval_and_parameters_work_with_lazy_modules(): + result = _run_import_script( + """ +import json + +from braintrust import Eval +from braintrust.parameters import validate_json_schema + + +def exact_match(input, output, expected): + return float(output == expected) + + +eval_result = Eval( + "lazy-imports-smoke", + data=[{"input": 1, "expected": 2}, {"input": 2, "expected": 3}], + task=lambda input: input + 1, + scores=[exact_match], + no_send_logs=True, +) + +schema = {"type": "object", "properties": {"n": {"type": "integer", "default": 7}}} +try: + validate_json_schema({"n": "not-an-int"}, schema) + invalid_error = None +except ValueError as exc: + invalid_error = str(exc) + +print(json.dumps({ + "outputs": sorted(r.output for r in eval_result.results), + "scores": sorted(r.scores["exact_match"] for r in eval_result.results), + "validated": validate_json_schema({}, schema), + "invalid_error": invalid_error, +})) +""" + ) + + assert result == { + "outputs": [2, 3], + "scores": [1.0, 1.0], + "validated": {"n": 7}, + "invalid_error": "Invalid parameters: n: 'not-an-int' is not of type 'integer'", + } diff --git a/py/src/braintrust/trace.py b/py/src/braintrust/trace.py index cdcb6455a..cb3b7d258 100644 --- a/py/src/braintrust/trace.py +++ b/py/src/braintrust/trace.py @@ -5,6 +5,9 @@ spans from the current evaluation task without making server round-trips. """ +# PEP 810: defer these imports until first use on Python 3.15+. Ignored on older versions. +__lazy_modules__ = {"asyncio"} + import asyncio import math from collections.abc import Awaitable, Callable, Mapping diff --git a/py/src/braintrust/workflow_eval.py b/py/src/braintrust/workflow_eval.py index b9ee968e4..6619a91e9 100644 --- a/py/src/braintrust/workflow_eval.py +++ b/py/src/braintrust/workflow_eval.py @@ -1,5 +1,8 @@ """Experimental workflow evaluations with asynchronous provider submissions.""" +# PEP 810: defer these imports until first use on Python 3.15+. Ignored on older versions. +__lazy_modules__ = {"asyncio"} + import asyncio import base64 import dataclasses