class_docstr: Scheduler extension for spans support
class_docstr_tok: ["Scheduler","extension","for","spans","support"]
docstr: Acknowledge the existence of runnable tasks on the scheduler. These may
  either be new tasks, tasks that were previously unrunnable, or tasks that were
  already fed into this method already.
  
  Attach newly observed tasks to either the desired span or to ("default", ).
  Update TaskGroup.span_id and wipe TaskState.annotations["span"].
  
  Returns
  -------
  Updated 'span' annotations: {key: {"name": (..., ...), "ids": (..., ...)}}
docstr_tok: ["Acknowledge","the","existence","of","runnable","tasks","on","the","scheduler",".","These","may","either","be","new","tasks",",","tasks","that","were","previously","unrunnable",",","or","tasks","that","were","already","fed","into","this","method","already",".","Attach","newly","observed","tasks","to","either","the","desired","span","or","to","(","``","default","''",",",")",".","Update","TaskGroup.span_id","and","wipe","TaskState.annotations","[","``","span","''","]",".","Returns","--","--","--","-","Updated","'span","'","annotations",":","{","key",":","{","``","name","''",":","(","...",",","...",")",",","``","ids","''",":","(","...",",","...",")","}","}"]
code: def observe_tasks(
      self,
      tss: Iterable[scheduler_module.TaskState],
      code: tuple[SourceCode, ...],
  ) -> dict[Key, dict]:
      """Acknowledge the existence of runnable tasks on the scheduler. These may
      either be new tasks, tasks that were previously unrunnable, or tasks that were
      already fed into this method already.
  
      Attach newly observed tasks to either the desired span or to ("default", ).
      Update TaskGroup.span_id and wipe TaskState.annotations["span"].
  
      Returns
      -------
      Updated 'span' annotations: {key: {"name": (..., ...), "ids": (..., ...)}}
      """
      out = {}
      default_span = None
  
      for ts in tss:
          if ts.annotations is None:
              ts.annotations = dict()
          # You may have different tasks belonging to the same TaskGroup but to
          # different spans. If that happens, arbitrarily force everything onto the
          # span of the earliest encountered TaskGroup.
          tg = ts.group
          if tg.span_id:
              span = self.spans[tg.span_id]
          else:
              ann = ts.annotations.get("span")
              if ann:
                  span = self._ensure_span(ann["name"], ann["ids"])
              else:
                  if not default_span:
                      default_span = self._ensure_default_span()
                  span = default_span
  
              tg.span_id = span.id
              span.groups.add(tg)
  
          if code:
              span._code[code] = None
  
          # The span may be completely different from the one referenced by the
          # annotation, due to the TaskGroup collision issue explained above.
          if ann := span.annotation:
              ts.annotations["span"] = out[ts.key] = ann
          else:
              ts.annotations.pop("span", None)
  
      return out
  
imports: from __future__ import annotations
  import uuid
  import weakref
  from collections import defaultdict
  from collections.abc import Hashable, Iterable, Iterator
  from contextlib import contextmanager
  from itertools import islice
  from typing import TYPE_CHECKING, Any
  import dask.config
  from dask.typing import Key
  from distributed.collections import sum_mappings
  from distributed.itertools import ffill
  from distributed.metrics import time
label_desc: Use image node_id 2 for calling the SpansSchedulerExtension obj's underlying member method code with example usage: obj.observe_tasks(tss, code) and returns: out