rigfl 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (89) hide show
  1. rigfl-0.1.0/DEVIATIONS.md +38 -0
  2. rigfl-0.1.0/LICENSE +21 -0
  3. rigfl-0.1.0/MANIFEST.in +6 -0
  4. rigfl-0.1.0/PKG-INFO +357 -0
  5. rigfl-0.1.0/README.md +330 -0
  6. rigfl-0.1.0/configs/datasets.yaml +16 -0
  7. rigfl-0.1.0/examples/smoke.py +129 -0
  8. rigfl-0.1.0/experiments/cifar10_run.yaml +13 -0
  9. rigfl-0.1.0/experiments/cifar10_tune.yaml +51 -0
  10. rigfl-0.1.0/experiments/cifar_baselines.yaml +21 -0
  11. rigfl-0.1.0/pyproject.toml +37 -0
  12. rigfl-0.1.0/rigfl/__init__.py +8 -0
  13. rigfl-0.1.0/rigfl/algorithms/__init__.py +1 -0
  14. rigfl-0.1.0/rigfl/algorithms/fedavg.py +164 -0
  15. rigfl-0.1.0/rigfl/algorithms/feddes.py +380 -0
  16. rigfl-0.1.0/rigfl/algorithms/fedgh.py +106 -0
  17. rigfl-0.1.0/rigfl/algorithms/fedkd.py +143 -0
  18. rigfl-0.1.0/rigfl/algorithms/fedproto.py +147 -0
  19. rigfl-0.1.0/rigfl/algorithms/fedprox.py +42 -0
  20. rigfl-0.1.0/rigfl/algorithms/fedtgp.py +157 -0
  21. rigfl-0.1.0/rigfl/algorithms/fml.py +83 -0
  22. rigfl-0.1.0/rigfl/algorithms/global_ensemble.py +52 -0
  23. rigfl-0.1.0/rigfl/algorithms/lgfedavg.py +70 -0
  24. rigfl-0.1.0/rigfl/algorithms/local.py +47 -0
  25. rigfl-0.1.0/rigfl/core/__init__.py +27 -0
  26. rigfl-0.1.0/rigfl/core/adapters.py +76 -0
  27. rigfl-0.1.0/rigfl/core/config.py +17 -0
  28. rigfl-0.1.0/rigfl/core/interfaces.py +115 -0
  29. rigfl-0.1.0/rigfl/core/model.py +52 -0
  30. rigfl-0.1.0/rigfl/core/round.py +399 -0
  31. rigfl-0.1.0/rigfl/data/README.md +107 -0
  32. rigfl-0.1.0/rigfl/data/__init__.py +61 -0
  33. rigfl-0.1.0/rigfl/data/biosilo.py +56 -0
  34. rigfl-0.1.0/rigfl/data/builder.py +147 -0
  35. rigfl-0.1.0/rigfl/data/config.py +265 -0
  36. rigfl-0.1.0/rigfl/data/flower.py +608 -0
  37. rigfl-0.1.0/rigfl/data/generate.py +38 -0
  38. rigfl-0.1.0/rigfl/data/partitions.py +254 -0
  39. rigfl-0.1.0/rigfl/eval/__init__.py +14 -0
  40. rigfl-0.1.0/rigfl/eval/metrics.py +240 -0
  41. rigfl-0.1.0/rigfl/eval/protocol.py +112 -0
  42. rigfl-0.1.0/rigfl/eval/report.py +269 -0
  43. rigfl-0.1.0/rigfl/eval/selection.py +274 -0
  44. rigfl-0.1.0/rigfl/experiment/__init__.py +11 -0
  45. rigfl-0.1.0/rigfl/experiment/artifacts.py +396 -0
  46. rigfl-0.1.0/rigfl/experiment/collect.py +380 -0
  47. rigfl-0.1.0/rigfl/experiment/config.py +160 -0
  48. rigfl-0.1.0/rigfl/experiment/device.py +23 -0
  49. rigfl-0.1.0/rigfl/experiment/env.py +110 -0
  50. rigfl-0.1.0/rigfl/experiment/launch.py +444 -0
  51. rigfl-0.1.0/rigfl/experiment/registry.py +160 -0
  52. rigfl-0.1.0/rigfl/experiment/run.py +402 -0
  53. rigfl-0.1.0/rigfl/experiment/tracking.py +78 -0
  54. rigfl-0.1.0/rigfl/experiment/tuning.py +887 -0
  55. rigfl-0.1.0/rigfl/models/__init__.py +1 -0
  56. rigfl-0.1.0/rigfl/models/cifar.py +185 -0
  57. rigfl-0.1.0/rigfl/models/eicu.py +96 -0
  58. rigfl-0.1.0/rigfl/models/registry.py +120 -0
  59. rigfl-0.1.0/rigfl/prediction.py +120 -0
  60. rigfl-0.1.0/rigfl.egg-info/PKG-INFO +357 -0
  61. rigfl-0.1.0/rigfl.egg-info/SOURCES.txt +87 -0
  62. rigfl-0.1.0/rigfl.egg-info/dependency_links.txt +1 -0
  63. rigfl-0.1.0/rigfl.egg-info/requires.txt +16 -0
  64. rigfl-0.1.0/rigfl.egg-info/top_level.txt +1 -0
  65. rigfl-0.1.0/scripts/run_grid.sh +35 -0
  66. rigfl-0.1.0/setup.cfg +4 -0
  67. rigfl-0.1.0/tests/test_artifacts.py +218 -0
  68. rigfl-0.1.0/tests/test_biosilo.py +170 -0
  69. rigfl-0.1.0/tests/test_builder.py +107 -0
  70. rigfl-0.1.0/tests/test_collect.py +318 -0
  71. rigfl-0.1.0/tests/test_config_gaps.py +120 -0
  72. rigfl-0.1.0/tests/test_contract.py +84 -0
  73. rigfl-0.1.0/tests/test_fedavg_fedprox.py +278 -0
  74. rigfl-0.1.0/tests/test_flower_data.py +292 -0
  75. rigfl-0.1.0/tests/test_grid.py +165 -0
  76. rigfl-0.1.0/tests/test_history.py +159 -0
  77. rigfl-0.1.0/tests/test_identity.py +77 -0
  78. rigfl-0.1.0/tests/test_loop.py +143 -0
  79. rigfl-0.1.0/tests/test_metrics.py +67 -0
  80. rigfl-0.1.0/tests/test_model_registry.py +209 -0
  81. rigfl-0.1.0/tests/test_partitions.py +236 -0
  82. rigfl-0.1.0/tests/test_pool_cache.py +233 -0
  83. rigfl-0.1.0/tests/test_prediction.py +567 -0
  84. rigfl-0.1.0/tests/test_provenance.py +255 -0
  85. rigfl-0.1.0/tests/test_report.py +223 -0
  86. rigfl-0.1.0/tests/test_run_experiment.py +79 -0
  87. rigfl-0.1.0/tests/test_selection.py +208 -0
  88. rigfl-0.1.0/tests/test_tuning.py +1011 -0
  89. rigfl-0.1.0/tests/test_validation_loss.py +392 -0
