Repository: PrefectHQ/prefect

Issue:
`dask_kubernetes` Restart feature is broken 
It seems like something in prefect's codebase broke `dask_kubernetes` ability to restart and respawn worker pods, if for some reason the worker pod is killed the default behavior of  `dask_kubernetes` is to restart this pod again however when used with prefect it does not, using the code below it will create dask's cluster for testing - try killing one or two of the pods -

```python
from dask_kubernetes import KubeCluster, make_pod_spec

pod_spec = make_pod_spec(image='daskdev/dask:latest')
cluster = KubeCluster(pod_spec)
cluster.adapt(minimum=4, maximum=4)
input("pause ...")
```

Using prefect dask Kubernetes run - try killing worker pods they just die -
```python
from dask_kubernetes import KubeCluster, make_pod_spec
from prefect.run_configs import KubernetesRun
from prefect.executors import DaskExecutor
from prefect.storage import Docker
from prefect import task, Flow
import prefect
import time

storage = Docker(
    base_image="prefecthq/prefect:1.0.0-python3.9",
    python_dependencies=[
        "dask-kubernetes==2021.3.1",
    ],
)


@task()
def transform(number: int)->int:
    time.sleep(30)


with Flow(
    "Dask Kubernetes Flow",
    storage=storage,
    executor=DaskExecutor(
        cluster_class=lambda: KubeCluster(make_pod_spec(image=prefect.context.image)),
        adapt_kwargs={"minimum": 4, "maximum": 4},
    ),
    run_config=KubernetesRun(),
) as flow:
    transform.map([1, 2, 3, 4, 5, 6, 7, 8, 9, 10])

```

## Expected Behavior
I expect to see workers restarting and respawning again



Discussion:
Hi @madkinsz 

Should we prioritize this? Dask is Prefect's main feature to parallelize and distributes workloads. here are some of my notes 

1) Scheduler pod itself doesn't have to ability to restart anything. restarting takes place on the initiator process in this case  the `prefect-job` Kubernetes job
    - Confirmed by spawning cluster on my local machine and killing the python process 

2) I think it has something to do with [dask.py#L291](https://github.com/PrefectHQ/prefect/blob/master/src/prefect/executors/dask.py#L291) I think the event subscription blocks `dask_kuberenetes` event loop


3) Orion `DaskTaskRunner` [task_runners.py#L206](https://github.com/PrefectHQ/prefect/blob/orion/src/prefect/task_runners.py#L206) is cleaner and doesn't suffer from the same issue however since Orion is not ready using Orion is not a suitable option

4) If possible can we get someone from `dask-kubernetes` to take a look? maybe @jacobtomlinson ?

Hi @madkinsz 

