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

dask_cuda/tests/test_local_cuda_cluster.py

@@ -2,6 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import asyncio
+import contextlib
 import os
 import pkgutil
 import sys
@@ -20,6 +21,7 @@ from dask_cuda.utils import (
     get_device_total_memory,
     get_gpu_count_mig,
     get_gpu_uuid,
+    has_device_memory_resource,
     print_cluster_config,
 )
 from dask_cuda.utils_test import MockWorker
@@ -28,7 +30,8 @@ from dask_cuda.utils_test import MockWorker
 @gen_test(timeout=20)
 async def test_local_cuda_cluster():
     async with LocalCUDACluster(
-        scheduler_port=0, asynchronous=True, device_memory_limit=1
+        scheduler_port=0,
+        asynchronous=True,
     ) as cluster:
         async with Client(cluster, asynchronous=True) as client:
             assert len(cluster.workers) == utils.get_n_gpus()
@@ -68,8 +71,8 @@ async def test_with_subset_of_cuda_visible_devices():
     async with LocalCUDACluster(
         scheduler_port=0,
         asynchronous=True,
-        device_memory_limit=1,
         worker_class=MockWorker,
+        data=dict,
     ) as cluster:
         async with Client(cluster, asynchronous=True) as client:
             assert len(cluster.workers) == 4
@@ -90,40 +93,27 @@ async def test_with_subset_of_cuda_visible_devices():
                 }
 
 
-@pytest.mark.parametrize(
-    "protocol",
-    ["ucx", "ucxx"],
-)
 @gen_test(timeout=20)
-async def test_ucx_protocol(protocol):
-    if protocol == "ucx":
-        pytest.importorskip("ucp")
-    elif protocol == "ucxx":
-        pytest.importorskip("ucxx")
+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", "ucxx"],
-)
 @gen_test(timeout=20)
