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: Clean up scheduler and worker state once a shuffle becomes inactive.
docstr_tok: ["Clean","up","scheduler","and","worker","state","once","a","shuffle","becomes","inactive","."]
code: def transition(
      self,
      key: Key,
      start: TaskStateState,
      finish: TaskStateState,
      *args: Any,
      stimulus_id: str,
      **kwargs: Any,
  ) -> None:
      """Clean up scheduler and worker state once a shuffle becomes inactive."""
      if finish not in ("released", "forgotten"):
          return
      shuffle_id = id_from_key(key)
      if not shuffle_id:
          return
  
      if shuffle := self.active_shuffles.get(shuffle_id):
          self._fail_on_workers(shuffle, message=f"{shuffle} forgotten")
          self._clean_on_scheduler(shuffle_id, stimulus_id=stimulus_id)
          logger.debug(
              "Shuffle %s forgotten because task %r transitioned to %s due to "
              "stimulus '%s'",
              shuffle_id,
              key,
              finish,
              stimulus_id,
          )
  
      if finish == "forgotten":
          shuffles = self._shuffles.pop(shuffle_id, set())
          for shuffle in shuffles:
              if shuffle._archived_by:
                  archived = self._archived_by_stimulus[
                      shuffle._archived_by
                  ]
                  archived.remove(shuffle)
                  if not archived:
                      del self._archived_by_stimulus[
                          shuffle._archived_by
                      ]
  
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 16 for calling the ShuffleSchedulerPlugin obj's underlying member method code with example usage: obj.transition(key, start, finish) without return types