Reject concurrent subprocess coordinator execution consistently

Bundle resolution does not need mutable instance state, while overwriting scan roots can make a command build against another task's artifacts.

Author: jason810496

task-sdk/src/airflow/sdk/coordinators/_subprocess.py

@@ -446,7 +446,6 @@ class SubprocessCoordinator(BaseCoordinator):
     # The subclass's explicit root, recorded at construction so the base can
     # resolve roots without knowing the subclass field name.
     _configured_roots: list[pathlib.Path] = attrs.field(init=False, factory=list)
-    _active_bundle_info: BundleInfo | None = attrs.field(init=False, default=None)
     _active_scan_roots: tuple[pathlib.Path, ...] | None = attrs.field(init=False, default=None)
 
     @property
@@ -500,7 +499,7 @@ def _classify_artifact_source(self) -> None:
         )
 
     def _init_root_source(
-        self, logger: FilteringBoundLogger
+        self, bundle_info: BundleInfo, logger: FilteringBoundLogger
     ) -> tuple[list[pathlib.Path], BaseDagBundle | None]:
         """
         Resolve the directories to scan for artifacts, dispatched on the classified mode.
@@ -517,9 +516,7 @@ def _init_root_source(
             # NAMED_BUNDLE implies dag_bundle_name is set.
             target = BundleInfo(name=cast("str", self.dag_bundle_name))
         else:
-            if self._active_bundle_info is None:
-                raise RuntimeError("_init_root_source requires an active task; call it during execute_task.")
-            target = self._active_bundle_info
+            target = bundle_info
 
         bundle = _initialize_pinned_bundle(target, logger)
         path = bundle.path
@@ -547,25 +544,11 @@ def _build_execute_task_command(self, *, what: TaskInstance) -> tuple[list[str],
         """
         raise NotImplementedError
 
-    @contextlib.contextmanager
-    def _set_current_bundle(self, bundle_info: BundleInfo):
-        """
-        Bind *bundle_info* as the active task for the duration of the block, clearing it on exit.
-
-        Rejects a second concurrent bind: this coordinator runs one blocking task
-        per process and is not re-entrant.
-        """
-        if self._active_bundle_info is not None:
-            raise RuntimeError("SubprocessCoordinator.execute_task is not re-entrant.")
-        self._active_bundle_info = bundle_info
-        try:
-            yield
-        finally:
-            self._active_bundle_info = None
-
     @contextlib.contextmanager
     def _set_scan_roots(self, roots: Sequence[pathlib.Path]):
         """Expose *roots* to the command builder for the duration of the task."""
+        if self._active_scan_roots is not None:
+            raise RuntimeError("SubprocessCoordinator.execute_task is not re-entrant.")
         self._active_scan_roots = tuple(roots)
         try:
             yield
@@ -586,8 +569,7 @@ def execute_task(
     ) -> BaseCoordinator.ExecutionResult:
         task_logger = logger or log
         with contextlib.ExitStack() as stack:
-            stack.enter_context(self._set_current_bundle(bundle_info))
-            roots, resolved_bundle = self._init_root_source(task_logger)
+            roots, resolved_bundle = self._init_root_source(bundle_info, task_logger)
             if resolved_bundle is not None:
                 # Hold the version lock across start()/wait() so bundle cleanup
                 # cannot rmtree a version this task is still reading from,

task-sdk/tests/task_sdk/coordinators/test_subprocess.py

@@ -918,29 +918,23 @@ class TestInitRootSource:
     def test_returns_configured_root_and_no_bundle_in_explicit_mode(self):
         configured = [pathlib.Path("/a"), pathlib.Path("/b")]
         coordinator = _StubSubprocessCoordinator(command=["x"], explicit_roots=configured)
-        roots, bundle = coordinator._init_root_source(log)
+        roots, bundle = coordinator._init_root_source(MagicMock(), log)
         assert roots == configured
         assert bundle is None
 
     @patch("airflow.sdk.coordinators._subprocess.initialize_ti_bundle")
-    def test_task_bundle_mode_uses_active_task_bundle(self, mock_initialize, tmp_path):
+    def test_task_bundle_mode_uses_passed_task_bundle(self, mock_initialize, tmp_path):
         resolved = MagicMock(path=tmp_path, version="v3")
         mock_initialize.return_value = resolved
         coordinator = _StubSubprocessCoordinator(command=["x"], explicit_roots=[])
         bundle_info = BundleInfo(name="dags", version="v3", version_data={"k": "v"})
-        coordinator._active_bundle_info = bundle_info
 
-        roots, bundle = coordinator._init_root_source(log)
+        roots, bundle = coordinator._init_root_source(bundle_info, log)
 
         assert roots == [tmp_path]
         assert bundle is resolved
         mock_initialize.assert_called_once_with(bundle_info)
 
