def _serialize_domain(dom, pb_dom):
for d in dom:
new_dom = pb_dom.domains.add()
if isinstance(d, Node):
new_dom.type = pb.Attribute.Type.NODE
new_dom.s = _to_bytes_or_false(d.name)
elif isinstance(d, np.ndarray):
new_dom.type = pb.Attribute.Type.NDARRAY
new_dom.nda.CopyFrom(ndarray_to_proto(d))
elif isinstance(d, Integral):
new_dom.type = pb.Attribute.Type.INT32
new_dom.i32 = d
elif isinstance(d, float):
new_dom.type = pb.Attribute.Type.DOUBLE
new_dom.d = d
elif isinstance(d, str):
new_dom.type = pb.Attribute.Type.STRING
new_dom.s = _to_bytes_or_false(d)
elif isinstance(d, bool):
new_dom.type = pb.Attribute.Type.BOOL
new_dom.b = d
elif isinstance(d, Iterable):
if all(isinstance(a, Node) for a in d):
new_dom.type = pb.Attribute.Type.NODES
new_dom.ss.extend([_to_bytes_or_false(a.name) for a in d])
elif all(isinstance(a, list) for a in d):
new_dom.type = pb.Attribute.Type.NDARRAYS
np_arr = [ndarray_to_proto(np.asarray(a)) for a in d]
new_dom.ndas.extend(np_arr)
elif all(isinstance(a, np.ndarray) for a in d):
new_dom.type = pb.Attribute.Type.NDARRAYS
new_dom.ndas.extend(ndarray_to_proto(a) for a in d)
elif all(isinstance(a, Integral) for a in d):
new_dom.type = pb.Attribute.Type.INT32S
new_dom.i32s.extend(d)
elif all(isinstance(a, float) for a in d):
new_dom.type = pb.Attribute.Type.DOUBLES
new_dom.ds.extend(d)
elif all(
map(lambda bytes_or_false: bytes_or_false is not False,
[_to_bytes_or_false(a) for a in d])):
new_dom.type = pb.Attribute.Type.STRINGS
new_dom.ss.extend([_to_bytes_or_false(a) for a in d])
elif all(isinstance(a, bool) for a in d):
new_dom.type = pb.Attribute.Type.BOOLS
new_dom.bs.extend(d)
else:
raise TypeError(
f"Cannot find serializable method for argument {d} with "
f"type {type(d)} in domain {dom.names}")
else:
raise TypeError(
f"Cannot find serializable method for domain {d} with type {type(d)}"
)
def StartTraining(self, request, context):
print(f"Received: {type(request)} from {context.peer()}")
theta_proto = [ndarray_to_proto(nda) for nda in self.theta]
# send reply
return coordinator_pb2.StartTrainingReply(theta=theta_proto,
epochs=self.epochs,
epoch_base=self.epoch_base)
def SayHelloNumProto(self, request, context):
nda = proto_to_ndarray(request.arr)
logger.info("NumProto server received", nda=nda)
nda *= 2
logger.info("NumProto server sent", nda=nda)
return hellonumproto_pb2.NumProtoReply(arr=ndarray_to_proto(nda))
def SayHelloNumProto(self, request, context):
nda = proto_to_ndarray(request.arr)
print("NumProto server received: {}".format(nda))
nda *= 2
print("NumProto server sent: {}".format(nda))
return hellonumproto_pb2.NumProtoReply(arr=ndarray_to_proto(nda))
def end_training(channel, theta_n: Tuple[Theta, int], history: History,
metrics: Metrics):
"""Starts a training completion exchange with Coordinator, sending a locally
trained model and metadata.
Args:
channel: gRPC channel to Coordinator.
theta_n (obj:`Tuple[Theta, int]`): Locally trained model.
history (obj:`History`): History metadata.
Metrics (obj:`Metrics`): Metrics metadata.
"""
# pylint: disable=no-member
stub = coordinator_pb2_grpc.CoordinatorStub(channel)
# build request starting with theta update
theta, num = theta_n
theta_n_proto = coordinator_pb2.EndTrainingRequest.ThetaUpdate(
theta_prime=[ndarray_to_proto(nda) for nda in theta], num_examples=num)
# history data
h = {
k: coordinator_pb2.EndTrainingRequest.HistoryValue(values=v)
for k, v in history.items()
}
# metrics
cid, vbc = metrics
m = coordinator_pb2.EndTrainingRequest.Metrics(cid=cid, vol_by_class=vbc)
# assemble req
req = coordinator_pb2.EndTrainingRequest(theta_update=theta_n_proto,
history=h,
metrics=m)
# send request to end training
reply = stub.EndTraining(req)
logger.info("Participant received", reply_type=type(reply))
def _handle_start_training(
self, _message: coordinator_pb2.StartTrainingRequest, participant_id: str
) -> coordinator_pb2.StartTrainingReply:
"""Handles a StartTraining request.
Args:
_message (:class:`~.coordinator_pb2.StartTrainingRequest`): The
request to handle. Currently not used.
participant_id (:obj:`str`): The id of the participant making the
request.
Returns:
:class:`~.coordinator_pb2.StartTrainingReply`: The reply to the participant.
"""
# The coordinator should only accept StartTraining requests if it is
# in the ROUND state and when the participant has been selected for the round.
coordinator_not_in_a_round = self.state != coordinator_pb2.State.ROUND
participant_not_selected = participant_id not in self.round.participant_ids
if coordinator_not_in_a_round or participant_not_selected:
raise InvalidRequestError(
f"Participant {participant_id} sent a "
"StartTrainingRequest outside of a round"
)
weights_proto = [ndarray_to_proto(nda) for nda in self.weights]
return coordinator_pb2.StartTrainingReply(
weights=weights_proto, epochs=self.epochs, epoch_base=self.epoch_base
)
def test_greeter_server(greeter_server):
with grpc.insecure_channel("localhost:50051") as channel:
stub = hellonumproto_pb2_grpc.NumProtoServerStub(channel)
nda = np.arange(10)
response = stub.SayHelloNumProto(
hellonumproto_pb2.NumProtoRequest(arr=ndarray_to_proto(nda)))
response_nda = proto_to_ndarray(response.arr)
assert np.array_equal(nda * 2, response_nda)
def run():
with grpc.insecure_channel("localhost:50051") as channel:
stub = hellonumproto_pb2_grpc.NumProtoServerStub(channel)
nda = np.arange(10)
print("NumProto client sent: {}".format(nda))
response = stub.SayHelloNumProto(
hellonumproto_pb2.NumProtoRequest(arr=ndarray_to_proto(nda))
)
print("NumProto client received: {}".format(proto_to_ndarray(response.arr)))
def end_training(
channel: Channel,
weights: List[ndarray],
number_samples: int,
metrics: Dict[str, List[ndarray]],
) -> None:
"""Start a training completion exchange with a coordinator.
The locally trained weights and the number of samples as well as metrics metadata is sent.
Args:
channel (~grpc.Channel): A gRPC channel to the coordinator.
weights (~typing.List[~numpy.ndarray]): The weights of the locally trained model.
number_samples (int): The number of samples in the training dataset.
metrics (~typing.Dict[str, ~numpy.ndarray]): Metrics metadata.
"""
coordinator: CoordinatorStub = CoordinatorStub(channel=channel)
# build request starting with weight update
weights_proto: List = [ndarray_to_proto(weight) for weight in weights]
# metric data containing the metric names mapped to Metrics as protobuf message
metrics_proto: Dict[str, EndTrainingRequest.Metrics] = {
key: EndTrainingRequest.Metrics(
metrics=[ndarray_to_proto(value) for value in values]
)
for key, values in metrics.items()
}
# assembling a request with the update of the weights and the metrics
request: EndTrainingRequest = EndTrainingRequest(
weights=weights_proto, number_samples=number_samples, metrics=metrics_proto
)
reply: EndTrainingReply = coordinator.EndTraining(request=request)
logger.info("Participant received reply", reply=type(reply))
def pointer_to_proto(pointer, pointer_len, nptype=np.uint8):
"""
Convert C u_int or float pointer to proto for RPC serialization
Parameters
----------
pointer : ctypes.POINTER
pointer_len : length of pointer
nptype : np type to cast to
if pointer is of type ctypes.c_uint, nptype should be np.uint32
if pointer is of type ctypes.c_float, nptype should be np.float32
Returns:
proto : proto.NDArray
"""
ndarray = ctypes2numpy(pointer, pointer_len, nptype)
proto = ndarray_to_proto(ndarray)
return proto
def from_polymath_arg(serialized_arg, arg):
if isinstance(arg, Node):
serialized_arg.type = pb.Attribute.Type.NODE
serialized_arg.s = _to_bytes_or_false(arg.name)
elif isinstance(arg, np.ndarray):
serialized_arg.type = pb.Attribute.Type.NDARRAY
serialized_arg.nda.CopyFrom(ndarray_to_proto(arg))
elif isinstance(arg, Integral):
serialized_arg.type = pb.Attribute.Type.INT32
serialized_arg.i32 = arg
elif isinstance(arg, float):
serialized_arg.type = pb.Attribute.Type.DOUBLE
serialized_arg.d = arg
elif isinstance(arg, str):
serialized_arg.type = pb.Attribute.Type.STRING
serialized_arg.s = _to_bytes_or_false(arg)
elif isinstance(arg, bool):
serialized_arg.type = pb.Attribute.Type.BOOL
serialized_arg.b = arg
else:
raise RuntimeError(f"Unable to find valid type for arg {arg}")
def end_training(channel, theta_update: Tuple[Theta, int], history: History,
metrics: Metrics):
stub = coordinator_pb2_grpc.CoordinatorStub(channel)
# build request starting with theta update
theta, num = theta_update
theta_update_p = coordinator_pb2.EndTrainingRequest.ThetaUpdate(
theta_prime=[ndarray_to_proto(nda) for nda in theta], num_examples=num)
# history data
h = {
k: coordinator_pb2.EndTrainingRequest.HistoryValue(values=v)
for k, v in history.items()
}
# metrics
cid, vbc = metrics
m = coordinator_pb2.EndTrainingRequest.Metrics(cid=cid, vol_by_class=vbc)
# assemble req
req = coordinator_pb2.EndTrainingRequest(theta_update=theta_update_p,
history=h,
metrics=m)
# send request to end training
reply = stub.EndTraining(req)
print(f"Participant received: {type(reply)}")
def test_numproto(nda):
"""Tests serialization and deserialization of numpy arrays."""
result = proto_to_ndarray(ndarray_to_proto(nda))
assert np.array_equal(nda, result)
def _add_client_key():
"""
Add private (symmetric) key to enclave.
This function encrypts the user's symmetric key using the enclave's public key,
and signs the ciphertext with the user's private key.
The signed message is sent to the enclave.
"""
# Convert key to serialized numpy array
enclave_public_key_size = _CONF["enclave_pk_size"].value
enclave_public_key = ctypes2numpy(_CONF["enclave_pk"],
enclave_public_key_size, np.uint8)
enclave_public_key = ndarray_to_proto(enclave_public_key)
# Convert nonce to serialized numpy array
nonce_size = _CONF["nonce_size"].value
nonce = ctypes2numpy(_CONF["nonce"], nonce_size, np.uint8)
nonce = ndarray_to_proto(nonce)
if _CONF.get("general_config") is None:
raise MC2ClientConfigError("Configuration not set")
user_config = yaml.safe_load(open(_CONF["general_config"]).read())["user"]
symm_key_path = user_config["symmetric_key"]
if not os.path.exists(symm_key_path):
raise FileNotFoundError(
"Symmetric key not found at {}".format(symm_key_path))
else:
_CONF["symm_key"] = symm_key_path
priv_key_path = user_config["private_key"]
if not os.path.exists(priv_key_path):
raise FileNotFoundError(
"Private key not found at {}".format(priv_key_path))
else:
_CONF["private_key"] = priv_key_path
cert_path = user_config["certificate"]
if not os.path.exists(cert_path):
raise FileNotFoundError(
"Certificate not found at {}".format(cert_path))
with open(symm_key_path, "rb") as symm_keyfile:
user_symm_key = symm_keyfile.read()
with open(cert_path, "rb") as cert_file:
cert = cert_file.read()
enc_sym_key, enc_sym_key_size = encrypt_data_with_pk(
user_symm_key, len(user_symm_key), enclave_public_key,
enclave_public_key_size)
# Sign the encrypted symmetric key
sig, sig_size = sign_data(priv_key_path, enc_sym_key, enc_sym_key_size)
# Send the encrypted key to the enclave
channel_addr = _CONF["remote_addr"]
if channel_addr is None:
raise MC2ClientConfigError(
"Remote orchestrator IP not set. Run oc.create_cluster() \
to launch VMs and configure IPs automatically or explicitly set it in the user YAML."
)
if channel_addr:
with grpc.insecure_channel(channel_addr) as channel:
stub = remote_pb2_grpc.RemoteStub(channel)
stub.rpc_add_client_key_with_certificate(
remote_pb2.DataMetadata(
certificate=cert,
enc_sym_key=enc_sym_key,
key_size=enc_sym_key_size,
signature=sig,
sig_len=sig_size,
))
def _serialize_node(node_instance):
graph_id = -1 if node_instance.graph is None else id(node_instance.graph)
pb_node = pb.Node(
name=node_instance.name,
op_name=node_instance.op_name,
uuid=id(node_instance),
module=
f"{node_instance.__class__.__module__}.{node_instance.__class__.__name__}",
graph_id=graph_id)
for shape in node_instance.shape:
pb_shape = pb_node.shape.add()
if isinstance(shape, Node):
pb_shape.shape_id = shape.name
elif not isinstance(shape, Integral):
raise TypeError(f"Invalid type for shape {shape} - {type(shape)}")
else:
pb_shape.shape_const = shape
pb_node.dependencies.extend(node_instance.dependencies)
for arg in node_instance.args:
new_arg = pb_node.args.add()
if isinstance(arg, Node):
new_arg.type = pb.Attribute.Type.NODE
new_arg.s = _to_bytes_or_false(arg.name)
elif isinstance(arg, np.ndarray):
new_arg.type = pb.Attribute.Type.NDARRAY
new_arg.nda.CopyFrom(ndarray_to_proto(arg))
elif isinstance(arg, Integral):
new_arg.type = pb.Attribute.Type.INT32
new_arg.i32 = arg
elif isinstance(arg, float):
new_arg.type = pb.Attribute.Type.DOUBLE
new_arg.d = arg
elif isinstance(arg, str):
new_arg.type = pb.Attribute.Type.STRING
new_arg.s = _to_bytes_or_false(arg)
elif isinstance(arg, bool):
new_arg.type = pb.Attribute.Type.BOOL
new_arg.b = arg
elif isinstance(arg, Mapping):
new_arg.type = pb.Attribute.Type.MAP
for name, value in arg.items():
mapped_arg = new_arg.mapping[name]
from_polymath_arg(mapped_arg, value)
elif isinstance(arg, Iterable):
# TODO: Fix this to be more generic
if isinstance(arg, list) and len(arg) == 1 and isinstance(
arg[0], tuple):
arg = arg[0]
if all(isinstance(a, Node) for a in arg):
new_arg.type = pb.Attribute.Type.NODES
new_arg.ss.extend([_to_bytes_or_false(a.name) for a in arg])
elif all(isinstance(a, list) for a in arg):
new_arg.type = pb.Attribute.Type.NDARRAYS
np_arr = [ndarray_to_proto(np.asarray(a)) for a in arg]
new_arg.ndas.extend(np_arr)
elif all(isinstance(a, np.ndarray) for a in arg):
new_arg.type = pb.Attribute.Type.NDARRAYS
new_arg.ndas.extend(ndarray_to_proto(a) for a in arg)
elif all(isinstance(a, Integral) for a in arg):
new_arg.type = pb.Attribute.Type.INT32S
new_arg.i32s.extend(arg)
elif all(isinstance(a, float) for a in arg):
new_arg.type = pb.Attribute.Type.DOUBLES
new_arg.ds.extend(arg)
elif all(
map(lambda bytes_or_false: bytes_or_false is not False,
[_to_bytes_or_false(a) for a in arg])):
new_arg.type = pb.Attribute.Type.STRINGS
new_arg.ss.extend([_to_bytes_or_false(a) for a in arg])
elif all(isinstance(a, bool) for a in arg):
new_arg.type = pb.Attribute.Type.BOOLS
new_arg.bs.extend(arg)
else:
raise TypeError(
f"Cannot find serializable method for argument {arg} with "
f"type {type(arg)} in node {node_instance.name} - {node_instance.op_name}\n"
f"All args: {node_instance.args}")
else:
raise TypeError(
f"Cannot find serializable method for argument {arg} with type {type(arg)}"
)
for name, arg in node_instance.kwargs.items():
if arg is None:
continue
new_arg = pb_node.kwargs[name]
if isinstance(arg, Domain):
_serialize_domain(arg, new_arg.dom)
new_arg.type = pb.Attribute.Type.DOM
elif isinstance(arg, Node):
new_arg.type = pb.Attribute.Type.NODE
new_arg.s = _to_bytes_or_false(arg.name)
elif isinstance(arg, np.ndarray):
new_arg.type = pb.Attribute.Type.NDARRAY
new_arg.nda.CopyFrom(ndarray_to_proto(arg))
elif isinstance(arg, Integral):
new_arg.type = pb.Attribute.Type.INT32
new_arg.i32 = arg
elif isinstance(arg, float):
new_arg.type = pb.Attribute.Type.DOUBLE
new_arg.d = arg
elif isinstance(arg, str):
new_arg.type = pb.Attribute.Type.STRING
new_arg.s = _to_bytes_or_false(arg)
elif isinstance(arg, bool):
new_arg.type = pb.Attribute.Type.BOOL
new_arg.b = _to_bytes_or_false(arg)
elif isinstance(arg, Iterable):
if all(isinstance(a, Node) for a in arg):
new_arg.type = pb.Attribute.Type.NODES
new_arg.ss.extend([_to_bytes_or_false(a.name) for a in arg])
elif all(isinstance(a, np.ndarray) for a in arg):
new_arg.type = pb.Attribute.Type.NDARRAYS
new_arg.ndas.extend(ndarray_to_proto(a) for a in arg)
elif all(isinstance(a, Integral) for a in arg):
new_arg.type = pb.Attribute.Type.INT32S
new_arg.i32s.extend(arg)
elif all(isinstance(a, float) for a in arg):
new_arg.type = pb.Attribute.Type.DOUBLES
new_arg.ds.extend(arg)
elif all(
map(lambda bytes_or_false: bytes_or_false is not False,
[_to_bytes_or_false(a) for a in arg])):
new_arg.type = pb.Attribute.Type.STRINGS
new_arg.ss.extend([_to_bytes_or_false(a) for a in arg])
elif all(isinstance(a, bool) for a in arg):
new_arg.type = pb.Attribute.Type.BOOLS
new_arg.bs.extend(arg)
else:
raise TypeError(
f"Cannot find serializable method for argument {name}={arg} with type {type(arg)} in {node_instance}"
)
serialized = []
for k, node in node_instance.nodes.items():
if not isinstance(node, Node):
raise RuntimeError(
f"Non-node object included in graph for {node_instance.name} with name {k}."
)
elif node.name != node_instance.name:
serialized.append(_serialize_node(node))
pb_node.nodes.extend(serialized)
return pb_node
