Docs

distributed/shuffle/_core.py

@@ -285,16 +285,16 @@ def create_new_run(
     def _pin_output_workers(
         self, plugin: ShuffleSchedulerPlugin
     ) -> dict[_T_partition_id, str]:
-        """TODO"""
+        """Pin output tasks to workers and return the mapping of partition ID to worker."""
 
     @abc.abstractmethod
-    def initialize_run_on_worker(
+    def create_run_on_worker(
         self,
         run_id: int,
         worker_for: dict[_T_partition_id, str],
         plugin: ShuffleWorkerPlugin,
     ) -> ShuffleRun:
-        """TODO"""
+        """Create the new shuffle run on the worker."""
 
 
 @dataclass(eq=False)

distributed/shuffle/_rechunk.py

@@ -463,7 +463,7 @@ def _pin_output_workers(self, plugin: ShuffleSchedulerPlugin) -> dict[NDIndex, s
             self.id, parts_out, _get_worker_for_hash_sharding
         )
 
-    def initialize_run_on_worker(
+    def create_run_on_worker(
         self,
         run_id: int,
         worker_for: dict[NDIndex, str],

distributed/shuffle/_shuffle.py

@@ -502,7 +502,7 @@ def _pin_output_workers(self, plugin: ShuffleSchedulerPlugin) -> dict[int, str]:
         pick_worker = partial(_get_worker_for_range_sharding, self.npartitions)
         return plugin._pin_output_workers(self.id, self.parts_out, pick_worker)
 
-    def initialize_run_on_worker(
+    def create_run_on_worker(
         self, run_id: int, worker_for: dict[int, str], plugin: ShuffleWorkerPlugin
     ) -> ShuffleRun:
         return DataFrameShuffleRun(

distributed/shuffle/_worker_plugin.py

@@ -288,7 +288,7 @@ async def _(
 
                 self.worker._ongoing_background_tasks.call_soon(_, self, existing)
 
-        shuffle: ShuffleRun = result.spec.initialize_run_on_worker(
+        shuffle: ShuffleRun = result.spec.create_run_on_worker(
             result.run_id, result.worker_for, self
         )
         self.shuffles[shuffle_id] = shuffle