class_docstr: Shuffle plugin for the scheduler
  This coordinates the individual worker plugins to ensure correctness
  and collects heartbeat messages for the dashboard.
  See Also
  --------
  ShuffleWorkerPlugin
class_docstr_tok: ["Shuffle","plugin","for","the","scheduler","This","coordinates","the","individual","worker","plugins","to","ensure","correctness","and","collects","heartbeat","messages","for","the","dashboard",".","See","Also","--","--","--","--","ShuffleWorkerPlugin"]
docstr: Pin the outputs of a P2P shuffle to specific workers.
  
  The P2P implementation of a hash join combines the loading of shuffled output
  partitions for the left and right side with the actual merge operation into a
  single output task. As a consequence, we need to make sure that shuffles with
  shared output tasks align on the output mapping.
  
  Parameters
  ----------
  id: ID of the shuffle to pin
  output_partitions: Output partition IDs to pin
  pick: Function that picks a worker given a partition ID and sequence of worker
  
  .. note:
      This function assumes that the barrier task and the output tasks share
      the same worker restrictions.
docstr_tok: ["Pin","the","outputs","of","a","P2P","shuffle","to","specific","workers",".","The","P2P","implementation","of","a","hash","join","combines","the","loading","of","shuffled","output","partitions","for","the","left","and","right","side","with","the","actual","merge","operation","into","a","single","output","task",".","As","a","consequence",",","we","need","to","make","sure","that","shuffles","with","shared","output","tasks","align","on","the","output","mapping",".","Parameters","--","--","--","--","--","id",":","ID","of","the","shuffle","to","pin","output_partitions",":","Output","partition","IDs","to","pin","pick",":","Function","that","picks","a","worker","given","a","partition","ID","and","sequence","of","worker","..","note",":","This","function","assumes","that","the","barrier","task","and","the","output","tasks","share","the","same","worker","restrictions","."]
code: def _calculate_worker_for(self, spec: ShuffleSpec) -> dict[Any, str]:
      """Pin the outputs of a P2P shuffle to specific workers.
  
      The P2P implementation of a hash join combines the loading of shuffled output
      partitions for the left and right side with the actual merge operation into a
      single output task. As a consequence, we need to make sure that shuffles with
      shared output tasks align on the output mapping.
  
      Parameters
      ----------
      id: ID of the shuffle to pin
      output_partitions: Output partition IDs to pin
      pick: Function that picks a worker given a partition ID and sequence of worker
  
      .. note:
          This function assumes that the barrier task and the output tasks share
          the same worker restrictions.
      """
      existing: dict[Any, str] = {}
      shuffle_id = spec.id
      barrier = self.scheduler.tasks[barrier_key(shuffle_id)]
  
      if barrier.worker_restrictions:
          workers = list(barrier.worker_restrictions)
      else:
          workers = list(self.scheduler.workers)
  
      # Ensure homogeneous cluster utilization when there are multiple small,
      # independent shuffles going on at the same time, e.g. due to partial rechunking
      shift_by = next(self._shift_counter) % len(workers)
      workers = workers[shift_by:] + workers[:shift_by]
  
      # Check if this shuffle shares output tasks with a different shuffle that has
      # already been initialized and needs to be taken into account when
      # mapping output partitions to workers.
      # Naively, you could delete this whole paragraph and just call
      # spec.pick_worker; it would return two identical sets of results on both calls
      # of this method... until the set of available workers changes between the two
      # calls, which would cause misaligned shuffle outputs and a deadlock.
      seen = {barrier}
      for dependent in barrier.dependents:
          for possible_barrier in dependent.dependencies:
              if possible_barrier in seen:
                  continue
              seen.add(possible_barrier)
              if not (
                  other_barrier_key := id_from_key(possible_barrier.key)
              ):
                  continue
              if not (
                  shuffle := self.active_shuffles.get(other_barrier_key)
              ):
                  continue
              current_worker_for = shuffle.run_spec.worker_for
              # This is a fail-safe for future three-ways merges. At the moment there
              # should only ever be at most one other shuffle that shares output
              # tasks, so existing will always be empty.
              if existing:  # pragma: nocover
                  for shared_key in (
                      existing.keys() & current_worker_for.keys()
                  ):
                      if (
                          existing[shared_key]
                          != current_worker_for[shared_key]
                      ):
                          raise RuntimeError(
                              f"Failed to initialize shuffle {spec.id} because "
                              "it cannot align output partition mappings between "
                              f"existing shuffles {seen}. "
                              f"Mismatch encountered for output partition {shared_key!r}: "
                              f"{existing[shared_key]} != {current_worker_for[shared_key]}."
                          )
              existing.update(current_worker_for)
  
      worker_for = {}
      for partition in spec.output_partitions:
          if (worker := existing.get(partition, None)) is None:
              worker = spec.pick_worker(partition, workers)
          worker_for[partition] = worker
      return worker_for
  
imports: from __future__ import annotations
  import contextlib
  import itertools
  import logging
  from collections import defaultdict
  from typing import TYPE_CHECKING, Any
  from dask.typing import Key
  from distributed.diagnostics.plugin import SchedulerPlugin
  from distributed.metrics import time
  from distributed.protocol.pickle import dumps
  from distributed.protocol.serialize import ToPickle
  from distributed.shuffle._core import SchedulerShuffleState, ShuffleId, ShuffleRunSpec, ShuffleSpec, barrier_key, id_from_key
  from distributed.shuffle._worker_plugin import ShuffleWorkerPlugin
  from distributed.utils import log_errors
label_desc: Use image node_id 9 for calling the ShuffleSchedulerPlugin obj's underlying member method code with example usage: obj._calculate_worker_for(spec) and returns: worker_for