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

dask_cuda/utils_test.py

@@ -1,14 +1,19 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+
 from typing import Literal
 
 import distributed
 from distributed import Nanny, Worker
 
+from .utils import _get_active_ucx_implementation_name
+
 
 class MockWorker(Worker):
     """Mock Worker class preventing NVML from getting used by SystemMonitor.
 
     By preventing the Worker from initializing NVML in the SystemMonitor, we can
-    mock test multiple devices in `CUDA_VISIBLE_DEVICES` behavior with single-GPU
+    mock test multiple devices in ``CUDA_VISIBLE_DEVICES`` behavior with single-GPU
     machines.
     """
 
@@ -26,20 +31,31 @@ class MockWorker(Worker):
 
 
 class IncreasedCloseTimeoutNanny(Nanny):
-    """Increase `Nanny`'s close timeout.
+    """Increase ``Nanny``'s close timeout.
 
-    The internal close timeout mechanism of `Nanny` recomputes the time left to kill
-    the `Worker` process based on elapsed time of the close task, which may leave
+    The internal close timeout mechanism of ``Nanny`` recomputes the time left to kill
+    the ``Worker`` process based on elapsed time of the close task, which may leave
     very little time for the subprocess to shutdown cleanly, which may cause tests
     to fail when the system is under higher load. This class increases the default
-    close timeout of 5.0 seconds that `Nanny` sets by default, which can be overriden
+    close timeout of 5.0 seconds that ``Nanny`` sets by default, which can be overriden
     via Distributed's public API.
 
