hiro/ray
1
0
Fork 0
mirror of https://github.com/vale981/ray synced 2025-03-06 10:31:39 -05:00
Code Issues Projects Releases Packages Wiki Activity Actions

[Dask] Re-enable scheduler on dask_shuffle example (#16405)

Browse source
This commit is contained in:
Amog Kamsetty 2021-06-15 17:47:57 -07:00 • committed by GitHub
parent d23494d25a
commit ca22df2367
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 12 additions and 10 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

17
python/ray/util/dask/BUILD
View file

@ -101,14 +101,15 @@ py_test(
deps = [":dask_lib"],
)
py_test(
name = "dask_ray_shuffle_optimization_client_mode",
size = "medium",
main = "dask_ray_shuffle_optimization.py",
srcs = ["examples/dask_ray_shuffle_optimization.py"],
tags = ["exclusive", "client"],
deps = [":dask_lib"],
)
# This is currently failing.
#py_test(
# name = "dask_ray_shuffle_optimization_client_mode",
# size = "medium",
# main = "dask_ray_shuffle_optimization.py",
# srcs = ["examples/dask_ray_shuffle_optimization.py"],
# tags = ["exclusive", "client"],
# deps = [":dask_lib"],
#)
# This is a dummy test dependency that causes the above tests to be
# re-run if any of these files changes.

5
python/ray/util/dask/examples/dask_ray_shuffle_optimization.py
View file

@ -1,5 +1,5 @@
import ray
from ray.util.dask import dataframe_optimize
from ray.util.dask import dataframe_optimize, ray_dask_get
import dask
import dask.dataframe as dd
import numpy as np
@ -12,7 +12,8 @@ ray.init()
# Set the Dask DataFrame optimizer to
# our custom optimization function, this time using the config setter as a
# context manager.
with dask.config.set(dataframe_optimize=dataframe_optimize):
with dask.config.set(
scheduler=ray_dask_get, dataframe_optimize=dataframe_optimize):
npartitions = 100
df = dd.from_pandas(
pd.DataFrame(