class_docstr: An subtask graph builder which build subtask graph for chunk graph based on passed band_resource.
  
  If push shuffle is used, this builder will validate predecessors orders consistency of ShuffleProxy between
  chunk graph and generated subtask graph.
docstr: Analyze chunk graph and generate subtask graph.
  
  Returns
  -------
  subtask_graph: SubtaskGraph
      Subtask graph.
docstr_tok: ["Analyze","chunk","graph","and","generate","subtask","graph",".","Returns","--","--","--","-","subtask_graph",":","SubtaskGraph","Subtask","graph","."]
code: def gen_subtask_graph(
      self, op_to_bands: Dict[str, BandType] = None
  ) -> SubtaskGraph:
      """
      Analyze chunk graph and generate subtask graph.
  
      Returns
      -------
      subtask_graph: SubtaskGraph
          Subtask graph.
      """
      # reassign worker when specified reassign_worker = True
      # or it's a reducer operands
      reassign_worker_ops = [
          chunk.op
          for chunk in self._chunk_graph
          if need_reassign_worker(chunk.op)
      ]
      start_ops = (
          list(self._iter_start_ops(self._chunk_graph))
          if len(self._chunk_graph) > 0
          else []
      )
  
      # assign start chunks
      to_assign_ops = start_ops + reassign_worker_ops
      assigner = self._graph_assigner_cls(
          self._chunk_graph, to_assign_ops, self._band_resource
      )
      # assign expect bands
      cur_assigns = {
          op.key: self._get_expect_band(op)
          for op in start_ops
          if op.expect_band is not None or op.expect_worker is not None
      }
      if op_to_bands:
          cur_assigns.update(op_to_bands)
      logger.debug(
          "Start to assign %s start chunks for task %s",
          len(start_ops),
          self._task.task_id,
      )
      chunk_to_bands = assigner.assign(cur_assigns=cur_assigns)
      logger.debug(
          "Assigned %s start chunks for task %s",
          len(start_ops),
          self._task.task_id,
      )
      # assign expect workers for those specified with `expect_worker` or `expect_band`
      # skip `start_ops`, which have been assigned before
      start_ops_set = set(start_ops)
      for chunk in self._chunk_graph:
          if chunk not in start_ops_set:
              if chunk.op.expect_band is not None:
                  chunk_to_bands[chunk] = chunk.op.expect_band
              elif chunk.op.expect_worker is not None:
                  chunk_to_bands[chunk] = self._to_band(
                      chunk.op.expect_worker
                  )
  
      # color nodes
      if self._fuse_enabled:
          logger.debug(
              "Start to fuse chunks for task %s", self._task.task_id
          )
          # sort start chunks in coloring as start_ops
          op_key_to_chunks = defaultdict(list)
          for chunk in self._chunk_graph:
              op_key_to_chunks[chunk.op.key].append(chunk)
          init_chunk_to_bands = dict()
          for start_op in start_ops:
              for start_chunk in op_key_to_chunks[start_op.key]:
                  init_chunk_to_bands[start_chunk] = chunk_to_bands[
                      start_chunk
                  ]
          if (
              self._has_shuffle
              and self._shuffle_fetch_type
              == ShuffleFetchType.FETCH_BY_INDEX
          ):
              # ensure no shuffle mapper chunks fused into same subtask.
              initial_same_color_num = 1
          else:
              initial_same_color_num = getattr(
                  self._config, "initial_same_color_num", None
              )
          coloring = Coloring(
              self._chunk_graph,
              list(self._band_resource),
              init_chunk_to_bands,
              initial_same_color_num=initial_same_color_num,
              as_broadcaster_successor_num=getattr(
                  self._config, "as_broadcaster_successor_num", None
              ),
          )
          chunk_to_colors = coloring.color()
      else:
          # if not fuse enabled, color all chunks with different colors
          op_to_colors = dict()
          chunk_to_colors = dict()
          color_gen = itertools.count()
          for c in self._chunk_graph.topological_iter():
              if c.op not in op_to_colors:
                  chunk_to_colors[c] = op_to_colors[c.op] = next(
                      color_gen
                  )
              else:
                  chunk_to_colors[c] = op_to_colors[c.op]
  
      color_to_chunks = defaultdict(list)
      subtask_fusion = SubtaskFusion()
      for chunk, color in chunk_to_colors.items():
          subtask_fusion.collect(chunk, color)
          if not isinstance(chunk.op, Fetch):
              color_to_chunks[color].append(chunk)
      subtask_fusion.fuse(
          self._chunk_graph,
          chunk_to_colors,
          color_to_chunks,
          chunk_to_bands,
      )
  
      # gen subtask graph
      subtask_graph = SubtaskGraph()
      chunk_to_fetch_chunk = dict()
      chunk_to_subtask = self._chunk_to_subtasks
      # states
      logic_key_to_subtasks = defaultdict(list)
      if self._shuffle_fetch_type == ShuffleFetchType.FETCH_BY_INDEX:
          for chunk in self._chunk_graph.topological_iter():
              if not isinstance(chunk.op, ShuffleProxy):
                  continue
              # Can't use `OperandStage.map` to find mappers directly, since `stage` of some operand
              # such as `DataFrameIndexAlign` are `OperandStage.map` but not a shuffle mapper sometimes.
              mapper_chunks = self._chunk_graph.predecessors(chunk)
              for mapper_chunk in mapper_chunks:
                  chunk_color = chunk_to_colors[mapper_chunk]
                  same_color_chunks = color_to_chunks[chunk_color]
                  mappers = [
                      c
                      for c in same_color_chunks
                      if c.op.stage == OperandStage.map
                      and any(
                          isinstance(succ.op, ShuffleProxy)
                          for succ in self._chunk_graph.iter_successors(
                              c
                          )
                      )
                  ]
                  if len(mappers) > 1:
                      # ensure every subtask contains only at most one mapper
                      for mapper in mappers:
                          same_color_chunks.remove(mapper)
                          mapper_color = coloring.next_color()
                          chunk_to_colors[mapper] = mapper_color
                          color_to_chunks[mapper_color] = [mapper]
  
      # build color dag
      color_set = set(chunk_to_colors.values())
      color_dag = DAG()
      for c in color_set:
          color_dag.add_node(c)
      for chunk, color in chunk_to_colors.items():
          for succ in self._chunk_graph.iter_successors(chunk):
              succ_color = chunk_to_colors[succ]
              if color != succ_color:
                  color_dag.add_edge(color, succ_color)
  
      # build subtasks in color topological order
      for color in color_dag.topological_iter():
          same_color_chunks = color_to_chunks[color]
          if all(isinstance(c.op, Fetch) for c in same_color_chunks):
              # all fetch ops, no need to gen subtask
              continue
          (
              subtask,
              inp_subtasks,
              is_shuffle_proxy,
          ) = self._gen_subtask_info(
              same_color_chunks,
              chunk_to_subtask,
              chunk_to_bands,
              chunk_to_fetch_chunk,
          )
          subtask_graph.add_node(subtask)
          if is_shuffle_proxy:
              subtask_graph.add_shuffle_proxy_subtask(subtask)
          logic_key_to_subtasks[subtask.logic_key].append(subtask)
          for inp_subtask in inp_subtasks:
              subtask_graph.add_edge(inp_subtask, subtask)
  
          for c in same_color_chunks:
              chunk_to_subtask[c] = subtask
              if (
                  self._map_reduce_id_to_infos is not None
                  and isinstance(c.op, ShuffleProxy)
              ):
                  self._gen_map_reduce_info(c, chunk_to_bands)
  
      for subtasks in logic_key_to_subtasks.values():
          for logic_index, subtask in enumerate(subtasks):
              subtask.logic_index = logic_index
              subtask.logic_parallelism = len(subtasks)
      return subtask_graph
  
imports: import itertools
  import logging
  from collections import deque, defaultdict
  from typing import Dict, List, Tuple, Type, Union
  from ....config import Config
  from ....core import ChunkGraph, ChunkType, enter_mode, DAG
  from ....core.operand import Fetch, VirtualOperand, LogicKeyGenerator, MapReduceOperand, OperandStage, ShuffleProxy, ShuffleFetchType
  from ....lib.ordered_set import OrderedSet
  from ....resource import Resource
  from ....typing import BandType, OperandType
  from ....utils import build_fetch, build_fetch_shuffle, tokenize
  from ...subtask import SubtaskGraph, Subtask
  from ..core import Task, new_task_id, MapReduceInfo
  from .assigner import AbstractGraphAssigner, GraphAssigner
  from .fusion import Coloring
label_desc: Use image node_id 10 for calling the GraphAnalyzer obj's underlying member method code with example usage: obj.gen_subtask_graph(op_to_bands) and returns: subtask_graph