test_execute_plan_module_remotely

def test_execute_plan_module_remotely(hook, start_proc):
    """Test plan execution remotely."""
    hook.local_worker.is_client_worker = False

    class Net(nn.Module):
        def __init__(self):
            super(Net, self).__init__()
            self.fc1 = nn.Linear(2, 3)
            self.fc2 = nn.Linear(3, 2)

        @sy.method2plan
        def forward(self, x):
            x = F.relu(self.fc1(x))
            x = self.fc2(x)
            return F.log_softmax(x, dim=0)

    net = Net()
    x = th.tensor([-1, 2.0])
    
    # local_res = net(x)

    kwargs = {"id": "test_plan_worker_2", "host": "localhost", "port": 8799, "hook": hook}
    _ = start_proc(WebsocketServerWorker, kwargs)

    time.sleep(0.1)
    socket_pipe = WebsocketClientWorker(**kwargs)

    plan_ptr = net.send(socket_pipe)
    x_ptr = x.send(socket_pipe)
    remote_res = plan_ptr(x_ptr).get()