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, SERVER_DIR  # pylint: disable=no-name-in-module)

from concrete.ml.deployment import FHEModelServer

# 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!"}


@app.post("/send_input")
def send_input(
    user_id: str = Form(),
    files: List[UploadFile] = File(),
):
    """Send the inputs to the server."""

    print("\nSend the data to the server ............\n")
    # Retrieve the encrypted input and the evaluation key paths
    evaluation_key_path = SERVER_DIR / f"{user_id}_valuation_key"
    encrypted_input_path = SERVER_DIR / f"{user_id}_encrypted_symptoms"

    # # Write the files using the above paths
    with encrypted_input_path.open("wb") as encrypted_input, evaluation_key_path.open(
        "wb"
    ) as evaluation_key:
        encrypted_input.write(files[0].file.read())
        evaluation_key.write(files[1].file.read())


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

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

    # Read the files 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()

    # Load the FHE server and the model
    fhe_server = FHEModelServer(DEPLOYMENT_DIR)

    # Run the FHE execution
    start = time.time()
    encrypted_output = fhe_server.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"

    # 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")
def get_output(user_id: str = Form()):
    """Retrieve the encrypted output."""

    print("\nGet the output from the server ............\n")
    # Retrieve the encrypted output path
    encrypted_output_path = SERVER_DIR / f"{user_id}_encrypted_output"

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

    time.sleep(1)

    return Response(encrypted_output)