0.36.1
Search
⌃K
Links

Specifying Step Resources

How to specify per-step resources
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
Class-based API
from zenml.steps import step, ResourceSettings
@step(settings={"resources": ResourceSettings(cpu_count=8, gpu_count=2)})
def training_step(...) -> ...:
# train a model
from zenml.steps import BaseStep, ResourceSettings
class TrainingStep(BaseStep):
...
step = TrainingStep(settings = {"resources": ResourceSettings(cpu_count=8, gpu_count=2)})
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.