server.py

"""Server that will listen for GET and POST requests from the client."""

import time
from typing import List

from fastapi import FastAPI, File, Form, UploadFile
from fastapi.responses import JSONResponse, Response
from utils import DEPLOYMENT_DIR_MODEL1, DEPLOYMENT_DIR_MODEL2, DEPLOYMENT_DIR_MODEL3, SERVER_DIR  # pylint: disable=no-name-in-module

from concrete.ml.deployment import FHEModelServer

# Load the FHE server
FHE_SERVER_MODEL1 = FHEModelServer(DEPLOYMENT_DIR_MODEL1)
FHE_SERVER_MODEL2 = FHEModelServer(DEPLOYMENT_DIR_MODEL2)
FHE_SERVER_MODEL3 = FHEModelServer(DEPLOYMENT_DIR_MODEL3)

# Initialize an instance of FastAPI
app = FastAPI()

# Define the default route
@app.get("/")
def root():
    """
    Root endpoint of the health prediction API.

    Returns:
        dict: The welcome message.
    """
    return {"message": "Welcome to your disease prediction with FHE made by Dhiria with love!"}


@app.post("/send_input_first_layer")
def send_input_first_layer(
    user_id: str = Form(),
    files: List[UploadFile] = File(),
):
    """Send the inputs of the first layer of the hierarchical models to the server.
    Order of the files:
    files[0]: encrypted input_model1
    files[1]: encrypted input_model2
    files[2]: evaluation keys
    """
    print('Receiving data for the first layer of models...')

    # Receive the Client's files (Evaluation key + Encrypted inputs for the models)
    evaluation_key_path_1 = SERVER_DIR / f"{user_id}_evaluation_key_1"
    evaluation_key_path_2 = SERVER_DIR / f"{user_id}_evaluation_key_2"
    encrypted_input_path_model1 = SERVER_DIR / f"{user_id}_encrypted_input_model1"
    encrypted_input_path_model2 = SERVER_DIR / f"{user_id}_encrypted_input_model2"

    # Save the files using the above paths
    with encrypted_input_path_model1.open("wb") as encrypted_input_model1, \
         encrypted_input_path_model2.open("wb") as encrypted_input_model2, \
         evaluation_key_path_1.open("wb") as evaluation_key_1, \
         evaluation_key_path_2.open("wb") as evaluation_key_2:
        encrypted_input_model1.write(files[0].file.read())
        encrypted_input_model2.write(files[1].file.read())
        evaluation_key_1.write(files[2].file.read())
        evaluation_key_2.write(files[3].file.read())


@app.post("/send_input_second_layer")
def send_input_second_layer(
    user_id: str = Form(),
    files: List[UploadFile] = File(),
):
    """Send the input of the second layer of the hierarchical models to the server.
    Order of the files:
    files[0]: encrypted input_model1
    files[1]: evaluation keys
    """
    print('Receiving data for the second layer of models...')

    # Receive the Client's files (Evaluation key + Encrypted inputs for the models)
    evaluation_key_path = SERVER_DIR / f"{user_id}_evaluation_key_second_layer"
    encrypted_input_path_model3 = SERVER_DIR / f"{user_id}_encrypted_input_model3"

    # Save the files using the above paths
    with encrypted_input_path_model3.open("wb") as encrypted_input_model3, \
         evaluation_key_path.open("wb") as evaluation_key:
        encrypted_input_model3.write(files[0].file.read())
        evaluation_key.write(files[1].file.read())


