Specify 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:

Functional API

from zenml.steps import step, ResourceConfiguration

@step(resource_configuration=ResourceConfiguration(cpu_count=8, gpu_count=2))
def training_step(...) -> ...:
    # train a model

Class-based API

from zenml.steps import BaseStep, ResourceConfiguration

class TrainingStep(BaseStep):
    ...

step = TrainingStep(resource_configuration=ResourceConfiguration(cpu_count=8, gpu_count=2))

If you're using an orchestrator which doesn't 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.