def test_generate_tasks_async_active_pipelines(self, mock_async_task_gen,
mock_sync_task_gen):
with self._mlmd_connection as m:
# One active pipeline.
pipeline1 = _test_pipeline('pipeline1')
pipeline_ops.initiate_pipeline_start(m, pipeline1)
# Another active pipeline (with previously completed execution).
pipeline2 = _test_pipeline('pipeline2')
execution2 = pipeline_ops.initiate_pipeline_start(m, pipeline2)
execution2.last_known_state = metadata_store_pb2.Execution.COMPLETE
m.store.put_executions([execution2])
execution2 = pipeline_ops.initiate_pipeline_start(m, pipeline2)
# Inactive pipelines should be ignored.
pipeline3 = _test_pipeline('pipeline3')
execution3 = pipeline_ops.initiate_pipeline_start(m, pipeline3)
execution3.last_known_state = metadata_store_pb2.Execution.COMPLETE
m.store.put_executions([execution3])
# For active pipelines pipeline1 and pipeline2, there are a couple of
# active executions.
def _exec_node_tasks():
for pipeline_id in ('pipeline1', 'pipeline2'):
yield [
test_utils.create_exec_node_task(
node_uid=task_lib.
NodeUid(pipeline_uid=task_lib.PipelineUid(
pipeline_id=pipeline_id, pipeline_run_id=None),
node_id='Transform')),
test_utils.create_exec_node_task(
node_uid=task_lib.
NodeUid(pipeline_uid=task_lib.PipelineUid(
pipeline_id=pipeline_id, pipeline_run_id=None),
node_id='Trainer'))
]
mock_async_task_gen.return_value.generate.side_effect = _exec_node_tasks(
)
task_queue = tq.TaskQueue()
pipeline_ops.generate_tasks(m, task_queue)
self.assertEqual(
2, mock_async_task_gen.return_value.generate.call_count)
mock_sync_task_gen.assert_not_called()
# Verify that tasks are enqueued in the expected order.
for node_id in ('Transform', 'Trainer'):
task = task_queue.dequeue()
task_queue.task_done(task)
self.assertTrue(task_lib.is_exec_node_task(task))
self.assertEqual(node_id, task.node_uid.node_id)
self.assertEqual('pipeline1',
task.node_uid.pipeline_uid.pipeline_id)
for node_id in ('Transform', 'Trainer'):
task = task_queue.dequeue()
task_queue.task_done(task)
self.assertTrue(task_lib.is_exec_node_task(task))
self.assertEqual(node_id, task.node_uid.node_id)
self.assertEqual('pipeline2',
task.node_uid.pipeline_uid.pipeline_id)
self.assertTrue(task_queue.is_empty())
def test_stop_initiated_async_pipelines(self, mock_gen_task_from_active,
mock_async_task_gen,
mock_sync_task_gen):
with self._mlmd_connection as m:
pipeline1 = _test_pipeline('pipeline1')
pipeline1.nodes.add().pipeline_node.node_info.id = 'Transform'
pipeline1.nodes.add().pipeline_node.node_info.id = 'Trainer'
pipeline1.nodes.add().pipeline_node.node_info.id = 'Evaluator'
pipeline_ops.initiate_pipeline_start(m, pipeline1)
pipeline1_execution = pipeline_ops._initiate_pipeline_stop(
m, task_lib.PipelineUid.from_pipeline(pipeline1))
task_queue = tq.TaskQueue()
# For the stop-initiated pipeline, "Transform" execution task is in queue,
# "Trainer" has an active execution in MLMD but no task in queue,
# "Evaluator" has no active execution.
task_queue.enqueue(
test_utils.create_exec_node_task(node_uid=task_lib.NodeUid(
pipeline_uid=task_lib.PipelineUid(pipeline_id='pipeline1',
pipeline_run_id=None),
node_id='Transform')))
transform_task = task_queue.dequeue(
) # simulates task being processed
mock_gen_task_from_active.side_effect = [
test_utils.create_exec_node_task(node_uid=task_lib.NodeUid(
pipeline_uid=task_lib.PipelineUid(pipeline_id='pipeline1',
pipeline_run_id=None),
node_id='Trainer'),
is_cancelled=True), None,
None, None, None
]
pipeline_ops.generate_tasks(m, task_queue)
# There are no active pipelines so these shouldn't be called.
mock_async_task_gen.assert_not_called()
mock_sync_task_gen.assert_not_called()
# Simulate finishing the "Transform" ExecNodeTask.
task_queue.task_done(transform_task)
# CancelNodeTask for the "Transform" ExecNodeTask should be next.
task = task_queue.dequeue()
task_queue.task_done(task)
self.assertTrue(task_lib.is_cancel_node_task(task))
self.assertEqual('Transform', task.node_uid.node_id)
# ExecNodeTask for "Trainer" is next.
task = task_queue.dequeue()
task_queue.task_done(task)
self.assertTrue(task_lib.is_exec_node_task(task))
self.assertEqual('Trainer', task.node_uid.node_id)
self.assertTrue(task_queue.is_empty())
mock_gen_task_from_active.assert_has_calls([
mock.call(m,
pipeline1,
pipeline1.nodes[1].pipeline_node,
mock.ANY,
is_cancelled=True),
mock.call(m,
pipeline1,
pipeline1.nodes[2].pipeline_node,
mock.ANY,
is_cancelled=True)
])
self.assertEqual(2, mock_gen_task_from_active.call_count)
# Pipeline execution should continue to be active since active node
# executions were found in the last call to `generate_tasks`.
[execution
] = m.store.get_executions_by_id([pipeline1_execution.id])
self.assertTrue(execution_lib.is_execution_active(execution))
# Call `generate_tasks` again; this time there are no more active node
# executions so the pipeline should be marked as cancelled.
pipeline_ops.generate_tasks(m, task_queue)
self.assertTrue(task_queue.is_empty())
[execution
] = m.store.get_executions_by_id([pipeline1_execution.id])
self.assertEqual(metadata_store_pb2.Execution.CANCELED,
execution.last_known_state)
def test_active_pipelines_with_stop_initiated_nodes(
self, mock_gen_task_from_active, mock_async_task_gen):
with self._mlmd_connection as m:
pipeline = _test_pipeline('pipeline')
pipeline.nodes.add().pipeline_node.node_info.id = 'Transform'
pipeline.nodes.add().pipeline_node.node_info.id = 'Trainer'
pipeline.nodes.add().pipeline_node.node_info.id = 'Evaluator'
transform_node_uid = task_lib.NodeUid.from_pipeline_node(
pipeline, pipeline.nodes[0].pipeline_node)
transform_task = test_utils.create_exec_node_task(
node_uid=transform_node_uid)
trainer_node_uid = task_lib.NodeUid.from_pipeline_node(
pipeline, pipeline.nodes[1].pipeline_node)
trainer_task = test_utils.create_exec_node_task(
node_uid=trainer_node_uid)
evaluator_node_uid = task_lib.NodeUid.from_pipeline_node(
pipeline, pipeline.nodes[2].pipeline_node)
evaluator_task = test_utils.create_exec_node_task(
node_uid=evaluator_node_uid)
cancelled_evaluator_task = test_utils.create_exec_node_task(
node_uid=evaluator_node_uid, is_cancelled=True)
pipeline_ops.initiate_pipeline_start(m, pipeline)
with pstate.PipelineState.load(
m, task_lib.PipelineUid.from_pipeline(
pipeline)) as pipeline_state:
# Stop trainer and evaluator.
pipeline_state.initiate_node_stop(trainer_node_uid)
pipeline_state.initiate_node_stop(evaluator_node_uid)
task_queue = tq.TaskQueue()
# Simulate a new transform execution being triggered.
mock_async_task_gen.return_value.generate.return_value = [
transform_task
]
# Simulate ExecNodeTask for trainer already present in the task queue.
task_queue.enqueue(trainer_task)
# Simulate Evaluator having an active execution in MLMD.
mock_gen_task_from_active.side_effect = [evaluator_task]
pipeline_ops.generate_tasks(m, task_queue)
self.assertEqual(
1, mock_async_task_gen.return_value.generate.call_count)
# Verify that tasks are enqueued in the expected order:
# Pre-existing trainer task.
task = task_queue.dequeue()
task_queue.task_done(task)
self.assertEqual(trainer_task, task)
# CancelNodeTask for trainer.
task = task_queue.dequeue()
task_queue.task_done(task)
self.assertTrue(task_lib.is_cancel_node_task(task))
self.assertEqual(trainer_node_uid, task.node_uid)
# ExecNodeTask with is_cancelled=True for evaluator.
task = task_queue.dequeue()
task_queue.task_done(task)
self.assertTrue(cancelled_evaluator_task, task)
# ExecNodeTask for newly triggered transform node.
task = task_queue.dequeue()
task_queue.task_done(task)
self.assertEqual(transform_task, task)
# No more tasks.
self.assertTrue(task_queue.is_empty())