Specify 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, ResourceConfiguration
@step(resource_configuration=ResourceConfiguration(cpu_count=8, gpu_count=2))
def training_step(...) -> ...:
# train a model
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.
Copy link