Develop a Custom Step Operator

Learning how to develop a custom step operator.

Before diving into the specifics of this component type, it is beneficial to familiarize yourself with our general guide to writing custom component flavors in ZenML. This guide provides an essential understanding of ZenML's component flavor concepts.

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:

from abc import ABC, abstractmethod

from typing import List, Type

from zenml.enums import StackComponentType

from zenml.stack import StackComponent, StackComponentConfig, Flavor

from zenml.config.step_run_info import StepRunInfo

class BaseStepOperatorConfig(StackComponentConfig):

"""Base config for step operators."""

class BaseStepOperator(StackComponent, ABC):

"""Base class for all ZenML step operators."""

@abstractmethod

def launch(

self,

info: StepRunInfo,

entrypoint_command: List[str],

) -> None:

"""Abstract method to execute a step.

Subclasses must implement this method and launch a **synchronous**

job that executes the `entrypoint_command`.

Args:

info: Information about the step run.

entrypoint_command: Command that executes the step.

"""

class BaseStepOperatorFlavor(Flavor):

"""Base class for all ZenML step operator flavors."""

@property

@abstractmethod

def name(self) -> str:

"""Returns the name of the flavor."""

@property

def type(self) -> StackComponentType:

"""Returns the flavor type."""

return StackComponentType.STEP_OPERATOR

@property

def config_class(self) -> Type[BaseStepOperatorConfig]:

"""Returns the config class for this flavor."""

return BaseStepOperatorConfig

@property

@abstractmethod

def implementation_class(self) -> Type[BaseStepOperator]: