diff --git a/workers/executor/__init__.py b/workers/executor/__init__.py index 7982e4d411..f96f99b32f 100644 --- a/workers/executor/__init__.py +++ b/workers/executor/__init__.py @@ -3,10 +3,33 @@ Celery worker for running extraction executors. Dispatches ExecutionContext to registered executors and returns ExecutionResult via the Celery result backend. + +``celery_app`` resolves lazily (PEP 562). Importing it eagerly made *every* +consumer of this package pay the executor worker's full bootstrap: ``.worker`` +built the Celery app, and registration then imported ``LegacyExecutor`` and the +adapter stack behind it. That was ~9s of work the file_processing worker never +needed — it imports ``ExecutorToolShim`` and some string constants from this +package and dispatches everything else over the PG queue (UN-4136). Attribute +access is unchanged, so ``from executor import celery_app`` still works for +callers that genuinely want the app. """ -from .worker import app as celery_app +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from celery import Celery + + celery_app: Celery __all__ = [ "celery_app", ] + + +def __getattr__(name: str) -> object: + """Resolve ``celery_app`` on first access instead of at import time.""" + if name == "celery_app": + from .worker import app + + return app + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/workers/executor/executors/__init__.py b/workers/executor/executors/__init__.py index cb2b54c980..80bb6c257b 100644 --- a/workers/executor/executors/__init__.py +++ b/workers/executor/executors/__init__.py @@ -1,16 +1,133 @@ """Executor implementations package. -Importing this module triggers ``@ExecutorRegistry.register`` for all -bundled executors and discovers cloud executors via entry points. +Registration is **explicit**: call :func:`register_all`. It used to run as a +side effect of importing this package, which meant that importing *any* +submodule — including ``executor.executors.constants``, which itself imports +nothing but ``enum`` — dragged in ``LegacyExecutor`` and the entire adapter +stack behind it. That cost ~9s, and the file_processing worker paid it once per +forked child just to read a few string constants (UN-4136). + +``executor/tasks.py`` is the only *production* caller: ``workers/worker.py`` +exec-loads that file by path for both the Celery and PG executor roles. +``executor/worker.py`` reaches it transitively, by importing ``executor.tasks`` +— that import binds the task definitions to its app and populates the registry, +so it is load-bearing despite its ``noqa: F401``. The app it builds is not on +any deployed path; today only the tests use it. The test suite calls +:func:`register_all` directly. """ -from executor.executors.legacy_executor import LegacyExecutor -from executor.executors.plugins.loader import ExecutorPluginLoader +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from executor.executors.legacy_executor import LegacyExecutor # noqa: TCH004 + +#: Cloud entry point names; None when discovery has not run, or ran and +#: raised. Doubles as the re-entrancy latch — see :func:`register_all`. +_cloud_executors: list[str] | None = None + + +def register_all() -> list[str]: + """Import the executor modules once per process, registering each of them. + + ``LegacyExecutor`` and every cloud executor carry + ``@ExecutorRegistry.register``, which fires when their module is first + imported. That is the whole mechanism, and it bounds the guarantee: this + populates ``ExecutorRegistry`` **in a fresh process**, and cannot repopulate + it if something empties it afterwards, because the second import is a + ``sys.modules`` hit and the decorator does not run again. Production never + clears the registry. Several test modules do, and most of them do not put it + back; ``tests/test_legacy_executor_scaffold.py`` restores what it cleared. + A module that re-registers must not call ``ExecutorRegistry.register`` while + the name is already present — either clear immediately before, or guard on + ``"legacy" not in ExecutorRegistry.list_executors()``. That rule is hand-copied + across the test modules in several spellings; folding it into a shared test + helper is deliberately left out of this change. + + Idempotent, and safe to re-enter: a cloud plugin whose own import graph + reaches this function during ``ep.load()`` will not restart discovery. + + Returns: + The cloud executor entry point names, as a fresh copy each time so a + caller cannot mutate the latched state. + + The list can under-report the registry at any length, including zero. + An empty one does not distinguish its causes: no cloud plugins are + installed (the OSS case); every one of them failed to import, because + ``ExecutorPluginLoader.discover_executors`` catches per-entry-point + failures and only logs a warning, so a broken plugin wheel boots clean + (a follow-up to make that aggregate loud is recorded on UN-4136); the + caller is *inside* discovery, having re-entered while the latch is still + the empty placeholder; or a failed discovery was re-run and the plugin + that raised is live in ``ExecutorRegistry`` but absent here — see the + handler below. + """ + global _cloud_executors + + from executor.executors.legacy_executor import LegacyExecutor # noqa: F401 + + if _cloud_executors is None: + from executor.executors.plugins.loader import ExecutorPluginLoader + + # Latch BEFORE discovering. ``ep.load()`` executes third-party code, and + # a plugin that reaches back into this function would otherwise find the + # latch still unset and restart the entry point loop, nesting once per + # level. The import-side-effect version this replaced got that safety + # free from ``sys.modules``. + # + try: + _cloud_executors = [] + _cloud_executors = ExecutorPluginLoader.discover_executors() + except BaseException: + # Re-raise unchanged, and put the latch back to its un-armed value, + # so a later call re-runs discovery instead of reporting "no cloud + # executors" as though this one had succeeded. That does re-arm a + # retry for a later caller — no production caller survives the + # re-raise to make one (see the module docstring), so today only + # tests reach it. + # + # ``BaseException`` rather than ``Exception``: see + # ``tests/test_plugin_loader.py::test_failed_discovery_un_arms_the_latch``, + # parametrised over every class that reaches here and pinning them + # executably. + # + # A re-run is not free. The plugin whose module was mid-import when + # the exception escaped is evicted from ``sys.modules`` but stays in + # the registry if its decorator had already fired, so re-importing + # it raises a duplicate-name ``ValueError`` that + # ``discover_executors`` swallows into a warning — live in the + # registry, absent from the returned list. Re-running discovery has + # always had that property, including in the import-side-effect + # version this replaced. + _cloud_executors = None + raise + + return list(_cloud_executors) + + +def __getattr__(name: str) -> object: + """Keep ``from executor.executors import LegacyExecutor`` working. + + The class used to sit in this namespace as a side effect of the eager + import. Resolving it on access preserves that spelling without reviving the + cost for everyone who imports a sibling submodule. + """ + if name == "LegacyExecutor": + from executor.executors.legacy_executor import LegacyExecutor + + return LegacyExecutor + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + +def _reset_discovery_for_tests() -> None: + """Re-arm entry point discovery so a test can observe it running again. + + Discovery only. This does **not** restore the bundled executor: that is + registered by module import, which cannot be made to happen twice in one + process. A test that clears ``ExecutorRegistry`` must re-register + explicitly, as ``tests/test_legacy_executor_scaffold.py`` does. + """ + global _cloud_executors + _cloud_executors = None -# Discover and register cloud executors installed via entry points. -# Each cloud executor class is decorated with @ExecutorRegistry.register, -# so importing it (via ep.load()) is enough to register it. -# If no cloud plugins are installed this returns an empty list. -_cloud_executors = ExecutorPluginLoader.discover_executors() -__all__ = ["LegacyExecutor"] +__all__ = ["LegacyExecutor", "register_all"] diff --git a/workers/executor/executors/plugins/loader.py b/workers/executor/executors/plugins/loader.py index 3f2a54c92a..2959b07b20 100644 --- a/workers/executor/executors/plugins/loader.py +++ b/workers/executor/executors/plugins/loader.py @@ -8,7 +8,7 @@ - ``unstract.executor.executors`` Executor classes that self-register via ``@ExecutorRegistry.register``. - Loaded eagerly at worker startup from ``executors/__init__.py``. + Loaded when ``executor.executors.register_all()`` runs at worker startup. """ import logging diff --git a/workers/executor/tasks.py b/workers/executor/tasks.py index 1244e60b9c..16b97c2c33 100644 --- a/workers/executor/tasks.py +++ b/workers/executor/tasks.py @@ -5,16 +5,9 @@ ExecutionOrchestrator, and returns an ExecutionResult dict. """ -# Import the executor implementations so their ``@ExecutorRegistry.register`` -# decorators run before ``execute_extraction`` can be invoked. Coupling this to -# the task module (not only the Celery ``executor/worker.py`` entrypoint) ensures -# the registry is populated wherever the task is registered — in particular the PG -# executor consumer, which bootstraps via the root-worker import that loads -# ``executor/tasks.py`` (this module) but NOT ``executor/worker.py`` where this -# import historically lived; without it the consumer hits "No executor -# registered". Import is idempotent (module cached), so the Celery entrypoint -# importing it again is harmless. -import executor.executors # noqa: E402, F401 +import logging + +from executor.executors import register_all from queue_backend import worker_task from shared.clients import UsageAPIClient from shared.enums.task_enums import TaskName @@ -25,8 +18,30 @@ from unstract.sdk1.execution.orchestrator import ExecutionOrchestrator from unstract.sdk1.execution.result import ExecutionResult +# Suppress Celery trace logging of task return values. +# The trace logger prints the full result dict on task success, which can +# contain sensitive customer data (extracted text, summaries, etc.). +# +# This lives here, not in ``executor/worker.py``, because that module is not on +# any deployed path: ``workers/worker.py`` exec-loads THIS file by path for both +# the Celery and PG executor roles, and no launcher runs ``celery -A executor``. +# It used to be reached only because ``executor/__init__.py`` eagerly imported +# ``.worker``; making that lazy (UN-4136) would otherwise have silently +# re-enabled result logging on every extraction. +logging.getLogger("celery.app.trace").setLevel(logging.WARNING) + logger = WorkerLogger.get_logger(__name__) +# Populate ``ExecutorRegistry`` so the executor implementations are registered +# before ``execute_extraction`` can be invoked. Coupling this to the task module +# (not only the Celery ``executor/worker.py`` entrypoint) ensures the registry is +# populated wherever the task is registered — in particular the PG executor +# consumer, which bootstraps via the root-worker import that loads +# ``executor/tasks.py`` (this module) but NOT ``executor/worker.py``; without it +# the consumer hits "No executor registered". This is the only call site; +# ``executor/worker.py`` reaches it by importing this module. +register_all() + _LLM_BEARING_OPS = frozenset( { "answer_prompt", diff --git a/workers/executor/worker.py b/workers/executor/worker.py index 611f074e5a..b32b0f4672 100644 --- a/workers/executor/worker.py +++ b/workers/executor/worker.py @@ -4,7 +4,6 @@ Routes execute_extraction tasks to registered executors. """ -import logging import os from queue_backend import worker_task @@ -17,11 +16,6 @@ logger = WorkerLogger.setup(WorkerType.EXECUTOR) app, config = WorkerBuilder.build_celery_app(WorkerType.EXECUTOR) -# Suppress Celery trace logging of task return values. -# The trace logger prints the full result dict on task success, which -# can contain sensitive customer data (extracted text, summaries, etc.). -logging.getLogger("celery.app.trace").setLevel(logging.WARNING) - def check_executor_health(): """Custom health check for executor worker.""" @@ -79,7 +73,8 @@ def healthcheck(self): } -# Import tasks so shared_task definitions bind to this app. -# Import executors to trigger @ExecutorRegistry.register at import time. -import executor.executors # noqa: E402, F401 +# Import tasks so shared_task definitions bind to this app. This also populates +# ExecutorRegistry: executor/tasks.py calls register_all() at module scope, and +# it is the site that does the work on every deployed path — workers/worker.py +# exec-loads that file by path, and nothing launches `celery -A executor`. import executor.tasks # noqa: E402, F401 diff --git a/workers/tests/test_executor_import_isolation.py b/workers/tests/test_executor_import_isolation.py new file mode 100644 index 0000000000..246d792edd --- /dev/null +++ b/workers/tests/test_executor_import_isolation.py @@ -0,0 +1,136 @@ +"""Regression: the file_processing worker must not boot the executor stack. + +``structure_tool_task`` dispatches all real extraction work to the executor +worker over the PG queue. From the ``executor`` package it needs only two cheap +things: ``ExecutorToolShim`` (a StreamMixin wrapper) and some string constants. + +Both used to drag in ``LegacyExecutor`` and every adapter behind it, because +``executor/__init__.py`` eagerly imported ``.worker``, which imported +``executor.executors``, which imported ``LegacyExecutor`` and ran entry-point +discovery. ``structure_tool_task`` makes those imports inside the task function, +so the cost landed on each forked child's FIRST task rather than at startup — +~9s per child, measured on staging (UN-4136). + +These pin the import graph, not a duration: a timing assertion would be flaky on +CI, while the thing that actually regresses is an eager import creeping back +into either ``__init__``. Each runs in a fresh interpreter so the rest of the +suite cannot pre-import the stack and mask the regression. +""" + +import os +import subprocess +import sys + +# Importing these must not pull the executor stack in behind them. +_FILE_PROCESSING_IMPORTS = ( + "from executor.executor_tool_shim import ExecutorToolShim", + "from executor.executors.constants import PromptServiceConstants", +) + +# The expensive modules: ``legacy_executor`` pulls the x2text adapter stack at +# module scope, and ``executor.worker`` builds the Celery app. +_MUST_NOT_LOAD = ( + "executor.worker", + "executor.executors.legacy_executor", +) + +_WORKERS_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + + +# Child snippets use ``raise SystemExit``, never a bare ``assert``: the child +# inherits the parent's environment, and under ``PYTHONOPTIMIZE``/``-O`` every +# assert is stripped — the snippet would print OK, exit 0, and the guard would +# pass having checked nothing. +def _run(snippet: str) -> subprocess.CompletedProcess: + return subprocess.run( + [sys.executable, "-c", snippet], + capture_output=True, + text=True, + env={**os.environ, "WORKER_TYPE": "file_processing"}, + cwd=_WORKERS_DIR, + ) + + +def test_file_processing_imports_do_not_load_the_executor_stack(): + """Neither import may leave the executor worker or LegacyExecutor loaded.""" + code = "\n".join( + [ + "import sys", + *_FILE_PROCESSING_IMPORTS, + f"loaded = [m for m in {_MUST_NOT_LOAD!r} if m in sys.modules]", + "if loaded: raise SystemExit(f'eagerly imported: {loaded}')", + "print('OK')", + ] + ) + result = _run(code) + assert result.returncode == 0, ( + "importing ExecutorToolShim / PromptServiceConstants pulled in the " + f"executor stack.\nstdout: {result.stdout}\nstderr: {result.stderr}" + ) + assert "OK" in result.stdout + + +def test_loading_executor_tasks_suppresses_celery_result_logging(): + """The trace logger must be muted by the module the executor actually loads. + + ``celery.app.trace`` logs the full result dict on task success, which for + ``execute_extraction`` is extracted customer document text. The + ``setLevel(WARNING)`` that mutes it used to live in ``executor/worker.py`` + and was reached only because ``executor/__init__.py`` eagerly imported + ``.worker`` — so making that import lazy silently un-suppressed it, and + every extraction would have logged ~1KB of the payload at INFO. + + Nothing else in the repo sets this level, and no launcher runs + ``celery -A executor``, so ``executor/tasks.py`` — which ``workers/worker.py`` + exec-loads by path for both executor roles — is the only module that can + carry it. A fresh interpreter is required: another test importing + ``executor.worker`` first would mask the regression. + """ + code = "\n".join( + [ + "import logging, sys", + "import executor.tasks", # exactly what the deployed executor loads + "lvl = logging.getLogger('celery.app.trace').level", + "if lvl != logging.WARNING:", + " raise SystemExit(f'trace logger not suppressed: {lvl}')", + # The suppression must not come back via the expensive import. + "if 'executor.worker' in sys.modules:", + " raise SystemExit('pulled in executor.worker')", + "print('OK')", + ] + ) + result = subprocess.run( + [sys.executable, "-c", code], + capture_output=True, + text=True, + env={**os.environ, "WORKER_TYPE": "executor"}, + cwd=_WORKERS_DIR, + ) + assert result.returncode == 0, ( + "executor.tasks did not suppress celery.app.trace result logging.\n" + f"stdout: {result.stdout}\nstderr: {result.stderr}" + ) + assert "OK" in result.stdout + + +def test_importing_the_executor_package_does_not_build_the_celery_app(): + """``import executor`` alone must stay cheap. + + ``celery_app`` is still reachable — it resolves on attribute access (PEP + 562) — so this pins laziness, not removal. + """ + code = "\n".join( + [ + "import sys, executor", + "if 'executor.worker' in sys.modules:", + " raise SystemExit('executor/__init__ built the app')", + "if executor.celery_app is None:", + " raise SystemExit('celery_app no longer resolves')", + "if 'executor.worker' not in sys.modules:", + " raise SystemExit('attribute access did not load it')", + "print('OK')", + ] + ) + result = _run(code) + assert result.returncode == 0, f"stdout: {result.stdout}\nstderr: {result.stderr}" + assert "OK" in result.stdout diff --git a/workers/tests/test_legacy_executor_extract.py b/workers/tests/test_legacy_executor_extract.py index 0711d2255a..48d7e5f114 100644 --- a/workers/tests/test_legacy_executor_extract.py +++ b/workers/tests/test_legacy_executor_extract.py @@ -44,9 +44,13 @@ def _clean_registry(): def _register_legacy(): - from executor.executors.legacy_executor import LegacyExecutor # noqa: F401 + # Guarded because the import below may be the FIRST in the process, in which + # case ``@ExecutorRegistry.register`` fires on it and registering again + # raises a duplicate-name ValueError. + from executor.executors.legacy_executor import LegacyExecutor - ExecutorRegistry.register(LegacyExecutor) + if "legacy" not in ExecutorRegistry.list_executors(): + ExecutorRegistry.register(LegacyExecutor) def _make_context(**overrides): diff --git a/workers/tests/test_legacy_executor_index.py b/workers/tests/test_legacy_executor_index.py index b8e5bdd413..95c6537ea0 100644 --- a/workers/tests/test_legacy_executor_index.py +++ b/workers/tests/test_legacy_executor_index.py @@ -36,9 +36,13 @@ def _clean_registry(): def _register_legacy(): - from executor.executors.legacy_executor import LegacyExecutor # noqa: F401 + # Guarded because the import below may be the FIRST in the process, in which + # case ``@ExecutorRegistry.register`` fires on it and registering again + # raises a duplicate-name ValueError. + from executor.executors.legacy_executor import LegacyExecutor - ExecutorRegistry.register(LegacyExecutor) + if "legacy" not in ExecutorRegistry.list_executors(): + ExecutorRegistry.register(LegacyExecutor) def _make_index_context(**overrides): diff --git a/workers/tests/test_legacy_executor_scaffold.py b/workers/tests/test_legacy_executor_scaffold.py index 4895e6ffe9..d99e9c2281 100644 --- a/workers/tests/test_legacy_executor_scaffold.py +++ b/workers/tests/test_legacy_executor_scaffold.py @@ -41,9 +41,13 @@ def _clean_registry(): def _register_legacy(): """Import executor.executors to trigger LegacyExecutor registration.""" - from executor.executors.legacy_executor import LegacyExecutor # noqa: F401 + # Guarded because the import below may be the FIRST in the process, in which + # case ``@ExecutorRegistry.register`` fires on it and registering again + # raises a duplicate-name ValueError. + from executor.executors.legacy_executor import LegacyExecutor - ExecutorRegistry.register(LegacyExecutor) + if "legacy" not in ExecutorRegistry.list_executors(): + ExecutorRegistry.register(LegacyExecutor) def _make_context(**overrides): diff --git a/workers/tests/test_line_item_extraction.py b/workers/tests/test_line_item_extraction.py index 5a46b47b1f..515f141f26 100644 --- a/workers/tests/test_line_item_extraction.py +++ b/workers/tests/test_line_item_extraction.py @@ -405,10 +405,11 @@ def test_celery_eager_chain_with_line_item_plugin( """Push a LINE_ITEM payload through the full Celery eager chain with a fake line_item plugin registered. """ - # Re-register LegacyExecutor since the autouse fixture cleared it - from executor.executors.legacy_executor import LegacyExecutor - - ExecutorRegistry.register(LegacyExecutor) + # Re-register LegacyExecutor since the autouse fixture cleared it. + # Via the module's guarded helper: a bare register() raises a + # duplicate-name ValueError when this test's import is the first in the + # process, because the decorator then fires after the fixture's clear. + _get_legacy_executor() # Register fake line_item plugin plugin_cls = _make_success_plugin( diff --git a/workers/tests/test_plugin_loader.py b/workers/tests/test_plugin_loader.py index 4ac720cd9e..b6af7e59c3 100644 --- a/workers/tests/test_plugin_loader.py +++ b/workers/tests/test_plugin_loader.py @@ -9,14 +9,18 @@ 6. text_processor.add_hex_line_numbers() 7. Queue-per-executor naming convention (QUEUE_PREFIX) 8. Protocol classes importable and runtime-checkable -9. executors/__init__.py triggers discover_executors() +9. executors.register_all() triggers discover_executors() """ +import os +import subprocess +import sys from unittest.mock import MagicMock, patch import pytest from executor.executors.plugins.loader import ExecutorPluginLoader from executor.executors.plugins.text_processor import add_hex_line_numbers + from unstract.workflow_execution.executor_rpc import QUEUE_PREFIX @@ -289,14 +293,216 @@ def run(self): assert isinstance(FakeChallenge(), ChallengeProtocol) -# ── 7. executors/__init__.py triggers discovery ───────────────────── +# ── 7. executors.register_all() triggers discovery ────────────────── class TestExecutorsInit: - def test_cloud_executors_list_exists(self): - """executors.__init__ populates _cloud_executors (empty in OSS).""" + """``register_all()``'s guarantees, each pinned against its own failure. + + These cases *write* process-global state — ``register_all()`` registers into + the ``ExecutorRegistry`` singleton, and they stamp on the module's + ``_cloud_executors`` latch — so the fixture snapshots and restores both. + The ``sys.modules`` entry for ``legacy_executor`` is not restorable and is + deliberately left. The sibling modules' guards are correct either way, so + leaving it is harmless. + + No case here asserts registry *contents* off an inherited state. Registration + is asserted in a subprocess, and the one case that needs ``legacy`` already + present arranges that itself. That is deliberate: in-process such an + assertion would pin whatever module ran before it rather than the code under + test, and the repo's test rig runs these under xdist with the default + per-test scheduler (``/tests/rig/cli.py``), so collection order is + not something to rely on. + """ + + @pytest.fixture(autouse=True) + def _isolate_global_state(self): + import executor.executors as mod + + from unstract.sdk1.execution.registry import ExecutorRegistry + + saved = dict(ExecutorRegistry._registry) + saved_cloud = mod._cloud_executors + yield + # Restore exactly, including removing keys this class added — importing + # ``legacy_executor`` registers as a side effect, and running real + # discovery would register the cloud executors too. Leaving those behind + # would collide with the sibling modules that register the same names + # themselves, since duplicate registration raises. Nothing depends on + # this class having registered ``legacy``: every consumer in the suite + # registers it explicitly or assigns it directly. + ExecutorRegistry._registry.clear() + ExecutorRegistry._registry.update(saved) + mod._cloud_executors = saved_cloud + + def test_register_all_registers_the_bundled_executor(self): + """In a fresh process, register_all() populates the registry. + + A fresh interpreter is the only honest way to assert this: registration + rides on importing ``legacy_executor``, so in-process the decorator has + already fired and the assertion would pass on whatever an earlier module + left behind rather than on the call under test. Subprocess isolation is + the same technique ``test_executor_registration.py`` uses, and for the + same reason. + """ + code = ( + "from executor.executors import register_all\n" + "from unstract.sdk1.execution.registry import ExecutorRegistry\n" + "if ExecutorRegistry.list_executors():\n" + " raise SystemExit('registry pre-populated')\n" + "register_all()\n" + "names = ExecutorRegistry.list_executors()\n" + "if names.count('legacy') != 1:\n" + " raise SystemExit(f'expected one legacy: {names}')\n" + "print('OK')\n" + ) + result = subprocess.run( + [sys.executable, "-c", code], + capture_output=True, + text=True, + env={**os.environ, "WORKER_TYPE": "executor"}, + cwd=os.path.dirname(os.path.dirname(os.path.abspath(__file__))), + ) + assert result.returncode == 0, ( + f"register_all() did not register in a fresh process.\n" + f"stdout: {result.stdout}\nstderr: {result.stderr}" + ) + assert "OK" in result.stdout + + def test_legacy_executor_resolves_through_the_package_namespace(self): + """``from executor.executors import LegacyExecutor`` must keep working. + + The class used to be a real module attribute; it is now resolved by + ``__getattr__``. No caller in this repo or the cloud plugins uses this + spelling today — they import the submodule directly — so nothing else + would catch a typo in the name comparison or the inner import path, and + the attribute it replaced could not fail this way. + """ + import executor.executors as mod + from executor.executors.legacy_executor import LegacyExecutor + + assert mod.LegacyExecutor is LegacyExecutor + + with pytest.raises(AttributeError): + mod.NoSuchExecutor + + def test_register_all_discovers_entry_points_only_once(self): + """Idempotency, asserted on the work done rather than the value returned. + + In OSS no cloud plugins are installed, so ``register_all()`` returns + ``[]`` on the first call as well as later ones — asserting on the return + value cannot tell a working latch from no latch at all. + """ + import executor.executors as mod + from executor.executors.plugins.loader import ExecutorPluginLoader + + mod._reset_discovery_for_tests() + with patch.object( + ExecutorPluginLoader, "discover_executors", return_value=["fake_cloud"] + ) as spy: + first = mod.register_all() + second = mod.register_all() + + assert spy.call_count == 1, f"discovery re-ran: {spy.call_count} calls" + assert first == ["fake_cloud"] + assert second == ["fake_cloud"], "the names must survive the latched call" + + def test_register_all_is_reentrant(self): + """A plugin that calls back into ``register_all()`` must not restart discovery. + + ``ep.load()`` executes third-party code. If the latch were set only + after discovery, a plugin whose import graph reaches this function would + re-enter with discovery un-latched and loop the entry points again, once + per level. + """ import executor.executors as mod + from executor.executors.plugins.loader import ExecutorPluginLoader + + mod._reset_discovery_for_tests() + loads = [] + + def _reentrant_discovery(): + loads.append(1) + if len(loads) < 5: + mod.register_all() # what a re-entrant ep.load() would do + return ["reentrant"] + + with patch.object( + ExecutorPluginLoader, "discover_executors", _reentrant_discovery + ): + mod.register_all() + + assert len(loads) == 1, f"discovery re-entered {len(loads)} times" + + def test_register_all_tolerates_legacy_executor_already_registered(self): + """``register_all()`` must not raise when ``legacy`` is already present. + + ``ExecutorRegistry.register`` raises ``ValueError`` on a duplicate name, + and this ordering is real: a cloud plugin imports ``LegacyExecutor`` at + module scope and is loaded by ``ep.load()`` *inside* discovery, so the + class can already be registered when the call reaches it. + + The precondition is arranged here rather than inherited — asserting on a + registry state some earlier module left behind would pin that module, + not this code. Discovery is patched for the same reason: unpatched, this + would run real ``entry_points()`` and, where the cloud wheels are + installed, register every cloud executor as a side effect of a test that + is not about them. + """ + import executor.executors as mod + from executor.executors.legacy_executor import LegacyExecutor + from executor.executors.plugins.loader import ExecutorPluginLoader + + from unstract.sdk1.execution.registry import ExecutorRegistry + + if "legacy" not in ExecutorRegistry.list_executors(): + ExecutorRegistry.register(LegacyExecutor) + + mod._reset_discovery_for_tests() + with patch.object(ExecutorPluginLoader, "discover_executors", return_value=[]): + mod.register_all() # must not raise + + assert "legacy" in ExecutorRegistry.list_executors() + + @pytest.mark.parametrize("exc", [RuntimeError, KeyboardInterrupt, SystemExit]) + def test_failed_discovery_un_arms_the_latch(self, exc): + """A discovery that raises must not leave ``[]`` latched. + + Otherwise every later call in the process reports "no cloud executors" + as though discovery had succeeded, which is indistinguishable from the + OSS case. + + Every class that reaches the handler is exercised: an ordinary + ``Exception``, which escapes ``entry_points()`` because that call runs + before any per-entry-point ``try``, and the ``BaseException``-not- + ``Exception`` classes, which are the only ones that escape + ``ep.load()``. Parametrised rather than looped so one failing arm cannot + stop the others from running. + """ + import executor.executors as mod + from executor.executors.plugins.loader import ExecutorPluginLoader + + mod._reset_discovery_for_tests() + with patch.object( + ExecutorPluginLoader, "discover_executors", side_effect=exc("boom") + ): + with pytest.raises(exc): + mod.register_all() + + assert ( + mod._cloud_executors is None + ), f"{exc.__name__} left the latch armed: {mod._cloud_executors!r}" + + def test_register_all_returns_a_copy_callers_cannot_corrupt(self): + """The latched list is module state; callers must not be able to edit it.""" + import executor.executors as mod + from executor.executors.plugins.loader import ExecutorPluginLoader + + mod._reset_discovery_for_tests() + with patch.object( + ExecutorPluginLoader, "discover_executors", return_value=["table"] + ): + first = mod.register_all() - assert hasattr(mod, "_cloud_executors") - # In pure OSS, no cloud executors are installed - assert isinstance(mod._cloud_executors, list) + first.append("MUTATED-BY-CALLER") + assert mod.register_all() == ["table"] diff --git a/workers/tests/test_plugin_migration_regression.py b/workers/tests/test_plugin_migration_regression.py index b2a5046b5b..8cdf144f9f 100644 --- a/workers/tests/test_plugin_migration_regression.py +++ b/workers/tests/test_plugin_migration_regression.py @@ -48,9 +48,13 @@ def eager_app(): def _register_legacy(): + # Guarded because the import below may be the FIRST in the process, in which + # case ``@ExecutorRegistry.register`` fires on it and registering again + # raises a duplicate-name ValueError. from executor.executors.legacy_executor import LegacyExecutor - ExecutorRegistry.register(LegacyExecutor) + if "legacy" not in ExecutorRegistry.list_executors(): + ExecutorRegistry.register(LegacyExecutor) # Mock cloud executors for multi-executor tests