Improve import startup with lazy top-level exports (refs openai#2819)

Author: fede-kamel

scripts/bench_import.py

@@ -0,0 +1,96 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import argparse
+import os
+import subprocess
+import sys
+import time
+from pathlib import Path
+
+
+def _pythonpath_for_repo() -> str | None:
+    src = Path(__file__).resolve().parents[1] / "src"
+    if not src.exists():
+        return None
+
+    existing = os.environ.get("PYTHONPATH")
+    if existing:
+        return f"{src}{os.pathsep}{existing}"
+    return str(src)
+
+
+def _cold_import_seconds(repeats: int, env: dict[str, str]) -> list[float]:
+    samples: list[float] = []
+    for _ in range(repeats):
+        start = time.perf_counter()
+        subprocess.run([sys.executable, "-c", "import openai"], check=True, env=env, stdout=subprocess.DEVNULL)
+        samples.append(time.perf_counter() - start)
+    return samples
+
+
+def _importtime_output(env: dict[str, str]) -> str:
+    proc = subprocess.run(
+        [sys.executable, "-X", "importtime", "-c", "import openai"],
+        check=True,
+        env=env,
+        stderr=subprocess.PIPE,
+        stdout=subprocess.DEVNULL,
+        text=True,
+    )
+    return proc.stderr
+
+
+def _parse_importtime(importtime_stderr: str) -> list[tuple[int, str]]:
+    rows: list[tuple[int, str]] = []
+    for line in importtime_stderr.splitlines():
+        if "| " not in line:
+            continue
+        if "import time:" not in line:
+            continue
+        _, _, payload = line.partition("import time:")
+        parts = [p.strip() for p in payload.split("|")]
+        if len(parts) != 3:
+            continue
+        cumulative_raw = parts[1]
+        module = parts[2]
+        if not module.startswith("openai"):
+            continue
+        try:
+            cumulative = int(cumulative_raw)
+        except ValueError:
+            continue
+        rows.append((cumulative, module))
+    rows.sort(reverse=True)
+    return rows
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description="Benchmark openai import time for this checkout.")
+    parser.add_argument("--repeats", type=int, default=5, help="Number of cold imports to sample")
+    parser.add_argument("--top", type=int, default=20, help="How many importtime rows to print")
+    args = parser.parse_args()
+
+    env = dict(os.environ)
+    pythonpath = _pythonpath_for_repo()
+    if pythonpath is not None:
+        env["PYTHONPATH"] = pythonpath
+
+    samples = _cold_import_seconds(repeats=args.repeats, env=env)
+    avg = sum(samples) / len(samples)
+
+    print(f"Python: {sys.executable}")
+    print(f"Samples (s): {[round(s, 4) for s in samples]}")
+    print(f"Average cold import (s): {avg:.4f}")
+    print()
+
+    rows = _parse_importtime(_importtime_output(env))
+    print(f"Top {min(args.top, len(rows))} cumulative importtime rows (us):")
+    for cumulative, module in rows[: args.top]:
+        print(f"{cumulative:>8}  {module}")
+
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

src/openai/__init__.py

@@ -6,7 +6,6 @@
 import typing as _t
 from typing_extensions import override
 
-from . import types
 from ._types import NOT_GIVEN, Omit, NoneType, NotGiven, Transport, ProxiesTypes, omit, not_given
 from ._utils import file_from_path
 from ._client import Client, OpenAI, Stream, Timeout, Transport, AsyncClient, AsyncOpenAI, AsyncStream, RequestOptions
@@ -87,14 +86,23 @@
 if not _t.TYPE_CHECKING:
     from ._utils._resources_proxy import resources as resources
 
-from .lib import azure as _azure, pydantic_function_tool as pydantic_function_tool
+if _t.TYPE_CHECKING:
+    from . import types as types
+    from .lib.azure import AzureADTokenProvider, AzureOpenAI, AsyncAzureOpenAI
+else:
+    from ._utils._proxy import LazyProxy as _LazyProxy
+
+    class _TypesProxy(_LazyProxy[_t.Any]):
+        @override
+        def __load__(self) -> _t.Any:
+            import importlib
+
+            return importlib.import_module("openai.types")
+
+    types = _TypesProxy().__as_proxied__()
+
 from .version import VERSION as VERSION
