Amazon SageMaker
Executing individual steps in SageMaker.
This is an older version of the ZenML documentation. To read and view the latest version please visit this up-to-date URL.
Amazon SageMaker
SageMaker offers specialized compute instances to run your training jobs and has a comprehensive UI to track and manage your models and logs. ZenML's SageMaker step operator allows you to submit individual steps to be run on Sagemaker compute instances.
When to use it
You should use the SageMaker step operator if:
one or more steps of your pipeline require computing resources (CPU, GPU, memory) that are not provided by your orchestrator.
How to deploy it
Create a role in the IAM console that you want the jobs running in SageMaker to assume. This role should at least have the
AmazonS3FullAccess
andAmazonSageMakerFullAccess
policies applied. Check here for a guide on how to set up this role.
Infrastructure Deployment
A Sagemaker step operator can be deployed directly from the ZenML CLI:
You can pass other configurations specific to the stack components as key-value arguments. If you don't provide a name, a random one is generated for you. For more information about how to work use the CLI for this, please refer to the dedicated documentation section.
How to use it
To use the SageMaker step operator, we need:
The ZenML
aws
integration installed. If you haven't done so, runDocker installed and running.
An IAM role with the correct permissions. See the deployment section for detailed instructions.
An AWS container registry as part of our stack. Take a look here for a guide on how to set that up.
A remote artifact store as part of your stack. This is needed so that both your orchestration environment and SageMaker can read and write step artifacts. Check out the documentation page of the artifact store you want to use for more information on how to set that up and configure authentication for it.
If using a local orchestrator: The
aws
cli set up and authenticated. Make sure you have the permissions to create and manage SageMaker runs.If using a remote orchestrator: The environment in which the orchestrator runs its containers needs to be able to assume the IAM role specified when registering the SageMaker step operator.
An instance type that we want to execute our steps on. See here for a list of available instance types.
(Optional) An experiment that is used to group SageMaker runs. Check this guide to see how to create an experiment.
We can then register the step operator and use it in our active stack:
Once you added the step operator to your active stack, you can use it to execute individual steps of your pipeline by specifying it in the @step
decorator as follows:
ZenML will build a Docker image called <CONTAINER_REGISTRY_URI>/zenml:<PIPELINE_NAME>
which includes your code and use it to run your steps in SageMaker. Check out this page if you want to learn more about how ZenML builds these images and how you can customize them.
Additional configuration
For additional configuration of the SageMaker step operator, you can pass SagemakerStepOperatorSettings
when defining or running your pipeline. Check out the SDK docs for a full list of available attributes and this docs page for more information on how to specify settings.
For more information and a full list of configurable attributes of the SageMaker step operator, check out the API Docs .
Enabling CUDA for GPU-backed hardware
Note that if you wish to use this step operator to run steps on a GPU, you will need to follow the instructions on this page to ensure that it works. It requires adding some extra settings customization and is essential to enable CUDA for the GPU to give its full acceleration.
Last updated