Docker settings on a step

You have the option to customize the Docker settings at a step level.

By default every step of a pipeline uses the same Docker image that is defined at the pipeline level. Sometimes your steps will have special requirements that make it necessary to define a different Docker image for one or many steps. This can easily be accomplished by adding the DockerSettings to the step decorator directly.

from zenml import step
from zenml.config import DockerSettings

@step(
  settings={
    "docker": DockerSettings(
      parent_image="pytorch/pytorch:1.12.1-cuda11.3-cudnn8-runtime"
    )
  }
)
def training(...):
	...

Alternatively, this can also be done within the configuration file.

steps:
  training:
    settings:
      docker:
        parent_image: pytorch/pytorch:2.2.0-cuda11.8-cudnn8-runtime
        required_integrations:
          - gcp
          - github
        requirements:
          - zenml  # Make sure to include ZenML for other parent images
          - numpy