dask-cuda 25.8.0__py3-none-any.whl → 25.10.0__py3-none-any.whl

dask_cuda/tests/test_initialize.py

@@ -2,8 +2,14 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import multiprocessing as mp
+import os
+import shutil
+import subprocess
 import sys
+import tempfile
+import textwrap
 
+import cuda.core.experimental
 import numpy
 import psutil
 import pytest
@@ -14,7 +20,7 @@ from distributed.deploy.local import LocalCluster
 
 from dask_cuda.initialize import initialize
 from dask_cuda.utils import get_ucx_config
-from dask_cuda.utils_test import IncreasedCloseTimeoutNanny, get_ucx_implementation
+from dask_cuda.utils_test import IncreasedCloseTimeoutNanny
 
 mp = mp.get_context("spawn")  # type: ignore
 
@@ -24,19 +30,19 @@ mp = mp.get_context("spawn")  # type: ignore
 # of UCX before retrieving the current config.
 
 
-def _test_initialize_ucx_tcp(protocol):
-    ucp = get_ucx_implementation(protocol)
+def _test_initialize_ucx_tcp():
+    ucxx = pytest.importorskip("ucxx")
 
     kwargs = {"enable_tcp_over_ucx": True}
-    initialize(protocol=protocol, **kwargs)
+    initialize(**kwargs)
     with LocalCluster(
-        protocol=protocol,
+        protocol="ucx",
         dashboard_address=None,
         n_workers=1,
         threads_per_worker=1,
         processes=True,
         worker_class=IncreasedCloseTimeoutNanny,
-        config={"distributed.comm.ucx": get_ucx_config(**kwargs)},
+        config={"distributed-ucxx": get_ucx_config(**kwargs)},
     ) as cluster:
         with Client(cluster) as client:
             res = da.from_array(numpy.arange(10000), chunks=(1000,))
@@ -44,7 +50,7 @@ def _test_initialize_ucx_tcp(protocol):
             assert res == 49995000
 
             def check_ucx_options():
-                conf = ucp.get_config()
+                conf = ucxx.get_config()
                 assert "TLS" in conf
                 assert "tcp" in conf["TLS"]
                 assert "cuda_copy" in conf["TLS"]
@@ -55,29 +61,28 @@ def _test_initialize_ucx_tcp(protocol):
             assert all(client.run(check_ucx_options).values())
 
 
-@pytest.mark.parametrize("protocol", ["ucx", "ucx-old"])
-def test_initialize_ucx_tcp(protocol):
-    get_ucx_implementation(protocol)
+def test_initialize_ucx_tcp():
+    pytest.importorskip("distributed_ucxx")
 
-    p = mp.Process(target=_test_initialize_ucx_tcp, args=(protocol,))
+    p = mp.Process(target=_test_initialize_ucx_tcp)
     p.start()
     p.join()
     assert not p.exitcode
 
 
-def _test_initialize_ucx_nvlink(protocol):
-    ucp = get_ucx_implementation(protocol)
+def _test_initialize_ucx_nvlink():
+    ucxx = pytest.importorskip("ucxx")
 
     kwargs = {"enable_nvlink": True}
-    initialize(protocol=protocol, **kwargs)
+    initialize(**kwargs)
     with LocalCluster(
-        protocol=protocol,
+        protocol="ucx",
         dashboard_address=None,
         n_workers=1,
         threads_per_worker=1,
         processes=True,
         worker_class=IncreasedCloseTimeoutNanny,
-        config={"distributed.comm.ucx": get_ucx_config(**kwargs)},
+        config={"distributed-ucxx": get_ucx_config(**kwargs)},
     ) as cluster:
         with Client(cluster) as client:
             res = da.from_array(numpy.arange(10000), chunks=(1000,))
@@ -85,7 +90,7 @@ def _test_initialize_ucx_nvlink(protocol):
             assert res == 49995000
 
             def check_ucx_options():
-                conf = ucp.get_config()
+                conf = ucxx.get_config()
                 assert "TLS" in conf
                 assert "cuda_ipc" in conf["TLS"]
                 assert "tcp" in conf["TLS"]
@@ -97,29 +102,28 @@ def _test_initialize_ucx_nvlink(protocol):
             assert all(client.run(check_ucx_options).values())
 
 
