dask-cuda 24.2.0a19__py3-none-any.whl → 24.4.0__py3-none-any.whl

dask_cuda/VERSION

@@ -1 +1 @@
-24.02.00a19
+24.04.00

dask_cuda/__init__.py

@@ -3,7 +3,6 @@ import sys
 if sys.platform != "linux":
     raise ImportError("Only Linux is supported by Dask-CUDA at this time")
 
-
 import dask
 import dask.utils
 import dask.dataframe.core

dask_cuda/_version.py

@@ -17,4 +17,4 @@ import importlib.resources
 __version__ = (
     importlib.resources.files("dask_cuda").joinpath("VERSION").read_text().strip()
 )
-__git_commit__ = "6de222f9763c0b267cc969ccc290d9f3a3a12b8a"
+__git_commit__ = "7ed67359e2264393e97ad57fc39829c3a69aa1f3"

dask_cuda/explicit_comms/dataframe/shuffle.py

@@ -577,7 +577,7 @@ def get_rearrange_by_column_wrapper(func):
             kw = kw.arguments
             # Notice, we only overwrite the default and the "tasks" shuffle
             # algorithm. The "disk" and "p2p" algorithm, we don't touch.
-            if kw["shuffle"] in ("tasks", None):
+            if kw["shuffle_method"] in ("tasks", None):
                 col = kw["col"]
                 if isinstance(col, str):
                     col = [col]

dask_cuda/get_device_memory_objects.py

@@ -124,6 +124,10 @@ def get_device_memory_objects_register_cudf():
     def get_device_memory_objects_cudf_multiindex(obj):
         return dispatch(obj._columns)
 
+    @dispatch.register(cudf.core.column.ColumnBase)
+    def get_device_memory_objects_cudf_column(obj):
+        return dispatch(obj.data) + dispatch(obj.children) + dispatch(obj.mask)
+
 
 @sizeof.register_lazy("cupy")
 def register_cupy():  # NB: this overwrites dask.sizeof.register_cupy()

dask_cuda/tests/test_dgx.py

