annotations

dask · Oct 25, 2023 · e1db838 · e1db838
1 parent d53392f
commit e1db838
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/distributed/shuffle/_core.py b/distributed/shuffle/_core.py
@@ -129,7 +129,7 @@ async def barrier(self, run_ids: Sequence[int]) -> int:
         return self.run_id
 
     async def _send(
-        self, address: str, shards: list[tuple[_T_partition_id, bytes]]
+        self, address: str, shards: list[tuple[_T_partition_id, Any]]
     ) -> None:
         self.raise_if_closed()
         return await self.rpc(address).shuffle_receive(
@@ -139,7 +139,7 @@ async def _send(
         )
 
     async def send(
-        self, address: str, shards: list[tuple[_T_partition_id, bytes]]
+        self, address: str, shards: list[tuple[_T_partition_id, Any]]
     ) -> None:
         retry_count = dask.config.get("distributed.p2p.comm.retry.count")
         retry_delay_min = parse_timedelta(

diff --git a/distributed/shuffle/_worker_plugin.py b/distributed/shuffle/_worker_plugin.py
@@ -299,7 +299,7 @@ async def shuffle_receive(
         self,
         shuffle_id: ShuffleId,
         run_id: int,
-        data: list[tuple[int, bytes]],
+        data: list[tuple[int, Any]],
     ) -> None:
         """
         Handler: Receive an incoming shard of data from a peer worker.