-    This class can be used with the `worker_class` argument of `LocalCluster` or
-    `LocalCUDACluster` to provide a much higher default of 30.0 seconds.
+    This class can be used with the ``worker_class`` argument of ``LocalCluster`` or
+    ``LocalCUDACluster`` to provide a much higher default of 30.0 seconds.
     """
 
     async def close(  # type:ignore[override]
         self, timeout: float = 30.0, reason: str = "nanny-close"
     ) -> Literal["OK"]:
         return await super().close(timeout=timeout, reason=reason)
+
+
+def get_ucx_implementation(protocol):
+    import pytest
+
+    protocol = _get_active_ucx_implementation_name(protocol)
+
+    if protocol == "ucxx":
+        return pytest.importorskip("ucxx")
+    else:
+        return pytest.importorskip("ucp")

dask_cuda/worker_common.py

@@ -0,0 +1,196 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+import warnings
+
+from .device_host_file import DeviceHostFile
+from .plugins import CPUAffinity, CUDFSetup, PreImport, RMMSetup
+from .proxify_host_file import ProxifyHostFile
+from .utils import (
+    get_cpu_affinity,
+    has_device_memory_resource,
+    parse_device_memory_limit,
+)
+
+
+def worker_data_function(
+    device_memory_limit=None,
+    memory_limit=None,
+    jit_unspill=False,
+    enable_cudf_spill=False,
+    shared_filesystem=None,
+):
+    """
+    Create a data function for CUDA workers based on memory configuration.
+
+    This function creates and returns a callable that generates data configuration
+    for CUDA workers. The returned callable takes a device index parameter and
+    returns the appropriate data configuration for that device.
+
+    Parameters
+    ----------
+    device_memory_limit : str or int, optional
+        Limit of device memory, defaults to None
+    memory_limit : str or int, optional
+        Limit of host memory, defaults to None
+    jit_unspill : bool, optional
+        Whether to enable JIT unspill functionality, defaults to False
+    enable_cudf_spill : bool, optional
+        Whether to enable cuDF spilling, defaults to False
+    shared_filesystem : str or bool, optional
+        Whether to use shared filesystem for spilling, defaults to None
+
+    Returns
+    -------
+    callable
+        A function that takes device index `device_index` and returns appropriate
+        data configuration based on the availability of an dedicated device memory
+        resource and arguments passed to the worker.
+    """
+
+    def data(device_index):
+        if int(os.environ.get("DASK_CUDA_TEST_DISABLE_DEVICE_SPECIFIC", "0")) != 0:
+            return {}
+
+        # First raise errors for invalid configurations
+        if not has_device_memory_resource(device_index):
+            if jit_unspill:
+                raise ValueError(
+                    "JIT-Unspill is not supported on devices without dedicated memory, "
+                    "such as system on a chip (SoC) devices."
+                )
+            elif enable_cudf_spill:
+                raise ValueError(
+                    "cuDF spilling is not supported on devices without dedicated "
+                    "memory, such as system on a chip (SoC) devices."
+                )
+            elif device_memory_limit not in [None, "default"]:
+                raise ValueError(
+                    "device_memory_limit is set but device has no dedicated memory."
+                )
+
+        if device_memory_limit is None and memory_limit is None:
+            # All spilling is disabled
+            return {}
+        elif not has_device_memory_resource(device_index):
+            if device_memory_limit == "default" and memory_limit is None:
+                # Devices without a dedicated memory resource only support default
+                # host<->disk spilling via Dask's default mechanism.
+                return {}
+            # Devices without a dedicated memory resource only support default
+            # host<->disk spilling via Dask's default mechanism.
+            return None
+        else:
+            if jit_unspill:
+                # JIT-Unspill is enabled
+                if enable_cudf_spill:
+                    warnings.warn(
+                        "Enabling cuDF spilling and JIT-Unspill together is not "
+                        "safe, consider disabling JIT-Unspill."
+                    )
+
+                return (
+                    ProxifyHostFile,
+                    {
+                        "device_memory_limit": parse_device_memory_limit(
+                            device_memory_limit, device_index=device_index
+                        ),
+                        "memory_limit": memory_limit,
+                        "shared_filesystem": shared_filesystem,
+                    },
+                )
+            else:
+                # Device has dedicated memory and host memory is limited
+                return (
+                    DeviceHostFile,
+                    {
+                        "device_memory_limit": parse_device_memory_limit(
+                            device_memory_limit, device_index=device_index
+                        ),
+                        "memory_limit": memory_limit,
+                    },
+                )
+
+    return data
+
+
+def worker_plugins(
+    *,
+    device_index,
+    rmm_initial_pool_size,
+    rmm_maximum_pool_size,
+    rmm_managed_memory,
+    rmm_async_alloc,
+    rmm_release_threshold,
+    rmm_log_directory,
+    rmm_track_allocations,
+    rmm_allocator_external_lib_list,
+    pre_import,
+    enable_cudf_spill,
+    cudf_spill_stats,
+):
+    """Create a set of plugins for CUDA workers with specified configurations.
+
+    This function creates and returns a set of plugins that configure various aspects
+    of CUDA worker behavior, including CPU affinity, RMM memory management, pre-import
+    modules and cuDF spilling functionality.
+
+    Parameters
+    ----------
+    device_index : int
+        The CUDA device index to configure
+    rmm_initial_pool_size : int or str
+        Initial size of the RMM memory pool
+    rmm_maximum_pool_size : int or str
+        Maximum size of the RMM memory pool
+    rmm_managed_memory : bool
+        Whether to use CUDA managed memory
+    rmm_async_alloc : bool
+        Whether to use asynchronous allocation
+    rmm_release_threshold : int
+        Memory threshold for releasing memory back to the system
+    rmm_log_directory : str
+        Directory for RMM logging
+    rmm_track_allocations : bool
+        Whether to track memory allocations
+    rmm_allocator_external_lib_list : list
+        List of external libraries to use with RMM allocator
+    pre_import : list
+        List of modules to pre-import
+    enable_cudf_spill : bool
+        Whether to enable cuDF spilling
+    cudf_spill_stats : bool
+        Whether to track cuDF spilling statistics
+
+    Returns
+    -------
+    set
+        A set of configured plugins including:
+        - CPUAffinity: Configures CPU affinity for the worker
+        - RMMSetup: Configures RMM memory management
+        - PreImport: Handles module pre-importing
+        - CUDFSetup: Configures cuDF functionality and spilling
+    """
+    if int(os.environ.get("DASK_CUDA_TEST_DISABLE_DEVICE_SPECIFIC", "0")) != 0:
+        return {
+            PreImport(pre_import),
+            CUDFSetup(spill=enable_cudf_spill, spill_stats=cudf_spill_stats),
+        }
+    return {
+        CPUAffinity(
+            get_cpu_affinity(device_index),
+        ),
+        RMMSetup(
+            initial_pool_size=rmm_initial_pool_size,
+            maximum_pool_size=rmm_maximum_pool_size,
+            managed_memory=rmm_managed_memory,
+            async_alloc=rmm_async_alloc,
+            release_threshold=rmm_release_threshold,
+            log_directory=rmm_log_directory,
+            track_allocations=rmm_track_allocations,
+            external_lib_list=rmm_allocator_external_lib_list,
+        ),
+        PreImport(pre_import),
+        CUDFSetup(spill=enable_cudf_spill, spill_stats=cudf_spill_stats),
+    }

dask_cuda/worker_spec.py

@@ -1,3 +1,6 @@
+# SPDX-FileCopyrightText: Copyright (c) 2019-2025, NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+
 import os
 
 from dask.distributed import Nanny
@@ -6,7 +9,7 @@ from distributed.system import MEMORY_LIMIT
 from .initialize import initialize
 from .local_cuda_cluster import cuda_visible_devices
 from .plugins import CPUAffinity
-from .utils import get_cpu_affinity, get_gpu_count
+from .utils import _get_active_ucx_implementation_name, get_cpu_affinity, get_gpu_count
 
 
 def worker_spec(
@@ -81,10 +84,14 @@ def worker_spec(
        'preload_argv': ['--create-cuda-context']}}}
 
     """