@@ -144,6 +144,10 @@ def _test_ucx_infiniband_nvlink(
     else:
         skip_queue.put("ok")
 
+    # `ucp.get_active_transports()` call above initializes UCX, we must reset it
+    # so that Dask doesn't try to initialize it again and raise an exception.
+    ucp.reset()
+
     if enable_infiniband is None and enable_nvlink is None and enable_rdmacm is None:
         enable_tcp_over_ucx = None
         cm_tls = ["all"]

dask_cuda/tests/test_explicit_comms.py

@@ -1,6 +1,9 @@
 import asyncio
 import multiprocessing as mp
 import os
+import signal
+import time
+from functools import partial
 from unittest.mock import patch
 
 import numpy as np
@@ -175,7 +178,7 @@ def test_dataframe_shuffle(backend, protocol, nworkers, _partitions):
 
 
 @pytest.mark.parametrize("in_cluster", [True, False])
-def test_dask_use_explicit_comms(in_cluster):
+def _test_dask_use_explicit_comms(in_cluster):
     def check_shuffle():
         """Check if shuffle use explicit-comms by search for keys named
         'explicit-comms-shuffle'
@@ -217,6 +220,31 @@ def test_dask_use_explicit_comms(in_cluster):
         check_shuffle()
 
 
+@pytest.mark.parametrize("in_cluster", [True, False])
+def test_dask_use_explicit_comms(in_cluster):
+    def _timeout(process, function, timeout):
+        if process.is_alive():
+            function()
+        timeout = time.time() + timeout
+        while process.is_alive() and time.time() < timeout:
+            time.sleep(0.1)
+
+    p = mp.Process(target=_test_dask_use_explicit_comms, args=(in_cluster,))
+    p.start()
+
+    # Timeout before killing process
+    _timeout(p, lambda: None, 60.0)
+
+    # Send SIGINT (i.e., KeyboardInterrupt) hoping we get a stack trace.
+    _timeout(p, partial(p._popen._send_signal, signal.SIGINT), 3.0)
+
+    # SIGINT didn't work, kill process.
+    _timeout(p, p.kill, 3.0)
+
+    assert not p.is_alive()
+    assert p.exitcode == 0
+
+
 def _test_dataframe_shuffle_merge(backend, protocol, n_workers):
     if backend == "cudf":
         cudf = pytest.importorskip("cudf")

dask_cuda/tests/test_proxify_host_file.py

@@ -302,13 +302,24 @@ def test_dataframes_share_dev_mem(root_dir):
 def test_cudf_get_device_memory_objects():
     cudf = pytest.importorskip("cudf")
     objects = [
-        cudf.DataFrame({"a": range(10), "b": range(10)}, index=reversed(range(10))),
+        cudf.DataFrame(
+            {"a": [0, 1, 2, 3, None, 5, 6, 7, 8, 9], "b": range(10)},
+            index=reversed(range(10)),
+        ),
         cudf.MultiIndex(
             levels=[[1, 2], ["blue", "red"]], codes=[[0, 0, 1, 1], [1, 0, 1, 0]]
         ),
     ]
     res = get_device_memory_ids(objects)
-    assert len(res) == 4, "We expect four buffer objects"
+    # Buffers are:
+    # 1. int data for objects[0].a
+    # 2. mask data for objects[0].a
+    # 3. int data for objects[0].b
+    # 4. int data for objects[0].index
+    # 5. int data for objects[1].levels[0]
+    # 6. char data for objects[1].levels[1]
+    # 7. offset data for objects[1].levels[1]
+    assert len(res) == 7, "We expect seven buffer objects"
 
 
 def test_externals(root_dir):
@@ -403,7 +414,7 @@ async def test_compatibility_mode_dataframe_shuffle(compatibility_mode, npartiti
                 ddf = dask.dataframe.from_pandas(
                     cudf.DataFrame({"key": np.arange(10)}), npartitions=npartitions
                 )
-                res = ddf.shuffle(on="key", shuffle="tasks").persist()
+                res = ddf.shuffle(on="key", shuffle_method="tasks").persist()
 
                 # With compatibility mode on, we shouldn't encounter any proxy objects
                 if compatibility_mode:

{dask_cuda-24.2.0a19.dist-info → dask_cuda-24.4.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: dask-cuda
-Version: 24.2.0a19
+Version: 24.4.0
 Summary: Utilities for Dask and CUDA interactions
 Author: NVIDIA Corporation
 License: Apache-2.0
@@ -14,15 +14,16 @@ Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
 Requires-Python: >=3.9
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: click >=8.1
 Requires-Dist: numba >=0.57
-Requires-Dist: numpy >=1.21
-Requires-Dist: pandas <1.6.0.dev0,>=1.3
+Requires-Dist: numpy <2.0a0,>=1.23
+Requires-Dist: pandas >=1.3
 Requires-Dist: pynvml <11.5,>=11.0.0
-Requires-Dist: rapids-dask-dependency ==24.2.*,>=0.0.0a0
+Requires-Dist: rapids-dask-dependency ==24.4.*
 Requires-Dist: zict >=2.0.0
 Provides-Extra: docs
 Requires-Dist: numpydoc >=1.1.0 ; extra == 'docs'
@@ -30,12 +31,12 @@ Requires-Dist: sphinx ; extra == 'docs'
 Requires-Dist: sphinx-click >=2.7.1 ; extra == 'docs'
 Requires-Dist: sphinx-rtd-theme >=0.5.1 ; extra == 'docs'
 Provides-Extra: test
-Requires-Dist: cudf ==24.2.* ; extra == 'test'
-Requires-Dist: dask-cudf ==24.2.* ; extra == 'test'
-Requires-Dist: kvikio ==24.2.* ; extra == 'test'
+Requires-Dist: cudf ==24.4.* ; extra == 'test'
+Requires-Dist: dask-cudf ==24.4.* ; extra == 'test'
+Requires-Dist: kvikio ==24.4.* ; extra == 'test'
 Requires-Dist: pytest ; extra == 'test'
 Requires-Dist: pytest-cov ; extra == 'test'
-Requires-Dist: ucx-py ==0.36.* ; extra == 'test'
+Requires-Dist: ucx-py ==0.37.* ; extra == 'test'
 
 Dask CUDA
 =========

{dask_cuda-24.2.0a19.dist-info → dask_cuda-24.4.0.dist-info}/RECORD

@@ -1,11 +1,11 @@
-dask_cuda/VERSION,sha256=f7OZVJVxMaXTPGW-tA4V_zsXVn3l2dBiDE0k_JXKdJQ,12
-dask_cuda/__init__.py,sha256=XnMTUi-SvoGn7g1Dj6XW97HnQzGQv0G3EnvSjcZ7vU4,1455
-dask_cuda/_version.py,sha256=wOhN09vJAkzZyQ9x-YTQQgxa5lgLkzhYxCsOMHFUInI,778
+dask_cuda/VERSION,sha256=cS6_wRwcl4ZhhE4gk3mZbC_4dD-r28YRdEQeaKdKu1U,9
+dask_cuda/__init__.py,sha256=Vt42yCT1WhjZgehiNLDLw2uvfxjqLThD8uKDMyKQYsw,1454
+dask_cuda/_version.py,sha256=OcuA1fDRuukb6OC8LzTwSO2aJ3ZtB1N2pm08Nrga8y8,778
 dask_cuda/cli.py,sha256=XNRH0bu-6jzRoyWJB5qSWuzePJSh3z_5Ng6rDCnz7lg,15970
 dask_cuda/cuda_worker.py,sha256=bIu-ESeIpJG_WaTYrv0z9z5juJ1qR5i_5Ng3CN1WK8s,8579
 dask_cuda/device_host_file.py,sha256=yS31LGtt9VFAG78uBBlTDr7HGIng2XymV1OxXIuEMtM,10272
 dask_cuda/disk_io.py,sha256=urSLKiPvJvYmKCzDPOUDCYuLI3r1RUiyVh3UZGRoF_Y,6626
-dask_cuda/get_device_memory_objects.py,sha256=zMSqWzm5rflRInbNMz7U2Ewv5nMcE-H8stMJeWHVWyc,3890
+dask_cuda/get_device_memory_objects.py,sha256=R3U2cq4fJZPgtsUKyIguy9161p3Q99oxmcCmTcg6BtQ,4075
 dask_cuda/initialize.py,sha256=Gjcxs_c8DTafgsHe5-2mw4lJdOmbFJJAZVOnxA8lTjM,6462
 dask_cuda/is_device_object.py,sha256=CnajvbQiX0FzFzwft0MqK1OPomx3ZGDnDxT56wNjixw,1046
 dask_cuda/is_spillable_object.py,sha256=CddGmg0tuSpXh2m_TJSY6GRpnl1WRHt1CRcdWgHPzWA,1457
@@ -28,26 +28,26 @@ dask_cuda/benchmarks/utils.py,sha256=baL5zK6VS6Mw_M4x9zJe8vMLUd2SZd1lS78JrL-h6oo
 dask_cuda/explicit_comms/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 dask_cuda/explicit_comms/comms.py,sha256=Su6PuNo68IyS-AwoqU4S9TmqWsLvUdNa0jot2hx8jQQ,10400
 dask_cuda/explicit_comms/dataframe/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-dask_cuda/explicit_comms/dataframe/shuffle.py,sha256=2f2wlPyqXpryIHgMpsZzs3pDE7eyslYam-jQh3ujszQ,20124
+dask_cuda/explicit_comms/dataframe/shuffle.py,sha256=YferHNWKsMea8tele-ynPVr_6RAZNZIR-VzK_uFuEQU,20131
 dask_cuda/tests/test_cudf_builtin_spilling.py,sha256=u3kW91YRLdHFycvpGfSQKrEucu5khMJ1k4sjmddO490,4910
 dask_cuda/tests/test_dask_cuda_worker.py,sha256=gViHaMCSfB6ip125OEi9D0nfKC-qBXRoHz6BRodEdb4,17729
 dask_cuda/tests/test_device_host_file.py,sha256=79ssUISo1YhsW_7HdwqPfsH2LRzS2bi5BjPym1Sdgqw,5882
-dask_cuda/tests/test_dgx.py,sha256=IDP5vxDgVx6n2Hm-7PhlZQFb_zlgn_3nBW9t7MeJcTM,6986
-dask_cuda/tests/test_explicit_comms.py,sha256=AZlBi16prk3hc0ydzqAGecMcYZeyUvu6Pecb7gY_0yY,12315
+dask_cuda/tests/test_dgx.py,sha256=Oh2vwL_CdUzSVQQoiIu6SPwXGRtmXwaW_Hh3ipXPUOc,7162
+dask_cuda/tests/test_explicit_comms.py,sha256=I4lSW-NQ0E08baEoG7cY4Ix3blGb1Auz88q2BNd1cPA,13136
 dask_cuda/tests/test_from_array.py,sha256=okT1B6UqHmLxoy0uER0Ylm3UyOmi5BAXwJpTuTAw44I,601
 dask_cuda/tests/test_gds.py,sha256=6jf0HPTHAIG8Mp_FC4Ai4zpn-U1K7yk0fSXg8He8-r8,1513
 dask_cuda/tests/test_initialize.py,sha256=Rba59ZbljEm1yyN94_sWZPEE_f7hWln95aiBVc49pmY,6960
 dask_cuda/tests/test_local_cuda_cluster.py,sha256=G3kR-4o-vCqWWfSuQLFKVEK0F243FaDSgRlDTUll5aU,18376
-dask_cuda/tests/test_proxify_host_file.py,sha256=cp-U1uNPhesQaHbftKV8ir_dt5fbs0ZXSIsL39oI0fE,18630
+dask_cuda/tests/test_proxify_host_file.py,sha256=Yiv0sDcUoWw0d2oiPeHGoHqqSSM4lfQ4rChCiaxb6EU,18994
 dask_cuda/tests/test_proxy.py,sha256=6iicSYYT2BGo1iKUQ7jM00mCjC4gtfwwxFXfGwH3QHc,23807
 dask_cuda/tests/test_spill.py,sha256=xN9PbVERBYMuZxvscSO0mAM22loq9WT3ltZVBFxlmM4,10239
 dask_cuda/tests/test_utils.py,sha256=JRIwXfemc3lWSzLJX0VcvR1_0wB4yeoOTsw7kB6z6pU,9176
 dask_cuda/tests/test_worker_spec.py,sha256=Bvu85vkqm6ZDAYPXKMJlI2pm9Uc5tiYKNtO4goXSw-I,2399
 examples/ucx/client_initialize.py,sha256=YN3AXHF8btcMd6NicKKhKR9SXouAsK1foJhFspbOn70,1262
 examples/ucx/local_cuda_cluster.py,sha256=7xVY3EhwhkY2L4VZin_BiMCbrjhirDNChoC86KiETNc,1983
-dask_cuda-24.2.0a19.dist-info/LICENSE,sha256=MjI3I-EgxfEvZlgjk82rgiFsZqSDXHFETd2QJ89UwDA,11348
-dask_cuda-24.2.0a19.dist-info/METADATA,sha256=xSO6BsLRFAAXfzjN7sqhEU7xUyY4tALSEG_uwpVHXww,2537
-dask_cuda-24.2.0a19.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
-dask_cuda-24.2.0a19.dist-info/entry_points.txt,sha256=UcRaKVEpywtxc6pF1VnfMB0UK4sJg7a8_NdZF67laPM,136
-dask_cuda-24.2.0a19.dist-info/top_level.txt,sha256=3kKxJxeM108fuYc_lwwlklP7YBU9IEmdmRAouzi397o,33
-dask_cuda-24.2.0a19.dist-info/RECORD,,
+dask_cuda-24.4.0.dist-info/LICENSE,sha256=MjI3I-EgxfEvZlgjk82rgiFsZqSDXHFETd2QJ89UwDA,11348
+dask_cuda-24.4.0.dist-info/METADATA,sha256=mjy5YntsqIKdldcPEj-jjR70Aphe-B0LXqjcZjVFS9U,2570
+dask_cuda-24.4.0.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+dask_cuda-24.4.0.dist-info/entry_points.txt,sha256=UcRaKVEpywtxc6pF1VnfMB0UK4sJg7a8_NdZF67laPM,136
+dask_cuda-24.4.0.dist-info/top_level.txt,sha256=3kKxJxeM108fuYc_lwwlklP7YBU9IEmdmRAouzi397o,33
+dask_cuda-24.4.0.dist-info/RECORD,,

{dask_cuda-24.2.0a19.dist-info → dask_cuda-24.4.0.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: bdist_wheel (0.42.0)
+Generator: bdist_wheel (0.43.0)
 Root-Is-Purelib: true
 Tag: py3-none-any