Performance [P3]: Lazily load sandbox runtime API exports

durabletask-azuremanaged/CHANGELOG.md

@@ -13,6 +13,10 @@ adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
   result, credential failures (for example an unavailable managed identity or a misconfigured
   `DefaultAzureCredential`) now surface from the first request rather than from the constructor.
   The exception type and message are unchanged; only the timing differs.
+- Importing `durabletask.azuremanaged.preview.sandboxes` no longer loads the sandbox worker
+  runtime or `azure-identity` up front. The package's public names are now resolved on first
+  use, which roughly halves import cost for callers that only declare sandbox worker profiles
+  or use the sandbox client. Exported names, `__all__`, and import paths are unchanged.
 - `FailureDetails.error_type` now carries the fully-qualified type name (e.g. `durabletask.task.TaskFailedError`) instead of the bare class name, and the new `FailureDetails.is_caused_by()` helper is available (both inherited from durabletask). See the core `durabletask` changelog for details, including the breaking-change notes.
 - Improved async access token refresh concurrency handling to avoid duplicate
   refresh operations under concurrent access, matching the existing sync

durabletask-azuremanaged/durabletask/azuremanaged/preview/sandboxes/__init__.py

@@ -13,14 +13,22 @@
         SandboxWorker,
         SandboxActivitiesClient,
     )
