Repository: PrefectHQ/prefect
Issue category: Bug Report

Issue:
DaskTaskRunner can not work together with `wait_for`
### Bug summary

When using the `DaskTaskRunner` using wait_for is broken. Using the `ThreadPoolTaskRunner` the same flow works just fine.

```py
from prefect import flow, task
from prefect.task_runners import ThreadPoolTaskRunner
from prefect_dask.task_runners import DaskTaskRunner

@task()
def task_a() -> bool:
    print("Task A")
    return True

@task()
def task_b() -> bool:
    
    print("Task B")
    return True


@flow(name='test_flow', log_prints=True, task_runner=DaskTaskRunner()) #ThreadPoolTaskRunner()
def main_flow():
    task_a_future = task_a.submit()
    task_b_future = task_b.submit(wait_for=[task_a_future])

    task_b_future.wait()


if __name__ == '__main__':
    main_flow()
```

generates the errors:


```
Exception ignored in: <function PrefectDaskFuture.__del__ at 0x00000199CFD7B4C0>
Traceback (most recent call last):
  File "C:\mambaforge\envs\prefect\Lib\site-packages\prefect_dask\task_runners.py", line 153, in __del__
    if self._final_state or self._wrapped_future.done():
                            ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\mambaforge\envs\prefect\Lib\site-packages\distributed\client.py", line 378, in done
    return self._state.done()
           ^^^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'done'
Exception ignored in: <function PrefectDaskFuture.__del__ at 0x00000199CFD7B4C0>
Traceback (most recent call last):
  File "C:\mambaforge\envs\prefect\Lib\site-packages\prefect_dask\task_runners.py", line 153, in __del__
    if self._final_state or self._wrapped_future.done():
                            ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\mambaforge\envs\prefect\Lib\site-packages\distributed\client.py", line 378, in done
    return self._state.done()
           ^^^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'done'
Exception ignored in: <function PrefectDaskFuture.__del__ at 0x00000199CFD7B4C0>
Traceback (most recent call last):
  File "C:\mambaforge\envs\prefect\Lib\site-packages\prefect_dask\task_runners.py", line 153, in __del__
    if self._final_state or self._wrapped_future.done():
                            ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\mambaforge\envs\prefect\Lib\site-packages\distributed\client.py", line 378, in done
    return self._state.done()
           ^^^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'done'
```

Seems like wrapping the normal future is broken when wait_for is used. If we use the return value of task_a in task_b and skip the wait_for the behaviour under both TaskRunners is correct.

### Version info

```Text
Version:             3.0.10
API version:         0.8.4
Python version:      3.12.7
Git commit:          3aa2d893
Built:               Tue, Oct 15, 2024 1:31 PM
OS/Arch:             win32/AMD64
Profile:             ephemeral
Server type:         server
Pydantic version:    2.9.2
Integrations:
  prefect-dask:      0.3.1
```


### Additional context

_No response_


Hints:
hi @EraYaN - thank you for the issue! will take a look at this

Patch:
diff --git a/src/integrations/prefect-dask/prefect_dask/task_runners.py b/src/integrations/prefect-dask/prefect_dask/task_runners.py
index f4edbab25295..60066faaf870 100644
--- a/src/integrations/prefect-dask/prefect_dask/task_runners.py
+++ b/src/integrations/prefect-dask/prefect_dask/task_runners.py
@@ -319,7 +319,7 @@ def submit(
         self,
         task: "Task[P, Coroutine[Any, Any, R]]",
         parameters: Dict[str, Any],
-        wait_for: Optional[Iterable[PrefectFuture]] = None,
+        wait_for: Optional[Iterable[PrefectDaskFuture[R]]] = None,
         dependencies: Optional[Dict[str, Set[TaskRunInput]]] = None,
     ) -> PrefectDaskFuture[R]:
         ...
@@ -329,26 +329,26 @@ def submit(
         self,
         task: "Task[Any, R]",
         parameters: Dict[str, Any],
-        wait_for: Optional[Iterable[PrefectFuture]] = None,
+        wait_for: Optional[Iterable[PrefectDaskFuture[R]]] = None,
         dependencies: Optional[Dict[str, Set[TaskRunInput]]] = None,
     ) -> PrefectDaskFuture[R]:
         ...
 
     def submit(
         self,
-        task: Task,
+        task: "Union[Task[P, R], Task[P, Coroutine[Any, Any, R]]]",
         parameters: Dict[str, Any],
-        wait_for: Optional[Iterable[PrefectFuture]] = None,
+        wait_for: Optional[Iterable[PrefectDaskFuture[R]]] = None,
         dependencies: Optional[Dict[str, Set[TaskRunInput]]] = None,
-    ) -> PrefectDaskFuture:
+    ) -> PrefectDaskFuture[R]:
         if not self._started:
             raise RuntimeError(
                 "The task runner must be started before submitting work."
             )
 
-        # unpack the upstream call in order to cast Prefect futures to Dask futures
-        # where possible to optimize Dask task scheduling
+        # Convert both parameters and wait_for futures to Dask futures
         parameters = self._optimize_futures(parameters)
+        wait_for = self._optimize_futures(wait_for) if wait_for else None
 
         future = self._client.submit(
             task,
@@ -357,7 +357,9 @@ def submit(
             dependencies=dependencies,
             return_type="state",
         )
-        return PrefectDaskFuture(wrapped_future=future, task_run_id=future.task_run_id)
+        return PrefectDaskFuture[R](
+            wrapped_future=future, task_run_id=future.task_run_id
+        )
 
     @overload
     def map(


Test patch:
diff --git a/src/integrations/prefect-dask/tests/test_task_runners.py b/src/integrations/prefect-dask/tests/test_task_runners.py
index 748de299c65e..d1bc5c8bec0b 100644
--- a/src/integrations/prefect-dask/tests/test_task_runners.py
+++ b/src/integrations/prefect-dask/tests/test_task_runners.py
@@ -452,3 +452,49 @@ def umbrella_flow():
             return future.result()
 
         assert umbrella_flow() == "nested task"
+
+    def test_state_dependencies_via_wait_for(self, task_runner):
+        @task
+        def task_a():
+            return time.time()
+
+        @task
+        def task_b():
+            return time.time()
+
+        @flow(task_runner=task_runner)
+        def test_flow() -> tuple[float, float]:
+            a = task_a.submit()
+            b = task_b.submit(wait_for=[a])
+            return a.result(), b.result()
+
+        a_time, b_time = test_flow()
+
+        assert b_time > a_time, "task_b timestamp should be after task_a timestamp"
+
+    def test_state_dependencies_via_wait_for_disparate_upstream_tasks(
+        self, task_runner
+    ):
+        @task
+        def task_a():
+            return time.time()
+
+        @task
+        def task_b():
+            return time.time()
+
+        @task
+        def task_c():
+            return time.time()
+
+        @flow(task_runner=task_runner)
+        def test_flow() -> tuple[float, float, float]:
+            a = task_a.submit()
+            b = task_b.submit()
+            c = task_c.submit(wait_for=[a, b])
+
+            return a.result(), b.result(), c.result()
+
+        a_time, b_time, c_time = test_flow()
+
+        assert c_time > a_time and c_time > b_time