Repository: PrefectHQ/prefect

Issue:
Investigate size of cloudpickle payloads passed to `executor.submit`
Currently, we submit a class method to `executor.submit`.  This ultimately has the effect of serializing large portions of the class that could actually be imported.

I believe that our unusually large submissions here are partially responsible for eating up Dask Scheduler CPU cycles, but we'll need to investigate further to decide if that's true.

Ultimately this ticket is for:
- determining the effect of our payload size on the dask scheduler
- reducing our payload sizes where possible


Discussion:

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
@@ -40,25 +40,18 @@ class CloudTaskRunner(TaskRunner):
             (current) state, with the following signature: `state_handler(TaskRunner, old_state, new_state) -> State`;
             If multiple functions are passed, then the `new_state` argument will be the
             result of the previous handler.
-        - result (Result, optional): the result instance used to retrieve and store task results during execution;
-            if not provided, will default to the one on the provided Task
-        - default_result (Result, optional): the fallback result type to use for retrieving and storing state results
-            during execution (to be used on upstream inputs if they don't provide their own results)
+        - flow_result: the result instance configured for the flow (if any)
     """
 
     def __init__(
         self,
         task: Task,
         state_handlers: Iterable[Callable] = None,
-        result: Result = None,
-        default_result: Result = None,
+        flow_result: Result = None,
     ) -> None:
         self.client = Client()
         super().__init__(
-            task=task,
-            state_handlers=state_handlers,
-            result=result,
-            default_result=default_result,
+            task=task, state_handlers=state_handlers, flow_result=flow_result,
         )
 
     def _heartbeat(self) -> bool:
@@ -269,11 +262,11 @@ def load_results(
         try:
             for edge, upstream_state in upstream_states.items():
                 upstream_states[edge] = upstream_state.load_result(
-                    edge.upstream_task.result or self.default_result
+                    edge.upstream_task.result or self.flow_result
                 )
                 if edge.key is not None:
                     upstream_results[edge.key] = (
-                        edge.upstream_task.result or self.default_result
+                        edge.upstream_task.result or self.flow_result
                     )
 
             state.load_cached_results(upstream_results)
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
@@ -88,7 +88,6 @@ def __init__(
         task_runner_cls: type = None,
         state_handlers: Iterable[Callable] = None,
     ):
-        self.context = prefect.context.to_dict()
         self.flow = flow
         if task_runner_cls is None:
             task_runner_cls = prefect.engine.get_default_task_runner_class()
@@ -465,24 +464,23 @@ def get_flow_run_state(
 
                     ## wait on upstream states to determine the width of the pipeline
                     ## this is the key to depth-first execution
-                    upstream_states.update(
-                        executor.wait(
-                            {e: state for e, state in upstream_states.items()}
-                        )
+                    upstream_states = executor.wait(
+                        {e: state for e, state in upstream_states.items()}
                     )
-
                     ## we submit the task to the task runner to determine if
                     ## we can proceed with mapping - if the new task state is not a Mapped
                     ## state then we don't proceed
                     task_states[task] = executor.wait(
                         executor.submit(
-                            self.run_task,
+                            run_task,
                             task=task,
                             state=task_state,  # original state
                             upstream_states=upstream_states,
                             context=dict(
                                 prefect.context, **task_contexts.get(task, {})
                             ),
+                            flow_result=self.flow.result,
+                            task_runner_cls=self.task_runner_cls,
                             task_runner_state_handlers=task_runner_state_handlers,
                             upstream_mapped_states=upstream_mapped_states,
                             is_mapped_parent=True,
@@ -518,7 +516,7 @@ def get_flow_run_state(
                         extra_context.update(task_name=f"{task.name}[{idx}]")
                         submitted_states.append(
                             executor.submit(
-                                self.run_task,
+                                run_task,
                                 task=task,
                                 state=current_state,
                                 upstream_states=states,
@@ -527,6 +525,8 @@ def get_flow_run_state(
                                     **task_contexts.get(task, {}),
                                     map_index=idx,
                                 ),
+                                flow_result=self.flow.result,
+                                task_runner_cls=self.task_runner_cls,
                                 task_runner_state_handlers=task_runner_state_handlers,
                                 upstream_mapped_states=upstream_mapped_states,
                                 extra_context=extra_context,
@@ -535,14 +535,15 @@ def get_flow_run_state(
                     if isinstance(task_states.get(task), Mapped):
                         mapped_children[task] = submitted_states  # type: ignore
 
-                # -- run the task
                 else:
                     task_states[task] = executor.submit(
-                        self.run_task,
+                        run_task,
                         task=task,
                         state=task_state,
                         upstream_states=upstream_states,
                         context=dict(prefect.context, **task_contexts.get(task, {})),
+                        flow_result=self.flow.result,
+                        task_runner_cls=self.task_runner_cls,
                         task_runner_state_handlers=task_runner_state_handlers,
                         upstream_mapped_states=upstream_mapped_states,
                         extra_context=extra_context,
@@ -638,66 +639,56 @@ def determine_final_state(
 
         return state
 
-    def run_task(
-        self,
-        task: Task,
-        state: State,
-        upstream_states: Dict[Edge, State],
-        context: Dict[str, Any],
-        task_runner_state_handlers: Iterable[Callable],
-        is_mapped_parent: bool = False,
-        upstream_mapped_states: Dict[Edge, list] = None,
-    ) -> State:
-        """
-
-        Runs a specific task. This method is intended to be called by submitting it to
-        an executor.
-
-        Args:
-            - task (Task): the task to run
-            - state (State): starting state for the Flow. Defaults to
-                `Pending`
-            - upstream_states (Dict[Edge, State]): dictionary of upstream states
-            - context (Dict[str, Any]): a context dictionary for the task run
-            - task_runner_state_handlers (Iterable[Callable]): A list of state change
-                handlers that will be provided to the task_runner, and called whenever a task changes
-                state.
-            - is_mapped_parent (bool): a boolean indicating whether this task run is the run of a parent
-                mapped task
-            - upstream_mapped_states (Dict[Edge, list]): dictionary of upstream states corresponding to
-                mapped children dependencies
-
-        Returns:
-            - State: `State` representing the final post-run state of the `Flow`.
-
-        """
-        upstream_mapped_states = upstream_mapped_states or {}
-
-        with prefect.context(self.context):
-            default_result = task.result or self.flow.result
-            task_runner = self.task_runner_cls(
-                task=task,
-                state_handlers=task_runner_state_handlers,
-                result=default_result or Result(),
-                default_result=self.flow.result,
-            )
 
-            # if this task reduces over a mapped state, make sure its children have finished
-            for edge, upstream_state in upstream_states.items():
-
-                # if the upstream state is Mapped, wait until its results are all available
-                if not edge.mapped and upstream_state.is_mapped():
-                    assert isinstance(upstream_state, Mapped)  # mypy assert
-                    upstream_state.map_states = upstream_mapped_states.get(
-                        edge, upstream_state.map_states
-                    )
-                    upstream_state.result = [
-                        s.result for s in upstream_state.map_states
-                    ]
+def run_task(
+    task: Task,
+    state: State,
+    upstream_states: Dict[Edge, State],
+    context: Dict[str, Any],
+    flow_result: Result,
+    task_runner_cls: Callable,
+    task_runner_state_handlers: Iterable[Callable],
+    upstream_mapped_states: Dict[Edge, list],
+    is_mapped_parent: bool = False,
+) -> State:
+    """
+    Runs a specific task. This method is intended to be called by submitting it to
+    an executor.
 
-            return task_runner.run(
-                state=state,
-                upstream_states=upstream_states,
-                context=context,
-                is_mapped_parent=is_mapped_parent,
-            )
+    Args:
+        - task (Task): the task to run
+        - state (State): starting state for the Flow. Defaults to `Pending`
+        - task_runner_cls (Callable): the `TaskRunner` class to use
+        - upstream_states (Dict[Edge, State]): dictionary of upstream states
+        - context (Dict[str, Any]): a context dictionary for the task run
+        - flow_result (Result): the `Result` associated with the flow (if any)
+        - task_runner_state_handlers (Iterable[Callable]): A list of state change
+            handlers that will be provided to the task_runner, and called
+            whenever a task changes state.
+        - upstream_mapped_states (Dict[Edge, list]): dictionary of upstream states
+            corresponding to mapped children dependencies
+        - is_mapped_parent (bool): a boolean indicating whether this task run is the
+            run of a parent mapped task
+
+    Returns:
+        - State: `State` representing the final post-run state of the `Flow`.
+    """
+    with prefect.context(context):
+        # Update upstream_states with info from upstream_mapped_states
+        for edge, upstream_state in upstream_states.items():
+            if not edge.mapped and upstream_state.is_mapped():
+                assert isinstance(upstream_state, Mapped)  # mypy assert
+                upstream_state.map_states = upstream_mapped_states.get(
+                    edge, upstream_state.map_states
+                )
+                upstream_state.result = [s.result for s in upstream_state.map_states]
+        task_runner = task_runner_cls(
+            task=task,
+            state_handlers=task_runner_state_handlers,
+            flow_result=flow_result,
+        )
+        return task_runner.run(
+            state=state,
+            upstream_states=upstream_states,
+            is_mapped_parent=is_mapped_parent,
+        )
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
@@ -69,18 +69,14 @@ class TaskRunner(Runner):
             (current) state, with the following signature: `state_handler(TaskRunner, old_state, new_state) -> Optional[State]`;
             If multiple functions are passed, then the `new_state` argument will be the
             result of the previous handler.
-        - result (Result, optional): the result type to use for retrieving and storing state results
-            during execution (if the Task doesn't already have one)
-        - default_result (Result, optional): the fallback result type to use for retrieving and storing state results
-            during execution (to be used on upstream inputs if they don't provide their own results)
+        - flow_result: the result instance configured for the flow (if any)
     """
 
     def __init__(
         self,
         task: Task,
         state_handlers: Iterable[Callable] = None,
-        result: Result = None,
-        default_result: Result = None,
+        flow_result: Result = None,
     ):
         self.context = prefect.context.to_dict()
         self.task = task
@@ -90,10 +86,10 @@ def __init__(
         if task.result:
             self.result = task.result
         else:
-            self.result = Result().copy() if result is None else result.copy()
+            self.result = Result() if flow_result is None else flow_result
             if self.task.target:
                 self.result.location = self.task.target
-        self.default_result = default_result or Result()
+        self.flow_result = flow_result
         super().__init__(state_handlers=state_handlers)
 
     def __repr__(self) -> str:
@@ -215,7 +211,7 @@ def run(
             - `State` object representing the final post-run state of the Task
         """
         upstream_states = upstream_states or {}
-        context = context or {}
+        context = context or prefect.context.to_dict()
         map_index = context.setdefault("map_index", None)
         context["task_full_name"] = "{name}{index}".format(
             name=self.task.name,


Test patch:
diff --git a/tests/core/test_task_map.py b/tests/core/test_task_map.py
--- a/tests/core/test_task_map.py
+++ b/tests/core/test_task_map.py
@@ -442,7 +442,7 @@ def register(x):
 @pytest.mark.parametrize(
     "executor", ["local", "sync", "mproc", "mthread"], indirect=True
 )
-def test_map_preserves_flowrunners_initial_context(executor):
+def test_map_preserves_flowrunners_run_context(executor):
     @task
     def whats_id():
         return prefect.context.get("special_id")
@@ -452,8 +452,8 @@ def whats_id():
 
     with prefect.context(special_id="FOOBAR"):
         runner = FlowRunner(flow=flow)
+        flow_state = runner.run(return_tasks=[result])
 
-    flow_state = runner.run(return_tasks=[result])
     assert flow_state.is_successful()
     assert flow_state.result[result].result == ["FOOBAR"] * 10
 
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
@@ -550,26 +550,6 @@ def downstream(l):
     assert flow_state.result[final].result == 100
 
 
-def test_cloud_flow_runner_can_successfully_initialize_cloud_task_runners():
-    """
-    After the context refactor wherein config settings were pulled from context.config,
-    there were various errors related to `prefect.context(self.context.to_dict())`
-    caused by the Context object not creating a nested DotDict structure.  The main
-    symptom of this was when a CloudFlowRunner submitted a job to a dask worker and an error
-    was raised: `dict has no attribute cloud`
-
-    Note: DotDicts have been replaced by Box objects
-    """
-    fr = CloudFlowRunner(flow=prefect.Flow(name="test"))
-    fr.run_task(
-        task=MagicMock(),
-        state=None,
-        upstream_states=dict(),
-        context=dict(),
-        task_runner_state_handlers=[],
-    )
-
-
 def test_cloud_task_runners_submitted_to_remote_machines_respect_original_config(
     monkeypatch,
 ):
@@ -579,22 +559,13 @@ def test_cloud_task_runners_submitted_to_remote_machines_respect_original_config
     settings which were present on the original machine are respected in the remote job, reflected
     here by having the CloudHandler called during logging and the special values present in context.
     """
+    from prefect.engine.flow_runner import run_task
 
-    class CustomFlowRunner(CloudFlowRunner):
-        def run_task(self, *args, **kwargs):
-            with prefect.utilities.configuration.set_temporary_config(
-                {"logging.log_to_cloud": False, "cloud.auth_token": ""}
-            ):
-                return super().run_task(*args, **kwargs)
-
-    @prefect.task(result=PrefectResult())
-    def log_stuff():
-        logger = prefect.context.get("logger")
-        logger.critical("important log right here")
-        return (
-            prefect.context.config.special_key,
-            prefect.context.config.cloud.auth_token,
-        )
+    def my_run_task(*args, **kwargs):
+        with prefect.utilities.configuration.set_temporary_config(
+            {"logging.log_to_cloud": False, "cloud.auth_token": ""}
+        ):
+            return run_task(*args, **kwargs)
 
     calls = []
 
@@ -614,11 +585,21 @@ def get_flow_run_info(self, *args, **kwargs):
                 task_runs=[MagicMock(task_slug=log_stuff.slug, id="TESTME")],
             )
 
+    monkeypatch.setattr("prefect.engine.flow_runner.run_task", my_run_task)
     monkeypatch.setattr("prefect.client.Client", Client)
     monkeypatch.setattr("prefect.engine.cloud.task_runner.Client", Client)
     monkeypatch.setattr("prefect.engine.cloud.flow_runner.Client", Client)
     prefect.utilities.logging.prefect_logger.handlers[-1].client = Client()
 
+    @prefect.task(result=PrefectResult())
+    def log_stuff():
+        logger = prefect.context.get("logger")
+        logger.critical("important log right here")
+        return (
+            prefect.context.config.special_key,
+            prefect.context.config.cloud.auth_token,
+        )
+
     with prefect.utilities.configuration.set_temporary_config(
         {
             "logging.log_to_cloud": True,
@@ -627,10 +608,8 @@ def get_flow_run_info(self, *args, **kwargs):
         }
     ):
         # captures config at init
-        runner = CustomFlowRunner(flow=prefect.Flow("test", tasks=[log_stuff]))
-        flow_state = runner.run(
-            return_tasks=[log_stuff], task_contexts={log_stuff: dict(special_key=99)}
-        )
+        flow = prefect.Flow("test", tasks=[log_stuff])
+        flow_state = flow.run(task_contexts={log_stuff: dict(special_key=99)})
 
     assert flow_state.is_successful()
     assert flow_state.result[log_stuff].result == (42, "original")
@@ -642,7 +621,7 @@ def get_flow_run_info(self, *args, **kwargs):
     loggers = [c["name"] for c in logs]
     assert set(loggers) == {
         "prefect.CloudTaskRunner",
-        "prefect.CustomFlowRunner",
+        "prefect.CloudFlowRunner",
         "prefect.log_stuff",
     }
 
diff --git a/tests/engine/test_flow_runner.py b/tests/engine/test_flow_runner.py
--- a/tests/engine/test_flow_runner.py
+++ b/tests/engine/test_flow_runner.py
@@ -5,7 +5,6 @@
 import time
 from unittest.mock import MagicMock
 
-import cloudpickle
 import pendulum
 import pytest
 
@@ -33,9 +32,8 @@
     TimedOut,
     TriggerFailed,
 )
-from prefect.engine.task_runner import TaskRunner
 from prefect.tasks.secrets import PrefectSecret
-from prefect.triggers import any_failed, manual_only
+from prefect.triggers import manual_only
 from prefect.utilities.debug import raise_on_exception
 
 
@@ -1321,46 +1319,18 @@ def test_paused_tasks_stay_paused_when_run():
 
 
 class TestContext:
-    def test_flow_runner_inits_with_current_context(self):
-        runner = FlowRunner(Flow(name="test"))
-        assert isinstance(runner.context, dict)
-        assert "chris" not in runner.context
-
-        with prefect.context(chris="foo"):
-            runner2 = prefect.engine.task_runner.TaskRunner(Task())
-            assert "chris" in runner2.context
-
-        assert "chris" not in prefect.context
-        assert runner2.context["chris"] == "foo"
-
-    def test_flow_runner_passes_along_its_init_context_to_tasks(self):
+    def test_flow_runner_passes_along_its_run_context_to_tasks(self):
         @prefect.task
         def grab_key():
             return prefect.context["THE_ANSWER"]
 
         with prefect.context(THE_ANSWER=42):
             runner = FlowRunner(Flow(name="test", tasks=[grab_key]))
+            flow_state = runner.run(return_tasks=[grab_key])
 
-        flow_state = runner.run(return_tasks=[grab_key])
         assert flow_state.is_successful()
         assert flow_state.result[grab_key].result == 42
 
-    def test_flow_runner_passes_along_its_init_context_to_tasks_after_serialization(
-        self,
-    ):
-        @prefect.task
-        def grab_key():
-            return prefect.context["THE_ANSWER"]
-
-        with prefect.context(THE_ANSWER=42):
-            prerunner = FlowRunner(Flow(name="test", tasks=[grab_key]))
-
-        runner = cloudpickle.loads(cloudpickle.dumps(prerunner))
-
-        flow_state = runner.run(return_tasks=list(runner.flow.tasks))
-        assert flow_state.is_successful()
-        assert flow_state.result[runner.flow.tasks.pop()].result == 42
-
     def test_flow_runner_provides_scheduled_start_time(self):
         @prefect.task
         def return_scheduled_start_time():
@@ -1474,12 +1444,13 @@ def test_task_runners_submitted_to_remote_machines_respect_original_config(monke
     here by having the CloudHandler called during logging and the special values present in context.
     """
 
-    class CustomFlowRunner(FlowRunner):
-        def run_task(self, *args, **kwargs):
-            with prefect.utilities.configuration.set_temporary_config(
-                {"logging.log_to_cloud": False, "cloud.auth_token": ""}
-            ):
-                return super().run_task(*args, **kwargs)
+    from prefect.engine.flow_runner import run_task
+
+    def my_run_task(*args, **kwargs):
+        with prefect.utilities.configuration.set_temporary_config(
+            {"logging.log_to_cloud": False, "cloud.auth_token": ""}
+        ):
+            return run_task(*args, **kwargs)
 
     calls = []
 
@@ -1487,6 +1458,7 @@ class Client:
         def write_run_logs(self, *args, **kwargs):
             calls.append(args)
 
+    monkeypatch.setattr("prefect.engine.flow_runner.run_task", my_run_task)
     monkeypatch.setattr("prefect.client.Client", Client)
 
     @prefect.task
@@ -1506,10 +1478,8 @@ def log_stuff():
         }
     ):
         # captures config at init
-        runner = CustomFlowRunner(flow=Flow("test", tasks=[log_stuff]))
-        flow_state = runner.run(
-            return_tasks=[log_stuff], task_contexts={log_stuff: dict(special_key=99)}
-        )
+        flow = Flow("test", tasks=[log_stuff])
+        flow_state = flow.run(task_contexts={log_stuff: dict(special_key=99)})
 
     assert flow_state.is_successful()
     assert flow_state.result[log_stuff].result == (42, "original")
@@ -1521,7 +1491,7 @@ def log_stuff():
     loggers = [log["name"] for call in calls for log in call[0]]
     assert set(loggers) == {
         "prefect.TaskRunner",
-        "prefect.CustomFlowRunner",
+        "prefect.FlowRunner",
         "prefect.log_stuff",
     }
 
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
@@ -1140,12 +1140,12 @@ def test_returns_success_with_correct_result_type(self):
         assert isinstance(state._result, PrefectResult)
 
     def test_success_state_without_checkpoint(self):
-        @prefect.task(checkpoint=False)
+        @prefect.task(checkpoint=False, result=PrefectResult())
         def fn(x):
             return x + 1
 
         with prefect.context(checkpointing=True):
-            new_state = TaskRunner(task=fn, result=PrefectResult()).get_task_run_state(
+            new_state = TaskRunner(task=fn).get_task_run_state(
                 state=Running(), inputs={"x": Result(1)},
             )
         assert new_state.is_successful()