Skip to content

Commit

Permalink
Use client.run for async function
Browse files Browse the repository at this point in the history
Using `client.submit` seems not allowed to resolve `get_worker`, but
`client.run` does.
  • Loading branch information
pentschev committed Mar 22, 2023
1 parent 8d5241f commit 884a595
Showing 1 changed file with 3 additions and 4 deletions.
7 changes: 3 additions & 4 deletions dask_cuda/tests/test_proxify_host_file.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,6 @@
from dask.utils import format_bytes
from distributed import Client
from distributed.utils_test import gen_test
from distributed.worker import get_worker

import dask_cuda
import dask_cuda.proxify_device_objects
Expand Down Expand Up @@ -429,9 +428,9 @@ async def test_worker_force_spill_to_disk():
ddf = dask.dataframe.from_pandas(df, npartitions=1).persist()
await ddf

async def f():
async def f(dask_worker):
"""Trigger a memory_monitor() and reset memory_limit"""
w = get_worker()
w = dask_worker
# Set a host memory limit that triggers spilling to disk
w.memory_manager.memory_pause_fraction = False
memory = w.monitor.proc.memory_info().rss
Expand All @@ -443,7 +442,7 @@ async def f():
assert w.monitor.proc.memory_info().rss < memory - 10**7
w.memory_manager.memory_limit = memory * 10 # Un-limit

await client.submit(f)
client.run(f)
log = str(await client.get_worker_logs())
# Check that the worker doesn't complain about unmanaged memory
assert "Unmanaged memory use is high" not in log
Expand Down

0 comments on commit 884a595

Please sign in to comment.