blazefl.contrib.FedAvgProcessPoolClientTrainer

class blazefl.contrib.FedAvgProcessPoolClientTrainer(model_selector: ModelSelector, model_name: str, state_dir: Path, dataset: PartitionedDataset[FedAvgPartitionType], device: str, num_clients: int, epochs: int, batch_size: int, lr: float, seed: int, num_parallels: int)

Bases: ProcessPoolClientTrainer[FedAvgUplinkPackage, FedAvgDownlinkPackage, FedAvgClientConfig, FedAvgProcessPoolUplinkPackage]

Parallel client trainer for the Federated Averaging (FedAvg) algorithm.

This trainer handles the parallelized training and evaluation of local models across multiple clients, distributing tasks to different processes or devices.

model_selector

Selector for initializing the local model.

Type:

ModelSelector

model_name

Name of the model to be used.

Type:

str

state_dir

Directory to save random states for reproducibility.

Type:

Path

dataset

Dataset partitioned across clients.

Type:

PartitionedDataset

device

Device to run the models on (‘cpu’ or ‘cuda’).

Type:

str

num_clients

Total number of clients in the federation.

Type:

int

epochs

Number of local training epochs per client.

Type:

int

batch_size

Batch size for local training.

Type:

int

lr

Learning rate for the optimizer.

Type:

float

seed

Seed for reproducibility.

Type:

int

num_parallels

Number of parallel processes for training.

Type:

int

device_count

Number of CUDA devices available (0 if not using GPU).

Type:

int

__init__(model_selector: ModelSelector, model_name: str, state_dir: Path, dataset: PartitionedDataset[FedAvgPartitionType], device: str, num_clients: int, epochs: int, batch_size: int, lr: float, seed: int, num_parallels: int) → None

Initialize the FedAvgParalleClientTrainer.

Parameters:

• model_selector (ModelSelector) – Selector for initializing the local model.

• model_name (str) – Name of the model to be used.

• state_dir (Path) – Directory to save random states for reproducibility.

• dataset (PartitionedDataset) – Dataset partitioned across clients.

• device (str) – Device to run the models on (‘cpu’ or ‘cuda’).

• num_clients (int) – Total number of clients in the federation.

• epochs (int) – Number of local training epochs per client.

• batch_size (int) – Batch size for local training.

• lr (float) – Learning rate for the optimizer.

• seed (int) – Seed for reproducibility.

• num_parallels (int) – Number of parallel processes for training.

Methods

__init__(model_selector, model_name, ...)	Initialize the FedAvgParalleClientTrainer.
convert_buffer_to_uplink(buffer)	Convert a reconstructed BufferPackage to an UplinkPackage.
get_client_config(cid)	Generate the client configuration for a specific client.
get_client_device(cid)	Retrieve the device to use for processing a given client.
local_process(payload, cid_list)	Manage the parallel processing of clients.
prepare_uplink_package_buffer()	Allocate a pre-initialized shared memory buffer for a single client's result.
progress_fn(it)	A no-op progress function that can be overridden to provide custom progress tracking.
shutdown()	Shut down process-shared coordination resources owned by the trainer.
train(model, model_parameters, train_loader, ...)
uplink_package()	Retrieve the uplink packages for transmission to the server.
worker(config, payload, device, stop_event, *)	Process a single client's local training and evaluation.

Attributes

num_parallels
device
device_count
cache
stop_event

convert_buffer_to_uplink(buffer: FedAvgProcessPoolUplinkPackage) → FedAvgUplinkPackage

Convert a reconstructed BufferPackage to an UplinkPackage.

Called by local_process after shared memory reconstruction. When BufferPackage and UplinkPackage are the same type, implement this as return buffer.

Parameters:

buffer (BufferPackage) – The reconstructed buffer from shared memory.

Returns:

The uplink package to be stored in cache.

Return type:

UplinkPackage

get_client_config(cid: int) → FedAvgClientConfig

Generate the client configuration for a specific client.

Parameters:

cid (int) – Client ID.

Returns:

Client configuration data structure.

Return type:

FedAvgClientConfig

prepare_uplink_package_buffer() → FedAvgProcessPoolUplinkPackage

Allocate a pre-initialized shared memory buffer for a single client’s result.

Returns:

A buffer object whose tensors are in shared memory.

Return type:

BufferPackage

progress_fn(it: list[ApplyResult]) → Iterable[ApplyResult]

A no-op progress function that can be overridden to provide custom progress tracking.

Parameters:

it (list[ApplyResult]) – A list of ApplyResult objects.

Returns:

The original iterable.

Return type:

Iterable[ApplyResult]

static train(model: Module, model_parameters: Tensor, train_loader: DataLoader, device: str, epochs: int, lr: float, stop_event: Event, cid: int) → FedAvgProcessPoolUplinkPackage

uplink_package() → list[FedAvgUplinkPackage]

Retrieve the uplink packages for transmission to the server.

Returns:

A list of uplink packages.

Return type:

list[FedAvgUplinkPackage]

static worker(config: FedAvgClientConfig, payload: FedAvgDownlinkPackage, device: str, stop_event: Event, *, shm_buffer: FedAvgProcessPoolUplinkPackage | None = None) → FedAvgProcessPoolUplinkPackage

Process a single client’s local training and evaluation.

This method is executed by a worker process and handles loading client configuration and payload, performing the client-specific training, and returning the result.

Parameters:

• config (FedAvgClientConfig) – The client’s configuration data.

• payload (FedAvgDownlinkPackage) – The downlink payload from the server.

• device (str) – Device to use for processing (e.g., “cpu”, “cuda:0”).

• shm_buffer (FedAvgProcessPoolUplinkPackage | None) – Optional shared memory buffer for the uplink package.

Returns:

The uplink package containing the client’s results.

Return type:

FedAvgUplinkPackage