fix: decouple oa-vm from the ML training stack via lazy package imports

openadapt_evals/__init__.py

@@ -68,61 +68,80 @@
 except PackageNotFoundError:
     __version__ = "0.0.0+unknown"
 
-# Import from canonical locations
-from openadapt_evals.agents import (
-    BenchmarkAgent,
-    RandomAgent,
-    ScriptedAgent,
-    SmartMockAgent,
-    ApiAgent,
-    RetrievalAugmentedAgent,
-    DemoGuidedAgent,
-    action_to_string,
-    format_accessibility_tree,
-    parse_action_response,
-)
-from openadapt_evals.demo_library import DemoLibrary
-from openadapt_evals.adapters import (
-    BenchmarkAction,
-    BenchmarkAdapter,
-    BenchmarkObservation,
-    BenchmarkResult,
-    BenchmarkTask,
-    StaticDatasetAdapter,
-    UIElement,
-    WAAAdapter,
-    WAAConfig,
-    WAAMockAdapter,
-    WAALiveAdapter,
-    WAALiveConfig,
-)
-from openadapt_evals.benchmarks import (
-    EvaluationConfig,
-    compute_domain_metrics,
-    compute_metrics,
-    evaluate_agent_on_benchmark,
-    generate_benchmark_viewer,
-    ExecutionTraceCollector,
-    LiveEvaluationTracker,
-    save_execution_trace,
-)
-from openadapt_evals.evaluation.verifier_registry import (
-    TaskVerifierRegistry,
-    VerificationResult,
-)
-from openadapt_evals.task_config import evaluate_milestones_screenshot
-
-# Lazy imports for optional dependencies
+# Public API is resolved lazily (PEP 562). Importing this package must stay
+# cheap: `oa-vm` (VM lifecycle) and other light entry points import
+# `openadapt_evals` transitively, and eagerly pulling in `agents` drags the
+# whole ML training stack (transformers/peft) into commands that never touch a
+# model — which also made `oa-vm` crash outright under a NumPy 2 / stale
+# transformers environment. Each name below imports its submodule only on first
+# access, then caches it as a module global so later lookups are free.
+_LAZY_IMPORTS = {
+    # agents
+    "BenchmarkAgent": "openadapt_evals.agents",
+    "RandomAgent": "openadapt_evals.agents",
+    "ScriptedAgent": "openadapt_evals.agents",
+    "SmartMockAgent": "openadapt_evals.agents",
+    "ApiAgent": "openadapt_evals.agents",
+    "RetrievalAugmentedAgent": "openadapt_evals.agents",
+    "DemoGuidedAgent": "openadapt_evals.agents",
+    "PolicyAgent": "openadapt_evals.agents",
+    "action_to_string": "openadapt_evals.agents",
+    "format_accessibility_tree": "openadapt_evals.agents",
+    "parse_action_response": "openadapt_evals.agents",
+    # demo library
+    "DemoLibrary": "openadapt_evals.demo_library",
+    # adapters
+    "BenchmarkAction": "openadapt_evals.adapters",
+    "BenchmarkAdapter": "openadapt_evals.adapters",
+    "BenchmarkObservation": "openadapt_evals.adapters",
+    "BenchmarkResult": "openadapt_evals.adapters",
+    "BenchmarkTask": "openadapt_evals.adapters",
+    "StaticDatasetAdapter": "openadapt_evals.adapters",
+    "UIElement": "openadapt_evals.adapters",
+    "WAAAdapter": "openadapt_evals.adapters",
+    "WAAConfig": "openadapt_evals.adapters",
+    "WAAMockAdapter": "openadapt_evals.adapters",
+    "WAALiveAdapter": "openadapt_evals.adapters",
+    "WAALiveConfig": "openadapt_evals.adapters",
+    # benchmarks
+    "EvaluationConfig": "openadapt_evals.benchmarks",
+    "compute_domain_metrics": "openadapt_evals.benchmarks",
+    "compute_metrics": "openadapt_evals.benchmarks",
+    "evaluate_agent_on_benchmark": "openadapt_evals.benchmarks",
+    "generate_benchmark_viewer": "openadapt_evals.benchmarks",
+    "ExecutionTraceCollector": "openadapt_evals.benchmarks",
+    "LiveEvaluationTracker": "openadapt_evals.benchmarks",
+    "save_execution_trace": "openadapt_evals.benchmarks",
+    # evaluation
+    "TaskVerifierRegistry": "openadapt_evals.evaluation.verifier_registry",
+    "VerificationResult": "openadapt_evals.evaluation.verifier_registry",
+    # task config
+    "evaluate_milestones_screenshot": "openadapt_evals.task_config",
+}
+
+_AZURE_LAZY = ("AzureConfig", "AzureWAAOrchestrator", "AzureMLClient", "estimate_cost")
+
+
 def __getattr__(name: str):
