Amazon SageMaker AI と FHE を用いたエンドツーエンド暗号化 ML 推論の紹介

Machine learning (ML) inference often requires processing sensitive data—medical records, proprietary business information, or personal communications. What if you could run ML inference in the cloud while hiding your data from the cloud itself? More specifically, what if you could enforce that your data stayed encrypted throughout the entire ML inference process? This post will show you how to use Amazon SageMaker AI with fully homomorphic encryption (FHE) to perform ML inference. Using FHE, we present an approach to ML inference that’s designed to keep queries, responses, and intermediate values encrypted and unreadable by observers—including SageMaker AI itself.

FHE is a form of encryption that allows encrypted data to be processed in encrypted form without decryption. In the ML inference setting, you can use it to apply a model to an encrypted query without decryption, producing an encrypted prediction. Consider these scenarios where such a capability would provide value:

• Healthcare: A health insurance company wants to provide doctors with an ML model that predicts medical procedure outcomes based on diagnostic data. Publishing the model in the cloud simplifies deployment, but doctors can’t expose patient medical information to third parties due to privacy regulations.

• Energy sector: An oil and gas corporation uses ML to evaluate satellite photos of potential drill sites and select photos for further expert evaluation. They want to host the model in the cloud for cost savings but can’t expose photographs of politically sensitive locations to third parties.

• Telecommunications: A telecom operator wants to process customer emails to detect spam and phishing. They need cloud-based ML for scalability, but data protection regulations require that customer messages remain encrypted at third parties.

This blog has previously discussed FHE for ML inference in the post Enable fully homomorphic encryption with Amazon SageMaker endpoints for secure, real-time inferencing, but this post goes a little further. That previous post showed how to implement FHE-based inference ‘from scratch’ by hand-crafting a linear-regression algorithm using a low-level library called SEAL. Instead, this post shows a much more flexible and higher-level approach based on concrete-ml, a high-level library built specifically for FHE-based inference. It supports several common types of models ‘out of the box’ and is even API compatible with the well-known ML library scikit-learn.

In this post, you will learn how to:

• Train a concrete-ml model in SageMaker AI using a custom container

• Deploy that model to a SageMaker AI inference endpoint

• Create a custom client for concrete-ml inference

• Use that client to make queries to your inference endpoint

When finished you will have a system that uses concrete-ml in SageMaker AI designed to perform end-to-end encrypted ML inference.

Solution overview

Using concrete-ml in SageMaker AI works as follows:

• The model owner prepares their data for training. Concrete-ml works well when all features have been normalized to the same scale, such as [-1, 1].

• The model owner uses this data to train an FHE-enabled version of their model. This model is designed to perform computations over encrypted data instead of plaintext.

• The model owner hosts this model in SageMaker AI.

• Clients encrypt their queries using the FHE scheme supported by the model.

• Clients send encrypted queries to the FHE-enabled model in the cloud.

• The model transforms the encrypted query into an encrypted prediction without decrypting values during the FHE computation.

• The model returns the encrypted response to the client, who decrypts it to retrieve the prediction.

This differs from, and complements, confidential computing environments like those provided by the Amazon Web Services (AWS) Nitro System in Amazon Elastic Compute Cloud (Amazon EC2). With AWS Nitro Enclaves, queries are decrypted and processed in plaintext within hardened, isolated environments that provide CPU and memory isolation. With FHE, queries remain encrypted throughout; security relies on mathematics rather than hardware or software.

Prerequisites

To implement this solution, you need:

• A local development environment with Python 3.12 installed, the ability to install packages using pip, and Docker or other container-building software installed locally. In addition, these instructions will recommend that you work in virtual environments, but this isn’t strictly necessary.

• An AWS account, containing:

Repositories in Amazon Elastic Container Registry (Amazon ECR) to hold the images for training and inference containers,

• Locations in Amazon Simple Storage Service (Amazon S3) to hold:

The model

• The training code (if you wish it to be stored in a separate bucket from the model)

• Keys and ciphertexts

We suggest you follow the security best practices for Amazon S3.

• Roles in AWS Identity and Access Management (IAM) for

The model creator

• The inference endpoint creator

• The inference endpoint itself

• The clients

Find IAM policies for these roles, along with a worked example for the MNIST corpus of handwritten digits, in the repository of sample code.

Before starting, note that at the time of writing, concrete-ml is available from Zama for prototyping or non-commercial use without requiring a paid license. However, you may require a commercial license for commercial use.

Training

Build and deploy the training container

To build the training container:

• Assume the model-trainer role.

• Create a Dockerfile.training file locally.

• Add the following content to Dockerfile.training:

FROM python:3.12
RUN apt-get update && apt-get upgrade -y && apt-get clean
RUN apt-get -y install --no-install-recommends cmake
RUN pip install sagemaker_training==5.1.1 concrete-ml==1.9.0 concrete-python==2.10.0 torch==2.3.1

Verify that the version numbers match across the entire system. The concrete-ml library requires version parity across the entire system for Python, the concrete-ml package, and the concrete-python package.

• Build the container image:

docker build -f ./Dockerfile.training

• Push the image to Amazon ECR:

Run the authentication command to log in Docker to your Amazon ECR registry:

aws ecr get-login-password --region  | docker login --username AWS --password-stdin .dkr.ecr..amazonaws.com

• Tag the image with your repository name:

