From 0979589c7cd8b6115e3294136b0461f4d75c00ee Mon Sep 17 00:00:00 2001 From: Clark Zinzow Date: Fri, 23 Oct 2020 17:39:22 -0600 Subject: [PATCH] [dask-on-ray] Convert tuple of object refs to list before ray.get() call. (#11582) --- python/ray/tests/test_dask_scheduler.py | 9 +++++++++ python/ray/util/dask/scheduler.py | 8 +++++--- 2 files changed, 14 insertions(+), 3 deletions(-) diff --git a/python/ray/tests/test_dask_scheduler.py b/python/ray/tests/test_dask_scheduler.py index 0321b77c9..28a98a76e 100644 --- a/python/ray/tests/test_dask_scheduler.py +++ b/python/ray/tests/test_dask_scheduler.py @@ -1,4 +1,6 @@ import dask +import numpy as np +import dask.array as da import pytest import ray @@ -30,6 +32,13 @@ def test_ray_dask_basic(ray_start_regular_shared): assert ans == "The answer is 6", ans +def test_ray_dask_persist(ray_start_regular_shared): + arr = da.ones(5) + 2 + result = arr.persist(scheduler=ray_dask_get) + np.testing.assert_array_equal(result.dask.values()[0], np.ones(5) + 2) + + if __name__ == "__main__": import sys + sys.exit(pytest.main(["-v", __file__])) diff --git a/python/ray/util/dask/scheduler.py b/python/ray/util/dask/scheduler.py index aeded7bca..0614d3564 100644 --- a/python/ray/util/dask/scheduler.py +++ b/python/ray/util/dask/scheduler.py @@ -346,9 +346,11 @@ def ray_get_unpack(object_refs): The input Python object with all contained Ray object references resolved with their concrete values. """ - if isinstance(object_refs, - (tuple, list)) and any(not isinstance(x, ray.ObjectRef) - for x in object_refs): + if isinstance(object_refs, tuple): + object_refs = list(object_refs) + + if isinstance(object_refs, list) and any(not isinstance(x, ray.ObjectRef) + for x in object_refs): # We flatten the object references before calling ray.get(), since Dask # loves to nest collections in nested tuples and Ray expects a flat # list of object references. We repack the results after ray.get()