Skip to content

Commit 84db665

Browse files
jaychiaJay Chia
andauthored
[BUG] Fix ray wait in RayPartitionSet (Eventual-Inc#3251)
Closes: Eventual-Inc#3249 Ray's `ray.wait` is supposed to: 1. Defaults to `fetch_local=True` which will supposedly fetch data to wherever the wait is called before returning 2. Defaults to `num_returns=1` which will wait until only the first item is ready before returning This seems to not be the intended behavior here, where `RayPartitionSet` is trying to wait on ALL the partitions to be ready, and does not want to pull any data down to the calling site. Co-authored-by: Jay Chia <[email protected]@users.noreply.github.com>
1 parent cf25ad4 commit 84db665

File tree

1 file changed

+1
-1
lines changed

1 file changed

+1
-1
lines changed

daft/runners/ray_runner.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -341,7 +341,7 @@ def num_partitions(self) -> int:
341341

342342
def wait(self) -> None:
343343
deduped_object_refs = {r.partition() for r in self._results.values()}
344-
ray.wait(list(deduped_object_refs))
344+
ray.wait(list(deduped_object_refs), fetch_local=False, num_returns=len(deduped_object_refs))
345345

346346

347347
def _from_arrow_type_with_ray_data_extensions(arrow_type: pa.lib.DataType) -> DataType:

0 commit comments

Comments
 (0)