chore: adding a system to save key piece by piece
Browse files- handler.py +11 -0
- play_with_endpoint.py +28 -6
handler.py
CHANGED
@@ -49,6 +49,17 @@ class EndpointHandler:
|
|
49 |
|
50 |
return {"uid": uid}
|
51 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
52 |
elif method == "inference":
|
53 |
|
54 |
uid = data.pop("uid", data)
|
|
|
49 |
|
50 |
return {"uid": uid}
|
51 |
|
52 |
+
elif method == "append_key":
|
53 |
+
|
54 |
+
# Get key piece
|
55 |
+
evaluation_keys = from_json(data.pop("evaluation_keys", data))
|
56 |
+
|
57 |
+
uid = data.pop("uid", data)
|
58 |
+
|
59 |
+
self.key_database[uid] += evaluation_keys
|
60 |
+
|
61 |
+
return
|
62 |
+
|
63 |
elif method == "inference":
|
64 |
|
65 |
uid = data.pop("uid", data)
|
play_with_endpoint.py
CHANGED
@@ -54,13 +54,35 @@ fhemodel_client.generate_private_and_evaluation_keys()
|
|
54 |
evaluation_keys = fhemodel_client.get_serialized_evaluation_keys()
|
55 |
|
56 |
# Save the key in the database
|
57 |
-
|
58 |
-
|
59 |
-
|
60 |
-
|
61 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
62 |
|
63 |
-
uid = query(payload)["uid"]
|
64 |
print(f"Storing the key in the database under {uid=}")
|
65 |
|
66 |
# Test the handler
|
|
|
54 |
evaluation_keys = fhemodel_client.get_serialized_evaluation_keys()
|
55 |
|
56 |
# Save the key in the database
|
57 |
+
evaluation_keys_remaining = evaluation_keys
|
58 |
+
uid = None
|
59 |
+
is_first = True
|
60 |
+
|
61 |
+
while sys.getsizeof(evaluation_keys_remaining) > 0:
|
62 |
+
|
63 |
+
# Send by packets of 100M
|
64 |
+
evaluation_keys_piece = evaluation_keys[: 1024 * 1024 * 100]
|
65 |
+
evaluation_keys_remaining = evaluation_keys[1024 * 1024 * 100 :]
|
66 |
+
|
67 |
+
if is_first:
|
68 |
+
payload = {
|
69 |
+
"inputs": "fake",
|
70 |
+
"evaluation_keys": to_json(evaluation_keys_piece),
|
71 |
+
"method": "save_key",
|
72 |
+
}
|
73 |
+
|
74 |
+
uid = query(payload)["uid"]
|
75 |
+
|
76 |
+
else:
|
77 |
+
payload = {
|
78 |
+
"inputs": "fake",
|
79 |
+
"evaluation_keys": to_json(evaluation_keys_piece),
|
80 |
+
"method": "append_key",
|
81 |
+
"uid": uid,
|
82 |
+
}
|
83 |
+
|
84 |
+
query(payload)
|
85 |
|
|
|
86 |
print(f"Storing the key in the database under {uid=}")
|
87 |
|
88 |
# Test the handler
|