[Python] Why does my flwr federated learning setup freeze at round 1 while on deployment...

I'm new to the federated learning systems and I'm trying to train a model using 2 supernodes (clients). I managed to start the superlink (server) succesfully and successfully connect the nodes to the server and when I run the project it does the initialization normally and gives me the init result accuracy and then on round 1 when the clients are supposed to work it just freezes on a specific moment. The terminal log is this on round 1.

INFO : \[ROUND 1\]
INFO : configure_fit: strategy sampled 2 clients (out of 2)
INFO : \[Fleet.PullMessages\] node_id=8592592299027874682
INFO : \[Fleet.GetRun\] Requesting `Run` for run_id=8240479815143473977
INFO : \[Fleet.GetFab\] Requesting FAB for fab_hash=3a3df80e63875d5b80523c401b4449796aa11fb2cfba09c8484e1d77742a9eb5
INFO : \[Fleet.PullMessages\] node_id=10322998543652892263
INFO : \[Fleet.GetRun\] Requesting `Run` for run_id=8240479815143473977
INFO : \[Fleet.GetFab\] Requesting FAB for fab_hash=3a3df80e63875d5b80523c401b4449796aa11fb2cfba09c8484e1d77742a9eb5


after that both client and server do nothing.

This is the client log on the terminal:

INFO : Received: train message a73708823075e2486f9ae5396b9d6be84332babc6cf619c26f9dab9ec71e007b
INFO : Start `flwr-clientapp` process
INFO : \[flwr-clientapp\] Pull `ClientAppInputs` for token 817f...bdb6


it doesn't do anything after this

This is my client code:

import torch
from flwr.client import ClientApp, NumPyClient
from flwr.common import Context
from fedsystem.task import Net, get_weights, load_data, set_weights, test, train


# Define Flower Client and client_fn
class FlowerClient(NumPyClient):
def __init__(self, net, trainloader, valloader, local_epochs):
self.net = net
self.trainloader = trainloader
self.valloader = valloader
self.local_epochs = local_epochs
self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
self.net.to(self.device)

def fit(self, parameters, config):
set_weights(self.net, parameters)
train_loss = train(
self.net,
self.trainloader,
self.local_epochs,
config['lr'],
self.device,
)
return (
get_weights(self.net),
len(self.trainloader.dataset),
{"train_loss": train_loss},
)

def evaluate(self, parameters, config):
set_weights(self.net, parameters)
loss, accuracy = test(self.net, self.valloader, self.device)
return loss, len(self.valloader.dataset), {"accuracy": accuracy}


def client_fn(context: Context):
# Load model and data
net = Net()
partition_id = context.node_config["partition-id"]
num_partitions = context.node_config["num-partitions"]
trainloader, valloader = load_data(partition_id, num_partitions)
local_epochs = context.run_config["local-epochs"]

# Return Client instance
return FlowerClient(net, trainloader, valloader, local_epochs).to_client()


# Flower ClientApp
app = ClientApp(
client_fn,
)


I've tried running it on both insecure mode and with certificates for TLS connections but the issue still persists.

Continue reading...