Failure and Success Hooks
How to run failure and success hooks after step execution
This is an older version of the ZenML documentation. To read and view the latest version please visit this up-to-date URL.
Hooks
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 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 event of a step failing.on_success
: This hook triggers in 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 this example, we define two hooks: on_failure
and on_success
, which print a message when the step fails or succeeds, respectively. We then use these hooks with two steps, my_failing_step
and my_successful_step
. When my_failing_step
is executed, it raises a ValueError
, which triggers the on_failure
hook. Similarly, when my_successful_step
is executed, it returns an integer successfully, which triggers the on_success hook.
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 steps 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.
Note, that step-level defined hooks take precedence over pipeline-level defined hooks.
Accessing step information inside a hook
A hook function signature can optionally take three type annotated arguments of the following types:
StepContext
: You can pass an object inside a hook of typeStepContext
to get access to information such as pipeline name, run name, step name etcBaseParameters
: You can pass aBaseParameters
inside the stepBaseException
: In case of failure hooks, if you add anBaseException
argument to the hook, then ZenML passes the BaseException that caused your step to fail.
Please note that in case of BaseParameters
and BaseException
the concrete class defined by the step will be passed. For example, if a step's parameters class is called MyParameters
, that will be the object that is passed into the hook. Also, if a ValueError
is raised inside a step, the exception would also be of type BaseException
.
Linking to the Alerter
Stack component
Alerter
Stack componentA common use-case is to use the Alerter component inside the failure or success hooks to notify relevant people. It is quite easy to do this:
ZenML provides standard failure and success hooks that use the alerter you have configured in your stack. Here's an example of how to use them in your pipelines:
Using the OpenAI ChatGPT Failure Hook
The OpenAI ChatGPT failure hook is a hook that uses the OpenAI integration to generate a possible fix for whatever exception caused the step to fail. It is quite easy to use. (You will need a valid OpenAI API key that has correctly set up billing for this.)
Note that using this integration will incur charges on your OpenAI account.
First ensure that you have the OpenAI integration installed and have stored your API key within a ZenML secret:
Then, you can use the hook in your pipeline:
If you had set up a Slack alerter as your alerter, for example, then you would see a message like this:
You can use the suggestions as input that can help you fix whatever is going wrong in your code. If you have GPT-4 enabled for your account, you can use the openai_gpt4_alerter_failure_hook
hook instead (imported from the same module).
Last updated