+
+The exports below are resolved lazily on first attribute access, so importing
+this package does not load the sandbox worker runtime (and its Azure Identity
+and gRPC dependencies) unless those APIs are actually used.
 """
 
-from durabletask.azuremanaged.preview.sandboxes.client import SandboxActivitiesClient
-from durabletask.azuremanaged.preview.sandboxes.helpers import SandboxActivity
-from durabletask.azuremanaged.preview.sandboxes.worker_profiles import SandboxWorkerProfile
-from durabletask.azuremanaged.preview.sandboxes.worker_profiles import SandboxWorkerProfileOptions
-from durabletask.azuremanaged.preview.sandboxes.worker_profiles import sandbox_worker_profile
-from durabletask.azuremanaged.preview.sandboxes.worker import SandboxWorker
+from importlib import import_module
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from durabletask.azuremanaged.preview.sandboxes.client import SandboxActivitiesClient
+    from durabletask.azuremanaged.preview.sandboxes.helpers import SandboxActivity
+    from durabletask.azuremanaged.preview.sandboxes.worker_profiles import SandboxWorkerProfile
+    from durabletask.azuremanaged.preview.sandboxes.worker_profiles import SandboxWorkerProfileOptions
+    from durabletask.azuremanaged.preview.sandboxes.worker_profiles import sandbox_worker_profile
+    from durabletask.azuremanaged.preview.sandboxes.worker import SandboxWorker
 
 __all__ = [
     "SandboxWorker",
@@ -30,3 +38,43 @@
     "SandboxActivitiesClient",
     "sandbox_worker_profile",
 ]
+
+# Public export name -> submodule of this package that defines it.
+_LAZY_EXPORTS: dict[str, str] = {
+    "SandboxWorker": "worker",
+    "SandboxActivity": "helpers",
+    "SandboxWorkerProfile": "worker_profiles",
+    "SandboxWorkerProfileOptions": "worker_profiles",
+    "SandboxActivitiesClient": "client",
+    "sandbox_worker_profile": "worker_profiles",
+}
+
+# Submodules that eager imports previously bound as attributes of this package.
+# They remain reachable through attribute access without an explicit import.
+_LAZY_SUBMODULES: frozenset[str] = frozenset({
+    "client",
+    "helpers",
+    "profile_builder",
+    "transport",
+    "worker",
+    "worker_messages",
+    "worker_profiles",
+})
+
+
+def __getattr__(name: str) -> Any:
+    """Import public sandbox exports on first access (PEP 562)."""
+    submodule = _LAZY_EXPORTS.get(name)
+    if submodule is not None:
+        value = getattr(import_module(f".{submodule}", __name__), name)
+    elif name in _LAZY_SUBMODULES:
+        value = import_module(f".{name}", __name__)
+    else:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+    globals()[name] = value
+    return value
+
+
+def __dir__() -> list[str]:
+    return sorted(set(globals()) | set(__all__) | _LAZY_SUBMODULES)

tests/durabletask-azuremanaged/test_sandboxes_extension.py

@@ -3,9 +3,12 @@
 
 import inspect
 import random
+import subprocess
+import sys
 import threading
 
 import grpc
+import pytest
 from azure.core.credentials import AccessToken
 
 import durabletask.azuremanaged.preview.sandboxes as sandbox
@@ -69,6 +72,113 @@ def test_public_sandbox_package_exports_customer_entrypoints_only() -> None:
     assert not hasattr(SandboxWorker, f"_configure_{legacy_prefix}_activity_filters")
 
 
+# Runs in a subprocess because this test module imports the sandbox worker at
+# module scope, which would otherwise hide the lazy-import behavior.
+_LAZY_EXPORT_PROBE = """
+import sys
+
+import durabletask.azuremanaged.preview.sandboxes as sandbox
+
+eager = [
+    name
+    for name in (
+        "durabletask.azuremanaged.preview.sandboxes.client",
+        "durabletask.azuremanaged.preview.sandboxes.transport",
+        "durabletask.azuremanaged.preview.sandboxes.worker",
+        "azure.identity",
+    )
+    if name in sys.modules
+]
+assert not eager, f"sandbox package eagerly imported: {eager}"
+
+expected = [
+    "SandboxWorker",
+    "SandboxActivity",
+    "SandboxWorkerProfile",
+    "SandboxWorkerProfileOptions",
+    "SandboxActivitiesClient",
+    "sandbox_worker_profile",
+]
+assert sandbox.__all__ == expected, sandbox.__all__
+
+listing = dir(sandbox)
+assert not [name for name in expected if name not in listing], listing
+assert not [name for name in ("client", "helpers", "worker", "worker_profiles") if name not in listing], listing
+
+from durabletask.azuremanaged.preview.sandboxes.client import SandboxActivitiesClient
+from durabletask.azuremanaged.preview.sandboxes.helpers import SandboxActivity
+from durabletask.azuremanaged.preview.sandboxes.worker import SandboxWorker
+from durabletask.azuremanaged.preview.sandboxes.worker_profiles import SandboxWorkerProfile
+from durabletask.azuremanaged.preview.sandboxes.worker_profiles import SandboxWorkerProfileOptions
+from durabletask.azuremanaged.preview.sandboxes.worker_profiles import sandbox_worker_profile
+
+assert sandbox.SandboxWorker is SandboxWorker
+assert sandbox.SandboxActivity is SandboxActivity
+assert sandbox.SandboxWorkerProfile is SandboxWorkerProfile
+assert sandbox.SandboxWorkerProfileOptions is SandboxWorkerProfileOptions
+assert sandbox.SandboxActivitiesClient is SandboxActivitiesClient
+assert sandbox.sandbox_worker_profile is sandbox_worker_profile
+
+# Submodules stay reachable as package attributes, as they were while the
+# package imported them eagerly.
+assert sandbox.worker is sys.modules["durabletask.azuremanaged.preview.sandboxes.worker"]
+assert sandbox.helpers is sys.modules["durabletask.azuremanaged.preview.sandboxes.helpers"]
+"""
+
+
+_STAR_IMPORT_PROBE = """
+from durabletask.azuremanaged.preview.sandboxes import *  # noqa: F403
+
+import durabletask.azuremanaged.preview.sandboxes as sandbox
+
+star_names = {name: value for name, value in globals().items() if not name.startswith("_")}
+for name in sandbox.__all__:
+    assert name in star_names, f"missing star export: {name}"
+    assert star_names[name] is getattr(sandbox, name), name
+
+# __all__ still gates the star import, so lazy-import plumbing stays private.
+assert "import_module" not in star_names
+assert "_LAZY_EXPORTS" not in globals()
+"""
+
+
+_ISOLATED_PROBE_TIMEOUT_SECONDS = 60
+
+
+def _run_isolated(source: str) -> None:
+    try:
+        result = subprocess.run(
+            [sys.executable, "-c", source],
+            capture_output=True,
+            text=True,
+            check=False,
+            timeout=_ISOLATED_PROBE_TIMEOUT_SECONDS)
+    except subprocess.TimeoutExpired as exc:
+        pytest.fail(
+            f"isolated import probe did not complete within "
+            f"{_ISOLATED_PROBE_TIMEOUT_SECONDS}s and was terminated. "
+            f"This usually means the import deadlocked.\n"
+            f"stdout: {exc.stdout or ''}\n"
+            f"stderr: {exc.stderr or ''}")
+    assert result.returncode == 0, result.stdout + result.stderr
+
+
+def test_sandbox_package_defers_worker_runtime_imports() -> None:
+    _run_isolated(_LAZY_EXPORT_PROBE)
+
+
+def test_sandbox_package_star_import_matches_all() -> None:
+    _run_isolated(_STAR_IMPORT_PROBE)
+
+
+def test_sandbox_package_raises_attribute_error_for_unknown_attributes() -> None:
+    with pytest.raises(AttributeError, match="has no attribute 'NotARealSandboxExport'"):
+        getattr(sandbox, "NotARealSandboxExport")
+
+    assert not hasattr(sandbox, "NotARealSandboxExport")
+    assert all(hasattr(sandbox, name) for name in sandbox.__all__)
+
+
 def _sandbox_image(
     image_ref: str,
     managed_identity_client_id: str = "image-pull-client-id",