nvflare.edge.tools.edge_fed_buff_recipe module
- class DeviceManagerConfig(device_selection_size: int = 100, min_hole_to_fill: int = 1, device_reuse: bool = True)[source]
Bases:
objectConfiguration class for the device manager in federated learning.
This class configures how the device manager selects and manages devices for participation in federated learning workflow.
- device_selection_size
Number of devices to select for each training round. Default: 100
- min_hole_to_fill
Minimum number of model updates to wait for before sampling the next batch of devices and dispatching the current global model. - If set to 1, the server immediately dispatch the current global model to a sampled device. - Higher values cause the server to wait for more updates before dispatching. - If set to device_selection_size, we will have synchronous training since all devices’ responses need to be collected before dispatching the next global model. This parameter works with num_updates_for_model from model manager to achieve trade-off between global model versioning and local execution. Default: 1 (immediately dispatch the current global model)
- device_reuse
Whether to allow devices to participate in multiple rounds. if False, devices will be selected only once, which could be realistic for real-world scenarios where the device pool is huge while participation is random. Default: True (always reuse / include the existing devices for further learning)
- class EdgeFedBuffRecipe(job_name: str, model: Any | Dict, model_manager_config: ModelManagerConfig, device_manager_config: DeviceManagerConfig, initial_ckpt: str | None = None, evaluator_config: EvaluatorConfig | None = None, simulation_config: SimulationConfig | None = None, custom_source_root: str | None = None, device_wait_timeout: float | None = None)[source]
Bases:
RecipeRecipe class for cross-edge federated learning using NVFlare’s hierarchical edge system.
This class provides a high-level interface for configuring cross-edge federated learning jobs. It configures the necessary components including model managers, device managers, evaluators, and device simulation settings.
The recipe supports both real device connections and simulated device training, making it suitable for both production deployment and prototyping/testing.
- Example usage:
```python # Basic configuration with model instance recipe = EdgeFedBuffRecipe(
job_name=”my_edge_job”, model=MyModel(), model_manager_config=ModelManagerConfig(…), device_manager_config=DeviceManagerConfig(…)
)
# With dict config and pre-trained checkpoint recipe = EdgeFedBuffRecipe(
job_name=”my_edge_job”, model={“class_path”: “my_module.MyModel”, “args”: {“num_classes”: 10}}, initial_ckpt=”/path/to/pretrained.pt”, model_manager_config=ModelManagerConfig(…), device_manager_config=DeviceManagerConfig(…)
- Parameters:
job_name – Name of the federated learning job.
model – PyTorch model to be trained. Can be: - nn.Module instance: e.g., MyModel() - Dict config: {“class_path”: “module.ClassName”, “args”: {“param”: value}}
initial_ckpt – Absolute path to a pre-trained checkpoint file (.pt, .pth). The file may not exist locally (server-side path). Used to resume training from pre-trained weights.
model_manager_config – Configuration for the model manager.
device_manager_config – Configuration for the device manager.
evaluator_config – Configuration for the global evaluator (optional).
simulation_config – Configuration for simulated devices settings (optional).
custom_source_root – Path to custom source code (optional).
device_wait_timeout – Timeout in seconds for waiting for sufficient devices to join before stopping the job. None means wait indefinitely. WARNING: when device_reuse=False with a finite device pool, leaving this as None can cause the job to hang indefinitely once the pool is exhausted. In that case, set an explicit timeout (e.g., 300.0 seconds). Default: None
This is base class of a recipe. Recipes are implemented by jobs. A concrete recipe must provide the job for recipe implementation.
- Parameters:
job – the job that implements the recipe.
- class EvaluatorConfig(eval_frequency: int = 1, torchvision_dataset: Dict | None = None, custom_dataset: Dict | None = None)[source]
Bases:
objectConfiguration class for the global evaluator.
This class configures how the global model is evaluated during training, including dataset selection and evaluation frequency.
- eval_frequency
Frequency of global model evaluation (every N new model versions). Default: 1
- torchvision_dataset
Configuration for torchvision datasets. Should be a dict with ‘name’ and ‘path’ keys. Default: None
- custom_dataset
Configuration for custom datasets. Default: None
- class ModelManagerConfig(max_num_active_model_versions: int = 3, max_model_version: int = 20, update_timeout: int = 5, num_updates_for_model: int = 100, max_model_history: int = 10, staleness_weight: bool = False, global_lr: float = 0.01)[source]
Bases:
objectConfiguration class for the model manager in federated learning.
This class configures how the model manager handles model updates, versioning, and aggregation strategies for federated learning workflow.
- max_num_active_model_versions
Maximum number of active model versions that can be processed for the current model version. Default: 3
- max_model_version
Maximum model version number before stopping training. We start with version 1 initial model, so the minimum for this arg is 2 to have at least one local update phase. Default: 20
- update_timeout
Timeout in seconds for waiting for model updates. Default: 5.0
- num_updates_for_model
Number of received updates required before generating a new global model. Default: 100
- max_model_history
Maximum number of model versions to keep in history for staleness calculations and update aggregation. Default: 10
- staleness_weight
Whether to apply staleness weighting to model updates. Default: False
- global_lr
Global learning rate for model aggregation. Default: 0.01
- class SimulationConfig(task_processor: DeviceTaskProcessor | None, job_timeout: float = 60.0, num_devices: int = 1000, num_workers: int = 10)[source]
Bases:
objectConfiguration class for simulation settings in federated learning.
This class configures the simulated devices for testing federated learning pipelines.
- task_processor
Task processor for handling device training simulation.
- job_timeout
Timeout in seconds for the entire job execution. Default: 60.0
- num_devices
Total number of simulated devices for each leaf node. Default: 1000
- num_workers
Number of worker processes for parallel device simulation on each leaf node. Default: 10