Handle shutdown crashes in executor and worker submission paths

Author: jellyfish0316

src/prefect/runner/_flow_run_executor.py

@@ -123,6 +123,14 @@ async def submit(
             # after this await runs only once the process is done.
             handle = await self._start_process(task_status)
 
+        except anyio.get_cancelled_exc_class():
+            if self._process_manager.get(self._flow_run.id) is not None:
+                with anyio.CancelScope(shield=True):
+                    await self._state_proposer.propose_crashed(
+                        self._flow_run,
+                        message="Flow run process exited due to worker shutdown.",
+                    )
+            raise
         except Exception as exc:
             self._logger.exception(
                 "Flow run '%s' could not start: %s",
@@ -136,7 +144,7 @@ async def submit(
             return
         finally:
             # Step 6: remove handle from process_manager
-            if handle is not None:
+            if handle is not None or self._process_manager.get(self._flow_run.id):
                 await self._process_manager.remove(self._flow_run.id)
 
         # Step 7: interpret exit code and propose terminal state

src/prefect/runner/runner.py

@@ -114,9 +114,7 @@ def fast_flow():
     PREFECT_RUNNER_SERVER_ENABLE,
     get_current_settings,
 )
-from prefect.states import (
-    AwaitingRetry,
-)
+from prefect.states import AwaitingRetry
 from prefect.types.entrypoint import EntrypointType
 from prefect.utilities._engine import get_hook_name
 from prefect.utilities._infrastructure_exit_codes import get_infrastructure_exit_info
@@ -256,7 +254,6 @@ def goodbye_flow(name):
         # --- Facade-owned mutable state (kept until methods are fully delegated) ---
         self._submitting_flow_run_ids: set[UUID] = set()
         self._rescheduling: bool = False
-        self._rescheduled_flow_run_ids: set[UUID] = set()
 
         # --- Facade-owned mutable containers (exposed via @property) ---
         self.__flow_run_process_map_internal: dict[UUID, ProcessMapEntry] = {}
@@ -794,39 +791,14 @@ async def execute_bundle(
             self._flow_run_bundle_map[flow_run.id] = bundle
 
             flow_run_logger = self._get_flow_run_logger(flow_run)
-            exit_code: int | None = None
             try:
-                await anyio.to_thread.run_sync(
-                    process.join,
-                    abandon_on_cancel=True,
-                )
-                exit_code = process.exitcode
-            except anyio.get_cancelled_exc_class():
-                with anyio.CancelScope(shield=True):
-                    flow_run_logger.info(
-                        "Runner shutdown interrupted flow run %r; terminating subprocess.",
-                        flow_run.name,
-                    )
-                    exit_code = await self._terminate_bundle_process(process)
-
-                    if (
-                        exit_code != 0
-                        and flow_run.id not in self._rescheduled_flow_run_ids
-                    ):
-                        terminal_state = await self._propose_crashed_state(
-                            flow_run,
-                            "Flow run process exited due to worker shutdown.",
-                        )
-                        if terminal_state:
-                            await self._run_on_crashed_hooks(
-                                flow_run=flow_run, state=terminal_state
-                            )
-                return
+                await anyio.to_thread.run_sync(process.join)
             finally:
                 with anyio.CancelScope(shield=True):
                     await self._remove_flow_run_process_map_entry(flow_run.id)
                 self._release_limit_slot(flow_run.id)
 
+            exit_code = process.exitcode
             if exit_code is None:
                 raise RuntimeError("Process has no exit code")
 
@@ -851,27 +823,6 @@ async def execute_bundle(
                     f"Process for flow run {flow_run.name!r} exited cleanly."
                 )
 
-    async def _terminate_bundle_process(
-        self,
-        process: multiprocessing.context.SpawnProcess,
-        grace_seconds: int = 30,
-    ) -> int | None:
-        if process.exitcode is not None:
-            return process.exitcode
-
-        if process.is_alive():
-            await anyio.to_thread.run_sync(process.terminate)
-
-            with anyio.move_on_after(grace_seconds):
-                while process.is_alive():
-                    await anyio.sleep(0.1)
-
-            if process.is_alive():
-                await anyio.to_thread.run_sync(process.kill)
-
-        await anyio.to_thread.run_sync(process.join)
-        return process.exitcode
-
     def _get_flow_run_logger(self, flow_run: "FlowRun") -> PrefectLogAdapter:
         return flow_run_logger(flow_run=flow_run).getChild(
             "runner",
@@ -1071,7 +1022,6 @@ def reschedule_current_flow_runs(
                 )
                 try:
                     propose_state_sync(client, AwaitingRetry(), flow_run_id=flow_run.id)
-                    self._rescheduled_flow_run_ids.add(flow_run.id)
                     os.kill(process_info["pid"], signal.SIGTERM)
                     run_logger.info("Rescheduled flow run for resubmission")
                 except ProcessLookupError:
@@ -1259,17 +1209,6 @@ async def _submit_run_and_capture_errors(
                 flow_run_logger.info(
                     f"Process for flow run {flow_run.name!r} exited cleanly."
                 )
-        except anyio.get_cancelled_exc_class():
-            if (
-                task_status._future.done()
-                and flow_run.id not in self._rescheduled_flow_run_ids
-            ):  # type: ignore[attr-defined]
-                with anyio.CancelScope(shield=True):
-                    await self._propose_crashed_state(
-                        flow_run,
-                        "Flow run process exited due to worker shutdown.",
-                    )
-            raise
         except Exception as exc:
             if not task_status._future.done():  # type: ignore
                 # This flow run was being submitted and did not start successfully
@@ -1287,12 +1226,20 @@ async def _submit_run_and_capture_errors(
                     "occurred."
                 )
             return exc
+        except anyio.get_cancelled_exc_class():
+            if self.stopping and task_status._future.done():  # type: ignore[attr-defined]
+                with anyio.CancelScope(shield=True):
+                    await self._propose_crashed_state(
+                        flow_run,
+                        "Flow run process exited due to worker shutdown.",
+                    )
+            raise
         finally:
             self._release_limit_slot(flow_run.id)
 
             await self._remove_flow_run_process_map_entry(flow_run.id)
 
-        if exit_code != 0 and flow_run.id not in self._rescheduled_flow_run_ids:
+        if exit_code != 0:
             await self._propose_crashed_state(
                 flow_run,
                 f"Flow run process exited with non-zero status code {exit_code}.",

src/prefect/workers/base.py

@@ -1472,6 +1472,14 @@ async def _submit_run_and_capture_errors(
                 task_status=task_status,
                 configuration=configuration,
             )
+        except anyio.get_cancelled_exc_class():
+            if task_status and getattr(task_status, "_future").done():
+                with anyio.CancelScope(shield=True):
+                    await self._propose_crashed_state(
+                        flow_run,
+                        "Flow run process exited due to worker shutdown.",
+                    )
+            raise
         except Exception as exc:
             if task_status and not getattr(task_status, "_future").done():
                 # This flow run was being submitted and did not start successfully

tests/runner/test__flow_run_executor.py

@@ -197,6 +197,23 @@ async def test_submit_proposes_crashed_on_nonzero_exit(self):
         assert m["flow_run"] == call_args[0][0]
         assert "1" in call_args[1]["message"]
 
+    async def test_submit_proposes_crashed_when_cancelled_after_start(self):
+        """Cancellation after the process has started should mark the run as crashed."""
+        executor, m = _make_executor()
+        cancelled_error = anyio.get_cancelled_exc_class()()
+        executor._start_process = AsyncMock(side_effect=cancelled_error)
+        m["process_manager"].get = MagicMock(return_value=m["handle"])
+
+        with pytest.raises(anyio.get_cancelled_exc_class()):
+            await executor.submit()
+
+        m["process_manager"].remove.assert_awaited_once_with(m["flow_run"].id)
+        m["state_proposer"].propose_crashed.assert_awaited_once_with(
+            m["flow_run"],
+            message="Flow run process exited due to worker shutdown.",
+        )
+        m["hook_runner"].run_crashed_hooks.assert_not_awaited()
+
     async def test_submit_crashed_message_includes_registry_explanation(self):
         """The crashed state message should include the explanation from
         the centralized exit code registry."""

tests/runner/test_runner.py

@@ -1414,33 +1414,6 @@ async def test_runner_logs_exit_code_help_message(
         else:
             assert record.levelname == "ERROR"
 
-    async def test_runner_does_not_mark_rescheduled_flow_run_crashed(
-        self,
-        monkeypatch: pytest.MonkeyPatch,
-    ):
-        runner = Runner()
-        runner._release_limit_slot = MagicMock()
-        runner._remove_flow_run_process_map_entry = AsyncMock()
-        runner._propose_crashed_state = AsyncMock()
-
-        cancelled_error = anyio.get_cancelled_exc_class()()
-        mock_run_process = AsyncMock(side_effect=cancelled_error)
-        monkeypatch.setattr(runner, "_run_process", mock_run_process)
-
-        task_status = MagicMock()
-        task_status._future = asyncio.get_running_loop().create_future()
-        task_status._future.set_result(None)
-
-        flow_run = MagicMock()
-        flow_run.id = uuid.uuid4()
-        flow_run.name = "test-flow-run"
-        runner._rescheduled_flow_run_ids.add(flow_run.id)
-
-        with pytest.raises(anyio.get_cancelled_exc_class()):
-            await runner._submit_run_and_capture_errors(flow_run, task_status)
-
-        runner._propose_crashed_state.assert_not_awaited()
-
     @pytest.mark.skipif(
         sys.platform != "win32",
         reason="subprocess.CREATE_NEW_PROCESS_GROUP is only defined in Windows",
@@ -1829,48 +1802,6 @@ def da_hook(
 
             assert "This flow was cancelled!" in caplog.text
 
-        async def test_cancelled_bundle_execution_is_marked_crashed(
-            self, prefect_client: PrefectClient
-        ):
-            runner = Runner()
-
-            @flow
-            def slow_flow():
-                sleep(30)
-
-            flow_run = await prefect_client.create_flow_run(slow_flow)
-
-            result = create_bundle_for_flow_run(slow_flow, flow_run)
-            bundle = result["bundle"]
-
-            async def execute_bundle() -> None:
-                await runner.execute_bundle(bundle)
-
-            async with anyio.create_task_group() as tg:
-                tg.start_soon(execute_bundle)
-
-                with anyio.fail_after(30):
-                    while True:
-                        flow_run = await prefect_client.read_flow_run(
-                            flow_run_id=flow_run.id
-                        )
-                        assert flow_run.state is not None
-                        if flow_run.state.is_running():
-                            break
-                        await anyio.sleep(0.5)
-
-                tg.cancel_scope.cancel()
-
-            with anyio.fail_after(30):
-                while True:
-                    flow_run = await prefect_client.read_flow_run(
-                        flow_run_id=flow_run.id
-                    )
-                    assert flow_run.state is not None
-                    if flow_run.state.is_crashed():
-                        break
-                    await anyio.sleep(0.5)
-
         async def test_crashed_bundle_execution(
             self, prefect_client: PrefectClient, caplog: pytest.LogCaptureFixture
         ):