diff --git a/scripts/bench_import.py b/scripts/bench_import.py new file mode 100644 index 0000000000..3d63cfb7f5 --- /dev/null +++ b/scripts/bench_import.py @@ -0,0 +1,96 @@ +#!/usr/bin/env python3 +from __future__ import annotations + +import os +import sys +import time +import argparse +import subprocess +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()) diff --git a/src/openai/__init__.py b/src/openai/__init__.py index 3786d106cb..528eca863e 100644 --- a/src/openai/__init__.py +++ b/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 @@ -39,7 +38,6 @@ from ._base_client import DefaultHttpxClient, DefaultAioHttpClient, DefaultAsyncHttpxClient from ._utils._logs import setup_logging as _setup_logging from ._legacy_response import HttpxBinaryResponseContent as HttpxBinaryResponseContent -from .types.websocket_reconnection import ReconnectingEvent, ReconnectingOverrides __all__ = [ "types", @@ -98,15 +96,15 @@ if not _t.TYPE_CHECKING: from ._utils._resources_proxy import resources as resources -from .lib import azure as _azure, bedrock as _bedrock, pydantic_function_tool as pydantic_function_tool +from . import types as types + +if _t.TYPE_CHECKING: + from .lib.azure import AzureOpenAI as AzureOpenAI, AsyncAzureOpenAI as AsyncAzureOpenAI, AzureADTokenProvider + from .lib.bedrock import BedrockOpenAI, AsyncBedrockOpenAI, BedrockTokenProvider + from .types.websocket_reconnection import ReconnectingEvent, ReconnectingOverrides + from .version import VERSION as VERSION -from .lib.azure import AzureOpenAI as AzureOpenAI, AsyncAzureOpenAI as AsyncAzureOpenAI -from .lib.bedrock import BedrockOpenAI as BedrockOpenAI, AsyncBedrockOpenAI as AsyncBedrockOpenAI from .lib._old_api import * -from .lib.streaming import ( - AssistantEventHandler as AssistantEventHandler, - AsyncAssistantEventHandler as AsyncAssistantEventHandler, -) _setup_logging() @@ -117,12 +115,108 @@ __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_bedrock_openai() -> object: + from .lib.bedrock import BedrockOpenAI + + return BedrockOpenAI + + +def _lazy_async_bedrock_openai() -> object: + from .lib.bedrock import AsyncBedrockOpenAI + + return AsyncBedrockOpenAI + + +def _lazy_reconnecting_event() -> object: + from .types.websocket_reconnection import ReconnectingEvent + + return ReconnectingEvent + + +def _lazy_reconnecting_overrides() -> object: + from .types.websocket_reconnection import ReconnectingOverrides + + return ReconnectingOverrides + + +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, + "BedrockOpenAI": _lazy_bedrock_openai, + "AsyncBedrockOpenAI": _lazy_async_bedrock_openai, + "ReconnectingEvent": _lazy_reconnecting_event, + "ReconnectingOverrides": _lazy_reconnecting_overrides, + "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 + + # Resolve all lazy exports up-front in eager mode to catch import failures in CI/dev. + for name in _LAZY_EXPORTS: + __getattr__(name) + + +_resolve_eager_imports() + # ------ Module level client ------ import typing as _t import typing_extensions as _te @@ -163,11 +257,11 @@ 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 _bedrock_api_key: str | None = None -bedrock_token_provider: _bedrock.BedrockTokenProvider | None = None +bedrock_token_provider: BedrockTokenProvider | None = None class _ModuleClient(OpenAI): @@ -297,21 +391,55 @@ 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 + ... -class _BedrockModuleClient(_ModuleClient, BedrockOpenAI): # type: ignore - @property # type: ignore - @override - def api_key(self) -> str | None: - return api_key if api_key is not None else _bedrock_api_key + return _AzureModuleClient - @api_key.setter # type: ignore - def api_key(self, value: str | None) -> None: # type: ignore - global _bedrock_api_key - _bedrock_api_key = value +_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 + + +def _create_bedrock_module_client_class() -> type[OpenAI]: + from .lib.bedrock import BedrockOpenAI + + class _BedrockModuleClient(_ModuleClient, BedrockOpenAI): # type: ignore + @property # type: ignore + @override + def api_key(self) -> str | None: + return api_key if api_key is not None else _bedrock_api_key + + @api_key.setter # type: ignore + def api_key(self, value: str | None) -> None: # type: ignore + global _bedrock_api_key + + _bedrock_api_key = value + + return _BedrockModuleClient + + +_BEDROCK_MODULE_CLIENT_CLASS: type[OpenAI] | None = None + + +def _bedrock_module_client_class() -> type[OpenAI]: + global _BEDROCK_MODULE_CLIENT_CLASS + + if _BEDROCK_MODULE_CLIENT_CLASS is None: + _BEDROCK_MODULE_CLIENT_CLASS = _create_bedrock_module_client_class() + + return _BEDROCK_MODULE_CLIENT_CLASS class _AmbiguousModuleClientUsageError(OpenAIError): @@ -374,7 +502,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, @@ -391,7 +519,7 @@ def _load_client() -> OpenAI: # type: ignore[reportUnusedFunction] return _client if api_type == "amazon-bedrock": - _client = _BedrockModuleClient( # type: ignore + _client = _bedrock_module_client_class()( # type: ignore api_key=api_key, bedrock_token_provider=bedrock_token_provider, organization=organization, diff --git a/tests/lib/test_import_surface_live.py b/tests/lib/test_import_surface_live.py new file mode 100644 index 0000000000..4eb589ac1d --- /dev/null +++ b/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) diff --git a/tests/test_import_surface.py b/tests/test_import_surface.py new file mode 100644 index 0000000000..197e50e222 --- /dev/null +++ b/tests/test_import_surface.py @@ -0,0 +1,54 @@ +from __future__ import annotations + +import sys +import importlib + + +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)