Develop a Custom Step Operator
How to develop a custom step operator
This is an older version of the ZenML documentation. To read and view the latest version please visit this up-to-date URL.
Base Abstraction
The BaseStepOperator
is the abstract base class that needs to be subclassed in order to run specific steps of your pipeline in a separate environment. As step operators can come in many shapes and forms, the base class exposes a deliberately basic and generic interface:
As it is the base class for a specific type of
StackComponent
, it inherits from theStackComponent
class. This sets theTYPE
variable to the specificStackComponentType
.The
FLAVOR
class variable needs to be set in subclasses as it is meant to identify the implementation flavor of the particular step operator.Lastly, the base class features one
abstractmethod
calledlaunch()
. In the implementation of every step operator flavor, thelaunch()
method is responsible for preparing the environment and executing the step.
This is a slimmed-down version of the base implementation which aims to highlight the abstraction layer. In order to see the full implementation and get the complete docstrings, please check the API docs.
Building your own custom step operator
If you want to create a custom step operator, you can follow these steps:
Create a class which inherits from the
BaseStepOperator
.Define the
FLAVOR
class variable.Implement the abstract
launch()
method. This method has two main responsibilities:Preparing a suitable execution environment (e.g. a Docker image): The general environment is highly dependent on the concrete step operator implementation, but for ZenML to be able to run the step it requires you to install some
pip
dependencies. The list of requirements needed to successfully execute the step is passed in via therequirements
argument of thelaunch()
method. Additionally, you'll have to make sure that all the source code of your ZenML step and pipeline are available within this execution environment.Running the entrypoint command: Actually running a single step of a pipeline requires knowledge of many ZenML internals and is implemented in the
zenml.step_operators.entrypoint
module. As long as your environment was set up correctly (see the previous bullet point), you can run the step using the command provided via theentrypoint_command
argument of thelaunch()
method.
If your step operator allows specification of per-step resources, make sure to handle the options in the
ResourceConfiguration
passed to thelaunch()
method.
Once you are done with the implementation, you can register it through the CLI as:
Last updated