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)