Repository: PrefectHQ/prefect

Issue:
Proposal: Create new `ClientFailed` state
This state will be a new `MetaState` which Tasks can enter into whenever `client.set_state` fails.

 This is designed specifically for the following situation: imagine a Flow with two tasks: Task A -> Task B.  Task B has an `all_failed` trigger.  Task A requests retries.

Task A tries to run but an error is raised; Task A enters a Failed state first, but for _some reason_ Task A fails to set its state in Cloud.  This produces an `ENDRUN` signal and Task A _ends in a `Failed` state_.  Task B now sees this and runs successfully.  I don't think this is intended behavior.

If instead, a bad client call produced a `ClientFailed` meta state which:
- is not considered finished
- _wraps_ the previous state

We could avoid the incorrect execution of Task B.

**cc:** @jlowin 


Discussion:
@cicdw Just for clarification, this state is only valid when in context of cloud?
Handling erroneous state handlers leads to recursive logic very quickly (who should handle the resulting state handler state?), so while this state is broadly applicable, at this time we will only introduce it for a set of known Cloud issues (like the client being unable to communicate via network, or a concurrency lock)

Resolving patch:
diff --git a/src/prefect/engine/cloud/task_runner.py b/src/prefect/engine/cloud/task_runner.py
--- a/src/prefect/engine/cloud/task_runner.py
+++ b/src/prefect/engine/cloud/task_runner.py
@@ -10,7 +10,7 @@
 from prefect.engine.result import NoResult, Result
 from prefect.engine.result_handlers import ResultHandler
 from prefect.engine.runner import ENDRUN, call_state_handlers
-from prefect.engine.state import Cached, Failed, Mapped, State
+from prefect.engine.state import Cached, ClientFailed, Failed, Mapped, State
 from prefect.engine.task_runner import TaskRunner, TaskRunnerInitializeResult
 from prefect.utilities.graphql import with_args
 
@@ -95,7 +95,7 @@ def call_runner_target_handlers(self, old_state: State, new_state: State) -> Sta
             self.logger.debug(
                 "Failed to set task state with error: {}".format(repr(exc))
             )
-            raise ENDRUN(state=new_state)
+            raise ENDRUN(state=ClientFailed(state=new_state))
 
         if version is not None:
             prefect.context.update(task_run_version=version + 1)  # type: ignore
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
@@ -310,6 +310,24 @@ def __init__(
         self.state = state
 
 
+class ClientFailed(_MetaState):
+    """
+    The `ClientFailed` state is used to indicate that the Prefect Client failed
+    to set a task run state, and thus this task run should exit, without triggering any
+    downstream task runs.
+
+    The `ClientFailed` state should be initialized with another state, which it wraps. The
+    wrapped state is the state which the client attempted to set in the database, but failed to
+    for some reason.
+
+    Args:
+        - message (string): a message for the state.
+        - result (Any, optional): Defaults to `None`.
+        - state (State): the `State` state that the task run ended in
+
+    """
+
+
 class Submitted(_MetaState):
     """
     The `Submitted` state is used to indicate that another state, usually a `Scheduled`
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
@@ -59,6 +59,11 @@ class MetaStateSchema(BaseStateSchema):
     state = fields.Nested("StateSchema", allow_none=True)
 
 
+class ClientFailedSchema(MetaStateSchema):
+    class Meta:
+        object_class = state.ClientFailed
+
+
 class SubmittedSchema(MetaStateSchema):
     class Meta:
         object_class = state.Submitted
@@ -185,6 +190,7 @@ class StateSchema(OneOfSchema):
     # map class name to schema
     type_schemas = {
         "Cached": CachedSchema,
+        "ClientFailed": ClientFailedSchema,
         "Failed": FailedSchema,
         "Finished": FinishedSchema,
         "Mapped": MappedSchema,


Test patch:
diff --git a/tests/engine/cloud/test_cloud_task_runner.py b/tests/engine/cloud/test_cloud_task_runner.py
--- a/tests/engine/cloud/test_cloud_task_runner.py
+++ b/tests/engine/cloud/test_cloud_task_runner.py
@@ -21,6 +21,7 @@
 from prefect.engine.runner import ENDRUN
 from prefect.engine.state import (
     Cached,
+    ClientFailed,
     Failed,
     Finished,
     Mapped,
@@ -127,7 +128,8 @@ def raise_error():
     ## an ENDRUN will cause the TaskRunner to return the most recently computed state
     res = CloudTaskRunner(task=raise_error).run(context={"map_index": 1})
     assert set_task_run_state.called
-    assert res.is_running()
+    assert isinstance(res, ClientFailed)
+    assert res.state.is_running()
 
 
 def test_task_runner_queries_for_cached_states_if_task_has_caching(client):
@@ -138,7 +140,7 @@ def cached_task():
     state = Cached(
         cached_result_expiration=datetime.datetime.utcnow()
         + datetime.timedelta(days=1),
-        result=99,
+        result=Result(99, JSONResultHandler()),
     )
     old_state = Cached(
         cached_result_expiration=datetime.datetime.utcnow()
@@ -155,19 +157,21 @@ def cached_task():
 
 
 def test_task_runner_validates_cached_states_if_task_has_caching(client):
-    @prefect.task(cache_for=datetime.timedelta(minutes=1))
+    @prefect.task(
+        cache_for=datetime.timedelta(minutes=1), result_handler=JSONResultHandler()
+    )
     def cached_task():
         return 42
 
     state = Cached(
         cached_result_expiration=datetime.datetime.utcnow()
         - datetime.timedelta(minutes=2),
-        result=99,
+        result=Result(99, JSONResultHandler()),
     )
     old_state = Cached(
         cached_result_expiration=datetime.datetime.utcnow()
         - datetime.timedelta(days=1),
-        result=13,
+        result=Result(13, JSONResultHandler()),
     )
     client.get_latest_cached_states = MagicMock(return_value=[state, old_state])
 
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 @@
 from prefect.engine.result_handlers import JSONResultHandler, LocalResultHandler
 from prefect.engine.state import (
     Cached,
+    ClientFailed,
     Failed,
     Finished,
     Mapped,
@@ -369,6 +370,21 @@ def test_state_type_methods_with_retry_state(self):
         assert not state.is_meta_state()
         assert not state.is_submitted()
 
+    def test_state_type_methods_with_clientfailed_state(self):
+        state = ClientFailed()
+        assert not state.is_cached()
+        assert not state.is_retrying()
+        assert not state.is_pending()
+        assert not state.is_running()
+        assert not state.is_finished()
+        assert not state.is_skipped()
+        assert not state.is_scheduled()
+        assert not state.is_successful()
+        assert not state.is_failed()
+        assert not state.is_mapped()
+        assert state.is_meta_state()
+        assert not state.is_submitted()
+
     def test_state_type_methods_with_submitted_state(self):
         state = Submitted()
         assert not state.is_cached()