def test_sort_with_progress_bar(ray_start_1_cpu):
from ray.util.dask import ray_dask_get
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="tasks",
max_branch=npartitions).compute(
scheduler=ray_dask_get,
_ray_enable_progress_bar=True)
sorted_without_pb = df.set_index(
["age"], shuffle="tasks",
max_branch=npartitions).compute(scheduler=ray_dask_get)
assert sorted_with_pb.equals(sorted_without_pb)
def call_add():
z = add(2, 4)
with ProgressBarCallback():
r = z.compute(scheduler=ray_dask_get)
return r
def call_add():
z = add(2, 4)
# Can call Dask graphs from inside Ray.
with ProgressBarCallback():
r = z.compute(scheduler=ray_dask_get)
return r