Repository: PrefectHQ/prefect

Issue:
Docker Agent should pull images in the background
## Current behavior
Right now, the docker agent pulls flow images synchronously which, depending on slow connections can take up to 5 minutes for large images.

## Proposed behavior
It would be nice if the Docker agent instead pulled these images "in the background" (maybe in a background thread if possible which can be cleaned up with each heartbeat?).  This would increase the speed at which flows can be submitted in parallel and prevent any edge cases with the proposed new "agent registration" hooks.


Discussion:
I'd recommend we expand the idea of background threads for pulling images for the Docker agent to instead make deployment of flows in general be done in background threads for all agents. That is, we allow the main thread to pull work as fast as possible (up to a configured capacity) and anything regarding deploying the flow and updating state is pushed to the background. The easiest way to accomplish this is to introduce a thread pool to submit the `deploy_flow` and `update_state` as a single wrapped call. The `agent_process` method would return the number of flow runs submitted to the pool. If the pool is at capacity, we block, similar as we do today.

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
@@ -1,11 +1,14 @@
+import os
 import ast
+import functools
 import logging
 import signal
 import sys
 import time
 import threading
 from contextlib import contextmanager
-from typing import Any, Callable, Generator, Iterable, Union
+from concurrent.futures import ThreadPoolExecutor, Future
+from typing import Any, Generator, Iterable, Union, Set
 
 import pendulum
 
