使用 Amazon SageMaker AI 和 FHE 实现端到端加密的机器学习推理

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

aws ecr describe-images --repository-name <repo-name>

import argparse
import os
import numpy
from concrete.ml.sklearn import <Model class to train>
from concrete.ml.deployment import FHEModelDev

def do_training(model_dir, train):
    # 从 train 目录加载你的数据
    # 训练你的模型实例，然后使用以下代码保存它。
    FHEModelDev(model_dir, model).save()

def model_fn(model_dir):
    # SageMaker AI 需要这个函数存在，但不会使用它
    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)

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 = <Training container location>
        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

boto3==1.37.38
sagemaker==2.243.2

from flask import Flask
import flask
import logging
import json
from concrete.ml.deployment import FHEModelServer
from sagemaker.s3 import S3Uploader, S3Downloader

# 加载模型
try:
    model = FHEModelServer("/opt/ml/model/")
except Exception:
    logging.exception("初始化 FHEModelServer 失败")
    raise

app = Flask(__name__)

@app.route('/ping', methods=['GET'])
def ping():
    return flask.Response(response='\n', status=200, mimetype='application/json')

@app.route('/invocations', methods=['POST'])
def transformation():
    try:
        input_json = flask.request.get_json()
        if not input_json or not isinstance(input_json, dict):
            return flask.Response(
                response=json.dumps({"error": "Invalid JSON"}),
                status=400,
                mimetype="application/json",
            )
        required_keys = [
            "evaluation_keys_uri",
            "encrypted_query_uri",
        ]
        for key in required_keys:
            if key not in input_json:
                return flask.Response(response=f'Missing required field: {key}',
                                      status=400)
            if (not isinstance(input_json[key], str)
                    or not input_json[key].startswith('s3://')):
                return flask.Response(response=f'Invalid Amazon S3 URI for {key}', status=400)
        evaluation_keys_uri = input_json["evaluation_keys_uri"]
        encrypted_query_uri = input_json["encrypted_query_uri"]
        downloader = S3Downloader()
        try:
            evaluation_keys = downloader.read_bytes(evaluation_keys_uri)
            encrypted_query = downloader.read_bytes(encrypted_query_uri)
        except Exception as e:
            logging.error(f"Failed to download from S3: {e}")
            return flask.Response(response='Failed to retrieve data from Amazon S3',
                                  status=500)
        prediction = model.run(encrypted_query, evaluation_keys)
        return flask.Response(
            response=prediction, status=200, mimetype="application/octet-stream"
        )
    except KeyError as e:
        return flask.Response(
            response=json.dumps({"error": f"Missing key: {str(e)}"}),
            status=400,
            mimetype="application/json",
        )
    except Exception as e:
        return flask.Response(
            response=json.dumps({"error": "Internal server error"}),
            status=500,
            mimetype="application/json",
        )

  FROM python:3.12
  RUN apt-get -y update && apt-get install -y --no-install-recommends \
      nginx \
      ca-certificates \
      cmake \
      && rm -rf /var/lib/apt/lists/*
  RUN pip install flask gevent gunicorn sagemaker sagemaker_training==5.1.1 concrete-ml==1.9.0 concrete-python==2.10.0
  RUN rm -rf /root/.cache
  # 设置环境变量
  ENV PYTHONUNBUFFERED=TRUE
  ENV PYTHONDONTWRITEBYTECODE=TRUE
  ENV PATH="/opt/program:${PATH}"
  COPY <包含容器文件的目录>/ /opt/program
  RUN chmod +x /opt/program/serve
  WORKDIR /opt/program

  docker build -f ./Dockerfile.inference

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

  docker tag <image-id> <account-id>.dkr.ecr.<region>.amazonaws.com/<repo-name>:latest

  docker push <account-id>.dkr.ecr.<region>.amazonaws.com/<repo-name>:latest

  aws ecr describe-images --repository-name <repo-name>

  pip install boto3==1.37.38 sagemaker==2.243.2

  from sagemaker.session import Session
  from sagemaker.model import Model
  from sagemaker.predictor import Predictor
  from sagemaker.async_inference.async_inference_config import AsyncInferenceConfig

  sagemaker_session = Session()
  model = Model(
      image_uri="123456789012.dkr.ecr.us-east-1.amazonaws.com/concrete-inference:latest",  # 使用前面构建步骤中的 ECR URI
      model_data="s3://my-model-bucket/concrete-ml/models/model.tar.gz",  # 训练作业保存模型的路径
      role="arn:aws:iam::123456789012:role/SageMakerEndpointRole",  # 使用 Prerequisites 中的端点角色 ARN
      sagemaker_session=sagemaker_session,
      predictor_cls=Predictor,
  )
  async_config = AsyncInferenceConfig(
      max_concurrent_invocations_per_instance=1,
      output_path=<Amazon S3 存放结果密文的位置>,
      failure_path=<Amazon S3 存放推理失败的位置>,
  )
  endpoint = model.deploy(
      initial_instance_count=1,  # 测试时从 1 个实例开始
      instance_type="ml.m5.xlarge",  # FHE 的最低推荐；使用 ml.m5.24xlarge 以获得更好的性能
      wait=True,
      endpoint_logging=True,
      async_inference_config=async_config,
  )
  print(f"Endpoint name: {endpoint.endpoint_name}")

  python start_inference_endpoint.py

  aws sagemaker describe-endpoint --endpoint-name <endpoint-name>

import tempfile
import tarfile
import os
import json
import sagemaker
from sagemaker.s3 import S3Uploader, S3Downloader
from sagemaker.base_deserializers import BytesDeserializer
from sagemaker.base_serializers import JSONSerializer
from sagemaker.predictor import Predictor
from sagemaker.predictor_async import AsyncPredictor
from sagemaker.async_inference.waiter_config import WaiterConfig
from concrete.ml.deployment import FHEModelClient

sagemaker_session = sagemaker.Session()

predictor = AsyncPredictor(
    Predictor(
        <name of the endpoint created above>,
        serializer=JSONSerializer(),
        deserializer=BytesDeserializer(),
        sagemaker_session=sagemaker_session,
    )
)

model_location = <model Amazon S3 location>

def get_query():
    # 返回需要加密的查询的代码
    ...

# 下载并提取客户端配置
with tempfile.TemporaryDirectory() as config_dir_name:
    try:
        S3Downloader().download(
            model_location,
            local_path=config_dir_name,
            sagemaker_session=sagemaker_session,
        )
        tf = tarfile.open(os.path.join(config_dir_name, "model.tar.gz"), mode="r:gz")
        tf.extract("client.zip", config_dir_name)
    except FileNotFoundError as e:
        <handle exception>
    except tarfile.TarError as e:
        <handle exception>
    except Exception as e:
        <handle exception>

with tempfile.TemporaryDirectory() as key_dir_name:
    concrete_client = FHEModelClient(config_dir_name, key_dir=key_dir_name)

    # 生成并上传评估密钥
    eval_keys_location = <eval keys Amazon S3 location>
    concrete_client.generate_private_and_evaluation_keys()
    eval_keys = concrete_client.get_serialized_evaluation_keys()
    uploader = S3Uploader()
    uploader.upload_bytes(
        eval_keys,
        eval_keys_location,
        sagemaker_session=sagemaker_session
    )

    # 加密并上传查询
    encrypted_query_location = <Amazon S3 location for encrypted query>
    plaintext_query = get_query()
    encrypted_query = concrete_client.quantize_encrypt_serialize(plaintext_query)
    uploader.upload_bytes(
        encrypted_query,
        encrypted_query_location,
        sagemaker_session=sagemaker_session
    )

    # 向端点发送请求
    query = {
        'evaluation_keys_uri': eval_keys_location,
        'encrypted_query_uri': encrypted_query_location,
    }
    query_json = json.dumps(query)
    try:
        async_response = predictor.predict_async(
            data=query_json,
            input_path="<Amazon S3 location for the async query>",
            initial_args={"ContentType": "application/json"},
        )
        # 等待端点返回结果
        encrypted_result = async_response.get_result(
            waiter_config=WaiterConfig("<configuration values of your choice>")
        )
        prediction = concrete_client.deserialize_decrypt(encrypted_result)
    except TimeoutError as e:
        <handle exception>
    except Exception as e:
        <handle exception>