def test_reconstruct_shared_model(
is_remote: bool, model_type: Type[sy.Module], get_clients: Callable[[int], Any]
):
net = model_type(torch)
clients = get_clients(2)
session = Session(parties=clients)
SessionManager.setup_mpc(session)
if is_remote:
model = net.send(clients[0])
else:
model = net
mpc_model = model.share(session=session)
res = mpc_model.reconstruct()
assert isinstance(res, sy.Module)
if is_remote:
# If the model is remote fetch it such that we could compare it
model = model.get()
for name_res, name_expected in zip(res.modules, model.modules):
assert name_res == name_expected
module_expected = model.modules[name_expected]
module_res = res.modules[name_res]
name_module = type(module_expected).__name__
assert MAP_TORCH_TO_SYMPC[name_module].eq_close(
module_expected, module_res, atol=1e-4
)
def run_conv_model(get_clients: Callable[[int], List[Any]]):
"""Create a convolutional network and do a feedforwad.
Arguments:
get_clients: Fixture that returns a list of clients
"""
model = ConvNet(torch)
clients = get_clients(2)
session = Session(parties=clients)
SessionManager.setup_mpc(session)
mpc_model = model.share(session=session)
x_secret = torch.randn((1, 1, 28, 28))
x_mpc = MPCTensor(secret=x_secret, session=session)
model.eval()
# For the moment we have only inference
expected = model(x_secret)
res_mpc = mpc_model(x_mpc)
res_mpc.reconstruct()
expected = expected.detach().numpy()
def test_rst_invalid_triple(get_clients) -> None:
parties = get_clients(3)
falcon = Falcon("malicious")
session = Session(parties, protocol=falcon)
SessionManager.setup_mpc(session)
shape_x = (1, )
shape_y = (1, )
# create an inconsistent sharing,invoke a prrs first
session.session_ptrs[0].prrs_generate_random_share(shape_x)
with pytest.raises(ValueError):
CryptoPrimitiveProvider.generate_primitives(
"beaver_mul",
session=session,
g_kwargs={
"session": session,
"a_shape": shape_x,
"b_shape": shape_y,
"nr_parties": session.nr_parties,
},
p_kwargs={
"a_shape": shape_x,
"b_shape": shape_y
},
)
def test_get_grad_input_padding(get_clients, common_args: List,
nr_parties) -> None:
clients = get_clients(2)
session = Session(parties=clients)
SessionManager.setup_mpc(session)
grad = torch.Tensor([[[[1.0, 1.0, 1.0], [1.0, 1.0, 1.0], [1.0, 1.0,
1.0]]]])
grad_mpc = MPCTensor(secret=grad, session=session)
input_size, stride, padding, kernel_size, dilation = common_args
expected_padding = torch.nn.functional.grad._grad_input_padding(
grad,
input_size,
(stride, stride),
(padding, padding),
kernel_size,
(dilation, dilation),
)
args = [[el] + common_args + [session] for el in grad_mpc.share_ptrs]
shares = parallel_execution(GradConv2d.get_grad_input_padding,
grad_mpc.session.parties)(args)
grad_input_padding = MPCTensor(shares=shares, session=grad_mpc.session)
output_padding_tensor = grad_input_padding.reconstruct()
output_padding_tensor /= grad_mpc.session.nr_parties
calculated_padding = tuple(output_padding_tensor.to(torch.int).tolist())
assert calculated_padding == expected_padding
def test_generate_and_transfer_primitive(
get_clients: Callable,
nr_parties: int,
nr_instances: int,
nr_instances_retrieve: int,
) -> None:
parties = get_clients(nr_parties)
session = Session(parties=parties)
SessionManager.setup_mpc(session)
g_kwargs = {"nr_parties": nr_parties, "nr_instances": nr_instances}
CryptoPrimitiveProvider.generate_primitives(
"test",
session=session,
g_kwargs=g_kwargs,
p_kwargs={},
)
for i in range(nr_parties):
remote_crypto_store = session.session_ptrs[i].crypto_store
primitives = remote_crypto_store.get_primitives_from_store(
op_str="test", nr_instances=nr_instances_retrieve).get()
assert primitives == [
tuple(ShareTensor(i) for _ in range(PRIMITIVE_NR_ELEMS))
for _ in range(nr_instances_retrieve)
]
def test_setup_przs(get_clients):
"""Test _setup_przs method for session."""
alice_client, bob_client = get_clients(2)
session = Session(parties=[alice_client, bob_client])
assert len(session.przs_generators) == 0
SessionManager._setup_przs(session)
assert len(session.przs_generators) == 2
def test_additional_attributes(get_clients):
class ConvNet_attr(sy.Module):
def __init__(self, torch_ref):
super(ConvNet_attr, self).__init__(torch_ref=torch_ref)
self.conv1 = self.torch_ref.nn.Conv2d(
in_channels=2,
out_channels=6,
kernel_size=(3, 5),
stride=(2, 1),
padding=(2, 1),
groups=1,
dilation=2,
)
def forward(self, x):
x = self.conv1(x)
return x
model = ConvNet_attr(torch)
clients = get_clients(2)
session = Session(parties=clients)
SessionManager.setup_mpc(session)
mpc_model = model.share(session=session)
model = mpc_model.reconstruct()
assert model.conv1.kernel_size == (3, 5)
assert model.conv1.padding == (2, 1)
assert model.conv1.stride == (2, 1)
assert model.conv1.groups == 1
assert model.conv1.dilation == (2, 2)
def test_setup_mpc(get_clients):
"""Test setup_mpc method for session."""
alice_client, bob_client = get_clients(2)
session = Session(parties=[alice_client, bob_client])
SessionManager.setup_mpc(session)
assert session.rank == 1
assert len(session.session_ptrs) == 2
def protobuf_session_deserializer(proto: MPCSession_PB) -> Session:
id_session = UUID(bytes=proto.uuid)
rank = proto.rank
conf_dict = Dict._proto2object(proto=proto.config)
_conf_dict = {key: value for key, value in conf_dict.items()}
conf = Config(**_conf_dict)
ring_size = int.from_bytes(proto.ring_size, "big")
nr_parties = int.from_bytes(proto.nr_parties, "big")
protocol_deserialized = proto.protocol.decode()
session = Session(
config=conf,
uuid=id_session,
ring_size=ring_size,
protocol=protocol_deserialized,
)
session.rank = rank
session.crypto_store = CryptoStore()
session.nr_parties = nr_parties
if "session" in globals():
warning("Overwritting session for MPC")
globals()["session"] = session
return session
def __init__(
self,
data: Optional[Union[float, int, torch.Tensor]] = None,
session: Optional[Session] = None,
encoder_base: int = 2,
encoder_precision: int = 16,
ring_size: int = 2**64,
) -> None:
"""Initializer for the ShareTensor."""
if session is None:
self.session = Session(ring_size=ring_size, )
self.session.config.encoder_precision = encoder_precision
self.session.config.encoder_base = encoder_base
else:
self.session = session
encoder_precision = self.session.config.encoder_precision
encoder_base = self.session.config.encoder_base
# TODO: It looks like the same logic as above
self.fp_encoder = FixedPointEncoder(base=encoder_base,
precision=encoder_precision)
self.tensor: Optional[torch.Tensor] = None
if data is not None:
tensor_type = self.session.tensor_type
self.tensor = self._encode(data).type(tensor_type)
def share(_self, **kwargs: Dict[Any, Any]) -> MPCTensor: # noqa
session = None
if "parties" not in kwargs and "session" not in kwargs:
raise ValueError("Parties or Session should be provided as a kwarg")
if "session" not in kwargs:
parties = frozenset({client.id for client in kwargs["parties"]})
if parties not in PARTIES_TO_SESSION:
from sympc.session import SessionManager
session = Session(kwargs["parties"])
PARTIES_TO_SESSION[parties] = session
SessionManager.setup_mpc(session)
for key, val in kwargs.items():
setattr(session, key, val)
else:
session = PARTIES_TO_SESSION[parties]
kwargs.pop("parties")
kwargs["session"] = session
return MPCTensor(secret=_self, **kwargs)
def set_up_model(
get_clients: Callable[[int], List[Any]]
) -> Tuple[sy.Module, sy.Module, MPCTensor]:
"""Create a convolutional network and do a feedforwad.
Args:
get_clients: Fixture that returns a list of clients
Returns:
A tuple with three items: (origin Model, Syft Model, MPCTensor)
"""
model = ConvNet(torch)
clients = get_clients(2)
session = Session(parties=clients)
SessionManager.setup_mpc(session)
mpc_model = model.share(session=session)
x_secret = torch.randn((1, 1, 28, 28))
x_mpc = MPCTensor(secret=x_secret, session=session)
model.eval()
expected = model(x_secret)
return (expected, mpc_model, x_mpc)
def test_generate_primitive_from_dict_beaver_mul(get_clients) -> None:
clients = get_clients(2)
session = Session(parties=clients)
SessionManager.setup_mpc(session)
primitive_log = {
"beaver_mul": [
(
{
"a_shape": (2, 10),
"b_shape": (2, 10)
},
{
"session": session,
"a_shape": (2, 10),
"b_shape": (2, 10),
"nr_parties": 2,
},
),
(
{
"a_shape": (1, 5),
"b_shape": (1, 5)
},
{
"session": session,
"a_shape": (1, 5),
"b_shape": (1, 5),
"nr_parties": 2,
},
),
]
}
CryptoPrimitiveProvider.generate_primitive_from_dict(
primitive_log=primitive_log, session=session)
args = primitive_log.popitem()[1]
for arg in args:
a_shape = arg[1].get("a_shape")
b_shape = arg[1].get("b_shape")
key = f"beaver_mul_{a_shape}_{b_shape}"
store_client_1 = session.session_ptrs[0].crypto_store.store.get()
store_client_2 = session.session_ptrs[1].crypto_store.store.get()
a_shape_client_1 = tuple(store_client_1.get(key)[0][0].shape)
b_shape_client_1 = tuple(store_client_1.get(key)[0][1].shape)
assert a_shape == a_shape_client_1
assert b_shape == b_shape_client_1
a_shape_client_2 = tuple(store_client_2.get(key)[0][0].shape)
b_shape_client_2 = tuple(store_client_2.get(key)[0][1].shape)
assert a_shape == a_shape_client_2
assert b_shape == b_shape_client_2
def test_invalid_mpc_pointer(get_clients) -> None:
parties = get_clients(3)
session = Session(parties=parties)
SessionManager.setup_mpc(session)
x = MPCTensor(secret=1, session=session)
# passing sharetensor pointer
with pytest.raises(ValueError):
ABY3.truncate(x, session, 2**32, None)
def test_select_shares_exception_ring(get_clients) -> None:
parties = get_clients(3)
falcon = Falcon()
session = Session(parties=parties, protocol=falcon, ring_size=2**32)
SessionManager.setup_mpc(session)
val = MPCTensor(secret=1, session=session)
with pytest.raises(ValueError):
Falcon.select_shares(val, val, val)
def test_setup_mpc(get_clients):
"""Test setup_mpc method for session."""
alice_client, bob_client = get_clients(2)
session = Session(parties=[alice_client, bob_client])
SessionManager.setup_mpc(session)
assert isinstance(session.uuid, UUID)
assert list(session.rank_to_uuid.keys()) == [0, 1]
assert all(isinstance(e, UUID) for e in session.rank_to_uuid.values())
def test_copy() -> None:
session = Session()
copy_session = session.copy()
assert session.nr_parties == copy_session.nr_parties
assert session.config == copy_session.config
assert session.protocol == copy_session.protocol
def test_remote_mpc_no_shape(get_clients) -> None:
alice_client, bob_client = get_clients(2)
session = Session(parties=[alice_client, bob_client])
SessionManager.setup_mpc(session)
x_remote = alice_client.torch.Tensor([1, -2, 0.3])
with pytest.raises(ValueError):
MPCTensor(secret=x_remote, session=session)
def test_setupmpc_nocall_exception(get_clients) -> None:
alice_client, bob_client = get_clients(2)
session = Session(parties=[alice_client, bob_client])
with pytest.raises(ValueError):
MPCTensor(secret=42, session=session)
with pytest.raises(ValueError):
MPCTensor(secret=torch.Tensor([1, -2]), session=session)
def test_invalid_share_class(get_clients) -> None:
clients = get_clients(2)
session = Session(parties=clients)
SessionManager.setup_mpc(session)
x = torch.tensor([1, 2, 3])
x_s = MPCTensor(secret=x, session=session)
x_s.session.protocol.share_class = "invalid"
with pytest.raises(TypeError):
x_s + x
def test_mpc_len(get_clients) -> None:
clients = get_clients(2)
session = Session(parties=clients)
SessionManager.setup_mpc(session)
for dim in range(1, 5):
x_secret = torch.arange(-6, 6).view(dim, -1)
x = MPCTensor(secret=x_secret, session=session)
assert len(x_secret) == len(x)
def test_eq() -> None:
"""Test __eq__ for Session."""
session = Session()
other1 = Session()
other2 = session
# Test different instances:
assert session != 1
# Test equal sessions:
assert session == other2
# Test same sessions (until we call setup mpc):
assert session == other1
SessionManager.setup_mpc(session)
assert session != other1
def test_max_multiple_max(get_clients) -> None:
clients = get_clients(2)
session = Session(parties=clients)
SessionManager.setup_mpc(session)
x = MPCTensor(secret=torch.Tensor([1, 2, 3, -1, 3]), session=session)
with pytest.raises(ValueError):
x.argmax()
def test_session_manager_init():
"""Test correct initialisation of the SessionManager class."""
# Test default init
session = SessionManager()
assert isinstance(session.uuid, UUID)
# Test custom init
uuid = uuid4()
session = Session(uuid=uuid)
assert session.uuid == uuid
def test_bit_injection_exception(get_clients) -> None:
parties = get_clients(3)
falcon = Falcon()
session = Session(parties=parties, protocol=falcon)
SessionManager.setup_mpc(session)
x = MPCTensor(secret=1, session=session)
with pytest.raises(ValueError):
ABY3.bit_injection(x, session, 2**64)
def test_remote_mpc_with_shape(get_clients) -> None:
alice_client, bob_client = get_clients(2)
session = Session(parties=[alice_client, bob_client])
SessionManager.setup_mpc(session)
x_remote = alice_client.torch.Tensor([1, -2, 0.3])
x = MPCTensor(secret=x_remote, shape=(1, 3), session=session)
result = x.reconstruct()
assert np.allclose(x_remote.get(), result, atol=1e-5)
def test_rst_distribute_reconstruct_float_secret(get_clients, parties,
security) -> None:
parties = get_clients(parties)
protocol = Falcon(security)
session = Session(protocol=protocol, parties=parties)
SessionManager.setup_mpc(session)
secret = 43.2
a = MPCTensor(secret=secret, session=session)
assert np.allclose(secret, a.reconstruct(), atol=1e-3)
def test_exception_value_error(get_clients) -> None:
clients = get_clients(2)
session_one = Session(parties=clients)
SessionManager.setup_mpc(session_one)
x_secret = torch.Tensor([-2.0, 6.0, 2.0, 3.0, -5.0, -0.5])
x = MPCTensor(secret=x_secret, session=session_one)
with pytest.raises(ValueError):
reciprocal(x, method="exp")
def test_share_distribution_number_shares(get_clients, parties):
parties = get_clients(parties)
protocol = Falcon("semi-honest")
session = Session(protocol=protocol, parties=parties)
SessionManager.setup_mpc(session)
shares = MPCTensor.generate_shares(100.42, len(parties))
share_ptrs = ReplicatedSharedTensor.distribute_shares(shares, session)
for RSTensor in share_ptrs:
assert len(RSTensor.get_shares().get()) == (len(parties) - 1)
def test_exception_conv2d_kernel_mismatch(get_clients):
clients = get_clients(2)
model = ConvNet(torch, kernel_size=(5, 4))
# Setup the session for the computation
session = Session(parties=clients)
SessionManager.setup_mpc(session)
with pytest.raises(ValueError):
model.share(session=session)
