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

dask_cuda/tests/test_rdd_ucx.py

@@ -0,0 +1,160 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+
+
+import importlib
+import io
+import multiprocessing as mp
+import sys
+
+import pytest
+
+from dask_cuda import LocalCUDACluster
+
+mp = mp.get_context("spawn")  # type: ignore
+
+
+def _has_distributed_ucxx() -> bool:
+    return bool(importlib.util.find_spec("distributed_ucxx"))
+
+
+def _test_protocol_ucx():
+    with LocalCUDACluster(protocol="ucx") as cluster:
+        assert cluster.scheduler_comm.address.startswith("ucx://")
+
+        if _has_distributed_ucxx():
+            import distributed_ucxx
+
+            assert all(
+                isinstance(batched_send.comm, distributed_ucxx.ucxx.UCXX)
+                for batched_send in cluster.scheduler.stream_comms.values()
+            )
+        else:
+            import rapids_dask_dependency
+
+            assert all(
+                isinstance(
+                    batched_send.comm,
+                    rapids_dask_dependency.patches.distributed.comm.__rdd_patch_ucx.UCX,
+                )
+                for batched_send in cluster.scheduler.stream_comms.values()
+            )
+
+
+def _test_protocol_ucxx():
+    if _has_distributed_ucxx():
+        with LocalCUDACluster(protocol="ucxx") as cluster:
+            assert cluster.scheduler_comm.address.startswith("ucxx://")
+            import distributed_ucxx
+
+            assert all(
+                isinstance(batched_send.comm, distributed_ucxx.ucxx.UCXX)
+                for batched_send in cluster.scheduler.stream_comms.values()
+            )
+    else:
+        with pytest.raises(RuntimeError, match="Cluster failed to start"):
+            LocalCUDACluster(protocol="ucxx")
+
+
+def _test_protocol_ucx_old():
+    with LocalCUDACluster(protocol="ucx-old") as cluster:
+        assert cluster.scheduler_comm.address.startswith("ucx-old://")
+
+        import rapids_dask_dependency
+
+        assert all(
+            isinstance(
+                batched_send.comm,
+                rapids_dask_dependency.patches.distributed.comm.__rdd_patch_ucx.UCX,
+            )
+            for batched_send in cluster.scheduler.stream_comms.values()
+        )
+
+
+def _run_test_with_output_capture(test_func_name, conn):
+    """Run a test function in a subprocess and capture stdout/stderr."""
+    # Redirect stdout and stderr to capture output
+    old_stdout = sys.stdout
+    old_stderr = sys.stderr
+    captured_output = io.StringIO()
+    sys.stdout = sys.stderr = captured_output
+
+    try:
+        # Import and run the test function
+        if test_func_name == "_test_protocol_ucx":
+            _test_protocol_ucx()
+        elif test_func_name == "_test_protocol_ucxx":
+            _test_protocol_ucxx()
+        elif test_func_name == "_test_protocol_ucx_old":
+            _test_protocol_ucx_old()
+        else:
+            raise ValueError(f"Unknown test function: {test_func_name}")
+
+        output = captured_output.getvalue()
+        conn.send((True, output))  # True = success
+    except Exception as e:
+        output = captured_output.getvalue()
+        output += f"\nException: {e}"
+        import traceback
+
+        output += f"\nTraceback:\n{traceback.format_exc()}"
+        conn.send((False, output))  # False = failure
+    finally:
+        # Restore original stdout/stderr
+        sys.stdout = old_stdout
+        sys.stderr = old_stderr
+        conn.close()
+
+
+@pytest.mark.parametrize("protocol", ["ucx", "ucxx", "ucx-old"])
+def test_rdd_protocol(protocol):
+    """Test rapids-dask-dependency protocol selection"""
+    if protocol == "ucx":
+        test_func_name = "_test_protocol_ucx"
+    elif protocol == "ucxx":
+        test_func_name = "_test_protocol_ucxx"
+    else:
+        test_func_name = "_test_protocol_ucx_old"
+
+    # Create a pipe for communication between parent and child processes
+    parent_conn, child_conn = mp.Pipe()
+    p = mp.Process(
+        target=_run_test_with_output_capture, args=(test_func_name, child_conn)
+    )
+
+    p.start()
+    p.join(timeout=60)
+
+    if p.is_alive():
+        p.kill()
+        p.close()
+        raise TimeoutError("Test process timed out")
+
+    # Get the result from the child process
+    success, output = parent_conn.recv()
+
+    # Check that the test passed
+    assert success, f"Test failed in subprocess. Output:\n{output}"
+
+    # For the ucx protocol, check if warnings are printed when distributed_ucxx is not
+    # available
+    if protocol == "ucx" and not _has_distributed_ucxx():
+        # Check if the warning about protocol='ucx' is printed
+        print(f"Output for {protocol} protocol:\n{output}")
+        assert (
+            "you have requested protocol='ucx'" in output
+        ), f"Expected warning not found in output: {output}"
+        assert (
+            "'distributed-ucxx' is not installed" in output
+        ), f"Expected warning about distributed-ucxx not found in output: {output}"
+    elif protocol == "ucx" and _has_distributed_ucxx():
+        # When distributed_ucxx is available, the warning should NOT be printed
+        assert "you have requested protocol='ucx'" not in output, (
+            "Warning should not be printed when distributed_ucxx is available: "
+            f"{output}"
+        )
+    elif protocol == "ucx-old":
+        # The ucx-old protocol should not generate warnings
+        assert (
+            "you have requested protocol='ucx'" not in output
+        ), f"Warning should not be printed for ucx-old protocol: {output}"

