Deactivate Plans in Plans

[LaRiffle]

Plans should not be executed as Plans when they are included in another Plans. This allows to easily extend Plan or to cut Plans in small modules which can also be run individually as Plans.

Examples:

@func2plan
def foo1(x):
     return x + 1

@func2plan
def foo2(x):
     y = foo1(x)
     return y * 2

Idea: create a variable global at worker scale when a plan is being built to force other plans to run as normal functions.

[gmuraru]

@LaRiffle tried the following code - in the Plans notebook:

@func2plan()
def plan_double_abs(x):
    x = x + x
    x = torch.abs(x)
    return x

@func2plan()
def plan_test(x):
    f = plan_double_abs(x)
    return f*2

and from what I checked, if I run plan_test.build(torch.tensor([1., -2.])) it will simply call the build method for the plan_test (will not call build for plan_double_abs even thought it is called inside plan_test)
and after that I ran:

pointer_plan = plan_test.send(device_1) # device_1 is a VirtualWorker
pointer_to_data = device_1.search('input_data')[0]
pointer_to_result = pointer_plan(pointer_to_data)

Also, added a print message in the execute_commands methpd from the Plan class:

for message in self.procedure.operations:
    print("message ", message)
    bin_message = sy.serde.serialize(message, simplified=True)
    _ = self.owner.recv_msg(bin_message)

The messages generated are:

message  (33, ((6, ((5, (b'__add__',)), (25, (98142542040, 78527494689, b'device_1', None, (11, (2,)), True)), (6, ((25, (98142542040, 78527494689, b'device_1', None, (11, (2,)), True)),)), (0, ()))), (70741324193,)))
message  (33, ((6, ((5, (b'torch.abs',)), None, (6, ((25, (25994186521, 70741324193, b'device_1', None, None, True)),)), (0, ()))), (91046882485,)))
message  (33, ((6, ((5, (b'__mul__',)), (25, (14581124267, 91046882485, b'device_1', None, None, True)), (6, (2,)), (0, ()))), (82141830512,)))

From what it seems __add__ and torch.abs are from plan_double_abs and the __mul__ from the plan_test (from the return).

What I think is happening is that the plan_double_abs plan is not building the plan inside.
But also, I might be wrong...

[LaRiffle]

Amazing! 🕺🏻
Thanks for this finding :)
Ok what happen now if you supply args_shape in the plan definition?
Example:

@sy.func2plan(args_shape=[(1,)])
    def plan_abs(data):
        return data.abs()

[gmuraru]

Tested with:

@func2plan(args_shape=[(1,)])
def plan_double_abs(x):
    x = x + x
    x = torch.abs(x)
    return x

@func2plan()
def plan_test(x):
    f = plan_double_abs(x)
    return f*2

It looks like if the plan_double_abs is built, then we will call the plan_double_abs in plan_test it would go through the plan_double_abs (which is a "plan class").

Added a building context at the owner (worker) level.