code: def test_task_scheduler_cancellations_are_prioritized(
      self, unused_mock
  ) -> None:
      collector = _Collector()
  
      # Register a fake task scheduler.
      ts.TaskSchedulerRegistry.register(
          self._type_url,
          functools.partial(
              _FakeTaskScheduler,
              block_nodes={"Trainer", "Transform"},
              collector=collector,
          ),
      )
  
      task_queue = tq.TaskQueue()
      with self._task_manager(
          task_queue, max_active_task_schedulers=2
      ) as task_manager:
  
          def _wait_for(
              num_pending, num_active, num_ts_futures, timeout=30.0
          ) -> None:
              start_time = time.time()
              while time.time() - start_time <= timeout:
                  if (
                      len(task_manager._pending_schedulers)
                      == num_pending
                      and task_manager._active_scheduler_counter.count()
                      == num_active
                      and len(task_manager._ts_futures)
                      == num_ts_futures
                  ):
                      return
                  time.sleep(0.1)
              raise TimeoutError(
                  f"Timeout waiting for {num_pending} pending and {num_active} task"
                  " schedulers."
              )
  
          # Enqueue 4 tasks.
          task_queue.enqueue(
              _test_exec_node_task(
                  "Trainer", "test-pipeline", pipeline=self._pipeline
              )
          )
          task_queue.enqueue(
              _test_exec_node_task(
                  "Transform", "test-pipeline", pipeline=self._pipeline
              )
          )
          task_queue.enqueue(
              _test_exec_node_task(
                  "Evaluator", "test-pipeline", pipeline=self._pipeline
              )
          )
          task_queue.enqueue(
              _test_exec_node_task(
                  "Pusher", "test-pipeline", pipeline=self._pipeline
              )
          )
  
          # Since max_active_task_schedulers=2, the first two tasks should be active
          # and the other two pending.
          _wait_for(num_pending=2, num_active=2, num_ts_futures=2)
  
          self.assertEqual(
              ["Evaluator", "Pusher"],
              [
                  s.task_scheduler.task.node_uid.node_id
                  for s in task_manager._pending_schedulers
              ],
          )
          task_queue.enqueue(
              _test_cancel_node_task("Evaluator", "test-pipeline")
          )
          task_queue.enqueue(
              _test_cancel_node_task("Pusher", "test-pipeline")
          )
  
          # Cancellations should be prioritized and go through even when
          # `max_active_task_schedulers` slots are occupied.
          _wait_for(num_pending=0, num_active=2, num_ts_futures=2)
  
          task_queue.enqueue(
              _test_cancel_node_task("Trainer", "test-pipeline")
          )
          task_queue.enqueue(
              _test_cancel_node_task("Transform", "test-pipeline")
          )
          _wait_for(num_pending=0, num_active=0, num_ts_futures=0)
  
      self.assertTrue(task_manager.done())
      self.assertIsNone(task_manager.exception())
      self.assertCountEqual(
          ["Trainer", "Transform", "Evaluator", "Pusher"],
          [task.node_uid.node_id for task in collector.scheduled_tasks],
      )
  
imports: import contextlib
  import functools
  import os
  import threading
  import time
  from absl import logging
  from absl.testing.absltest import mock
  import tensorflow
  from tfx.orchestration import data_types_utils
  from tfx.orchestration import metadata
  from tfx.orchestration.experimental.core import async_pipeline_task_gen
  from tfx.orchestration.experimental.core import constants
  from tfx.orchestration.experimental.core import pipeline_state
  from tfx.orchestration.experimental.core import post_execution_utils
  from tfx.orchestration.experimental.core import service_jobs
  from tfx.orchestration.experimental.core import task
  from tfx.orchestration.experimental.core import task_manager
  from tfx.orchestration.experimental.core import task_queue
  from tfx.orchestration.experimental.core import task_scheduler
  from tfx.orchestration.experimental.core import test_utils
  from tfx.orchestration.experimental.core.testing import test_async_pipeline
  from tfx.orchestration import mlmd_connection_manager
  from tfx.proto.orchestration import execution_result_pb2
  from tfx.proto.orchestration import pipeline_pb2
  from tfx.utils import status
  from ml_metadata.proto import metadata_store_pb2
label_desc: Use image node_id 5 for calling the TaskManagerTest obj's underlying member method code with example usage: obj.test_task_scheduler_cancellations_are_prioritized(unused_mock) without return types