dask_cuda/tests/test_spill.py

@@ -1,14 +1,18 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.
+
 import gc
 import os
 from time import sleep
+from typing import TypedDict
 
 import pytest
 
 import dask
 from dask import array as da
-from distributed import Client, wait
+from distributed import Client, Worker, wait
 from distributed.metrics import time
 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
@@ -16,6 +20,13 @@ import dask_cudf
 from dask_cuda import LocalCUDACluster, utils
 from dask_cuda.utils_test import IncreasedCloseTimeoutNanny
 
+if not utils.has_device_memory_resource():
+    pytest.skip(
+        "No spilling tests supported for devices without memory resources. "
+        "See https://github.com/rapidsai/dask-cuda/issues/1510",
+        allow_module_level=True,
+    )
+
 if utils.get_device_total_memory() < 1e10:
     pytest.skip("Not enough GPU memory", allow_module_level=True)
 
@@ -72,24 +83,66 @@ def cudf_spill(request):
 
 
 def device_host_file_size_matches(
-    dhf, total_bytes, device_chunk_overhead=0, serialized_chunk_overhead=1024
+    dask_worker: Worker,
+    total_bytes,
+    device_chunk_overhead=0,
+    serialized_chunk_overhead=1024,
 ):
-    byte_sum = dhf.device_buffer.fast.total_weight
+    worker_data_sizes = collect_device_host_file_size(
+        dask_worker,
+        device_chunk_overhead=device_chunk_overhead,
+        serialized_chunk_overhead=serialized_chunk_overhead,
+    )
+    byte_sum = (
+        worker_data_sizes["device_fast"]
+        + worker_data_sizes["host_fast"]
+        + worker_data_sizes["host_buffer"]
+        + worker_data_sizes["disk"]
+    )
+    return (
+        byte_sum >= total_bytes
+        and byte_sum
+        <= total_bytes
+        + worker_data_sizes["device_overhead"]
+        + worker_data_sizes["host_overhead"]
+        + worker_data_sizes["disk_overhead"]
+    )
+
+
+class WorkerDataSizes(TypedDict):
+    device_fast: int
+    host_fast: int
+    host_buffer: int
+    disk: int
+    device_overhead: int
+    host_overhead: int
+    disk_overhead: int
+
 
-    # `dhf.host_buffer.fast` is only available when Worker's `memory_limit != 0`
+def collect_device_host_file_size(
+    dask_worker: Worker,
+    device_chunk_overhead: int,
+    serialized_chunk_overhead: int,
+) -> WorkerDataSizes:
+    dhf = dask_worker.data
+
+    device_fast = dhf.device_buffer.fast.total_weight or 0
     if hasattr(dhf.host_buffer, "fast"):
-        byte_sum += dhf.host_buffer.fast.total_weight
+        host_fast = dhf.host_buffer.fast.total_weight or 0
+        host_buffer = 0
     else:
-        byte_sum += sum([sizeof(b) for b in dhf.host_buffer.values()])
+        host_buffer = sum([sizeof(b) for b in dhf.host_buffer.values()])
+        host_fast = 0
 
-    # `dhf.disk` is only available when Worker's `memory_limit != 0`
     if dhf.disk is not None:
         file_path = [
             os.path.join(dhf.disk.directory, fname)
             for fname in dhf.disk.filenames.values()
         ]
         file_size = [os.path.getsize(f) for f in file_path]
