Streamline Machine Learning Workflows with Amazon SageMaker Pipelines and AWS Controllers for Kubernetes

Introduction

Kubernetes and Amazon SageMaker are two powerful tools in the realm of machine learning (ML) and DevOps. Kubernetes is an orchestration platform for managing containers, while Amazon SageMaker is a fully managed service that provides capabilities for building, training, and deploying ML models. When used together, these two tools can help streamline the ML workflow, making it easier for ML engineers to focus on developing and training models, and for DevOps engineers to manage the underlying infrastructure. In this post, we’ll explore how to integrate Kubernetes and Amazon SageMaker to create a seamless ML workflow.

Solution Overview

We’ll consider a use case in which an ML engineer configures a SageMaker model building pipeline using a Jupyter notebook. This configuration takes the form of a Directed Acyclic Graph (DAG) represented as a JSON pipeline definition. The JSON document can be stored and versioned in an Amazon Simple Storage Service (Amazon S3) bucket. If encryption is required, it can be implemented using an AWS Key Management Service (AWS KMS) managed key for Amazon S3. A DevOps engineer with access to fetch this definition file from Amazon S3 can load the pipeline definition into an ACK service controller for SageMaker, which is running as part of an Amazon Elastic Kubernetes Service (Amazon EKS) cluster.

Prerequisites

To follow along, you should have the following prerequisites:

• An EKS cluster where the ML pipeline will be created.
• A user with access to an AWS Identity and Access Management (IAM) role that has IAM permissions (iam:CreateRole, iam:AttachRolePolicy, and iam:PutRolePolicy) to allow creating roles and attaching policies to roles.
• The following command line tools on the local machine or cloud-based development environment used to access the Kubernetes cluster:

Install the SageMaker ACK Service Controller

The SageMaker ACK service controller makes it straightforward for DevOps engineers to use Kubernetes as their control plane to create and manage ML pipelines. To install the controller in your EKS cluster, complete the following steps:

1. Configure IAM permissions to make sure the controller has access to the appropriate AWS resources.
2. Install the controller using a SageMaker Helm Chart to make it available on the client machine.

The following tutorial provides step-by-step instructions with the required commands to install the ACK service controller for SageMaker.

Generate a Pipeline JSON Definition

In most companies, ML engineers are responsible for creating the ML pipeline in their organization. They often work with DevOps engineers to operate those pipelines. In SageMaker, ML engineers can use the SageMaker Python SDK to generate a pipeline definition in JSON format. A SageMaker pipeline definition must follow the provided schema, which includes base images, dependencies, steps, and instance types and sizes that are needed to fully define the pipeline. This definition then gets retrieved by the DevOps engineer for deploying and maintaining the infrastructure needed for the pipeline.

With SageMaker, ML model artifacts and other system artifacts are encrypted in transit and at rest. SageMaker encrypts these by default using the AWS managed key for Amazon S3. You can optionally specify a custom key using the KmsKeyId property of the OutputDataConfig argument. For more information on how SageMaker protects data, see Data Protection in Amazon SageMaker.

Submit the Pipeline to SageMaker

To submit your prepared pipeline specification, apply the specification to your Kubernetes cluster as follows:

kubectl apply -f my-pipeline.yaml

Create and Submit a Pipeline Execution YAML Specification

Refer to the following Kubernetes YAML specification for a SageMaker pipeline. Prepare the pipeline execution YAML specification (pipeline-execution.yaml) as follows:

apiVersion: sagemaker.services.k8s.aws/v1alpha1
kind: PipelineExecution
metadata:
  name: my-kubernetes-pipeline-execution
spec:
  parallelismConfiguration:
  	maxParallelExecutionSteps: 2
  pipelineExecutionDescription: "My first pipeline execution via Amazon EKS cluster."
  pipelineName: my-kubernetes-pipeline

To start a run of the pipeline, use the following code:

kubectl apply -f pipeline-execution.yaml

Review and Troubleshoot the Pipeline Run

To list all pipelines created using the ACK controller, use the following command:

kubectl get pipelineexecution

To list all pipeline runs, use the following command:

kubectl get pipelineexecution

To get more details about the pipeline after it’s submitted, like checking the status, errors, or parameters of the pipeline, use the following command:

kubectl describe pipeline my-kubernetes-pipeline

To troubleshoot a pipeline run by reviewing more details about the run, use the following command:

kubectl describe pipelineexecution my-kubernetes-pipeline-execution

Clean Up

Use the following command to delete any pipelines you created:

kubectl delete pipelineexecution

Use the following command to cancel any pipeline runs you started:

kubectl delete pipelineexecution

Conclusion

In this post, we presented an example of how ML engineers familiar with Jupyter notebooks and SageMaker environments can efficiently work with DevOps engineers familiar with Kubernetes and related tools to design and maintain an ML pipeline with the right infrastructure for their organization. This enables DevOps engineers to manage all the steps of the ML lifecycle with the same set of tools and environment they are used to, which enables organizations to innovate faster and more efficiently.

Frequently Asked Questions

Q1: What is the benefit of using Kubernetes and Amazon SageMaker together?

Kubernetes and Amazon SageMaker provide a powerful combination for managing machine learning (ML) pipelines. Kubernetes provides orchestration and management capabilities for containers, while Amazon SageMaker provides a fully managed service for building, training, and deploying ML models. Together, they enable ML engineers to focus on developing and training models, and DevOps engineers to manage the underlying infrastructure.

Q2: How do I install the SageMaker ACK service controller?

To install the SageMaker ACK service controller, you need to configure IAM permissions to ensure the controller has access to the appropriate AWS resources, and then install the controller using a SageMaker Helm Chart to make it available on the client machine.

Q3: What is the difference between a pipeline definition and a pipeline execution?

A pipeline definition is a JSON document that defines the steps and resources needed to create an ML pipeline. A pipeline execution is a specific instance of a pipeline that is submitted to SageMaker for execution.

Q4: How do I troubleshoot a pipeline run?

You can troubleshoot a pipeline run by using the kubectl describe command to get more details about the pipeline after it’s submitted, or by using the kubectl describe pipelineexecution command to get more details about the pipeline execution.

Q5: Can I use Kubernetes and Amazon SageMaker together in my organization?

Yes, you can use Kubernetes and Amazon SageMaker together in your organization to create a seamless ML workflow. SageMaker provides a fully managed service for building, training, and deploying ML models, while Kubernetes provides orchestration and management capabilities for containers. By integrating these two services, you can enable ML engineers to focus on developing and training models, and DevOps engineers to manage the underlying infrastructure.