-@pytest.mark.parametrize("protocol", ["ucx", "ucx-old"])
-def test_initialize_ucx_nvlink(protocol):
-    get_ucx_implementation(protocol)
+def test_initialize_ucx_nvlink():
+    pytest.importorskip("distributed_ucxx")
 
-    p = mp.Process(target=_test_initialize_ucx_nvlink, args=(protocol,))
+    p = mp.Process(target=_test_initialize_ucx_nvlink)
     p.start()
     p.join()
     assert not p.exitcode
 
 
-def _test_initialize_ucx_infiniband(protocol):
-    ucp = get_ucx_implementation(protocol)
+def _test_initialize_ucx_infiniband():
+    ucxx = pytest.importorskip("ucxx")
 
     kwargs = {"enable_infiniband": True}
-    initialize(protocol=protocol, **kwargs)
+    initialize(**kwargs)
     with LocalCluster(
-        protocol=protocol,
+        protocol="ucx",
         dashboard_address=None,
         n_workers=1,
         threads_per_worker=1,
         processes=True,
         worker_class=IncreasedCloseTimeoutNanny,
-        config={"distributed.comm.ucx": get_ucx_config(**kwargs)},
+        config={"distributed-ucxx": get_ucx_config(**kwargs)},
     ) as cluster:
         with Client(cluster) as client:
             res = da.from_array(numpy.arange(10000), chunks=(1000,))
@@ -127,7 +131,7 @@ def _test_initialize_ucx_infiniband(protocol):
             assert res == 49995000
 
             def check_ucx_options():
-                conf = ucp.get_config()
+                conf = ucxx.get_config()
                 assert "TLS" in conf
                 assert "rc" in conf["TLS"]
                 assert "tcp" in conf["TLS"]
@@ -142,28 +146,27 @@ def _test_initialize_ucx_infiniband(protocol):
 @pytest.mark.skipif(
     "ib0" not in psutil.net_if_addrs(), reason="Infiniband interface ib0 not found"
 )
-@pytest.mark.parametrize("protocol", ["ucx", "ucx-old"])
-def test_initialize_ucx_infiniband(protocol):
-    get_ucx_implementation(protocol)
+def test_initialize_ucx_infiniband():
+    pytest.importorskip("distributed_ucxx")
 
-    p = mp.Process(target=_test_initialize_ucx_infiniband, args=(protocol,))
+    p = mp.Process(target=_test_initialize_ucx_infiniband)
     p.start()
     p.join()
     assert not p.exitcode
 
 
-def _test_initialize_ucx_all(protocol):
-    ucp = get_ucx_implementation(protocol)
+def _test_initialize_ucx_all():
+    ucxx = pytest.importorskip("ucxx")
 
-    initialize(protocol=protocol)
+    initialize()
     with LocalCluster(
-        protocol=protocol,
+        protocol="ucx",
         dashboard_address=None,
         n_workers=1,
         threads_per_worker=1,
         processes=True,
         worker_class=IncreasedCloseTimeoutNanny,
-        config={"distributed.comm.ucx": get_ucx_config()},
+        config={"distributed-ucxx": get_ucx_config()},
     ) as cluster:
         with Client(cluster) as client:
             res = da.from_array(numpy.arange(10000), chunks=(1000,))
@@ -171,7 +174,7 @@ def _test_initialize_ucx_all(protocol):
             assert res == 49995000
 
             def check_ucx_options():
