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"]
code: def get_or_create(
      self,
      # FIXME: This should never be ToPickle[ShuffleSpec]
      spec: ShuffleSpec | ToPickle[ShuffleSpec],
      key: Key,
      worker: str,
  ) -> ToPickle[ShuffleRunSpec]:
      # FIXME: Sometimes, this doesn't actually get pickled
      if isinstance(spec, ToPickle):
          spec = spec.data
      try:
          return self.get(spec.id, worker)
      except KeyError:
          # FIXME: The current implementation relies on the barrier task to be
          # known by its name. If the name has been mangled, we cannot guarantee
          # that the shuffle works as intended and should fail instead.
          self._raise_if_barrier_unknown(spec.id)
          self._raise_if_task_not_processing(key)
          worker_for = self._calculate_worker_for(spec)
          self._ensure_output_tasks_are_non_rootish(spec)
          state = spec.create_new_run(
              worker_for=worker_for,
              span_id=self.scheduler.tasks[key].group.span_id,
          )
          self.active_shuffles[spec.id] = state
          self._shuffles[spec.id].add(state)
          state.participating_workers.add(worker)
          logger.warning(
              "Shuffle %s initialized by task %r executed on worker %s",
              spec.id,
              key,
              worker,
          )
          return ToPickle(state.run_spec)
  
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 6 for calling the ShuffleSchedulerPlugin obj's underlying member method code with example usage: obj.get_or_create(spec, key, worker) and returns: self, ToPickle