perf: defer heavy imports on Python 3.15 with __lazy_modules__

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

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']}")

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"

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,

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,

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

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

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
 
 

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

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

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

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'",
+    }

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