-                conf = ucp.get_config()
+                conf = ucxx.get_config()
                 assert "TLS" in conf
                 assert conf["TLS"] == "all"
                 assert all(
@@ -186,11 +189,10 @@ def _test_initialize_ucx_all(protocol):
             assert all(client.run(check_ucx_options).values())
 
 
-@pytest.mark.parametrize("protocol", ["ucx", "ucx-old"])
-def test_initialize_ucx_all(protocol):
-    get_ucx_implementation(protocol)
+def test_initialize_ucx_all():
+    pytest.importorskip("distributed_ucxx")
 
-    p = mp.Process(target=_test_initialize_ucx_all, args=(protocol,))
+    p = mp.Process(target=_test_initialize_ucx_all)
     p.start()
     p.join()
     assert not p.exitcode
@@ -229,3 +231,190 @@ def test_dask_cuda_import():
     p.start()
     p.join()
     assert not p.exitcode
+
+
+def _test_cuda_context_warning_with_subprocess_warnings(protocol):
+    """Test CUDA context warnings from both parent and worker subprocesses.
+
+    This test creates a standalone script that imports a problematic library
+    and creates LocalCUDACluster with processes=True. This should generate
+    warnings from both the parent process and each worker subprocess, since
+    they all inherit the CUDA context created at import time.
+    """
+    # Create temporary directory for our test files
+    temp_dir = tempfile.mkdtemp()
+
+    # Create the problematic library that creates CUDA context at import
+    problematic_library_code = textwrap.dedent(
+        """
+        # Problematic library that creates CUDA context at import time
+        import os
+
+        import cuda.core.experimental
+
+        try:
+            # Create CUDA context at import time, this will be inherited by subprocesses
+            cuda.core.experimental.Device().set_current()
+            print("Problematic library: Created CUDA context at import time")
+            os.environ['SUBPROCESS_CUDA_CONTEXT_CREATED'] = '1'
+        except Exception as e:
+            raise RuntimeError(
+                f"Problematic library: Failed to create CUDA context({e})"
+            )
+            os.environ['SUBPROCESS_CUDA_CONTEXT_CREATED'] = '0'
+        """
+    )
+
+    problematic_lib_path = os.path.join(temp_dir, "problematic_cuda_library.py")
+    with open(problematic_lib_path, "w") as f:
+        f.write(problematic_library_code)
+
+    # Create the main test script that imports the problematic library
+    # and creates LocalCUDACluster - this will run in a subprocess
+    main_script_code = textwrap.dedent(
+        f"""
+        # Main script that demonstrates the real-world problematic scenario
+        import os
+        import sys
+        import logging
+
+        # Add the temp directory to path so we can import our problematic library
+        sys.path.insert(0, '{temp_dir}')
+
+        print("=== Starting subprocess warnings test ===")
+
+        # This is the key part: import the problematic library BEFORE creating
+        # LocalCUDACluster. This creates a CUDA context that will be inherited
+        # by all worker subprocesses
+        print("Importing problematic library...")
+        import problematic_cuda_library
+
+        context_mode = os.environ.get('SUBPROCESS_CUDA_CONTEXT_CREATED', None)
+        if context_mode == "1":
+            print(f"Context creation successful")
+        else:
+            raise RuntimeError("Context creation failed")
+
+        if __name__ == "__main__":
+            try:
+                from dask_cuda import LocalCUDACluster
+                from dask_cuda.utils_test import IncreasedCloseTimeoutNanny
+
+                cluster = LocalCUDACluster(
+                    dashboard_address=None,
+                    worker_class=IncreasedCloseTimeoutNanny,
+                    protocol=f"{protocol}",
+                )
+                print("LocalCUDACluster created successfully!")
+
+                cluster.close()
+                print("Cluster closed successfully")
+
+            except Exception as e:
+                raise RuntimeError(f"Cluster setup error: {{e}}")
+
+        print("=== Subprocess warnings test completed ===")
+    """
+    )
+
+    main_script_path = os.path.join(temp_dir, "test_subprocess_warnings.py")
+    with open(main_script_path, "w") as f:
+        f.write(main_script_code)
+
+    try:
+        # Run the main script in a subprocess
+        result = subprocess.run(
+            [sys.executable, main_script_path],
+            capture_output=True,
+            text=True,
+            timeout=30,  # Reduced timeout for simpler test
+            cwd=os.getcwd(),
+        )
+
+        # Check for successful test execution regardless of warnings
+        assert (
+            "Context creation successful" in result.stdout
+        ), "Test did not create a CUDA context"
+        assert (
+            "Creating LocalCUDACluster" in result.stdout
+            or "LocalCUDACluster created successfully" in result.stdout
+        ), "LocalCUDACluster was not created"
+
+        # Check the log file for warnings from multiple processes
+        warnings_found = []
+        warnings_assigned_device_found = []
+
+        # Look for CUDA context warnings from different processes
+        lines = result.stderr.split("\n")
+        for line in lines:
+            if "A CUDA context for device" in line and "already exists" in line:
+                warnings_found.append(line)
+            if (
+                "should have a CUDA context assigned to device" in line
+                and "but instead the CUDA context is on device" in line
+            ):
+                warnings_assigned_device_found.append(line)
+
+        num_devices = cuda.core.experimental.system.num_devices
+
+        # Every worker raises the warning once. With protocol="ucx" the warning is
+        # raised once more by the parent process.
+        expected_warnings = num_devices if protocol == "tcp" else num_devices + 1
+        assert len(warnings_found) == expected_warnings, (
+            f"Expected {expected_warnings} CUDA context warnings, "
+            f"but found {len(warnings_assigned_device_found)}"
+        )
+
+        # Can only be tested in multi-GPU test environment, device 0 can never raise
+        # this warning (because it's where all CUDA contexts are created), thus one
+        # warning is raised by every device except 0.
+        expected_assigned_device_warnings = num_devices - 1
+        assert (
+            len(warnings_assigned_device_found) == expected_assigned_device_warnings
+        ), (
+            f"Expected {expected_assigned_device_warnings} warnings assigned to "
+            f"device, but found {len(warnings_assigned_device_found)}"
+        )
+
+        # Verify warnings contents
+        for warning in warnings_found:
+            assert (
+                "This is often the result of a CUDA-enabled library calling a "
+                "CUDA runtime function before Dask-CUDA" in warning
+            ), f"Warning missing explanatory text: {warning}"
+        for warning in warnings_assigned_device_found:
+            assert (
+                "This is often the result of a CUDA-enabled library calling a "
+                "CUDA runtime function before Dask-CUDA" in warning
+            ), f"Warning missing explanatory text: {warning}"
+    finally:
+        # Clean up temporary files
+        try:
+            if os.path.exists(temp_dir):
+                shutil.rmtree(temp_dir)
+        except Exception as e:
+            print(f"Cleanup error: {e}")
+
+
+@pytest.mark.parametrize("protocol", ["tcp", "ucx"])
+def test_cuda_context_warning_with_subprocess_warnings(protocol):
+    """Test CUDA context warnings from parent and worker subprocesses.
+
+    This test creates a standalone script that imports a problematic library at the top
+    level and then creates LocalCUDACluster with processes=True. This replicates the
+    exact real-world scenario where:
+
+    1. User imports a problematic library that creates CUDA context at import time
+    2. User creates LocalCUDACluster with multiple workers
+    3. Each worker subprocess inherits the CUDA context and emits warnings
+    4. Multiple warnings are generated (parent process + each worker subprocess)
+
+    This is the ultimate test as it demonstrates the distributed warning scenario
+    that users actually encounter in production.
+    """
+    p = mp.Process(
+        target=_test_cuda_context_warning_with_subprocess_warnings, args=(protocol,)
+    )
+    p.start()
+    p.join()
+    assert not p.exitcode

dask_cuda/tests/test_local_cuda_cluster.py

@@ -24,7 +24,7 @@ from dask_cuda.utils import (
     has_device_memory_resource,
     print_cluster_config,
 )
-from dask_cuda.utils_test import MockWorker, get_ucx_implementation
+from dask_cuda.utils_test import MockWorker
 
 
 @gen_test(timeout=20)
@@ -93,53 +93,39 @@ async def test_with_subset_of_cuda_visible_devices():
                 }
 
 
