class_docstr: TaskManager acts on the tasks fetched from the task queues.
  
  TaskManager instance can be used as a context manager:
class_docstr_tok: ["TaskManager","acts","on","the","tasks","fetched","from","the","task","queues",".","TaskManager","instance","can","be","used","as","a","context","manager",":"]
docstr: Constructs `TaskManager`.
  
  Args:
    mlmd_handle: ML metadata db connection.
    task_queue: Task queue.
    max_active_task_schedulers: Maximum number of task schedulers that can be
      active at once.
    max_dequeue_wait_secs: Maximum time to wait when dequeuing if the queue is
      empty.
    process_all_queued_tasks_before_exit: All existing items in the queues are
      processed before exiting the context manager. This is useful for
      deterministic behavior in tests.
docstr_tok: ["Constructs","`","TaskManager","`",".","Args",":","mlmd_handle",":","ML","metadata","db","connection",".","task_queue",":","Task","queue",".","max_active_task_schedulers",":","Maximum","number","of","task","schedulers","that","can","be","active","at","once",".","max_dequeue_wait_secs",":","Maximum","time","to","wait","when","dequeuing","if","the","queue","is","empty",".","process_all_queued_tasks_before_exit",":","All","existing","items","in","the","queues","are","processed","before","exiting","the","context","manager",".","This","is","useful","for","deterministic","behavior","in","tests","."]
code: def __init__(
      self,
      mlmd_handle: metadata.Metadata,
      task_queue: tq.TaskQueue,
      max_active_task_schedulers: int,
      max_dequeue_wait_secs: float = _MAX_DEQUEUE_WAIT_SECS,
      process_all_queued_tasks_before_exit: bool = False,
  ):
      """Constructs `TaskManager`.
  
      Args:
        mlmd_handle: ML metadata db connection.
        task_queue: Task queue.
        max_active_task_schedulers: Maximum number of task schedulers that can be
          active at once.
        max_dequeue_wait_secs: Maximum time to wait when dequeuing if the queue is
          empty.
        process_all_queued_tasks_before_exit: All existing items in the queues are
          processed before exiting the context manager. This is useful for
          deterministic behavior in tests.
      """
      self._mlmd_handle = mlmd_handle
      self._task_queue = task_queue
      self._max_dequeue_wait_secs = max_dequeue_wait_secs
      self._process_all_queued_tasks_before_exit = (
          process_all_queued_tasks_before_exit
      )
  
      self._tm_lock = threading.Lock()
      self._stop_event = threading.Event()
      self._scheduler_by_node_uid: Dict[
          task_lib.NodeUid, _SchedulerWrapper
      ] = {}
      self._active_scheduler_counter = _ActiveSchedulerCounter()
  
      # Async executor for the main task management thread.
      self._main_executor = futures.ThreadPoolExecutor(
          max_workers=1,
          thread_name_prefix="orchestrator_task_manager_main",
      )
      self._main_future = None
      self._max_active_task_schedulers = max_active_task_schedulers
  
      self._pending_schedulers: List[_SchedulerWrapper] = []
  
      # Async executor for task schedulers. We have 1 extra worker so that task
      # schedulers being canceled can be run without being blocked by active ones.
      self._ts_executor = futures.ThreadPoolExecutor(
          max_workers=self._max_active_task_schedulers + 1,
          thread_name_prefix="orchestrator_active_task_schedulers",
      )
      self._ts_futures = set()
  
imports: from concurrent import futures
  import datetime
  import sys
  import textwrap
  import threading
  import time
  import typing
  from typing import Dict, List, Optional
  from absl import logging
  import pytz
  from tfx.orchestration import data_types_utils
  from tfx.orchestration import metadata
  from tfx.orchestration.experimental.core import constants
  from tfx.orchestration.experimental.core import mlmd_state
  from tfx.orchestration.experimental.core import pipeline_state
  from tfx.orchestration.experimental.core import post_execution_utils
  from tfx.orchestration.experimental.core import task
  from tfx.orchestration.experimental.core import task_queue
  from tfx.orchestration.experimental.core import task_scheduler
  from tfx.utils import status
  from ml_metadata.proto import metadata_store_pb2
label_desc: Use image node_id 1 to create a new TaskManager object  with example: obj = TaskManager(mlmd_handle, task_queue, max_active_task_schedulers, max_dequeue_wait_secs, process_all_queued_tasks_before_exit)