@@ -0,0 +1,38 @@
1
+ # Deviations from the original algorithms
2
+
3
+ RigFL aims to follow each algorithm's published specification. Departures that
4
+ affect interpretation or comparability are listed here.
5
+
6
+ | Algorithm | Departure | Why / consequence |
7
+ |---|---|---|
8
+ | **FedAvg / FedProx** | Every client participates in every round. | The current shared round loop has no client-sampling policy. The update is faithful to full-participation FedAvg/FedProx, but experiments do not reproduce partial-participation settings from the papers. |
9
+ | **FedAvg / FedProx** | Non-floating model state is copied from the client with the largest local sample count (first client wins ties). | Integer buffers such as BatchNorm's `num_batches_tracked` cannot be sample-weight averaged without inventing a non-client integer value. Floating parameters and buffers remain sample-count weighted. |
10
+ | **FedProx** | Clients perform a fixed configured number of local epochs. | This implements the proximal objective but not the paper's systems-heterogeneity experiments with variable work or its abstract gamma-inexact local solver. |
11
+ | **FedProto / FedTGP** | Predictive probabilities are `softmax(-d)` over Euclidean prototype distances. | The papers define nearest-prototype labels but not probabilities. This preserves the paper's decision rule and supplies a predictive loss for evaluation and early stopping. Its scale follows the learned representation, so compare this loss within a run rather than as a calibrated score across algorithms or architectures. |
12
+ | **FedProto** | Prototypes are computed in a clean pass after local training. | Some implementations accumulate them during training, which averages features from a model that was still moving. |
13
+ | **FedGH** | The global header is trained on the server, per Algorithm 1 / Eq. 4. | Comparisons should verify that the server optimizer updates the header. |
14
+ | **LG-FedAvg** | Trains from scratch. | The original's released scripts warm-start from an 800–1800-round FedAvg checkpoint and then run 500 LG rounds. From-scratch is a weaker configuration, so this number is not comparable to the paper's. |
15
+ | **LG-FedAvg** | Shares exactly one `nn.Linear` (the head). | The original keeps two layers local. RigFL's split point is fixed by `ClientModel`'s structure rather than configurable. |
16
+ | **FedKD** | Averages mentee *parameters*, not gradients. | The paper transmits gradients and applies `Θ_s -= η_s · ḡ`. The two coincide at `local_epochs = 1` and diverge above it. |
17
+ | **FedKD** | Distils only the final shared representation. | The original pairs 4 student to 12 teacher layers on a fixed stride and also distils self-attention maps. Neither has an analogue in a CNN pool. |
18
+ | **FedKD** | 4-D convolution tensors are not SVD-compressed. | Only 2-D weights are factorized. **Accuracy is faithful; communication cost is not — do not quote a communication-saving number from this implementation.** |
19
+ | **FedKD** | SGD rather than Adam. | The paper's learning rates (2e-6 / 5e-6) are tuned for a Transformer on NLP tasks and do not transfer to a CNN pool. |
20
+ | **FML / FedKD** | Applied to a heterogeneous architecture pool. | Both originals assume an identical shared component across clients (a meme model, a mentee). Using them across genuinely different backbones is a generalization beyond their original setting. |
21
+ | **LG-FedAvg** | Applied to a heterogeneous architecture pool. | The original's main experiments use one architecture for every client. |
22
+
23
+ ## Configuration
24
+
25
+ Optimizer hyperparameters are set per experiment and may be overridden by a
26
+ sweep. The resolved values are stored with each result.
27
+
28
+ ## Prototype loss interpretation
29
+
30
+ FedProto and FedTGP predict the nearest prototype. RigFL converts those distances
31
+ into `softmax(-d)` probabilities, whose argmax preserves the nearest-prototype
32
+ decision.
33
+
34
+ Distance scale is learned and can change across rounds. Prototype loss is useful
35
+ for monitoring that prediction rule within one run, but it is not a calibrated
36
+ quantity for comparing unrelated algorithms or representation spaces. RigFL does not
37
+ add normalization or fitted temperature calibration, because either would change
38
+ the evaluated algorithm and introduce another experimental choice.
rigfl-0.1.0/LICENSE ADDED
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 Brianna Mueller
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,6 @@
1
+ include DEVIATIONS.md
2
+ recursive-include configs *.yaml
3
+ recursive-include examples *.py
4
+ recursive-include experiments *.yaml
5
+ recursive-include scripts *.sh
6
+ include rigfl/data/README.md
rigfl-0.1.0/PKG-INFO ADDED
@@ -0,0 +1,357 @@
1
+ Metadata-Version: 2.4
2
+ Name: rigfl
3
+ Version: 0.1.0
4
+ Summary: A modular research framework for rigorous federated learning experimentation
5
+ Author: Brianna Mueller
6
+ License-Expression: MIT
7
+ Project-URL: Repository, https://github.com/briannamueller/RigFL
8
+ Project-URL: Issues, https://github.com/briannamueller/RigFL/issues
9
+ Requires-Python: <3.13,>=3.10
10
+ Description-Content-Type: text/markdown
11
+ License-File: LICENSE
12
+ Requires-Dist: numpy<2
13
+ Requires-Dist: torch
14
+ Requires-Dist: flwr-datasets[vision]<0.7,>=0.6
15
+ Requires-Dist: pydantic>=2
16
+ Requires-Dist: pyyaml
17
+ Requires-Dist: datasets<2.19
18
+ Requires-Dist: pyarrow<15
19
+ Requires-Dist: pillow<11
20
+ Requires-Dist: pandas<2.3
21
+ Requires-Dist: graphroute<0.2,>=0.1.0
22
+ Provides-Extra: wandb
23
+ Requires-Dist: wandb; extra == "wandb"
24
+ Provides-Extra: test
25
+ Requires-Dist: pytest; extra == "test"
26
+ Dynamic: license-file
27
+
28
+ # RigFL
29
+
30
+ RigFL is a modular framework for rigorous federated learning experimentation.
31
+ Algorithm-specific behavior is isolated behind a common interface, so that algorithms use the same orchestration, evaluation, configuration, and reporting
32
+ machinery.
33
+
34
+ - [Key Features](#key-features)
35
+ - [Installation](#installation)
36
+ - [Example workflow](#example-workflow)
37
+ - [Algorithms](#algorithms)
38
+ - [Sweeps and tuning](#sweeps-and-tuning)
39
+ - [Adding an algorithm](#adding-an-algorithm)
40
+ - [Experiment tracking with Weights & Biases](#experiment-tracking-with-weights--biases)
41
+ - [Development and testing](#development-and-testing)
42
+
43
+ ## Key Features
44
+
45
+ - **Stable experiment and partition identity.** RigFL derives two separate
46
+ fingerprints: one identifying a result from its distinct experiment
47
+ configuration, the other identifying a partitioned dataset from its data
48
+ configuration. A change to either produces a new identity, so earlier generated
49
+ results and partitions are never overwritten. An experiment whose result
50
+ already exists is not rerun—expanding or changing a sweep will only execute new
51
+ combinations.
52
+
53
+ - **Support for model-heterogeneous algorithms.** RigFL supports algorithms
54
+ designed for clients with different model architectures. The architecture
55
+ selection can be configured as a named family or an explicit ordered list.
56
+
57
+ - **Joint hyperparameter tuning across multiple seeds.** Support for evaluating
58
+ combinations of hyperparameters across several random seeds.
59
+
60
+ - **Client-centered performance reporting.** Evaluation metrics that reveal
61
+ whether the benefits of collaboration are broadly shared across clients,
62
+ exposing disparities and uneven benefits that commonly reported averages
63
+ obscure. See [Client-centered metrics](#client-centered-metrics).
64
+
65
+ - **Traceable result files.** Each completed experiment produces a result file
66
+ containing its full evaluation history, resolved configuration, Git commit and
67
+ uncommitted-change status, software versions, and client data-partition
68
+ information.
69
+
70
+ - **Documented fidelity to the source papers.** Algorithms follow their published
71
+ specifications; where a paper leaves a detail unspecified or its released code
72
+ diverges from the text, the resolution is recorded in
73
+ [DEVIATIONS.md](https://github.com/briannamueller/RigFL/blob/main/DEVIATIONS.md).
74
+
75
+ - **Optional W&B tracking.** Weights & Biases can be enabled to log experiment
76
+ settings and validation performance during training.
77
+
78
+ ## Installation
79
+
80
+ ```bash
81
+ pip install rigfl
82
+ ```
83
+
84
+ RigFL requires Python 3.10–3.12.
85
+
86
+ ## Example workflow
87
+
88
+ The following CIFAR-10 example walks you through generating client data
89
+ partitions, running an experiment, and reporting the results.
90
+
91
+ ### Generate client data partitions
92
+
93
+ First define the data source and data partitioning configuration in
94
+ [`configs/datasets.yaml`](https://github.com/briannamueller/RigFL/blob/main/configs/datasets.yaml):
95
+
96
+ ```yaml
97
+ datasets:
98
+ cifar10:
99
+ backend: flower
100
+ source_dataset: uoft-cs/cifar10
101
+ partition:
102
+ scheme: dirichlet
103
+ num_clients: 3
104
+ alpha: 0.5
105
+ ```
106
+
107
+ Generate client datasets by running:
108
+
109
+ ```bash
110
+ python -m rigfl.data.generate --dataset cifar10
111
+ ```
112
+
113
+ RigFL passes the settings to Flower, derives a stable fingerprint from the data
114
+ partitioning configuration, and saves the generated files under:
115
+
116
+ ```text
117
+ data/cifar10/partition_<fingerprint>/
118
+ ├── manifest.json
119
+ └── clients/
120
+ ├── client_0/
121
+ │ ├── train.pt
122
+ │ ├── validation.pt
123
+ │ └── test.pt
124
+ └── ...
125
+ ```
126
+
127
+ Running the command again with the same data configuration reuses the existing
128
+ partition. Changing a partitioning entry produces a different fingerprint and a
129
+ separate directory instead of replacing the previous partition.
130
+
131
+ To add another dataset, create another entry in
132
+ [`configs/datasets.yaml`](https://github.com/briannamueller/RigFL/blob/main/configs/datasets.yaml).
133
+ See the
134
+ [data configuration guide](https://github.com/briannamueller/RigFL/blob/main/rigfl/data/README.md)
135
+ for the available settings and guidance for datasets with multiple
136
+ configurations, nonstandard splits, or ambiguous input and target columns.
137
+
138
+ ### Define and run the experiment
139
+
140
+ YAML files define experiment configurations:
141
+ [`experiments/cifar10_run.yaml`](https://github.com/briannamueller/RigFL/blob/main/experiments/cifar10_run.yaml):
142
+
143
+ ```yaml
144
+ experiment:
145
+ dataset: cifar10
146
+ model_architectures: [fedavg_cnn]
147
+ rounds: 2
148
+ seed: 0
149
+ shared_dim: 128
150
+ eval_gap: 1
151
+ device: cpu
152
+ out_dir: results/cifar10_run
153
+
154
+ algorithm:
155
+ local_epochs: 1
156
+ lr: 0.01
157
+ ```
158
+
159
+ The YAML has two sections. Entries under `experiment` define the overarching
160
+ configuration for the execution of RigFL’s shared workflow. Entries under
161
+ `algorithm` specify how individual algorithms operate. An algorithm entry may be
162
+ supported by one or several algorithms. In a multi-algorithm sweep, each entry is
163
+ applied only to algorithms that support it.
164
+
165
+ Run the experiment with:
166
+
167
+ ```bash
168
+ python -m rigfl.experiment.run \
169
+ --algorithm fedavg \
170
+ --config experiments/cifar10_run.yaml
171
+ ```
172
+
173
+ This trains FedAvg for two communication rounds and writes results to
174
+ `results/cifar10_run`.
175
+
176
+ ### Report results
177
+
178
+ Summarize the results with:
179
+
180
+ ```bash
181
+ python -m rigfl.experiment.collect \
182
+ --results-dir results/cifar10_run
183
+ ```
184
+
185
+ ### Client-centered metrics
186
+
187
+ Aggregate performance metrics can signal that collaborative learning improves
188
+ upon local training on average, even though collaboration worsens performance at
189
+ some individual clients. RigFL provides evaluation metrics that surface unevenly
190
+ distributed benefits.
191
+
192
+ - **Win rate:** the fraction of matched client-and-seed pairs in which an algorithm
193
+ results in improved performance over the Local baseline.
194
+
195
+ - **Performance among the worst-served clients:** reports the average performance
196
+ of the lowest-scoring 10% of clients and the 10th-percentile score, which marks
197
+ the lower tail of the client-performance distribution.
198
+
199
+ - **Standard deviation:** the spread in performance across clients.
200
+
201
+ ## Algorithms
202
+
203
+ Implemented algorithms:
204
+
205
+ - [FedAvg](https://proceedings.mlr.press/v54/mcmahan17a.html)
206
+ - [FedProx](https://arxiv.org/abs/1812.06127)
207
+ - [FedProto](https://ojs.aaai.org/index.php/AAAI/article/view/20819)
208
+ - [FedGH](https://arxiv.org/abs/2303.13137)
209
+ - [LG-FedAvg](https://arxiv.org/abs/2001.01523)
210
+ - [FML](https://arxiv.org/abs/2006.16765)
211
+ - [FedKD](https://www.nature.com/articles/s41467-022-29763-x)
212
+ - [FedTGP](https://ojs.aaai.org/index.php/AAAI/article/view/29617)
213
+ - [FedDES](https://arxiv.org/abs/2603.28006)
214
+
215
+ Local training and Global Ensemble are available as reference baselines.
216
+
217
+ Algorithm-specific departures from the original papers are documented in
218
+ [DEVIATIONS.md](https://github.com/briannamueller/RigFL/blob/main/DEVIATIONS.md).
219
+
220
+ ## Sweeps and tuning
221
+
222
+ [`experiments/cifar10_tune.yaml`](https://github.com/briannamueller/RigFL/blob/main/experiments/cifar10_tune.yaml)
223
+ provides a multi-algorithm, multi-seed tuning example. A sweep expands the values
224
+ defined along each axis. Algorithm entries are applied only to algorithms that
225
+ support them, so options belonging to different algorithms are not unnecessarily
226
+ cross-multiplied.
227
+
228
+ Expand the sweep and print its cluster submission command with:
229
+
230
+ ```bash
231
+ python -m rigfl.experiment.launch \
232
+ --config experiments/cifar10_tune.yaml \
233
+ --queue <queue>
234
+ ```
235
+
236
+ Each complete hyperparameter combination is treated as one candidate, with its
237
+ seeds aggregated as replicates. Rank the completed candidates and write runnable
238
+ selected configurations with:
239
+
240
+ ```bash
241
+ python -m rigfl.experiment.collect \
242
+ --results-dir results/cifar10_tune \
243
+ --selection-metric accuracy \
244
+ --selection-view both \
245
+ --rank \
246
+ --select-out results/cifar10_tune_selected
247
+ ```
248
+
249
+ ## Adding an algorithm
250
+
251
+ Extend RigFL by adding a module under `rigfl/algorithms/` containing:
252
+
253
+ - A configuration class that inherits from `AlgorithmConfig`.
254
+ - An algorithm class that inherits from `Algorithm`.
255
+
256
+ The algorithm class must define four operations:
257
+
258
+ 1. `init_globals()` initializes the shared state, which represents the
259
+ information the server maintains and distributes to clients at the start of
260
+ each round. The shared state may take the form of a global model, model
261
+ parameters, prototypes, a classifier head, or another algorithm-specific
262
+ structure.
263
+ 2. `local_train(...)` is called once per client per round. It receives the client
264
+ and shared state, performs the client-side computation, and returns the
265
+ client's upload, which represents the information the client sends to the
266
+ server. The upload may have the same form as the shared state, be a different
267
+ structure entirely, or carry additional information required for server-side
268
+ computation.
269
+ 3. `aggregate(...)` receives all client uploads, performs the server-side
270
+ computation, and returns the shared state for the next round. This may involve
271
+ averaging parameters, combining prototypes, or training a server-side
272
+ component.
273
+ 4. `predict(...)` performs inference for the supplied inputs and returns a
274
+ `Predictions` object.
275
+
276
+ Declare all of the relevant arguments for the algorithm in its configuration
277
+ class.
278
+
279
+ ```python
280
+ from rigfl.core import Algorithm, Predictions
281
+ from rigfl.core.config import AlgorithmConfig
282
+
283
+
284
+ class NewAlgorithmConfig(AlgorithmConfig):
285
+ local_epochs: int = 1
286
+ lr: float = 0.01
287
+ # ...additional arguments
288
+
289
+
290
+ class NewAlgorithm(Algorithm):
291
+ def init_globals(self):
292
+ ...
293
+
294
+ def local_train(self, client, shared_state):
295
+ ...
296
+
297
+ def aggregate(self, client_uploads, shared_state):
298
+ ...
299
+
300
+ def predict(self, client, x, shared_state) -> Predictions:
301
+ ...
302
+ ```
303
+
304
+ In `local_train(...)` and `predict(...)`, `client` refers to the
305
+ `Client` instance being processed. The client's local model and training data loader are accessed
306
+ through `client.model` and `client.train_loader`, respectively. `client.state` is
307
+ a dictionary that can carry any additional client-specific information that must
308
+ persist across rounds.
309
+
310
+ Access the arguments defined in the algorithm’s configuration class through self.config, such as self.config.lr.
311
+
312
+ Register both classes in `rigfl/experiment/registry.py`:
313
+
314
+ ```python
315
+ REGISTRY = {
316
+ "local": AlgorithmSpec(Local, LocalConfig),
317
+ "fedavg": AlgorithmSpec(FedAvg, FedAvgConfig),
318
+ # ...other algorithms
319
+ "new_algorithm": AlgorithmSpec(NewAlgorithm, NewAlgorithmConfig),
320
+ }
321
+ ```
322
+
323
+ > **Runner note:** `AlgorithmSpec` uses the `iterative` runner by default. If an
324
+ > algorithm genuinely cannot be expressed as repeated local training followed
325
+ > by aggregation, define a different runner and matching operation protocol
326
+ > instead of changing the meaning of the standard operations. FedDES is one
327
+ > such exception and uses `p2p_one_shot`.
328
+
329
+ ## Experiment tracking with Weights & Biases
330
+
331
+ Install the optional W&B dependency with:
332
+
333
+ ```bash
334
+ pip install "rigfl[wandb]"
335
+ ```
336
+
337
+ Enable tracking by setting `wandb: true` under `experiment` in the YAML
338
+ configuration file, or pass `--wandb` when running experiments from the command
339
+ line.
340
+
341
+ ## Development and testing
342
+
343
+ Clone the repository and install RigFL in editable mode with its testing
344
+ dependency, then run the test suite:
345
+
346
+ ```bash
347
+ git clone https://github.com/briannamueller/RigFL.git
348
+ cd RigFL
349
+ python -m venv .venv
350
+ source .venv/bin/activate
351
+ pip install -e ".[test]"
352
+ pytest -q
353
+ ```
354
+
355
+ ## License
356
+
357
+ MIT. See [LICENSE](https://github.com/briannamueller/RigFL/blob/main/LICENSE).