-@pytest.mark.parametrize(
-    "protocol",
-    ["ucx", "ucx-old"],
-)
 @gen_test(timeout=20)
-async def test_ucx_protocol(protocol):
-    get_ucx_implementation(protocol)
+async def test_ucx_protocol():
+    pytest.importorskip("distributed_ucxx")
 
     async with LocalCUDACluster(
-        protocol=protocol, asynchronous=True, data=dict
+        protocol="ucx", asynchronous=True, data=dict
     ) as cluster:
         assert all(
-            ws.address.startswith(f"{protocol}://")
-            for ws in cluster.scheduler.workers.values()
+            ws.address.startswith("ucx://") for ws in cluster.scheduler.workers.values()
         )
 
 
-@pytest.mark.parametrize(
-    "protocol",
-    ["ucx", "ucx-old"],
-)
 @gen_test(timeout=20)
-async def test_explicit_ucx_with_protocol_none(protocol):
-    get_ucx_implementation(protocol)
+async def test_explicit_ucx_with_protocol_none():
+    pytest.importorskip("distributed_ucxx")
 
-    initialize(protocol=protocol, enable_tcp_over_ucx=True)
+    initialize(enable_tcp_over_ucx=True)
     async with LocalCUDACluster(
         protocol=None,
         enable_tcp_over_ucx=True,
         asynchronous=True,
     ) as cluster:
         assert all(
-            ws.address.startswith(f"{protocol}://")
-            for ws in cluster.scheduler.workers.values()
+            ws.address.startswith("ucx://") for ws in cluster.scheduler.workers.values()
         )
 
 
 @pytest.mark.filterwarnings("ignore:Exception ignored in")
