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 restrict_task(
      self, id: ShuffleId, run_id: int, key: Key, worker: str
  ) -> dict:
      shuffle = self.active_shuffles[id]
      if shuffle.run_id > run_id:
          return {
              "status": "error",
              "message": f"Request stale, expected {run_id=} for {shuffle}",
          }
      elif shuffle.run_id < run_id:
          return {
              "status": "error",
              "message": f"Request invalid, expected {run_id=} for {shuffle}",
          }
      ts = self.scheduler.tasks[key]
      self._set_restriction(ts, worker)
      return {"status": "OK"}
  
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 3 for calling the ShuffleSchedulerPlugin obj's underlying member method code with example usage: obj.restrict_task(id, run_id, key, worker) and returns: dict, dict, dict