Add Actor Init Callback

[jafermarq]

This extends the content of #1969 adding support for Actors to accept a user-defined callback to execute upon object creation. This callback is optional.

One of the most obvious reasons for using this callback is when running Flower with TF. In those cases, due to the default behaviour of TF of mapping the entire VRAM, it's likely that users will encounter OOM immediately (even if their workloads don't actually need all VRAM). Lukily, we can design a workaround by enabling gpu growth. This is exactly what the built-in enable_tf_gpu_growth callback in src/.../simulation/ray_transport/utils.py allows you to do.

You can easily test this by editing the lines around start_simulation in the TF simulation example:

# keep rest of sim.py as is.
....

def main() -> None:

    from flwr.simulation.ray_transport.utils import enable_tf_gpu_growth

    # Enable GPU growth in the main thread (the one used by the server to quite likely
    # run global evaluation using GPU)
    enable_tf_gpu_growth()

    # Start Flower simulation
    fl.simulation.start_simulation(
        client_fn=client_fn,
        num_clients=NUM_CLIENTS,
        client_resources={"num_cpus": 2, "num_gpus": 0.25}, #<----- let's pack 4 actors in each GPU.
        config=fl.server.ServerConfig(num_rounds=5),
        strategy=fl.server.strategy.FedAvg(
            fraction_fit=0.1,
            fraction_evaluate=0.1,
            min_fit_clients=10,
            min_evaluate_clients=10,
            min_available_clients=NUM_CLIENTS,
        ),
        actor_kwargs={
            "on_actor_init_fn": enable_tf_gpu_growth # <--------------- pass the same func to execute upon actor init.
        },  
    )

[danieljanes]

We can log this before the try by checking if TF is None (we should also return from the function in that case).

An exception on this line could be many things (whatever the experimental TensorFlow function decides to raise).

[jafermarq]

true. i moved the TF check at the top of the function.

[danieljanes]

Lgtm!