diff --git a/dask_cuda/tests/test_dask_cuda_worker.py b/dask_cuda/tests/test_dask_cuda_worker.py index 449fdba7e..974ad1319 100644 --- a/dask_cuda/tests/test_dask_cuda_worker.py +++ b/dask_cuda/tests/test_dask_cuda_worker.py @@ -40,7 +40,7 @@ def test_cuda_visible_devices_and_memory_limit_and_nthreads(loop): # noqa: F811 str(nthreads), "--no-dashboard", "--worker-class", - "dask_cuda.utils.MockWorker", + "dask_cuda.utils_test.MockWorker", ] ): with Client("127.0.0.1:9359", loop=loop) as client: @@ -329,7 +329,7 @@ def test_cuda_mig_visible_devices_and_memory_limit_and_nthreads(loop): # noqa: str(nthreads), "--no-dashboard", "--worker-class", - "dask_cuda.utils.MockWorker", + "dask_cuda.utils_test.MockWorker", ] ): with Client("127.0.0.1:9359", loop=loop) as client: @@ -364,7 +364,7 @@ def test_cuda_visible_devices_uuid(loop): # noqa: F811 "127.0.0.1", "--no-dashboard", "--worker-class", - "dask_cuda.utils.MockWorker", + "dask_cuda.utils_test.MockWorker", ] ): with Client("127.0.0.1:9359", loop=loop) as client: diff --git a/dask_cuda/tests/test_local_cuda_cluster.py b/dask_cuda/tests/test_local_cuda_cluster.py index 845759dfd..5d7762579 100644 --- a/dask_cuda/tests/test_local_cuda_cluster.py +++ b/dask_cuda/tests/test_local_cuda_cluster.py @@ -13,13 +13,13 @@ from dask_cuda import CUDAWorker, LocalCUDACluster, utils from dask_cuda.initialize import initialize from dask_cuda.utils import ( - MockWorker, get_cluster_configuration, get_device_total_memory, get_gpu_count_mig, get_gpu_uuid_from_index, print_cluster_config, ) +from dask_cuda.utils_test import MockWorker @gen_test(timeout=20) diff --git a/dask_cuda/worker_spec.py b/dask_cuda/worker_spec.py index 6a61fa8f8..84ce51725 100644 --- a/dask_cuda/worker_spec.py +++ b/dask_cuda/worker_spec.py @@ -5,7 +5,8 @@ from .initialize import initialize from .local_cuda_cluster import cuda_visible_devices -from .utils import CPUAffinity, get_cpu_affinity, get_gpu_count +from .plugins import CPUAffinity +from .utils import get_cpu_affinity, get_gpu_count def worker_spec(