Use failure/success hooks

Running failure and success hooks after step execution.

Hooks are a way to perform an action after a step has completed execution. They can be useful in a variety of scenarios, such as sending notifications, logging, or cleaning up resources after a step has been completed.

A hook executes right after step execution, within the same environment as the step, therefore it has access to all the dependencies that a step has. Currently, there are two sorts of hooks that can be defined: on_failure and on_success .

on_failure: This hook triggers in the event of a step failing.

on_success: This hook triggers in the event of a step succeeding.

Here is a short demo for hooks in ZenML:

Defining hooks

A hook can be defined as a callback function, and must be accessible within the repository where the pipeline and steps are located.

In case of failure hooks, you can optionally add a BaseException argument to the hook, allowing you to access the concrete Exception that caused your step to fail:

from zenml import step

def on_failure(exception: BaseException):

print(f"Step failed: {str(exception)}")

def on_success():

print("Step succeeded!")

@step(on_failure=on_failure)

def my_failing_step() -> int:

"""Returns an integer."""

raise ValueError("Error")

@step(on_success=on_success)

def my_successful_step() -> int:

"""Returns an integer."""

return 1

A step can also be specified as a local user-defined function path (of the form mymodule.myfile.my_function). This is particularly useful when defining the hooks via a YAML Config.

Defining hooks on a pipeline level

In some cases, there is a need to define a hook on all steps of a given pipeline. Rather than having to define it on all steps individually, you can also specify any hook on the pipeline level.

@pipeline(on_failure=on_failure, on_success=on_success)

def my_pipeline(...):

...

Note, that step-level defined hooks take precedence over pipeline-level defined hooks.