-async def test_explicit_ucx_with_protocol_none(protocol):
-    if protocol == "ucx":
-        pytest.importorskip("ucp")
-    elif protocol == "ucxx":
-        pytest.importorskip("ucxx")
+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, data=dict
+        protocol=None,
+        enable_tcp_over_ucx=True,
+        asynchronous=True,
     ) as cluster:
         assert all(
             ws.address.startswith("ucx://") for ws in cluster.scheduler.workers.values()
@@ -131,18 +121,11 @@ async def test_explicit_ucx_with_protocol_none(protocol):
 
 
 @pytest.mark.filterwarnings("ignore:Exception ignored in")
-@pytest.mark.parametrize(
-    "protocol",
-    ["ucx", "ucxx"],
-)
 @gen_test(timeout=20)
-async def test_ucx_protocol_type_error(protocol):
-    if protocol == "ucx":
-        pytest.importorskip("ucp")
-    elif protocol == "ucxx":
-        pytest.importorskip("ucxx")
+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
@@ -153,7 +136,10 @@ async def test_ucx_protocol_type_error(protocol):
 @gen_test(timeout=20)
 async def test_n_workers():
     async with LocalCUDACluster(
-        CUDA_VISIBLE_DEVICES="0,1", worker_class=MockWorker, asynchronous=True
+        CUDA_VISIBLE_DEVICES="0,1",
+        worker_class=MockWorker,
+        asynchronous=True,
+        data=dict,
     ) as cluster:
         assert len(cluster.workers) == 2
         assert len(cluster.worker_spec) == 2
@@ -208,10 +194,13 @@ async def test_no_memory_limits_cudaworker():
 @gen_test(timeout=20)
 async def test_all_to_all():
     async with LocalCUDACluster(
-        CUDA_VISIBLE_DEVICES="0,1", worker_class=MockWorker, asynchronous=True
+        CUDA_VISIBLE_DEVICES="0,1",
+        worker_class=MockWorker,
+        asynchronous=True,
+        data=dict,
     ) as cluster:
         async with Client(cluster, asynchronous=True) as client:
-            workers = list(client.scheduler_info()["workers"])
+            workers = list(client.scheduler_info(n_workers=-1)["workers"])
             n_workers = len(workers)
             await utils.all_to_all(client)
             # assert all to all has resulted in all data on every worker
@@ -263,11 +252,6 @@ async def test_rmm_managed():
 async def test_rmm_async():
     rmm = pytest.importorskip("rmm")
 
-    driver_version = rmm._cuda.gpu.driverGetVersion()
-    runtime_version = rmm._cuda.gpu.runtimeGetVersion()
-    if driver_version < 11020 or runtime_version < 11020:
-        pytest.skip("cudaMallocAsync not supported")
-
     async with LocalCUDACluster(
         rmm_async=True,
         rmm_pool_size="2GB",
@@ -290,11 +274,6 @@ async def test_rmm_async():
 async def test_rmm_async_with_maximum_pool_size():
     rmm = pytest.importorskip("rmm")
 
-    driver_version = rmm._cuda.gpu.driverGetVersion()
-    runtime_version = rmm._cuda.gpu.runtimeGetVersion()
-    if driver_version < 11020 or runtime_version < 11020:
-        pytest.skip("cudaMallocAsync not supported")
-
     async with LocalCUDACluster(
         rmm_async=True,
         rmm_pool_size="2GB",
@@ -381,7 +360,6 @@ async def test_cluster_worker():
     async with LocalCUDACluster(
         scheduler_port=0,
         asynchronous=True,
-        device_memory_limit=1,
         n_workers=1,
     ) as cluster:
         assert len(cluster.workers) == 1
@@ -464,7 +442,7 @@ async def test_get_cluster_configuration():
     async with LocalCUDACluster(
         rmm_pool_size="2GB",
         rmm_maximum_pool_size="3GB",
-        device_memory_limit="30B",
+        device_memory_limit="30B" if has_device_memory_resource() else None,
         CUDA_VISIBLE_DEVICES="0",
         scheduler_port=0,
         asynchronous=True,
@@ -474,10 +452,14 @@ async def test_get_cluster_configuration():
             assert ret["[plugin] RMMSetup"]["initial_pool_size"] == 2000000000
             assert ret["[plugin] RMMSetup"]["maximum_pool_size"] == 3000000000
             assert ret["jit-unspill"] is False
-            assert ret["device-memory-limit"] == 30
+            if has_device_memory_resource():
+                assert ret["device-memory-limit"] == 30
 
 
 @gen_test(timeout=20)
+@pytest.mark.skip_if_no_device_memory(
+    "Devices without dedicated memory resources do not support fractional limits"
+)
 async def test_worker_fraction_limits():
     async with LocalCUDACluster(
         dashboard_address=None,
@@ -503,6 +485,40 @@ async def test_worker_fraction_limits():
             )
 
 
+# Intentionally not using @gen_test to skip cleanup checks
+@pytest.mark.parametrize(
+    "argument", ["pool_size", "maximum_pool_size", "release_threshold"]
+)
+@pytest.mark.xfail(reason="https://github.com/rapidsai/dask-cuda/issues/1265")
+@pytest.mark.skip_if_device_memory(
+    "Devices with dedicated memory resources cannot test error"
+)
+def test_worker_fraction_limits_no_dedicated_memory(argument):
+    async def _test_worker_fraction_limits_no_dedicated_memory():
+        if argument == "pool_size":
+            kwargs = {"rmm_pool_size": "0.1"}
+        elif argument == "maximum_pool_size":
+            kwargs = {"rmm_pool_size": "1 GiB", "rmm_maximum_pool_size": "0.1"}
+        else:
+            kwargs = {"rmm_async": True, "rmm_release_threshold": "0.1"}
+
+        with raises_with_cause(
+            RuntimeError,
+            "Nanny failed to start",
+            RuntimeError,
+            "Worker failed to start",
+            ValueError,
+            "Fractional of total device memory not supported in devices without a "
+            "dedicated memory resource",
+        ):
+            await LocalCUDACluster(
+                asynchronous=True,
+                **kwargs,
+            )
+
+    asyncio.run(_test_worker_fraction_limits_no_dedicated_memory())
+
+
 @gen_test(timeout=20)
 async def test_cudf_spill_disabled():
     cudf = pytest.importorskip("cudf")
@@ -527,6 +543,9 @@ async def test_cudf_spill_disabled():
 
 
 @gen_test(timeout=20)
+@pytest.mark.skip_if_no_device_memory(
+    "Devices without dedicated memory resources cannot enable cuDF spill"
+)
 async def test_cudf_spill():
     cudf = pytest.importorskip("cudf")
 
@@ -551,27 +570,69 @@ async def test_cudf_spill():
                 assert v == 2
 
 
+@pytest.mark.skip_if_device_memory(
+    "Devices with dedicated memory resources cannot test error"
+)
+@gen_test(timeout=20)
+async def test_cudf_spill_no_dedicated_memory():
+    cudf = pytest.importorskip("cudf")  # noqa: F841
+
+    with pytest.raises(
+        ValueError,
+        match="cuDF spilling is not supported on devices without dedicated memory",
+    ):
+        await LocalCUDACluster(
+            enable_cudf_spill=True,
+            cudf_spill_stats=2,
+            asynchronous=True,
+        )
+
+
+@pytest.mark.parametrize(
+    "jit_unspill",
+    [False, True],
+)
 @pytest.mark.parametrize(
-    "protocol",
-    ["ucx", "ucxx"],
+    "device_memory_limit",
+    [None, "1B"],
 )
-def test_print_cluster_config(capsys, protocol):
-    if protocol == "ucx":
-        pytest.importorskip("ucp")
-    elif protocol == "ucxx":
-        pytest.importorskip("ucxx")
+def test_print_cluster_config(capsys, jit_unspill, device_memory_limit):
+    pytest.importorskip("distributed_ucxx")
 
     pytest.importorskip("rich")
-    with LocalCUDACluster(
-        n_workers=1, device_memory_limit="1B", jit_unspill=True, protocol=protocol
-    ) 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 "1 B" in captured.out
-            assert "[plugin]" in captured.out
+
+    ctx = contextlib.nullcontext()
+    if not has_device_memory_resource():
+        if device_memory_limit:
+            ctx = pytest.raises(
+                ValueError,
+                match="device_memory_limit is set but device has no dedicated memory.",
+            )
+        if jit_unspill:
+            # JIT-Unspill exception has precedence, thus overwrite ctx if both are
+            # enabled
+            ctx = pytest.raises(
+                ValueError,
+                match="JIT-Unspill is not supported on devices without dedicated "
+                "memory",
+            )
+
+    with ctx:
+        with LocalCUDACluster(
+            n_workers=1,
+            device_memory_limit=device_memory_limit,
+            jit_unspill=jit_unspill,
+            protocol="ucx",
+        ) as cluster:
+            with Client(cluster) as client:
+                print_cluster_config(client)
+                captured = capsys.readouterr()
+                assert "Dask Cluster Configuration" 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()
 
 
 @pytest.mark.xfail(reason="https://github.com/rapidsai/dask-cuda/issues/1265")

dask_cuda/tests/test_proxify_host_file.py

@@ -1,4 +1,5 @@
-# Copyright (c) 2025, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2021-2025, NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
 
 from typing import Iterable
 from unittest.mock import patch
@@ -219,6 +220,9 @@ def test_one_item_host_limit(capsys, root_dir):
     assert len(dhf.manager) == 0
 
 
+@pytest.mark.skip_if_no_device_memory(
+    "Devices without dedicated memory resources do not support spilling"
+)
 def test_spill_on_demand(root_dir):
     """
     Test spilling on demand by disabling the device_memory_limit
@@ -241,6 +245,9 @@ def test_spill_on_demand(root_dir):
 
 
 @pytest.mark.parametrize("jit_unspill", [True, False])
+@pytest.mark.skip_if_no_device_memory(
+    "Devices without dedicated memory resources do not support spilling"
+)
 @gen_test(timeout=20)
 async def test_local_cuda_cluster(jit_unspill):
     """Testing spilling of a proxied cudf dataframe in a local cuda cluster"""
@@ -377,9 +384,9 @@ def test_externals(root_dir):
 
 @patch("dask_cuda.proxify_device_objects.incompatible_types", (cupy.ndarray,))
 def test_incompatible_types(root_dir):
-    """Check that ProxifyHostFile unproxifies `cupy.ndarray` on retrieval
+    """Check that ProxifyHostFile unproxifies ``cupy.ndarray`` on retrieval
 
-    Notice, in this test we add `cupy.ndarray` to the incompatible_types temporarily.
+    Notice, in this test we add ``cupy.ndarray`` to the incompatible_types temporarily.
     """
     cupy = pytest.importorskip("cupy")
     cudf = pytest.importorskip("cudf")
@@ -398,6 +405,9 @@ def test_incompatible_types(root_dir):
 
 @pytest.mark.parametrize("npartitions", [1, 2, 3])
 @pytest.mark.parametrize("compatibility_mode", [True, False])
+@pytest.mark.skip_if_no_device_memory(
+    "Devices without dedicated memory resources do not support JIT-Unspill"
+)
 @gen_test(timeout=30)
 async def test_compatibility_mode_dataframe_shuffle(compatibility_mode, npartitions):
     cudf = pytest.importorskip("cudf")
@@ -430,19 +440,28 @@ async def test_compatibility_mode_dataframe_shuffle(compatibility_mode, npartiti
                     assert all(res)  # Only proxy objects
 
 
+@pytest.mark.skip_if_no_device_memory(
+    "Devices without dedicated memory resources do not support JIT-Unspill"
+)
 @gen_test(timeout=60)
 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):
@@ -453,18 +472,20 @@ 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
 
 
+@pytest.mark.skip_if_no_device_memory(
+    "Devices without dedicated memory resources do not support JIT-Unspill"
+)
 def test_on_demand_debug_info():
     """Test worker logging when on-demand-spilling fails"""
     rmm = pytest.importorskip("rmm")

dask_cuda/tests/test_proxy.py

@@ -1,3 +1,6 @@
+# SPDX-FileCopyrightText: Copyright (c) 2020-2025, NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+
 import operator
 import os
 import pickle
@@ -242,7 +245,7 @@ def test_serialize_of_proxied_cudf(proxy_serializers, dask_serializers):
 
 @pytest.mark.parametrize("backend", ["numpy", "cupy"])
 def test_fixed_attribute_length(backend):
-    """Test fixed attribute `x.__len__` access
+    """Test fixed attribute ``x.__len__`` access
 
     Notice, accessing fixed attributes shouldn't de-serialize the proxied object
     """
@@ -263,7 +266,7 @@ def test_fixed_attribute_length(backend):
 
 
 def test_fixed_attribute_name():
-    """Test fixed attribute `x.name` access
+    """Test fixed attribute ``x.name`` access
 
     Notice, accessing fixed attributes shouldn't de-serialize the proxied object
     """
@@ -284,6 +287,9 @@ def test_fixed_attribute_name():
 
 
 @pytest.mark.parametrize("jit_unspill", [True, False])
+@pytest.mark.skip_if_no_device_memory(
+    "Spilling not supported in devices without dedicated memory resource"
+)
 @gen_test(timeout=20)
 async def test_spilling_local_cuda_cluster(jit_unspill):
     """Testing spilling of a proxied cudf dataframe in a local cuda cluster"""
@@ -386,8 +392,8 @@ def test_serializing_array_to_disk(backend, serializers, size):
 class _PxyObjTest(proxy_object.ProxyObject):
     """
     A class that:
-        - defines `__dask_tokenize__` in order to avoid deserialization when
-          calling `client.scatter()`
+        - defines ``__dask_tokenize__`` in order to avoid deserialization when
+          calling ``client.scatter()``
         - Asserts that no deserialization is performaned when communicating.
     """
 
@@ -401,14 +407,12 @@ class _PxyObjTest(proxy_object.ProxyObject):
 
 
 @pytest.mark.parametrize("send_serializers", [None, ("dask", "pickle"), ("cuda",)])
-@pytest.mark.parametrize("protocol", ["tcp", "ucx", "ucxx"])
+@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 == "ucx":
-        pytest.importorskip("ucp")
-    elif protocol == "ucxx":
-        pytest.importorskip("ucxx")
+        pytest.importorskip("distributed_ucxx")
     cudf = pytest.importorskip("cudf")
 
     def task(x):
@@ -417,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", "ucxx"]:
+        if protocol == "ucx":
             if send_serializers is None:
                 assert serializers_used == "cuda"
             else:
@@ -448,15 +452,13 @@ async def test_communicating_proxy_objects(protocol, send_serializers):
             await client.submit(task, df)
 
 
-@pytest.mark.parametrize("protocol", ["tcp", "ucx", "ucxx"])
+@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 == "ucx":
-        pytest.importorskip("ucp")
-    elif protocol == "ucxx":
-        pytest.importorskip("ucxx")
+        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,17 @@ 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(
+        "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)