Class Based API
Use the class-based API for Steps and Pipelines.
The class-based ZenML API is defined by the base classes BaseStep and BasePipeline. You'll be subclassing these instead of using the @step and @pipeline decorators that you have used in the previous sections. These interfaces allow our users to maintain a higher level of control while they are creating a step definition and using it within the context of a pipeline.
We'll be using the code for the Runtime Configuration as a point of comparison.

Subclassing the BaseStep

In order to create a step, you will need to create a subclass of the BaseStep and implement its entrypoint() method. This entrypoint contains the logic of the step.
Class-based API
Functional API
1
from zenml.steps import BaseStep, BaseStepConfig, Output
2
3
class SecondStepConfig(BaseStepConfig):
4
"""Trainer params"""
5
multiplier: int = 4
6
7
8
class SecondStep(BaseStep):
9
def entrypoint(
10
self,
11
config: SecondStepConfig,
12
input_int: int,
13
input_float: float) -> Output(output_int=int, output_float=float):
14
"""Step that multiply the inputs"""
15
return config.multiplier * input_int, config.multiplier * input_float
Copied!
1
from zenml.steps import BaseStepConfig, step, Output
2
class SecondStepConfig(BaseStepConfig):
3
"""Trainer params"""
4
multiplier: int = 4
5
6
@step
7
def my_second_step(config: SecondStepConfig, input_int: int,
8
input_float: float
9
) -> Output(output_int=int, output_float=float):
10
"""Step that multiply the inputs"""
11
return config.multiplier * input_int, config.multiplier * input_float
Copied!

Use decorators to enhance your BaseStep subclasses

ZenML allows you to easily add functionality like automated experiment tracking to your steps. To do so with the class-based API, simply decorate your BaseStep subclass as follows:
1
from zenml.steps import BaseStep
2
from zenml.integrations.mlflow.mlflow_step_decorator import enable_mlflow
3
4
@enable_mlflow
5
class TFTrainer(BaseStep):
6
def entrypoint(
7
self,
8
x_train: np.ndarray,
9
y_train: np.ndarray,
10
) -> tf.keras.Model:
11
mlflow.tensorflow.autolog()
12
...
Copied!
Check out our MLflow, Wandb and whylogs examples for more information on how to use the specific decorators.

Subclassing the BasePipeline

Class-based API
Functional API
1
from zenml.steps import BaseStep
2
from zenml.pipelines import BasePipeline
3
4
5
class FirstPipeline(BasePipeline):
6
"""Pipeline to show off the class-based API"""
7
8
def connect(self,
9
step_1: BaseStep,
10
step_2: BaseStep) -> None:
11
output_1, output_2 = step_1()
12
step_2(output_1, output_2)
13
14
FirstPipeline(step_1=my_first_step(),
15
step_2=SecondStep(SecondStepConfig(multiplier=3))).run()
Copied!
1
from zenml.pipelines import pipeline
2
3
@pipeline
4
def first_pipeline(
5
step_1,
6
step_2
7
):
8
output_1, output_2 = step_1()
9
step_2(output_1, output_2)
10
11
first_pipeline(step_1=my_first_step(),
12
step_2=my_second_step(SecondStepConfig(multiplier=3))
13
).run()
Copied!
You can also mix and match the two APIs to your hearts content. Check out the full Code example below to see how.

Summary in Code

Code Example for this Section
Export as PDF
Copy link
Edit on GitHub
Contents