Revolutionizing Inference Workflows with Amazon SageMaker Python SDK

Enhanced Capabilities for Generative AI and Machine Learning Models

Key Improvements and User Experience

Solution Overview

Python Inference Orchestration

Build and Deploy the Workflow

Invoke the Endpoint

Customer Story: Amazon Search

Clean Up

Conclusion

About the Authors

Unlocking the Power of AI Inference Workflows with Amazon SageMaker

As the demand for advanced machine learning (ML) and generative AI applications grows, businesses need robust tools for deploying these technologies effectively. Amazon SageMaker Inference has emerged as a leading solution, enabling seamless deployment of multiple models to handle inference requests at scale. Recognizing this evolution, we’re excited to introduce a groundbreaking capability in the SageMaker Python SDK that transforms how inference workflows are built and deployed.

The Rise of Multi-Model Inference Workflows

Today’s AI applications often require complex interactions among multiple models. From processing interconnected inference requests to orchestrating workflows that involve generative AI, businesses need solutions that go beyond simple model deployment.

Addressing Complexity with SageMaker Python SDK

To cater to the growing demand for sophisticated inference capabilities, our new enhancements in the SageMaker Python SDK simplify the process of managing interconnected models. Using Amazon Search as a case study, this post will demonstrate how the new features empower businesses to build streamlined inference workflows effectively.

Overview of the User Experience

Imagine building a coordinated ensemble of models that work together to enhance user experiences. This new SDK capability provides a user-friendly interface to create and manage inference workflows. You can deploy multiple models directly within a single SageMaker endpoint, significantly reducing complexity and improving resource utilization.

Key Improvements and User Experience

Here’s a closer look at the improvements introduced in the SageMaker Python SDK:

• Deployment of Multiple Models: With the new tooling, ML teams can now deploy several models as inference components in one endpoint, simplifying operational management and enhancing cost-effectiveness.

• Workflow Definition with Workflow Mode: This new feature allows users to define complex inference workflows directly in Python, making it easier to connect models and specify how data flows between them.

• Development and Deployment Options: A new quick-deployment option enables rapid testing and refinement of workflows—ideal for teams experimenting with different configurations.

• Invocation Flexibility: Users can invoke specific models or entire workflows based on their needs, offering granular control over how and when to execute inference requests.

• Dependency Management: The SDK supports various model-serving libraries, reducing the complexity associated with environment setup.

Building Complex Inference Workflows

To get started, you can use the SageMaker Python SDK to deploy your models as inference components and then create a workflow using Python code. With this streamlined approach, developers can focus on business logic and model integrations.

Solution Architecture

The new SDK introduces powerful components and classes, such as:

• ModelBuilder: Automates the packaging process for individual models as inference components, handling everything from model loading to dependency management.

• CustomOrchestrator: A standardized template class that allows users to define custom inference logic and orchestrate multiple models in a workflow seamlessly.

Let’s see how these components come together to build an inference workflow with Amazon SageMaker.

Example: IT Customer Service Workflow

1. Define Custom Orchestration Class: Extend the capabilities of CustomOrchestrator to process and pass data between models:

   class PythonCustomInferenceEntryPoint(CustomOrchestrator):
      def preprocess(self, data):
          payload = {"inputs": data.decode("utf-8")}
          return json.dumps(payload)
   
      def handle(self, data, context=None):
          return self._invoke_workflow(data)

2. Build and Deploy the Workflow: Use ModelBuilder instances for each model and consolidate them into a workflow for deployment.

3. Invoke the Endpoint: Once deployed, you can easily maintain and test all components using the predictor functionality from the SDK.

Customer Example: Amazon Search

Amazon Search is leveraging the enhanced SageMaker Python SDK to refine its sophisticated search ranking workflows. By optimizing model integration and management, Amazon Search aims to deliver more relevant results tailored to user queries, whether the customer is browsing electronics, fashion, or other categories.

Vaclav Petricek, Sr. Manager of Applied Science at Amazon Search, emphasizes the benefit: “These capabilities represent a significant advancement in our ability to develop and deploy sophisticated inference workflows.”

Conclusion

The enhancements in the SageMaker Python SDK mark a pivotal moment for businesses aiming to deploy advanced AI applications effectively. By streamlining the process of building and managing inference workflows, developers can focus on delivering value through innovation rather than getting bogged down by infrastructure challenges.

Whether you are deploying classic ML models, building complex AI applications, or creating multi-step inference workflows, the newly enhanced SDK provides the flexibility, ease of use, and scalability necessary to bring your vision to life.

We encourage all SageMaker users to explore these new capabilities, which empower businesses to evolve alongside the rapidly changing landscape of AI applications.

Ready to transform your AI inference workflows? Start building with the SageMaker Python SDK today!

About the Authors

Melanie Li, PhD, is a Senior Generative AI Specialist Solutions Architect at AWS.
Saurabh Trikande is a Senior Product Manager for Amazon Bedrock.
Osho Gupta is a Senior Software Developer at AWS SageMaker.
Joseph Zhang is a software engineer at AWS.
Gary Wang is a Software Developer at AWS SageMaker.
James Park is a Solutions Architect at Amazon Web Services.
Vaclav Petricek is a Senior Applied Science Manager at Amazon Search.
Wei Li is a Senior Software Dev Engineer in Amazon Search.
Brian Granger is a Senior Principal Technologist at Amazon Web Services.

By leveraging advanced capabilities in the SageMaker Python SDK, these professionals are committed to paving the way for the next generation of AI applications.