-    """Lazy import for PolicyAgent and Azure modules."""
-    if name == "PolicyAgent":
-        from openadapt_evals.agents import PolicyAgent
-        return PolicyAgent
-    if name in ("AzureConfig", "AzureWAAOrchestrator", "AzureMLClient", "estimate_cost"):
+    """Resolve public names lazily so importing the package stays light."""
+    import importlib
+
+    module_path = _LAZY_IMPORTS.get(name)
+    if module_path is not None:
+        value = getattr(importlib.import_module(module_path), name)
+        globals()[name] = value  # cache: subsequent lookups skip __getattr__
+        return value
+    if name in _AZURE_LAZY:
         from openadapt_evals.benchmarks import azure
-        return getattr(azure, name)
+        value = getattr(azure, name)
+        globals()[name] = value
+        return value
     raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
 
+
+def __dir__():
+    return sorted(__all__)
+
 __all__ = [
     # Version
     "__version__",

openadapt_evals/benchmarks/__init__.py

@@ -51,62 +51,72 @@
     ```
 """
 
-# Import from canonical locations (agents/ and adapters/)
-from openadapt_evals.agents import (
-    BenchmarkAgent,
-    RandomAgent,
-    ScriptedAgent,
-    SmartMockAgent,
-    ApiAgent,
-    action_to_string,
-    format_accessibility_tree,
-    parse_action_response,
-)
-from openadapt_evals.adapters import (
-    BenchmarkAction,
-    BenchmarkAdapter,
-    BenchmarkObservation,
-    BenchmarkResult,
-    BenchmarkTask,
-    StaticDatasetAdapter,
-    UIElement,
-    WAAAdapter,
-    WAAConfig,
-    WAAMockAdapter,
-    WAALiveAdapter,
-    WAALiveConfig,
-)
-
-# Import evaluation utilities from runner
-from openadapt_evals.benchmarks.runner import (
-    EvaluationConfig,
-    compute_domain_metrics,
-    compute_metrics,
-    evaluate_agent_on_benchmark,
-)
-
-# Import data collection utilities
-from openadapt_evals.benchmarks.data_collection import (
-    ExecutionTraceCollector,
-    save_execution_trace,
-)
-from openadapt_evals.benchmarks.live_tracker import LiveEvaluationTracker
-
-# Import viewers
-from openadapt_evals.benchmarks.viewer import generate_benchmark_viewer
-from openadapt_evals.benchmarks.comparison_viewer import generate_comparison_viewer
-
-# Lazy imports for optional dependencies
+# Re-exports are resolved lazily (PEP 562). `openadapt_evals.benchmarks.vm_cli`
+# is the `oa-vm` entry point, so importing this package must not eagerly pull in
+# `agents` (which drags the transformers/peft training stack — and crashes
+# outright under NumPy 2 with a stale transformers). Each name maps to the
+# submodule it lives in and is imported + cached on first access.
+_LAZY_IMPORTS = {
+    # agents
+    "BenchmarkAgent": "openadapt_evals.agents",
+    "RandomAgent": "openadapt_evals.agents",
+    "ScriptedAgent": "openadapt_evals.agents",
+    "SmartMockAgent": "openadapt_evals.agents",
+    "ApiAgent": "openadapt_evals.agents",
+    "PolicyAgent": "openadapt_evals.agents",
+    "action_to_string": "openadapt_evals.agents",
+    "format_accessibility_tree": "openadapt_evals.agents",
+    "parse_action_response": "openadapt_evals.agents",
+    # adapters
+    "BenchmarkAction": "openadapt_evals.adapters",
+    "BenchmarkAdapter": "openadapt_evals.adapters",
+    "BenchmarkObservation": "openadapt_evals.adapters",
+    "BenchmarkResult": "openadapt_evals.adapters",
+    "BenchmarkTask": "openadapt_evals.adapters",
+    "StaticDatasetAdapter": "openadapt_evals.adapters",
+    "UIElement": "openadapt_evals.adapters",
+    "WAAAdapter": "openadapt_evals.adapters",
+    "WAAConfig": "openadapt_evals.adapters",
+    "WAAMockAdapter": "openadapt_evals.adapters",
+    "WAALiveAdapter": "openadapt_evals.adapters",
+    "WAALiveConfig": "openadapt_evals.adapters",
+    # evaluation runner
+    "EvaluationConfig": "openadapt_evals.benchmarks.runner",
+    "compute_domain_metrics": "openadapt_evals.benchmarks.runner",
+    "compute_metrics": "openadapt_evals.benchmarks.runner",
+    "evaluate_agent_on_benchmark": "openadapt_evals.benchmarks.runner",
+    # data collection
+    "ExecutionTraceCollector": "openadapt_evals.benchmarks.data_collection",
+    "save_execution_trace": "openadapt_evals.benchmarks.data_collection",
+    "LiveEvaluationTracker": "openadapt_evals.benchmarks.live_tracker",
+    # viewers
+    "generate_benchmark_viewer": "openadapt_evals.benchmarks.viewer",
+    "generate_comparison_viewer": "openadapt_evals.benchmarks.comparison_viewer",
+}
+
+_AZURE_LAZY = ("AzureConfig", "AzureWAAOrchestrator", "AzureMLClient", "estimate_cost")
+
+
 def __getattr__(name: str):
-    """Lazy import Azure modules (requires azure-ai-ml, azure-identity) and PolicyAgent."""
-    if name in ("AzureConfig", "AzureWAAOrchestrator", "AzureMLClient", "estimate_cost"):
+    """Resolve re-exports lazily so importing this package stays light."""
+    import importlib
+
+    module_path = _LAZY_IMPORTS.get(name)
+    if module_path is not None:
+        value = getattr(importlib.import_module(module_path), name)
+        globals()[name] = value  # cache for subsequent lookups
+        return value
+    if name in _AZURE_LAZY:
         from openadapt_evals.benchmarks import azure
-        return getattr(azure, name)
-    if name == "PolicyAgent":
-        from openadapt_evals.agents import PolicyAgent
-        return PolicyAgent
+        value = getattr(azure, name)
+        globals()[name] = value
+        return value
     raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
 
+
+def __dir__():
+    return sorted(__all__)
+
 __all__ = [
     # Base classes (from adapters)
     "BenchmarkAdapter",

tests/test_lightweight_install.py

@@ -179,3 +179,60 @@ def test_demo_library_import(self):
     def test_benchmark_action_import(self):
         from openadapt_evals.adapters.base import BenchmarkAction
         assert callable(BenchmarkAction)
+
+
+# ---------------------------------------------------------------------------
+# The oa-vm entry point must not import the ML training stack
+# ---------------------------------------------------------------------------
+
+
+class TestVmCliStaysLight:
+    """`oa-vm` (VM lifecycle) must import without the ML/vision stack.
+
+    Regression guard: the ``oa-vm`` console script imports
+    ``openadapt_evals.benchmarks.vm_cli``, which triggers the package
+    ``__init__`` modules. Those used to eagerly import ``agents`` -> transformers
+    / peft, which (a) is dead weight for VM management and (b) crashed ``oa-vm``
+    outright under a NumPy 2 / stale-transformers environment. The package
+    ``__init__``s now resolve those names lazily (PEP 562); this test pins that
+    importing the CLI path pulls in none of the heavy modules.
+
+    Run in a subprocess with a fresh interpreter so that heavy modules imported
+    by other tests in this process cannot mask a regression.
+    """
+
+    def test_importing_vm_cli_does_not_import_ml_stack(self):
+        import subprocess
+
+        heavy = ["transformers", "peft", "torch", "openadapt_evals.agents"]
+        code = (
+            "import sys\n"
+            "import openadapt_evals\n"
+            "import openadapt_evals.benchmarks\n"
+            "import openadapt_evals.benchmarks.vm_cli\n"
+            f"heavy = {heavy!r}\n"
+            "leaked = [m for m in heavy if m in sys.modules]\n"
+            "print('LEAKED:' + ','.join(leaked))\n"
+            "sys.exit(1 if leaked else 0)\n"
+        )
+        result = subprocess.run(
+            [sys.executable, "-c", code],
+            capture_output=True,
+            text=True,
+        )
+        assert result.returncode == 0, (
+            "oa-vm import path pulled in heavy modules: "
+            f"{result.stdout.strip()} {result.stderr.strip()}"
+        )
+
+    def test_lazy_public_api_still_resolves(self):
+        # The laziness must not break the re-exported public API.
+        import openadapt_evals as oe
+        from openadapt_evals.adapters import WAAAdapter as direct_adapter
+
+        assert oe.WAAAdapter is direct_adapter  # lazy value is the real object
+        assert callable(oe.compute_metrics)
+        from openadapt_evals.benchmarks import WAAMockAdapter, EvaluationConfig
+
+        assert WAAMockAdapter is not None
+        assert EvaluationConfig is not None