-from .lib.azure import AzureOpenAI as AzureOpenAI, AsyncAzureOpenAI as AsyncAzureOpenAI
 from .lib._old_api import *
-from .lib.streaming import (
-    AssistantEventHandler as AssistantEventHandler,
-    AsyncAssistantEventHandler as AsyncAssistantEventHandler,
-)
 
 _setup_logging()
 
@@ -105,12 +113,83 @@
 __locals = locals()
 for __name in __all__:
     if not __name.startswith("__"):
+        __obj = __locals.get(__name)
+        if __obj is None:
+            continue
         try:
-            __locals[__name].__module__ = "openai"
+            __obj.__module__ = "openai"
         except (TypeError, AttributeError):
             # Some of our exported symbols are builtins which we can't set attributes for.
             pass
 
+
+def _is_truthy_env_var(name: str) -> bool:
+    value = _os.environ.get(name, "")
+    return value not in ("", "0", "false", "False")
+
+
+def _lazy_azure_openai() -> object:
+    from .lib.azure import AzureOpenAI
+
+    return AzureOpenAI
+
+
+def _lazy_async_azure_openai() -> object:
+    from .lib.azure import AsyncAzureOpenAI
+
+    return AsyncAzureOpenAI
+
+
+def _lazy_pydantic_function_tool() -> object:
+    from .lib._tools import pydantic_function_tool
+
+    return pydantic_function_tool
+
+
+def _lazy_assistant_event_handler() -> object:
+    from .lib.streaming import AssistantEventHandler
+
+    return AssistantEventHandler
+
+
+def _lazy_async_assistant_event_handler() -> object:
+    from .lib.streaming import AsyncAssistantEventHandler
+
+    return AsyncAssistantEventHandler
+
+
+_LAZY_EXPORTS: dict[str, _t.Callable[[], object]] = {
+    "AzureOpenAI": _lazy_azure_openai,
+    "AsyncAzureOpenAI": _lazy_async_azure_openai,
+    "pydantic_function_tool": _lazy_pydantic_function_tool,
+    "AssistantEventHandler": _lazy_assistant_event_handler,
+    "AsyncAssistantEventHandler": _lazy_async_assistant_event_handler,
+}
+
+
+def __getattr__(name: str) -> object:
+    if name in _LAZY_EXPORTS:
+        value = _LAZY_EXPORTS[name]()
+        globals()[name] = value
+        return value
+
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+def _resolve_eager_imports() -> None:
+    if not _is_truthy_env_var("OPENAI_EAGER_IMPORT"):
+        return
+
+    import importlib
+
+    # Resolve all lazy exports up-front in eager mode to catch import failures in CI/dev.
+    globals()["types"] = importlib.import_module("openai.types")
+    for name in _LAZY_EXPORTS:
+        __getattr__(name)
+
+
+_resolve_eager_imports()
+
 # ------ Module level client ------
 import typing as _t
 import typing_extensions as _te
@@ -149,7 +228,7 @@
 
 azure_ad_token: str | None = _os.environ.get("AZURE_OPENAI_AD_TOKEN")
 
-azure_ad_token_provider: _azure.AzureADTokenProvider | None = None
+azure_ad_token_provider: AzureADTokenProvider | None = None
 
 
 class _ModuleClient(OpenAI):
@@ -268,8 +347,25 @@ def _client(self, value: _httpx.Client) -> None:  # type: ignore
         http_client = value
 
 
-class _AzureModuleClient(_ModuleClient, AzureOpenAI):  # type: ignore
-    ...
+def _create_azure_module_client_class() -> type[OpenAI]:
+    from .lib.azure import AzureOpenAI
+
+    class _AzureModuleClient(_ModuleClient, AzureOpenAI):  # type: ignore
+        ...
+
+    return _AzureModuleClient
+
+
+_AZURE_MODULE_CLIENT_CLASS: type[OpenAI] | None = None
+
+
+def _azure_module_client_class() -> type[OpenAI]:
+    global _AZURE_MODULE_CLIENT_CLASS
+
+    if _AZURE_MODULE_CLIENT_CLASS is None:
+        _AZURE_MODULE_CLIENT_CLASS = _create_azure_module_client_class()
+
+    return _AZURE_MODULE_CLIENT_CLASS
 
 
 class _AmbiguousModuleClientUsageError(OpenAIError):
