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: Output tasks are created without worker restrictions and run once with the
  only purpose of setting the worker restriction and then raising Reschedule, and
  then running again properly on the correct worker. It would be non-trivial to
  set the worker restriction before they're first run due to potential task
  fusion.
  
  Most times, this lack of initial restrictions would cause output tasks to be
  labelled as rootish on their first (very fast) run, which in turn would break
  the design assumption that the worker-side queue of rootish tasks will last long
  enough to cover the round-trip to the scheduler to receive more tasks, which in
  turn would cause a measurable slowdown on the overall runtime of the shuffle
  operation.
  
  This method ensures that, given M output tasks and N workers, each worker-side
  queue is pre-loaded with M/N output tasks which can be flushed very fast as
  they all raise Reschedule() in quick succession.
  
  See Also
  --------
  ShuffleRun._ensure_output_worker
docstr_tok: ["Output","tasks","are","created","without","worker","restrictions","and","run","once","with","the","only","purpose","of","setting","the","worker","restriction","and","then","raising","Reschedule",",","and","then","running","again","properly","on","the","correct","worker",".","It","would","be","non-trivial","to","set","the","worker","restriction","before","they","'re","first","run","due","to","potential","task","fusion",".","Most","times",",","this","lack","of","initial","restrictions","would","cause","output","tasks","to","be","labelled","as","rootish","on","their","first","(","very","fast",")","run",",","which","in","turn","would","break","the","design","assumption","that","the","worker-side","queue","of","rootish","tasks","will","last","long","enough","to","cover","the","round-trip","to","the","scheduler","to","receive","more","tasks",",","which","in","turn","would","cause","a","measurable","slowdown","on","the","overall","runtime","of","the","shuffle","operation",".","This","method","ensures","that",",","given","M","output","tasks","and","N","workers",",","each","worker-side","queue","is","pre-loaded","with","M\/N","output","tasks","which","can","be","flushed","very","fast","as","they","all","raise","Reschedule","(",")","in","quick","succession",".","See","Also","--","--","--","--","ShuffleRun._ensure_output_worker"]
code: def _ensure_output_tasks_are_non_rootish(
      self, spec: ShuffleSpec
  ) -> None:
      """Output tasks are created without worker restrictions and run once with the
      only purpose of setting the worker restriction and then raising Reschedule, and
      then running again properly on the correct worker. It would be non-trivial to
      set the worker restriction before they're first run due to potential task
      fusion.
  
      Most times, this lack of initial restrictions would cause output tasks to be
      labelled as rootish on their first (very fast) run, which in turn would break
      the design assumption that the worker-side queue of rootish tasks will last long
      enough to cover the round-trip to the scheduler to receive more tasks, which in
      turn would cause a measurable slowdown on the overall runtime of the shuffle
      operation.
  
      This method ensures that, given M output tasks and N workers, each worker-side
      queue is pre-loaded with M/N output tasks which can be flushed very fast as
      they all raise Reschedule() in quick succession.
  
      See Also
      --------
      ShuffleRun._ensure_output_worker
      """
      barrier = self.scheduler.tasks[barrier_key(spec.id)]
      for dependent in barrier.dependents:
          dependent._rootish = False
  
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 10 for calling the ShuffleSchedulerPlugin obj's underlying member method code with example usage: obj._ensure_output_tasks_are_non_rootish(spec) without return types