@@ -93,6 +96,7 @@ def __init__(
             logger.addHandler(ch)
 
         self.logger = logger
+        self.submitting_flow_runs = set()  # type: Set[str]
 
     def _verify_token(self, token: str) -> None:
         """
@@ -134,20 +138,26 @@ def start(self) -> None:
                 }
 
                 index = 0
-                while not exit_event.wait(timeout=loop_intervals[index]):
-                    self.heartbeat()
-
-                    runs = self.agent_process(tenant_id)
-                    if runs:
-                        index = 0
-                    elif index < max(loop_intervals.keys()):
-                        index += 1
-
-                    self.logger.debug(
-                        "Next query for flow runs in {} seconds".format(
-                            loop_intervals[index]
+
+                # the max workers default has changed in 3.5 and 3.8. For stable results the
+                # default 3.8 behavior is elected here.
+                max_workers = min(32, (os.cpu_count() or 1) + 4)
+
+                with ThreadPoolExecutor(max_workers=max_workers) as executor:
+                    self.logger.debug("Max Workers: {}".format(max_workers))
+                    while not exit_event.wait(timeout=loop_intervals[index]):
+                        self.heartbeat()
+
+                        if self.agent_process(executor, tenant_id):
+                            index = 0
+                        elif index < max(loop_intervals.keys()):
+                            index += 1
+
+                        self.logger.debug(
+                            "Next query for flow runs in {} seconds".format(
+                                loop_intervals[index]
+                            )
                         )
-                    )
         finally:
             self.on_shutdown()
 
@@ -187,11 +197,57 @@ def agent_connect(self) -> str:
 
         return tenant_id
 
-    def agent_process(self, tenant_id: str) -> bool:
+    def deploy_and_update_flow_run(self, flow_run: "GraphQLResult") -> None:
+        """
+        Deploy a flow run and update Cloud with the resulting deployment info.
+        If any errors occur when submitting the flow run, capture the error and log to Cloud.
+
+        Args:
+            - flow_run (GraphQLResult): The specific flow run to deploy
+        """
+        # Deploy flow run and mark failed if any deployment error
+        try:
+            deployment_info = self.deploy_flow(flow_run)
+            self.update_state(flow_run, deployment_info)
+        except Exception as exc:
+            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",
+                    )
+                ]
+            )
+            self.mark_failed(flow_run=flow_run, exc=exc)
+
+    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 
+        attempting to deploy the flow run if the flow run id is still in the Cloud run queue.
+
+        Args:
+            - fut (Future): a callback requirement, the future which has completed or been cancelled.
+            - flow_run_id (str): the id of the flow run that the future represents.
+        """
+        self.submitting_flow_runs.remove(flow_run_id)
+        self.logger.debug("Completed flow run submission (id: {})".format(flow_run_id))
+
+    def agent_process(self, executor: "ThreadPoolExecutor", tenant_id: str) -> bool:
         """
         Full process for finding flow runs, updating states, and deploying.
 
         Args:
+            - executor (ThreadPoolExecutor): the interface to submit flow deployments in background threads
             - tenant_id (str): The tenant id to use in the query
 
         Returns:
@@ -208,22 +264,17 @@ def agent_process(self, tenant_id: str) -> bool:
                     )
                 )
 
-                # Iterate over flow runs
-                for flow_run in flow_runs:
-
-                    # Deploy flow run and mark failed if any deployment error
-                    try:
-                        deployment_info = self.deploy_flow(flow_run)
-                        self.update_state(flow_run, deployment_info)
-                    except Exception as exc:
-                        self.mark_failed(flow_run=flow_run, exc=exc)
-
-                self.logger.info(
-                    "Submitted {} flow run(s) for execution.".format(len(flow_runs))
+            for flow_run in flow_runs:
+                fut = executor.submit(self.deploy_and_update_flow_run, flow_run)
+                self.submitting_flow_runs.add(flow_run.id)
+                fut.add_done_callback(
+                    functools.partial(
+                        self.on_flow_run_deploy_attempt, flow_run_id=flow_run.id
+                    )
                 )
+
         except Exception as exc:
             self.logger.error(exc)
-            self._log_flow_run_exceptions(flow_runs or [], exc)  # type: ignore
 
         return bool(flow_runs)
 
@@ -272,8 +323,25 @@ def query_flow_runs(self, tenant_id: str) -> list:
                 }
             },
         )
-        flow_run_ids = result.data.getRunsInQueue.flow_run_ids  # type: ignore
-        self.logger.debug("Found flow runs {}".format(flow_run_ids))
+
+        # we queried all of the available flow runs, however, some may have already been pulled
+        # 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
+        if already_submitting:
+            msg += " ({} already submitting)".format(len(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 = {
@@ -283,7 +351,7 @@ def query_flow_runs(self, tenant_id: str) -> list:
                     {
                         # match flow runs in the flow_run_ids list
                         "where": {
-                            "id": {"_in": flow_run_ids},
+                            "id": {"_in": list(target_flow_run_ids)},
                             "_or": [
                                 # who are EITHER scheduled...
                                 {"state": {"_eq": "Scheduled"}},
@@ -393,33 +461,6 @@ def mark_failed(self, flow_run: GraphQLResult, exc: Exception) -> None:
         )
         self.logger.error("Error while deploying flow: {}".format(repr(exc)))
 
-    def _log_flow_run_exceptions(self, flow_runs: list, exc: Exception) -> None:
-        """
-        Log platform issues to Prefect Cloud, attached to each flow run which
-        could not start.
-
-        Args:
-            - flow_runs (list): A list of GraphQLResult flow run objects
-            - exc (Exception): A caught exception to log
-        """
-        for flow_run in flow_runs:
-
-            self.logger.debug(
-                "Logging platform error for flow run {}".format(
-                    flow_run.id  # type: ignore
-                )
-            )
-            self.client.write_run_logs(
-                [
-                    dict(
-                        flowRunId=flow_run.id,  # type: ignore
-                        name="agent",
-                        message=str(exc),
-                        level="ERROR",
-                    )
-                ]
-            )
-
     def deploy_flow(self, flow_run: GraphQLResult) -> str:
         """
         Meant to be overridden by a platform specific deployment option


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
@@ -156,6 +156,30 @@ def test_query_flow_runs(monkeypatch, runner_token):
     assert flow_runs == [{"id": "id"}]
 
 
+def test_query_flow_runs_ignores_currently_submitting_runs(monkeypatch, runner_token):
+    gql_return = MagicMock(
+        return_value=MagicMock(
+            data=MagicMock(
+                getRunsInQueue=MagicMock(flow_run_ids=["id1", "id2"]),
+                flow_run=[{"id1": "id1"}],
+            )
+        )
+    )
+    client = MagicMock()
+    client.return_value.graphql = gql_return
+    monkeypatch.setattr("prefect.agent.agent.Client", client)
+
+    agent = Agent()
+    agent.submitting_flow_runs.add("id2")
+    agent.query_flow_runs(tenant_id="id")
+
+    assert len(gql_return.call_args_list) == 2
+    assert (
+        'id: { _in: ["id1"] }'
+        in list(gql_return.call_args_list[1][0][0]["query"].keys())[0]
+    )
+
+
 def test_update_states_passes_no_task_runs(monkeypatch, runner_token):
     gql_return = MagicMock(
         return_value=MagicMock(
@@ -276,6 +300,13 @@ def test_agent_connect_no_tenant_id(monkeypatch, runner_token):
         assert agent.agent_connect()
 
 
+def test_on_flow_run_deploy_attempt_removes_id(monkeypatch, runner_token):
+    agent = Agent()
+    agent.submitting_flow_runs.add("id")
+    agent.on_flow_run_deploy_attempt(None, "id")
+    assert len(agent.submitting_flow_runs) == 0
+
+
 def test_agent_process(monkeypatch, runner_token):
     gql_return = MagicMock(
         return_value=MagicMock(
@@ -308,9 +339,14 @@ def test_agent_process(monkeypatch, runner_token):
     client.return_value.graphql = gql_return
     monkeypatch.setattr("prefect.agent.agent.Client", client)
 
-    # Assert it doesn't return everything but all functions are called properly
+    executor = MagicMock()
+    future_mock = MagicMock()
+    executor.submit = MagicMock(return_value=future_mock)
+
     agent = Agent()
-    assert agent.agent_process("id")
+    assert agent.agent_process(executor, "id")
+    assert executor.submit.called
+    assert future_mock.add_done_callback.called
 
 
 def test_agent_process_no_runs_found(monkeypatch, runner_token):
@@ -328,12 +364,14 @@ def test_agent_process_no_runs_found(monkeypatch, runner_token):
     client.return_value.graphql = gql_return
     monkeypatch.setattr("prefect.agent.agent.Client", client)
 
-    # Assert it doesn't return everything but all functions are called properly
+    executor = MagicMock()
+
     agent = Agent()
-    assert not agent.agent_process("id")
+    assert not agent.agent_process(executor, "id")
+    assert not executor.submit.called
 
 
-def test_agent_logs_flow_run_exceptions(monkeypatch, runner_token):
+def test_agent_logs_flow_run_exceptions(monkeypatch, runner_token, caplog):
     gql_return = MagicMock(
         return_value=MagicMock(data=MagicMock(writeRunLogs=MagicMock(success=True)))
     )
@@ -342,46 +380,31 @@ def test_agent_logs_flow_run_exceptions(monkeypatch, runner_token):
     monkeypatch.setattr("prefect.agent.agent.Client", MagicMock(return_value=client))
 
     agent = Agent()
-    agent._log_flow_run_exceptions(
-        flow_runs=[
-            GraphQLResult(
-                {
-                    "id": "id",
-                    "serialized_state": Scheduled().serialize(),
-                    "version": 1,
-                    "task_runs": [
-                        GraphQLResult(
-                            {
-                                "id": "id",
-                                "version": 1,
-                                "serialized_state": Scheduled().serialize(),
-                            }
-                        )
-                    ],
-                }
-            )
-        ],
-        exc=ValueError("Error Here"),
+    agent.deploy_flow = MagicMock(side_effect=Exception("Error Here"))
+    agent.deploy_and_update_flow_run(
+        flow_run=GraphQLResult(
+            {
+                "id": "id",
+                "serialized_state": Scheduled().serialize(),
+                "version": 1,
+                "task_runs": [
+                    GraphQLResult(
+                        {
+                            "id": "id",
+                            "version": 1,
+                            "serialized_state": Scheduled().serialize(),
+                        }
+                    )
+                ],
+            }
+        )
     )
 
     assert client.write_run_logs.called
     client.write_run_logs.assert_called_with(
         [dict(flowRunId="id", level="ERROR", message="Error Here", name="agent")]
     )
-
-
-def test_agent_logs_flow_run_exceptions_no_flow_runs(monkeypatch, runner_token):
-    gql_return = MagicMock(
-        return_value=MagicMock(data=MagicMock(writeRunLog=MagicMock(success=True)))
-    )
-    client = MagicMock()
-    client.return_value.write_run_log = gql_return
-    monkeypatch.setattr("prefect.agent.agent.Client", MagicMock(return_value=client))
-
-    agent = Agent()
-    agent._log_flow_run_exceptions(flow_runs=[], exc=ValueError("Error Here"))
-
-    assert not client.write_run_log.called
+    assert "Logging platform error for flow run" in caplog.text
 
 
 def test_agent_process_raises_exception_and_logs(monkeypatch, runner_token):
@@ -389,7 +412,9 @@ def test_agent_process_raises_exception_and_logs(monkeypatch, runner_token):
     client.return_value.graphql.side_effect = ValueError("Error")
     monkeypatch.setattr("prefect.agent.agent.Client", client)
 
+    executor = MagicMock()
+
     agent = Agent()
     with pytest.raises(Exception):
-        agent.agent_process("id")
+        agent.agent_process(executor, "id")
         assert client.write_run_log.called