@@ -332,7 +428,7 @@ def _load_client() -> OpenAI:  # type: ignore[reportUnusedFunction]
                 api_type = "openai"
 
         if api_type == "azure":
-            _client = _AzureModuleClient(  # type: ignore
+            _client = _azure_module_client_class()(  # type: ignore
                 api_version=api_version,
                 azure_endpoint=azure_endpoint,
                 api_key=api_key,

tests/lib/test_import_surface_live.py

@@ -0,0 +1,45 @@
+from __future__ import annotations
+
+import os
+import sys
+import importlib
+
+import pytest
+
+
+def _openai_modules() -> dict[str, object]:
+    return {name: mod for name, mod in sys.modules.items() if name == "openai" or name.startswith("openai.")}
+
+
+def _restore_openai_modules(original_modules: dict[str, object]) -> None:
+    for name in list(sys.modules):
+        if name == "openai" or name.startswith("openai."):
+            sys.modules.pop(name, None)
+    sys.modules.update(original_modules)
+
+
+@pytest.mark.skipif(os.environ.get("OPENAI_LIVE") != "1", reason="requires OPENAI_LIVE=1")
+@pytest.mark.skipif(not os.environ.get("OPENAI_API_KEY"), reason="requires OPENAI_API_KEY")
+def test_eager_import_with_live_token_allows_real_request(monkeypatch) -> None:
+    # Exercise eager mode in a real SDK flow behind explicit live-test flags.
+    monkeypatch.setenv("OPENAI_EAGER_IMPORT", "1")
+    original_modules = _openai_modules()
+
+    for name in original_modules:
+        sys.modules.pop(name, None)
+
+    client = None
+    try:
+        openai = importlib.import_module("openai")
+
+        assert "openai.types" in sys.modules
+        assert "openai.lib.azure" in sys.modules
+        assert "AzureOpenAI" in openai.__dict__
+
+        client = openai.OpenAI(timeout=20.0)
+        page = client.models.list()
+        assert page.data is not None
+    finally:
+        if client is not None:
+            client.close()
+        _restore_openai_modules(original_modules)

tests/test_import_surface.py

@@ -0,0 +1,54 @@
+from __future__ import annotations
+
+import importlib
+import sys
+
+
+def _openai_modules() -> dict[str, object]:
+    return {name: mod for name, mod in sys.modules.items() if name == "openai" or name.startswith("openai.")}
+
+
+def _restore_openai_modules(original_modules: dict[str, object]) -> None:
+    for name in list(sys.modules):
+        if name == "openai" or name.startswith("openai."):
+            sys.modules.pop(name, None)
+    sys.modules.update(original_modules)
+
+
+def test_openai_azure_is_lazy_imported(monkeypatch) -> None:
+    monkeypatch.delenv("OPENAI_EAGER_IMPORT", raising=False)
+    original_modules = _openai_modules()
+
+    for name in original_modules:
+        sys.modules.pop(name, None)
+
+    try:
+        openai = importlib.import_module("openai")
+
+        assert "openai.lib.azure" not in sys.modules
+
+        assert openai.AzureOpenAI is not None
+        assert "openai.lib.azure" in sys.modules
+    finally:
+        _restore_openai_modules(original_modules)
+
+
+def test_openai_eager_import_resolves_lazy_exports(monkeypatch) -> None:
+    original_modules = _openai_modules()
+    monkeypatch.setenv("OPENAI_EAGER_IMPORT", "1")
+
+    for name in original_modules:
+        sys.modules.pop(name, None)
+
+    try:
+        openai = importlib.import_module("openai")
+
+        assert "openai.types" in sys.modules
+        assert "openai.lib.azure" in sys.modules
+        assert "AzureOpenAI" in openai.__dict__
+        assert "AsyncAzureOpenAI" in openai.__dict__
+        assert "pydantic_function_tool" in openai.__dict__
+        assert "AssistantEventHandler" in openai.__dict__
+        assert "AsyncAssistantEventHandler" in openai.__dict__
+    finally:
+        _restore_openai_modules(original_modules)