Pipeline Configuration

Setting step parameters using Config

Sometimes you want to flexibly adjust parameters when you run your pipeline. This is where the step configurations come into play. In the following example we want to be able to change the learning rate for each pipeline run. For this, we create a TrainerConfig that contains all the parameters that concern the trainer step.
1
import tensorflow as tf
2
3
from zenml.pipelines import pipeline
4
from zenml.steps import step, Output, BaseStepConfig
5
6
@step
7
def importer_func() -> Output(
8
X_train=np.ndarray, y_train=np.ndarray, X_test=np.ndarray, y_test=np.ndarray
9
):
10
"""Importing data"""
11
...
12
13
class TrainerConfig(BaseStepConfig):
14
"""Trainer params"""
15
lr: float = 0.001
16
17
@step
18
def trainer_func(
19
config: TrainerConfig, # not an artifact, passed in when pipeline is instantiated
20
X_train: np.ndarray,
21
y_train: np.ndarray
22
):
23
"""Training model"""
24
optimizer = tf.keras.optimizers.Adam(config.lr)
25
...
26
27
@pipeline
28
def my_pipeline(
29
importer,
30
trainer,
31
):
32
"""Links all the steps together in a pipeline"""
33
X_train, y_train, X_test, y_test = importer()
34
model = trainer(X_train=X_train, y_train=y_train)
35
36
pipeline_instance = my_pipeline(
37
importer=importer_func(),
38
trainer=trainer_func(config=TrainerConfig(lr=0.003))
39
)
40
41
pipeline_instance.run()
Copied!

Setting step parameters using a config file

In addition to setting parameters for your pipeline steps in code as seen above, ZenML also allows you to use a configuration yaml file. This configuration file must follow the following structure:
1
steps:
2
step_name:
3
parameters:
4
parameter_name: parameter_value
5
some_other_parameter_name: 2
6
some_other_step_name:
7
...
Copied!
For our example from above this results in the following configuration yaml.
1
steps:
2
trainer:
3
parameters:
4
lr: 0.005
Copied!
Use the configuration file by calling the pipeline method with_config(...):
1
pipeline_instance = my_pipeline(
2
importer=importer_func(),
3
trainer=trainer_func()
4
).with_config("path_to_config.yaml")
5
6
pipeline_instance.run()
Copied!
Last modified 23h ago