Repository: PrefectHQ/prefect
File: src/prefect/runner/runner.py (Python)

Code before:
        """
        will_cancel = get_current_settings().runner.crash_on_cancellation_failure
        cancel_message = (
            "Cancellation observing failed - both websocket and polling mechanisms "
            "are unavailable. Marking flow run as crashed to prevent indefinite execution."
        )
        continue_message = (
            "Cancellation observing failed - both websocket and polling mechanisms "
            "are unavailable. Flow run will continue executing but cannot be "
            "cancelled. Set PREFECT_RUNNER_CRASH_ON_CANCELLATION_FAILURE=true to "
            "crash flow runs and shut down the runner when this occurs."
        )
        for flow_run_id in flow_run_ids:
            process_entry = self._flow_run_process_map.get(flow_run_id)
            if process_entry:
                flow_run = process_entry["flow_run"]
                run_logger = self._get_flow_run_logger(flow_run)
                run_logger.error(cancel_message if will_cancel else continue_message)
                if will_cancel:
                    await self._propose_crashed_state(flow_run, cancel_message)
        if will_cancel:
            # Raise to trigger runner shutdown - the task group will handle cleanup
            raise RuntimeError(
                "Cancellation observer failed. Runner cannot safely continue without "
                "the ability to respond to cancellation requests."
            )

    async def _mark_flow_run_as_cancelled(
        self, flow_run: "FlowRun", state_updates: Optional[dict[str, Any]] = None
    ) -> None:
        state_updates = state_updates or {}
        state_updates.setdefault("name", "Cancelled")
        state_updates.setdefault("type", StateType.CANCELLED)
        state = (
            flow_run.state.model_copy(update=state_updates) if flow_run.state else None
        )
        if not state:
            self._logger.warning(
                f"Could not find state for flow run {flow_run.id} and cancellation cannot be guaranteed."
            )
            return

        try:
            await self._client.set_flow_run_state(flow_run.id, state, force=True)
        except ObjectNotFound:
            # Flow run was deleted - log it but don't crash the runner
            run_logger = self._get_flow_run_logger(flow_run)
            run_logger.debug(
                f"Flow run '{flow_run.id}' was deleted before it could be marked as cancelled"
            )


Review comment (suggestion):
Does raising here stop flow run execution? I know the task run group will wait until execution completes during a normal exit, but I can't remember if the same is true when an exception is raised.

Code after:
        state_updates.setdefault("type", StateType.CANCELLED)
        state = (
            flow_run.state.model_copy(update=state_updates) if flow_run.state else None
        )
        if not state:
            self._logger.warning(
                f"Could not find state for flow run {flow_run.id} and cancellation cannot be guaranteed."
            )
            return

        try:
            await self._client.set_flow_run_state(flow_run.id, state, force=True)
        except ObjectNotFound:
            # Flow run was deleted - log it but don't crash the runner
            run_logger = self._get_flow_run_logger(flow_run)
            run_logger.debug(
                f"Flow run '{flow_run.id}' was deleted before it could be marked as cancelled"
            )

    async def _run_on_cancellation_hooks(
        self,
        flow_run: "FlowRun",
        state: State,
    ) -> None:
        """
        Run the hooks for a flow.
        """
        run_logger = self._get_flow_run_logger(flow_run)
        if state.is_cancelling():
            try:
                if flow_run.id in self._flow_run_bundle_map:
                    flow = extract_flow_from_bundle(
                        self._flow_run_bundle_map[flow_run.id]
                    )
                elif flow_run.deployment_id and self._deployment_flow_map.get(
                    flow_run.deployment_id
                ):
                    flow = self._deployment_flow_map[flow_run.deployment_id]
                else:
                    run_logger.info("Loading flow to check for on_cancellation hooks")
                    flow = await load_flow_from_flow_run(
                        flow_run, storage_base_path=str(self._tmp_dir)
                    )
                hooks = flow.on_cancellation_hooks or []

                await _run_hooks(hooks, flow_run, flow, state)
            except Exception:
                run_logger.warning(
                    f"Runner failed to retrieve flow to execute on_cancellation hooks for flow run {flow_run.id!r}.",
                    exc_info=True,
                )