-    if (
-        enable_tcp_over_ucx or enable_infiniband or enable_nvlink
-    ) and protocol != "ucx":
-        raise TypeError("Enabling InfiniBand or NVLink requires protocol='ucx'")
+    if enable_tcp_over_ucx or enable_infiniband or enable_nvlink:
+        try:
+            _get_active_ucx_implementation_name(protocol)
+        except ValueError:
+            raise TypeError(
+                "Enabling InfiniBand or NVLink requires protocol='ucx', "
+                "protocol='ucxx' or protocol='ucx-old'"
+            ) from None
 
     if CUDA_VISIBLE_DEVICES is None:
         CUDA_VISIBLE_DEVICES = os.environ.get(

{dask_cuda-25.6.0.dist-info → dask_cuda-25.8.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dask-cuda
-Version: 25.6.0
+Version: 25.8.0
 Summary: Utilities for Dask and CUDA interactions
 Author: NVIDIA Corporation
 License: Apache-2.0
@@ -24,7 +24,7 @@ Requires-Dist: numba<0.62.0a0,>=0.59.1
 Requires-Dist: numpy<3.0a0,>=1.23
 Requires-Dist: pandas>=1.3
 Requires-Dist: pynvml<13.0.0a0,>=12.0.0
-Requires-Dist: rapids-dask-dependency==25.6.*
+Requires-Dist: rapids-dask-dependency==25.8.*
 Requires-Dist: zict>=2.0.0
 Provides-Extra: docs
 Requires-Dist: numpydoc>=1.1.0; extra == "docs"

dask_cuda-25.8.0.dist-info/RECORD

@@ -0,0 +1,63 @@
+dask_cuda/GIT_COMMIT,sha256=RAxzl5uQmWcKe5o7s9iHwueNfKc91R_-kniXtpRqmng,41
+dask_cuda/VERSION,sha256=mZz9G1Ul4kEOksaMu07UE-AVtGzT1t91nQu3CK9KUAk,8
+dask_cuda/__init__.py,sha256=Wbc7R0voN4vsQkb7SKuVXH0YXuXtfnAxrupxfM4lT10,1933
+dask_cuda/_compat.py,sha256=AG2lKGAtZitDPBjHeFDKLTN_B5HKodrhZ2kHlk1Z-D0,498
+dask_cuda/_version.py,sha256=cHDO9AzNtxkCVhwYu7hL3H7RPAkQnxpKBjElOst3rkI,964
+dask_cuda/cli.py,sha256=VRYuryhViVWkCH7H7fDDTMToSOC17nAUMIPbd3K2jRs,17490
+dask_cuda/cuda_worker.py,sha256=7ZLZ3AY3l1fLumx2XynUOej5Sx6bwZQRomT4Nj9pbyA,8590
+dask_cuda/device_host_file.py,sha256=wTiTyYthUrR8l2WM7iV5lvjQDzeYxnnK7GUlPsHI6p4,11042
+dask_cuda/disk_io.py,sha256=IpD2hA-AjTw3SEu3w9pT9ELAUpKW0XcphPh8WwlKH70,6757
+dask_cuda/get_device_memory_objects.py,sha256=Nk0f5kv2f0e6JSd5hrwenAOARgF4NS21Zv_x3zF3ONI,4577
+dask_cuda/initialize.py,sha256=JBgXNzt52OzcusC89k0eVn2ivbCoCzZxOPNvQgj3YfE,7385
+dask_cuda/is_device_object.py,sha256=x9klFdeQzLcug7wZMxN3GK2AS121tlDe-LQ2uznm5yo,1179
+dask_cuda/is_spillable_object.py,sha256=8gj6QgtKcmzrpQwy8rE-pS1R8tjaJOeD-Fzr6LumjJg,1596
+dask_cuda/local_cuda_cluster.py,sha256=pocJgHbs8h2z_hfChU2_s7kwYKuYTgFZtmrEgYHjWwc,20735
+dask_cuda/plugins.py,sha256=u4gWQy8DgaTAsd59KCNXLG6zmdmGg5qhaI8ha2rMEFs,7085
+dask_cuda/proxify_device_objects.py,sha256=hGHK2gBkuaMYyrZqcitRGUiH_up56R1hsC7b_6YcCBU,8325
+dask_cuda/proxify_host_file.py,sha256=TbdszPvAgV4CRPMz99tumnCWwiCTqMBZqHQua3RRpHE,31031
+dask_cuda/proxy_object.py,sha256=eII-S0vdpFa6NLebpHFVQuJ2eyXJ5cBg8sucgFEa91g,30307
+dask_cuda/utils.py,sha256=smUv6DGJlHgXjfASVJMlRL1NgECiwwQ2zgn_FCWxMQ4,33454
+dask_cuda/utils_test.py,sha256=CKRMB5KUAg-7VSf21AU1TQoqGpJXm9ftgRV7mGIGQ3s,2132
+dask_cuda/worker_common.py,sha256=uXoYZ1IZanAbHddfpL3NbVDV5WqCF0m94nOiqqvEnxc,7168
+dask_cuda/worker_spec.py,sha256=cI4vS08gyrIU3PKJIjjWZNXChUm5Pv9LwaShPqYYMUQ,4698
+dask_cuda/benchmarks/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+dask_cuda/benchmarks/common.py,sha256=4AHgWkomo1RsPwH6eogz4vbE9vg_Dg-krScn9B1BRuw,7057
+dask_cuda/benchmarks/local_cudf_groupby.py,sha256=ehKOJqnHgQoEyNWuebIWdJP8f_bRiwBd4pax3hkQP_8,8646
+dask_cuda/benchmarks/local_cudf_merge.py,sha256=D2yXcEj1dLvfdd-ugV6kLgcUoiHzLmN3jNNs_BjaWME,12341
+dask_cuda/benchmarks/local_cudf_shuffle.py,sha256=MxUHzif_7TFPj96yE-Y3-U_NMXD6VJLNB1fB39LXgrk,8386
+dask_cuda/benchmarks/local_cupy.py,sha256=1riE9_hVkDxfSP2pWU0h2VF9HTXhF3C2ArKb19KMnVo,10457
+dask_cuda/benchmarks/local_cupy_map_overlap.py,sha256=27OtGOhmHIQhHSAmyOm0GgUX8u10npmwBo_q5fLAqUM,6151
+dask_cuda/benchmarks/read_parquet.py,sha256=spKu6RLWYngPZq9hnaoU0mz7INIaJnErfqjBG2wH8Zc,7614
+dask_cuda/benchmarks/utils.py,sha256=pl0VG8BQnIrVvpdVQtP8ybX1SOcv1E01jZzhkfOIPw4,30235
+dask_cuda/explicit_comms/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+dask_cuda/explicit_comms/comms.py,sha256=FbSDPQUz0gwQNa8EQMNEGEwWCM2KhuDk4CiNNjZy6qo,11611
+dask_cuda/explicit_comms/dataframe/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+dask_cuda/explicit_comms/dataframe/shuffle.py,sha256=vq-NBPYtKvKcfo7-eikPYaeWDf-V_DTHR5PVE_E-dyE,23863
+dask_cuda/tests/conftest.py,sha256=NjTnsgu7UewTpoTakt-2bpe0rtarwni25anjZPWbbb0,1501
+dask_cuda/tests/test_cudf_builtin_spilling.py,sha256=qVN9J0Hdv66A9COFArLIdRriyyxEKpS3lEZGHbVHaq8,4903
+dask_cuda/tests/test_dask_cuda_worker.py,sha256=iOkYAGSYfql7Wxxnb1hZJ7oedQyLmweXAHpQhckaTgY,23232
+dask_cuda/tests/test_device_host_file.py,sha256=79ssUISo1YhsW_7HdwqPfsH2LRzS2bi5BjPym1Sdgqw,5882
+dask_cuda/tests/test_dgx.py,sha256=Y7OgFlrq_BzoF4L8qoA5byKX1kDA1foaFzT-gEfhrsI,7294
+dask_cuda/tests/test_explicit_comms.py,sha256=ULrVhOcVavEWyZBxVSXDOFT_nznF8mRrU1M3fq_DUjo,20622
+dask_cuda/tests/test_from_array.py,sha256=-mXVwrw3MjB32vFfdUvxiwojqu9mEYNyjZjGPg1h1sE,733
+dask_cuda/tests/test_gds.py,sha256=j1Huud6UGm1fbkyRLQEz_ysrVw__5AimwSn_M-2GEvs,1513
+dask_cuda/tests/test_initialize.py,sha256=2Fm84sAN0A9Kxhm7PiWgrwppNCaJLA2yj65Y3LPOrNw,7542
+dask_cuda/tests/test_local_cuda_cluster.py,sha256=wYhY4OzQcYDtevL2peT6k5bAmZs2MDpQMrwPs2To7AY,23083
+dask_cuda/tests/test_proxify_host_file.py,sha256=6YPQia-V5tcVOy2ZNwL9ZpD6qKxKCcaNZZK-8f4aqpM,19996
+dask_cuda/tests/test_proxy.py,sha256=3uREA-d4y0lmGHKGJ5z9OiAYTBie1Y3bpB8djQRrluE,23996
+dask_cuda/tests/test_rdd_ucx.py,sha256=nMbEm7kxAATHfVyZOkfpexyk-ZKtCbLzvPHdspK877M,5433
+dask_cuda/tests/test_spill.py,sha256=QmBdFBvmZsDbOo0IAeUAJQfMyx-BBxokosch4gZHa00,15704
+dask_cuda/tests/test_utils.py,sha256=bTpiNQyKTxLHzHTLn-G0SWVgetq_tqUGq7rbafvdwgg,12297
+dask_cuda/tests/test_version.py,sha256=vK2HjlRLX0nxwvRsYxBqhoZryBNZklzA-vdnyuWDxVg,365
+dask_cuda/tests/test_worker_spec.py,sha256=dlZ4OIOl7CVm4euMtvCUif1QetCWcntSRdwzAmkH6ec,2550
+dask_cuda-25.8.0.dist-info/licenses/LICENSE,sha256=MjI3I-EgxfEvZlgjk82rgiFsZqSDXHFETd2QJ89UwDA,11348
+examples/ucx/client_initialize.py,sha256=YN3AXHF8btcMd6NicKKhKR9SXouAsK1foJhFspbOn70,1262
+examples/ucx/local_cuda_cluster.py,sha256=7xVY3EhwhkY2L4VZin_BiMCbrjhirDNChoC86KiETNc,1983
+shared-actions/check_nightly_success/check-nightly-success/check.py,sha256=lBhwgJALfUXdk4B9IhYf1AV0OUu0dQol3GtUe-CnUgY,5379
+shared-actions/telemetry-impls/summarize/bump_time.py,sha256=FXcGDqjCabsL6vlqd5RIV2PCWi2ns_ju8Ul0ERM21cA,2033
+shared-actions/telemetry-impls/summarize/send_trace.py,sha256=i30O_cKZ1OtXIbqaTE3R1JLJSw2XoYdmoVpM2Nc9wj8,16555
+dask_cuda-25.8.0.dist-info/METADATA,sha256=euY_vTYPba5yD4W8FtKDhCAjF4qs0ZU6-z5cs2lxYQg,2345
+dask_cuda-25.8.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+dask_cuda-25.8.0.dist-info/entry_points.txt,sha256=UcRaKVEpywtxc6pF1VnfMB0UK4sJg7a8_NdZF67laPM,136
+dask_cuda-25.8.0.dist-info/top_level.txt,sha256=SaQxjNeXM9ZhEBjU6CY25OIOt2UojSREGiVxg9UIc08,63
+dask_cuda-25.8.0.dist-info/RECORD,,

dask_cuda-25.8.0.dist-info/top_level.txt

@@ -0,0 +1,6 @@
+ci
+conda
+dask_cuda
+examples
+shared-actions
+telemetry-artifacts

shared-actions/check_nightly_success/check-nightly-success/check.py

@@ -0,0 +1,148 @@
+# Copyright (c) 2024-2025, NVIDIA CORPORATION.
+
+"""Check whether a GHA workflow has run successfully in the last N days."""
+# ruff: noqa: INP001
+
+import argparse
+import os
+import re
+import sys
+from collections import defaultdict
+from datetime import datetime
+
+import requests
+
+# Constants
+GITHUB_TOKEN = os.environ["RAPIDS_GH_TOKEN"]
+GOOD_STATUSES = {"success"}
+
+
+def main(
+    repo: str,
+    repo_owner: str,
+    workflow_id: str,
+    max_days_without_success: int,
+    num_attempts: int = 5,
+) -> bool:
+    """Check whether a GHA workflow has run successfully in the last N days.
+
+    Returns True if the workflow has not run successfully in the last N days, False
+    otherwise (values are inverted for use as a return code).
+    """
+    headers = {"Authorization": f"token {GITHUB_TOKEN}"}
+    url = f"https://api.github.com/repos/{repo_owner}/{repo}/actions/workflows/{workflow_id}/runs"
+    exceptions = []
+    for _ in range(num_attempts):
+        try:
+            response = requests.get(url, headers=headers, timeout=10)
+            response.raise_for_status()
+            break
+        except requests.RequestException as e:
+            exceptions.append(e)
+    else:
+        sep = "\n\t"
+        msg = (
+            f"Failed to fetch {url} after {num_attempts} attempts with the following "
+            f"errors: {sep}{'{sep}'.join(exceptions)}"
+        )
+        raise RuntimeError(msg)
+
+    runs = response.json()["workflow_runs"]
+    tz = datetime.fromisoformat(runs[0]["run_started_at"]).tzinfo
+    now = datetime.now(tz=tz)
+
+    latest_success = {}
+    workflow_active_for_max_days = {}
+    # Rather frustratingly, the workflow runs returned from the GitHub API can
+    # have alternating ordering of `head_branch`
+    # e.g.
+    #   run[0]['head_branch'] == "branch-25.02"
+    #   run[1]['head_branch'] == "branch-25.04"
+    #   run[2]['head_branch'] == "branch-25.02"
+    #
+    # In this situation, the behavior of `itertools.groupby` (previously used
+    # here) is to only group _consecutive_ runs, so the results of the
+    # subsequent branch match (i.e.  the second group of `branch-25.02` runs)
+    # will overwrite the results of the first one, potentially overwriting a
+    # previous success. The snippet below unifies the groups so it's more like a
+    # SQL groupby and there is no chance of overwriting.
+    branch_dict = defaultdict(list)
+    for run in runs:
+        branch_dict[run["head_branch"]].append(run)
+
+    for branch, branch_runs in branch_dict.items():
+        # Only consider RAPIDS release branches, which have versions like
+        # '25.02' (RAPIDS) or '0.42' (ucxx, ucx-py).
+        if not re.match("branch-[0-9]{1,2}.[0-9]{2}", branch):
+            continue
+
+        latest_success[branch] = None
+        runs = sorted(branch_runs, key=lambda r: r["run_started_at"], reverse=True)
+        for run in runs:
+            days_since_run = (now - datetime.fromisoformat(run["run_started_at"])).days
+            if days_since_run > max_days_without_success:
+                break
+            if run["conclusion"] in GOOD_STATUSES:
+                latest_success[branch] = run
+                break
+
+        workflow_active_for_max_days[branch] = False
+        if len(runs) > 0:
+            run = runs[-1]
+            days_since_run = (now - datetime.fromisoformat(run["run_started_at"])).days
+            if days_since_run > max_days_without_success:
+                workflow_active_for_max_days[branch] = True
+
+    latest_branch = max(latest_success)
+    has_latest_success = latest_success[latest_branch] is not None
+
+    # We are producing Unix return codes so success/failure is inverted from the
+    # expected Python boolean values.
+    if has_latest_success:
+        print(  # noqa: T201
+            f"The most recent successful run of the {workflow_id} workflow on "
+            f"{latest_branch} was "
+            f"{datetime.fromisoformat(latest_success[latest_branch]['run_started_at'])}, "
+            f"which is within the last {max_days_without_success} days. View logs:"
+            f"\n  - {latest_success[latest_branch]['html_url']}"
+        )
+        return 0
+    elif not workflow_active_for_max_days[latest_branch]:
+        print(  # noqa: T201
+            f"The oldest run of the {workflow_id} workflow on {latest_branch} was less "
+            f"than {max_days_without_success} days ago. This exempts the workflow from "
+            "check-nightly-success because the workflow has not been running for very long."
+        )
+        return 0
+
+    print(  # noqa: T201
+        f"{latest_branch} has no successful runs of {workflow_id} in the last {max_days_without_success} days"
+    )
+    return 1
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("repo", type=str, help="Repository name")
+    parser.add_argument(
+        "--repo-owner",
+        default="rapidsai",
+        help="Repository organization/owner",
+    )
+    parser.add_argument("--workflow-id", default="test.yaml", help="Workflow ID")
+    parser.add_argument(
+        "--max-days-without-success",
+        type=int,
+        default=7,
+        help="Maximum number of days without a successful run",
+    )
+    args = parser.parse_args()
+
+    sys.exit(
+        main(
+            args.repo,
+            args.repo_owner,
+            args.workflow_id,
+            args.max_days_without_success,
+        ),
+    )

shared-actions/telemetry-impls/summarize/bump_time.py

@@ -0,0 +1,54 @@
+#!/usr/bin/env python
+# Copyright (c) 2024-2025, NVIDIA CORPORATION.
+
+# This script is meant to act on an 'all_jobs.json' file that comes from
+# the summarize job when debug info is enabled. Bumping the time makes
+# it easier to re-run the span-sending python script and check results
+# in either Jaeger or Grafana
+
+import datetime
+import json
+
+with open("all_jobs.json") as f:
+    jobs = json.load(f)
+
+
+def _parse_time(x: str) -> int:
+    return int(datetime.datetime.strptime(x, "%Y-%m-%dT%H:%M:%SZ").timestamp() * 1e9)
+
+
+start_time = _parse_time(jobs[0]["created_at"])
+needed_time = _parse_time(jobs[-3]["completed_at"]) - _parse_time(jobs[0]["created_at"])
+new_start_time = datetime.datetime.utcnow() - datetime.timedelta(minutes=60)
+
+for idx, job in enumerate(jobs):
+    if job["created_at"]:
+        job["created_at"] = (
+            new_start_time + datetime.timedelta(seconds=(_parse_time(job["created_at"]) - start_time) / 1e9)
+        ).strftime("%Y-%m-%dT%H:%M:%SZ")
+    if job["started_at"]:
+        job["started_at"] = (
+            new_start_time + datetime.timedelta(seconds=(_parse_time(job["started_at"]) - start_time) / 1e9)
+        ).strftime("%Y-%m-%dT%H:%M:%SZ")
+    if job["completed_at"]:
+        job["completed_at"] = (
+            new_start_time + datetime.timedelta(seconds=(_parse_time(job["completed_at"]) - start_time) / 1e9)
+        ).strftime("%Y-%m-%dT%H:%M:%SZ")
+    steps = []
+    for step in job["steps"]:
+        if step["started_at"]:
+            step["started_at"] = (
+                new_start_time + datetime.timedelta(seconds=(_parse_time(step["started_at"]) - start_time) / 1e9)
+            ).strftime("%Y-%m-%dT%H:%M:%SZ")
+        if step["completed_at"]:
+            step["completed_at"] = (
+                new_start_time + datetime.timedelta(seconds=(_parse_time(step["completed_at"]) - start_time) / 1e9)
+            ).strftime("%Y-%m-%dT%H:%M:%SZ")
+        steps.append(step)
+    job["steps"] = steps
+
+    jobs[idx] = job
+
+
+with open("all_jobs.json", "w") as f:
+    json.dump(jobs, f)