wehub-resource-sync
99 lines
2.8 KiB
Python
99 lines
2.8 KiB
Python
import sys
|
|
|
|
import dask
|
|
import dask.array as da
|
|
import dask.dataframe as dd
|
|
import numpy as np
|
|
import pandas as pd
|
|
import pytest
|
|
|
|
import ray
|
|
from ray.tests.conftest import * # noqa: F403, F401
|
|
from ray.util.client.common import ClientObjectRef
|
|
from ray.util.dask import disable_dask_on_ray, enable_dask_on_ray, ray_dask_get
|
|
from ray.util.dask.callbacks import ProgressBarCallback
|
|
|
|
|
|
@pytest.fixture
|
|
def ray_enable_dask_on_ray():
|
|
with enable_dask_on_ray():
|
|
yield
|
|
|
|
|
|
@pytest.mark.skipif(sys.platform == "win32", reason="Failing on Windows.")
|
|
def test_ray_dask_basic(ray_start_regular_shared):
|
|
@ray.remote
|
|
def stringify(x):
|
|
return "The answer is {}".format(x)
|
|
|
|
zero_id = ray.put(0)
|
|
|
|
def add(x, y):
|
|
# Can retrieve ray objects from inside Dask.
|
|
zero = ray.get(zero_id)
|
|
# Can call Ray methods from inside Dask.
|
|
return ray.get(stringify.remote(x + y + zero))
|
|
|
|
add = dask.delayed(add)
|
|
|
|
expected = "The answer is 6"
|
|
# Test with explicit scheduler argument.
|
|
assert add(2, 4).compute(scheduler=ray_dask_get) == expected
|
|
|
|
# Test with config setter.
|
|
enable_dask_on_ray()
|
|
assert add(2, 4).compute() == expected
|
|
disable_dask_on_ray()
|
|
|
|
# Test with config setter as context manager.
|
|
with enable_dask_on_ray():
|
|
assert add(2, 4).compute() == expected
|
|
|
|
# Test within Ray task.
|
|
|
|
@ray.remote
|
|
def call_add():
|
|
z = add(2, 4)
|
|
with ProgressBarCallback():
|
|
r = z.compute(scheduler=ray_dask_get)
|
|
return r
|
|
|
|
ans = ray.get(call_add.remote())
|
|
assert ans == "The answer is 6", ans
|
|
|
|
|
|
@pytest.mark.skipif(sys.platform == "win32", reason="Failing on Windows.")
|
|
def test_ray_dask_persist(ray_start_regular_shared):
|
|
arr = da.ones(5) + 2
|
|
result = arr.persist(scheduler=ray_dask_get)
|
|
assert isinstance(
|
|
next(iter(result.dask.values())), (ray.ObjectRef, ClientObjectRef)
|
|
)
|
|
|
|
|
|
def test_sort_with_progress_bar(ray_start_regular_shared):
|
|
npartitions = 10
|
|
df = dd.from_pandas(
|
|
pd.DataFrame(
|
|
np.random.randint(0, 100, size=(100, 2)), columns=["age", "grade"]
|
|
),
|
|
npartitions=npartitions,
|
|
)
|
|
# We set max_branch=npartitions in order to ensure that the task-based
|
|
# shuffle happens in a single stage, which is required in order for our
|
|
# optimization to work.
|
|
sorted_with_pb = None
|
|
sorted_without_pb = None
|
|
with ProgressBarCallback():
|
|
sorted_with_pb = df.set_index(
|
|
["age"], shuffle_method="tasks", max_branch=npartitions
|
|
).compute(scheduler=ray_dask_get, _ray_enable_progress_bar=True)
|
|
sorted_without_pb = df.set_index(
|
|
["age"], shuffle_method="tasks", max_branch=npartitions
|
|
).compute(scheduler=ray_dask_get)
|
|
assert sorted_with_pb.equals(sorted_without_pb)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main(["-v", __file__]))
|