binoua commited on
Commit
df7187e
1 Parent(s): eb807a4

chore: adding a system to save key piece by piece

Browse files
Files changed (2) hide show
  1. handler.py +11 -0
  2. 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
- payload = {
58
- "inputs": "fake",
59
- "evaluation_keys": to_json(evaluation_keys),
60
- "method": "save_key",
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