chore: using a database for keys

handler.py

@@ -20,6 +20,9 @@ class EndpointHandler:
         # For server
         self.fhemodel_server = FHEModelServer(path + "/compiled_model")
 
+        # Simulate a database of keys
+        self.key_database = {}
+
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
         """
          data args:
@@ -29,11 +32,32 @@ class EndpointHandler:
               A :obj:`list` | `dict`: will be serialized and returned
         """
 
+        # Get method
+        method = data.pop("method", data)
+
+        if method is "save_key":
+
+            # Get keys
+            evaluation_keys = from_json(data.pop("evaluation_keys", data))
+
+            uid = np.random.randint()
+
+            while uid is in self.key_database.keys():
+                uid = np.random.randint()
+
+            self.key_database[uid] = evaluation_keys
+
+            return uid
+
+        assert method == "inference":
+
+        uid = data.pop("uid", data)
+
         # Get inputs
         encrypted_inputs = from_json(data.pop("encrypted_inputs", data))
 
-        # Get keys
-        evaluation_keys = from_json(data.pop("evaluation_keys", data))
+        # Find key in the database
+        evaluation_keys = self.key_database[uid]
 
         # Run CML prediction
         encrypted_prediction = self.fhemodel_server.run(encrypted_inputs, evaluation_keys)

play_with_endpoint.py

@@ -46,6 +46,15 @@ fhemodel_client = FHEModelClient(path_to_model)
 fhemodel_client.generate_private_and_evaluation_keys()
 evaluation_keys = fhemodel_client.get_serialized_evaluation_keys()
 
+# Save the key in the database
+payload = {
+    "inputs": "fake",
+    "evaluation_keys": to_json(evaluation_keys),
+    "method": "save_key",
+}
+
+uid = query(payload)
+
 # Test the handler
 nb_good = 0
 nb_samples = len(X_test)
@@ -58,15 +67,12 @@ for i in range(nb_samples):
     # Quantize the input and encrypt it
     encrypted_inputs = fhemodel_client.quantize_encrypt_serialize([X_test[i]])
 
-    if verbose:
-        print(f"Size of encrypted input: {sys.getsizeof(encrypted_inputs) / 1024 / 1024} megabytes")
-        print(f"Size of keys: {sys.getsizeof(evaluation_keys) / 1024 / 1024} megabytes")
-
-    # Prepare the payload, including the evaluation keys which are needed server side
+    # Prepare the payload
     payload = {
         "inputs": "fake",
         "encrypted_inputs": to_json(encrypted_inputs),
-        "evaluation_keys": to_json(evaluation_keys),
+        "method": "inference",
+        "uid": uid,
     }
 
     # Run the inference on HF servers