Run Inference on a Sensitive Model and Data using Python and Cages

Evervault Cages make it easy to build, deploy, and scale apps running in Docker containers inside secure enclaves — a cloud based Trusted Execution Environment (TEE). In this guide you will use Python and Cages to build an app that takes in a sensitive medical image, in this case a brain scan, run it against a model to determine whether areas of the scan might be cancerous, and returns the results.

Using this process, users can upload sensitive images to the Cage, while proving that their images cannot be leaked from the Cage, and in the case of running a sensitive model you can also ensure that the model stays protected. For this example, we will use a publicly available brain tumor segmentation model from Mateusz Buda that can be found here.

When run, this will be the outcome of the result.

Sample Image

Result

Prerequisites

• An AWS account (sign up or log in here)
• An Evervault account (create a free account here)
• Docker installed (get it here)

Set up

Clone the working GitHub repo.

1
git clone https://github.com/evervault/examples
2
cd cages-medical

To run Cages you will need to install the Cages CLI by running the following command.

1
curl https://cage-build-assets.evervault.com/cli/install -sL | sh

Add the model to S3

Because the Cage is protected from public access, you will need to make the model available using a data store of your choice — this example uses S3.

Download the model weights [tumor-model.pt](http://tumor-model.pt) from the repo and add this file to an S3 bucket. This will allow your Cage to access the model in order to run inference on it.

While in your AWS account, grab your access key and secret access key: You will need to add these to your Evervault account in a few steps.

The Python app

The back end of the app takes a model downloaded from S3 and runs inference on the model using a sample image supplied from the client. It will check to see which areas of the image may contain a tumor, then will generate a new image with the identified areas overlayed with a red highlight.

1
# Load the model
2
s3 = boto3.client('s3', aws_access_key_id=os.environ.get('ACCESS_KEY'), aws_secret_access_key=os.environ.get('SECRET_ACCESS_KEY'), region_name=os.environ.get('S3_REGION'))
3
s3.download_file(os.environ.get('BUCKET_NAME'),'tumor-model.pt','./tumor-model.pt')
4
model = torch.jit.load('./tumor-model.pt')
5
# Set the model to evaluation mode
6
model.eval()
7
app.logger.info('model downloaded')
8
9
@app.route('/upload', methods=['GET', 'POST'])
10
def upload():
11
    app.logger.info('upload received')
12
13
    pil_image = Image.open(request.files['upload'].stream)
14
    transform = transforms.Compose([...])
15
    input_tensor = transform(pil_image)
16
    with torch.no_grad():
17
        output_tensor = model(input_tensor.unsqueeze(0))
18
19
    output_image = output_image.resize(pil_image.size)
20
    mask_image = Image.new(mode='L', size=pil_image.size, color=0)
21
    mask_image.paste(output_image, box=(0, 0))
22
    overlay_image = pil_image.copy().convert('RGBA')
23
    overlay_image.paste((255, 0, 0, 128), box=(0, 0), mask=mask_image)
24
    img_io = io.BytesIO()
25
    
26
    return send_file(img_io, mimetype='image/png')

Open up the Dockerfile. You will use a virtual environment to install the required libraries needed to run the app. You’ll also tell Docker your webserver will listen on port 8008, which matches the Flask server port defined in app.py.

1
FROM python:3.8-slim
2
3
# create a virtual environment
4
RUN python3 -m venv /opt/venv
5
# install the requirements to the virtual environment
6
COPY requirements.txt requirements.txt
7
RUN /opt/venv/bin/pip install -r requirements.txt -f https://download.pytorch.org/whl/torch_stable.html
8
9
COPY app.py /app.py
10
11
# this must match the port your Flask server in app.py is running on
12
EXPOSE 8008
13
# use the python binary in the virtual environment we've set up to run our server
14
ENTRYPOINT ["/opt/venv/bin/python", "/app.py"]

Initialize the Cage

First, make sure that you have Docker running. Then, in your terminal run the following command to initalize the Cage. You can use the suggested name below or change it to one of your choosing.

1
ev-cage init -f ./Dockerfile \
2
--name med-cage
3
--disable-api-key-auth

You should see that a cert.pem, key.pem and cage.toml are generated. Open up the cage.toml. You can see that important details are generated that will help with the deployment of your Cage.

Add Environment Variables

Your Cage will need to access your S3 credentials as well as your Cage name and App ID. Add all of these variables by going to Cages > Environment in your Evervault Dashboard. Be sure to check the “secret” box on any sensitive credentials (like the access keys). This will encrypt them and they will only be decrypted in the Cage.

Build the Cage

Now build the Cage using the following command. This will also generate a Dockerfile used to generate the EIF file which will run inside the enclave (it may take a few minutes).

1
ev-cage build --write --output .

It will also provide a list of PCRs, cryptographic measurements that are required for the attestation process for the enclave. These PCRs will be injected into your cage.toml. You will need to add these to your client in order to make an attested request to the Cage (you can still make a request to the cage without these, it just won’t be attested).

Deploy the Cage

Finally, you can deploy the Cage

1
ev-cage deploy --eif-path ./enclave.eif

Setting up the Client

Copy the PCR values generated and add them into the empty JSON values that they correspond with in client_script.py.

1
cage_name = os.environ.get('CAGE_NAME')
2
app_id = os.environ.get('APP_ID')
3
4
attested_session = evervault.cage_requests_session({ 
5
    cage_name: {
6
        "PCR0": "",
7
        "PCR1": "",
8
        "PCR2": "",
9
        "PCR8": ""
10
  } 
11
})

The rest of this code will take the sample image provided and send it to the back end. It will also calculate the time elapsed to run the Cage so that you can monitor it to evaluate for performance.

1
filename = "sample-image.png"
2
with open(filename, "rb") as file:
3
    file_bytes = file.read()
4
5
fields = {
6
    "upload": (filename, file_bytes),
7
}
8
9
body, header = urllib3.encode_multipart_formdata(fields)
10
11
start = time.time()
12
13
response = attested_session.post(
14
    'https://{cage_name}.{app_id}.cages.evervault.com/upload',
15
    headers={'Content-Type': header},
16
	data=body
17
)
18
19
print(f'Time: {time.time() - start}')
20
print(response)
21
22
with open('response-image.png', 'wb') as output_file:
23
    output_file.write(response.content)

Run the Client

In the Client directory, open up .env.example and rename it to .env, then add your Cage name and app ID. — you can also find these in the cage.toml.

Now you may run the client script, which will pass the sample image into the Cage and run inference on the ML model.

1
cd client
2
python3 -m venv venv
3
source venv/bin/activate
4
pip3 install -r requirements.txt
5
python3 client_script.py

When run successfully, you should see response-image.png in your file tree. When opened, it will look like the below:

Conclusion

In this guide, you used Python to process a sensitive image on a model running inside a Cage. If you ran into an issue or have any questions about this guide, feel free to raise them on GitHub or drop them in the Evervault Discord. Let us know what you’re interested in using with secure enclaves!