Repository: PrefectHQ/prefect

Issue:
Suport for Dask Performance Report on Task/Flow Runs
## Current behavior

It's very hard to get a Dask performance report for a Flow. These performance reports are very useful for memory profiling and checking if tasks cause memory issues. More on the performance report can be found [here](https://distributed.dask.org/en/latest/api.html#distributed.performance_report). It generates an HTML doc. This would also help debug flows that end up in a stuck RUNNING state.

Without Prefect, using it looks like:

```python
with performance_report(filename="dask-report.html"):
    return dd.read_csv(filepath)
```

But with Prefect, it becomes very hard to use it. Trying 

```python
@task(checkpoint=False)
def using_checkpoint_false(filepath):
    with worker_client():
        with performance_report(filename="dask-report.html"):
            return dd.read_csv(filepath)
```

only gives the report for one task. This means that for a mapped task with 10 elements, you get 10 HTML files.

The other attempt

```python
with Flow("example") as flow:
    with performance_report(...):
        some_task()
        other_task()
```

just throws an error.


## Proposed behavior

I may be wrong but this needs to be done at the `executor.submit()` level.

```python
    with performance_report(...):
        executor.submit()
```

which means that we need changes in the FlowRunner to support this. The FlowRunner takes a mapped task and calls `executor.submit()` for each map iterator. This means that to get the report for a mapped task, it needs to be done at the FlowRunner level.

## Open Questions
Do we support this at the task level or Flow level?
How can it be done that doesn't make the solution Executor specific? (the FlowRunner is currently executor agnostic)



Discussion:

Resolving patch:
diff --git a/src/prefect/executors/dask.py b/src/prefect/executors/dask.py
--- a/src/prefect/executors/dask.py
+++ b/src/prefect/executors/dask.py
@@ -6,6 +6,8 @@
 from contextlib import contextmanager
 from typing import Any, Callable, Iterator, TYPE_CHECKING, Union, Optional, Dict
 
+from prefect.utilities.compatibility import nullcontext
+
 from prefect import context
 from prefect.executors.base import Executor
 from prefect.utilities.importtools import import_object
@@ -96,6 +98,8 @@ class name (e.g. `"distributed.LocalCluster"`), or the class itself.
             `debug=True` will increase dask's logging level, providing
             potentially useful debug info. Defaults to the `debug` value in
             your Prefect configuration.
+        - performance_report_path (str, optional): An optional path for the [dask performance
+            report](https://distributed.dask.org/en/latest/api.html#distributed.performance_report).
 
     Examples:
 
@@ -134,6 +138,7 @@ def __init__(
         adapt_kwargs: dict = None,
         client_kwargs: dict = None,
         debug: bool = None,
+        performance_report_path: str = None,
     ):
         if address is None:
             address = context.config.engine.executor.dask.address or None
@@ -186,6 +191,8 @@ def __init__(
         # A ref to a background task subscribing to dask cluster events
         self._watch_dask_events_task = None  # type: Optional[concurrent.futures.Future]
 
+        self.performance_report_path = performance_report_path
+
         super().__init__()
 
     @contextmanager
@@ -198,7 +205,13 @@ def start(self) -> Iterator[None]:
         if sys.platform != "win32":
             # Fix for https://github.com/dask/distributed/issues/4168
             import multiprocessing.popen_spawn_posix  # noqa
-        from distributed import Client
+        from distributed import Client, performance_report
+
+        performance_report_context = (
+            performance_report(self.performance_report_path)
+            if self.performance_report_path
+            else nullcontext()
+        )
 
         try:
             if self.address is not None:
@@ -206,12 +219,14 @@ def start(self) -> Iterator[None]:
                     "Connecting to an existing Dask cluster at %s", self.address
                 )
                 with Client(self.address, **self.client_kwargs) as client:
-                    self.client = client
-                    try:
-                        self._pre_start_yield()
-                        yield
-                    finally:
-                        self._post_start_yield()
+
+                    with performance_report_context:
+                        self.client = client
+                        try:
+                            self._pre_start_yield()
+                            yield
+                        finally:
+                            self._post_start_yield()
             else:
                 assert callable(self.cluster_class)  # mypy
                 assert isinstance(self.cluster_kwargs, dict)  # mypy
@@ -229,12 +244,13 @@ def start(self) -> Iterator[None]:
                     if self.adapt_kwargs:
                         cluster.adapt(**self.adapt_kwargs)
                     with Client(cluster, **self.client_kwargs) as client:
-                        self.client = client
-                        try:
-                            self._pre_start_yield()
-                            yield
-                        finally:
-                            self._post_start_yield()
+                        with performance_report_context:
+                            self.client = client
+                            try:
+                                self._pre_start_yield()
+                                yield
+                            finally:
+                                self._post_start_yield()
         finally:
             self.client = None
 
@@ -423,6 +439,27 @@ def wait(self, futures: Any) -> Any:
 
         return self.client.gather(futures)
 
+    @property
+    def performance_report(self) -> str:
+        """The performance report html string."""
+        if self.performance_report_path is None:
+            self.logger.warning(
+                "Executor was not configured to generate performance report"
+            )
+            return ""
+        self.logger.debug(
+            f"Retreiving dask performance report from {self.performance_report_path!r}"
+        )
+        try:
+            with open(self.performance_report_path, "r", encoding="utf-8") as f:
+                report = f.read()
+                return report
+        except Exception as exc:
+            self.logger.error(
+                f"Failed to get dask performance report with exception {exc}"
+            )
+            return ""
+
 
 def _multiprocessing_pool_initializer() -> None:
     """Initialize a process used in a `multiprocssing.Pool`.


Test patch:
diff --git a/tests/executors/test_executors.py b/tests/executors/test_executors.py
--- a/tests/executors/test_executors.py
+++ b/tests/executors/test_executors.py
@@ -2,6 +2,7 @@
 import os
 import random
 import sys
+import tempfile
 import threading
 import time
 import uuid
@@ -389,6 +390,27 @@ def test_debug_is_converted_to_silence_logs(self, debug):
         level = logging.WARNING if debug else logging.CRITICAL
         assert executor.cluster_kwargs["silence_logs"] == level
 
+    def test_performance_report(self):
+        # should not error
+        assert DaskExecutor().performance_report == ""
+        assert (
+            DaskExecutor(
+                performance_report_path="not a readable path"
+            ).performance_report
+            == ""
+        )
+
+        with tempfile.TemporaryDirectory() as report_dir:
+            with open(f"{report_dir}/report.html", "w") as fp:
+                fp.write("very advanced report")
+
+            assert (
+                DaskExecutor(
+                    performance_report_path=f"{report_dir}/report.html"
+                ).performance_report
+                == "very advanced report"
+            )
+
     def test_cant_specify_both_address_and_cluster_class(self):
         with pytest.raises(ValueError):
             DaskExecutor(