def _restore_placeholders(action: ComputationAction):
"""Converts PlaceholderId's to PlaceHolder in an Action"""
def wrap_in_placeholder(ph_id):
return PlaceHolder(id=ph_id)
action.target = Role.nested_object_traversal(action.target,
wrap_in_placeholder,
PlaceholderId)
action.args = Role.nested_object_traversal(action.args,
wrap_in_placeholder,
PlaceholderId)
action.kwargs = Role.nested_object_traversal(action.kwargs,
wrap_in_placeholder,
PlaceholderId)
action.return_ids = Role.nested_object_traversal(
action.return_ids, wrap_in_placeholder, PlaceholderId)
return action
def build(self, *args):
"""Builds the plan.
First, run the function to be converted in a plan in a context which
activates the tracing and record the actions in trace.logs
Second, store the result ids temporarily to helper ordering the output
placeholders at return time
Third, loop through the trace logs and replace the tensors found in the
actions logged by PlaceHolders. Record those actions in
plan.actions
Args:
args: Input arguments to run the plan
"""
self.owner.init_plan = self
self._tmp_args_ids = [
t.id for t in args if isinstance(t, FrameworkTensor)
]
with sy.hook.trace.enabled():
# We usually have include_state==True for functions converted to plan
# using @func2plan and we need therefore to add the state manually
if self.include_state:
results = self.forward(*args, self.state)
else:
results = self.forward(*args)
results = (results, ) if not isinstance(results, tuple) else results
self._tmp_result_ids = [
t.id for t in results if isinstance(t, FrameworkTensor)
]
for arg in args:
self.replace_with_placeholders(arg, node_type="input")
for log in sy.hook.trace.logs:
command, response = log
command_placeholders = self.replace_with_placeholders(
command, node_type="input")
return_placeholders = self.replace_with_placeholders(
response, node_type="output")
# We're cheating a bit here because we put placeholders instead of return_ids
action = ComputationAction(*command_placeholders,
return_ids=return_placeholders)
self.actions.append(action)
sy.hook.trace.clear()
del self._tmp_result_ids
del self._tmp_args_ids
self.is_built = True
self.owner.init_plan = None
def computation(name,
target,
args_,
kwargs_,
return_ids,
return_value=False):
""" Helper function to build a TensorCommandMessage containing a ComputationAction
directly from the action arguments.
"""
action = ComputationAction(name, target, args_, kwargs_, return_ids,
return_value)
return TensorCommandMessage(action)
def __init__(self, name, target, args_, kwargs_, return_ids):
"""Initialize an action message
Args:
message (Tuple): this is typically the args and kwargs of a method call on the client, but it
can be any information necessary to execute the action properly.
return_ids (Tuple): primarily for our async infrastructure (Plan, Protocol, etc.), the id of
action results are set by the client. This allows the client to be able to predict where
the results will be ahead of time. Importantly, this allows the client to pre-initalize the
pointers to the future data, regardless of whether the action has yet executed. It also
reduces the size of the response from the action (which is very often empty).
"""
self.action = ComputationAction(name, target, args_, kwargs_,
return_ids)
def bufferize(worker: AbstractWorker,
action_message: "CommandMessage") -> "CommandMessagePB":
"""
This function takes the attributes of a CommandMessage and saves them in Protobuf
Args:
worker (AbstractWorker): a reference to the worker doing the serialization
action_message (CommandMessage): an CommandMessage
Returns:
protobuf_obj: a Protobuf message holding the unique attributes of the message
Examples:
data = bufferize(message)
"""
protobuf_op_msg = CommandMessagePB()
protobuf_op = ComputationAction.bufferize(worker,
action_message.action)
protobuf_op_msg.action.CopyFrom(protobuf_op)
return protobuf_op_msg
def unbufferize(worker: AbstractWorker,
protobuf_obj: "CommandMessagePB") -> "CommandMessage":
"""
This function takes the Protobuf version of this message and converts
it into an CommandMessage. The bufferize() method runs the inverse of this method.
Args:
worker (AbstractWorker): a reference to the worker necessary for detailing. Read
syft/serde/serde.py for more information on why this is necessary.
protobuf_obj (CommandMessagePB): the Protobuf message
Returns:
obj (CommandMessage): an CommandMessage
Examples:
message = unbufferize(sy.local_worker, protobuf_msg)
"""
detailed = ComputationAction.unbufferize(worker, protobuf_obj.action)
return CommandMessage(detailed.name, detailed.target, detailed.args,
detailed.kwargs, detailed.return_ids)