Good news. I found the cause of this [dask.py#L253-L259](https://github.com/PrefectHQ/prefect/blob/master/src/prefect/executors/dask.py#L253-L259) - confirming my secound point - changing these lines to:-

```python
with Client(cluster, **self.client_kwargs) as client:
    self.client = client
    yield
    #with performance_report_context:
    #    self.client = client
    #    try:
    #        self._pre_start_yield()
    #        yield
    #    finally:
    #        self._post_start_yield()
```

This solves the problem and it seems like it was introduced with the performance report feature. really don't know how to reimplement this but for the time being, I think maybe we should toggle this option?

I'll be trying to narrow it further in the few next hours 
Found the root cause [dask.py#L291](https://github.com/PrefectHQ/prefect/blob/master/src/prefect/executors/dask.py#L291) I think `dask_kubernetes` is using the same event subscription to restart workers so by writing to it here we caused this behavior this feature is only for logging workers status to console. I think it's not that important as the logging is not shown anywhere else except `prefect-job` k8s job stdout. an easy workaround is just removing the logging. I'll keep the final decision to you
Hi @madkinsz 

Do we have an update on this?
Yeah, we can't just tear out event subscription, people are using it. I talked to the Dask team and can confirm that new subscribers to worker status will cause the existing subscriber to be disconnected. I might take the time to contribute a fix upstream, but I'd like to fix the behavior here first. I think the simplest fix may be adding a toggle to disable watching for worker events and add a warning if we see a user using `KubeCluster`.
We are in the process of replacing `KubeCluster` with an operator that handles scaling separately from the cluster manager. So the new cluster manager shouldn't have this problem.

It does sound like this will affect any cluster manager that uses adaptive scaling. Upstream fixes would be much appreciated!

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
@@ -104,6 +104,14 @@ class name (e.g. `"distributed.LocalCluster"`), or the class itself.
             can cause strain on the scheduler as each task needs to retrieve a client
             to check the status of the cancellation event. If set to `False`, we will
             skip this check.
+        - watch_worker_status (bool, optional): By default, Prefect subscribes to Dask
+            worker events and logs when a worker is added or removed. This provides a
+            hook for users to extend behavior on worker changes. This setting is `None`
+            by default and will be enabled unless `adapt_kwargs` is set, in which case
+            it will be disabled. Adaptive clusters often require this feature to be
+            disabled as they use the worker status events for scaling and only one
+            subscriber is allowed. If you set the value to `True` or `False`, it will be
+            respected regardless of the value of `adapt_kwargs`.
 
     Examples:
 
@@ -144,6 +152,7 @@ def __init__(
         debug: bool = None,
         performance_report_path: str = None,
         disable_cancellation_event: bool = False,
+        watch_worker_status: bool = None,
     ):
         if address is None:
             address = context.config.engine.executor.dask.address or None
@@ -187,6 +196,8 @@ def __init__(
         self.adapt_kwargs = adapt_kwargs
         self.client_kwargs = client_kwargs
         self.disable_cancellation_event = disable_cancellation_event
+        self.watch_worker_status = watch_worker_status
+
         # Runtime attributes
         self.client = None
         # These are coupled - they're either both None, or both non-None.
@@ -264,6 +275,11 @@ async def on_worker_status_changed(self, op: str, message: dict) -> None:
         """
         This method is triggered when a worker is added or removed from the cluster.
 
+        This method will not be called if `watch_worker_status` is not set.
+
+        We recommend not relying on this method since worker status subscription is used
+        by Dask cluster implementations to manage worker state.
+
         Args:
             - op (str): Either "add" or "remove"
             - message (dict): Information about the event that the scheduler has sent
@@ -325,9 +341,12 @@ def _pre_start_yield(self) -> None:
             )
             self._should_run_event.set()
 
-        self._watch_dask_events_task = asyncio.run_coroutine_threadsafe(
-            self._watch_dask_events(), self.client.loop.asyncio_loop  # type: ignore
-        )
+        if self.watch_worker_status is True or (
+            self.watch_worker_status is None and not self.adapt_kwargs
+        ):
+            self._watch_dask_events_task = asyncio.run_coroutine_threadsafe(
+                self._watch_dask_events(), self.client.loop.asyncio_loop  # type: ignore
+            )
 
     def _post_start_yield(self) -> None:
         from distributed import wait


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
@@ -478,6 +478,9 @@ def test_executor_logs_worker_events(self, caplog):
         ) as client:
             executor = DaskExecutor(address=client.scheduler.address)
             with executor.start():
+                assert executor.watch_worker_status is None
+                assert executor._watch_dask_events_task is not None
+
                 time.sleep(0.1)
                 client.cluster.scale(4)
                 while len(client.scheduler_info()["workers"]) < 4:
@@ -489,6 +492,42 @@ def test_executor_logs_worker_events(self, caplog):
         assert any(re.match("Worker .+ added", rec.msg) for rec in caplog.records)
         assert any(re.match("Worker .+ removed", rec.msg) for rec in caplog.records)
 
+    def test_executor_disables_watch_worker_events_with_false(self):
+        with distributed.Client(
+            n_workers=1, processes=False, set_as_default=False
+        ) as client:
+            executor = DaskExecutor(
+                address=client.scheduler.address, watch_worker_status=False
+            )
+            with executor.start():
+                assert executor.watch_worker_status is False
+                assert executor._watch_dask_events_task is None
+
+    def test_executor_disables_watch_worker_events_with_adapt(self):
+        with distributed.Client(
+            n_workers=1, processes=False, set_as_default=False
+        ) as client:
+            executor = DaskExecutor(
+                address=client.scheduler.address,
+                adapt_kwargs={"maximum": 4},
+            )
+            with executor.start():
+                assert executor.watch_worker_status is None
+                assert executor._watch_dask_events_task is None
+
+    def test_executor_enables_watch_worker_events_with_true(self):
+        with distributed.Client(
+            n_workers=1, processes=False, set_as_default=False
+        ) as client:
+            executor = DaskExecutor(
+                address=client.scheduler.address,
+                watch_worker_status=True,
+                adapt_kwargs={"maximum": 4},
+            )
+            with executor.start():
+                assert executor.watch_worker_status is True
+                assert executor._watch_dask_events_task is not None
+
     @pytest.mark.parametrize("kind", ["external", "inproc"])
     @pytest.mark.flaky
     def test_exit_early_with_external_or_inproc_cluster_waits_for_pending_futures(