Specifying Step Resources

How to specify per-step resources

This is an older version of the ZenML documentation. To read and view the latest version please visit this up-to-date URL.

Some steps of your machine learning pipeline might be more resource-intensive and require special hardware to execute. In such cases, you can specify the required resources for steps as follows:

from zenml.steps import step, ResourceSettings

@step(settings={"resources": ResourceSettings(cpu_count=8, gpu_count=2)})
def training_step(...) -> ...:
    # train a model

If you're using an orchestrator which does not support this feature or its underlying infrastructure doesn't cover your requirements, you can also take a look at step operators which allow you to execute individual steps of your pipeline in environments independent of your orchestrator.

Last updated