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

dask_cuda/tests/test_dgx.py

@@ -1,3 +1,6 @@
+# SPDX-FileCopyrightText: Copyright (c) 2019-2025, NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+
 import multiprocessing as mp
 import os
 from enum import Enum, auto
@@ -10,6 +13,7 @@ from distributed import Client
 
 from dask_cuda import LocalCUDACluster
 from dask_cuda.initialize import initialize
+from dask_cuda.utils_test import get_ucx_implementation
 
 mp = mp.get_context("spawn")  # type: ignore
 psutil = pytest.importorskip("psutil")
@@ -78,10 +82,7 @@ def test_default():
 
 
 def _test_tcp_over_ucx(protocol):
-    if protocol == "ucx":
-        ucp = pytest.importorskip("ucp")
-    elif protocol == "ucxx":
-        ucp = pytest.importorskip("ucxx")
+    ucp = get_ucx_implementation(protocol)
 
     with LocalCUDACluster(protocol=protocol, enable_tcp_over_ucx=True) as cluster:
         with Client(cluster) as client:
@@ -102,13 +103,10 @@ def _test_tcp_over_ucx(protocol):
 
 @pytest.mark.parametrize(
     "protocol",
-    ["ucx", "ucxx"],
+    ["ucx", "ucx-old"],
 )
 def test_tcp_over_ucx(protocol):
-    if protocol == "ucx":
-        ucp = pytest.importorskip("ucp")
-    elif protocol == "ucxx":
-        ucp = pytest.importorskip("ucxx")
+    ucp = get_ucx_implementation(protocol)
     if _is_ucx_116(ucp):
         pytest.skip("https://github.com/rapidsai/ucx-py/issues/1037")
 
@@ -137,10 +135,7 @@ def _test_ucx_infiniband_nvlink(
     skip_queue, protocol, enable_infiniband, enable_nvlink, enable_rdmacm
 ):
     cupy = pytest.importorskip("cupy")
