Privacy Preserving Machine Learning Demo using Tensorflow

This demo uses Gramine to run a Tensorflow Model Server in an SGX enclave, and MarbleRun to take care of attestation and secret provisioning.

How it works

1. The model owner encrypts the model and uploads the encrypted model to a cloud storage
2. The administrator deploys MarbleRun with a manifest defining the topology and components of the confidential ML deployment
3. The administrator deploys the confidential ML application.
4. MarbleRun takes care of authentication and bootstrapping procedures.
5. The model owner verifies the deployment via MarbleRun and uploads the encryption key securely to the TensorFlow Serving application via MarbleRun’s secret distribution.
6. The application can decrypt the model inside the enclave via the provisioned key.
7. Clients can verify the deployment via MarbleRun and connect securely to the inference service, knowing that their data is only accessible inside the enclave and their predictions are made by the integrity-protected TensorFlow Serving application.

Install dependencies

To run the python scripts we need python3 and some extra libraries. Make sure pip is up to date and run:

pip3 install --upgrade pip
pip3 install -r ./client/requirements.txt
pip3 install grpcio~=1.34.0

To encrypt the model we need Gramine's pf-crypt tool. For installation instructions refer to the Gramine GitHub

Running the demo

We provide a docker image to run TensorFlow Serving with Gramine and MarbleRun. You can also build it yourself.

On Kubernetes

This tutorial will show you how to run the demo on Kubernetes. A running cluster is required.

Make sure your cluster supports SGX and out-of-process attestation. You can follow the guide by Microsoft to create a AKS cluster with all the needed resources.

If you built your own image you will have to change the image name in kubernetes/templates/tf-server.yaml.

1. Start the MarbleRun coordinator

   marblerun install

2. Wait for MarbleRun to set-up

   marblerun check

3. Port-forward the client API service to localhost

   kubectl -n marblerun port-forward svc/coordinator-client-api 4433:4433 --address localhost >/dev/null &
   export MARBLERUN=localhost:4433

4. Download and convert the model

   ./tools/download_model.sh
   models_abs_dir=`pwd -P`
   python3 ./tools/model_graph_to_saved_model.py --import_path ${models_abs_dir}/models/resnet50-v15-fp32/resnet50-v15-fp32.pb --export_dir ${models_abs_dir}/models/resnet50-v15-fp32 --model_version 1 --inputs input --outputs predict
   mkdir plain
   mv models/resnet50-v15-fp32/1/saved_model.pb plain/

5. Use Gramine's pf-crypt to generate a key and encrypt the model.

   gramine-sgx-pf-crypt gen-key --wrap-key model_key
   gramine-sgx-pf-crypt encrypt --input plain/saved_model.pb --output models/resnet50-v15-fp32/1/saved_model.pb --wrap-key model_key

6. Generate a user certificate and key.

   openssl req -x509 -newkey rsa:4096 -sha256 -nodes -keyout user_credentials.key -out user_credentials.crt

7. Insert the output of the following command as Certificate for user tf-admin in manifest.json (replacing USER_CERT)

   awk 'NF {sub(/\r/, ""); printf "%s\\n",$0;}' user_credentials.crt

8. Upload the manifest:

   marblerun manifest set manifest.json $MARBLERUN

9. Upload the model key to MarbleRun.

   sed -i "s|KEY_DATA|$(cat model_key | base64)|g" pfKey.json
   marblerun secret set pfKey.json $MARBLERUN --key user_credentials.key --cert user_credentials.crt

10. Start the Tensorflow Model Server

    helm install -f ./kubernetes/values.yaml tensorflow-demo ./kubernetes --create-namespace -n tensorflow

11. Upload the model to Kubernetes

    kubectl cp ./models/resnet50-v15-fp32/1/saved_model.pb tensorflow/`kubectl -n tensorflow get pods --template '{{range .items}}{{.metadata.name}}{{"\n"}}{{end}}'`:/tensorflow-marblerun/models/resnet50-v15-fp32/1/saved_model.pb

12. Get MarbleRun's certificate

    marblerun certificate intermediate $MARBLERUN -o tensorflow.crt

13. Get TensorFlow's domain name

    Usually, one would set up DNS resolution for the cluster. To keep things simple, we will create a mapping of the TensorFlow Model Server IP to a domain name using /etc/hosts

    • First get the IP Adress:

      tf_ip_addr=`kubectl get svc -n tensorflow -o jsonpath='{.items[0].status.loadBalancer.ingress[0].ip}'`

    • Set the mapping in /etc/hosts:

      echo "${tf_ip_addr} grpc.tensorflow-serving.com" >> /etc/hosts

14. Submit a request using encrypted traffic over gRPC

    python3 ./client/resnet_client_grpc.py --url grpc.tensorflow-serving.com:8500 --crt ./tensorflow.crt --batch 1 --cnum 1 --loop 10

Cleaning up

1. Remove tensorflow from the cluster

   helm uninstall tensorflow-demo -n tensorflow
   kubectl delete namespace tensorflow

2. Uninstall MarbleRun

   marblerun uninstall

Standalone

You can run the demo with MarbleRun in standalone mode as follows:

1. Start MarbleRun

   erthost ${marblerun_dir}/build/coordinator-enclave.signed
   export MARBLERUN=localhost:4433

2. Download and convert the model

   ./tools/download_model.sh
   models_abs_dir=`pwd -P`
   python3 ./tools/model_graph_to_saved_model.py --import_path ${models_abs_dir}/models/resnet50-v15-fp32/resnet50-v15-fp32.pb --export_dir ${models_abs_dir}/models/resnet50-v15-fp32 --model_version 1 --inputs input --outputs predict
   mkdir plain
   mv models/resnet50-v15-fp32/1/saved_model.pb plain/

3. Use Gramine's pf-crypt to generate a key and encrypt the model.

   gramine-sgx-pf-crypt gen-key --wrap-key model_key
   gramine-sgx-pf-crypt encrypt --input plain/saved_model.pb --output models/resnet50-v15-fp32/1/saved_model.pb --wrap-key model_key

4. Generate a user certificate and key.

   openssl req -x509 -newkey rsa:4096 -sha256 -nodes -keyout user_credentials.key -out user_credentials.crt

5. Insert the output of the following command as Certificate for user tf-admin in manifest.json(replacing USER_CERT)

   awk 'NF {sub(/\r/, ""); printf "%s\\n",$0;}' user_credentials.crt

6. Upload the manifest:

   marblerun manifest set manifest.json $MARBLERUN

7. Upload the model key to MarbleRun.

   sed -i "s|KEY_DATA|$(cat model_key | base64)|g" pfKey.json
   marblerun secret set pfKey.json $MARBLERUN --key user_credentials.key --cert user_credentials.crt

8. Start the Tensorflow Model Server This will pull our docker image. If you wish to use your own, specify the name of your image instead.

   ./tools/run_tf_image.sh

9. Get MarbleRun's intermediate certificate to connect to the model server.

   marblerun certificate intermediate $MARBLERUN -o tensorflow.crt

10. Test the model server using the gRPC client

    python3 ./client/resnet_client_grpc.py --url localhost:8500 --crt tensorflow.crt --batch 1 --cnum 1 --loop 10

Building the Docker Image

1. Generate a signing key

   openssl genrsa -3 -out enclave-key.pem 3072

2. Next we can build the Docker image:

   DOCKER_BUILDKIT=1 docker build --secret id=signingkey,src=<path to private.pem> --tag ghcr.io/edgelesssys/tensorflow-gramine-marble:latest .