def test_stop_pipeline_wait_for_inactivation_timeout(self, pipeline):
with self._mlmd_connection as m:
pipeline_ops.initiate_pipeline_start(m, pipeline)
with self.assertRaisesRegex(
status_lib.StatusNotOkError,
'Timed out.*waiting for execution inactivation.'
) as exception_context:
pipeline_ops.stop_pipeline(
m,
task_lib.PipelineUid.from_pipeline(pipeline),
timeout_secs=1.0)
self.assertEqual(status_lib.Code.DEADLINE_EXCEEDED,
exception_context.exception.code)
def test_stop_pipeline_wait_for_inactivation(self, pipeline):
with self._mlmd_connection as m:
execution = pipeline_ops.initiate_pipeline_start(m, pipeline).execution
def _inactivate(execution):
time.sleep(2.0)
with pipeline_ops._PIPELINE_OPS_LOCK:
execution.last_known_state = metadata_store_pb2.Execution.COMPLETE
m.store.put_executions([execution])
thread = threading.Thread(
target=_inactivate, args=(copy.deepcopy(execution),))
thread.start()
pipeline_ops.stop_pipeline(
m, task_lib.PipelineUid.from_pipeline(pipeline), timeout_secs=10.0)
thread.join()
def test_stop_pipeline_wait_for_inactivation(self, pipeline):
with self._mlmd_connection as m:
pipeline_state = pipeline_ops.initiate_pipeline_start(m, pipeline)
def _inactivate(pipeline_state):
time.sleep(2.0)
with pipeline_ops._PIPELINE_OPS_LOCK:
with pipeline_state:
pipeline_state.set_pipeline_execution_state(
metadata_store_pb2.Execution.COMPLETE)
thread = threading.Thread(target=_inactivate,
args=(pipeline_state, ))
thread.start()
pipeline_ops.stop_pipeline(
m,
task_lib.PipelineUid.from_pipeline(pipeline),
timeout_secs=10.0)
thread.join()
def test_stop_pipeline_non_existent_or_inactive(self, pipeline):
with self._mlmd_connection as m:
# Stop pipeline without creating one.
with self.assertRaises(status_lib.StatusNotOkError) as exception_context:
pipeline_ops.stop_pipeline(m,
task_lib.PipelineUid.from_pipeline(pipeline))
self.assertEqual(status_lib.Code.NOT_FOUND,
exception_context.exception.code)
# Initiate pipeline start and mark it completed.
execution = pipeline_ops.initiate_pipeline_start(m, pipeline).execution
pipeline_uid = task_lib.PipelineUid.from_pipeline(pipeline)
with pstate.PipelineState.load(m, pipeline_uid) as pipeline_state:
pipeline_state.initiate_stop(status_lib.Status(code=status_lib.Code.OK))
execution.last_known_state = metadata_store_pb2.Execution.COMPLETE
m.store.put_executions([execution])
# Try to initiate stop again.
with self.assertRaises(status_lib.StatusNotOkError) as exception_context:
pipeline_ops.stop_pipeline(m, pipeline_uid)
self.assertEqual(status_lib.Code.NOT_FOUND,
exception_context.exception.code)