-    if protocol == "ucx":
-        ucp = pytest.importorskip("ucp")
-    elif protocol == "ucxx":
-        ucp = pytest.importorskip("ucxx")
+    ucp = get_ucx_implementation(protocol)
 
     if enable_infiniband and not any(
         [at.startswith("rc") for at in ucp.get_active_transports()]
@@ -206,7 +201,7 @@ def _test_ucx_infiniband_nvlink(
             assert all(client.run(check_ucx_options).values())
 
 
-@pytest.mark.parametrize("protocol", ["ucx", "ucxx"])
+@pytest.mark.parametrize("protocol", ["ucx", "ucx-old"])
 @pytest.mark.parametrize(
     "params",
     [
@@ -222,10 +217,7 @@ def _test_ucx_infiniband_nvlink(
     reason="Automatic InfiniBand device detection Unsupported for %s" % _get_dgx_name(),
 )
 def test_ucx_infiniband_nvlink(protocol, params):
-    if protocol == "ucx":
-        ucp = pytest.importorskip("ucp")
-    elif protocol == "ucxx":
-        ucp = pytest.importorskip("ucxx")
+    ucp = get_ucx_implementation(protocol)
     if _is_ucx_116(ucp) and params["enable_infiniband"] is False:
         pytest.skip("https://github.com/rapidsai/ucx-py/issues/1037")
 

dask_cuda/tests/test_explicit_comms.py

@@ -1,4 +1,5 @@
-# Copyright (c) 2021-2025 NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2021-2025, NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
 
 import asyncio
 import multiprocessing as mp
@@ -25,7 +26,7 @@ from dask_cuda.explicit_comms.dataframe.shuffle import (
     _contains_shuffle_expr,
     shuffle as explicit_comms_shuffle,
 )
-from dask_cuda.utils_test import IncreasedCloseTimeoutNanny
+from dask_cuda.utils_test import IncreasedCloseTimeoutNanny, get_ucx_implementation
 
 mp = mp.get_context("spawn")  # type: ignore
 ucp = pytest.importorskip("ucp")
@@ -53,8 +54,10 @@ def _test_local_cluster(protocol):
             assert sum(c.run(my_rank, 0)) == sum(range(4))
 
 
-@pytest.mark.parametrize("protocol", ["tcp", "ucx", "ucxx"])
+@pytest.mark.parametrize("protocol", ["tcp", "ucx", "ucx-old"])
 def test_local_cluster(protocol):
+    if protocol.startswith("ucx"):
+        get_ucx_implementation(protocol)
     p = mp.Process(target=_test_local_cluster, args=(protocol,))
     p.start()
     p.join()
@@ -97,7 +100,7 @@ def test_dataframe_merge_empty_partitions():
 
 
 def check_partitions(df, npartitions):
-    """Check that all values in `df` hashes to the same"""
+    """Check that all values in ``df`` hashes to the same"""
     dtypes = {}
     for col, dtype in df.dtypes.items():
         if pd.api.types.is_numeric_dtype(dtype):
@@ -199,11 +202,13 @@ def _test_dataframe_shuffle(backend, protocol, n_workers, _partitions):
 
 @pytest.mark.parametrize("nworkers", [1, 2, 3])
 @pytest.mark.parametrize("backend", ["pandas", "cudf"])
-@pytest.mark.parametrize("protocol", ["tcp", "ucx", "ucxx"])
+@pytest.mark.parametrize("protocol", ["tcp", "ucx", "ucx-old"])
 @pytest.mark.parametrize("_partitions", [True, False])
 def test_dataframe_shuffle(backend, protocol, nworkers, _partitions):
     if backend == "cudf":
         pytest.importorskip("cudf")
+    if protocol.startswith("ucx"):
+        get_ucx_implementation(protocol)
 
     p = mp.Process(
         target=_test_dataframe_shuffle, args=(backend, protocol, nworkers, _partitions)
@@ -320,10 +325,12 @@ def _test_dataframe_shuffle_merge(backend, protocol, n_workers):
 
 @pytest.mark.parametrize("nworkers", [1, 2, 4])
 @pytest.mark.parametrize("backend", ["pandas", "cudf"])
-@pytest.mark.parametrize("protocol", ["tcp", "ucx", "ucxx"])
+@pytest.mark.parametrize("protocol", ["tcp", "ucx", "ucx-old"])
 def test_dataframe_shuffle_merge(backend, protocol, nworkers):
     if backend == "cudf":
         pytest.importorskip("cudf")
+    if protocol.startswith("ucx"):
+        get_ucx_implementation(protocol)
     p = mp.Process(
         target=_test_dataframe_shuffle_merge, args=(backend, protocol, nworkers)
     )
@@ -357,9 +364,14 @@ def _test_jit_unspill(protocol):
             assert_eq(got, expected)
 
 
-@pytest.mark.parametrize("protocol", ["tcp", "ucx", "ucxx"])
+@pytest.mark.parametrize("protocol", ["tcp", "ucx", "ucx-old"])
+@pytest.mark.skip_if_no_device_memory(
+    "JIT-Unspill not supported in devices without dedicated memory resource"
+)
 def test_jit_unspill(protocol):
     pytest.importorskip("cudf")
+    if protocol.startswith("ucx"):
+        get_ucx_implementation(protocol)
 
     p = mp.Process(target=_test_jit_unspill, args=(protocol,))
     p.start()
@@ -384,7 +396,7 @@ def _test_lock_workers(scheduler_address, ranks):
 
 def test_lock_workers():
     """
-    Testing `run(...,lock_workers=True)` by spawning 30 runs with overlapping
+    Testing ``run(...,lock_workers=True)`` by spawning 30 runs with overlapping
     and non-overlapping worker sets.
     """
     try:
@@ -423,7 +435,9 @@ def test_create_destroy_create():
     with LocalCluster(n_workers=1) as cluster:
         with Client(cluster) as client:
             context = comms.default_comms()
-            scheduler_addresses_old = list(client.scheduler_info()["workers"].keys())
+            scheduler_addresses_old = list(
+                client.scheduler_info(n_workers=-1)["workers"].keys()
+            )
             comms_addresses_old = list(comms.default_comms().worker_addresses)
             assert comms.default_comms() is context
             assert len(comms._comms_cache) == 1
@@ -444,7 +458,9 @@ def test_create_destroy_create():
     # because we referenced the old cluster's addresses.
     with LocalCluster(n_workers=1) as cluster:
         with Client(cluster) as client:
-            scheduler_addresses_new = list(client.scheduler_info()["workers"].keys())
+            scheduler_addresses_new = list(
+                client.scheduler_info(n_workers=-1)["workers"].keys()
+            )
             comms_addresses_new = list(comms.default_comms().worker_addresses)
 
     assert scheduler_addresses_new == comms_addresses_new
@@ -485,7 +501,8 @@ def test_scaled_cluster_gets_new_comms_context():
                 "n_workers": 2,
             }
             expected_1 = {
-                k: expected_values for k in client.scheduler_info()["workers"]
+                k: expected_values
+                for k in client.scheduler_info(n_workers=-1)["workers"]
             }
             assert result_1 == expected_1
 
@@ -513,7 +530,8 @@ def test_scaled_cluster_gets_new_comms_context():
                 "n_workers": 3,
             }
             expected_2 = {
-                k: expected_values for k in client.scheduler_info()["workers"]
+                k: expected_values
+                for k in client.scheduler_info(n_workers=-1)["workers"]
             }
             assert result_2 == expected_2
 

dask_cuda/tests/test_from_array.py

@@ -1,19 +1,21 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+
 import pytest
 
 import dask.array as da
 from distributed import Client
 
 from dask_cuda import LocalCUDACluster
+from dask_cuda.utils_test import get_ucx_implementation
 
 cupy = pytest.importorskip("cupy")
 
 
-@pytest.mark.parametrize("protocol", ["ucx", "ucxx", "tcp"])
+@pytest.mark.parametrize("protocol", ["ucx", "ucx-old", "tcp"])
 def test_ucx_from_array(protocol):
-    if protocol == "ucx":
-        pytest.importorskip("ucp")
-    elif protocol == "ucxx":
-        pytest.importorskip("ucxx")
+    if protocol.startswith("ucx"):
+        get_ucx_implementation(protocol)
 
     N = 10_000
     with LocalCUDACluster(protocol=protocol) as cluster:

dask_cuda/tests/test_initialize.py

@@ -1,3 +1,6 @@
+# SPDX-FileCopyrightText: Copyright (c) 2019-2025, NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+
 import multiprocessing as mp
 import sys
 
@@ -11,7 +14,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
+from dask_cuda.utils_test import IncreasedCloseTimeoutNanny, get_ucx_implementation
 
 mp = mp.get_context("spawn")  # type: ignore
 
@@ -22,10 +25,7 @@ mp = mp.get_context("spawn")  # type: ignore
 
 
 def _test_initialize_ucx_tcp(protocol):
-    if protocol == "ucx":
-        ucp = pytest.importorskip("ucp")
-    elif protocol == "ucxx":
-        ucp = pytest.importorskip("ucxx")
+    ucp = get_ucx_implementation(protocol)
 
     kwargs = {"enable_tcp_over_ucx": True}
     initialize(protocol=protocol, **kwargs)
@@ -55,12 +55,9 @@ def _test_initialize_ucx_tcp(protocol):
             assert all(client.run(check_ucx_options).values())
 
 
-@pytest.mark.parametrize("protocol", ["ucx", "ucxx"])
+@pytest.mark.parametrize("protocol", ["ucx", "ucx-old"])
 def test_initialize_ucx_tcp(protocol):
-    if protocol == "ucx":
-        pytest.importorskip("ucp")
-    elif protocol == "ucxx":
-        pytest.importorskip("ucxx")
+    get_ucx_implementation(protocol)
 
     p = mp.Process(target=_test_initialize_ucx_tcp, args=(protocol,))
     p.start()
@@ -69,10 +66,7 @@ def test_initialize_ucx_tcp(protocol):
 
 
 def _test_initialize_ucx_nvlink(protocol):
-    if protocol == "ucx":
-        ucp = pytest.importorskip("ucp")
-    elif protocol == "ucxx":
-        ucp = pytest.importorskip("ucxx")
+    ucp = get_ucx_implementation(protocol)
 
     kwargs = {"enable_nvlink": True}
     initialize(protocol=protocol, **kwargs)
@@ -103,12 +97,9 @@ def _test_initialize_ucx_nvlink(protocol):
             assert all(client.run(check_ucx_options).values())
 
 
-@pytest.mark.parametrize("protocol", ["ucx", "ucxx"])
+@pytest.mark.parametrize("protocol", ["ucx", "ucx-old"])
 def test_initialize_ucx_nvlink(protocol):
-    if protocol == "ucx":
-        pytest.importorskip("ucp")
-    elif protocol == "ucxx":
-        pytest.importorskip("ucxx")
+    get_ucx_implementation(protocol)
 
     p = mp.Process(target=_test_initialize_ucx_nvlink, args=(protocol,))
     p.start()
@@ -117,10 +108,7 @@ def test_initialize_ucx_nvlink(protocol):
 
 
 def _test_initialize_ucx_infiniband(protocol):
-    if protocol == "ucx":
-        ucp = pytest.importorskip("ucp")
-    elif protocol == "ucxx":
-        ucp = pytest.importorskip("ucxx")
+    ucp = get_ucx_implementation(protocol)
 
     kwargs = {"enable_infiniband": True}
     initialize(protocol=protocol, **kwargs)
@@ -154,12 +142,9 @@ 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", "ucxx"])
+@pytest.mark.parametrize("protocol", ["ucx", "ucx-old"])
 def test_initialize_ucx_infiniband(protocol):
-    if protocol == "ucx":
-        pytest.importorskip("ucp")
-    elif protocol == "ucxx":
-        pytest.importorskip("ucxx")
+    get_ucx_implementation(protocol)
 
     p = mp.Process(target=_test_initialize_ucx_infiniband, args=(protocol,))
     p.start()
@@ -168,10 +153,7 @@ def test_initialize_ucx_infiniband(protocol):
 
 
 def _test_initialize_ucx_all(protocol):
-    if protocol == "ucx":
-        ucp = pytest.importorskip("ucp")
-    elif protocol == "ucxx":
-        ucp = pytest.importorskip("ucxx")
+    ucp = get_ucx_implementation(protocol)
 
     initialize(protocol=protocol)
     with LocalCluster(
@@ -204,12 +186,9 @@ def _test_initialize_ucx_all(protocol):
             assert all(client.run(check_ucx_options).values())
 
 
-@pytest.mark.parametrize("protocol", ["ucx", "ucxx"])
+@pytest.mark.parametrize("protocol", ["ucx", "ucx-old"])
 def test_initialize_ucx_all(protocol):
-    if protocol == "ucx":
-        pytest.importorskip("ucp")
-    elif protocol == "ucxx":
-        pytest.importorskip("ucxx")
+    get_ucx_implementation(protocol)
 
     p = mp.Process(target=_test_initialize_ucx_all, args=(protocol,))
     p.start()