-    def test_task_bundle_mode_without_active_task_raises(self):
-        coordinator = _StubSubprocessCoordinator(command=["x"], explicit_roots=[])
-        with pytest.raises(RuntimeError, match="requires an active task"):
-            coordinator._init_root_source(log)
-
     @patch("airflow.sdk.coordinators._subprocess.initialize_ti_bundle")
     def test_version_less_bundle_is_rematerialized_at_its_current_version(self, mock_initialize, tmp_path):
         """A bundle resolved without a version is re-resolved at the version current now.
@@ -962,7 +956,7 @@ def test_version_less_bundle_is_rematerialized_at_its_current_version(self, mock
         coordinator.dag_bundle_name = "artifacts"
         coordinator._artifact_source = _ArtifactSource.NAMED_BUNDLE
 
-        roots, bundle = coordinator._init_root_source(log)
+        roots, bundle = coordinator._init_root_source(MagicMock(), log)
 
         assert roots == [pinned_tree]
         assert bundle is pinned
@@ -981,7 +975,7 @@ def test_unversioned_bundle_keeps_its_single_path(self, mock_initialize, tmp_pat
         coordinator.dag_bundle_name = "artifacts"
         coordinator._artifact_source = _ArtifactSource.NAMED_BUNDLE
 
-        roots, bundle = coordinator._init_root_source(log)
+        roots, bundle = coordinator._init_root_source(MagicMock(), log)
 
         assert roots == [tmp_path]
         assert bundle is resolved
@@ -996,11 +990,11 @@ def test_missing_resolved_path_raises(self, mock_initialize, tmp_path):
         coordinator._artifact_source = _ArtifactSource.NAMED_BUNDLE
 
         with pytest.raises(FileNotFoundError, match="does not exist"):
-            coordinator._init_root_source(log)
+            coordinator._init_root_source(MagicMock(), log)
 
 
 class TestExecuteTaskBundleWiring:
-    """execute_task binds the task bundle, forwards resolved roots, and holds the version lock."""
+    """execute_task passes the task bundle, forwards resolved roots, and holds the version lock."""
 
     @patch("airflow.sdk.coordinators._subprocess.BundleVersionLock")
     @patch("airflow.sdk.coordinators._subprocess.initialize_ti_bundle")
@@ -1024,11 +1018,8 @@ def test_task_bundle_mode_binds_forwards_roots_and_locks(
             subprocess_logs_to_stdout=False,
         )
 
-        # Resolution runs against the bound task bundle.
         mock_initialize.assert_called_once_with(bundle_info)
-        # The resolved root is forwarded to the subclass command builder.
         assert coordinator.recorded_roots == [[tmp_path]]
-        # The version lock is held around the subprocess for the resolved bundle.
         mock_lock.assert_called_once_with(bundle_name="dags", bundle_version="v9")
         mock_lock.return_value.__enter__.assert_called_once()
         mock_lock.return_value.__exit__.assert_called_once()
@@ -1095,34 +1086,21 @@ def test_returns_bound_roots_and_clears_them_on_exit(self, tmp_path):
         with pytest.raises(RuntimeError, match="requires an active task"):
             coordinator._get_scan_roots()
 
-
-class TestSetCurrentBundle:
-    def test_binds_for_the_block_and_clears_on_exit_allowing_reuse(self):
+    def test_rejects_nested_reentry(self, tmp_path):
         coordinator = _StubSubprocessCoordinator(command=["/bin/true"])
-        bundle = MagicMock()
-        with coordinator._set_current_bundle(bundle):
-            assert coordinator._active_bundle_info is bundle
-        assert coordinator._active_bundle_info is None
-        # A subsequent task on the same (reused) coordinator binds cleanly.
-        with coordinator._set_current_bundle(MagicMock()):
-            assert coordinator._active_bundle_info is not None
-        assert coordinator._active_bundle_info is None
-
-    def test_rejects_nested_reentry(self):
-        coordinator = _StubSubprocessCoordinator(command=["/bin/true"])
-        with coordinator._set_current_bundle(MagicMock()):
+        with coordinator._set_scan_roots([tmp_path]):
             with pytest.raises(RuntimeError, match="not re-entrant"):
-                with coordinator._set_current_bundle(MagicMock()):
+                with coordinator._set_scan_roots([tmp_path]):
                     pass
 
-    def test_execute_task_is_not_reentrant(self, mock_client):
+    def test_execute_task_is_not_reentrant(self, mock_client, tmp_path):
         coordinator = _StubSubprocessCoordinator(command=["/bin/true"])
-        coordinator._active_bundle_info = MagicMock()  # simulate an in-flight task
-        with pytest.raises(RuntimeError, match="not re-entrant"):
-            coordinator.execute_task(
-                what=_make_ti(),
-                dag_rel_path="dag.py",
-                bundle_info=MagicMock(),
-                client=mock_client,
-                subprocess_logs_to_stdout=False,
-            )
+        with coordinator._set_scan_roots([tmp_path]):
+            with pytest.raises(RuntimeError, match="not re-entrant"):
+                coordinator.execute_task(
+                    what=_make_ti(),
+                    dag_rel_path="dag.py",
+                    bundle_info=MagicMock(),
+                    client=mock_client,
+                    subprocess_logs_to_stdout=False,
+                )