def run_party(cid, func, rank, world_size, master_addr, master_port, func_args,
func_kwargs):
"""Start crypten party localy and run computation.
Args:
cid (int): CrypTen computation id.
func (function): computation to be done.
rank (int): rank of the crypten party.
world_size (int): number of crypten parties involved in the computation.
master_addr (str): IP address of the master party (party with rank 0).
master_port (int or str): port of the master party (party with rank 0).
func_args (list): arguments to be passed to func.
func_kwargs (dict): keyword arguments to be passed to func.
Returns:
The return value of func.
"""
process, queue = _new_party(cid, func, rank, world_size, master_addr,
master_port, func_args, func_kwargs)
was_initialized = DistributedCommunicator.is_initialized()
if was_initialized:
crypten.uninit()
process.start()
# wait for response
res = queue.get()
if was_initialized:
crypten.init()
return res
def test_is_initialized(self):
"""Tests that the is_initialized flag is set properly"""
comm = crypten.communicator
self.assertTrue(crypten.is_initialized())
self.assertTrue(comm.is_initialized())
crypten.uninit()
self.assertFalse(crypten.is_initialized())
self.assertFalse(comm.is_initialized())
crypten.init()
self.assertTrue(crypten.is_initialized())
self.assertTrue(comm.is_initialized())
def _launch(func, rank, world_size, rendezvous_file, queue, func_args,
func_kwargs):
communicator_args = {
"WORLD_SIZE": world_size,
"RANK": rank,
"RENDEZVOUS": "file://%s" % rendezvous_file,
"DISTRIBUTED_BACKEND": "gloo",
}
for key, val in communicator_args.items():
os.environ[key] = str(val)
crypten.init()
return_value = func(*func_args, **func_kwargs)
crypten.uninit()
queue.put((rank, return_value))
def test_name(self):
# Test default name is correct
self.assertEqual(comm.get().get_name(), f"rank{comm.get().get_rank()}")
# Test name set / get
comm.get().set_name(f"{comm.get().get_rank()}")
self.assertEqual(comm.get().get_name(), f"{comm.get().get_rank()}")
# Test initialization using crypten.init()
name = f"init_{comm.get().get_rank()}"
crypten.uninit()
crypten.init(party_name=name)
self.assertEqual(comm.get().get_name(), f"init_{comm.get().get_rank()}")
# Test failure on bad input
for improper_input in [0, None, ["name"], ("name",)]:
with self.assertRaises(AssertionError):
comm.get().set_name(improper_input)
def test_is_initialized(self):
"""Tests that the is_initialized flag is set properly"""
comm = crypten.communicator
self.assertTrue(crypten.is_initialized())
self.assertTrue(comm.is_initialized())
crypten.uninit()
self.assertFalse(crypten.is_initialized())
self.assertFalse(comm.is_initialized())
# note that uninit() kills the TTP process, so we need to restart it:
if self.rank == self.MAIN_PROCESS_RANK and crypten.mpc.ttp_required():
self.processes += [self._spawn_ttp()]
crypten.init()
self.assertTrue(crypten.is_initialized())
self.assertTrue(comm.is_initialized())
def wrapper(*args, **kwargs):
rendezvous_file = tempfile.NamedTemporaryFile(delete=True).name
queue = multiprocessing.Queue()
processes = [
multiprocessing.Process(
target=_launch,
args=(func, rank, world_size, rendezvous_file, queue, args, kwargs),
)
for rank in range(world_size)
]
# This process will be forked and we need to re-initialize the
# communicator in the children. If the parent process happened to
# call crypten.init(), which might be valid in a Jupyter notebook
# for instance, then the crypten.init() call on the children
# process will not do anything. The call to uninit here makes sure
# we actually get to initialize the communicator on the child
# process. An alternative fix for this issue would be to use spawn
# instead of fork, but we run into issues serializing the function
# in that case.
was_initialized = DistributedCommunicator.is_initialized()
if was_initialized:
crypten.uninit()
for process in processes:
process.start()
for process in processes:
process.join()
if was_initialized:
crypten.init()
successful = [process.exitcode == 0 for process in processes]
if not all(successful):
logging.error("One of the parties failed. Check past logs")
return None
return_values = []
while not queue.empty():
return_values.append(queue.get())
return [value for _, value in sorted(return_values, key=itemgetter(0))]
def _launch(func, rank, world_size, master_addr, master_port, queue, func_args,
func_kwargs):
communicator_args = {
"RANK": rank,
"WORLD_SIZE": world_size,
"RENDEZVOUS": "env://",
"MASTER_ADDR": master_addr,
"MASTER_PORT": master_port,
"DISTRIBUTED_BACKEND": "gloo",
}
for key, val in communicator_args.items():
os.environ[key] = str(val)
crypten.init()
return_value = func(*func_args, **func_kwargs)
crypten.uninit()
return_value = utils.pack_values(return_value)
queue.put(return_value)
def wrapper(*args, **kwargs):
# TODO:
# - check if workers are reachable / they can handle the computation
# - check return code of processes for possible failure
if len(workers) != len(set(worker.id
for worker in workers)): # noqa: C401
raise RuntimeError(
"found workers with same ID but IDs must be unique")
if model is not None:
if not isinstance(model, th.nn.Module):
raise TypeError("model must be a torch.nn.Module")
if dummy_input is None:
raise ValueError(
"must provide dummy_input when model is set")
if not isinstance(dummy_input, th.Tensor):
raise TypeError("dummy_input must be a torch.Tensor")
onnx_model = utils.pytorch_to_onnx(model, dummy_input)
else:
onnx_model = None
crypten_model = None if onnx_model is None else utils.onnx_to_crypten(
onnx_model)
world_size = len(workers)
manager = multiprocessing.Manager()
return_values = manager.dict(
{rank: None
for rank in range(world_size)})
rank_to_worker_id = dict(
zip(range(0, len(workers)), [worker.id for worker in workers]))
# TODO: run ttp in a specified worker
# if crypten.mpc.ttp_required():
# ttp_process, _ = _new_party(
# crypten.mpc.provider.TTPServer,
# world_size,
# world_size,
# master_addr,
# master_port,
# (),
# {},
# )
# ttp_process.start()
if isinstance(func, sy.Plan):
plan = func
# This is needed because at building we use a set of methods defined in syft
# (ex: load)
hook_plan_building()
was_initialized = DistributedCommunicator.is_initialized()
if not was_initialized:
crypten.init()
# We can build the plan only using a crypten model such that the actions
# traced inside the plan would know about it's existance
if crypten_model is None:
plan.build()
else:
plan.build(crypten_model)
if not was_initialized:
crypten.uninit()
unhook_plan_building()
# Mark the plan so the other workers will use that tag to retrieve the plan
plan.tags = ["crypten_plan"]
for worker in workers:
plan.send(worker)
msg = CryptenInitPlan(
(rank_to_worker_id, world_size, master_addr, master_port),
onnx_model)
else: # func
jail_runner = jail.JailRunner(func=func)
ser_jail_runner = jail.JailRunner.simplify(jail_runner)
msg = CryptenInitJail(
(rank_to_worker_id, world_size, master_addr, master_port),
ser_jail_runner,
onnx_model,
)
# Send messages to other workers so they start their parties
threads = []
for i in range(len(workers)):
rank = i
thread = multiprocessing.Process(
target=_send_party_info,
args=(workers[i], rank, msg, return_values, crypten_model),
)
thread.start()
threads.append(thread)
# wait for workers running the parties return a response
for thread in threads:
thread.join()
return return_values
def wrapper(*args, **kwargs):
# TODO:
# - check if workers are reachable / they can handle the computation
# - check return code of processes for possible failure
if model is not None:
if not isinstance(model, th.nn.Module):
raise TypeError("model must be a torch.nn.Module")
if dummy_input is None:
raise ValueError(
"must provide dummy_input when model is set")
if not isinstance(dummy_input, th.Tensor):
raise TypeError("dummy_input must be a torch.Tensor")
onnx_model = utils.pytorch_to_onnx(model, dummy_input)
else:
onnx_model = None
crypten_model = None if onnx_model is None else utils.onnx_to_crypten(
onnx_model)
world_size = len(workers) + 1
return_values = {rank: None for rank in range(world_size)}
if isinstance(func, sy.Plan):
using_plan = True
plan = func
# This is needed because at building we use a set of methods defined in syft (ex: load)
hook_plan_building()
crypten.init()
plan.build()
crypten.uninit()
unhook_plan_building()
# Mark the plan so the other workers will use that tag to retrieve the plan
plan.tags = ["crypten_plan"]
for worker in workers:
plan.send(worker)
jail_or_plan = plan
else: # func
using_plan = False
jail_runner = jail.JailRunner(func=func, model=crypten_model)
ser_jail_runner = jail.JailRunner.simplify(jail_runner)
jail_or_plan = jail_runner
rank_to_worker_id = dict(
zip(range(1,
len(workers) + 1),
[worker.id for worker in workers]))
sy.local_worker._set_rank_to_worker_id(rank_to_worker_id)
# Start local party
process, queue = _new_party(jail_or_plan, 0, world_size,
master_addr, master_port, (), {})
was_initialized = DistributedCommunicator.is_initialized()
if was_initialized:
crypten.uninit()
process.start()
# Run TTP if required
# TODO: run ttp in a specified worker
if crypten.mpc.ttp_required():
ttp_process, _ = _new_party(
crypten.mpc.provider.TTPServer,
world_size,
world_size,
master_addr,
master_port,
(),
{},
)
ttp_process.start()
# Send messages to other workers so they start their parties
threads = []
for i in range(len(workers)):
rank = i + 1
if using_plan:
msg = CryptenInitPlan((rank_to_worker_id, world_size,
master_addr, master_port))
else: # jail
msg = CryptenInitJail(
(rank_to_worker_id, world_size, master_addr,
master_port),
ser_jail_runner,
onnx_model,
)
thread = threading.Thread(target=_send_party_info,
args=(workers[i], rank, msg,
return_values))
thread.start()
threads.append(thread)
# Wait for local party and sender threads
# Joining the process blocks! But queue.get() can also wait for the party
# and it works fine.
# process.join() -> blocks
local_party_result = queue.get()
return_values[0] = utils.unpack_values(local_party_result,
crypten_model)
for thread in threads:
thread.join()
if was_initialized:
crypten.init()
return return_values
def test_encrypted_model(model):
crypten.init()
model.encrypt()
with pytest.raises(RuntimeError):
model._check_encrypted()
crypten.uninit()
