Using Kubernetes Operator for Amazon SageMaker's new inference capabilities reduces LLM deployment costs by an average of 50%

We are excited to announce a new version of Amazon SageMaker Operators for Kubernetes with AWS Controllers for Kubernetes (ACK). ACK is a framework for building Kubernetes custom controllers, and each controller communicates with his AWS service API. These controllers allow a Kubernetes user to provision her AWS resources, such as buckets, databases, and message queues, simply by using her Kubernetes API.

Release v1.2.9 of SageMaker ACK Operators adds support for inference components that were previously available only through the SageMaker API and AWS Software Development Kit (SDK). Inference components help optimize deployment costs and reduce latency. The new Inference Components feature allows you to deploy one or more foundation models (FMs) on the same Amazon SageMaker endpoint and control the number of accelerators and the amount of memory reserved for each FM. This improves resource utilization, reduces model deployment costs by an average of 50%, and allows you to scale endpoints to suit your use case. For more information, see Amazon SageMaker adds new inference capabilities to help reduce deployment costs and latency for underlying models.

Inference components are available through the SageMaker controller, allowing customers who use Kubernetes as their control plane to leverage inference components when deploying models to SageMaker.

This post shows how to deploy SageMaker inference components using the SageMaker ACK Operator.

How ACK works

To demonstrate how ACK works, let's look at an example using Amazon Simple Storage Service (Amazon S3). In the following diagram, Alice is her Kubernetes user. Her application relies on the existence of an S3 bucket named S3. my-bucket.

The workflow consists of the following steps:

1. Alice calls the next destination. kubectl applypass the file that describes the Kubernetes custom resource that describes your S3 bucket. kubectl apply passes this file, called a manifest, to the Kubernetes API server running on the Kubernetes controller node.
2. The Kubernetes API server receives the manifest that describes the S3 bucket and determines whether Alice has permission to create the following types of custom resources: s3.services.k8s.aws/Bucketmake sure your custom resources are properly formatted.
3. If Alice is authorized and the custom resource is enabled, the Kubernetes API server writes the custom resource to the server. etcd data store.
4. Next, respond to Alice that the custom resource has been created.
5. At this point, the Amazon S3 ACK service controller running on a Kubernetes worker node within the context of a regular Kubernetes pod is notified that a new custom resource of the following type has been added: s3.services.k8s.aws/Bucket has been created.
6. Amazon S3's ACK service controller then communicates with the Amazon S3 API and calls the S3 CreateBucket API to create a bucket in AWS.
7. After communicating with the Amazon S3 API, the ACK service controller calls the Kubernetes API server and updates the status of the custom resource with the information received from Amazon S3.

Main components

The new inference capabilities are built on SageMaker's real-time inference endpoint. As before, create a SageMaker endpoint with an endpoint configuration that defines the endpoint's instance type and initial number of instances. The model consists of a new building block, the inference component. Here you specify the number of accelerators and the amount of memory to allocate to each copy of your model, along with the number of model artifacts, container images, and model copies to deploy.

You can use new inference features in Amazon SageMaker Studio, SageMaker Python SDK, AWS SDK, and AWS Command Line Interface (AWS CLI). These are also supported by AWS CloudFormation. Now also available for SageMaker Operators for Kubernetes.

Solution overview

This demo uses the SageMaker controller to deploy a copy of the Dolly v2 7B model and a copy of the FLAN-T5 XXL model from the Hugging Face Model Hub to a SageMaker real-time endpoint using new inference capabilities.

Prerequisites

To proceed, you need a Kubernetes cluster with SageMaker ACK Controller v1.2.9 or later installed. To learn how to use eksctl to provision an Amazon Elastic Kubernetes Service (Amazon EKS) cluster with Amazon Elastic Compute Cloud (Amazon EC2) Linux managed nodes, see Getting Started with Amazon EKS – eksctl . For instructions on installing the SageMaker controller, see Machine Learning with ACK SageMaker Controller.

To host LLM, you need access to an accelerated instance (GPU). This solution uses one instance of ml.g5.12xlarge. You can check the availability of these instances in your AWS account and request them if needed through a service quota increase request, as shown in the following screenshot.

Create an inference component

To create an inference component, EndpointConfig, Endpoint, Modeland InferenceComponent A YAML file similar to the one shown in this section.use kubectl apply -f <yaml file> Create Kubernetes resources.

To list the status of your resources, kubectl describe <resource-type>; for example, kubectl describe inferencecomponent.

You can also create inference components without using model resources. For more information, please refer to the guidance provided in the API documentation.

EndpointConfig YAML

The code for the EndpointConfig file is as follows:

apiVersion: sagemaker.services.k8s.aws/v1alpha1
kind: EndpointConfig
metadata:
  name: inference-component-endpoint-config
spec:
  endpointConfigName: inference-component-endpoint-config
  executionRoleARN: <EXECUTION_ROLE_ARN>
  productionVariants:
  - variantName: AllTraffic
    instanceType: ml.g5.12xlarge
    initialInstanceCount: 1
    routingConfig:
      routingStrategy: LEAST_OUTSTANDING_REQUESTS

Endpoint YAML

The code for the endpoint file is as follows:

apiVersion: sagemaker.services.k8s.aws/v1alpha1
kind: Endpoint
metadata:
  name: inference-component-endpoint
spec:
  endpointName: inference-component-endpoint
  endpointConfigName: inference-component-endpoint-config

Model YAML

The code for the model file is as follows:

apiVersion: sagemaker.services.k8s.aws/v1alpha1
kind: Model
metadata:
  name: dolly-v2-7b
