Repository: PrefectHQ/prefect

Issue:
Agents Should not Fail flow runs due to version mismatches
Suppose I have two agents running, and they both pick up the same flow run.  Whichever one submits _second_ will experience an error like:
```
[{'message': 'State update failed for flow run ID 5073a4ea-61dd-400c-892a-294b93ac3633: 
provided version 1 but current version is 2.', 'locations': [{'line': 4, 'column': 13}], 'path': 
['setFlowRunStates', 'states', 0, 'id'], 'extensions': {'code': 'INTERNAL_SERVER_ERROR'}}]
```
and then [mark the Flow Run as Failed](https://github.com/PrefectHQ/prefect/blob/master/src/prefect/agent/agent.py#L209-L228).  This is incorrect behavior though, because another agent did in fact successfully submit the flow run for work, which will then proceed shortly.  

We should do a check on the exception to match against `"State update failed"` and in that case, _not_ update the state at all.

**cc:** @jlowin @zdhughes 


Discussion:

Resolving patch:
diff --git a/src/prefect/agent/agent.py b/src/prefect/agent/agent.py
--- a/src/prefect/agent/agent.py
+++ b/src/prefect/agent/agent.py
@@ -163,8 +163,8 @@ def start(self) -> None:
 
     def on_shutdown(self) -> None:
         """
-        Invoked when the event loop is exiting and the agent is shutting down. Intended 
-        as a hook for child classes to optionally implement. 
+        Invoked when the event loop is exiting and the agent is shutting down. Intended
+        as a hook for child classes to optionally implement.
         """
         pass
 
@@ -207,24 +207,42 @@ def deploy_and_update_flow_run(self, flow_run: "GraphQLResult") -> None:
         """
         # Deploy flow run and mark failed if any deployment error
         try:
+            self.update_state(flow_run)
             deployment_info = self.deploy_flow(flow_run)
-            self.update_state(flow_run, deployment_info)
+            if getattr(flow_run, "id", None):
+                self.client.write_run_logs(
+                    [
+                        dict(
+                            flowRunId=getattr(flow_run, "id"),  # type: ignore
+                            name=self.name,
+                            message="Submitted for execution: {}".format(
+                                deployment_info
+                            ),
+                            level="INFO",
+                        )
+                    ]
+                )
         except Exception as exc:
+            ## if the state update failed, we don't want to follow up with another state update
+            if "State update failed" in str(exc):
+                self.logger.debug("Updating Flow Run state failed: {}".format(str(exc)))
+                return
             self.logger.error(
                 "Logging platform error for flow run {}".format(
                     getattr(flow_run, "id", "UNKNOWN")  # type: ignore
                 )
             )
-            self.client.write_run_logs(
-                [
-                    dict(
-                        flowRunId=getattr(flow_run, "id", "UNKNOWN"),  # type: ignore
-                        name="agent",
-                        message=str(exc),
-                        level="ERROR",
-                    )
-                ]
-            )
+            if getattr(flow_run, "id", None):
+                self.client.write_run_logs(
+                    [
+                        dict(
+                            flowRunId=getattr(flow_run, "id"),  # type: ignore
+                            name=self.name,
+                            message=str(exc),
+                            level="ERROR",
+                        )
+                    ]
+                )
             self.mark_failed(flow_run=flow_run, exc=exc)
 
     def on_flow_run_deploy_attempt(self, fut: "Future", flow_run_id: str) -> None:
@@ -232,7 +250,7 @@ def on_flow_run_deploy_attempt(self, fut: "Future", flow_run_id: str) -> None:
         Indicates that a flow run deployment has been deployed (sucessfully or otherwise).
         This is intended to be a future callback hook, called in the agent's main thread
         when the background thread has completed the deploy_and_update_flow_run() call, either
-        successfully, in error, or cancelled. In all cases the agent should be open to 
+        successfully, in error, or cancelled. In all cases the agent should be open to
         attempting to deploy the flow run if the flow run id is still in the Cloud run queue.
 
         Args:
@@ -304,6 +322,8 @@ def query_flow_runs(self, tenant_id: str) -> list:
             - list: A list of GraphQLResult flow run objects
         """
         self.logger.debug("Querying for flow runs")
+        # keep a copy of what was curringly running before the query (future callbacks may be updating this set)
+        currently_submitting_flow_runs = self.submitting_flow_runs.copy()
 
         # Get scheduled flow runs from queue
         mutation = {
@@ -328,21 +348,23 @@ def query_flow_runs(self, tenant_id: str) -> list:
         # by this agent and are in the process of being submitted in the background. We do not
         # want to act on these "duplicate" flow runs until we've been assured that the background
         # thread has attempted to submit the work (successful or otherwise).
-
         flow_run_ids = set(result.data.getRunsInQueue.flow_run_ids)  # type: ignore
 
         if flow_run_ids:
             msg = "Found flow runs {}".format(result.data.getRunsInQueue.flow_run_ids)
         else:
             msg = "No flow runs found"
-        already_submitting = flow_run_ids & self.submitting_flow_runs
+
+        already_submitting = flow_run_ids & currently_submitting_flow_runs
+        target_flow_run_ids = flow_run_ids - already_submitting
+
         if already_submitting:
-            msg += " ({} already submitting)".format(len(already_submitting))
+            msg += " ({} already submitting: {})".format(
+                len(already_submitting), list(already_submitting)
+            )
 
         self.logger.debug(msg)
 
-        target_flow_run_ids = flow_run_ids - self.submitting_flow_runs
-
         # Query metadata fow flow runs found in queue
         query = {
             "query": {
@@ -389,22 +411,20 @@ def query_flow_runs(self, tenant_id: str) -> list:
             }
         }
 
-        if flow_run_ids:
+        if target_flow_run_ids:
             self.logger.debug("Querying flow run metadata")
             result = self.client.graphql(query)
             return result.data.flow_run  # type: ignore
         else:
             return []
 
-    def update_state(self, flow_run: GraphQLResult, deployment_info: str) -> None:
+    def update_state(self, flow_run: GraphQLResult) -> None:
         """
         After a flow run is grabbed this function sets the state to Submitted so it
         won't be picked up by any other processes
 
         Args:
             - flow_run (GraphQLResult): A GraphQLResult flow run object
-            - deployment_info (str): Identifier information related to the Flow Run
-                deployment
         """
         self.logger.debug(
             "Updating states for flow run {}".format(flow_run.id)  # type: ignore
@@ -422,7 +442,7 @@ def update_state(self, flow_run: GraphQLResult, deployment_info: str) -> None:
                 flow_run_id=flow_run.id,
                 version=flow_run.version,
                 state=Submitted(
-                    message="Submitted for execution. {}".format(deployment_info),
+                    message="Submitted for execution",
                     state=state.StateSchema().load(flow_run.serialized_state),
                 ),
             )
@@ -440,7 +460,7 @@ def update_state(self, flow_run: GraphQLResult, deployment_info: str) -> None:
                     task_run_id=task_run.id,
                     version=task_run.version,
                     state=Submitted(
-                        message="Submitted for execution. {}".format(deployment_info),
+                        message="Submitted for execution.",
                         state=state.StateSchema().load(task_run.serialized_state),
                     ),
                 )


Test patch:
diff --git a/tests/agent/test_agent.py b/tests/agent/test_agent.py
--- a/tests/agent/test_agent.py
+++ b/tests/agent/test_agent.py
@@ -1,4 +1,5 @@
 from unittest.mock import MagicMock
+import logging
 
 import pytest
 
@@ -180,6 +181,33 @@ def test_query_flow_runs_ignores_currently_submitting_runs(monkeypatch, runner_t
     )
 
 
+def test_query_flow_runs_does_not_use_submitting_flow_runs_directly(
+    monkeypatch, runner_token, caplog
+):
+    gql_return = MagicMock(
+        return_value=MagicMock(
+            data=MagicMock(
+                getRunsInQueue=MagicMock(flow_run_ids=["already-submitted-id"]),
+                flow_run=[{"id": "id"}],
+            )
+        )
+    )
+    client = MagicMock()
+    client.return_value.graphql = gql_return
+    monkeypatch.setattr("prefect.agent.agent.Client", client)
+
+    agent = Agent()
+    agent.logger.setLevel(logging.DEBUG)
+    copy_mock = MagicMock(return_value=set(["already-submitted-id"]))
+    agent.submitting_flow_runs = MagicMock(copy=copy_mock)
+
+    flow_runs = agent.query_flow_runs(tenant_id="id")
+
+    assert flow_runs == []
+    assert "1 already submitting: ['already-submitted-id']" in caplog.text
+    copy_mock.assert_called_once_with()
+
+
 def test_update_states_passes_no_task_runs(monkeypatch, runner_token):
     gql_return = MagicMock(
         return_value=MagicMock(
@@ -199,8 +227,7 @@ def test_update_states_passes_no_task_runs(monkeypatch, runner_token):
                 "version": 1,
                 "task_runs": [],
             }
-        ),
-        deployment_info="test",
+        )
     )
 
 
@@ -231,8 +258,7 @@ def test_update_states_passes_task_runs(monkeypatch, runner_token):
                     )
                 ],
             }
-        ),
-        deployment_info="test",
+        )
     )