chore: using a database for keys
Browse files- handler.py +26 -2
- play_with_endpoint.py +12 -6
handler.py
CHANGED
@@ -20,6 +20,9 @@ class EndpointHandler:
|
|
20 |
# For server
|
21 |
self.fhemodel_server = FHEModelServer(path + "/compiled_model")
|
22 |
|
|
|
|
|
|
|
23 |
def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
|
24 |
"""
|
25 |
data args:
|
@@ -29,11 +32,32 @@ class EndpointHandler:
|
|
29 |
A :obj:`list` | `dict`: will be serialized and returned
|
30 |
"""
|
31 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
32 |
# Get inputs
|
33 |
encrypted_inputs = from_json(data.pop("encrypted_inputs", data))
|
34 |
|
35 |
-
#
|
36 |
-
evaluation_keys =
|
37 |
|
38 |
# Run CML prediction
|
39 |
encrypted_prediction = self.fhemodel_server.run(encrypted_inputs, evaluation_keys)
|
|
|
20 |
# For server
|
21 |
self.fhemodel_server = FHEModelServer(path + "/compiled_model")
|
22 |
|
23 |
+
# Simulate a database of keys
|
24 |
+
self.key_database = {}
|
25 |
+
|
26 |
def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
|
27 |
"""
|
28 |
data args:
|
|
|
32 |
A :obj:`list` | `dict`: will be serialized and returned
|
33 |
"""
|
34 |
|
35 |
+
# Get method
|
36 |
+
method = data.pop("method", data)
|
37 |
+
|
38 |
+
if method is "save_key":
|
39 |
+
|
40 |
+
# Get keys
|
41 |
+
evaluation_keys = from_json(data.pop("evaluation_keys", data))
|
42 |
+
|
43 |
+
uid = np.random.randint()
|
44 |
+
|
45 |
+
while uid is in self.key_database.keys():
|
46 |
+
uid = np.random.randint()
|
47 |
+
|
48 |
+
self.key_database[uid] = evaluation_keys
|
49 |
+
|
50 |
+
return uid
|
51 |
+
|
52 |
+
assert method == "inference":
|
53 |
+
|
54 |
+
uid = data.pop("uid", data)
|
55 |
+
|
56 |
# Get inputs
|
57 |
encrypted_inputs = from_json(data.pop("encrypted_inputs", data))
|
58 |
|
59 |
+
# Find key in the database
|
60 |
+
evaluation_keys = self.key_database[uid]
|
61 |
|
62 |
# Run CML prediction
|
63 |
encrypted_prediction = self.fhemodel_server.run(encrypted_inputs, evaluation_keys)
|
play_with_endpoint.py
CHANGED
@@ -46,6 +46,15 @@ fhemodel_client = FHEModelClient(path_to_model)
|
|
46 |
fhemodel_client.generate_private_and_evaluation_keys()
|
47 |
evaluation_keys = fhemodel_client.get_serialized_evaluation_keys()
|
48 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
49 |
# Test the handler
|
50 |
nb_good = 0
|
51 |
nb_samples = len(X_test)
|
@@ -58,15 +67,12 @@ for i in range(nb_samples):
|
|
58 |
# Quantize the input and encrypt it
|
59 |
encrypted_inputs = fhemodel_client.quantize_encrypt_serialize([X_test[i]])
|
60 |
|
61 |
-
|
62 |
-
print(f"Size of encrypted input: {sys.getsizeof(encrypted_inputs) / 1024 / 1024} megabytes")
|
63 |
-
print(f"Size of keys: {sys.getsizeof(evaluation_keys) / 1024 / 1024} megabytes")
|
64 |
-
|
65 |
-
# Prepare the payload, including the evaluation keys which are needed server side
|
66 |
payload = {
|
67 |
"inputs": "fake",
|
68 |
"encrypted_inputs": to_json(encrypted_inputs),
|
69 |
-
"
|
|
|
70 |
}
|
71 |
|
72 |
# Run the inference on HF servers
|
|
|
46 |
fhemodel_client.generate_private_and_evaluation_keys()
|
47 |
evaluation_keys = fhemodel_client.get_serialized_evaluation_keys()
|
48 |
|
49 |
+
# Save the key in the database
|
50 |
+
payload = {
|
51 |
+
"inputs": "fake",
|
52 |
+
"evaluation_keys": to_json(evaluation_keys),
|
53 |
+
"method": "save_key",
|
54 |
+
}
|
55 |
+
|
56 |
+
uid = query(payload)
|
57 |
+
|
58 |
# Test the handler
|
59 |
nb_good = 0
|
60 |
nb_samples = len(X_test)
|
|
|
67 |
# Quantize the input and encrypt it
|
68 |
encrypted_inputs = fhemodel_client.quantize_encrypt_serialize([X_test[i]])
|
69 |
|
70 |
+
# Prepare the payload
|
|
|
|
|
|
|
|
|
71 |
payload = {
|
72 |
"inputs": "fake",
|
73 |
"encrypted_inputs": to_json(encrypted_inputs),
|
74 |
+
"method": "inference",
|
75 |
+
"uid": uid,
|
76 |
}
|
77 |
|
78 |
# Run the inference on HF servers
|