Repository: PrefectHQ/prefect

Issue:
Add RESUME state
Following some discussion on reacting to states, we will implement a `Resume` state that subclasses `Scheduled` (and should have a default `start_time` of `pendulum.now('utc')`)

When the `TaskRunner` sees a `Resume` state, it should set the task context appropriately and then run the task. 

This will make pause/resume super clean: want to pause a task? Put it in a pause state! Want to resume the task? Put it in a resume state!

*NB: some relevance to #88, in that we no longer need dedicated endpoints -- resuming is just a state change mutation. However we may still want ways to easily generate the API calls.* 


Discussion:

Resolving patch:
diff --git a/src/prefect/engine/state.py b/src/prefect/engine/state.py
--- a/src/prefect/engine/state.py
+++ b/src/prefect/engine/state.py
@@ -284,6 +284,22 @@ def __init__(
         self.start_time = ensure_tz_aware(start_time or pendulum.now("utc"))
 
 
+class Resume(Scheduled):
+    """
+    Resume state indicating the object can resume execution (presumably from a `Paused` state).
+
+    Args:
+        - message (str or Exception, optional): Defaults to `None`. A message about the
+            state, which could be an `Exception` (or [`Signal`](signals.html)) that caused it.
+        - result (Any, optional): Defaults to `None`. A data payload for the state.
+        - start_time (datetime): time at which the task is scheduled to run
+        - cached_inputs (dict): Defaults to `None`. A dictionary of input
+            keys to values.  Used / set if the Task requires Retries.
+    """
+
+    color = "#20B2AA"
+
+
 class Retrying(Scheduled):
     """
     Pending state indicating the object has been scheduled to be retried.
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
@@ -19,6 +19,7 @@
     Mapped,
     Pending,
     Scheduled,
+    Resume,
     Retrying,
     Running,
     Skipped,
@@ -211,8 +212,9 @@ def run(
         with prefect.context(context, _task_name=self.task.name, _map_index=map_index):
 
             try:
-                # retrieve the run number and place in context
-                state = self.get_run_count(state=state)
+                # retrieve the run number and place in context,
+                # or put resume in context if needed
+                state = self.update_context_from_state(state=state)
 
                 # check if all upstream tasks have finished
                 state = self.check_upstream_finished(
@@ -298,12 +300,16 @@ def run(
         return state
 
     @call_state_handlers
-    def get_run_count(self, state: State) -> State:
+    def update_context_from_state(self, state: State) -> State:
         """
+        Updates context with information contained in the task state:
+
         If the task is being retried, then we retrieve the run count from the initial Retry
         state. Otherwise, we assume the run count is 1. The run count is stored in context as
         _task_run_count.
 
+        Also, if the task is being resumed through a `Resume` state, updates context to have `resume=True`.
+
         Args:
             - state (State): the current state of the task
 
@@ -314,6 +320,8 @@ def get_run_count(self, state: State) -> State:
             run_count = state.run_count + 1
         else:
             run_count = 1
+        if isinstance(state, Resume):
+            prefect.context.update(resume=True)
         prefect.context.update(_task_run_count=run_count)
         return state
 
diff --git a/src/prefect/serialization/state.py b/src/prefect/serialization/state.py
--- a/src/prefect/serialization/state.py
+++ b/src/prefect/serialization/state.py
@@ -67,6 +67,12 @@ class Meta:
     start_time = DateTime(allow_none=True)
 
 
+@version("0.3.3")
+class ResumeSchema(ScheduledSchema):
+    class Meta:
+        object_class = state.Resume
+
+
 @version("0.3.3")
 class RetryingSchema(ScheduledSchema):
     class Meta:
@@ -145,6 +151,7 @@ class StateSchema(OneOfSchema):
         "CachedState": CachedStateSchema,
         "Scheduled": ScheduledSchema,
         "Retrying": RetryingSchema,
+        "Resume": ResumeSchema,
         "Running": RunningSchema,
         "Finished": FinishedSchema,
         "Success": SuccessSchema,


Test patch:
diff --git a/tests/engine/test_state.py b/tests/engine/test_state.py
--- a/tests/engine/test_state.py
+++ b/tests/engine/test_state.py
@@ -12,6 +12,7 @@
     Mapped,
     Paused,
     Pending,
+    Resume,
     Retrying,
     Running,
     Scheduled,
@@ -230,6 +231,9 @@ class TestStateHierarchy:
     def test_scheduled_is_pending(self):
         assert issubclass(Scheduled, Pending)
 
+    def test_resume_is_scheduled(self):
+        assert issubclass(Resume, Scheduled)
+
     def test_mapped_is_success(self):
         assert issubclass(Mapped, Success)
 
@@ -298,6 +302,16 @@ def test_state_type_methods_with_scheduled_state(self):
         assert not state.is_successful()
         assert not state.is_failed()
 
+    def test_state_type_methods_with_resume_state(self):
+        state = Resume()
+        assert state.is_pending()
+        assert not state.is_running()
+        assert not state.is_finished()
+        assert not state.is_skipped()
+        assert state.is_scheduled()
+        assert not state.is_successful()
+        assert not state.is_failed()
+
     def test_state_type_methods_with_retry_state(self):
         state = Retrying()
         assert state.is_pending()
diff --git a/tests/engine/test_task_runner.py b/tests/engine/test_task_runner.py
--- a/tests/engine/test_task_runner.py
+++ b/tests/engine/test_task_runner.py
@@ -26,6 +26,7 @@
     Mapped,
     Paused,
     Pending,
+    Resume,
     Retrying,
     Running,
     Scheduled,
@@ -345,14 +346,14 @@ def test_task_that_starts_failed_doesnt_get_retried():
     assert state.is_failed()
 
 
-class TestGetRunCount:
+class TestUpdateContextFromState:
     @pytest.mark.parametrize(
         "state", [Success(), Failed(), Pending(), Scheduled(), Skipped(), CachedState()]
     )
     def test_states_without_run_count(self, state):
         with prefect.context() as ctx:
             assert "_task_run_count" not in ctx
-            new_state = TaskRunner(Task()).get_run_count(state)
+            new_state = TaskRunner(Task()).update_context_from_state(state)
             assert ctx._task_run_count == 1
             assert new_state is state
 
@@ -368,10 +369,27 @@ def test_states_without_run_count(self, state):
     def test_states_with_run_count(self, state):
         with prefect.context() as ctx:
             assert "_task_run_count" not in ctx
-            new_state = TaskRunner(Task()).get_run_count(state)
+            new_state = TaskRunner(Task()).update_context_from_state(state)
             assert ctx._task_run_count == state.run_count + 1
             assert new_state is state
 
+    def test_task_runner_puts_resume_in_context_if_state_is_resume(self):
+        with prefect.context() as ctx:
+            assert "resume" not in ctx
+            new_state = TaskRunner(Task()).update_context_from_state(Resume())
+            assert ctx.resume is True
+
+    @pytest.mark.parametrize(
+        "state", [Success(), Failed(), Pending(), Scheduled(), Skipped(), CachedState()]
+    )
+    def test_task_runner_doesnt_put_resume_in_context_if_state_is_not_resume(
+        self, state
+    ):
+        with prefect.context() as ctx:
+            assert "resume" not in ctx
+            new_state = TaskRunner(Task()).update_context_from_state(state)
+            assert "resume" not in ctx
+
 
 class TestCheckUpstreamFinished:
     def test_with_empty_set(self):
diff --git a/tests/serialization/test_states.py b/tests/serialization/test_states.py
--- a/tests/serialization/test_states.py
+++ b/tests/serialization/test_states.py
@@ -40,6 +40,8 @@ def complex_states():
         state.Retrying(start_time=naive_dt, run_count=3),
         state.Scheduled(start_time=utc_dt),
         state.Scheduled(start_time=naive_dt),
+        state.Resume(start_time=utc_dt),
+        state.Resume(start_time=naive_dt),
         cached_state,
         cached_state_naive,
         state.Success(result=complex_result, cached=cached_state),