Shuffle p2p

distributed/shuffle/_core.py

@@ -340,7 +340,7 @@ async def _receive(self, data: list[tuple[_T_partition_id, Any]]) -> None:
         """Receive shards belonging to output partitions of this shuffle run"""
 
     def add_partition(
-        self, data: _T_partition_type, partition_id: _T_partition_id
+        self, data: _T_partition_type, partition_id: _T_partition_id, **kwargs: Any
     ) -> int:
         self.raise_if_closed()
         if self.transferred:
@@ -351,13 +351,13 @@ def add_partition(
                 context_meter.meter("p2p-shard-partition-noncpu"),
                 context_meter.meter("p2p-shard-partition-cpu", func=thread_time),
             ):
-                shards = self._shard_partition(data, partition_id)
+                shards = self._shard_partition(data, partition_id, **kwargs)
             sync(self._loop, self._write_to_comm, shards)
         return self.run_id
 
     @abc.abstractmethod
     def _shard_partition(
-        self, data: _T_partition_type, partition_id: _T_partition_id
+        self, data: _T_partition_type, partition_id: _T_partition_id, **kwargs: Any
     ) -> dict[str, tuple[_T_partition_id, Any]]:
         """Shard an input partition by the assigned output workers"""
 

distributed/shuffle/_rechunk.py

@@ -907,7 +907,7 @@ async def _receive(
             raise
 
     def _shard_partition(
-        self, data: np.ndarray, partition_id: NDIndex
+        self, data: np.ndarray, partition_id: NDIndex, **kwargs: Any
     ) -> dict[str, tuple[NDIndex, list[tuple[NDIndex, tuple[NDIndex, np.ndarray]]]]]:
         out: dict[str, list[tuple[NDIndex, tuple[NDIndex, np.ndarray]]]] = defaultdict(
             list