spec:
  modelName: dolly-v2-7b
  executionRoleARN: <EXECUTION_ROLE_ARN>
  containers:
  - image: 763104351884.dkr.ecr.us-east-1.amazonaws.com/huggingface-pytorch-tgi-inference:2.0.1-tgi0.9.3-gpu-py39-cu118-ubuntu20.04
    environment:
      HF_MODEL_ID: databricks/dolly-v2-7b
      HF_TASK: text-generation
---
apiVersion: sagemaker.services.k8s.aws/v1alpha1
kind: Model
metadata:
  name: flan-t5-xxl
spec:
  modelName: flan-t5-xxl
  executionRoleARN: <EXECUTION_ROLE_ARN>
  containers:
  - image: 763104351884.dkr.ecr.us-east-1.amazonaws.com/huggingface-pytorch-tgi-inference:2.0.1-tgi0.9.3-gpu-py39-cu118-ubuntu20.04
    environment:
      HF_MODEL_ID: google/flan-t5-xxl
      HF_TASK: text-generation

YAML for inference components

The following YAML file allocates 2 GPUs, 2 CPUs, and 1,024 MB of memory to each model, considering that the ml.g5.12xlarge instance comes with 4 GPUs.

apiVersion: sagemaker.services.k8s.aws/v1alpha1
kind: InferenceComponent
metadata:
  name: inference-component-dolly
spec:
  inferenceComponentName: inference-component-dolly
  endpointName: inference-component-endpoint
  variantName: AllTraffic
  specification:
    modelName: dolly-v2-7b
    computeResourceRequirements:
      numberOfAcceleratorDevicesRequired: 2
      numberOfCPUCoresRequired: 2
      minMemoryRequiredInMb: 1024
  runtimeConfig:
    copyCount: 1

apiVersion: sagemaker.services.k8s.aws/v1alpha1
kind: InferenceComponent
metadata:
  name: inference-component-flan
spec:
  inferenceComponentName: inference-component-flan
  endpointName: inference-component-endpoint
  variantName: AllTraffic
  specification:
    modelName: flan-t5-xxl
    computeResourceRequirements:
      numberOfAcceleratorDevicesRequired: 2
      numberOfCPUCoresRequired: 2
      minMemoryRequiredInMb: 1024
  runtimeConfig:
    copyCount: 1

Calling the model

Now you can call the model using the following code:

import boto3
import json

sm_runtime_client = boto3.client(service_name="sagemaker-runtime")
payload = {"inputs": "Why is California a great place to live?"}

response_dolly = sm_runtime_client.invoke_endpoint(
    EndpointName="inference-component-endpoint",
    InferenceComponentName="inference-component-dolly",
    ContentType="application/json",
    Accept="application/json",
    Body=json.dumps(payload),
)
result_dolly = json.loads(response_dolly['Body'].read().decode())
print(result_dolly)

response_flan = sm_runtime_client.invoke_endpoint(
    EndpointName="inference-component-endpoint",
    InferenceComponentName="inference-component-flan",
    ContentType="application/json",
    Accept="application/json",
    Body=json.dumps(payload),
)
result_flan = json.loads(response_flan['Body'].read().decode())
print(result_flan)

Update the inference component

To update an existing inference component, update the YAML file, then kubectl apply -f <yaml file>. Below is an example of an updated file.

apiVersion: sagemaker.services.k8s.aws/v1alpha1
kind: InferenceComponent
metadata:
  name: inference-component-dolly
spec:
  inferenceComponentName: inference-component-dolly
  endpointName: inference-component-endpoint
  variantName: AllTraffic
  specification:
    modelName: dolly-v2-7b
    computeResourceRequirements:
      numberOfAcceleratorDevicesRequired: 2
      numberOfCPUCoresRequired: 4 # Update the numberOfCPUCoresRequired.
      minMemoryRequiredInMb: 1024
  runtimeConfig:
    copyCount: 1

Delete the inference component

To remove an existing inference component, use the following command: kubectl delete -f <yaml file>.

Availability and price

New SageMaker inference capabilities include US East (Ohio, Northern Virginia), US West (Oregon), Asia Pacific (Jakarta, Mumbai, Seoul, Singapore, Sydney, Tokyo), Canada (Central), Europe (Frankfurt, Ireland, London, Stockholm), the Middle East (UAE), and South America (São Paulo). For pricing information, see Amazon SageMaker Pricing.

conclusion

In this post, you learned how to deploy SageMaker inference components using the SageMaker ACK Operator. Launch a Kubernetes cluster today and deploy FM using new SageMaker inference capabilities.

About the author

Rajesh Ramchander I am a Principal ML Engineer in Professional Services at AWS. He helps customers at various stages of their AI/ML and GenAI journeys, from those just beginning their journey to those leading their businesses with AI-first strategies.

Amit Arora is an AI and ML Specialist Architect at Amazon Web Services, helping enterprise customers rapidly scale their innovations using cloud-based machine learning services. He is also an adjunct instructor in the MSc Data Science and Analytics program at Georgetown University in Washington, DC.

Suryanshu Singh is a software development engineer at AWS SageMaker, where he works on developing large-scale ML distributed infrastructure solutions for AWS customers.

Saurabh Trikhande I am a senior product manager for Amazon SageMaker Inference. He is passionate about collaborating with customers and is motivated by the goal of democratizing machine learning. He focuses on key challenges related to complex ML applications, multi-tenant ML models, cost optimization, and making the deployment of deep learning models more accessible. In his free time, he enjoys hiking, learning about innovative technology, following TechCrunch, and spending time with his family.

Jonah Liu I'm a software development engineer on the Amazon SageMaker team. Her current work focuses on helping developers efficiently host machine learning models and improve inference performance. She is passionate about spatial data analysis and using her AI to solve social problems.