docstr: Tests async pipeline task generation for backfill.
docstr_tok: ["Tests","async","pipeline","task","generation","for","backfill","."]
code: def test_backfill(self, throw_error):
      """Tests async pipeline task generation for backfill."""
      use_task_queue = True
      # Simulate that ExampleGen has already completed successfully.
      test_utils.fake_example_gen_run(
          self._mlmd_connection, self._example_gen, 1, 1
      )
  
      # Generate once.
      [
          update_example_gen_task,
          update_transform_task,
          exec_transform_task,
          update_trainer_task,
      ] = self._generate_and_test(
          use_task_queue,
          num_initial_executions=1,
          num_tasks_generated=4,
          num_new_executions=1,
          num_active_executions=1,
          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 transform execution complete.
      self._finish_node_execution(use_task_queue, exec_transform_task)
  
      # Trainer execution task should be generated next.
      [
          update_transform_task,
          update_trainer_task,
          exec_trainer_task,
      ] = self._generate_and_test(
          use_task_queue,
          num_initial_executions=2,
          num_tasks_generated=3,
          num_new_executions=1,
          num_active_executions=1,
          expected_exec_nodes=[self._trainer],
      )
      self.assertIsInstance(
          update_transform_task, task_lib.UpdateNodeStateTask
      )
      self.assertEqual(
          pstate.NodeState.STARTED, update_transform_task.state
      )
      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 trainer execution complete.
      self._finish_node_execution(use_task_queue, exec_trainer_task)
  
      # Only UpdateNodeStateTask are generated as there are no new inputs.
      [update_trainer_task] = self._generate_and_test(
          use_task_queue,
          num_initial_executions=3,
          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
      )
  
      # Put Transform in backfill mode.
      with pstate.PipelineState.load(
          self._mlmd_connection,
          task_lib.PipelineUid.from_pipeline(self._pipeline),
      ) as pipeline_state:
          transform_node = task_lib.NodeUid.from_node(
              self._pipeline, node_proto_view.get_view(self._transform)
          )
          with pipeline_state.node_state_update_context(
              transform_node
          ) as node_state:
              node_state.update(
                  pstate.NodeState.STARTING,
                  backfill_token="backfill-20221215-180505-123456",
              )
      if throw_error:
          # Mock the InputResolutionError when generate_resolved_info is called.
          with mock.patch.object(
              task_gen_utils, "generate_resolved_info", autospec=True
          ) as mock_generate_resolved_info:
              mock_generate_resolved_info.side_effect = (
                  exceptions.InputResolutionError()
              )
              expected_error_msg = (
                  "Backfill of node my_transform failed Error: failure to resolve"
                  " inputs; node uid:"
                  " NodeUid(pipeline_uid=PipelineUid(pipeline_id='my_pipeline',"
                  " pipeline_run_id=None), node_id='my_transform'); error:"
                  " ['tfx.orchestration.portable.input_resolution.exceptions.InputResolutionError\\n']"
              )
  
              [
                  update_transform_task,
                  failed_transform_task,
                  update_trainer_task,
              ] = self._generate_and_test(
                  use_task_queue,
                  num_initial_executions=3,
                  num_tasks_generated=3,
                  num_new_executions=0,
                  num_active_executions=0,
                  expected_exec_nodes=[],
              )
              self.assertIsInstance(
                  update_transform_task, task_lib.UpdateNodeStateTask
              )
              self.assertEqual(
                  pstate.NodeState.STARTED, update_transform_task.state
              )
              self.assertEqual(
                  "backfill-20221215-180505-123456",
                  update_transform_task.backfill_token,
              )
              self.assertIsInstance(
                  failed_transform_task, task_lib.UpdateNodeStateTask
              )
              self.assertEqual(
                  pstate.NodeState.FAILED, failed_transform_task.state
              )
              self.assertEqual(
                  status_lib.Code.FAILED_PRECONDITION,
                  failed_transform_task.status.code,
              )
              self.assertEqual(
                  expected_error_msg,
                  failed_transform_task.status.message,
              )
              self.assertEqual(
                  "",
                  failed_transform_task.backfill_token,
              )
              self.assertIsInstance(
                  update_trainer_task, task_lib.UpdateNodeStateTask
              )
              self.assertEqual(
                  pstate.NodeState.STARTED, update_trainer_task.state
              )
          return
      # Transform tasks should be generated as it will start a backfill.
      # Trainer will just be updated to STARTED state, since there are no new
      # inputs.
      [
          update_transform_to_started_task,
          update_transform_to_running_task,
          exec_transform_task,
      ] = self._generate_and_test(
          use_task_queue,
          num_initial_executions=3,
          num_tasks_generated=3,
          num_new_executions=1,
          num_active_executions=1,
          expected_exec_nodes=[self._transform],
      )
      self.assertIsInstance(
          update_transform_to_started_task, task_lib.UpdateNodeStateTask
      )
      self.assertEqual(
          pstate.NodeState.STARTED,
          update_transform_to_started_task.state,
      )
      self.assertEqual(
          "backfill-20221215-180505-123456",
          update_transform_to_started_task.backfill_token,
      )
      self.assertIsInstance(
          update_transform_to_running_task, task_lib.UpdateNodeStateTask
      )
      self.assertEqual(
          pstate.NodeState.RUNNING,
          update_transform_to_running_task.state,
      )
      self.assertEqual(
          "backfill-20221215-180505-123456",
          update_transform_to_running_task.backfill_token,
      )
      self.assertIsInstance(exec_transform_task, task_lib.ExecNodeTask)
  
      # Mark transform execution complete.
      self._finish_node_execution(use_task_queue, exec_transform_task)
  
      # Transform should be stopped, since the backfill is complete.
      # Trainer should be triggered again due to transform producing new output.
      [
          update_transform_task,
          update_trainer_task,
          exec_trainer_task,
      ] = self._generate_and_test(
          use_task_queue,
          num_initial_executions=4,
          num_tasks_generated=3,
          num_new_executions=1,
          num_active_executions=1,
          expected_exec_nodes=[self._trainer],
      )
      self.assertIsInstance(
          update_transform_task, task_lib.UpdateNodeStateTask
      )
      self.assertEqual(
          pstate.NodeState.STOPPED, update_transform_task.state
      )
      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)
  
      # Trainer completes, goes back into STARTED state.
      self._finish_node_execution(use_task_queue, exec_trainer_task)
      [update_trainer_task] = self._generate_and_test(
          use_task_queue,
          num_initial_executions=5,
          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
      )
  
      # Put Transform in backfill mode with the same token as before.
      with pstate.PipelineState.load(
          self._mlmd_connection,
          task_lib.PipelineUid.from_pipeline(self._pipeline),
      ) as pipeline_state:
          transform_node = task_lib.NodeUid.from_node(
              self._pipeline, node_proto_view.get_view(self._transform)
          )
          with pipeline_state.node_state_update_context(
              transform_node
          ) as node_state:
              node_state.update(
                  pstate.NodeState.STARTING,
                  backfill_token="backfill-20221215-180505-123456",
              )
  
      # Transform should stop immediately, since it sees the previous backfill
      # execution.
      [
          update_transform_to_started_task,
          update_transform_to_stopped_task,
      ] = self._generate_and_test(
          use_task_queue,
          num_initial_executions=5,
          num_tasks_generated=2,
          num_new_executions=0,
          num_active_executions=0,
      )
      self.assertIsInstance(
          update_transform_to_started_task, task_lib.UpdateNodeStateTask
      )
      self.assertEqual(
          pstate.NodeState.STARTED,
          update_transform_to_started_task.state,
      )
      self.assertIsInstance(
          update_transform_to_stopped_task, task_lib.UpdateNodeStateTask
      )
      self.assertEqual(
          pstate.NodeState.STOPPED,
          update_transform_to_stopped_task.state,
      )
  
      # Put Transform in backfill mode with a new token.
      with pstate.PipelineState.load(
          self._mlmd_connection,
          task_lib.PipelineUid.from_pipeline(self._pipeline),
      ) as pipeline_state:
          transform_node = task_lib.NodeUid.from_node(
              self._pipeline, node_proto_view.get_view(self._transform)
          )
          with pipeline_state.node_state_update_context(
              transform_node
          ) as node_state:
              node_state.update(
                  pstate.NodeState.STARTING,
                  backfill_token="backfill-20221215-192233-234567",
              )
  
      # Transform tasks should be generated as it will start a new backfill.
      [
          update_transform_to_started_task,
          update_transform_to_running_task,
          exec_transform_task,
      ] = self._generate_and_test(
          use_task_queue,
          num_initial_executions=5,
          num_tasks_generated=3,
          num_new_executions=1,
          num_active_executions=1,
          expected_exec_nodes=[self._transform],
      )
      self.assertIsInstance(
          update_transform_to_started_task, task_lib.UpdateNodeStateTask
      )
      self.assertEqual(
          pstate.NodeState.STARTED,
          update_transform_to_started_task.state,
      )
      self.assertEqual(
          "backfill-20221215-192233-234567",
          update_transform_to_started_task.backfill_token,
      )
      self.assertIsInstance(
          update_transform_to_running_task, task_lib.UpdateNodeStateTask
      )
      self.assertEqual(
          pstate.NodeState.RUNNING,
          update_transform_to_running_task.state,
      )
      self.assertEqual(
          "backfill-20221215-192233-234567",
          update_transform_to_running_task.backfill_token,
      )
      self.assertIsInstance(exec_transform_task, task_lib.ExecNodeTask)
  
      # Mark transform execution complete, but FAILED.
      self._finish_node_execution(
          use_task_queue, exec_transform_task, success=False
      )
  
      # In backfill mode, we don't retry failed executions, so Transform should
      # be stopped, since the backfill is complete.
      [update_transform_task] = self._generate_and_test(
          use_task_queue,
          num_initial_executions=6,
          num_tasks_generated=1,
          num_new_executions=0,
          num_active_executions=0,
          expected_exec_nodes=[],
      )
      self.assertIsInstance(
          update_transform_task, task_lib.UpdateNodeStateTask
      )
      self.assertEqual(
          pstate.NodeState.STOPPED, update_transform_task.state
      )
  
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 11 for calling the AsyncPipelineTaskGeneratorTest obj's underlying member method code with example usage: obj.test_backfill(throw_error) without return types