Refactoring to enable on_failure for imperative workflow