[BUG] Separate PartitionTask done from results

[jaychia]

This PR marks PartitionTasks as done only after they have been explicitly marked as done by the runner.

Previously, we used the existence of the .results on a PartitionTask to determine whether or not it is done. However, this is not quite correct in the case of the RayRunner, which will attach a result containing a Ray ObjectRef, which is a future. This future may not (and is likely not) be completed yet at the time of PartitionTask creation.

[codspeed-hq]

CodSpeed Performance Report

Merging #3155 will not alter performance

_{Comparing jay/fix-rayrunner-tasks-done (be1fbfd) with main (8817a08)}

Summary

✅ 17 untouched benchmarks

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 79.00%. Comparing base (b4332ec) to head (be1fbfd).
Report is 6 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3155      +/-   ##
==========================================
+ Coverage   78.87%   79.00%   +0.12%     
==========================================
  Files         624      634      +10     
  Lines       75925    76940    +1015     
==========================================
+ Hits        59888    60788     +900     
- Misses      16037    16152     +115     

Files with missing lines	Coverage Δ
daft/execution/execution_step.py	89.22% <100.00%> (+0.18%)	⬆️
daft/runners/pyrunner.py	87.62% <100.00%> (+0.08%)	⬆️
daft/runners/ray_runner.py	81.25% <100.00%> (+0.05%)	⬆️

... and 37 files with indirect coverage changes

[kevinzwang]

Changes look fine to me. Wondering if this will cause any performance regressions though?

[jaychia]

Changes look fine to me. Wondering if this will cause any performance regressions though?

I think that is unlikely, given that the ray wait is performed somewhere else (after this PR it is done in the _run_plan loop. Before this PR, it is done somewhere, perhaps in the physical plan, perhaps in the consumer of the data when showing the data etc).

[jaychia]

Also seeing now that the Ray trace no longer shows a long duration during next(tasks) iteration on the Physical Plan after this change.

Before:

After:

This is much more expected... I think in the past next(tasks) was taking so long because we called ray.get on some partition metadata in there. Because the task wasn't actually done, the physical plan was actually blocking on retrieving the metadata until the task completed.

Notice how the trace for tasks are also much nicer now... Much less time between when tasks are created compared when they actually run...