Enhancing configure_fit and configure_evaluate in FedAvg

[oabuhamdan]

Type

Feature

Description

Hello,
In FedAvg strategy, creating custom configs with on_fit_config is called before sampling.

    def configure_fit(
        self, server_round: int, parameters: Parameters, client_manager: ClientManager
    ) -> list[tuple[ClientProxy, FitIns]]:
        """Configure the next round of training."""
        config = {}
        if self.on_fit_config_fn is not None:
            # Custom fit config function provided
            config = self.on_fit_config_fn(server_round)
        fit_ins = FitIns(parameters, config)

        # Sample clients
        sample_size, min_num_clients = self.num_fit_clients(
            client_manager.num_available()
        )
        clients = client_manager.sample(
            num_clients=sample_size, min_num_clients=min_num_clients
        )

        # Return client/config pairs
        return [(client, fit_ins) for client in clients]

This make it so hard to pass time-ralted configs using the cutomized on_fit_config function. The function will not proceed until it has the minimum number of clients.

in ClientManager sample()

        if min_num_clients is None:
            min_num_clients = num_clients
        self.wait_for(min_num_clients)

For example, I can measure the time a client sepnds computing, but If I want to measure the time the client spends in the network or computing, I need to override the whole Server or FedAvg class to make this possible.
Moving the configs creation to below sampling make this much easier, passing the time before going in the network as a customized config does it all.

Planned Implementation

fedavg.py

Additional Context

No response

[oabuhamdan]

One more fix,
In many FL experiments, we need to change the batch size or the number of local rounds specifically for a client.
In the current implementation, configs are the same for all clients.
This is actually an easy fix, similar for both fit and eval

def configure_fit(
    self, server_round: int, parameters: Parameters, client_manager: ClientManager
) -> list[tuple[ClientProxy, FitIns]]:
    """Configure the next round of training."""
    # Sample clients
    sample_size, min_num_clients = self.num_fit_clients(client_manager.num_available())
    clients = client_manager.sample(
        num_clients=sample_size, min_num_clients=min_num_clients
    )

    client_fitins_pairs = []
    for client in clients:
        if self.on_fit_config_fn is not None:
            # Pass client to config function for per-client config
            config = self.on_fit_config_fn(server_round, client)
        else:
            config = {}
        fit_ins = FitIns(parameters, config)
        client_fitins_pairs.append((client, fit_ins))

    return client_fitins_pairs

With fixing the type hints to

on_fit_config_fn: Optional[Callable[[int, ClientProxy], dict[str, Scalar]]] = None,

[WilliamLindskog]

hi @oabuhamdan,

thanks for raising this. We'll take a look at it and get back to you soon!