docker tag  .dkr.ecr..amazonaws.com/:latest

• Push the tagged image:

docker push .dkr.ecr..amazonaws.com/:latest

Verify that the container is available

aws ecr describe-images --repository-name 

You should see JSON output containing your image with a non-empty imageDigest field and the latest tag.

Train the model

To train the model, complete the following.

Note: in these steps, concrete-ml is no different from any other ML framework and the training container is no different from any other custom training container. Note that training occurs over *plaintext* data. That is, concrete-ml doesn’t require pre-processing of this data beyond normalization. But if additional pre-processing is necessary for regular training, it remains necessary here (and must occur before, or as part of, the training job).

Create the training script

• Create a file named training_script.py.

• Add the following template code to training_script.py:

import argparse
import os
import numpy
from concrete.ml.sklearn import 
from concrete.ml.deployment import FHEModelDev

def do_training(model_dir, train):
    # Load your data from the train directory
    # Train your model instance, then save it
    # with the following line.
    FHEModelDev(model_dir, model).save()

def model_fn(model_dir):
    # SageMaker AI requires this function exist but doesn't use it
    raise NotImplementedError

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--model-dir', type=str, default=os.environ['SM_MODEL_DIR'])
    parser.add_argument('--train', type=str, default=os.environ['SM_CHANNEL_TRAINING'])
    args = parser.parse_args()
    do_training(args.model_dir, args.train)

• Implement the data loading logic in the do_training function.

• Implement the model training logic in the do_training function.

Create a custom framework

For convenience, we recommend that you create a custom framework to integrate your training container into SageMaker AI. To do so:

• Create a file named framework.py.

• Add the following content to framework.py:

from sagemaker.estimator import Framework

class Concrete(Framework):
    def __init__(
        self,
        entry_point,
        source_dir=None,
        hyperparameters=None,
        py_version="py312",
        framework_version="1.9.0",
        distributions=None,
        **kwargs,
    ):
        self.image_uri = 
        super(Concrete, self).__init__(
            entry_point, source_dir, hyperparameters,
            image_uri=self.image_uri,
            **kwargs
        )
        self.framework_version = framework_version
        self.py_version = py_version

    def training_image_uri(self, region=None):
        return self.image_uri

    def create_model(
        self,
        model_server_workers=None,
        role=None,
        vpc_config_override=None,
        entry_point=None,
        source_dir=None,
        dependencies=None,
        image_name=None,
        **kwargs,
    ):
        return None

• Update the image_uri value with your Amazon ECR training container location.

Launch the training job

This section will show how to launch the training job with a python script, but it can also be done using the console or the AWS Command Line Interface (AWS CLI). (Note: training jobs incur charges based on instance type and duration.)

• Create a virtual environment for Python 3.12.

• Activate the virtual environment.

• Install the following packages using pip:

boto3==1.37.38
sagemaker==2.243.2

• Create a file named start_training.py.

• Add the following content to start_training.py:

from sagemaker import session
from framework import Concrete

sagemaker_session = session.Session()

concrete = Concrete(
    entry_point="training_script.py",
    instance_count=1,
    instance_type="ml.m5.xlarge",  # Use ml.m5.xlarge for small models, ml.m5.4xlarge for larger models
    role="arn:aws:iam::123456789012:role/SageMakerModelTrainerRole",  # Use the model-trainer role ARN from Prerequisites
    sagemaker_session=sagemaker_session,
    hyperparameters={},
    output_path="s3://my-model-bucket/concrete-ml/models/",  # Use the model bucket from Prerequisites
    code_location="s3://my-model-bucket/concrete-ml/scripts/",  # S3 path for training script storage
)

concrete.fit(inputs=)

• Update the instance_type, role, output_path, code_location, and inputs values with your specific configuration.

• Execute this file:

python start_training.py

• Verify that the training completed successfully by checking the training job status:

aws sagemaker describe-training-job --training-job-name 

Look for TrainingJobStatus: Completed. Then verify that the output files exist:

aws s3 ls s3://my-model-bucket/concrete-ml/models/

Confirm server.zip and client.zip are present.

After training completes, the training container saves two files to the model bucket: server.zip (used by the inference endpoint) and client.zip (used by clients to encrypt queries).

Inference

Build and deploy the inference container

FHE-based ML inference will be more complex than standard ML inference because of some new technical constraints:

• Clients need model-specific information from client.zip to generate cryptographic keys.

• FHE ciphertexts can exceed SageMaker AI query size limits, so the client and service need to communicate them outside of SageMaker AI API calls.

• FHE evaluation might take longer than SageMaker AI timeouts, and so inference will use the SageMaker AI mechanisms for asynchronous inference.

• The endpoint needs an evaluation key (a type of public key) from the client to perform FHE evaluation.

To accommodate these new requirements and to streamline the user’s experience, we show you how to build a system in which

• A custom client encrypts queries and attaches evaluation keys to them

• A custom training endpoint retrieves client.zip when needed, and uses it to evaluate the FHE model

• The same custom client decrypts predictions from the training endpoint

• The client and endpoint communicate ciphertexts and keys to each other using Amazon S3

To deploy and use this system, complete the following sections.

Write your predictor

Create a file named [predictor.py](https://github.com/aws-samples/sample-end-to-end-encrypted-ml-inference-with-amazon-sagemaker-ai-and-fhe/blob/main/inference/endpoint/container_files/predictor.py) with the follo