chore: adding a DB system

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,13 +32,40 @@ class EndpointHandler:
               A :obj:`list` | `dict`: will be serialized and returned
         """
 
-        # Get inputs
-        encrypted_inputs = from_json(data.pop("encrypted_inputs", data))
+        # Get method
+        method = data.pop("method", data)
+
+        if method == "save_key":
+
+            # Get keys
+            evaluation_keys = from_json(data.pop("evaluation_keys", data))
+
+            uid = np.random.randint(2**32)
+
+            while uid in self.key_database.keys():
+                uid = np.random.randint(2**32)
+
+            self.key_database[uid] = evaluation_keys
+
+            return {"uid": uid}
+
+        elif method == "inference":
+
+            uid = data.pop("uid", data)
+
+            assert uid in self.key_database.keys(), f"{uid} not in DB, {self.key_database.keys()=}"
+
+            # Get inputs
+            encrypted_inputs = from_json(data.pop("encrypted_inputs", 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)
 
-        # Get keys
-        evaluation_keys = from_json(data.pop("evaluation_keys", data))
+            return to_json(encrypted_prediction)
 
-        # Run CML prediction
-        encrypted_prediction = self.fhemodel_server.run(encrypted_inputs, evaluation_keys)
+        else:
 
-        return to_json(encrypted_prediction)
+            return

play_with_endpoint.py

@@ -29,6 +29,10 @@ headers = {
 
 def query(payload):
     response = requests.post(API_URL, headers=headers, json=payload)
+
+    if "error" in response:
+        assert False, f"Got an error: {response=}"
+
     return response.json()
 
 
@@ -60,6 +64,16 @@ 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)["uid"]
+print(f"Storing the key in the database under {uid=}")
+
 # Test the handler
 nb_good = 0
 nb_samples = len(X_test)
@@ -72,17 +86,17 @@ for i in range(nb_samples):
     # Quantize the input and encrypt it
     encrypted_inputs = fhemodel_client.quantize_encrypt_serialize(X_test[i].reshape(1, -1))
 
-    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,
     }
 
+    if verbose or True:
+        print(f"Size of the payload: {sys.getsizeof(payload) / 1024} kilobytes")
+
     # Run the inference on HF servers
     duration -= time.time()
     duration_inference = -time.time()
@@ -98,7 +112,7 @@ for i in range(nb_samples):
 
     if verbose or True:
         print(
-            f"for {i}-th input, {prediction=} with expected {Y_test[i]} in {duration_inference} seconds"
+            f"for {i}-th input, {prediction=} with expected {Y_test[i]} in {duration_inference:.3f} seconds"
         )
 
     # Measure accuracy