Skip to content

Commit

Permalink
Add timeout to test_dask_use_explicit_comms
Browse files Browse the repository at this point in the history
Add timeout to `test_dask_use_explicit_comms` with SIGINT (i.e.,
KeyboardInterrupt) hoping that we can get a stacktrace that can help
identifying the cause of the test deadlock in CI.
  • Loading branch information
pentschev committed Dec 12, 2023
1 parent 1eecb1b commit a641e4e
Showing 1 changed file with 29 additions and 1 deletion.
30 changes: 29 additions & 1 deletion dask_cuda/tests/test_explicit_comms.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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'
Expand Down Expand Up @@ -217,6 +220,31 @@ def check_shuffle():
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")
Expand Down

0 comments on commit a641e4e

Please sign in to comment.