Fix excessive logging on P2P retry

Author: hendrikmakait

distributed/shuffle/_core.py

@@ -6,7 +6,15 @@
 import itertools
 import pickle
 import time
-from collections.abc import Callable, Generator, Hashable, Iterable, Iterator, Sequence
+from collections.abc import (
+    Callable,
+    Coroutine,
+    Generator,
+    Hashable,
+    Iterable,
+    Iterator,
+    Sequence,
+)
 from concurrent.futures import ThreadPoolExecutor
 from dataclasses import dataclass, field
 from enum import Enum
@@ -212,8 +220,11 @@ async def send(
         else:
             shards_or_bytes = shards
 
+        def _send() -> Coroutine[Any, Any, None]:
+            return self._send(address, shards_or_bytes)
+
         return await retry(
-            partial(self._send, address, shards_or_bytes),
+            _send,
             count=self.RETRY_COUNT,
             delay_min=self.RETRY_DELAY_MIN,
             delay_max=self.RETRY_DELAY_MAX,

distributed/shuffle/tests/test_shuffle.py

@@ -62,6 +62,7 @@
 from distributed.utils import Deadline
 from distributed.utils_test import (
     async_poll_for,
+    captured_logger,
     cluster,
     gen_cluster,
     gen_test,
@@ -2642,10 +2643,14 @@ async def test_flaky_connect_recover_with_retry(c, s, a, b):
         x = dd.shuffle.shuffle(df, "x")
 
     rpc = await FlakyConnectionPool(failing_connects=1)
-
-    with mock.patch.object(a, "rpc", rpc):
-        await c.compute(x)
-    assert rpc.failed_attempts == 1
+    with captured_logger("distributed.utils_comm") as caplog:
+        with mock.patch.object(a, "rpc", rpc):
+            await c.compute(x)
+        assert rpc.failed_attempts == 1
+    # Assert that we do not log the binary payload (or any other excessive amount of data)
+    logs = caplog.getvalue()
+    assert len(logs) < 200
+    assert "Retrying" in logs
 
     await check_worker_cleanup(a)
     await check_worker_cleanup(b)