fl4health.servers.nnunet_server module

class NnunetServer(client_manager, fl_config, on_init_parameters_config_fn, strategy=None, reporters=None, checkpoint_and_state_module=None, server_name=None, accept_failures=True, nnunet_trainer_class=<class 'nnunetv2.training.nnUNetTrainer.nnUNetTrainer.nnUNetTrainer'>, global_deep_supervision=False)[source]

Bases: FlServer

__init__(client_manager, fl_config, on_init_parameters_config_fn, strategy=None, reporters=None, checkpoint_and_state_module=None, server_name=None, accept_failures=True, nnunet_trainer_class=<class 'nnunetv2.training.nnUNetTrainer.nnUNetTrainer.nnUNetTrainer'>, global_deep_supervision=False)[source]

A Basic FlServer with added functionality to ask a client to initialize the global nnunet plans if one was not provided in the config. Intended for use with NnUNetClient.

Parameters:

  • client_manager (ClientManager) – Determines the mechanism by which clients are sampled by the server, if they are to be sampled at all.

  • fl_config (Config) –

    This should be the configuration that was used to setup the federated training. In most cases it should be the “source of truth” for how FL training/evaluation should proceed. For example, the config used to produce the on_fit_config_fn and on_evaluate_config_fn for the strategy.

    NOTE: This config is DISTINCT from the Flwr server config, which is extremely minimal.

  • on_init_parameters_config_fn (Callable[[int], dict[str, Scalar]]) – Function used to configure how one asks a client to provide parameters from which to initialize all other clients by providing a Config dictionary. For NnunetServers this is a required function to provide the additional information necessary to a client for parameter initialization

  • strategy (Strategy | None, optional) – The aggregation strategy to be used by the server to handle client updates and other information potentially sent by the participating clients. If None the strategy is FedAvg as set by the flwr Server. Defaults to None.

  • reporters (Sequence[BaseReporter] | None, optional) – A sequence of FL4Health reporters which the client should send data to. Defaults to None.

  • checkpoint_and_state_module (NnUnetServerCheckpointAndStateModule | None, optional) –

    This module is used to handle both model checkpointing and state checkpointing. The former is aimed at saving model artifacts to be used or evaluated after training. The latter is used to preserve training state (including models) such that if FL training is interrupted, the process may be restarted. If no module is provided, no checkpointing or state preservation will happen. Defaults to None.

    NOTE: For NnUnet, this module is allowed to have all components defined other than the model, as it may be set later when the server asks the clients to provide the architecture.

  • server_name (str | None, optional) – An optional string name to uniquely identify server. This name is also used as part of any state checkpointing done by the server. Defaults to None.

  • accept_failures (bool, optional) – Determines whether the server should accept failures during training or evaluation from clients or not. If set to False, this will cause the server to shutdown all clients and throw an exception. Defaults to True.

  • nnunet_trainer_class (type[nnUNetTrainer]) – nnUNetTrainer class. Useful for passing custom nnUNetTrainer. Defaults to the standard nnUNetTrainer class. Must match the nnunet_trainer_class passed to the NnunetClient.

  • global_deep_supervision (bool) – Whether or not the global model should use deep supervision. Does not affect the model architecture just the output during inference. This argument applies only to the global model, not local client models. Defaults to False.

initialize_server_model()[source]

Initializes the global server model so that it can be checkpointed.

Return type:

None

update_before_fit(num_rounds, timeout)[source]

Hook method to allow the server to do some additional initialization prior to fitting.

NunetServer uses this method to sample a client for properties for one of two reasons

  1. If a global nnunet_plans file is not provided in the config, this method will request that a random client which generate a plans file from it local dataset and return it to the server through the get_properties RPC. The server then distributes the nnunet_plans to the other clients by including it in the config for subsequent FL rounds.

    AND/OR

  2. If server side state or model checkpointing is being used, then server will poll a client in order to have the required properties to instantiate the model architecture on the server side. These properties include num_segmentation_heads and num_input_channels, essentially the number of input and output channels (which are not specified in nnunet plans for some reason).

Parameters:

  • num_rounds (int) – The number of server rounds of FL to be performed

  • timeout (float | None, optional) – The server’s timeout parameter. Useful if one is requesting information from a client. Defaults to None, which indicates indefinite timeout.

Return type:

None

add_items_to_config_fn(fn, items)[source]

Accepts a flwr Strategy configure function (either configure_fit or configure_evaluate) and returns a new function that returns the same thing except the dictionary items in the items argument have been added to the config that is returned by the original function

Parameters:

  • fn (CFG_FN) – The Strategy configure function to wrap

  • items (Config) – A Config containing additional items to update the original config with

Returns:

The wrapped function. Argument and return type is the same

Return type:

CFG_FN