Merge pull request #388 from rapidsai/branch-0.43

Forward-merge branch-0.43 into branch-0.44

Author: GPUtester

python/ucxx/ucxx/_lib_async/tests/conftest.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2022-2023, NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES.
 # SPDX-License-Identifier: BSD-3-Clause
 
 import asyncio
@@ -70,12 +70,20 @@ def pytest_pyfunc_call(pyfuncitem: pytest.Function):
     `pytest.mark.rerun_on_failure(reruns)`. This is similar to `pytest-rerunfailures`,
     but that module closes the event loop before this function has awaited, making the
     two incompatible.
+
+    The timeout value is made available to the test functions via `pytestconfig`. This
+    can be used to determine internal timeouts, for example to ensure subprocesses
+    timeout before the test timeout hits and thus prints internal information, such as
+    the call stack. The timeout value may be retrieved by calling
+    `pytestconfig.cache.get("asyncio_timeout", {})["timeout"]`, for that the test must
+    include the `pytestconfig` fixture as argument.
     """
     timeout_marker = pyfuncitem.get_closest_marker("asyncio_timeout")
     slow_marker = pyfuncitem.get_closest_marker("slow")
     rerun_marker = pyfuncitem.get_closest_marker("rerun_on_failure")
     default_timeout = 600.0 if slow_marker else 60.0
     timeout = float(timeout_marker.args[0]) if timeout_marker else default_timeout
+    pyfuncitem.config.cache.set("asyncio_timeout", {"timeout": timeout})
     if timeout <= 0.0:
         raise ValueError("The `pytest.mark.asyncio_timeout` value must be positive.")
 

python/ucxx/ucxx/_lib_async/tests/test_disconnect.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2022-2023, NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES.
 # SPDX-License-Identifier: BSD-3-Clause
 
 import asyncio
@@ -12,7 +12,10 @@
 
 import ucxx
 from ucxx._lib_async.utils import get_event_loop
-from ucxx._lib_async.utils_test import wait_listener_client_handlers
+from ucxx._lib_async.utils_test import (
+    compute_timeouts,
+    wait_listener_client_handlers,
+)
 from ucxx.testing import terminate_process
 
 mp = mp.get_context("spawn")
@@ -28,7 +31,7 @@ async def mp_queue_get_nowait(queue):
 
 
 def _test_shutdown_unexpected_closed_peer_server(
-    client_queue, server_queue, endpoint_error_handling
+    client_queue, server_queue, endpoint_error_handling, timeout
 ):
     global ep_alive
     ep_alive = None
@@ -61,20 +64,25 @@ async def server_node(ep):
 
     log_stream = StringIO()
     logging.basicConfig(stream=log_stream, level=logging.DEBUG)
-    get_event_loop().run_until_complete(run())
-    log = log_stream.getvalue()
 
-    if endpoint_error_handling is True:
-        assert ep_alive is False
-    else:
-        assert ep_alive
-        assert log.find("""UCXError('<[Send shutdown]""") != -1
+    loop = get_event_loop()
+    try:
+        loop.run_until_complete(asyncio.wait_for(run(), timeout=timeout))
+        log = log_stream.getvalue()
 
-    ucxx.stop_notifier_thread()
+        if endpoint_error_handling is True:
+            assert ep_alive is False
+        else:
+            assert ep_alive
+            assert log.find("""UCXError('<[Send shutdown]""") != -1
+    finally:
+        ucxx.stop_notifier_thread()
+
+        loop.close()
 
 
 def _test_shutdown_unexpected_closed_peer_client(
-    client_queue, server_queue, endpoint_error_handling
+    client_queue, server_queue, endpoint_error_handling, timeout
 ):
     async def run():
         server_port = client_queue.get()
@@ -86,20 +94,25 @@ async def run():
         msg = np.empty(100, dtype=np.int64)
         await ep.recv(msg)
 
-    get_event_loop().run_until_complete(run())
+    loop = get_event_loop()
+    try:
+        loop.run_until_complete(asyncio.wait_for(run(), timeout=timeout))
+    finally:
+        ucxx.stop_notifier_thread()
 
-    ucxx.stop_notifier_thread()
+        loop.close()
 
 
 @pytest.mark.parametrize("endpoint_error_handling", [True, False])
-def test_shutdown_unexpected_closed_peer(caplog, endpoint_error_handling):
+def test_shutdown_unexpected_closed_peer(pytestconfig, caplog, endpoint_error_handling):
     """
     Test clean server shutdown after unexpected peer close
 
     This will causes some UCX warnings to be issued, but this as expected.
     The main goal is to assert that the processes exit without errors
     despite a somewhat messy initial state.
     """
+    async_timeout, join_timeout = compute_timeouts(pytestconfig)
     if endpoint_error_handling is False:
         pytest.xfail(
             "Temporarily xfailing, due to https://github.com/rapidsai/ucxx/issues/21"
@@ -120,17 +133,20 @@ def test_shutdown_unexpected_closed_peer(caplog, endpoint_error_handling):
     server_queue = mp.Queue()
     p1 = mp.Process(
         target=_test_shutdown_unexpected_closed_peer_server,
-        args=(client_queue, server_queue, endpoint_error_handling),
+        args=(client_queue, server_queue, endpoint_error_handling, async_timeout),
     )
     p1.start()
     p2 = mp.Process(
         target=_test_shutdown_unexpected_closed_peer_client,
-        args=(client_queue, server_queue, endpoint_error_handling),
+        args=(client_queue, server_queue, endpoint_error_handling, async_timeout),
     )
     p2.start()
-    p2.join(timeout=30)
+
+    # Increase timeout by an additional 5s to give subprocesses a chance to
+    # timeout before being forcefully terminated.
+    p2.join(timeout=join_timeout)
     server_queue.put("client is down")
-    p1.join(timeout=30)
+    p1.join(timeout=join_timeout)
 
     terminate_process(p2)
     terminate_process(p1)

python/ucxx/ucxx/_lib_async/tests/test_from_worker_address.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2022-2023, NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES.
 # SPDX-License-Identifier: BSD-3-Clause
 
 import asyncio
@@ -11,12 +11,13 @@
 
 import ucxx
 from ucxx._lib_async.utils import get_event_loop, hash64bits
+from ucxx._lib_async.utils_test import compute_timeouts
 from ucxx.testing import join_processes, terminate_process
 
 mp = mp.get_context("spawn")
 
 
-def _test_from_worker_address_server(queue):
+def _test_from_worker_address_server(queue, timeout):
     async def run():
         # Send worker address to client process via multiprocessing.Queue
         address = ucxx.get_worker_address()
@@ -40,14 +41,15 @@ async def run():
         await ep.close()
 
     loop = get_event_loop()
-    loop.run_until_complete(run())
+    try:
+        loop.run_until_complete(asyncio.wait_for(run(), timeout=timeout))
+    finally:
+        ucxx.stop_notifier_thread()
 
-    ucxx.stop_notifier_thread()
+        loop.close()
 
-    loop.close()
 
-
-def _test_from_worker_address_client(queue):
+def _test_from_worker_address_client(queue, timeout):
     async def run():
         # Read local worker address
         address = ucxx.get_worker_address()
@@ -69,29 +71,32 @@ async def run():
         np.testing.assert_array_equal(recv_msg, np.arange(10, dtype=np.int64))
 
     loop = get_event_loop()
-    loop.run_until_complete(run())
+    try:
+        loop.run_until_complete(asyncio.wait_for(run(), timeout=timeout))
+    finally:
+        ucxx.stop_notifier_thread()
 
-    ucxx.stop_notifier_thread()
+        loop.close()
 
-    loop.close()
 
+def test_from_worker_address(pytestconfig):
+    async_timeout, join_timeout = compute_timeouts(pytestconfig)
 
-def test_from_worker_address():
     queue = mp.Queue()
 
     server = mp.Process(
         target=_test_from_worker_address_server,
-        args=(queue,),
+        args=(queue, async_timeout),
     )
     server.start()
 
     client = mp.Process(
         target=_test_from_worker_address_client,
-        args=(queue,),
+        args=(queue, async_timeout),
     )
     client.start()
 
-    join_processes([client, server], timeout=30)
+    join_processes([client, server], timeout=join_timeout)
     terminate_process(client)
     terminate_process(server)
 
@@ -157,7 +162,7 @@ def _unpack_address_and_tag(address_packed):
     }
 
 
-def _test_from_worker_address_server_fixedsize(num_nodes, queue):
+def _test_from_worker_address_server_fixedsize(num_nodes, queue, timeout):
     async def run():
         async def _handle_client(packed_remote_address):
             # Unpack the fixed-size address+tag buffer
@@ -198,14 +203,15 @@ async def _handle_client(packed_remote_address):
         await asyncio.gather(*server_tasks)
 
     loop = get_event_loop()
-    loop.run_until_complete(run())
-
-    ucxx.stop_notifier_thread()
+    try:
+        loop.run_until_complete(asyncio.wait_for(run(), timeout=timeout))
+    finally:
+        ucxx.stop_notifier_thread()
 
-    loop.close()
+        loop.close()
 
 
-def _test_from_worker_address_client_fixedsize(queue):
+def _test_from_worker_address_client_fixedsize(queue, timeout):
     async def run():
         # Read local worker address
         address = ucxx.get_worker_address()
@@ -232,33 +238,37 @@ async def run():
         await ep.send(send_msg, tag=send_tag, force_tag=True)
 
     loop = get_event_loop()
-    loop.run_until_complete(run())
+    try:
+        loop.run_until_complete(asyncio.wait_for(run(), timeout=timeout))
+    finally:
+        ucxx.stop_notifier_thread()
 
-    ucxx.stop_notifier_thread()
-
-    loop.close()
+        loop.close()
 
 
+@pytest.mark.slow
 @pytest.mark.parametrize("num_nodes", [1, 2, 4, 8])
-def test_from_worker_address_multinode(num_nodes):
+def test_from_worker_address_multinode(pytestconfig, num_nodes):
+    async_timeout, join_timeout = compute_timeouts(pytestconfig)
+
     queue = mp.Queue()
 
     server = mp.Process(
         target=_test_from_worker_address_server_fixedsize,
-        args=(num_nodes, queue),
+        args=(num_nodes, queue, async_timeout),
     )
     server.start()
 
     clients = []
     for i in range(num_nodes):
         client = mp.Process(
             target=_test_from_worker_address_client_fixedsize,
-            args=(queue,),
+            args=(queue, async_timeout),
         )
         client.start()
         clients.append(client)
 
-    join_processes(clients + [server], timeout=30)
+    join_processes(clients + [server], timeout=join_timeout)
     for client in clients:
         terminate_process(client)
     terminate_process(server)

python/ucxx/ucxx/_lib_async/tests/test_from_worker_address_error.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2022-2023, NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES.
 # SPDX-License-Identifier: BSD-3-Clause
 
 import asyncio
@@ -12,12 +12,13 @@
 
 import ucxx
 from ucxx._lib_async.utils import get_event_loop
+from ucxx._lib_async.utils_test import compute_timeouts
 from ucxx.testing import join_processes, terminate_process
 
 mp = mp.get_context("spawn")
 
 
-def _test_from_worker_address_error_server(q1, q2, error_type):
+def _test_from_worker_address_error_server(q1, q2, error_type, timeout):
     async def run():
         address = bytearray(ucxx.get_worker_address())
 
@@ -39,14 +40,15 @@ async def run():
             # q1.put("disconnected")
 
     loop = get_event_loop()
-    loop.run_until_complete(run())
-
-    ucxx.stop_notifier_thread()
+    try:
+        loop.run_until_complete(asyncio.wait_for(run(), timeout=timeout))
+    finally:
+        ucxx.stop_notifier_thread()
 
-    loop.close()
+        loop.close()
 
 
-def _test_from_worker_address_error_client(q1, q2, error_type):
+def _test_from_worker_address_error_client(q1, q2, error_type, timeout):
     async def run():
         # Receive worker address from server via multiprocessing.Queue
         remote_address = ucxx.get_ucx_address_from_buffer(q1.get())
@@ -138,11 +140,12 @@ async def run():
                     await task
 
     loop = get_event_loop()
-    loop.run_until_complete(run())
-
-    ucxx.stop_notifier_thread()
+    try:
+        loop.run_until_complete(asyncio.wait_for(run(), timeout=timeout))
+    finally:
+        ucxx.stop_notifier_thread()
 
-    loop.close()
+        loop.close()
 
 
 @pytest.mark.parametrize(
@@ -164,27 +167,29 @@ async def run():
         "UCX_UD_TIMEOUT": "100ms",
     },
 )
-def test_from_worker_address_error(error_type):
+def test_from_worker_address_error(pytestconfig, error_type):
+    async_timeout, join_timeout = compute_timeouts(pytestconfig)
+
     q1 = mp.Queue()
     q2 = mp.Queue()
 
     server = mp.Process(
         target=_test_from_worker_address_error_server,
-        args=(q1, q2, error_type),
+        args=(q1, q2, error_type, async_timeout),
     )
     server.start()
 
     client = mp.Process(
         target=_test_from_worker_address_error_client,
-        args=(q1, q2, error_type),
+        args=(q1, q2, error_type, async_timeout),
     )
     client.start()
 
     if error_type == "unreachable":
         server.join()
         q1.put("Server closed")
 
-    join_processes([client, server], timeout=30)
+    join_processes([client, server], timeout=join_timeout)
     terminate_process(server)
     try:
         terminate_process(client)