Re-raise P2PConsistencyError from failed P2P tasks. (#8748)

distributed/shuffle/_core.py

@@ -506,6 +506,8 @@ def handle_transfer_errors(id: ShuffleId) -> Iterator[None]:
         yield
     except ShuffleClosedError:
         raise Reschedule()
+    except P2PConsistencyError:
+        raise
     except Exception as e:
         raise RuntimeError(f"P2P shuffling {id} failed during transfer phase") from e
 
@@ -518,6 +520,8 @@ def handle_unpack_errors(id: ShuffleId) -> Iterator[None]:
         raise e
     except ShuffleClosedError:
         raise Reschedule()
+    except P2PConsistencyError:
+        raise
     except Exception as e:
         raise RuntimeError(f"P2P shuffling {id} failed during unpack phase") from e
 

distributed/shuffle/_shuffle.py

@@ -48,7 +48,7 @@
     handle_transfer_errors,
     handle_unpack_errors,
 )
-from distributed.shuffle._exceptions import DataUnavailable
+from distributed.shuffle._exceptions import DataUnavailable, P2PConsistencyError
 from distributed.shuffle._limiter import ResourceLimiter
 from distributed.shuffle._worker_plugin import ShuffleWorkerPlugin
 from distributed.sizeof import sizeof
@@ -105,6 +105,8 @@ def shuffle_barrier(id: ShuffleId, run_ids: list[int]) -> int:
         return get_worker_plugin().barrier(id, run_ids)
     except Reschedule as e:
         raise e
+    except P2PConsistencyError:
+        raise
     except Exception as e:
         raise RuntimeError(f"shuffle_barrier failed during shuffle {id}") from e
 

distributed/shuffle/tests/test_shuffle.py

@@ -254,9 +254,7 @@ async def test_shuffle_with_array_conversion(c, s, a, b, npartitions):
 
     if npartitions == 1:
         # FIXME: distributed#7816
-        with raises_with_cause(
-            RuntimeError, "failed during transfer", RuntimeError, "Barrier task"
-        ):
+        with pytest.raises(P2PConsistencyError, match="Barrier task"):
             await c.compute(out)
     else:
         await c.compute(out)