|
from typing import Dict, List, Any |
|
import numpy as np |
|
from concrete.ml.deployment import FHEModelServer |
|
|
|
|
|
def from_json(python_object): |
|
if "__class__" in python_object: |
|
return bytes(python_object["__value__"]) |
|
|
|
|
|
def to_json(python_object): |
|
if isinstance(python_object, bytes): |
|
return {"__class__": "bytes", "__value__": list(python_object)} |
|
raise TypeError(repr(python_object) + " is not JSON serializable") |
|
|
|
|
|
class EndpointHandler: |
|
def __init__(self, path=""): |
|
|
|
|
|
self.fhemodel_server = FHEModelServer(path + "/compiled_model") |
|
|
|
|
|
self.key_database = {} |
|
|
|
def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]: |
|
""" |
|
data args: |
|
inputs (:obj: `str`) |
|
date (:obj: `str`) |
|
Return: |
|
A :obj:`list` | `dict`: will be serialized and returned |
|
""" |
|
|
|
|
|
method = data.pop("method", data) |
|
|
|
if method == "save_key": |
|
|
|
|
|
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 == "append_key": |
|
|
|
|
|
evaluation_keys = from_json(data.pop("evaluation_keys", data)) |
|
|
|
uid = data.pop("uid", data) |
|
|
|
self.key_database[uid] += evaluation_keys |
|
|
|
return |
|
|
|
elif method == "inference": |
|
|
|
uid = data.pop("uid", data) |
|
|
|
assert uid in self.key_database.keys(), f"{uid} not in DB, {self.key_database.keys()=}" |
|
|
|
|
|
encrypted_inputs = from_json(data.pop("encrypted_inputs", data)) |
|
|
|
|
|
evaluation_keys = self.key_database[uid] |
|
|
|
|
|
encrypted_prediction = self.fhemodel_server.run(encrypted_inputs, evaluation_keys) |
|
|
|
return to_json(encrypted_prediction) |
|
|
|
else: |
|
|
|
return |
|
|
