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 operations 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
operations logged by PlaceHolders. Record those operations in
plan.operations
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 log in sy.hook.trace.logs:
command, response = log
command_placeholders, return_placeholders = (
self.replace_with_placeholders(command, node_type="input"),
self.replace_with_placeholders(response, node_type="output"),
)
# We're cheating a bit here because we put placeholders instead of return_ids
operation = Operation(*command_placeholders,
return_ids=return_placeholders)
self.operations.append(operation)
sy.hook.trace.clear()
del self._tmp_result_ids
del self._tmp_args_ids
self.is_built = True
self.owner.init_plan = None
def send_command(self,
recipient: "BaseWorker",
message: tuple,
return_ids: str = None
) -> Union[List[PointerTensor], PointerTensor]:
"""
Sends a command through a message to a recipient worker.
Args:
recipient: A recipient worker.
message: A tuple representing the message being sent.
return_ids: A list of strings indicating the ids of the
tensors that should be returned as response to the command execution.
Returns:
A list of PointerTensors or a single PointerTensor if just one response is expected.
"""
if return_ids is None:
return_ids = tuple([sy.ID_PROVIDER.pop()])
cmd_name = message[0]
cmd_owner = message[1]
cmd_args = message[2]
cmd_kwargs = message[3]
try:
ret_val = self.send_msg(Operation(cmd_name, cmd_owner, cmd_args,
cmd_kwargs, return_ids),
location=recipient)
except ResponseSignatureError as e:
ret_val = None
return_ids = e.ids_generated
if ret_val is None or type(ret_val) == bytes:
responses = []
for return_id in return_ids:
response = PointerTensor(
location=recipient,
id_at_location=return_id,
owner=self,
id=sy.ID_PROVIDER.pop(),
)
responses.append(response)
if len(return_ids) == 1:
responses = responses[0]
else:
responses = ret_val
return responses
def create_message_execute_command(
command_name: str, command_owner=None, return_ids=None, *args, **kwargs
):
"""helper function creating a message tuple for the execute_command call
Args:
command_name: name of the command that shall be called
command_owner: owner of the function (None for torch functions, "self" for classes derived from
workers.base or ptr_id for remote objects
return_ids: optionally set the ids of the return values (for remote objects)
*args: will be passed to the call of command_name
**kwargs: will be passed to the call of command_name
Returns:
tuple: (command_name, command_owner, args, kwargs), return_ids
"""
if return_ids is None:
return_ids = []
return Operation((command_name, command_owner, args, kwargs), return_ids)