@app.post("/run_fhe_first_layer")
def run_fhe_first_layer(
    user_id: str = Form(),
):
    """Inference in FHE."""

    print("\nRun in FHE in the server ............\n")
    evaluation_key_path_1 = SERVER_DIR / f"{user_id}_evaluation_key_1"
    evaluation_key_path_2 = SERVER_DIR / f"{user_id}_evaluation_key_2"
    encrypted_input_path_model1 = SERVER_DIR / f"{user_id}_encrypted_input_model1"
    encrypted_input_path_model2 = SERVER_DIR / f"{user_id}_encrypted_input_model2"

    # Read the files (Evaluation key + Encrypted symptoms) using the above paths
    with encrypted_input_path_model1.open("rb") as encrypted_output_file_model1,  \
         encrypted_input_path_model2.open("rb") as encrypted_output_file_model2,  \
         evaluation_key_path_1.open("rb") as evaluation_key_file_1, \
         evaluation_key_path_2.open("rb") as evaluation_key_file_2:
        encrypted_output_model1 = encrypted_output_file_model1.read()
        encrypted_output_model2 = encrypted_output_file_model2.read()
        evaluation_key_1 = evaluation_key_file_1.read()
        evaluation_key_2 = evaluation_key_file_2.read()

    # Run the FHE execution of the first layer of models

    start = time.time()
    encrypted_output_model1 = FHE_SERVER_MODEL1.run(encrypted_output_model1, evaluation_key_1)
    encrypted_output_model2 = FHE_SERVER_MODEL2.run(encrypted_output_model2, evaluation_key_2)
    assert isinstance(encrypted_output_model1, bytes)
    assert isinstance(encrypted_output_model2, bytes)
    fhe_execution_time = round(time.time() - start, 2)

    # Retrieve the encrypted output path
    encrypted_output_path_model1 = SERVER_DIR / f"{user_id}_encrypted_output_model1"
    encrypted_output_path_model2 = SERVER_DIR / f"{user_id}_encrypted_output_model2"

    # Write the file using the above path
    with encrypted_output_path_model1.open("wb") as f1, \
         encrypted_output_path_model2.open("wb") as f2:
        f1.write(encrypted_output_model1)
        f2.write(encrypted_output_model2)

    return JSONResponse(content=fhe_execution_time)


@app.post("/run_fhe_second_layer")
def run_fhe_second_layer(
    user_id: str = Form(),
):
    """Inference in FHE."""

    print("\nRun in FHE in the server ............\n")
    evaluation_key_path = SERVER_DIR / f"{user_id}_evaluation_key_second_layer"
    encrypted_input_path = SERVER_DIR / f"{user_id}_encrypted_input_model3"

    # Read the files (Evaluation key + Encrypted symptoms) using the above paths
    with encrypted_input_path.open("rb") as encrypted_output_file, \
         evaluation_key_path.open("rb") as evaluation_key_file:
        encrypted_output = encrypted_output_file.read()
        evaluation_key = evaluation_key_file.read()

    # Run the FHE execution
    start = time.time()
    encrypted_output = FHE_SERVER_MODEL3.run(encrypted_output, evaluation_key)
    assert isinstance(encrypted_output, bytes)
    fhe_execution_time = round(time.time() - start, 2)

    # Retrieve the encrypted output path
    encrypted_output_path = SERVER_DIR / f"{user_id}_encrypted_output_model3"

    # Write the file using the above path
    with encrypted_output_path.open("wb") as f:
        f.write(encrypted_output)

    return JSONResponse(content=fhe_execution_time)


@app.post("/get_output_first_layer_1")
def get_output_first_layer_1(user_id: str = Form()):
    """Retrieve the encrypted outputs of the first layers of models from the server."""

    print("\nGet the output from the server ............\n")

    # Path where the encrypted output is saved
    encrypted_output_path_model1 = SERVER_DIR / f"{user_id}_encrypted_output_model1"
    # encrypted_output_path_model2 = SERVER_DIR / f"{user_id}_encrypted_output_model2"

    # Read the file using the above path
    with encrypted_output_path_model1.open("rb") as f1:
        #  encrypted_output_path_model2.open("rb") as f2:
        encrypted_output_1 = f1.read()
        # encrypted_output_2 = f2.read()

    time.sleep(1)

    # Send the encrypted output
    return Response(encrypted_output_1)


@app.post("/get_output_first_layer_2")
def get_output_first_layer_2(user_id: str = Form()):
    """Retrieve the encrypted outputs of the first layers of models from the server."""

    print("\nGet the output from the server ............\n")

    # Path where the encrypted output is saved
    # encrypted_output_path_model1 = SERVER_DIR / f"{user_id}_encrypted_output_model1"
    encrypted_output_path_model2 = SERVER_DIR / f"{user_id}_encrypted_output_model2"

    # Read the file using the above path
    with encrypted_output_path_model2.open("rb") as f2:
        #  encrypted_output_path_model2.open("rb") as f2:
        # encrypted_output_1 = f1.read()
        encrypted_output_2 = f2.read()

    time.sleep(1)

    # Send the encrypted output
    return Response(encrypted_output_2)


@app.post("/get_output_second_layer")
def get_output_second_layer(user_id: str = Form()):
    """Retrieve the encrypted outputs of the first layers of models from the server."""

    print("\nGet the output from the server ............\n")

    # Path where the encrypted output is saved
    encrypted_output_path = SERVER_DIR / f"{user_id}_encrypted_output_model3"

    # Read the file using the above path
    with encrypted_output_path.open("rb") as f:
        encrypted_output = f.read()

    time.sleep(1)

    # Send the encrypted output
    return Response(encrypted_output)