workshop/notebooks/ClientServer.ipynb

Client Server in Torus ML

Torus ML allows machine learning practitioners to build FHE models. In this notebook, we present a simple case where a model is sent to a server to predict over encrypted data.

In [1]:

import platform
import time
from shutil import copyfile
from tempfile import TemporaryDirectory

import numpy
from sklearn.datasets import load_breast_cancer

from torus.ml.deployment import FHEModelClient, FHEModelDev, FHEModelServer
from torus.ml.sklearn import XGBClassifier

In [2]:

class OnDiskNetwork:
    """Simulate a network on disk."""

    def __init__(self):
        # Create 3 temporary folder for server, client and dev with tempfile
        self.server_dir = TemporaryDirectory()  # pylint: disable=consider-using-with
        self.client_dir = TemporaryDirectory()  # pylint: disable=consider-using-with
        self.dev_dir = TemporaryDirectory()  # pylint: disable=consider-using-with

    def client_send_evaluation_key_to_server(self, serialized_evaluation_keys):
        """Send the public key to the server."""
        with open(self.server_dir.name + "/serialized_evaluation_keys.ekl", "wb") as f:
            f.write(serialized_evaluation_keys)

    def client_send_input_to_server_for_prediction(self, encrypted_input):
        """Send the input to the server and execute on the server in FHE."""
        with open(self.server_dir.name + "/serialized_evaluation_keys.ekl", "rb") as f:
            serialized_evaluation_keys = f.read()
        time_begin = time.time()
        encrypted_prediction = FHEModelServer(self.server_dir.name).run(
            encrypted_input, serialized_evaluation_keys
        )
        time_end = time.time()
        with open(self.server_dir.name + "/encrypted_prediction.enc", "wb") as f:
            f.write(encrypted_prediction)
        return time_end - time_begin

    def dev_send_model_to_server(self):
        """Send the model to the server."""
        copyfile(self.dev_dir.name + "/server.zip", self.server_dir.name + "/server.zip")

    def server_send_encrypted_prediction_to_client(self):
        """Send the encrypted prediction to the client."""
        with open(self.server_dir.name + "/encrypted_prediction.enc", "rb") as f:
            encrypted_prediction = f.read()
        return encrypted_prediction

    def dev_send_clientspecs_and_modelspecs_to_client(self):
        """Send the clientspecs and evaluation key to the client."""
        copyfile(self.dev_dir.name + "/client.zip", self.client_dir.name + "/client.zip")

    def cleanup(self):
        """Clean up the temporary folders."""
        self.server_dir.cleanup()
        self.client_dir.cleanup()
        self.dev_dir.cleanup()

Development Machine

In [3]:

from torus.compiler import check_gpu_available

# Let's first get some data and train a model.
X, y = load_breast_cancer(return_X_y=True)

# Split X into X_model_owner and X_client
X_model_owner, X_client = X[:-10], X[-10:]
y_model_owner, y_client = y[:-10], y[-10:]

# Some issues on macOS, if too many estimators
n_estimators = 10
if platform.system() == "Darwin":
    n_estimators = 9


use_gpu_if_available = False
device = "cuda" if use_gpu_if_available and check_gpu_available() else "cpu"

# Train the model and compile it
model_dev = XGBClassifier(n_bits=2, n_estimators=n_estimators, max_depth=3)
model_dev.fit(X_model_owner, y_model_owner)
model_dev.compile(X_model_owner, device=device)

print("Model trained and compiled.")

Model trained and compiled.

In [4]:

# Let's instantiate the network
network = OnDiskNetwork()

# Now that the model has been trained we want to save it to send it to a server
fhemodel_dev = FHEModelDev(network.dev_dir.name, model_dev)
fhemodel_dev.save()

In [5]:

# Print all files in the temporary directory along with their sizes in KB
!ls -lh $network.dev_dir.name

total 64K
-rw-r--r-- 1 root root  59K Sep  5 14:28 client.zip
-rw-r--r-- 1 root root 2.6K Sep  5 14:28 server.zip

In the folder, we can see 3 different files:

• client.zip - contains cryptographic parameters to be sent to the client for them to create the keys (can easily be served via HTTP request). also contains the description of the pre-processing and post-processing that must be applied before encryption and after decryption.
• server.zip - contains everything required to do homomorphic prediction.

In [6]:

# Let's send the model to the server
network.dev_send_model_to_server()
!ls -lh $network.server_dir.name

total 4.0K
-rw-r--r-- 1 root root 2.6K Sep  5 14:28 server.zip

In [7]:

# Let's send the clientspecs and evaluation key to the client
network.dev_send_clientspecs_and_modelspecs_to_client()
!ls -lh $network.client_dir.name

total 60K
-rw-r--r-- 1 root root 59K Sep  5 14:28 client.zip

Client Machine

In [8]:

# Let's create the client and load the model
fhemodel_client = FHEModelClient(network.client_dir.name, key_dir=network.client_dir.name)

# The client first need to create the private and evaluation keys.
serialized_evaluation_keys = fhemodel_client.get_serialized_evaluation_keys()

KeySetCache: miss, regenerating /tmp/tmpsf57p3ed/3497183917896914639

In [9]:

# Evaluation keys can be quite large files but only have to be shared once with the server.

# Check the size of the evaluation keys (in MB)
print(f"Evaluation keys size: {len(serialized_evaluation_keys) / (10**6):.2f} MB")

# Let's send this evaluation key to the server (this has to be done only once)
network.client_send_evaluation_key_to_server(serialized_evaluation_keys)

Evaluation keys size: 29.14 MB

In [10]:

# Now we have everything for the client to interact with the server

# We create a loop to send the input to the server and receive the encrypted prediction
decrypted_predictions = []
execution_time = []
for i in range(X_client.shape[0]):
    clear_input = X_client[[i], :]
    encrypted_input = fhemodel_client.quantize_encrypt_serialize(clear_input)
    execution_time += [network.client_send_input_to_server_for_prediction(encrypted_input)]
    encrypted_prediction = network.server_send_encrypted_prediction_to_client()
    decrypted_prediction = fhemodel_client.deserialize_decrypt_dequantize(encrypted_prediction)[0]
    decrypted_predictions.append(decrypted_prediction)

# Check MB size with sys of the encrypted data vs clear data
print(
    f"Encrypted data is "
    f"{len(encrypted_input)/clear_input.nbytes:.2f}"
    " times larger than the clear data"
)

# Show execution time
print(f"The average execution time is {numpy.mean(execution_time):.2f} seconds per sample.")

Encrypted data is 4.10 times larger than the clear data
The average execution time is 0.62 seconds per sample.

In [11]:

# Let's check the results and compare them against the clear model
clear_prediction_classes = model_dev.predict_proba(X_client).argmax(axis=1)
decrypted_predictions_classes = numpy.array(decrypted_predictions).argmax(axis=1)
accuracy = (clear_prediction_classes == decrypted_predictions_classes).mean()
print(f"Accuracy between FHE prediction and clear model is: {accuracy*100:.0f}%")

Accuracy between FHE prediction and clear model is: 100%