Skip to content

Commit

Permalink
Update test_spill.py to avoid FutureWarnings (#1243)
Browse files Browse the repository at this point in the history
Distributed has now deprecated parameters related to worker's fractional spilling in favor of Dask configs, update spilling tests to use Dask configs and avoid `FutureWarning`s.

Authors:
  - Peter Andreas Entschev (https://github.com/pentschev)

Approvers:
  - Mads R. B. Kristensen (https://github.com/madsbk)

URL: #1243
  • Loading branch information
pentschev authored Sep 28, 2023
1 parent 7400f95 commit 93b6677
Showing 1 changed file with 23 additions and 14 deletions.
37 changes: 23 additions & 14 deletions dask_cuda/tests/test_spill.py
Original file line number Diff line number Diff line change
Expand Up @@ -103,11 +103,12 @@ def delayed_worker_assert(
},
{
# This test setup differs from the one above as Distributed worker
# pausing is enabled and thus triggers `DeviceHostFile.evict()`
# spilling fraction is very low and thus forcefully triggers
# `DeviceHostFile.evict()`
"device_memory_limit": int(200e6),
"memory_limit": int(200e6),
"host_target": None,
"host_spill": None,
"host_target": False,
"host_spill": 0.01,
"host_pause": False,
"spills_to_disk": True,
},
Expand All @@ -124,7 +125,14 @@ def delayed_worker_assert(
@gen_test(timeout=120)
async def test_cupy_cluster_device_spill(params):
cupy = pytest.importorskip("cupy")
with dask.config.set({"distributed.worker.memory.terminate": False}):
with dask.config.set(
{
"distributed.worker.memory.terminate": False,
"distributed.worker.memory.pause": params["host_pause"],
"distributed.worker.memory.spill": params["host_spill"],
"distributed.worker.memory.target": params["host_target"],
}
):
async with LocalCUDACluster(
n_workers=1,
scheduler_port=0,
Expand All @@ -133,9 +141,6 @@ async def test_cupy_cluster_device_spill(params):
asynchronous=True,
device_memory_limit=params["device_memory_limit"],
memory_limit=params["memory_limit"],
memory_target_fraction=params["host_target"],
memory_spill_fraction=params["host_spill"],
memory_pause_fraction=params["host_pause"],
) as cluster:
async with Client(cluster, asynchronous=True) as client:

Expand Down Expand Up @@ -194,11 +199,12 @@ async def test_cupy_cluster_device_spill(params):
},
{
# This test setup differs from the one above as Distributed worker
# pausing is enabled and thus triggers `DeviceHostFile.evict()`
# spilling fraction is very low and thus forcefully triggers
# `DeviceHostFile.evict()`
"device_memory_limit": int(50e6),
"memory_limit": int(50e6),
"host_target": None,
"host_spill": None,
"host_target": False,
"host_spill": 0.01,
"host_pause": False,
"spills_to_disk": True,
},
Expand All @@ -221,16 +227,19 @@ async def test_cudf_cluster_device_spill(params):
"distributed.comm.compression": False,
"distributed.worker.memory.terminate": False,
"distributed.worker.memory.spill-compression": False,
"distributed.worker.memory.pause": params["host_pause"],
"distributed.worker.memory.spill": params["host_spill"],
"distributed.worker.memory.target": params["host_target"],
}
):
async with LocalCUDACluster(
n_workers=1,
scheduler_port=0,
silence_logs=False,
dashboard_address=None,
asynchronous=True,
device_memory_limit=params["device_memory_limit"],
memory_limit=params["memory_limit"],
memory_target_fraction=params["host_target"],
memory_spill_fraction=params["host_spill"],
memory_pause_fraction=params["host_pause"],
asynchronous=True,
) as cluster:
async with Client(cluster, asynchronous=True) as client:

Expand Down

0 comments on commit 93b6677

Please sign in to comment.