-@pytest.mark.parametrize(
-    "protocol",
-    ["ucx", "ucx-old"],
-)
 @gen_test(timeout=20)
-async def test_ucx_protocol_type_error(protocol):
-    get_ucx_implementation(protocol)
+async def test_ucx_protocol_type_error():
+    pytest.importorskip("distributed_ucxx")
 
-    initialize(protocol=protocol, enable_tcp_over_ucx=True)
+    initialize(enable_tcp_over_ucx=True)
     with pytest.raises(TypeError):
         async with LocalCUDACluster(
             protocol="tcp", enable_tcp_over_ucx=True, asynchronous=True, data=dict
@@ -602,10 +588,6 @@ async def test_cudf_spill_no_dedicated_memory():
         )
 
 
-@pytest.mark.parametrize(
-    "protocol",
-    ["ucx", "ucx-old"],
-)
 @pytest.mark.parametrize(
     "jit_unspill",
     [False, True],
@@ -614,8 +596,8 @@ async def test_cudf_spill_no_dedicated_memory():
     "device_memory_limit",
     [None, "1B"],
 )
-def test_print_cluster_config(capsys, protocol, jit_unspill, device_memory_limit):
-    get_ucx_implementation(protocol)
+def test_print_cluster_config(capsys, jit_unspill, device_memory_limit):
+    pytest.importorskip("distributed_ucxx")
 
     pytest.importorskip("rich")
 
@@ -640,46 +622,18 @@ def test_print_cluster_config(capsys, protocol, jit_unspill, device_memory_limit
             n_workers=1,
             device_memory_limit=device_memory_limit,
             jit_unspill=jit_unspill,
-            protocol=protocol,
+            protocol="ucx",
         ) as cluster:
             with Client(cluster) as client:
                 print_cluster_config(client)
                 captured = capsys.readouterr()
                 assert "Dask Cluster Configuration" in captured.out
-                assert protocol in captured.out
+                assert "ucx" in captured.out
                 if device_memory_limit == "1B":
                     assert "1 B" in captured.out
                 assert "[plugin]" in captured.out
                 client.shutdown()
 
-    def ucxpy_reset(timeout=20):
-        """Reset UCX-Py with a timeout.
-
-        Attempt to reset UCX-Py, not doing so may cause a deadlock because UCX-Py is
-        not thread-safe and the Dask cluster may still be alive while a new cluster
-        and UCX-Py instances are initalized.
-        """
-        import time
-
-        import ucp
-
-        start = time.monotonic()
-        while True:
-            try:
-                ucp.reset()
-            except ucp._libs.exceptions.UCXError as e:
-                if time.monotonic() - start > timeout:
-                    raise RuntimeError(
-                        f"Could not reset UCX-Py in {timeout} seconds, this may result "
-                        f"in a deadlock. Failure:\n{e}"
-                    )
-                continue
-            else:
-                break
-
-    if protocol == "ucx-old":
-        ucxpy_reset()
-
 
 @pytest.mark.xfail(reason="https://github.com/rapidsai/dask-cuda/issues/1265")
 def test_death_timeout_raises():

dask_cuda/tests/test_proxify_host_file.py

@@ -448,14 +448,20 @@ async def test_worker_force_spill_to_disk():
     """Test Dask triggering CPU-to-Disk spilling"""
     cudf = pytest.importorskip("cudf")
 
+    def create_dataframe():
+        return cudf.DataFrame({"key": np.arange(10**8)})
+
     with dask.config.set({"distributed.worker.memory.terminate": False}):
         async with dask_cuda.LocalCUDACluster(
             n_workers=1, device_memory_limit="1MB", jit_unspill=True, asynchronous=True
         ) as cluster:
             async with Client(cluster, asynchronous=True) as client:
                 # Create a df that are spilled to host memory immediately
