Repository: PrefectHQ/prefect

Issue:
Support cancellation of active flow runs
Currently flow runs can be transitioned into a `Cancelled` state, but that only prevents new tasks in the flow run from starting - it does nothing to stop already running tasks (it's also not exposed nicely in the UI).

In practice this is hard to do effectively in all situations, there is no straightforward way to interrupt an arbitrary python function in all contexts. If running in process, we can interrupt the main thread, but not other threads. If running in a subprocess we can kill the process. If running in a dask cluster we can cancel pending work, but not stop running work - but with #2667 the `DaskExecutor` can now manage the lifetime of a dask cluster, so we could have the executor shutdown the cluster as a hard stop.

Given the above restrictions, it seems we could add more support for cancelling flow runs with the following semantics:

- Cancellation of a flow run starts when flow runs are transitioned into a `Cancelling` state, where `Cancelling` is a subclass of `Running`. This indicates that the flow run is in the process of cancelling, but still is technically `Running`. Transitioning directly into `Cancelled` from `Running` skips a necessary intermediate state and should be deprecated.
- Cancellation is still "best effort":
  - If a task is started, we'll cancel it if possible, otherwise it may run until termination
  - If a task hasn't started, it won't be started

Focussing on how this works with orchestration, the following takes place after a user cancels a running flow:

- The flow run is moved to `Cancelling`
- Once in `Cancelling` , `TaskRunner` interactions with cloud/server (either heartbeats or state updates) will let the task runner know that the task should be cancelled. If possible, a running task will be interrupted, otherwise it will be allowed to finish. Tasks that are allowed to finish will transition to their final state as if cancellation hadn't occurred (e.g. `Success` or `Failed`). This provides a more accurate representation of the cancelled flow run. Tasks that are successfully cancelled are marked as  `Cancelled`.
- Likewise, `FlowRunner` interactions with cloud (either heartbeats or state updates) will let the flow runner know that the flow should be cancelled. Any tasks that haven't been submitted yet will transition directly to `Cancelled`.
  - When the `FlowRunner` notices that a flow run has been cancelled, it will call a new method on the `Executor` to ask the executor to stop any ongoing work if possible. Tentatively I’ll call this method `executor.cancel`. For the `DaskExecutor`, if the executor is managing a temporary cluster (newly enabled in #2667) the cluster will be shutdown, effectively cancelling any ongoing tasks. Otherwise a warning will be output (that pending tasks can’t be stopped) and any ongoing work completed. For the local executor we might try to interrupt the main thread, or we may just wait for the task to finish. As stated above, cancellation is best-effort, and doesn’t guarantee that ongoing work will be interrupted.
- Once all tasks in the flow have transitioned to a `Finished` state (`Cancelled` being a subclass of `Finished`), the flow run will transition from `Cancelling` to `Cancelled`.


Discussion:
@jcrist For Environments/Executors using AWS Fargate or Kubernetes Jobs, do you think this API would also be able to stop Flow Runs by shutting down the underlying infrastructure? It seems to be a common feature request. See: #2766
Good question. I think that cancelling a flow run shouldn't hard-kill the `FlowRunner` - I think we want a more graceful shutdown. IMO the flow runner should:

- Stop submitting new work
- Ask the `Executor` to cancel the whole execution if possible

In the case of a temporary dask cluster, this would shutdown the dask cluster. When the flow runner exits, the initial job (e.g. a fargate task) would exit. Hard killing the `FlowRunner` may not allow sufficient time to cleanup resources, and could leave things in a bad state, I think this should be handled at the `Executor` level.
@jcrist +1 to @dylanbhughes comment - we're using/about to use Fargate fairly heavily and would love to know that when we cancel a Flow, through whatever means that we don't have to chase down 100 running Fargate instances. 

Your comment above - "If a task is started, we'll cancel it if possible, otherwise it may run until termination" , I'd love to see hard kill option here as well. Some of the subprocesses we run in sub flows are long running and expensive in terms of storage and computation, so an extra level of hard kill on any given Task [as an explicit override of the default soft kill you mention] would be a big help
As long as you're either:

- Creating a new dask cluster for your flow, managed by the executor
- Using a `LocalExecutor`

Then having the executor handle closing tasks should be sufficient, in both cases we can stop a running task efficiently. The hard case is if you're connecting to an existing dask cluster. We don't want to kill at the `Environment` level, since we want to make sure the `FlowRunner` and `Executor` have time to clean themselves up.
Dask isn't in the picture for our case, it's local executor on Fargate. I can test the Flow/Task cancellation  in about an hour and see what it does, once the current flows are done 
>  I can test the Flow/Task cancellation in about an hour and see what it does, once the current flows are done

To be clear, currently no tasks are interrupted on cancellation. This is a proposal for new work. Currently if you cancel a flow run, no new tasks will be started, but no tasks will be interrupted either.

In your case with the proposal above, running tasks should be interrupted and cancellation should be fairly quick (provided your tasks are interruptible and don't silence SIGINT).
Yeah sounds good! Once there's the ability to distinguish between the 2 then that works fine

Resolving patch:
diff --git a/src/prefect/engine/cloud/flow_runner.py b/src/prefect/engine/cloud/flow_runner.py
--- a/src/prefect/engine/cloud/flow_runner.py
+++ b/src/prefect/engine/cloud/flow_runner.py
@@ -1,5 +1,9 @@
-import time
-from typing import Any, Callable, Dict, Iterable, Optional
+import os
+import signal
+import threading
+from time import sleep as time_sleep
+from typing import Any, Callable, Dict, Iterable, Optional, Iterator
+from contextlib import contextmanager
 
 import pendulum
 
@@ -9,7 +13,7 @@
 from prefect.engine.cloud import CloudTaskRunner
 from prefect.engine.flow_runner import FlowRunner, FlowRunnerInitializeResult
 from prefect.engine.runner import ENDRUN
-from prefect.engine.state import Failed, Queued, State
+from prefect.engine.state import Failed, Queued, State, Cancelling, Cancelled
 from prefect.utilities.graphql import with_args
 
 
@@ -134,6 +138,75 @@ def call_runner_target_handlers(self, old_state: State, new_state: State) -> Sta
 
         return new_state
 
+    @contextmanager
+    def check_for_cancellation(self) -> Iterator:
+        """Contextmanager used to wrap a cancellable section of a flow run."""
+
+        cancelling = False
+        done = threading.Event()
+        flow_run_version = None
+        flow_run_id = prefect.context.get("flow_run_id")
+
+        def interrupt_if_cancelling() -> None:
+            while True:
+                exiting_context = done.wait(
+                    prefect.config.cloud.check_cancellation_interval
+                )
+                try:
+                    self.logger.debug("Checking flow run state...")
+                    flow_run_info = self.client.get_flow_run_info(flow_run_id)
+                except Exception:
+                    self.logger.warning("Error getting flow run info", exc_info=True)
+                    continue
+                else:
+                    self.logger.debug(
+                        "Successfully queried server for flow run state: %r",
+                        flow_run_info.state,
+                    )
+                if isinstance(flow_run_info.state, Cancelling):
+                    self.logger.info(
+                        "Flow run has been cancelled, cancelling active tasks"
+                    )
+                    nonlocal cancelling
+                    nonlocal flow_run_version
+                    cancelling = True
+                    flow_run_version = flow_run_info.version
+                    # If not already leaving context, raise KeyboardInterrupt in the main thread
+                    if not exiting_context:
+                        if hasattr(signal, "raise_signal"):
+                            # New in python 3.8
+                            signal.raise_signal(signal.SIGINT)  # type: ignore
+                        else:
+                            if os.name == "nt":
+                                # This doesn't actually send a signal, so it will only
+                                # interrupt the next Python bytecode instruction - if the
+                                # main thread is blocked in a c extension the interrupt
+                                # won't be seen until that returns.
+                                from _thread import interrupt_main
+
+                                interrupt_main()
+                            else:
+                                signal.pthread_kill(
+                                    threading.main_thread().ident, signal.SIGINT  # type: ignore
+                                )
+                    break
+                elif exiting_context:
+                    break
+
+        thread = threading.Thread(target=interrupt_if_cancelling, daemon=True)
+        thread.start()
+        try:
+            yield
+        except KeyboardInterrupt:
+            if not cancelling:
+                raise
+        finally:
+            done.set()
+            thread.join()
+            if cancelling:
+                prefect.context.update(flow_run_version=flow_run_version)
+                raise ENDRUN(state=Cancelled("Flow run is cancelled"))
+
     def run(
         self,
         state: State = None,
@@ -190,21 +263,30 @@ def run(
             end_state.start_time <= pendulum.now("utc").add(minutes=10)  # type: ignore
         ):
             assert isinstance(end_state, Queued)
-            naptime = max(
+            time_remaining = max(
                 (end_state.start_time - pendulum.now("utc")).total_seconds(), 0
             )
             self.logger.info(
                 (
-                    "Flow run is in a Queued state."
-                    f" Sleeping for {naptime:.2f} seconds and attempting to run again."
+                    f"Flow run is in a Queued state. Sleeping for at most {time_remaining:.2f} "
+                    f"seconds and attempting to run again."
                 )
             )
-            time.sleep(naptime)
+            # Sleep until not in a queued state, then attempt to re-run
+            while time_remaining > 0:
+                delay = min(
+                    prefect.config.cloud.check_cancellation_interval, time_remaining
+                )
+                time_remaining -= delay
+                # Imported `time.sleep` as `time_sleep` to allow monkeypatching in tests
+                time_sleep(delay)
 
-            flow_run_info = self.client.get_flow_run_info(
-                flow_run_id=prefect.context.get("flow_run_id")
-            )
-            context.update(flow_run_version=flow_run_info.version)
+                flow_run_info = self.client.get_flow_run_info(
+                    flow_run_id=prefect.context.get("flow_run_id")
+                )
+                context.update(flow_run_version=flow_run_info.version)
+                if not isinstance(flow_run_info.state, Queued):
+                    break
 
             # When concurrency slots become free, this will eventually result
             # in a non queued state, but will result in more or less just waiting
diff --git a/src/prefect/engine/flow_runner.py b/src/prefect/engine/flow_runner.py
--- a/src/prefect/engine/flow_runner.py
+++ b/src/prefect/engine/flow_runner.py
@@ -3,10 +3,12 @@
     Callable,
     Dict,
     Iterable,
+    Iterator,
     NamedTuple,
     Optional,
     Set,
 )
+from contextlib import contextmanager
 
 import pendulum
 import prefect
@@ -15,7 +17,6 @@
 from prefect.engine.results import ConstantResult
 from prefect.engine.runner import ENDRUN, Runner, call_state_handlers
 from prefect.engine.state import (
-    Cancelled,
     Failed,
     Mapped,
     Pending,
@@ -223,7 +224,6 @@ def run(
             - State: `State` representing the final post-run state of the `Flow`.
 
         """
-
         self.logger.info("Beginning Flow run for '{}'".format(self.flow.name))
 
         # make copies to avoid modifying user inputs
@@ -259,10 +259,6 @@ def run(
         except ENDRUN as exc:
             state = exc.state
 
-        except KeyboardInterrupt:
-            self.logger.debug("Interrupt signal raised, cancelling Flow run.")
-            state = Cancelled(message="Interrupt signal raised, cancelling flow run.")
-
         # All other exceptions are trapped and turned into Failed states
         except Exception as exc:
             self.logger.exception(
@@ -278,6 +274,14 @@ def run(
 
         return state
 
+    @contextmanager
+    def check_for_cancellation(self) -> Iterator:
+        """Contextmanager used to wrap a cancellable section of a flow run.
+
+        No-op for the default `FlowRunner` class.
+        """
+        yield
+
     @call_state_handlers
     def check_flow_reached_start_time(self, state: State) -> State:
         """
@@ -410,7 +414,7 @@ def extra_context(task: Task, task_index: int = None) -> dict:
 
         # -- process each task in order
 
-        with executor.start():
+        with self.check_for_cancellation(), executor.start():
 
             for task in self.flow.sorted_tasks():
                 task_state = task_states.get(task)
diff --git a/src/prefect/engine/task_runner.py b/src/prefect/engine/task_runner.py
--- a/src/prefect/engine/task_runner.py
+++ b/src/prefect/engine/task_runner.py
@@ -21,7 +21,6 @@
 from prefect.engine.runner import ENDRUN, Runner, call_state_handlers
 from prefect.engine.state import (
     Cached,
-    Cancelled,
     Failed,
     Looped,
     Mapped,
@@ -802,11 +801,6 @@ def get_task_run_state(self, state: State, inputs: Dict[str, Result],) -> State:
                     self.task.run, timeout=self.task.timeout, **raw_inputs
                 )
 
-        except KeyboardInterrupt:
-            self.logger.debug("Interrupt signal raised, cancelling task run.")
-            state = Cancelled(message="Interrupt signal raised, cancelling task run.")
-            return state
-
         # inform user of timeout
         except TimeoutError as exc:
             if prefect.context.get("raise_on_exception"):


Test patch:
diff --git a/tests/core/test_flow.py b/tests/core/test_flow.py
--- a/tests/core/test_flow.py
+++ b/tests/core/test_flow.py
@@ -2418,20 +2418,6 @@ def record_start_time():
         f.run()
         assert REPORTED_START_TIMES == start_times
 
-    def test_flow_dot_run_handles_keyboard_signals_gracefully(self):
-        class BadExecutor(LocalExecutor):
-            def submit(self, *args, **kwargs):
-                raise KeyboardInterrupt
-
-        @task
-        def do_something():
-            pass
-
-        f = Flow("test", tasks=[do_something])
-        state = f.run(executor=BadExecutor())
-        assert isinstance(state, Cancelled)
-        assert "interrupt" in state.message.lower()
-
 
 class TestFlowDiagnostics:
     def test_flow_diagnostics(self, monkeypatch):
diff --git a/tests/engine/cloud/test_cloud_flow_runner.py b/tests/engine/cloud/test_cloud_flow_runner.py
--- a/tests/engine/cloud/test_cloud_flow_runner.py
+++ b/tests/engine/cloud/test_cloud_flow_runner.py
@@ -1,5 +1,7 @@
 import datetime
+import itertools
 import time
+import threading
 from datetime import timedelta
 from unittest.mock import MagicMock
 
@@ -20,6 +22,7 @@
 from prefect.engine.signals import LOOP
 from prefect.engine.state import (
     Cancelled,
+    Cancelling,
     Failed,
     Finished,
     Pending,
@@ -78,13 +81,13 @@ def test_task_runner_cls_is_cloud_task_runner():
     assert fr.task_runner_cls is CloudTaskRunner
 
 
-def test_flow_runner_calls_client_the_approriate_number_of_times(client):
+def test_flow_runner_calls_client_the_appropriate_number_of_times(client):
     flow = prefect.Flow(name="test")
 
     res = CloudFlowRunner(flow=flow).run()
 
     ## assertions
-    assert client.get_flow_run_info.call_count == 1  # one time to pull latest state
+    assert client.get_flow_run_info.call_count == 2  # initial state & cancel check
     assert client.set_flow_run_state.call_count == 2  # Pending -> Running -> Success
 
     states = [call[1]["state"] for call in client.set_flow_run_state.call_args_list]
@@ -100,7 +103,7 @@ def test_flow_runner_doesnt_set_running_states_twice(client):
     )
 
     ## assertions
-    assert client.get_flow_run_info.call_count == 1  # one time to pull latest state
+    assert client.get_flow_run_info.call_count == 2  # initial state & cancel check
     assert client.set_flow_run_state.call_count == 1  # Pending -> Running
 
 
@@ -169,7 +172,7 @@ def handler(task, old, new):
     res = flow.run(state=Pending())
 
     ## assertions
-    assert client.get_flow_run_info.call_count == 1  # one time to pull latest state
+    assert client.get_flow_run_info.call_count == 1  # initial state, no cancel check
     assert client.set_flow_run_state.call_count == 1  # Failed
 
     flow_states = [
@@ -218,7 +221,7 @@ def test_flow_runner_respects_the_db_state(monkeypatch, state):
     res = CloudFlowRunner(flow=flow).run()
 
     ## assertions
-    assert get_flow_run_info.call_count == 1  # one time to pull latest state
+    assert get_flow_run_info.call_count == 1  # initial state, no cancel check
     assert set_flow_run_state.call_count == 0  # never needs to update state
     assert res == db_state
 
@@ -240,7 +243,7 @@ def test_flow_runner_prioritizes_kwarg_states_over_db_states(
     res = CloudFlowRunner(flow=flow).run(state=Pending("let's do this"))
 
     ## assertions
-    assert get_flow_run_info.call_count == 1  # one time to pull latest state
+    assert get_flow_run_info.call_count == 2  # initial state & cancel check
     assert client.set_flow_run_state.call_count == 2  # Pending -> Running -> Success
 
     states = [call[1]["state"] for call in client.set_flow_run_state.call_args_list]
@@ -357,7 +360,7 @@ def raise_me(x, y):
     res = flow.run(state=Pending())
 
     ## assertions
-    assert client.get_flow_run_info.call_count == 1  # one time to pull latest state
+    assert client.get_flow_run_info.call_count == 2  # initial state & cancel check
     assert client.set_flow_run_state.call_count == 2  # Pending -> Running -> Failed
 
     flow_states = [
@@ -550,42 +553,161 @@ def log_stuff():
     assert set(task_run_ids) == {"TESTME"}
 
 
-@pytest.mark.parametrize("num_attempts", [5, 10, 15])
-def test_flow_runner_retries_forever_on_queued_state(client, monkeypatch, num_attempts):
+class TestCloudFlowRunnerQueuedState:
+    queue_time = 55
+    check_cancellation_interval = 8
 
-    mock_sleep = MagicMock()
-    monkeypatch.setattr("prefect.engine.cloud.flow_runner.time.sleep", mock_sleep)
+    def do_mocked_run(
+        self, client, monkeypatch, n_attempts=None, n_queries=None, query_end_state=None
+    ):
+        """Mock out a cloud flow run that starts in a queued state and either
+        succeeds or exits early due to a state change."""
+        mock_sleep = MagicMock()
+
+        def run(*args, **kwargs):
+            if n_attempts is None or mock_run.call_count < n_attempts:
+                info = get_flow_run_info()
+                if info.state.is_queued():
+                    return Queued(
+                        start_time=pendulum.now("UTC").add(seconds=self.queue_time)
+                    )
+                return info.state
+            return Success()
+
+        mock_run = MagicMock(side_effect=run)
+
+        def get_flow_run_info(*args, **kwargs):
+            if n_queries is None or mock_get_flow_run_info.call_count < n_queries:
+                state = Queued()
+            else:
+                state = query_end_state
+            return MagicMock(version=mock_get_flow_run_info.call_count, state=state)
+
+        mock_get_flow_run_info = MagicMock(side_effect=get_flow_run_info)
+
+        client.get_flow_run_info = mock_get_flow_run_info
+        monkeypatch.setattr("prefect.engine.cloud.flow_runner.FlowRunner.run", mock_run)
+        monkeypatch.setattr("prefect.engine.cloud.flow_runner.time_sleep", mock_sleep)
+
+        @prefect.task
+        def return_one():
+            return 1
+
+        with prefect.Flow("test-cloud-flow-runner-with-queues") as flow:
+            return_one()
+
+        with set_temporary_config(
+            {"cloud.check_cancellation_interval": self.check_cancellation_interval}
+        ):
+            state = CloudFlowRunner(flow=flow).run()
+        return state, mock_sleep, mock_run
 
-    run_states = [
-        Queued(start_time=pendulum.now("UTC").add(seconds=i))
-        for i in range(num_attempts - 1)
-    ]
-    run_states.append(Success())
+    @pytest.mark.parametrize("n_attempts", [5, 10])
+    def test_rety_queued_state_until_success(self, client, monkeypatch, n_attempts):
+        state, mock_sleep, mock_run = self.do_mocked_run(
+            client, monkeypatch, n_attempts=n_attempts
+        )
 
-    mock_run = MagicMock(side_effect=run_states)
+        assert state.is_successful()
+        assert mock_run.call_count == n_attempts
+        sleep_times = [i[0][0] for i in mock_sleep.call_args_list]
+        assert max(sleep_times) == self.check_cancellation_interval
+        total_sleep_time = sum(sleep_times)
+        expected_sleep_time = (n_attempts - 1) * self.queue_time
+        # Slept for approximately the right amount of time. Due to processing time,
+        # the amount of time spent in sleep may be slightly less.
+        assert expected_sleep_time - 2 < total_sleep_time < expected_sleep_time + 2
+
+    @pytest.mark.parametrize("n_queries", [5, 10])
+    @pytest.mark.parametrize("final_state", [Cancelled(), Success()])
+    def test_exit_queued_loop_early_if_no_longer_queued(
+        self, client, monkeypatch, n_queries, final_state
+    ):
+        state, mock_sleep, mock_run = self.do_mocked_run(
+            client, monkeypatch, n_queries=n_queries, query_end_state=final_state
+        )
 
-    client.get_flow_run_info = MagicMock(
-        side_effect=[MagicMock(version=i) for i in range(num_attempts)]
-    )
+        assert type(state) == type(final_state)
+        sleep_times = [i[0][0] for i in mock_sleep.call_args_list]
+        assert max(sleep_times) == self.check_cancellation_interval
+        total_sleep_time = sum(sleep_times)
+        expected_sleep_time = n_queries * self.check_cancellation_interval
+        # Slept for approximately the right amount of time. Due to processing time,
+        # the amount of time spent in sleep may be slightly less.
+        assert expected_sleep_time - 2 < total_sleep_time < expected_sleep_time + 2
 
-    # Mock out the actual flow execution
-    monkeypatch.setattr("prefect.engine.cloud.flow_runner.FlowRunner.run", mock_run)
 
-    @prefect.task
-    def return_one():
-        return 1
-
-    with prefect.Flow("test-cloud-flow-runner-with-queues") as flow:
-        one = return_one()
-
-    # Without these (actual, not mocked) sleep calls, when running full test suite this
-    # test can fail for no reason.
-    final_state = CloudFlowRunner(flow=flow).run()
-    assert final_state.is_successful()
-
-    assert mock_run.call_count == num_attempts
-    # Not called on the initial run attempt
-    assert client.get_flow_run_info.call_count == num_attempts - 1
-    # Not checking the value of the `time.sleep` mock, since it appears to
-    # be getting called thousands of times, likely due to a pytest
-    # plugin or something of the like.
+class TestCloudFlowRunnerCancellation:
+    def test_cancelling_mid_flow_run_exits_early(self, client, monkeypatch):
+        trigger = threading.Event()
+
+        def get_flow_run_info(*args, _version=itertools.count(), **kwargs):
+            state = Cancelling() if trigger.is_set() else Running()
+            return MagicMock(version=next(_version), state=state)
+
+        client.get_flow_run_info = get_flow_run_info
+
+        @prefect.task
+        def inc(x):
+            time.sleep(0.5)
+            return x + 1
+
+        ran_longer_than_expected = False
+
+        @prefect.task
+        def set_trigger(x):
+            trigger.set()
+            time.sleep(10)
+            nonlocal ran_longer_than_expected
+            ran_longer_than_expected = True
+            return x + 1
+
+        with prefect.Flow("test") as flow:
+            a = inc(1)
+            b = set_trigger(a)
+            c = inc(b)
+
+        with set_temporary_config({"cloud.check_cancellation_interval": 0.1}):
+            res = CloudFlowRunner(flow=flow).run()
+
+        assert isinstance(res, Cancelled)
+        assert not ran_longer_than_expected
+
+    def test_check_interrupt_loop_robust_to_api_errors(self, client, monkeypatch):
+        trigger = threading.Event()
+
+        error_was_raised = False
+
+        def get_flow_run_info(*args, _call_count=itertools.count(), **kwargs):
+            call_count = next(_call_count)
+            import inspect
+
+            caller_name = inspect.currentframe().f_back.f_code.co_name
+            if caller_name == "interrupt_if_cancelling" and call_count % 2:
+                nonlocal error_was_raised
+                error_was_raised = True
+                raise ValueError("Woops!")
+            state = Cancelling() if trigger.is_set() else Running()
+            return MagicMock(version=call_count, state=state)
+
+        client.get_flow_run_info = get_flow_run_info
+
+        ran_longer_than_expected = False
+
+        @prefect.task
+        def set_trigger(x):
+            trigger.set()
+            time.sleep(10)
+            nonlocal ran_longer_than_expected
+            ran_longer_than_expected = True
+            return x + 1
+
+        with prefect.Flow("test") as flow:
+            set_trigger(1)
+
+        with set_temporary_config({"cloud.check_cancellation_interval": 0.1}):
+            res = CloudFlowRunner(flow=flow).run()
+
+        assert isinstance(res, Cancelled)
+        assert error_was_raised
+        assert not ran_longer_than_expected