-        byte_sum += sum(file_size)
+        disk = sum(file_size)
+    else:
+        disk = 0
 
     # Allow up to chunk_overhead bytes overhead per chunk
     device_overhead = len(dhf.device) * device_chunk_overhead
@@ -98,17 +151,25 @@ def device_host_file_size_matches(
         len(dhf.disk) * serialized_chunk_overhead if dhf.disk is not None else 0
     )
 
-    return (
-        byte_sum >= total_bytes
-        and byte_sum <= total_bytes + device_overhead + host_overhead + disk_overhead
+    return WorkerDataSizes(
+        device_fast=device_fast,
+        host_fast=host_fast,
+        host_buffer=host_buffer,
+        disk=disk,
+        device_overhead=device_overhead,
+        host_overhead=host_overhead,
+        disk_overhead=disk_overhead,
     )
 
 
 def assert_device_host_file_size(
-    dhf, total_bytes, device_chunk_overhead=0, serialized_chunk_overhead=1024
+    dask_worker: Worker,
+    total_bytes,
+    device_chunk_overhead=0,
+    serialized_chunk_overhead=1024,
 ):
     assert device_host_file_size_matches(
-        dhf, total_bytes, device_chunk_overhead, serialized_chunk_overhead
+        dask_worker, total_bytes, device_chunk_overhead, serialized_chunk_overhead
     )
 
 
@@ -119,7 +180,7 @@ def worker_assert(
     dask_worker=None,
 ):
     assert_device_host_file_size(
-        dask_worker.data, total_size, device_chunk_overhead, serialized_chunk_overhead
+        dask_worker, total_size, device_chunk_overhead, serialized_chunk_overhead
     )
 
 
@@ -131,12 +192,12 @@ def delayed_worker_assert(
 ):
     start = time()
     while not device_host_file_size_matches(
-        dask_worker.data, total_size, device_chunk_overhead, serialized_chunk_overhead
+        dask_worker, total_size, device_chunk_overhead, serialized_chunk_overhead
     ):
         sleep(0.01)
         if time() < start + 3:
             assert_device_host_file_size(
-                dask_worker.data,
+                dask_worker,
                 total_size,
                 device_chunk_overhead,
                 serialized_chunk_overhead,
@@ -224,8 +285,8 @@ async def test_cupy_cluster_device_spill(params):
                 x = rs.random(int(50e6), chunks=2e6)
                 await wait(x)
 
-                xx = x.persist()
-                await wait(xx)
+                [xx] = client.persist([x])
+                await xx
 
                 # Allow up to 1024 bytes overhead per chunk serialized
                 await client.run(
@@ -344,19 +405,38 @@ async def test_cudf_cluster_device_spill(params, cudf_spill):
                 sizes = sizes.to_arrow().to_pylist()
                 nbytes = sum(sizes)
 
-                cdf2 = cdf.persist()
-                await wait(cdf2)
+                [cdf2] = client.persist([cdf])
+                await cdf2
 
                 del cdf
                 gc.collect()
 
                 if enable_cudf_spill:
-                    await client.run(
-                        worker_assert,
-                        0,
-                        0,
-                        0,
+                    expected_data = WorkerDataSizes(
+                        device_fast=0,
+                        host_fast=0,
+                        host_buffer=0,
+                        disk=0,
+                        device_overhead=0,
+                        host_overhead=0,
+                        disk_overhead=0,
                     )
+
+                    deadline = Deadline.after(duration=3)
+                    while not deadline.expired:
+                        data = await client.run(
+                            collect_device_host_file_size,
+                            device_chunk_overhead=0,
+                            serialized_chunk_overhead=0,
+                        )
+                        expected = {k: expected_data for k in data}
+                        if data == expected:
+                            break
+                        sleep(0.01)
+
+                    # final assertion for pytest to reraise with a nice traceback
+                    assert data == expected
+
                 else:
                     await client.run(
                         assert_host_chunks,
@@ -419,8 +499,8 @@ async def test_cudf_spill_cluster(cudf_spill):
                 }
             )
 
-            ddf = dask_cudf.from_cudf(cdf, npartitions=2).sum().persist()
-            await wait(ddf)
+            [ddf] = client.persist([dask_cudf.from_cudf(cdf, npartitions=2).sum()])
+            await ddf
 
             await client.run(_assert_cudf_spill_stats, enable_cudf_spill)
             _assert_cudf_spill_stats(enable_cudf_spill)

dask_cuda/tests/test_utils.py

@@ -1,3 +1,6 @@
+# SPDX-FileCopyrightText: Copyright (c) 2019-2025, NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+
 import os
 from unittest.mock import patch
 
@@ -15,11 +18,13 @@ from dask_cuda.utils import (
     get_n_gpus,
     get_preload_options,
     get_ucx_config,
+    has_device_memory_resource,
     nvml_device_index,
     parse_cuda_visible_device,
     parse_device_memory_limit,
     unpack_bitmask,
 )
+from dask_cuda.utils_test import get_ucx_implementation
 
 
 @patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "0,1,2"})
@@ -76,19 +81,19 @@ def test_get_device_total_memory():
     for i in range(get_n_gpus()):
         with cuda.gpus[i]:
             total_mem = get_device_total_memory(i)
-            assert type(total_mem) is int
-            assert total_mem > 0
+            if has_device_memory_resource():
+                assert type(total_mem) is int
+                assert total_mem > 0
+            else:
+                assert total_mem is None
 
 
 @pytest.mark.parametrize(
     "protocol",
-    ["ucx", "ucxx"],
+    ["ucx", "ucx-old"],
 )
 def test_get_preload_options_default(protocol):
-    if protocol == "ucx":
-        pytest.importorskip("ucp")
-    elif protocol == "ucxx":
-        pytest.importorskip("ucxx")
+    get_ucx_implementation(protocol)
 
     opts = get_preload_options(
         protocol=protocol,
@@ -103,16 +108,13 @@ def test_get_preload_options_default(protocol):
 
 @pytest.mark.parametrize(
     "protocol",
-    ["ucx", "ucxx"],
+    ["ucx", "ucx-old"],
 )
 @pytest.mark.parametrize("enable_tcp", [True, False])
 @pytest.mark.parametrize("enable_infiniband", [True, False])
 @pytest.mark.parametrize("enable_nvlink", [True, False])
 def test_get_preload_options(protocol, enable_tcp, enable_infiniband, enable_nvlink):
-    if protocol == "ucx":
-        pytest.importorskip("ucp")
-    elif protocol == "ucxx":
-        pytest.importorskip("ucxx")
+    get_ucx_implementation(protocol)
 
     opts = get_preload_options(
         protocol=protocol,
@@ -135,11 +137,17 @@ def test_get_preload_options(protocol, enable_tcp, enable_infiniband, enable_nvl
         assert "--enable-nvlink" in opts["preload_argv"]
 
 
+@pytest.mark.parametrize(
+    "protocol",
+    ["ucx", "ucx-old"],
+)
 @pytest.mark.parametrize("enable_tcp_over_ucx", [True, False, None])
 @pytest.mark.parametrize("enable_nvlink", [True, False, None])
 @pytest.mark.parametrize("enable_infiniband", [True, False, None])
-def test_get_ucx_config(enable_tcp_over_ucx, enable_infiniband, enable_nvlink):
-    pytest.importorskip("ucp")
+def test_get_ucx_config(
+    protocol, enable_tcp_over_ucx, enable_infiniband, enable_nvlink
+):
+    get_ucx_implementation(protocol)
 
     kwargs = {
         "enable_tcp_over_ucx": enable_tcp_over_ucx,
@@ -234,20 +242,98 @@ def test_parse_visible_devices():
         parse_cuda_visible_device([])
 
 
+def test_parse_device_bytes():
+    total = get_device_total_memory(0)
+
+    assert parse_device_memory_limit(None) is None
+    assert parse_device_memory_limit(0) is None
+    assert parse_device_memory_limit("0") is None
+    assert parse_device_memory_limit("0.0") is None
+    assert parse_device_memory_limit("0 GiB") is None
+
+    assert parse_device_memory_limit(1) == 1
+    assert parse_device_memory_limit("1") == 1
+
+    assert parse_device_memory_limit(1000000000) == 1000000000
+    assert parse_device_memory_limit("1GB") == 1000000000
+
+    if has_device_memory_resource(0):
+        assert parse_device_memory_limit(1.0) == total
+        assert parse_device_memory_limit("1.0") == total
+
+        assert parse_device_memory_limit(0.8) == int(total * 0.8)
+        assert parse_device_memory_limit(0.8, alignment_size=256) == int(
+            total * 0.8 // 256 * 256
+        )
+
+        assert parse_device_memory_limit("default") == parse_device_memory_limit(0.8)
+    else:
+        assert parse_device_memory_limit("default") is None
+
+        with pytest.raises(ValueError):
+            assert parse_device_memory_limit(1.0) == total
+        with pytest.raises(ValueError):
+            assert parse_device_memory_limit("1.0") == total
+        with pytest.raises(ValueError):
+            assert parse_device_memory_limit(0.8) == int(total * 0.8)
+        with pytest.raises(ValueError):
+            assert parse_device_memory_limit(0.8, alignment_size=256) == int(
+                total * 0.8 // 256 * 256
+            )
+
+
 def test_parse_device_memory_limit():
     total = get_device_total_memory(0)
 
-    assert parse_device_memory_limit(None) == total
-    assert parse_device_memory_limit(0) == total
+    assert parse_device_memory_limit(None) is None
+    assert parse_device_memory_limit(0) is None
+    assert parse_device_memory_limit("0") is None
+    assert parse_device_memory_limit(0.0) is None
+    assert parse_device_memory_limit("0 GiB") is None
+
+    assert parse_device_memory_limit(1) == 1
+    assert parse_device_memory_limit("1") == 1
+
     assert parse_device_memory_limit("auto") == total
 
-    assert parse_device_memory_limit(0.8) == int(total * 0.8)
-    assert parse_device_memory_limit(0.8, alignment_size=256) == int(
-        total * 0.8 // 256 * 256
-    )
     assert parse_device_memory_limit(1000000000) == 1000000000
     assert parse_device_memory_limit("1GB") == 1000000000
 
+    if has_device_memory_resource(0):
+        assert parse_device_memory_limit(1.0) == total
+        assert parse_device_memory_limit("1.0") == total
+
+        assert parse_device_memory_limit(0.8) == int(total * 0.8)
+        assert parse_device_memory_limit(0.8, alignment_size=256) == int(
+            total * 0.8 // 256 * 256
+        )
+        assert parse_device_memory_limit("default") == parse_device_memory_limit(0.8)
+    else:
+        assert parse_device_memory_limit("default") is None
+
+        with pytest.raises(ValueError):
+            assert parse_device_memory_limit(1.0) == total
+        with pytest.raises(ValueError):
+            assert parse_device_memory_limit("1.0") == total
+        with pytest.raises(ValueError):
+            assert parse_device_memory_limit(0.8) == int(total * 0.8)
+        with pytest.raises(ValueError):
+            assert parse_device_memory_limit(0.8, alignment_size=256) == int(
+                total * 0.8 // 256 * 256
+            )
+
+
+def test_has_device_memory_resoure():
+    has_memory_resource = has_device_memory_resource()
+    total = get_device_total_memory(0)
+
+    if has_memory_resource:
+        # Tested only in devices with a memory resource
+        assert total == parse_device_memory_limit("auto")
+    else:
+        # Tested only in devices without a memory resource
+        assert total is None
+
 
 def test_parse_visible_mig_devices():
     pynvml.nvmlInit()

dask_cuda/tests/test_worker_spec.py

@@ -1,3 +1,6 @@
+# SPDX-FileCopyrightText: Copyright (c) 2019-2025, NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+
 import pytest
 
 from distributed import Nanny
@@ -28,7 +31,7 @@ def _check_env_value(spec, k, v):
 @pytest.mark.parametrize("num_devices", [1, 4])
 @pytest.mark.parametrize("cls", [Nanny])
 @pytest.mark.parametrize("interface", [None, "eth0", "enp1s0f0"])
-@pytest.mark.parametrize("protocol", [None, "tcp", "ucx"])
+@pytest.mark.parametrize("protocol", [None, "tcp", "ucx", "ucx-old"])
 @pytest.mark.parametrize("dashboard_address", [None, ":0", ":8787"])
 @pytest.mark.parametrize("threads_per_worker", [1, 8])
 @pytest.mark.parametrize("silence_logs", [False, True])
@@ -58,7 +61,7 @@ def test_worker_spec(
             enable_nvlink=enable_nvlink,
         )
 
-    if (enable_infiniband or enable_nvlink) and protocol != "ucx":
+    if (enable_infiniband or enable_nvlink) and protocol not in ("ucx", "ucx-old"):
         with pytest.raises(
             TypeError, match="Enabling InfiniBand or NVLink requires protocol='ucx'"
         ):