File size: 3,047 Bytes
13fb76e
 
 
 
 
 
 
853622a
13fb76e
 
 
e5b8aae
 
 
13fb76e
 
 
 
 
 
58df7f1
 
 
 
 
 
 
13fb76e
 
 
 
 
 
 
 
58df7f1
 
49a1dd4
e3e1dc8
58df7f1
e3e1dc8
13fb76e
e3e1dc8
f5aa6c7
13fb76e
 
f5aa6c7
13fb76e
 
 
58df7f1
 
 
 
 
 
 
 
e3e1dc8
58df7f1
e3e1dc8
58df7f1
 
 
 
 
 
 
 
e5b8aae
58df7f1
 
13fb76e
58df7f1
 
13fb76e
58df7f1
 
 
13fb76e
58df7f1
13fb76e
 
58df7f1
 
e3e1dc8
13fb76e
58df7f1
49a1dd4
e3e1dc8
58df7f1
13fb76e
58df7f1
 
 
13fb76e
58df7f1
13fb76e
e3e1dc8
58df7f1
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
"""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

# Load the FHE server
FHE_SERVER = FHEModelServer(DEPLOYMENT_DIR)

# 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")

    # Receive the Client's files (Evaluation key + Encrypted symptoms)
    evaluation_key_path = SERVER_DIR / f"{user_id}_valuation_key"
    encrypted_input_path = SERVER_DIR / f"{user_id}_encrypted_input"

    # Save 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_input"

    # 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.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 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"

    # 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)