From 4369a2e2fcdc2aaca7f1d40e18265676c93b02b5 Mon Sep 17 00:00:00 2001 From: Thomas Grainger Date: Mon, 4 Apr 2022 10:30:57 +0100 Subject: [PATCH 01/12] simple find and replace for pytest.mark.asyncio/gen_test --- distributed/cli/tests/test_dask_spec.py | 1 - distributed/comm/tests/test_comms.py | 84 +++++++++---------- distributed/comm/tests/test_ucx.py | 30 +++---- distributed/comm/tests/test_ucx_config.py | 2 +- distributed/comm/tests/test_ws.py | 4 +- .../deploy/tests/test_adaptive_core.py | 14 ++-- distributed/deploy/tests/test_cluster.py | 8 +- distributed/deploy/tests/test_local.py | 12 +-- .../deploy/tests/test_slow_adaptive.py | 6 +- distributed/deploy/tests/test_spec_cluster.py | 2 +- distributed/deploy/tests/test_ssh.py | 28 +++---- distributed/http/tests/test_routing.py | 2 +- distributed/protocol/tests/test_serialize.py | 2 +- distributed/tests/test_batched.py | 22 ++--- distributed/tests/test_client.py | 12 +-- distributed/tests/test_core.py | 66 +++++++-------- distributed/tests/test_nanny.py | 2 +- distributed/tests/test_preload.py | 6 +- distributed/tests/test_scheduler.py | 2 +- distributed/tests/test_utils.py | 4 +- distributed/tests/test_utils_test.py | 8 +- distributed/tests/test_worker.py | 14 ++-- 22 files changed, 165 insertions(+), 166 deletions(-) diff --git a/distributed/cli/tests/test_dask_spec.py b/distributed/cli/tests/test_dask_spec.py index eb1a3716da9..d415df9f3a5 100644 --- a/distributed/cli/tests/test_dask_spec.py +++ b/distributed/cli/tests/test_dask_spec.py @@ -39,7 +39,6 @@ async def test_text(): @pytest.mark.skipif(COMPILED, reason="Fails with cythonized scheduler") -@pytest.mark.asyncio @gen_cluster(client=True, nthreads=[]) async def test_file(c, s, tmp_path): fn = str(tmp_path / "foo.yaml") diff --git a/distributed/comm/tests/test_comms.py b/distributed/comm/tests/test_comms.py index f9ecf5e0726..da1376feed6 100644 --- a/distributed/comm/tests/test_comms.py +++ b/distributed/comm/tests/test_comms.py @@ -210,7 +210,7 @@ def test_get_local_address_for(tcp): # -@pytest.mark.asyncio +@gen_test() async def test_tcp_listener_does_not_call_handler_on_handshake_error(tcp): handle_comm_called = False @@ -232,7 +232,7 @@ async def handle_comm(comm): await writer.wait_closed() -@pytest.mark.asyncio +@gen_test() async def test_tcp_specific(tcp): """ Test concrete TCP API. @@ -275,7 +275,7 @@ async def client_communicate(key, delay=0): assert set(l) == {1234} | set(range(N)) -@pytest.mark.asyncio +@gen_test() async def test_tls_specific(tcp): """ Test concrete TLS API. @@ -321,7 +321,7 @@ async def client_communicate(key, delay=0): assert set(l) == {1234} | set(range(N)) -@pytest.mark.asyncio +@gen_test() async def test_comm_failure_threading(tcp): """ When we fail to connect, make sure we don't make a lot @@ -443,17 +443,17 @@ def run(): return fut -@pytest.mark.asyncio +@gen_test() async def test_inproc_specific_same_thread(): await check_inproc_specific(run_coro) -@pytest.mark.asyncio +@gen_test() async def test_inproc_specific_different_threads(): await check_inproc_specific(run_coro_in_thread) -@pytest.mark.asyncio +@gen_test() async def test_inproc_continues_listening_after_handshake_error(): async def handle_comm(): pass @@ -470,7 +470,7 @@ async def handle_comm(): await comm.close() -@pytest.mark.asyncio +@gen_test() async def test_inproc_handshakes_concurrently(): async def handle_comm(): pass @@ -567,7 +567,7 @@ async def client_communicate(key, delay=0): @pytest.mark.gpu -@pytest.mark.asyncio +@gen_test() async def test_ucx_client_server(): pytest.importorskip("distributed.comm.ucx") ucp = pytest.importorskip("ucp") @@ -603,7 +603,7 @@ def checker(loc): return checker -@pytest.mark.asyncio +@gen_test() async def test_default_client_server_ipv4(tcp): # Default scheme is (currently) TCP await check_client_server("127.0.0.1", tcp_eq("127.0.0.1")) @@ -620,7 +620,7 @@ async def test_default_client_server_ipv4(tcp): @requires_ipv6 -@pytest.mark.asyncio +@gen_test() async def test_default_client_server_ipv6(tcp): await check_client_server("[::1]", tcp_eq("::1")) await check_client_server("[::1]:3211", tcp_eq("::1", 3211)) @@ -630,7 +630,7 @@ async def test_default_client_server_ipv6(tcp): ) -@pytest.mark.asyncio +@gen_test() async def test_tcp_client_server_ipv4(tcp): await check_client_server("tcp://127.0.0.1", tcp_eq("127.0.0.1")) await check_client_server("tcp://127.0.0.1:3221", tcp_eq("127.0.0.1", 3221)) @@ -645,7 +645,7 @@ async def test_tcp_client_server_ipv4(tcp): @requires_ipv6 -@pytest.mark.asyncio +@gen_test() async def test_tcp_client_server_ipv6(tcp): await check_client_server("tcp://[::1]", tcp_eq("::1")) await check_client_server("tcp://[::1]:3231", tcp_eq("::1", 3231)) @@ -655,7 +655,7 @@ async def test_tcp_client_server_ipv6(tcp): ) -@pytest.mark.asyncio +@gen_test() async def test_tls_client_server_ipv4(tcp): await check_client_server("tls://127.0.0.1", tls_eq("127.0.0.1"), **tls_kwargs) await check_client_server( @@ -667,12 +667,12 @@ async def test_tls_client_server_ipv4(tcp): @requires_ipv6 -@pytest.mark.asyncio +@gen_test() async def test_tls_client_server_ipv6(tcp): await check_client_server("tls://[::1]", tls_eq("::1"), **tls_kwargs) -@pytest.mark.asyncio +@gen_test() async def test_inproc_client_server(): await check_client_server("inproc://", inproc_check()) await check_client_server(inproc.new_address(), inproc_check()) @@ -683,7 +683,7 @@ async def test_inproc_client_server(): # -@pytest.mark.asyncio +@gen_test() async def test_tls_reject_certificate(tcp): cli_ctx = get_client_ssl_context() serv_ctx = get_server_ssl_context() @@ -755,17 +755,17 @@ async def handle_comm(comm): await comm.read() -@pytest.mark.asyncio +@gen_test() async def test_tcp_comm_closed_implicit(tcp): await check_comm_closed_implicit("tcp://127.0.0.1") -@pytest.mark.asyncio +@gen_test() async def test_tls_comm_closed_implicit(tcp): await check_comm_closed_implicit("tls://127.0.0.1", **tls_kwargs) -@pytest.mark.asyncio +@gen_test() async def test_inproc_comm_closed_implicit(): await check_comm_closed_implicit(inproc.new_address()) @@ -793,22 +793,22 @@ async def check_comm_closed_explicit(addr, listen_args={}, connect_args={}): await b.close() -@pytest.mark.asyncio +@gen_test() async def test_tcp_comm_closed_explicit(tcp): await check_comm_closed_explicit("tcp://127.0.0.1") -@pytest.mark.asyncio +@gen_test() async def test_tls_comm_closed_explicit(tcp): await check_comm_closed_explicit("tls://127.0.0.1", **tls_kwargs) -@pytest.mark.asyncio +@gen_test() async def test_inproc_comm_closed_explicit(): await check_comm_closed_explicit(inproc.new_address()) -@pytest.mark.asyncio +@gen_test() async def test_inproc_comm_closed_explicit_2(): listener_errors = [] @@ -858,7 +858,7 @@ async def handle_comm(comm): await comm.close() -@pytest.mark.asyncio +@gen_test() async def test_comm_closed_on_buffer_error(tcp): # Internal errors from comm.stream.write, such as # BufferError should lead to the stream being closed @@ -890,7 +890,7 @@ async def echo(comm): await comm.write(message) -@pytest.mark.asyncio +@gen_test() async def test_retry_connect(tcp, monkeypatch): async def echo(comm): message = await comm.read() @@ -925,7 +925,7 @@ class UnreliableBackend(tcp.TCPBackend): listener.stop() -@pytest.mark.asyncio +@gen_test() async def test_handshake_slow_comm(tcp, monkeypatch): class SlowComm(tcp.TCP): def __init__(self, *args, delay_in_comm=0.5, **kwargs): @@ -976,12 +976,12 @@ async def check_connect_timeout(addr): assert 1 >= dt >= 0.1 -@pytest.mark.asyncio +@gen_test() async def test_tcp_connect_timeout(tcp): await check_connect_timeout("tcp://127.0.0.1:44444") -@pytest.mark.asyncio +@gen_test() async def test_inproc_connect_timeout(): await check_connect_timeout(inproc.new_address()) @@ -1004,14 +1004,14 @@ async def handle_comm(comm): listener.stop() -@pytest.mark.asyncio +@gen_test() async def test_tcp_many_listeners(tcp): await check_many_listeners("tcp://127.0.0.1") await check_many_listeners("tcp://0.0.0.0") await check_many_listeners("tcp://") -@pytest.mark.asyncio +@gen_test() async def test_inproc_many_listeners(): await check_many_listeners("inproc://") @@ -1160,12 +1160,12 @@ def check_out(deserialize_flag, out_value): await check_connector_deserialize(addr, True, msg, partial(check_out, True)) -@pytest.mark.asyncio +@gen_test() async def test_tcp_deserialize(tcp): await check_deserialize("tcp://") -@pytest.mark.asyncio +@gen_test() async def test_inproc_deserialize(): await check_deserialize("inproc://") @@ -1203,12 +1203,12 @@ async def check_deserialize_roundtrip(addr): assert isinstance(got["ser"], Serialized) -@pytest.mark.asyncio +@gen_test() async def test_inproc_deserialize_roundtrip(): await check_deserialize_roundtrip("inproc://") -@pytest.mark.asyncio +@gen_test() async def test_tcp_deserialize_roundtrip(tcp): await check_deserialize_roundtrip("tcp://") @@ -1238,7 +1238,7 @@ async def handle_comm(comm): await comm.read() -@pytest.mark.asyncio +@gen_test() async def test_tcp_deserialize_eoferror(tcp): await check_deserialize_eoferror("tcp://") @@ -1261,7 +1261,7 @@ async def check_repr(a, b): assert b.peer_address in repr(b) -@pytest.mark.asyncio +@gen_test() async def test_tcp_repr(tcp): a, b = await get_tcp_comm_pair() assert a.local_address in repr(b) @@ -1269,7 +1269,7 @@ async def test_tcp_repr(tcp): await check_repr(a, b) -@pytest.mark.asyncio +@gen_test() async def test_tls_repr(tcp): a, b = await get_tls_comm_pair() assert a.local_address in repr(b) @@ -1277,7 +1277,7 @@ async def test_tls_repr(tcp): await check_repr(a, b) -@pytest.mark.asyncio +@gen_test() async def test_inproc_repr(): a, b = await get_inproc_comm_pair() assert a.local_address in repr(b) @@ -1292,19 +1292,19 @@ async def check_addresses(a, b): b.abort() -@pytest.mark.asyncio +@gen_test() async def test_tcp_adresses(tcp): a, b = await get_tcp_comm_pair() await check_addresses(a, b) -@pytest.mark.asyncio +@gen_test() async def test_tls_adresses(tcp): a, b = await get_tls_comm_pair() await check_addresses(a, b) -@pytest.mark.asyncio +@gen_test() async def test_inproc_adresses(): a, b = await get_inproc_comm_pair() await check_addresses(a, b) diff --git a/distributed/comm/tests/test_ucx.py b/distributed/comm/tests/test_ucx.py index 8e6e9cb4ffe..14350f0a338 100644 --- a/distributed/comm/tests/test_ucx.py +++ b/distributed/comm/tests/test_ucx.py @@ -61,7 +61,7 @@ async def handle_comm(comm): return (comm, serv_comm) -@pytest.mark.asyncio +@gen_test() async def test_ping_pong(): com, serv_com = await get_comm_pair() msg = {"op": "ping"} @@ -79,7 +79,7 @@ async def test_ping_pong(): await serv_com.close() -@pytest.mark.asyncio +@gen_test() async def test_comm_objs(): comm, serv_comm = await get_comm_pair() @@ -92,7 +92,7 @@ async def test_comm_objs(): assert comm.peer_address == serv_comm.local_address -@pytest.mark.asyncio +@gen_test() async def test_ucx_specific(): """ Test concrete UCX API. @@ -146,7 +146,7 @@ async def client_communicate(key, delay=0): listener.stop() -@pytest.mark.asyncio +@gen_test() async def test_ping_pong_data(): np = pytest.importorskip("numpy") @@ -169,7 +169,7 @@ async def test_ping_pong_data(): await serv_com.close() -@pytest.mark.asyncio +@gen_test() async def test_ucx_deserialize(): # Note we see this error on some systems with this test: # `socket.gaierror: [Errno -5] No address associated with hostname` @@ -179,7 +179,7 @@ async def test_ucx_deserialize(): await check_deserialize("tcp://") -@pytest.mark.asyncio +@gen_test() @pytest.mark.parametrize( "g", [ @@ -219,7 +219,7 @@ async def test_ping_pong_cudf(g): await serv_com.close() -@pytest.mark.asyncio +@gen_test() @pytest.mark.parametrize("shape", [(100,), (10, 10), (4947,)]) async def test_ping_pong_cupy(shape): cupy = pytest.importorskip("cupy") @@ -238,7 +238,7 @@ async def test_ping_pong_cupy(shape): @pytest.mark.slow -@pytest.mark.asyncio +@gen_test() @pytest.mark.parametrize("n", [int(1e9), int(2.5e9)]) async def test_large_cupy(n, cleanup): cupy = pytest.importorskip("cupy") @@ -256,7 +256,7 @@ async def test_large_cupy(n, cleanup): await serv_com.close() -@pytest.mark.asyncio +@gen_test() async def test_ping_pong_numba(): np = pytest.importorskip("numpy") numba = pytest.importorskip("numba") @@ -274,7 +274,7 @@ async def test_ping_pong_numba(): assert result["op"] == "ping" -@pytest.mark.asyncio +@gen_test() @pytest.mark.parametrize("processes", [True, False]) async def test_ucx_localcluster(processes, cleanup): async with LocalCluster( @@ -296,7 +296,7 @@ async def test_ucx_localcluster(processes, cleanup): @pytest.mark.slow -@pytest.mark.asyncio +@gen_test() async def test_stress(): da = pytest.importorskip("dask.array") @@ -321,7 +321,7 @@ async def test_stress(): await wait(x) -@pytest.mark.asyncio +@gen_test() async def test_simple(): async with LocalCluster(protocol="ucx", asynchronous=True) as cluster: async with Client(cluster, asynchronous=True) as client: @@ -329,7 +329,7 @@ async def test_simple(): assert await client.submit(lambda x: x + 1, 10) == 11 -@pytest.mark.asyncio +@gen_test() async def test_cuda_context(): with dask.config.set({"distributed.comm.ucx.create-cuda-context": True}): async with LocalCluster( @@ -343,7 +343,7 @@ async def test_cuda_context(): assert list(worker_cuda_context.values())[0] == 0 -@pytest.mark.asyncio +@gen_test() async def test_transpose(): da = pytest.importorskip("dask.array") @@ -356,7 +356,7 @@ async def test_transpose(): await y -@pytest.mark.asyncio +@gen_test() @pytest.mark.parametrize("port", [0, 1234]) async def test_ucx_protocol(cleanup, port): async with Scheduler(protocol="ucx", port=port, dashboard_address=":0") as s: diff --git a/distributed/comm/tests/test_ucx_config.py b/distributed/comm/tests/test_ucx_config.py index 7f41320c229..7049b192250 100644 --- a/distributed/comm/tests/test_ucx_config.py +++ b/distributed/comm/tests/test_ucx_config.py @@ -21,7 +21,7 @@ rmm = pytest.importorskip("rmm") -@pytest.mark.asyncio +@gen_test() async def test_ucx_config(cleanup): ucx = { "nvlink": True, diff --git a/distributed/comm/tests/test_ws.py b/distributed/comm/tests/test_ws.py index c75c1a80b7e..4324b170a2e 100644 --- a/distributed/comm/tests/test_ws.py +++ b/distributed/comm/tests/test_ws.py @@ -125,7 +125,7 @@ async def test_large_transfer_with_no_compression(): await c.scatter(np.random.random(1_500_000)) -@pytest.mark.asyncio +@gen_test() @pytest.mark.parametrize( "dashboard,protocol,security,port", [ @@ -157,7 +157,7 @@ async def test_http_and_comm_server(cleanup, dashboard, protocol, security, port assert result == 11 -@pytest.mark.asyncio +@gen_test() @pytest.mark.parametrize("protocol", ["ws://", "wss://"]) async def test_connection_made_with_extra_conn_args(cleanup, protocol): if protocol == "ws://": diff --git a/distributed/deploy/tests/test_adaptive_core.py b/distributed/deploy/tests/test_adaptive_core.py index 7de0ffe87cc..0c44200e0b6 100644 --- a/distributed/deploy/tests/test_adaptive_core.py +++ b/distributed/deploy/tests/test_adaptive_core.py @@ -25,7 +25,7 @@ async def scale_down(self, workers=()): collection.discard(w) -@pytest.mark.asyncio +@gen_test() async def test_safe_target(): adapt = MyAdaptive(minimum=1, maximum=4) assert await adapt.safe_target() == 1 @@ -33,7 +33,7 @@ async def test_safe_target(): assert await adapt.safe_target() == 4 -@pytest.mark.asyncio +@gen_test() async def test_scale_up(): adapt = MyAdaptive(minimum=1, maximum=4) await adapt.adapt() @@ -46,7 +46,7 @@ async def test_scale_up(): assert adapt.plan == {0, 1, 2, 3} -@pytest.mark.asyncio +@gen_test() async def test_scale_down(): adapt = MyAdaptive(minimum=1, maximum=4, wait_count=2) adapt._target = 10 @@ -72,7 +72,7 @@ async def test_scale_down(): assert list(adapt.log) == old -@pytest.mark.asyncio +@gen_test() async def test_interval(): adapt = MyAdaptive(interval="5 ms") assert not adapt.plan @@ -92,7 +92,7 @@ async def test_interval(): assert len(adapt.plan) == 1 # last value from before, unchanged -@pytest.mark.asyncio +@gen_test() async def test_adapt_oserror_safe_target(): class BadAdaptive(MyAdaptive): """AdaptiveCore subclass which raises an OSError when attempting to adapt @@ -113,7 +113,7 @@ def safe_target(self): assert not adapt.periodic_callback -@pytest.mark.asyncio +@gen_test() async def test_adapt_oserror_scale(): """ FIXME: @@ -147,7 +147,7 @@ async def scale_down(self, workers=None): adapt.stop() -@pytest.mark.asyncio +@gen_test() async def test_adapt_stop_del(): adapt = MyAdaptive(interval="100ms") pc = adapt.periodic_callback diff --git a/distributed/deploy/tests/test_cluster.py b/distributed/deploy/tests/test_cluster.py index 48d8bb65336..a5a14a2a8bc 100644 --- a/distributed/deploy/tests/test_cluster.py +++ b/distributed/deploy/tests/test_cluster.py @@ -3,7 +3,7 @@ from distributed.deploy.cluster import Cluster -@pytest.mark.asyncio +@gen_test() async def test_eq(cleanup): clusterA = Cluster(asynchronous=True, name="A") clusterA2 = Cluster(asynchronous=True, name="A2") @@ -18,7 +18,7 @@ async def test_eq(cleanup): assert not (clusterA == clusterB) -@pytest.mark.asyncio +@gen_test() async def test_repr(cleanup): cluster = Cluster(asynchronous=True, name="A") assert cluster.scheduler_address == "" @@ -27,14 +27,14 @@ async def test_repr(cleanup): assert res == expected -@pytest.mark.asyncio +@gen_test() async def test_logs_deprecated(cleanup): cluster = Cluster(asynchronous=True) with pytest.warns(FutureWarning, match="get_logs"): cluster.logs() -@pytest.mark.asyncio +@gen_test() async def test_cluster_info(): class FooCluster(Cluster): def __init__(self): diff --git a/distributed/deploy/tests/test_local.py b/distributed/deploy/tests/test_local.py index d6ef2298c8e..4da82ed4d25 100644 --- a/distributed/deploy/tests/test_local.py +++ b/distributed/deploy/tests/test_local.py @@ -931,7 +931,7 @@ async def test_scale_memory_cores(): assert len(cluster.worker_spec) == 4 -@pytest.mark.asyncio +@gen_test() @pytest.mark.parametrize("memory_limit", ["2 GiB", None]) async def test_repr(memory_limit, cleanup): async with LocalCluster( @@ -968,7 +968,7 @@ async def test_threads_per_worker_set_to_0(): assert all(w.nthreads < CPU_COUNT for w in cluster.workers.values()) -@pytest.mark.asyncio +@gen_test() @pytest.mark.parametrize("temporary", [True, False]) async def test_capture_security(cleanup, temporary): if temporary: @@ -1039,7 +1039,7 @@ async def test_cluster_names(): assert unnamed_cluster2 != unnamed_cluster -@pytest.mark.asyncio +@gen_test() @pytest.mark.parametrize("nanny", [True, False]) async def test_local_cluster_redundant_kwarg(nanny): with pytest.raises(TypeError, match="unexpected keyword argument"): @@ -1060,7 +1060,7 @@ async def test_local_cluster_redundant_kwarg(nanny): await f -@pytest.mark.asyncio +@gen_test() async def test_cluster_info_sync(): async with LocalCluster( processes=False, asynchronous=True, scheduler_sync_interval="1ms" @@ -1087,7 +1087,7 @@ async def test_cluster_info_sync(): assert info["foo"] == "bar" -@pytest.mark.asyncio +@gen_test() async def test_cluster_info_sync_is_robust_to_network_blips(monkeypatch): async with LocalCluster( processes=False, asynchronous=True, scheduler_sync_interval="1ms" @@ -1123,7 +1123,7 @@ async def error(*args, **kwargs): assert info["foo"] == "bar" -@pytest.mark.asyncio +@gen_test() @pytest.mark.parametrize("host", [None, "127.0.0.1"]) @pytest.mark.parametrize("use_nanny", [True, False]) async def test_cluster_host_used_throughout_cluster(host, use_nanny): diff --git a/distributed/deploy/tests/test_slow_adaptive.py b/distributed/deploy/tests/test_slow_adaptive.py index c6697d1fa75..a55e4151acc 100644 --- a/distributed/deploy/tests/test_slow_adaptive.py +++ b/distributed/deploy/tests/test_slow_adaptive.py @@ -35,7 +35,7 @@ async def close(self): scheduler = {"cls": Scheduler, "options": {"dashboard_address": ":0"}} -@pytest.mark.asyncio +@gen_test() async def test_startup(cleanup): start = time() async with SpecCluster( @@ -55,7 +55,7 @@ async def test_startup(cleanup): await client.wait_for_workers(n_workers=2) -@pytest.mark.asyncio +@gen_test() @pytest.mark.flaky(reruns=10, reruns_delay=5) async def test_scale_up_down(cleanup): start = time() @@ -77,7 +77,7 @@ async def test_scale_up_down(cleanup): assert not cluster.worker_spec -@pytest.mark.asyncio +@gen_test() async def test_adaptive(cleanup): start = time() async with SpecCluster( diff --git a/distributed/deploy/tests/test_spec_cluster.py b/distributed/deploy/tests/test_spec_cluster.py index fe566c94d90..4b6ab24c894 100644 --- a/distributed/deploy/tests/test_spec_cluster.py +++ b/distributed/deploy/tests/test_spec_cluster.py @@ -213,7 +213,7 @@ async def test_restart(): # FIXME cleanup fails: # some RPCs left active by test: {} # @gen_test() -@pytest.mark.asyncio +@gen_test() async def test_broken_worker(): with pytest.raises(Exception) as info: async with SpecCluster( diff --git a/distributed/deploy/tests/test_ssh.py b/distributed/deploy/tests/test_ssh.py index 61993573ec7..be03c316466 100644 --- a/distributed/deploy/tests/test_ssh.py +++ b/distributed/deploy/tests/test_ssh.py @@ -26,7 +26,7 @@ def test_ssh_hosts_empty_list(): SSHCluster(hosts=[]) -@pytest.mark.asyncio +@gen_test() async def test_ssh_cluster_raises_if_asyncssh_not_installed(monkeypatch, cleanup): monkeypatch.setitem(sys.modules, "asyncssh", None) with pytest.raises(ImportError, match="SSHCluster requires the `asyncssh` package"): @@ -40,7 +40,7 @@ async def test_ssh_cluster_raises_if_asyncssh_not_installed(monkeypatch, cleanup assert not cluster -@pytest.mark.asyncio +@gen_test() async def test_basic(): async with SSHCluster( ["127.0.0.1"] * 3, @@ -58,7 +58,7 @@ async def test_basic(): assert "SSH" in repr(cluster) -@pytest.mark.asyncio +@gen_test() async def test_n_workers(): async with SSHCluster( ["127.0.0.1"] * 3, @@ -77,7 +77,7 @@ async def test_n_workers(): assert "SSH" in repr(cluster) -@pytest.mark.asyncio +@gen_test() async def test_nprocs_attribute_is_deprecated(): async with SSHCluster( ["127.0.0.1"] * 2, @@ -97,7 +97,7 @@ async def test_nprocs_attribute_is_deprecated(): assert worker.n_workers == 3 -@pytest.mark.asyncio +@gen_test() async def test_ssh_nprocs_renamed_to_n_workers(): with pytest.warns(FutureWarning, match="renamed to n_workers"): async with SSHCluster( @@ -112,7 +112,7 @@ async def test_ssh_nprocs_renamed_to_n_workers(): await client.wait_for_workers(4) -@pytest.mark.asyncio +@gen_test() async def test_ssh_n_workers_with_nprocs_is_an_error(): with pytest.raises(ValueError, match="Both nprocs and n_workers"): async with SSHCluster( @@ -125,7 +125,7 @@ async def test_ssh_n_workers_with_nprocs_is_an_error(): assert not cluster -@pytest.mark.asyncio +@gen_test() async def test_keywords(): async with SSHCluster( ["127.0.0.1"] * 3, @@ -179,7 +179,7 @@ def test_old_ssh_with_local_dir(loop): assert result == 11 -@pytest.mark.asyncio +@gen_test() async def test_config_inherited_by_subprocess(loop): def f(x): return dask.config.get("foo") + 1 @@ -197,7 +197,7 @@ def f(x): assert result == 101 -@pytest.mark.asyncio +@gen_test() async def test_unimplemented_options(): with pytest.raises(Exception): async with SSHCluster( @@ -215,7 +215,7 @@ async def test_unimplemented_options(): assert cluster -@pytest.mark.asyncio +@gen_test() async def test_list_of_connect_options(): async with SSHCluster( ["127.0.0.1"] * 3, @@ -233,7 +233,7 @@ async def test_list_of_connect_options(): assert "SSH" in repr(cluster) -@pytest.mark.asyncio +@gen_test() async def test_list_of_connect_options_raises(): with pytest.raises(RuntimeError): async with SSHCluster( @@ -246,7 +246,7 @@ async def test_list_of_connect_options_raises(): pass -@pytest.mark.asyncio +@gen_test() async def test_remote_python(): async with SSHCluster( ["127.0.0.1"] * 3, @@ -259,7 +259,7 @@ async def test_remote_python(): assert cluster.workers[0].remote_python == sys.executable -@pytest.mark.asyncio +@gen_test() async def test_remote_python_as_dict(): async with SSHCluster( ["127.0.0.1"] * 3, @@ -272,7 +272,7 @@ async def test_remote_python_as_dict(): assert cluster.workers[0].remote_python == sys.executable -@pytest.mark.asyncio +@gen_test() async def test_list_of_remote_python_raises(): with pytest.raises(RuntimeError): async with SSHCluster( diff --git a/distributed/http/tests/test_routing.py b/distributed/http/tests/test_routing.py index 764475ce4ed..a6a2a381f1d 100644 --- a/distributed/http/tests/test_routing.py +++ b/distributed/http/tests/test_routing.py @@ -15,7 +15,7 @@ def get(self): self.write("two") -@pytest.mark.asyncio +@gen_test() async def test_basic(): application = RoutingApplication([(r"/one", OneHandler)]) two = web.Application([(r"/two", TwoHandler)]) diff --git a/distributed/protocol/tests/test_serialize.py b/distributed/protocol/tests/test_serialize.py index f144259ba02..8c4eb51c02b 100644 --- a/distributed/protocol/tests/test_serialize.py +++ b/distributed/protocol/tests/test_serialize.py @@ -415,7 +415,7 @@ def dumps(f): assert "Hello-123" in str(info.value) -@pytest.mark.asyncio +@gen_test() async def test_profile_nested_sizeof(): # https://github.com/dask/distributed/issues/1674 n = 500 diff --git a/distributed/tests/test_batched.py b/distributed/tests/test_batched.py index ee84ec3224e..4efc71de15b 100644 --- a/distributed/tests/test_batched.py +++ b/distributed/tests/test_batched.py @@ -37,7 +37,7 @@ async def __aexit__(self, exc, typ, tb): self.stop() -@pytest.mark.asyncio +@gen_test() async def test_BatchedSend(): async with EchoServer() as e: comm = await connect(e.address) @@ -64,7 +64,7 @@ async def test_BatchedSend(): assert b.byte_count > 1 -@pytest.mark.asyncio +@gen_test() async def test_send_before_start(): async with EchoServer() as e: comm = await connect(e.address) @@ -79,7 +79,7 @@ async def test_send_before_start(): assert result == ("hello", "world") -@pytest.mark.asyncio +@gen_test() async def test_send_after_stream_start(): async with EchoServer() as e: comm = await connect(e.address) @@ -95,7 +95,7 @@ async def test_send_after_stream_start(): assert result == ("hello", "world") -@pytest.mark.asyncio +@gen_test() async def test_send_before_close(): async with EchoServer() as e: comm = await connect(e.address) @@ -117,7 +117,7 @@ async def test_send_before_close(): b.send("123") -@pytest.mark.asyncio +@gen_test() async def test_close_closed(): async with EchoServer() as e: comm = await connect(e.address) @@ -133,13 +133,13 @@ async def test_close_closed(): assert "closed" in str(b) -@pytest.mark.asyncio +@gen_test() async def test_close_not_started(): b = BatchedSend(interval=10) await b.close() -@pytest.mark.asyncio +@gen_test() async def test_close_twice(): async with EchoServer() as e: comm = await connect(e.address) @@ -151,7 +151,7 @@ async def test_close_twice(): @pytest.mark.slow -@pytest.mark.asyncio +@gen_test() async def test_stress(): async with EchoServer() as e: comm = await connect(e.address) @@ -215,18 +215,18 @@ async def run_traffic_jam(nsends, nbytes): await b.close() -@pytest.mark.asyncio +@gen_test() async def test_sending_traffic_jam(): await run_traffic_jam(50, 300000) @pytest.mark.slow -@pytest.mark.asyncio +@gen_test() async def test_large_traffic_jam(): await run_traffic_jam(500, 1500000) -@pytest.mark.asyncio +@gen_test() async def test_serializers(): async with EchoServer() as e: comm = await connect(e.address) diff --git a/distributed/tests/test_client.py b/distributed/tests/test_client.py index 5044c487bba..164d5e88121 100644 --- a/distributed/tests/test_client.py +++ b/distributed/tests/test_client.py @@ -6127,7 +6127,7 @@ async def test_wait_for_workers(c, s, a, b): @pytest.mark.skipif(WINDOWS, reason="num_fds not supported on windows") -@pytest.mark.asyncio +@gen_test() @pytest.mark.parametrize("Worker", [Worker, Nanny]) async def test_file_descriptors_dont_leak(Worker): pytest.importorskip("pandas") @@ -6173,7 +6173,7 @@ async def test_shutdown(): assert w.status == Status.closed -@pytest.mark.asyncio +@gen_test() async def test_shutdown_localcluster(cleanup): async with LocalCluster( n_workers=1, asynchronous=True, processes=False, dashboard_address=":0" @@ -7402,7 +7402,7 @@ def config_loader(self, monkeypatch, loader): ): yield - @pytest.mark.asyncio + @gen_test() async def test_security_loader(self, monkeypatch): security = tls_only_security() @@ -7418,7 +7418,7 @@ def loader(info): async with Client(scheduler.address, asynchronous=True) as client: assert client.security is security - @pytest.mark.asyncio + @gen_test() async def test_security_loader_ignored_if_explicit_security_provided( self, monkeypatch ): @@ -7436,7 +7436,7 @@ def loader(info): ) as client: assert client.security is security - @pytest.mark.asyncio + @gen_test() async def test_security_loader_ignored_if_returns_none(self, monkeypatch): """Test that if a security loader is configured, but it returns `None`, then the default security configuration is used""" @@ -7469,7 +7469,7 @@ def loader(info): assert loader.called - @pytest.mark.asyncio + @gen_test() async def test_security_loader_import_failed(self): security = tls_only_security() diff --git a/distributed/tests/test_core.py b/distributed/tests/test_core.py index 3598a00bbc9..90d43dee10c 100644 --- a/distributed/tests/test_core.py +++ b/distributed/tests/test_core.py @@ -145,7 +145,7 @@ class MyServer(Server): default_port = 8756 -@pytest.mark.asyncio +@gen_test() async def test_server_listen(): """ Test various Server.listen() arguments and their effect. @@ -291,7 +291,7 @@ async def check_rpc(listen_addr, rpc_addr=None, listen_args={}, connection_args= await asyncio.sleep(0) -@pytest.mark.asyncio +@gen_test() async def test_rpc_default(): await check_rpc(8883, "127.0.0.1:8883") await check_rpc(8883) @@ -303,7 +303,7 @@ async def test_rpc_tcp(): await check_rpc("tcp://") -@pytest.mark.asyncio +@gen_test() async def test_rpc_tls(): sec = tls_security() await check_rpc( @@ -314,12 +314,12 @@ async def test_rpc_tls(): ) -@pytest.mark.asyncio +@gen_test() async def test_rpc_inproc(): await check_rpc("inproc://", None) -@pytest.mark.asyncio +@gen_test() async def test_rpc_inputs(): L = [rpc("127.0.0.1:8884"), rpc(("127.0.0.1", 8884)), rpc("tcp://127.0.0.1:8884")] @@ -363,17 +363,17 @@ async def check_rpc_message_lifetime(*listen_args): server.stop() -@pytest.mark.asyncio +@gen_test() async def test_rpc_message_lifetime_default(): await check_rpc_message_lifetime() -@pytest.mark.asyncio +@gen_test() async def test_rpc_message_lifetime_tcp(): await check_rpc_message_lifetime("tcp://") -@pytest.mark.asyncio +@gen_test() async def test_rpc_message_lifetime_inproc(): await check_rpc_message_lifetime("inproc://") @@ -396,12 +396,12 @@ async def g(): assert all(comm.closed() for comm in remote.comms) -@pytest.mark.asyncio +@gen_test() async def test_rpc_with_many_connections_tcp(): await check_rpc_with_many_connections("tcp://") -@pytest.mark.asyncio +@gen_test() async def test_rpc_with_many_connections_inproc(): await check_rpc_with_many_connections("inproc://") @@ -424,12 +424,12 @@ async def check_large_packets(listen_arg): @pytest.mark.slow -@pytest.mark.asyncio +@gen_test() async def test_large_packets_tcp(): await check_large_packets("tcp://") -@pytest.mark.asyncio +@gen_test() async def test_large_packets_inproc(): await check_large_packets("inproc://") @@ -447,17 +447,17 @@ async def check_identity(listen_arg): server.stop() -@pytest.mark.asyncio +@gen_test() async def test_identity_tcp(): await check_identity("tcp://") -@pytest.mark.asyncio +@gen_test() async def test_identity_inproc(): await check_identity("inproc://") -@pytest.mark.asyncio +@gen_test() async def test_ports(loop): for port in range(9877, 9887): server = Server({}, io_loop=loop) @@ -491,7 +491,7 @@ def stream_div(comm=None, x=None, y=None): return x / y -@pytest.mark.asyncio +@gen_test() async def test_errors(): server = Server({"div": stream_div}) await server.listen(0) @@ -501,13 +501,13 @@ async def test_errors(): await r.div(x=1, y=0) -@pytest.mark.asyncio +@gen_test() async def test_connect_raises(): with pytest.raises(IOError): await connect("127.0.0.1:58259", timeout=0.01) -@pytest.mark.asyncio +@gen_test() async def test_send_recv_args(): server = Server({}) await server.listen(0) @@ -553,7 +553,7 @@ def test_coerce_to_address(): assert coerce_to_address(arg) == "tcp://127.0.0.1:8786" -@pytest.mark.asyncio +@gen_test() async def test_connection_pool(): async def ping(comm, delay=0.1): await asyncio.sleep(delay) @@ -602,7 +602,7 @@ async def ping(comm, delay=0.1): await rpc.close() -@pytest.mark.asyncio +@gen_test() async def test_connection_pool_close_while_connecting(monkeypatch): """ Ensure a closed connection pool guarantees to have no connections left open @@ -646,7 +646,7 @@ async def connect_to_server(): assert not pool._n_connecting -@pytest.mark.asyncio +@gen_test() async def test_connection_pool_outside_cancellation(monkeypatch): # Ensure cancellation errors are properly reraised from distributed.comm.registry import backends @@ -684,7 +684,7 @@ async def connect_to_server(): assert all(t.cancelled() for t in tasks) -@pytest.mark.asyncio +@gen_test() async def test_connection_pool_respects_limit(): limit = 5 @@ -707,7 +707,7 @@ async def do_ping(pool, port): await asyncio.gather(*(do_ping(pool, s.port) for s in servers)) -@pytest.mark.asyncio +@gen_test() async def test_connection_pool_tls(): """ Make sure connection args are supported. @@ -734,7 +734,7 @@ async def ping(comm, delay=0.01): await rpc.close() -@pytest.mark.asyncio +@gen_test() async def test_connection_pool_remove(): async def ping(comm, delay=0.01): await asyncio.sleep(delay) @@ -773,7 +773,7 @@ async def ping(comm, delay=0.01): await rpc.close() -@pytest.mark.asyncio +@gen_test() async def test_counters(): server = Server({"div": stream_div}) await server.listen("tcp://") @@ -832,7 +832,7 @@ async def f(): loop.run_sync(f) -@pytest.mark.asyncio +@gen_test() async def test_rpc_serialization(): server = Server({"echo": echo_serialize}) await server.listen("tcp://") @@ -853,7 +853,7 @@ async def test_thread_id(s, a, b): assert s.thread_id == a.thread_id == b.thread_id == threading.get_ident() -@pytest.mark.asyncio +@gen_test() async def test_deserialize_error(): server = Server({"throws": throws}) await server.listen(0) @@ -867,7 +867,7 @@ async def test_deserialize_error(): assert c.isalpha() or c in "(',!)" # no crazy bytestrings -@pytest.mark.asyncio +@gen_test() async def test_connection_pool_detects_remote_close(): server = Server({"ping": pingpong}) await server.listen("tcp://") @@ -899,7 +899,7 @@ async def test_connection_pool_detects_remote_close(): await p.close() -@pytest.mark.asyncio +@gen_test() async def test_close_properly(): """ If the server is closed we should cancel all still ongoing coros and close @@ -947,13 +947,13 @@ async def sleep(comm=None): assert not len(server._ongoing_coroutines) -@pytest.mark.asyncio +@gen_test() async def test_server_redundant_kwarg(): with pytest.raises(TypeError, match="unexpected keyword argument"): await Server({}, typo_kwarg="foo") -@pytest.mark.asyncio +@gen_test() async def test_server_comms_mark_active_handlers(): """Whether handlers are active can be read off of the self._comms values. ensure this is properly reflected and released. The sentinel for @@ -980,7 +980,7 @@ async def long_handler(comm): await asyncio.sleep(0.01) -@pytest.mark.asyncio +@gen_test() async def test_close_fast_without_active_handlers(): async def very_fast(comm): return "done" @@ -998,7 +998,7 @@ async def very_fast(comm): await asyncio.wait_for(fut, 0.1) -@pytest.mark.asyncio +@gen_test() async def test_close_grace_period_for_handlers(): async def long_handler(comm, delay=10): await asyncio.sleep(delay) diff --git a/distributed/tests/test_nanny.py b/distributed/tests/test_nanny.py index 33cffc8d72a..6312a4dfe9f 100644 --- a/distributed/tests/test_nanny.py +++ b/distributed/tests/test_nanny.py @@ -499,7 +499,7 @@ def raise_err(): @pytest.mark.parametrize("protocol", ["tcp", "ucx"]) -@pytest.mark.asyncio +@gen_test() async def test_nanny_closed_by_keyboard_interrupt(cleanup, protocol): if protocol == "ucx": # Skip if UCX isn't available pytest.importorskip("ucp") diff --git a/distributed/tests/test_preload.py b/distributed/tests/test_preload.py index 6d4c8b1fce8..0290da691a9 100644 --- a/distributed/tests/test_preload.py +++ b/distributed/tests/test_preload.py @@ -142,7 +142,7 @@ async def dask_setup(worker): assert w.foo == "setup" -@pytest.mark.asyncio +@gen_test() async def test_preload_import_time(cleanup): text = """ from distributed.comm.registry import backends @@ -207,7 +207,7 @@ def scheduler_preload(): @pytest.mark.skipif( MACOS and PY_VERSION == (3, 7), reason="HTTP Server doesn't come up" ) -@pytest.mark.asyncio +@gen_test() async def test_web_preload(cleanup, scheduler_preload): with captured_logger("distributed.preloading") as log: async with Scheduler( @@ -284,7 +284,7 @@ def worker_preload(): @pytest.mark.skipif( MACOS and PY_VERSION == (3, 7), reason="HTTP Server doesn't come up" ) -@pytest.mark.asyncio +@gen_test() async def test_web_preload_worker(cleanup, worker_preload): async with Scheduler(port=8786, host="localhost") as s: async with Nanny(preload_nanny=["http://127.0.0.1:12346/preload"]) as nanny: diff --git a/distributed/tests/test_scheduler.py b/distributed/tests/test_scheduler.py index 2750868dc01..7ee7e2fe0d9 100644 --- a/distributed/tests/test_scheduler.py +++ b/distributed/tests/test_scheduler.py @@ -1350,7 +1350,7 @@ async def test_non_existent_worker(c, s): ("127.0.0.1:0", ("127.0.0.1",)), ], ) -@pytest.mark.asyncio +@gen_test() async def test_dashboard_host(host, dashboard_address, expect): """Dashboard is accessible from any host by default, but it can be also bound to localhost. diff --git a/distributed/tests/test_utils.py b/distributed/tests/test_utils.py index c920366c4c2..51fd315a26c 100644 --- a/distributed/tests/test_utils.py +++ b/distributed/tests/test_utils.py @@ -563,13 +563,13 @@ def test_lru(): assert list(l.keys()) == ["c", "a", "d"] -@pytest.mark.asyncio +@gen_test() async def test_offload(): assert (await offload(inc, 1)) == 2 assert (await offload(lambda x, y: x + y, 1, y=2)) == 3 -@pytest.mark.asyncio +@gen_test() async def test_offload_preserves_contextvars(): var = contextvars.ContextVar("var") diff --git a/distributed/tests/test_utils_test.py b/distributed/tests/test_utils_test.py index dc481aae17f..9cdbadf2379 100755 --- a/distributed/tests/test_utils_test.py +++ b/distributed/tests/test_utils_test.py @@ -263,7 +263,7 @@ def test_tls_cluster(tls_client): assert tls_client.security -@pytest.mark.asyncio +@gen_test() async def test_tls_scheduler(security, cleanup): async with Scheduler( security=security, host="localhost", dashboard_address=":0" @@ -289,7 +289,7 @@ def pong(self, comm): return "pong" -@pytest.mark.asyncio +@gen_test() async def test_locked_comm_drop_in_replacement(loop): async with MyServer({}) as a, await MyServer({}) as b: @@ -318,7 +318,7 @@ async def test_locked_comm_drop_in_replacement(loop): assert await read_queue.get() == (b.address, "pong") -@pytest.mark.asyncio +@gen_test() async def test_locked_comm_intercept_read(loop): async with MyServer({}) as a, MyServer({}) as b: @@ -347,7 +347,7 @@ async def ping_pong(): assert await fut == "pong" -@pytest.mark.asyncio +@gen_test() async def test_locked_comm_intercept_write(loop): async with MyServer({}) as a, MyServer({}) as b: diff --git a/distributed/tests/test_worker.py b/distributed/tests/test_worker.py index aa8549b07ba..5a1fad15f91 100644 --- a/distributed/tests/test_worker.py +++ b/distributed/tests/test_worker.py @@ -389,7 +389,7 @@ def __str__(self): assert "Bar" in str(e.__cause__) -@pytest.mark.asyncio +@gen_test() async def test_plugin_exception(cleanup): class MyPlugin: def setup(self, worker=None): @@ -406,7 +406,7 @@ def setup(self, worker=None): pass -@pytest.mark.asyncio +@gen_test() async def test_plugin_multiple_exceptions(cleanup): class MyPlugin1: def setup(self, worker=None): @@ -434,7 +434,7 @@ def setup(self, worker=None): assert "MyPlugin2 Error" in text -@pytest.mark.asyncio +@gen_test() async def test_plugin_internal_exception(cleanup): async with Scheduler(port=0) as s: with pytest.raises(UnicodeDecodeError, match="codec can't decode"): @@ -1323,7 +1323,7 @@ async def test_host_address(c, s): await n.close() -@pytest.mark.asyncio +@gen_test() @pytest.mark.parametrize("Worker", [Worker, Nanny]) async def test_interface_async(cleanup, loop, Worker): from distributed.utils import get_ip_interface @@ -1356,7 +1356,7 @@ async def test_interface_async(cleanup, loop, Worker): @pytest.mark.gpu -@pytest.mark.asyncio +@gen_test() @pytest.mark.parametrize("Worker", [Worker, Nanny]) async def test_protocol_from_scheduler_address(cleanup, Worker): pytest.importorskip("ucp") @@ -1370,7 +1370,7 @@ async def test_protocol_from_scheduler_address(cleanup, Worker): assert info["address"].startswith("ucx://") -@pytest.mark.asyncio +@gen_test() async def test_host_uses_scheduler_protocol(cleanup, monkeypatch): # Ensure worker uses scheduler's protocol to determine host address, not the default scheme # See https://github.com/dask/distributed/pull/4883 @@ -1390,7 +1390,7 @@ def get_address_host(self, loc): pass -@pytest.mark.asyncio +@gen_test() @pytest.mark.parametrize("Worker", [Worker, Nanny]) async def test_worker_listens_on_same_interface_by_default(cleanup, Worker): async with Scheduler(host="localhost", dashboard_address=":0") as s: From aca8cec9a1bae149c1e895caf32720b63ba91463 Mon Sep 17 00:00:00 2001 From: Thomas Grainger Date: Mon, 4 Apr 2022 10:34:10 +0100 Subject: [PATCH 02/12] fix gen_test related flake8 import errors --- distributed/comm/tests/test_comms.py | 1 + distributed/comm/tests/test_ucx.py | 2 +- distributed/comm/tests/test_ucx_config.py | 2 +- distributed/deploy/tests/test_adaptive_core.py | 4 +--- distributed/deploy/tests/test_cluster.py | 1 + distributed/deploy/tests/test_slow_adaptive.py | 2 +- distributed/deploy/tests/test_ssh.py | 1 + distributed/http/tests/test_routing.py | 1 + distributed/tests/test_batched.py | 2 +- 9 files changed, 9 insertions(+), 7 deletions(-) diff --git a/distributed/comm/tests/test_comms.py b/distributed/comm/tests/test_comms.py index da1376feed6..c8f993de742 100644 --- a/distributed/comm/tests/test_comms.py +++ b/distributed/comm/tests/test_comms.py @@ -29,6 +29,7 @@ from distributed.protocol import Serialized, deserialize, serialize, to_serialize from distributed.utils import get_ip, get_ipv6, mp_context from distributed.utils_test import ( + gen_test, get_cert, get_client_ssl_context, get_server_ssl_context, diff --git a/distributed/comm/tests/test_ucx.py b/distributed/comm/tests/test_ucx.py index 14350f0a338..ac369d6c918 100644 --- a/distributed/comm/tests/test_ucx.py +++ b/distributed/comm/tests/test_ucx.py @@ -14,7 +14,7 @@ from distributed.deploy.local import LocalCluster from distributed.diagnostics.nvml import has_cuda_context from distributed.protocol import to_serialize -from distributed.utils_test import inc +from distributed.utils_test import gen_test, inc try: HOST = ucp.get_address() diff --git a/distributed/comm/tests/test_ucx_config.py b/distributed/comm/tests/test_ucx_config.py index 7049b192250..f1f3f08a3ab 100644 --- a/distributed/comm/tests/test_ucx_config.py +++ b/distributed/comm/tests/test_ucx_config.py @@ -10,7 +10,7 @@ from distributed import Client from distributed.comm.ucx import _scrub_ucx_config from distributed.utils import get_ip -from distributed.utils_test import popen +from distributed.utils_test import gen_test, popen try: HOST = get_ip() diff --git a/distributed/deploy/tests/test_adaptive_core.py b/distributed/deploy/tests/test_adaptive_core.py index 0c44200e0b6..3e1fa11bee9 100644 --- a/distributed/deploy/tests/test_adaptive_core.py +++ b/distributed/deploy/tests/test_adaptive_core.py @@ -1,10 +1,8 @@ import asyncio -import pytest - from distributed.deploy.adaptive_core import AdaptiveCore from distributed.metrics import time -from distributed.utils_test import captured_logger +from distributed.utils_test import captured_logger, gen_test class MyAdaptive(AdaptiveCore): diff --git a/distributed/deploy/tests/test_cluster.py b/distributed/deploy/tests/test_cluster.py index a5a14a2a8bc..b408ee1f435 100644 --- a/distributed/deploy/tests/test_cluster.py +++ b/distributed/deploy/tests/test_cluster.py @@ -1,6 +1,7 @@ import pytest from distributed.deploy.cluster import Cluster +from distributed.utils_test import gen_test @gen_test() diff --git a/distributed/deploy/tests/test_slow_adaptive.py b/distributed/deploy/tests/test_slow_adaptive.py index a55e4151acc..549ee40375c 100644 --- a/distributed/deploy/tests/test_slow_adaptive.py +++ b/distributed/deploy/tests/test_slow_adaptive.py @@ -5,7 +5,7 @@ from dask.distributed import Client, Scheduler, SpecCluster, Worker from distributed.metrics import time -from distributed.utils_test import slowinc +from distributed.utils_test import gen_test, slowinc class SlowWorker: diff --git a/distributed/deploy/tests/test_ssh.py b/distributed/deploy/tests/test_ssh.py index be03c316466..e7f415c5511 100644 --- a/distributed/deploy/tests/test_ssh.py +++ b/distributed/deploy/tests/test_ssh.py @@ -9,6 +9,7 @@ from distributed import Client from distributed.compatibility import MACOS, WINDOWS from distributed.deploy.ssh import SSHCluster +from distributed.utils_test import gen_test pytestmark = [ pytest.mark.xfail(MACOS, reason="very high flakiness; see distributed/issues/4543"), diff --git a/distributed/http/tests/test_routing.py b/distributed/http/tests/test_routing.py index a6a2a381f1d..9179061d29e 100644 --- a/distributed/http/tests/test_routing.py +++ b/distributed/http/tests/test_routing.py @@ -3,6 +3,7 @@ from tornado.httpclient import AsyncHTTPClient, HTTPClientError from distributed.http.routing import RoutingApplication +from distributed.utils_test import gen_test class OneHandler(web.RequestHandler): diff --git a/distributed/tests/test_batched.py b/distributed/tests/test_batched.py index 4efc71de15b..c1188ab79dd 100644 --- a/distributed/tests/test_batched.py +++ b/distributed/tests/test_batched.py @@ -9,7 +9,7 @@ from distributed.metrics import time from distributed.protocol import to_serialize from distributed.utils import All -from distributed.utils_test import captured_logger +from distributed.utils_test import captured_logger, gen_test class EchoServer: From 85b1720009ab80e71fe22f33d187f6a0ea9c6fbc Mon Sep 17 00:00:00 2001 From: Thomas Grainger Date: Mon, 4 Apr 2022 15:45:50 +0100 Subject: [PATCH 03/12] fix gen_test related test failures gen_test doesn't transport marks correctly and so must be the first decorator applied to a test function also gen_test is stricter than pytest-asyncio it checks for leaky comms and threads, so we need to make sure resources are closed properly --- distributed/comm/tests/test_ws.py | 4 +- distributed/deploy/tests/test_local.py | 41 +++-- distributed/deploy/tests/test_spec_cluster.py | 41 +++-- distributed/tests/test_batched.py | 3 + distributed/tests/test_client.py | 2 +- distributed/tests/test_core.py | 8 +- distributed/tests/test_preload.py | 172 ++++++------------ distributed/tests/test_worker.py | 12 +- 8 files changed, 121 insertions(+), 162 deletions(-) diff --git a/distributed/comm/tests/test_ws.py b/distributed/comm/tests/test_ws.py index 4324b170a2e..acfabaa2bb8 100644 --- a/distributed/comm/tests/test_ws.py +++ b/distributed/comm/tests/test_ws.py @@ -125,7 +125,6 @@ async def test_large_transfer_with_no_compression(): await c.scatter(np.random.random(1_500_000)) -@gen_test() @pytest.mark.parametrize( "dashboard,protocol,security,port", [ @@ -139,6 +138,7 @@ async def test_large_transfer_with_no_compression(): (False, "wss://", True, 8786), ], ) +@gen_test() async def test_http_and_comm_server(cleanup, dashboard, protocol, security, port): if security: xfail_ssl_issue5601() @@ -157,8 +157,8 @@ async def test_http_and_comm_server(cleanup, dashboard, protocol, security, port assert result == 11 -@gen_test() @pytest.mark.parametrize("protocol", ["ws://", "wss://"]) +@gen_test() async def test_connection_made_with_extra_conn_args(cleanup, protocol): if protocol == "ws://": security = Security( diff --git a/distributed/deploy/tests/test_local.py b/distributed/deploy/tests/test_local.py index 4da82ed4d25..135aabdd2d4 100644 --- a/distributed/deploy/tests/test_local.py +++ b/distributed/deploy/tests/test_local.py @@ -931,8 +931,8 @@ async def test_scale_memory_cores(): assert len(cluster.worker_spec) == 4 -@gen_test() @pytest.mark.parametrize("memory_limit", ["2 GiB", None]) +@gen_test() async def test_repr(memory_limit, cleanup): async with LocalCluster( n_workers=2, @@ -968,8 +968,8 @@ async def test_threads_per_worker_set_to_0(): assert all(w.nthreads < CPU_COUNT for w in cluster.workers.values()) -@gen_test() @pytest.mark.parametrize("temporary", [True, False]) +@gen_test() async def test_capture_security(cleanup, temporary): if temporary: xfail_ssl_issue5601() @@ -1039,25 +1039,28 @@ async def test_cluster_names(): assert unnamed_cluster2 != unnamed_cluster -@gen_test() @pytest.mark.parametrize("nanny", [True, False]) +@gen_test() async def test_local_cluster_redundant_kwarg(nanny): - with pytest.raises(TypeError, match="unexpected keyword argument"): - # Extra arguments are forwarded to the worker class. Depending on - # whether we use the nanny or not, the error treatment is quite - # different and we should assert that an exception is raised - async with await LocalCluster( - typo_kwarg="foo", - processes=nanny, - n_workers=1, - dashboard_address=":0", - ) as cluster: + cluster = LocalCluster( + typo_kwarg="foo", + processes=nanny, + n_workers=1, + dashboard_address=":0", + asynchronous=True, + ) + try: + with pytest.raises(TypeError, match="unexpected keyword argument") as exc_info: + # Extra arguments are forwarded to the worker class. Depending on + # whether we use the nanny or not, the error treatment is quite + # different and we should assert that an exception is raised + async with cluster: + pass + finally: + # FIXME: LocalCluster leaks if LocalCluster.__aenter__ raises + await cluster.close() - # This will never work but is a reliable way to block without hard - # coding any sleep values - async with Client(cluster) as c: - f = c.submit(sleep, 0) - await f + print(exc_info) @gen_test() @@ -1123,9 +1126,9 @@ async def error(*args, **kwargs): assert info["foo"] == "bar" -@gen_test() @pytest.mark.parametrize("host", [None, "127.0.0.1"]) @pytest.mark.parametrize("use_nanny", [True, False]) +@gen_test() async def test_cluster_host_used_throughout_cluster(host, use_nanny): """Ensure that the `host` kwarg is propagated through scheduler, nanny, and workers""" async with LocalCluster(host=host, asynchronous=True) as cluster: diff --git a/distributed/deploy/tests/test_spec_cluster.py b/distributed/deploy/tests/test_spec_cluster.py index 4b6ab24c894..4731e0fe9a2 100644 --- a/distributed/deploy/tests/test_spec_cluster.py +++ b/distributed/deploy/tests/test_spec_cluster.py @@ -21,14 +21,6 @@ class MyWorker(Worker): pass -class BrokenWorker(Worker): - def __await__(self): - async def _(): - raise Exception("Worker Broken") - - return _().__await__() - - worker_spec = { 0: {"cls": "dask.distributed.Worker", "options": {"nthreads": 1}}, 1: {"cls": Worker, "options": {"nthreads": 2}}, @@ -210,20 +202,31 @@ async def test_restart(): @pytest.mark.skipif(WINDOWS, reason="HTTP Server doesn't close out") -# FIXME cleanup fails: -# some RPCs left active by test: {} -# @gen_test() @gen_test() async def test_broken_worker(): - with pytest.raises(Exception) as info: - async with SpecCluster( - asynchronous=True, - workers={"good": {"cls": Worker}, "bad": {"cls": BrokenWorker}}, - scheduler=scheduler, - ): - pass + class BrokenWorkerException(Exception): + pass + + class BrokenWorker(Worker): + def __await__(self): + async def _(): + self.status = Status.closed + raise BrokenWorkerException("Worker Broken") + + return _().__await__() - assert "Broken" in str(info.value) + cluster = SpecCluster( + asynchronous=True, + workers={"good": {"cls": Worker}, "bad": {"cls": BrokenWorker}}, + scheduler=scheduler, + ) + try: + with pytest.raises(BrokenWorkerException, match=r"Worker Broken"): + async with cluster: + pass + finally: + # FIXME: SpecCluster leaks if SpecCluster.__aenter__ raises + await cluster.close() @pytest.mark.skipif(WINDOWS, reason="HTTP Server doesn't close out") diff --git a/distributed/tests/test_batched.py b/distributed/tests/test_batched.py index c1188ab79dd..e5cd1037669 100644 --- a/distributed/tests/test_batched.py +++ b/distributed/tests/test_batched.py @@ -62,6 +62,7 @@ async def test_BatchedSend(): assert result == ("HELLO", "HELLO") assert b.byte_count > 1 + await comm.close() @gen_test() @@ -77,6 +78,7 @@ async def test_send_before_start(): b.start(comm) result = await comm.read() assert result == ("hello", "world") + await comm.close() @gen_test() @@ -93,6 +95,7 @@ async def test_send_after_stream_start(): if len(result) < 2: result += await comm.read() assert result == ("hello", "world") + await comm.close() @gen_test() diff --git a/distributed/tests/test_client.py b/distributed/tests/test_client.py index 164d5e88121..2dd848fc80a 100644 --- a/distributed/tests/test_client.py +++ b/distributed/tests/test_client.py @@ -6127,8 +6127,8 @@ async def test_wait_for_workers(c, s, a, b): @pytest.mark.skipif(WINDOWS, reason="num_fds not supported on windows") -@gen_test() @pytest.mark.parametrize("Worker", [Worker, Nanny]) +@gen_test() async def test_file_descriptors_dont_leak(Worker): pytest.importorskip("pandas") df = dask.datasets.timeseries(freq="10s", dtypes={"x": int, "y": float}) diff --git a/distributed/tests/test_core.py b/distributed/tests/test_core.py index 90d43dee10c..16395e08f02 100644 --- a/distributed/tests/test_core.py +++ b/distributed/tests/test_core.py @@ -5,6 +5,7 @@ import weakref import pytest +from tornado.ioloop import IOLoop import dask @@ -458,7 +459,8 @@ async def test_identity_inproc(): @gen_test() -async def test_ports(loop): +async def test_ports(): + loop = IOLoop.current() for port in range(9877, 9887): server = Server({}, io_loop=loop) try: @@ -859,13 +861,15 @@ async def test_deserialize_error(): await server.listen(0) comm = await connect(server.address, deserialize=False) - with pytest.raises(Exception) as info: + with pytest.raises(Exception, match=r"RuntimeError\('hello!'\)") as info: await send_recv(comm, op="throws", x="foo") assert type(info.value) == Exception for c in str(info.value): assert c.isalpha() or c in "(',!)" # no crazy bytestrings + await comm.close() + @gen_test() async def test_connection_pool_detects_remote_close(): diff --git a/distributed/tests/test_preload.py b/distributed/tests/test_preload.py index 0290da691a9..6ab1637edb2 100644 --- a/distributed/tests/test_preload.py +++ b/distributed/tests/test_preload.py @@ -1,26 +1,20 @@ -import multiprocessing +import asyncio import os +import re import shutil +import socket import sys import tempfile -import urllib.error -import urllib.request from textwrap import dedent -from time import sleep import pytest -import tornado -from tornado import web import dask from distributed import Client, Nanny, Scheduler, Worker -from distributed.compatibility import MACOS -from distributed.metrics import time +from distributed.compatibility import to_thread from distributed.utils_test import captured_logger, cluster, gen_cluster, gen_test -PY_VERSION = sys.version_info[:2] - PRELOAD_TEXT = """ _worker_info = {} @@ -161,61 +155,40 @@ async def test_preload_import_time(cleanup): del backends["foo"] -class MyHandler(web.RequestHandler): - def get(self): - self.write( - """ -def dask_setup(dask_server): - dask_server.foo = 1 -""".strip() +@gen_test() +async def test_web_preload(): + def preload(): + with server_sock.accept()[0] as client_sock, client_sock.makefile("rwb") as f: + assert f.readline() == b"GET /preload HTTP/1.1\r\n" + f.write( + b"HTTP/1.1 200 OK" + b"\r\nContent-Length: 53" + b"\r\n\r\n" + b"def dask_setup(dask_server):" + b"\n dask_server.foo = 1" + b"\n" + ) + + async def test(): + port = server_sock.getsockname()[1] + with captured_logger("distributed.preloading") as log: + async with Scheduler( + host="localhost", + preload=[f"http://127.0.0.1:{port}/preload"], + ) as s: + assert s.foo == 1 + assert ( + re.match( + r"(?s).*Downloading preload at http://127.0.0.1:\d+/preload\n" + r".*Run preload setup function: http://127.0.0.1:\d+/preload\n" + r".*", + log.getvalue(), + ) + is not None ) - -def create_preload_application(): - app = web.Application([(r"/preload", MyHandler)]) - server = app.listen(12345, address="127.0.0.1") - tornado.ioloop.IOLoop.instance().start() - - -@pytest.fixture -def scheduler_preload(): - p = multiprocessing.Process(target=create_preload_application) - p.start() - start = time() - while not p.is_alive(): - if time() > start + 5: - raise AssertionError("Process didn't come up") - sleep(0.5) - # Make sure we can query the server - start = time() - request = urllib.request.Request("http://127.0.0.1:12345/preload", method="GET") - while True: - try: - response = urllib.request.urlopen(request) - if response.status == 200: - break - except urllib.error.URLError as e: - if time() > start + 10: - raise AssertionError("Webserver didn't come up", e) - sleep(0.5) - - yield - p.kill() - p.join(timeout=5) - - -@pytest.mark.skipif( - MACOS and PY_VERSION == (3, 7), reason="HTTP Server doesn't come up" -) -@gen_test() -async def test_web_preload(cleanup, scheduler_preload): - with captured_logger("distributed.preloading") as log: - async with Scheduler( - host="localhost", - preload=["http://127.0.0.1:12345/preload"], - ) as s: - assert s.foo == 1 - assert "12345/preload" in log.getvalue() + with socket.create_server(("127.0.0.1", 0)) as server_sock: + await asyncio.gather(to_thread(preload), test()) @gen_cluster(nthreads=[]) @@ -238,57 +211,30 @@ async def test_scheduler_startup_nanny(s): assert w.scheduler.address == s.address -class WorkerPreloadHandler(web.RequestHandler): - def get(self): - self.write( - """ -import dask -dask.config.set(scheduler_address="tcp://127.0.0.1:8786") -""".strip() - ) - - -def create_worker_preload_application(): - application = web.Application([(r"/preload", WorkerPreloadHandler)]) - server = application.listen(12346, address="127.0.0.1") - tornado.ioloop.IOLoop.instance().start() - - -@pytest.fixture -def worker_preload(): - p = multiprocessing.Process(target=create_worker_preload_application) - p.start() - start = time() - while not p.is_alive(): - if time() > start + 5: - raise AssertionError("Process didn't come up") - sleep(0.5) - # Make sure we can query the server - request = urllib.request.Request("http://127.0.0.1:12346/preload", method="GET") - start = time() - while True: - try: - response = urllib.request.urlopen(request) - if response.status == 200: - break - except urllib.error.URLError as e: - if time() > start + 10: - raise AssertionError("Webserver didn't come up", e) - sleep(0.5) - - yield - p.kill() - p.join(timeout=5) - - -@pytest.mark.skipif( - MACOS and PY_VERSION == (3, 7), reason="HTTP Server doesn't come up" -) @gen_test() -async def test_web_preload_worker(cleanup, worker_preload): - async with Scheduler(port=8786, host="localhost") as s: - async with Nanny(preload_nanny=["http://127.0.0.1:12346/preload"]) as nanny: - assert nanny.scheduler_addr == s.address +async def test_web_preload_worker(): + def preload(): + with server_sock.accept()[0] as client_sock, client_sock.makefile("rwb") as f: + assert f.readline() == b"GET /preload HTTP/1.1\r\n" + f.write( + b"HTTP/1.1 200 OK" + b"\r\nContent-Length: 70" + b"\r\n\r\n" + b"import dask" + b'\ndask.config.set(scheduler_address="tcp://127.0.0.1:8786")' + b"\n" + ) + + async def test(): + port = server_sock.getsockname()[1] + async with Scheduler(port=8786, host="localhost") as s: + async with Nanny( + preload_nanny=[f"http://127.0.0.1:{port}/preload"] + ) as nanny: + assert nanny.scheduler_addr == s.address + + with socket.create_server(("127.0.0.1", 0)) as server_sock: + await asyncio.gather(to_thread(preload), test()) # This test is blocked on https://github.com/dask/distributed/issues/5819 diff --git a/distributed/tests/test_worker.py b/distributed/tests/test_worker.py index 5a1fad15f91..357c34e5603 100644 --- a/distributed/tests/test_worker.py +++ b/distributed/tests/test_worker.py @@ -1323,9 +1323,9 @@ async def test_host_address(c, s): await n.close() -@gen_test() @pytest.mark.parametrize("Worker", [Worker, Nanny]) -async def test_interface_async(cleanup, loop, Worker): +@gen_test() +async def test_interface_async(Worker): from distributed.utils import get_ip_interface psutil = pytest.importorskip("psutil") @@ -1356,9 +1356,9 @@ async def test_interface_async(cleanup, loop, Worker): @pytest.mark.gpu -@gen_test() @pytest.mark.parametrize("Worker", [Worker, Nanny]) -async def test_protocol_from_scheduler_address(cleanup, Worker): +@gen_test() +async def test_protocol_from_scheduler_address(Worker): pytest.importorskip("ucp") async with Scheduler(protocol="ucx", dashboard_address=":0") as s: @@ -1390,9 +1390,9 @@ def get_address_host(self, loc): pass -@gen_test() @pytest.mark.parametrize("Worker", [Worker, Nanny]) -async def test_worker_listens_on_same_interface_by_default(cleanup, Worker): +@gen_test() +async def test_worker_listens_on_same_interface_by_default(Worker): async with Scheduler(host="localhost", dashboard_address=":0") as s: assert s.ip in {"127.0.0.1", "localhost"} async with Worker(s.address) as w: From ae62801bbfdb74620feaf54e5a8b56a279bec4cf Mon Sep 17 00:00:00 2001 From: Thomas Grainger Date: Mon, 4 Apr 2022 15:56:34 +0100 Subject: [PATCH 04/12] remove pytest-asyncio dep --- continuous_integration/environment-3.10.yaml | 1 - continuous_integration/environment-3.8.yaml | 1 - continuous_integration/environment-3.9.yaml | 1 - 3 files changed, 3 deletions(-) diff --git a/continuous_integration/environment-3.10.yaml b/continuous_integration/environment-3.10.yaml index f48a2df915d..28ccbe31c87 100644 --- a/continuous_integration/environment-3.10.yaml +++ b/continuous_integration/environment-3.10.yaml @@ -47,4 +47,3 @@ dependencies: - git+https://github.com/fsspec/filesystem_spec - git+https://github.com/joblib/joblib - keras - - pytest-asyncio<0.14.0 # `pytest-asyncio<0.14.0` isn't available on conda-forge for Python 3.10 diff --git a/continuous_integration/environment-3.8.yaml b/continuous_integration/environment-3.8.yaml index 715c95a90c3..aba501fb2f1 100644 --- a/continuous_integration/environment-3.8.yaml +++ b/continuous_integration/environment-3.8.yaml @@ -26,7 +26,6 @@ dependencies: - prometheus_client - psutil - pytest - - pytest-asyncio<0.14.0 - pytest-cov - pytest-faulthandler - pytest-repeat diff --git a/continuous_integration/environment-3.9.yaml b/continuous_integration/environment-3.9.yaml index acbb4626b10..8813552c9d7 100644 --- a/continuous_integration/environment-3.9.yaml +++ b/continuous_integration/environment-3.9.yaml @@ -28,7 +28,6 @@ dependencies: - psutil - pynvml # Only tested here - pytest - - pytest-asyncio<0.14.0 - pytest-cov - pytest-faulthandler - pytest-repeat From 48cfd4b0b576685a27b780bc34b2fd684111eddc Mon Sep 17 00:00:00 2001 From: Thomas Grainger Date: Mon, 4 Apr 2022 16:04:49 +0100 Subject: [PATCH 05/12] fix gen_test related ucx test failures --- distributed/comm/tests/test_ucx.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/distributed/comm/tests/test_ucx.py b/distributed/comm/tests/test_ucx.py index ac369d6c918..147915eae2a 100644 --- a/distributed/comm/tests/test_ucx.py +++ b/distributed/comm/tests/test_ucx.py @@ -179,7 +179,6 @@ async def test_ucx_deserialize(): await check_deserialize("tcp://") -@gen_test() @pytest.mark.parametrize( "g", [ @@ -196,6 +195,7 @@ async def test_ucx_deserialize(): lambda cudf: cudf.DataFrame({"a": ["Check", "str"], "b": ["Sup", "port"]}), ], ) +@gen_test() async def test_ping_pong_cudf(g): # if this test appears after cupy an import error arises # *** ImportError: /usr/lib/x86_64-linux-gnu/libstdc++.so.6: version `CXXABI_1.3.11' @@ -219,8 +219,8 @@ async def test_ping_pong_cudf(g): await serv_com.close() -@gen_test() @pytest.mark.parametrize("shape", [(100,), (10, 10), (4947,)]) +@gen_test() async def test_ping_pong_cupy(shape): cupy = pytest.importorskip("cupy") com, serv_com = await get_comm_pair() @@ -238,8 +238,8 @@ async def test_ping_pong_cupy(shape): @pytest.mark.slow -@gen_test() @pytest.mark.parametrize("n", [int(1e9), int(2.5e9)]) +@gen_test() async def test_large_cupy(n, cleanup): cupy = pytest.importorskip("cupy") com, serv_com = await get_comm_pair() @@ -274,8 +274,8 @@ async def test_ping_pong_numba(): assert result["op"] == "ping" -@gen_test() @pytest.mark.parametrize("processes", [True, False]) +@gen_test() async def test_ucx_localcluster(processes, cleanup): async with LocalCluster( protocol="ucx", @@ -356,8 +356,8 @@ async def test_transpose(): await y -@gen_test() @pytest.mark.parametrize("port", [0, 1234]) +@gen_test() async def test_ucx_protocol(cleanup, port): async with Scheduler(protocol="ucx", port=port, dashboard_address=":0") as s: assert s.address.startswith("ucx://") From b29d0e2cc69ac7521cf6e09c0a61bdeffe312d91 Mon Sep 17 00:00:00 2001 From: Thomas Grainger Date: Mon, 4 Apr 2022 18:17:09 +0100 Subject: [PATCH 06/12] fix gen_test related test_ssh failures --- distributed/deploy/tests/test_ssh.py | 23 ++++++++++++++--------- 1 file changed, 14 insertions(+), 9 deletions(-) diff --git a/distributed/deploy/tests/test_ssh.py b/distributed/deploy/tests/test_ssh.py index e7f415c5511..d5d55477941 100644 --- a/distributed/deploy/tests/test_ssh.py +++ b/distributed/deploy/tests/test_ssh.py @@ -115,15 +115,20 @@ async def test_ssh_nprocs_renamed_to_n_workers(): @gen_test() async def test_ssh_n_workers_with_nprocs_is_an_error(): - with pytest.raises(ValueError, match="Both nprocs and n_workers"): - async with SSHCluster( - ["127.0.0.1"] * 3, - connect_options=dict(known_hosts=None), - asynchronous=True, - scheduler_options={}, - worker_options={"n_workers": 2, "nprocs": 2}, - ) as cluster: - assert not cluster + cluster = SSHCluster( + ["127.0.0.1"] * 3, + connect_options=dict(known_hosts=None), + asynchronous=True, + scheduler_options={}, + worker_options={"n_workers": 2, "nprocs": 2}, + ) + try: + with pytest.raises(ValueError, match="Both nprocs and n_workers"): + async with cluster: + pass + finally: + # FIXME: SSHCluster leaks if SSHCluster.__aenter__ raises + await cluster.close() @gen_test() From 1f322c883f729063777623cd295ed42e52893463 Mon Sep 17 00:00:00 2001 From: Thomas Grainger Date: Tue, 5 Apr 2022 10:08:08 +0100 Subject: [PATCH 07/12] fix gen_test related @pytest.mark ordering --- distributed/deploy/tests/test_slow_adaptive.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/distributed/deploy/tests/test_slow_adaptive.py b/distributed/deploy/tests/test_slow_adaptive.py index 549ee40375c..eea54621a2f 100644 --- a/distributed/deploy/tests/test_slow_adaptive.py +++ b/distributed/deploy/tests/test_slow_adaptive.py @@ -55,9 +55,9 @@ async def test_startup(cleanup): await client.wait_for_workers(n_workers=2) -@gen_test() @pytest.mark.flaky(reruns=10, reruns_delay=5) -async def test_scale_up_down(cleanup): +@gen_test() +async def test_scale_up_down(): start = time() async with SpecCluster( scheduler=scheduler, From fc9b51de889e27379531d2290dd15515f0c43f5f Mon Sep 17 00:00:00 2001 From: Thomas Grainger Date: Tue, 5 Apr 2022 10:25:49 +0100 Subject: [PATCH 08/12] close comms in test_retry_connect --- distributed/comm/tests/test_comms.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/distributed/comm/tests/test_comms.py b/distributed/comm/tests/test_comms.py index c8f993de742..fe1f4855f5f 100644 --- a/distributed/comm/tests/test_comms.py +++ b/distributed/comm/tests/test_comms.py @@ -896,6 +896,7 @@ async def test_retry_connect(tcp, monkeypatch): async def echo(comm): message = await comm.read() await comm.write(message) + await comm.close() class UnreliableConnector(tcp.TCPConnector): def __init__(self): @@ -922,6 +923,7 @@ class UnreliableBackend(tcp.TCPBackend): await comm.write(b"test") msg = await comm.read() assert msg == b"test" + await comm.close() finally: listener.stop() From d64ce70b69e2caf9fe4ea25251aa9162931552c2 Mon Sep 17 00:00:00 2001 From: Thomas Grainger Date: Tue, 5 Apr 2022 12:09:46 +0100 Subject: [PATCH 09/12] close comms in test_connection_pool_respects_limit --- distributed/tests/test_core.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/distributed/tests/test_core.py b/distributed/tests/test_core.py index 16395e08f02..a32b7432929 100644 --- a/distributed/tests/test_core.py +++ b/distributed/tests/test_core.py @@ -1,4 +1,5 @@ import asyncio +import contextlib import os import socket import threading @@ -700,13 +701,16 @@ async def do_ping(pool, port): await pool(ip="127.0.0.1", port=port).ping() assert pool.open <= limit - servers = [Server({"ping": ping}) for i in range(10)] - for server in servers: - await server.listen(0) - - pool = await ConnectionPool(limit=limit) + async with contextlib.AsyncExitStack() as stack: + servers = [ + await stack.enter_async_context(Server({"ping": ping})) for i in range(10) + ] + for server in servers: + await server.listen(0) - await asyncio.gather(*(do_ping(pool, s.port) for s in servers)) + pool = await ConnectionPool(limit=limit) + await asyncio.gather(*(do_ping(pool, s.port) for s in servers)) + await pool.close() @gen_test() From 0d55a2bfb5d2e883e580335fd407aef71d3e8970 Mon Sep 17 00:00:00 2001 From: Thomas Grainger Date: Tue, 5 Apr 2022 13:46:58 +0100 Subject: [PATCH 10/12] double test_ucx.test_stress timeout --- distributed/comm/tests/test_ucx.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/distributed/comm/tests/test_ucx.py b/distributed/comm/tests/test_ucx.py index 147915eae2a..f4a5729826a 100644 --- a/distributed/comm/tests/test_ucx.py +++ b/distributed/comm/tests/test_ucx.py @@ -296,7 +296,7 @@ async def test_ucx_localcluster(processes, cleanup): @pytest.mark.slow -@gen_test() +@gen_test(timeout=60) async def test_stress(): da = pytest.importorskip("dask.array") From f7c465d1eddbd3ec049d17c54b4b205c73c4d4ef Mon Sep 17 00:00:00 2001 From: Thomas Grainger Date: Wed, 6 Apr 2022 10:10:11 +0100 Subject: [PATCH 11/12] remove stray print(exc_info) --- distributed/deploy/tests/test_local.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/distributed/deploy/tests/test_local.py b/distributed/deploy/tests/test_local.py index 135aabdd2d4..5c0e755f891 100644 --- a/distributed/deploy/tests/test_local.py +++ b/distributed/deploy/tests/test_local.py @@ -1050,7 +1050,7 @@ async def test_local_cluster_redundant_kwarg(nanny): asynchronous=True, ) try: - with pytest.raises(TypeError, match="unexpected keyword argument") as exc_info: + with pytest.raises(TypeError, match="unexpected keyword argument"): # Extra arguments are forwarded to the worker class. Depending on # whether we use the nanny or not, the error treatment is quite # different and we should assert that an exception is raised @@ -1060,8 +1060,6 @@ async def test_local_cluster_redundant_kwarg(nanny): # FIXME: LocalCluster leaks if LocalCluster.__aenter__ raises await cluster.close() - print(exc_info) - @gen_test() async def test_cluster_info_sync(): From 70ca3206c4a7213c299416a0f84b7e6441e7183c Mon Sep 17 00:00:00 2001 From: Thomas Grainger Date: Wed, 6 Apr 2022 15:25:43 +0100 Subject: [PATCH 12/12] use a mock.patch instead of a raw socket in test_web_preload(_worker) --- distributed/tests/test_preload.py | 78 ++++++++++++------------------- 1 file changed, 30 insertions(+), 48 deletions(-) diff --git a/distributed/tests/test_preload.py b/distributed/tests/test_preload.py index 6ab1637edb2..073c9a2d959 100644 --- a/distributed/tests/test_preload.py +++ b/distributed/tests/test_preload.py @@ -1,18 +1,16 @@ -import asyncio import os import re import shutil -import socket import sys import tempfile from textwrap import dedent +from unittest import mock import pytest import dask from distributed import Client, Nanny, Scheduler, Worker -from distributed.compatibility import to_thread from distributed.utils_test import captured_logger, cluster, gen_cluster, gen_test PRELOAD_TEXT = """ @@ -157,38 +155,30 @@ async def test_preload_import_time(cleanup): @gen_test() async def test_web_preload(): - def preload(): - with server_sock.accept()[0] as client_sock, client_sock.makefile("rwb") as f: - assert f.readline() == b"GET /preload HTTP/1.1\r\n" - f.write( - b"HTTP/1.1 200 OK" - b"\r\nContent-Length: 53" - b"\r\n\r\n" - b"def dask_setup(dask_server):" - b"\n dask_server.foo = 1" - b"\n" - ) - - async def test(): - port = server_sock.getsockname()[1] - with captured_logger("distributed.preloading") as log: - async with Scheduler( - host="localhost", - preload=[f"http://127.0.0.1:{port}/preload"], - ) as s: - assert s.foo == 1 + with mock.patch( + "urllib3.PoolManager.request", + **{ + "return_value.data": b"def dask_setup(dask_server):" + b"\n dask_server.foo = 1" + b"\n" + }, + ) as request, captured_logger("distributed.preloading") as log: + async with Scheduler( + host="localhost", preload=["http://example.com/preload"] + ) as s: + assert s.foo == 1 assert ( re.match( - r"(?s).*Downloading preload at http://127.0.0.1:\d+/preload\n" - r".*Run preload setup function: http://127.0.0.1:\d+/preload\n" + r"(?s).*Downloading preload at http://example.com/preload\n" + r".*Run preload setup function: http://example.com/preload\n" r".*", log.getvalue(), ) is not None ) - - with socket.create_server(("127.0.0.1", 0)) as server_sock: - await asyncio.gather(to_thread(preload), test()) + assert request.mock_calls == [ + mock.call(method="GET", url="http://example.com/preload", retries=mock.ANY) + ] @gen_cluster(nthreads=[]) @@ -213,28 +203,20 @@ async def test_scheduler_startup_nanny(s): @gen_test() async def test_web_preload_worker(): - def preload(): - with server_sock.accept()[0] as client_sock, client_sock.makefile("rwb") as f: - assert f.readline() == b"GET /preload HTTP/1.1\r\n" - f.write( - b"HTTP/1.1 200 OK" - b"\r\nContent-Length: 70" - b"\r\n\r\n" - b"import dask" - b'\ndask.config.set(scheduler_address="tcp://127.0.0.1:8786")' - b"\n" - ) - - async def test(): - port = server_sock.getsockname()[1] + with mock.patch( + "urllib3.PoolManager.request", + **{ + "return_value.data": b"import dask" + b'\ndask.config.set(scheduler_address="tcp://127.0.0.1:8786")' + b"\n" + }, + ) as request: async with Scheduler(port=8786, host="localhost") as s: - async with Nanny( - preload_nanny=[f"http://127.0.0.1:{port}/preload"] - ) as nanny: + async with Nanny(preload_nanny=["http://example.com/preload"]) as nanny: assert nanny.scheduler_addr == s.address - - with socket.create_server(("127.0.0.1", 0)) as server_sock: - await asyncio.gather(to_thread(preload), test()) + assert request.mock_calls == [ + mock.call(method="GET", url="http://example.com/preload", retries=mock.ANY) + ] # This test is blocked on https://github.com/dask/distributed/issues/5819