docstr: Tests async pipeline task generation.
  
  Args:
    use_task_queue: If task queue is enabled, new tasks are only generated if
      a task with the same task_id does not already exist in the queue.
      `use_task_queue=False` is useful to test the case of task generation
      when task queue is empty (for eg: due to orchestrator restart).
docstr_tok: ["Tests","async","pipeline","task","generation",".","Args",":","use_task_queue",":","If","task","queue","is","enabled",",","new","tasks","are","only","generated","if","a","task","with","the","same","task_id","does","not","already","exist","in","the","queue",".","`","use_task_queue=False","`","is","useful","to","test","the","case","of","task","generation","when","task","queue","is","empty","(","for","eg",":","due","to","orchestrator","restart",")","."]
code: def test_task_generation_for_each(self, use_task_queue):
      """Tests async pipeline task generation.
  
      Args:
        use_task_queue: If task queue is enabled, new tasks are only generated if
          a task with the same task_id does not already exist in the queue.
          `use_task_queue=False` is useful to test the case of task generation
          when task queue is empty (for eg: due to orchestrator restart).
      """
      # Simulate that ExampleGen run twice for 2 spans.
      test_utils.fake_example_gen_run(
          self._mlmd_connection, self._example_gen, 1, 1
      )
      test_utils.fake_example_gen_run(
          self._mlmd_connection, self._example_gen, 2, 1
      )
  
      # Generate once, two executions for Transform is generated.
      [
          update_example_gen_task,
          update_transform_task,
          exec_transform_task,
          update_trainer_task,
      ] = self._generate_and_test(
          use_task_queue,
          num_initial_executions=2,
          num_tasks_generated=4,
          num_new_executions=2,
          num_active_executions=2,
          expected_exec_nodes=[self._transform],
      )
      self.assertIsInstance(
          update_example_gen_task, task_lib.UpdateNodeStateTask
      )
      self.assertEqual(
          pstate.NodeState.RUNNING, update_example_gen_task.state
      )
      self.assertIsInstance(
          update_transform_task, task_lib.UpdateNodeStateTask
      )
      self.assertEqual(
          pstate.NodeState.RUNNING, update_transform_task.state
      )
      self.assertIsInstance(exec_transform_task, task_lib.ExecNodeTask)
      self.assertIsInstance(
          update_trainer_task, task_lib.UpdateNodeStateTask
      )
      self.assertEqual(
          pstate.NodeState.STARTED, update_trainer_task.state
      )
  
      self._mock_service_job_manager.ensure_node_services.assert_has_calls(
          [
              mock.call(mock.ANY, self._example_gen.node_info.id, ""),
              mock.call(mock.ANY, self._transform.node_info.id),
          ]
      )
  
      # Mark one of the Transform executions complete.
      self._finish_node_execution(use_task_queue, exec_transform_task)
  
      # Generate again, an execution for Trainer is generated.
      [
          update_transform_task,
          exec_transform_task,
          update_trainer_task,
          exec_trainer_task,
      ] = self._generate_and_test(
          use_task_queue,
          num_initial_executions=4,
          num_tasks_generated=4,
          num_new_executions=1,
          num_active_executions=2,
          expected_exec_nodes=[self._transform, self._trainer],
      )
      self.assertIsInstance(
          update_transform_task, task_lib.UpdateNodeStateTask
      )
      self.assertEqual(
          pstate.NodeState.RUNNING, update_transform_task.state
      )
      self.assertIsInstance(exec_transform_task, task_lib.ExecNodeTask)
      self.assertIsInstance(
          update_trainer_task, task_lib.UpdateNodeStateTask
      )
      self.assertEqual(
          pstate.NodeState.RUNNING, update_trainer_task.state
      )
      self.assertIsInstance(exec_trainer_task, task_lib.ExecNodeTask)
  
      # Mark the Transform execution complete.
      self._finish_node_execution(use_task_queue, exec_transform_task)
      # Mark the Trainer execution complete.
      self._finish_node_execution(use_task_queue, exec_trainer_task)
  
      # Generate again, another execution for Trainer is generated.
      [
          update_transform_task,
          update_trainer_task_1,
          update_trainer_task_2,
          exec_trainer_task,
      ] = self._generate_and_test(
          use_task_queue,
          num_initial_executions=5,
          num_tasks_generated=4,
          num_new_executions=1,
          num_active_executions=1,
      )
      self.assertIsInstance(
          update_transform_task, task_lib.UpdateNodeStateTask
      )
      self.assertEqual(
          pstate.NodeState.STARTED, update_transform_task.state
      )
      self.assertIsInstance(
          update_trainer_task_1, task_lib.UpdateNodeStateTask
      )
      self.assertEqual(
          pstate.NodeState.STARTED, update_trainer_task_1.state
      )
      self.assertIsInstance(
          update_trainer_task_2, task_lib.UpdateNodeStateTask
      )
      self.assertEqual(
          pstate.NodeState.RUNNING, update_trainer_task_2.state
      )
      self.assertIsInstance(exec_trainer_task, task_lib.ExecNodeTask)
  
      # Mark the trainer execution complete.
      self._finish_node_execution(use_task_queue, exec_trainer_task)
  
      # Finally, update Trainer's state to STARTED.
      [update_trainer_task] = self._generate_and_test(
          use_task_queue,
          num_initial_executions=6,
          num_tasks_generated=1,
          num_new_executions=0,
          num_active_executions=0,
      )
      self.assertIsInstance(
          update_trainer_task, task_lib.UpdateNodeStateTask
      )
      self.assertEqual(
          pstate.NodeState.STARTED, update_trainer_task.state
      )
  
      if use_task_queue:
          self.assertTrue(self._task_queue.is_empty())
  
imports: import os
  from absl.testing import parameterized
  from absl.testing.absltest import mock
  import tensorflow
  from tfx.orchestration import node_proto_view
  from tfx.orchestration.experimental.core import async_pipeline_task_gen
  from tfx.orchestration.experimental.core import pipeline_state
  from tfx.orchestration.experimental.core import service_jobs
  from tfx.orchestration.experimental.core import task
  from tfx.orchestration.experimental.core import task_gen_utils
  from tfx.orchestration.experimental.core import task_queue
  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.orchestration.portable.input_resolution import exceptions
  from tfx.utils import status
label_desc: Use image node_id 6 for calling the AsyncPipelineTaskGeneratorTest obj's underlying member method code with example usage: obj.test_task_generation_for_each(use_task_queue) without return types