Skip to content

Step Operators

Step operators allow you to run steps on custom infrastructure.

While an orchestrator defines how and where your entire pipeline runs, a step operator defines how and where an individual step runs. This can be useful in a variety of scenarios. An example could be if one step within a pipeline should run on a separate environment equipped with a GPU (like a trainer step).

BaseStepOperator

Bases: StackComponent, ABC

Base class for all ZenML step operators.

Source code in src/zenml/step_operators/base_step_operator.py 
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
class BaseStepOperator(StackComponent, ABC):
    """Base class for all ZenML step operators."""

    @property
    def config(self) -> BaseStepOperatorConfig:
        """Returns the config of the step operator.

        Returns:
            The config of the step operator.
        """
        return cast(BaseStepOperatorConfig, self._config)

    @property
    def entrypoint_config_class(
        self,
    ) -> Type[StepOperatorEntrypointConfiguration]:
        """Returns the entrypoint configuration class for this step operator.

        Concrete step operator implementations may override this property
        to return a custom entrypoint configuration class if they need to
        customize the entrypoint configuration.

        Returns:
            The entrypoint configuration class for this step operator.
        """
        return StepOperatorEntrypointConfiguration

    @abstractmethod
    def launch(
        self,
        info: "StepRunInfo",
        entrypoint_command: List[str],
        environment: Dict[str, 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.
            environment: Environment variables to set in the step operator
                environment.
        """

config property

Returns the config of the step operator.

Returns:

Type Description
BaseStepOperatorConfig

The config of the step operator.

entrypoint_config_class property

Returns the entrypoint configuration class for this step operator.

Concrete step operator implementations may override this property to return a custom entrypoint configuration class if they need to customize the entrypoint configuration.

Returns:

Type Description
Type[StepOperatorEntrypointConfiguration]

The entrypoint configuration class for this step operator.

launch(info, entrypoint_command, environment) abstractmethod

Abstract method to execute a step.

Subclasses must implement this method and launch a synchronous job that executes the entrypoint_command.

Parameters:

Name Type Description Default
info StepRunInfo

Information about the step run.

 required
entrypoint_command List[str]

Command that executes the step.

 required
environment Dict[str, str]

Environment variables to set in the step operator environment.

 required
Source code in src/zenml/step_operators/base_step_operator.py 
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
@abstractmethod
def launch(
    self,
    info: "StepRunInfo",
    entrypoint_command: List[str],
    environment: Dict[str, 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.
        environment: Environment variables to set in the step operator
            environment.
    """

BaseStepOperatorConfig

Bases: StackComponentConfig

Base config for step operators.

Source code in src/zenml/step_operators/base_step_operator.py 
33
34
class BaseStepOperatorConfig(StackComponentConfig):
    """Base config for step operators."""

BaseStepOperatorFlavor

Bases: Flavor

Base class for all ZenML step operator flavors.

Source code in src/zenml/step_operators/base_step_operator.py 
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
class BaseStepOperatorFlavor(Flavor):
    """Base class for all ZenML step operator flavors."""

    @property
    def type(self) -> StackComponentType:
        """Returns the flavor type.

        Returns:
            The type of the flavor.
        """
        return StackComponentType.STEP_OPERATOR

    @property
    def config_class(self) -> Type[BaseStepOperatorConfig]:
        """Returns the config class for this flavor.

        Returns:
            The config class for this flavor.
        """
        return BaseStepOperatorConfig

    @property
    @abstractmethod
    def implementation_class(self) -> Type[BaseStepOperator]:
        """Returns the implementation class for this flavor.

        Returns:
            The implementation class for this flavor.
        """

config_class property

Returns the config class for this flavor.

Returns:

Type Description
Type[BaseStepOperatorConfig]

The config class for this flavor.

implementation_class abstractmethod property

Returns the implementation class for this flavor.

Returns:

Type Description
Type[BaseStepOperator]

The implementation class for this flavor.

type property

Returns the flavor type.

Returns:

Type Description
StackComponentType

The type of the flavor.