-                df = cudf.DataFrame({"key": np.arange(10**8)})
-                [ddf] = client.persist([dask.dataframe.from_pandas(df, npartitions=1)])
+                ddf = dask.dataframe.from_delayed(
+                    dask.delayed(create_dataframe)(),
+                    meta=cudf.DataFrame({"key": cupy.arange(0)}),
+                )
+                [ddf] = client.persist([ddf])
                 await ddf
 
                 async def f(dask_worker):
@@ -466,13 +472,12 @@ async def test_worker_force_spill_to_disk():
                     memory = w.monitor.proc.memory_info().rss
                     w.memory_manager.memory_limit = memory - 10**8
                     w.memory_manager.memory_target_fraction = 1
-                    print(w.memory_manager.data)
                     await w.memory_manager.memory_monitor(w)
                     # Check that host memory are freed
                     assert w.monitor.proc.memory_info().rss < memory - 10**7
                     w.memory_manager.memory_limit = memory * 10  # Un-limit
 
-                client.run(f)
+                await 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

dask_cuda/tests/test_proxy.py

@@ -26,7 +26,7 @@ from dask_cuda import LocalCUDACluster, proxy_object
 from dask_cuda.disk_io import SpillToDiskFile
 from dask_cuda.proxify_device_objects import proxify_device_objects
 from dask_cuda.proxify_host_file import ProxifyHostFile
-from dask_cuda.utils_test import IncreasedCloseTimeoutNanny, get_ucx_implementation
+from dask_cuda.utils_test import IncreasedCloseTimeoutNanny
 
 # Make the "disk" serializer available and use a directory that are
 # remove on exit.
@@ -407,12 +407,12 @@ class _PxyObjTest(proxy_object.ProxyObject):
 
 
 @pytest.mark.parametrize("send_serializers", [None, ("dask", "pickle"), ("cuda",)])
-@pytest.mark.parametrize("protocol", ["tcp", "ucx", "ucx-old"])
+@pytest.mark.parametrize("protocol", ["tcp", "ucx"])
 @gen_test(timeout=120)
 async def test_communicating_proxy_objects(protocol, send_serializers):
     """Testing serialization of cuDF dataframe when communicating"""
-    if protocol.startswith("ucx"):
-        get_ucx_implementation(protocol)
+    if protocol == "ucx":
+        pytest.importorskip("distributed_ucxx")
     cudf = pytest.importorskip("cudf")
 
     def task(x):
@@ -421,7 +421,7 @@ async def test_communicating_proxy_objects(protocol, send_serializers):
         serializers_used = x._pxy_get().serializer
 
         # Check that `x` is serialized with the expected serializers
-        if protocol in ["ucx", "ucx-old"]:
+        if protocol == "ucx":
             if send_serializers is None:
                 assert serializers_used == "cuda"
             else:
@@ -452,13 +452,13 @@ async def test_communicating_proxy_objects(protocol, send_serializers):
             await client.submit(task, df)
 
 
-@pytest.mark.parametrize("protocol", ["tcp", "ucx", "ucx-old"])
+@pytest.mark.parametrize("protocol", ["tcp", "ucx"])
 @pytest.mark.parametrize("shared_fs", [True, False])
 @gen_test(timeout=20)
 async def test_communicating_disk_objects(protocol, shared_fs):
     """Testing disk serialization of cuDF dataframe when communicating"""
-    if protocol.startswith("ucx"):
-        get_ucx_implementation(protocol)
+    if protocol == "ucx":
+        pytest.importorskip("distributed_ucxx")
     cudf = pytest.importorskip("cudf")
     ProxifyHostFile._spill_to_disk.shared_filesystem = shared_fs
 

dask_cuda/tests/test_spill.py

@@ -15,10 +15,10 @@ from distributed.sizeof import sizeof
 from distributed.utils import Deadline
 from distributed.utils_test import gen_cluster, gen_test, loop  # noqa: F401
 
-import dask_cudf
+dask_cudf = pytest.importorskip("dask_cudf")
 
-from dask_cuda import LocalCUDACluster, utils
-from dask_cuda.utils_test import IncreasedCloseTimeoutNanny
+from dask_cuda import LocalCUDACluster, utils  # noqa: E402
+from dask_cuda.utils_test import IncreasedCloseTimeoutNanny  # noqa: E402
 
 if not utils.has_device_memory_resource():
     pytest.skip(