flwr-nightly 1.8.0.dev20240315__py3-none-any.whl → 1.15.0.dev20250115__py3-none-any.whl

Sign up to get free protection for your applications and to get access to all the features.
Files changed (312) hide show
  1. flwr/cli/app.py +16 -2
  2. flwr/cli/build.py +181 -0
  3. flwr/cli/cli_user_auth_interceptor.py +90 -0
  4. flwr/cli/config_utils.py +343 -0
  5. flwr/cli/example.py +4 -1
  6. flwr/cli/install.py +253 -0
  7. flwr/cli/log.py +182 -0
  8. flwr/{server/superlink/state → cli/login}/__init__.py +4 -10
  9. flwr/cli/login/login.py +88 -0
  10. flwr/cli/ls.py +327 -0
  11. flwr/cli/new/__init__.py +1 -0
  12. flwr/cli/new/new.py +210 -66
  13. flwr/cli/new/templates/app/.gitignore.tpl +163 -0
  14. flwr/cli/new/templates/app/LICENSE.tpl +202 -0
  15. flwr/cli/new/templates/app/README.baseline.md.tpl +127 -0
  16. flwr/cli/new/templates/app/README.flowertune.md.tpl +66 -0
  17. flwr/cli/new/templates/app/README.md.tpl +16 -32
  18. flwr/cli/new/templates/app/code/__init__.baseline.py.tpl +1 -0
  19. flwr/cli/new/templates/app/code/__init__.py.tpl +1 -1
  20. flwr/cli/new/templates/app/code/client.baseline.py.tpl +58 -0
  21. flwr/cli/new/templates/app/code/client.huggingface.py.tpl +55 -0
  22. flwr/cli/new/templates/app/code/client.jax.py.tpl +50 -0
  23. flwr/cli/new/templates/app/code/client.mlx.py.tpl +73 -0
  24. flwr/cli/new/templates/app/code/client.numpy.py.tpl +7 -7
  25. flwr/cli/new/templates/app/code/client.pytorch.py.tpl +30 -21
  26. flwr/cli/new/templates/app/code/client.sklearn.py.tpl +63 -0
  27. flwr/cli/new/templates/app/code/client.tensorflow.py.tpl +57 -1
  28. flwr/cli/new/templates/app/code/dataset.baseline.py.tpl +36 -0
  29. flwr/cli/new/templates/app/code/flwr_tune/__init__.py +15 -0
  30. flwr/cli/new/templates/app/code/flwr_tune/client_app.py.tpl +126 -0
  31. flwr/cli/new/templates/app/code/flwr_tune/dataset.py.tpl +87 -0
  32. flwr/cli/new/templates/app/code/flwr_tune/models.py.tpl +78 -0
  33. flwr/cli/new/templates/app/code/flwr_tune/server_app.py.tpl +94 -0
  34. flwr/cli/new/templates/app/code/flwr_tune/strategy.py.tpl +83 -0
  35. flwr/cli/new/templates/app/code/model.baseline.py.tpl +80 -0
  36. flwr/cli/new/templates/app/code/server.baseline.py.tpl +46 -0
  37. flwr/cli/new/templates/app/code/server.huggingface.py.tpl +38 -0
  38. flwr/cli/new/templates/app/code/server.jax.py.tpl +26 -0
  39. flwr/cli/new/templates/app/code/server.mlx.py.tpl +31 -0
  40. flwr/cli/new/templates/app/code/server.numpy.py.tpl +22 -9
  41. flwr/cli/new/templates/app/code/server.pytorch.py.tpl +21 -18
  42. flwr/cli/new/templates/app/code/server.sklearn.py.tpl +36 -0
  43. flwr/cli/new/templates/app/code/server.tensorflow.py.tpl +29 -1
  44. flwr/cli/new/templates/app/code/strategy.baseline.py.tpl +1 -0
  45. flwr/cli/new/templates/app/code/task.huggingface.py.tpl +102 -0
  46. flwr/cli/new/templates/app/code/task.jax.py.tpl +57 -0
  47. flwr/cli/new/templates/app/code/task.mlx.py.tpl +102 -0
  48. flwr/cli/new/templates/app/code/task.numpy.py.tpl +7 -0
  49. flwr/cli/new/templates/app/code/task.pytorch.py.tpl +29 -24
  50. flwr/cli/new/templates/app/code/task.sklearn.py.tpl +67 -0
  51. flwr/cli/new/templates/app/code/task.tensorflow.py.tpl +53 -0
  52. flwr/cli/new/templates/app/code/utils.baseline.py.tpl +1 -0
  53. flwr/cli/new/templates/app/pyproject.baseline.toml.tpl +138 -0
  54. flwr/cli/new/templates/app/pyproject.flowertune.toml.tpl +68 -0
  55. flwr/cli/new/templates/app/pyproject.huggingface.toml.tpl +46 -0
  56. flwr/cli/new/templates/app/pyproject.jax.toml.tpl +35 -0
  57. flwr/cli/new/templates/app/pyproject.mlx.toml.tpl +39 -0
  58. flwr/cli/new/templates/app/pyproject.numpy.toml.tpl +25 -12
  59. flwr/cli/new/templates/app/pyproject.pytorch.toml.tpl +29 -14
  60. flwr/cli/new/templates/app/pyproject.sklearn.toml.tpl +35 -0
  61. flwr/cli/new/templates/app/pyproject.tensorflow.toml.tpl +29 -14
  62. flwr/cli/run/__init__.py +1 -0
  63. flwr/cli/run/run.py +212 -34
  64. flwr/cli/stop.py +130 -0
  65. flwr/cli/utils.py +240 -5
  66. flwr/client/__init__.py +3 -2
  67. flwr/client/app.py +432 -255
  68. flwr/client/client.py +1 -11
  69. flwr/client/client_app.py +74 -13
  70. flwr/client/clientapp/__init__.py +22 -0
  71. flwr/client/clientapp/app.py +259 -0
  72. flwr/client/clientapp/clientappio_servicer.py +244 -0
  73. flwr/client/clientapp/utils.py +115 -0
  74. flwr/client/dpfedavg_numpy_client.py +7 -8
  75. flwr/client/grpc_adapter_client/__init__.py +15 -0
  76. flwr/client/grpc_adapter_client/connection.py +98 -0
  77. flwr/client/grpc_client/connection.py +21 -7
  78. flwr/client/grpc_rere_client/__init__.py +1 -1
  79. flwr/client/grpc_rere_client/client_interceptor.py +176 -0
  80. flwr/client/grpc_rere_client/connection.py +163 -56
  81. flwr/client/grpc_rere_client/grpc_adapter.py +167 -0
  82. flwr/client/heartbeat.py +74 -0
  83. flwr/client/message_handler/__init__.py +1 -1
  84. flwr/client/message_handler/message_handler.py +10 -11
  85. flwr/client/mod/__init__.py +5 -5
  86. flwr/client/mod/centraldp_mods.py +4 -2
  87. flwr/client/mod/comms_mods.py +5 -4
  88. flwr/client/mod/localdp_mod.py +10 -5
  89. flwr/client/mod/secure_aggregation/__init__.py +1 -1
  90. flwr/client/mod/secure_aggregation/secaggplus_mod.py +26 -26
  91. flwr/client/mod/utils.py +2 -4
  92. flwr/client/nodestate/__init__.py +26 -0
  93. flwr/client/nodestate/in_memory_nodestate.py +38 -0
  94. flwr/client/nodestate/nodestate.py +31 -0
  95. flwr/client/nodestate/nodestate_factory.py +38 -0
  96. flwr/client/numpy_client.py +8 -31
  97. flwr/client/rest_client/__init__.py +1 -1
  98. flwr/client/rest_client/connection.py +199 -176
  99. flwr/client/run_info_store.py +112 -0
  100. flwr/client/supernode/__init__.py +24 -0
  101. flwr/client/supernode/app.py +321 -0
  102. flwr/client/typing.py +1 -0
  103. flwr/common/__init__.py +17 -11
  104. flwr/common/address.py +47 -3
  105. flwr/common/args.py +153 -0
  106. flwr/common/auth_plugin/__init__.py +24 -0
  107. flwr/common/auth_plugin/auth_plugin.py +121 -0
  108. flwr/common/config.py +243 -0
  109. flwr/common/constant.py +135 -1
  110. flwr/common/context.py +32 -2
  111. flwr/common/date.py +22 -4
  112. flwr/common/differential_privacy.py +2 -2
  113. flwr/common/dp.py +2 -4
  114. flwr/common/exit_handlers.py +3 -3
  115. flwr/common/grpc.py +164 -5
  116. flwr/common/logger.py +230 -12
  117. flwr/common/message.py +191 -106
  118. flwr/common/object_ref.py +179 -44
  119. flwr/common/pyproject.py +1 -0
  120. flwr/common/record/__init__.py +2 -1
  121. flwr/common/record/configsrecord.py +58 -18
  122. flwr/common/record/metricsrecord.py +57 -17
  123. flwr/common/record/parametersrecord.py +88 -20
  124. flwr/common/record/recordset.py +153 -30
  125. flwr/common/record/typeddict.py +30 -55
  126. flwr/common/recordset_compat.py +31 -12
  127. flwr/common/retry_invoker.py +123 -30
  128. flwr/common/secure_aggregation/__init__.py +1 -1
  129. flwr/common/secure_aggregation/crypto/__init__.py +1 -1
  130. flwr/common/secure_aggregation/crypto/shamir.py +11 -11
  131. flwr/common/secure_aggregation/crypto/symmetric_encryption.py +68 -4
  132. flwr/common/secure_aggregation/ndarrays_arithmetic.py +17 -17
  133. flwr/common/secure_aggregation/quantization.py +8 -8
  134. flwr/common/secure_aggregation/secaggplus_constants.py +1 -1
  135. flwr/common/secure_aggregation/secaggplus_utils.py +10 -12
  136. flwr/common/serde.py +304 -23
  137. flwr/common/telemetry.py +65 -29
  138. flwr/common/typing.py +120 -19
  139. flwr/common/version.py +17 -3
  140. flwr/proto/clientappio_pb2.py +45 -0
  141. flwr/proto/clientappio_pb2.pyi +132 -0
  142. flwr/proto/clientappio_pb2_grpc.py +135 -0
  143. flwr/proto/clientappio_pb2_grpc.pyi +53 -0
  144. flwr/proto/exec_pb2.py +62 -0
  145. flwr/proto/exec_pb2.pyi +212 -0
  146. flwr/proto/exec_pb2_grpc.py +237 -0
  147. flwr/proto/exec_pb2_grpc.pyi +93 -0
  148. flwr/proto/fab_pb2.py +31 -0
  149. flwr/proto/fab_pb2.pyi +65 -0
  150. flwr/proto/fab_pb2_grpc.py +4 -0
  151. flwr/proto/fab_pb2_grpc.pyi +4 -0
  152. flwr/proto/fleet_pb2.py +42 -23
  153. flwr/proto/fleet_pb2.pyi +123 -1
  154. flwr/proto/fleet_pb2_grpc.py +170 -0
  155. flwr/proto/fleet_pb2_grpc.pyi +61 -0
  156. flwr/proto/grpcadapter_pb2.py +32 -0
  157. flwr/proto/grpcadapter_pb2.pyi +43 -0
  158. flwr/proto/grpcadapter_pb2_grpc.py +66 -0
  159. flwr/proto/grpcadapter_pb2_grpc.pyi +24 -0
  160. flwr/proto/log_pb2.py +29 -0
  161. flwr/proto/log_pb2.pyi +39 -0
  162. flwr/proto/log_pb2_grpc.py +4 -0
  163. flwr/proto/log_pb2_grpc.pyi +4 -0
  164. flwr/proto/message_pb2.py +41 -0
  165. flwr/proto/message_pb2.pyi +128 -0
  166. flwr/proto/message_pb2_grpc.py +4 -0
  167. flwr/proto/message_pb2_grpc.pyi +4 -0
  168. flwr/proto/node_pb2.py +2 -2
  169. flwr/proto/node_pb2.pyi +1 -4
  170. flwr/proto/recordset_pb2.py +35 -33
  171. flwr/proto/recordset_pb2.pyi +40 -14
  172. flwr/proto/run_pb2.py +64 -0
  173. flwr/proto/run_pb2.pyi +268 -0
  174. flwr/proto/run_pb2_grpc.py +4 -0
  175. flwr/proto/run_pb2_grpc.pyi +4 -0
  176. flwr/proto/serverappio_pb2.py +52 -0
  177. flwr/proto/{driver_pb2.pyi → serverappio_pb2.pyi} +62 -20
  178. flwr/proto/serverappio_pb2_grpc.py +410 -0
  179. flwr/proto/serverappio_pb2_grpc.pyi +160 -0
  180. flwr/proto/simulationio_pb2.py +38 -0
  181. flwr/proto/simulationio_pb2.pyi +65 -0
  182. flwr/proto/simulationio_pb2_grpc.py +239 -0
  183. flwr/proto/simulationio_pb2_grpc.pyi +94 -0
  184. flwr/proto/task_pb2.py +7 -8
  185. flwr/proto/task_pb2.pyi +8 -5
  186. flwr/proto/transport_pb2.py +8 -8
  187. flwr/proto/transport_pb2.pyi +9 -6
  188. flwr/server/__init__.py +2 -10
  189. flwr/server/app.py +579 -402
  190. flwr/server/client_manager.py +8 -6
  191. flwr/server/compat/app.py +6 -62
  192. flwr/server/compat/app_utils.py +14 -9
  193. flwr/server/compat/driver_client_proxy.py +25 -59
  194. flwr/server/compat/legacy_context.py +5 -4
  195. flwr/server/driver/__init__.py +2 -0
  196. flwr/server/driver/driver.py +36 -131
  197. flwr/server/driver/grpc_driver.py +220 -81
  198. flwr/server/driver/inmemory_driver.py +183 -0
  199. flwr/server/history.py +28 -29
  200. flwr/server/run_serverapp.py +15 -126
  201. flwr/server/server.py +50 -44
  202. flwr/server/server_app.py +59 -10
  203. flwr/server/serverapp/__init__.py +22 -0
  204. flwr/server/serverapp/app.py +256 -0
  205. flwr/server/serverapp_components.py +52 -0
  206. flwr/server/strategy/__init__.py +2 -2
  207. flwr/server/strategy/aggregate.py +37 -23
  208. flwr/server/strategy/bulyan.py +9 -9
  209. flwr/server/strategy/dp_adaptive_clipping.py +25 -25
  210. flwr/server/strategy/dp_fixed_clipping.py +23 -22
  211. flwr/server/strategy/dpfedavg_adaptive.py +8 -8
  212. flwr/server/strategy/dpfedavg_fixed.py +13 -12
  213. flwr/server/strategy/fault_tolerant_fedavg.py +11 -11
  214. flwr/server/strategy/fedadagrad.py +9 -9
  215. flwr/server/strategy/fedadam.py +20 -10
  216. flwr/server/strategy/fedavg.py +16 -16
  217. flwr/server/strategy/fedavg_android.py +17 -17
  218. flwr/server/strategy/fedavgm.py +9 -9
  219. flwr/server/strategy/fedmedian.py +5 -5
  220. flwr/server/strategy/fedopt.py +6 -6
  221. flwr/server/strategy/fedprox.py +7 -7
  222. flwr/server/strategy/fedtrimmedavg.py +8 -8
  223. flwr/server/strategy/fedxgb_bagging.py +12 -12
  224. flwr/server/strategy/fedxgb_cyclic.py +10 -10
  225. flwr/server/strategy/fedxgb_nn_avg.py +6 -6
  226. flwr/server/strategy/fedyogi.py +9 -9
  227. flwr/server/strategy/krum.py +9 -9
  228. flwr/server/strategy/qfedavg.py +16 -16
  229. flwr/server/strategy/strategy.py +10 -10
  230. flwr/server/superlink/driver/__init__.py +2 -2
  231. flwr/server/superlink/driver/serverappio_grpc.py +61 -0
  232. flwr/server/superlink/driver/serverappio_servicer.py +361 -0
  233. flwr/server/superlink/ffs/__init__.py +24 -0
  234. flwr/server/superlink/ffs/disk_ffs.py +108 -0
  235. flwr/server/superlink/ffs/ffs.py +79 -0
  236. flwr/server/superlink/ffs/ffs_factory.py +47 -0
  237. flwr/server/superlink/fleet/__init__.py +1 -1
  238. flwr/server/superlink/fleet/grpc_adapter/__init__.py +15 -0
  239. flwr/server/superlink/fleet/grpc_adapter/grpc_adapter_servicer.py +162 -0
  240. flwr/server/superlink/fleet/grpc_bidi/__init__.py +1 -1
  241. flwr/server/superlink/fleet/grpc_bidi/flower_service_servicer.py +4 -2
  242. flwr/server/superlink/fleet/grpc_bidi/grpc_bridge.py +3 -2
  243. flwr/server/superlink/fleet/grpc_bidi/grpc_client_proxy.py +1 -1
  244. flwr/server/superlink/fleet/grpc_bidi/grpc_server.py +5 -154
  245. flwr/server/superlink/fleet/grpc_rere/__init__.py +1 -1
  246. flwr/server/superlink/fleet/grpc_rere/fleet_servicer.py +120 -13
  247. flwr/server/superlink/fleet/grpc_rere/server_interceptor.py +228 -0
  248. flwr/server/superlink/fleet/message_handler/__init__.py +1 -1
  249. flwr/server/superlink/fleet/message_handler/message_handler.py +156 -13
  250. flwr/server/superlink/fleet/rest_rere/__init__.py +1 -1
  251. flwr/server/superlink/fleet/rest_rere/rest_api.py +119 -81
  252. flwr/server/superlink/fleet/vce/__init__.py +1 -0
  253. flwr/server/superlink/fleet/vce/backend/__init__.py +4 -4
  254. flwr/server/superlink/fleet/vce/backend/backend.py +8 -9
  255. flwr/server/superlink/fleet/vce/backend/raybackend.py +87 -68
  256. flwr/server/superlink/fleet/vce/vce_api.py +208 -146
  257. flwr/server/superlink/linkstate/__init__.py +28 -0
  258. flwr/server/superlink/linkstate/in_memory_linkstate.py +569 -0
  259. flwr/server/superlink/linkstate/linkstate.py +376 -0
  260. flwr/server/superlink/{state/state_factory.py → linkstate/linkstate_factory.py} +19 -10
  261. flwr/server/superlink/linkstate/sqlite_linkstate.py +1196 -0
  262. flwr/server/superlink/linkstate/utils.py +399 -0
  263. flwr/server/superlink/simulation/__init__.py +15 -0
  264. flwr/server/superlink/simulation/simulationio_grpc.py +65 -0
  265. flwr/server/superlink/simulation/simulationio_servicer.py +186 -0
  266. flwr/server/superlink/utils.py +65 -0
  267. flwr/server/typing.py +2 -0
  268. flwr/server/utils/__init__.py +1 -1
  269. flwr/server/utils/tensorboard.py +5 -5
  270. flwr/server/utils/validator.py +40 -45
  271. flwr/server/workflow/default_workflows.py +70 -26
  272. flwr/server/workflow/secure_aggregation/secagg_workflow.py +1 -0
  273. flwr/server/workflow/secure_aggregation/secaggplus_workflow.py +40 -27
  274. flwr/simulation/__init__.py +12 -5
  275. flwr/simulation/app.py +247 -315
  276. flwr/simulation/legacy_app.py +404 -0
  277. flwr/simulation/ray_transport/__init__.py +1 -1
  278. flwr/simulation/ray_transport/ray_actor.py +42 -67
  279. flwr/simulation/ray_transport/ray_client_proxy.py +37 -17
  280. flwr/simulation/ray_transport/utils.py +1 -0
  281. flwr/simulation/run_simulation.py +306 -163
  282. flwr/simulation/simulationio_connection.py +89 -0
  283. flwr/superexec/__init__.py +15 -0
  284. flwr/superexec/app.py +59 -0
  285. flwr/superexec/deployment.py +188 -0
  286. flwr/superexec/exec_grpc.py +80 -0
  287. flwr/superexec/exec_servicer.py +231 -0
  288. flwr/superexec/exec_user_auth_interceptor.py +101 -0
  289. flwr/superexec/executor.py +96 -0
  290. flwr/superexec/simulation.py +124 -0
  291. {flwr_nightly-1.8.0.dev20240315.dist-info → flwr_nightly-1.15.0.dev20250115.dist-info}/METADATA +33 -26
  292. flwr_nightly-1.15.0.dev20250115.dist-info/RECORD +328 -0
  293. flwr_nightly-1.15.0.dev20250115.dist-info/entry_points.txt +12 -0
  294. flwr/cli/flower_toml.py +0 -140
  295. flwr/cli/new/templates/app/flower.toml.tpl +0 -13
  296. flwr/cli/new/templates/app/requirements.numpy.txt.tpl +0 -2
  297. flwr/cli/new/templates/app/requirements.pytorch.txt.tpl +0 -4
  298. flwr/cli/new/templates/app/requirements.tensorflow.txt.tpl +0 -4
  299. flwr/client/node_state.py +0 -48
  300. flwr/client/node_state_tests.py +0 -65
  301. flwr/proto/driver_pb2.py +0 -44
  302. flwr/proto/driver_pb2_grpc.py +0 -169
  303. flwr/proto/driver_pb2_grpc.pyi +0 -66
  304. flwr/server/superlink/driver/driver_grpc.py +0 -54
  305. flwr/server/superlink/driver/driver_servicer.py +0 -129
  306. flwr/server/superlink/state/in_memory_state.py +0 -230
  307. flwr/server/superlink/state/sqlite_state.py +0 -630
  308. flwr/server/superlink/state/state.py +0 -154
  309. flwr_nightly-1.8.0.dev20240315.dist-info/RECORD +0 -211
  310. flwr_nightly-1.8.0.dev20240315.dist-info/entry_points.txt +0 -9
  311. {flwr_nightly-1.8.0.dev20240315.dist-info → flwr_nightly-1.15.0.dev20250115.dist-info}/LICENSE +0 -0
  312. {flwr_nightly-1.8.0.dev20240315.dist-info → flwr_nightly-1.15.0.dev20250115.dist-info}/WHEEL +0 -0
@@ -0,0 +1,74 @@
1
+ # Copyright 2024 Flower Labs GmbH. All Rights Reserved.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+ # ==============================================================================
15
+ """Heartbeat utility functions."""
16
+
17
+
18
+ import threading
19
+ from typing import Callable
20
+
21
+ import grpc
22
+
23
+ from flwr.common.constant import PING_CALL_TIMEOUT
24
+ from flwr.common.retry_invoker import RetryInvoker, RetryState, exponential
25
+
26
+
27
+ def _ping_loop(ping_fn: Callable[[], None], stop_event: threading.Event) -> None:
28
+ def wait_fn(wait_time: float) -> None:
29
+ if not stop_event.is_set():
30
+ stop_event.wait(wait_time)
31
+
32
+ def on_backoff(state: RetryState) -> None:
33
+ err = state.exception
34
+ if not isinstance(err, grpc.RpcError):
35
+ return
36
+ status_code = err.code()
37
+ # If ping call timeout is triggered
38
+ if status_code == grpc.StatusCode.DEADLINE_EXCEEDED:
39
+ # Avoid long wait time.
40
+ if state.actual_wait is None:
41
+ return
42
+ state.actual_wait = max(state.actual_wait - PING_CALL_TIMEOUT, 0.0)
43
+
44
+ def wrapped_ping() -> None:
45
+ if not stop_event.is_set():
46
+ ping_fn()
47
+
48
+ retrier = RetryInvoker(
49
+ exponential,
50
+ grpc.RpcError,
51
+ max_tries=None,
52
+ max_time=None,
53
+ on_backoff=on_backoff,
54
+ wait_function=wait_fn,
55
+ )
56
+ while not stop_event.is_set():
57
+ retrier.invoke(wrapped_ping)
58
+
59
+
60
+ def start_ping_loop(
61
+ ping_fn: Callable[[], None], stop_event: threading.Event
62
+ ) -> threading.Thread:
63
+ """Start a ping loop in a separate thread.
64
+
65
+ This function initializes a new thread that runs a ping loop, allowing for
66
+ asynchronous ping operations. The loop can be terminated through the provided stop
67
+ event.
68
+ """
69
+ thread = threading.Thread(
70
+ target=_ping_loop, args=(ping_fn, stop_event), daemon=True
71
+ )
72
+ thread.start()
73
+
74
+ return thread
@@ -1,4 +1,4 @@
1
- # Copyright 2020 Flower Labs GmbH. All Rights Reserved.
1
+ # Copyright 2022 Flower Labs GmbH. All Rights Reserved.
2
2
  #
3
3
  # Licensed under the Apache License, Version 2.0 (the "License");
4
4
  # you may not use this file except in compliance with the License.
@@ -1,4 +1,4 @@
1
- # Copyright 2020 Flower Labs GmbH. All Rights Reserved.
1
+ # Copyright 2022 Flower Labs GmbH. All Rights Reserved.
2
2
  #
3
3
  # Licensed under the Apache License, Version 2.0 (the "License");
4
4
  # you may not use this file except in compliance with the License.
@@ -16,7 +16,7 @@
16
16
 
17
17
 
18
18
  from logging import WARN
19
- from typing import Optional, Tuple, cast
19
+ from typing import Optional, cast
20
20
 
21
21
  from flwr.client.client import (
22
22
  maybe_call_evaluate,
@@ -25,7 +25,7 @@ from flwr.client.client import (
25
25
  maybe_call_get_properties,
26
26
  )
27
27
  from flwr.client.numpy_client import NumPyClient
28
- from flwr.client.typing import ClientFn
28
+ from flwr.client.typing import ClientFnExt
29
29
  from flwr.common import ConfigsRecord, Context, Message, Metadata, RecordSet, log
30
30
  from flwr.common.constant import MessageType, MessageTypeLegacy
31
31
  from flwr.common.recordset_compat import (
@@ -53,7 +53,7 @@ class UnknownServerMessage(Exception):
53
53
  """Exception indicating that the received message is unknown."""
54
54
 
55
55
 
56
- def handle_control_message(message: Message) -> Tuple[Optional[Message], int]:
56
+ def handle_control_message(message: Message) -> tuple[Optional[Message], int]:
57
57
  """Handle control part of the incoming message.
58
58
 
59
59
  Parameters
@@ -81,7 +81,7 @@ def handle_control_message(message: Message) -> Tuple[Optional[Message], int]:
81
81
  reason = cast(int, disconnect_msg.disconnect_res.reason)
82
82
  recordset = RecordSet()
83
83
  recordset.configs_records["config"] = ConfigsRecord({"reason": reason})
84
- out_message = message.create_reply(recordset, ttl="")
84
+ out_message = message.create_reply(recordset)
85
85
  # Return TaskRes and sleep duration
86
86
  return out_message, sleep_duration
87
87
 
@@ -90,10 +90,10 @@ def handle_control_message(message: Message) -> Tuple[Optional[Message], int]:
90
90
 
91
91
 
92
92
  def handle_legacy_message_from_msgtype(
93
- client_fn: ClientFn, message: Message, context: Context
93
+ client_fn: ClientFnExt, message: Message, context: Context
94
94
  ) -> Message:
95
95
  """Handle legacy message in the inner most mod."""
96
- client = client_fn(str(message.metadata.partition_id))
96
+ client = client_fn(context)
97
97
 
98
98
  # Check if NumPyClient is returend
99
99
  if isinstance(client, NumPyClient):
@@ -105,8 +105,6 @@ def handle_legacy_message_from_msgtype(
105
105
  "Please use `NumPyClient.to_client()` method to convert it to `Client`.",
106
106
  )
107
107
 
108
- client.set_context(context)
109
-
110
108
  message_type = message.metadata.message_type
111
109
 
112
110
  # Handle GetPropertiesIns
@@ -143,12 +141,12 @@ def handle_legacy_message_from_msgtype(
143
141
  raise ValueError(f"Invalid message type: {message_type}")
144
142
 
145
143
  # Return Message
146
- return message.create_reply(out_recordset, ttl="")
144
+ return message.create_reply(out_recordset)
147
145
 
148
146
 
149
147
  def _reconnect(
150
148
  reconnect_msg: ServerMessage.ReconnectIns,
151
- ) -> Tuple[ClientMessage, int]:
149
+ ) -> tuple[ClientMessage, int]:
152
150
  # Determine the reason for sending DisconnectRes message
153
151
  reason = Reason.ACK
154
152
  sleep_duration = None
@@ -172,6 +170,7 @@ def validate_out_message(out_message: Message, in_message_metadata: Metadata) ->
172
170
  and out_meta.reply_to_message == in_meta.message_id
173
171
  and out_meta.group_id == in_meta.group_id
174
172
  and out_meta.message_type == in_meta.message_type
173
+ and out_meta.created_at > in_meta.created_at
175
174
  ):
176
175
  return True
177
176
  return False
@@ -1,4 +1,4 @@
1
- # Copyright 2023 Flower Labs GmbH. All Rights Reserved.
1
+ # Copyright 2024 Flower Labs GmbH. All Rights Reserved.
2
2
  #
3
3
  # Licensed under the Apache License, Version 2.0 (the "License");
4
4
  # you may not use this file except in compliance with the License.
@@ -12,7 +12,7 @@
12
12
  # See the License for the specific language governing permissions and
13
13
  # limitations under the License.
14
14
  # ==============================================================================
15
- """Mods."""
15
+ """Flower Built-in Mods."""
16
16
 
17
17
 
18
18
  from .centraldp_mods import adaptiveclipping_mod, fixedclipping_mod
@@ -22,12 +22,12 @@ from .secure_aggregation import secagg_mod, secaggplus_mod
22
22
  from .utils import make_ffn
23
23
 
24
24
  __all__ = [
25
+ "LocalDpMod",
25
26
  "adaptiveclipping_mod",
26
27
  "fixedclipping_mod",
27
- "LocalDpMod",
28
28
  "make_ffn",
29
- "secagg_mod",
30
- "secaggplus_mod",
31
29
  "message_size_mod",
32
30
  "parameters_size_mod",
31
+ "secagg_mod",
32
+ "secaggplus_mod",
33
33
  ]
@@ -82,7 +82,9 @@ def fixedclipping_mod(
82
82
  clipping_norm,
83
83
  )
84
84
 
85
- log(INFO, "fixedclipping_mod: parameters are clipped by value: %s.", clipping_norm)
85
+ log(
86
+ INFO, "fixedclipping_mod: parameters are clipped by value: %.4f.", clipping_norm
87
+ )
86
88
 
87
89
  fit_res.parameters = ndarrays_to_parameters(client_to_server_params)
88
90
  out_msg.content = compat.fitres_to_recordset(fit_res, keep_input=True)
@@ -146,7 +148,7 @@ def adaptiveclipping_mod(
146
148
  )
147
149
  log(
148
150
  INFO,
149
- "adaptiveclipping_mod: parameters are clipped by value: %s.",
151
+ "adaptiveclipping_mod: parameters are clipped by value: %.4f.",
150
152
  clipping_norm,
151
153
  )
152
154
 
@@ -14,6 +14,7 @@
14
14
  # ==============================================================================
15
15
  """Mods that report statistics about message communication."""
16
16
 
17
+
17
18
  from logging import INFO
18
19
 
19
20
  import numpy as np
@@ -29,7 +30,7 @@ def message_size_mod(
29
30
  ) -> Message:
30
31
  """Message size mod.
31
32
 
32
- This mod logs the size in Bytes of the message being transmited.
33
+ This mod logs the size in bytes of the message being transmited.
33
34
  """
34
35
  message_size_in_bytes = 0
35
36
 
@@ -42,7 +43,7 @@ def message_size_mod(
42
43
  for m_record in msg.content.metrics_records.values():
43
44
  message_size_in_bytes += m_record.count_bytes()
44
45
 
45
- log(INFO, "Message size: %i Bytes", message_size_in_bytes)
46
+ log(INFO, "Message size: %i bytes", message_size_in_bytes)
46
47
 
47
48
  return call_next(msg, ctxt)
48
49
 
@@ -53,7 +54,7 @@ def parameters_size_mod(
53
54
  """Parameters size mod.
54
55
 
55
56
  This mod logs the number of parameters transmitted in the message as well as their
56
- size in Bytes.
57
+ size in bytes.
57
58
  """
58
59
  model_size_stats = {}
59
60
  parameters_size_in_bytes = 0
@@ -74,6 +75,6 @@ def parameters_size_mod(
74
75
  if model_size_stats:
75
76
  log(INFO, model_size_stats)
76
77
 
77
- log(INFO, "Total parameters transmited: %i Bytes", parameters_size_in_bytes)
78
+ log(INFO, "Total parameters transmitted: %i bytes", parameters_size_in_bytes)
78
79
 
79
80
  return call_next(msg, ctxt)
@@ -128,20 +128,25 @@ class LocalDpMod:
128
128
  self.clipping_norm,
129
129
  )
130
130
  log(
131
- INFO, "LocalDpMod: parameters are clipped by value: %s.", self.clipping_norm
131
+ INFO,
132
+ "LocalDpMod: parameters are clipped by value: %.4f.",
133
+ self.clipping_norm,
132
134
  )
133
135
 
134
136
  fit_res.parameters = ndarrays_to_parameters(client_to_server_params)
135
137
 
136
138
  # Add noise to model params
137
- add_localdp_gaussian_noise_to_params(
139
+ fit_res.parameters = add_localdp_gaussian_noise_to_params(
138
140
  fit_res.parameters, self.sensitivity, self.epsilon, self.delta
139
141
  )
142
+
143
+ noise_value_sd = (
144
+ self.sensitivity * np.sqrt(2 * np.log(1.25 / self.delta)) / self.epsilon
145
+ )
140
146
  log(
141
147
  INFO,
142
- "LocalDpMod: local DP noise with "
143
- "standard deviation: %s added to parameters.",
144
- self.sensitivity * np.sqrt(2 * np.log(1.25 / self.delta)) / self.epsilon,
148
+ "LocalDpMod: local DP noise with %.4f stedv added to parameters",
149
+ noise_value_sd,
145
150
  )
146
151
 
147
152
  out_msg.content = compat.fitres_to_recordset(fit_res, keep_input=True)
@@ -1,4 +1,4 @@
1
- # Copyright 2023 Flower Labs GmbH. All Rights Reserved.
1
+ # Copyright 2024 Flower Labs GmbH. All Rights Reserved.
2
2
  #
3
3
  # Licensed under the Apache License, Version 2.0 (the "License");
4
4
  # you may not use this file except in compliance with the License.
@@ -18,7 +18,7 @@
18
18
  import os
19
19
  from dataclasses import dataclass, field
20
20
  from logging import DEBUG, WARNING
21
- from typing import Any, Dict, List, Tuple, cast
21
+ from typing import Any, cast
22
22
 
23
23
  from flwr.client.typing import ClientAppCallable
24
24
  from flwr.common import (
@@ -91,11 +91,11 @@ class SecAggPlusState:
91
91
  # Random seed for generating the private mask
92
92
  rd_seed: bytes = b""
93
93
 
94
- rd_seed_share_dict: Dict[int, bytes] = field(default_factory=dict)
95
- sk1_share_dict: Dict[int, bytes] = field(default_factory=dict)
94
+ rd_seed_share_dict: dict[int, bytes] = field(default_factory=dict)
95
+ sk1_share_dict: dict[int, bytes] = field(default_factory=dict)
96
96
  # The dict of the shared secrets from sk2
97
- ss2_dict: Dict[int, bytes] = field(default_factory=dict)
98
- public_keys_dict: Dict[int, Tuple[bytes, bytes]] = field(default_factory=dict)
97
+ ss2_dict: dict[int, bytes] = field(default_factory=dict)
98
+ public_keys_dict: dict[int, tuple[bytes, bytes]] = field(default_factory=dict)
99
99
 
100
100
  def __init__(self, **kwargs: ConfigsRecordValues) -> None:
101
101
  for k, v in kwargs.items():
@@ -104,8 +104,8 @@ class SecAggPlusState:
104
104
  new_v: Any = v
105
105
  if k.endswith(":K"):
106
106
  k = k[:-2]
107
- keys = cast(List[int], v)
108
- values = cast(List[bytes], kwargs[f"{k}:V"])
107
+ keys = cast(list[int], v)
108
+ values = cast(list[bytes], kwargs[f"{k}:V"])
109
109
  if len(values) > len(keys):
110
110
  updated_values = [
111
111
  tuple(values[i : i + 2]) for i in range(0, len(values), 2)
@@ -115,17 +115,17 @@ class SecAggPlusState:
115
115
  new_v = dict(zip(keys, values))
116
116
  self.__setattr__(k, new_v)
117
117
 
118
- def to_dict(self) -> Dict[str, ConfigsRecordValues]:
118
+ def to_dict(self) -> dict[str, ConfigsRecordValues]:
119
119
  """Convert the state to a dictionary."""
120
120
  ret = vars(self)
121
121
  for k in list(ret.keys()):
122
122
  if isinstance(ret[k], dict):
123
123
  # Replace dict with two lists
124
- v = cast(Dict[str, Any], ret.pop(k))
124
+ v = cast(dict[str, Any], ret.pop(k))
125
125
  ret[f"{k}:K"] = list(v.keys())
126
126
  if k == "public_keys_dict":
127
- v_list: List[bytes] = []
128
- for b1_b2 in cast(List[Tuple[bytes, bytes]], v.values()):
127
+ v_list: list[bytes] = []
128
+ for b1_b2 in cast(list[tuple[bytes, bytes]], v.values()):
129
129
  v_list.extend(b1_b2)
130
130
  ret[f"{k}:V"] = v_list
131
131
  else:
@@ -187,7 +187,7 @@ def secaggplus_mod(
187
187
 
188
188
  # Return message
189
189
  out_content.configs_records[RECORD_KEY_CONFIGS] = ConfigsRecord(res, False)
190
- return msg.create_reply(out_content, ttl="")
190
+ return msg.create_reply(out_content)
191
191
 
192
192
 
193
193
  def check_stage(current_stage: str, configs: ConfigsRecord) -> None:
@@ -276,7 +276,7 @@ def check_configs(stage: str, configs: ConfigsRecord) -> None:
276
276
  )
277
277
  if not isinstance(configs[key], list) or any(
278
278
  elm
279
- for elm in cast(List[Any], configs[key])
279
+ for elm in cast(list[Any], configs[key])
280
280
  # pylint: disable-next=unidiomatic-typecheck
281
281
  if type(elm) is not expected_type
282
282
  ):
@@ -299,7 +299,7 @@ def check_configs(stage: str, configs: ConfigsRecord) -> None:
299
299
  )
300
300
  if not isinstance(configs[key], list) or any(
301
301
  elm
302
- for elm in cast(List[Any], configs[key])
302
+ for elm in cast(list[Any], configs[key])
303
303
  # pylint: disable-next=unidiomatic-typecheck
304
304
  if type(elm) is not expected_type
305
305
  ):
@@ -314,7 +314,7 @@ def check_configs(stage: str, configs: ConfigsRecord) -> None:
314
314
 
315
315
  def _setup(
316
316
  state: SecAggPlusState, configs: ConfigsRecord
317
- ) -> Dict[str, ConfigsRecordValues]:
317
+ ) -> dict[str, ConfigsRecordValues]:
318
318
  # Assigning parameter values to object fields
319
319
  sec_agg_param_dict = configs
320
320
  state.sample_num = cast(int, sec_agg_param_dict[Key.SAMPLE_NUMBER])
@@ -350,8 +350,8 @@ def _setup(
350
350
  # pylint: disable-next=too-many-locals
351
351
  def _share_keys(
352
352
  state: SecAggPlusState, configs: ConfigsRecord
353
- ) -> Dict[str, ConfigsRecordValues]:
354
- named_bytes_tuples = cast(Dict[str, Tuple[bytes, bytes]], configs)
353
+ ) -> dict[str, ConfigsRecordValues]:
354
+ named_bytes_tuples = cast(dict[str, tuple[bytes, bytes]], configs)
355
355
  key_dict = {int(sid): (pk1, pk2) for sid, (pk1, pk2) in named_bytes_tuples.items()}
356
356
  log(DEBUG, "Node %d: starting stage 1...", state.nid)
357
357
  state.public_keys_dict = key_dict
@@ -361,7 +361,7 @@ def _share_keys(
361
361
  raise ValueError("Available neighbours number smaller than threshold")
362
362
 
363
363
  # Check if all public keys are unique
364
- pk_list: List[bytes] = []
364
+ pk_list: list[bytes] = []
365
365
  for pk1, pk2 in state.public_keys_dict.values():
366
366
  pk_list.append(pk1)
367
367
  pk_list.append(pk2)
@@ -415,11 +415,11 @@ def _collect_masked_vectors(
415
415
  configs: ConfigsRecord,
416
416
  num_examples: int,
417
417
  updated_parameters: Parameters,
418
- ) -> Dict[str, ConfigsRecordValues]:
418
+ ) -> dict[str, ConfigsRecordValues]:
419
419
  log(DEBUG, "Node %d: starting stage 2...", state.nid)
420
- available_clients: List[int] = []
421
- ciphertexts = cast(List[bytes], configs[Key.CIPHERTEXT_LIST])
422
- srcs = cast(List[int], configs[Key.SOURCE_LIST])
420
+ available_clients: list[int] = []
421
+ ciphertexts = cast(list[bytes], configs[Key.CIPHERTEXT_LIST])
422
+ srcs = cast(list[int], configs[Key.SOURCE_LIST])
423
423
  if len(ciphertexts) + 1 < state.threshold:
424
424
  raise ValueError("Not enough available neighbour clients.")
425
425
 
@@ -467,7 +467,7 @@ def _collect_masked_vectors(
467
467
 
468
468
  quantized_parameters = factor_combine(q_ratio, quantized_parameters)
469
469
 
470
- dimensions_list: List[Tuple[int, ...]] = [a.shape for a in quantized_parameters]
470
+ dimensions_list: list[tuple[int, ...]] = [a.shape for a in quantized_parameters]
471
471
 
472
472
  # Add private mask
473
473
  private_mask = pseudo_rand_gen(state.rd_seed, state.mod_range, dimensions_list)
@@ -499,11 +499,11 @@ def _collect_masked_vectors(
499
499
 
500
500
  def _unmask(
501
501
  state: SecAggPlusState, configs: ConfigsRecord
502
- ) -> Dict[str, ConfigsRecordValues]:
502
+ ) -> dict[str, ConfigsRecordValues]:
503
503
  log(DEBUG, "Node %d: starting stage 3...", state.nid)
504
504
 
505
- active_nids = cast(List[int], configs[Key.ACTIVE_NODE_ID_LIST])
506
- dead_nids = cast(List[int], configs[Key.DEAD_NODE_ID_LIST])
505
+ active_nids = cast(list[int], configs[Key.ACTIVE_NODE_ID_LIST])
506
+ dead_nids = cast(list[int], configs[Key.DEAD_NODE_ID_LIST])
507
507
  # Send private mask seed share for every avaliable client (including itself)
508
508
  # Send first private key share for building pairwise mask for every dropped client
509
509
  if len(active_nids) < state.threshold:
flwr/client/mod/utils.py CHANGED
@@ -1,4 +1,4 @@
1
- # Copyright 2023 Flower Labs GmbH. All Rights Reserved.
1
+ # Copyright 2024 Flower Labs GmbH. All Rights Reserved.
2
2
  #
3
3
  # Licensed under the Apache License, Version 2.0 (the "License");
4
4
  # you may not use this file except in compliance with the License.
@@ -15,13 +15,11 @@
15
15
  """Utility functions for mods."""
16
16
 
17
17
 
18
- from typing import List
19
-
20
18
  from flwr.client.typing import ClientAppCallable, Mod
21
19
  from flwr.common import Context, Message
22
20
 
23
21
 
24
- def make_ffn(ffn: ClientAppCallable, mods: List[Mod]) -> ClientAppCallable:
22
+ def make_ffn(ffn: ClientAppCallable, mods: list[Mod]) -> ClientAppCallable:
25
23
  """."""
26
24
 
27
25
  def wrap_ffn(_ffn: ClientAppCallable, _mod: Mod) -> ClientAppCallable:
@@ -0,0 +1,26 @@
1
+ # Copyright 2024 Flower Labs GmbH. All Rights Reserved.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+ # ==============================================================================
15
+ """Flower NodeState."""
16
+
17
+
18
+ from .in_memory_nodestate import InMemoryNodeState as InMemoryNodeState
19
+ from .nodestate import NodeState as NodeState
20
+ from .nodestate_factory import NodeStateFactory as NodeStateFactory
21
+
22
+ __all__ = [
23
+ "InMemoryNodeState",
24
+ "NodeState",
25
+ "NodeStateFactory",
26
+ ]
@@ -0,0 +1,38 @@
1
+ # Copyright 2024 Flower Labs GmbH. All Rights Reserved.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+ # ==============================================================================
15
+ """In-memory NodeState implementation."""
16
+
17
+
18
+ from typing import Optional
19
+
20
+ from flwr.client.nodestate.nodestate import NodeState
21
+
22
+
23
+ class InMemoryNodeState(NodeState):
24
+ """In-memory NodeState implementation."""
25
+
26
+ def __init__(self) -> None:
27
+ # Store node_id
28
+ self.node_id: Optional[int] = None
29
+
30
+ def set_node_id(self, node_id: Optional[int]) -> None:
31
+ """Set the node ID."""
32
+ self.node_id = node_id
33
+
34
+ def get_node_id(self) -> int:
35
+ """Get the node ID."""
36
+ if self.node_id is None:
37
+ raise ValueError("Node ID not set")
38
+ return self.node_id
@@ -0,0 +1,31 @@
1
+ # Copyright 2024 Flower Labs GmbH. All Rights Reserved.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+ # ==============================================================================
15
+ """Abstract base class NodeState."""
16
+
17
+
18
+ import abc
19
+ from typing import Optional
20
+
21
+
22
+ class NodeState(abc.ABC):
23
+ """Abstract NodeState."""
24
+
25
+ @abc.abstractmethod
26
+ def set_node_id(self, node_id: Optional[int]) -> None:
27
+ """Set the node ID."""
28
+
29
+ @abc.abstractmethod
30
+ def get_node_id(self) -> int:
31
+ """Get the node ID."""
@@ -0,0 +1,38 @@
1
+ # Copyright 2024 Flower Labs GmbH. All Rights Reserved.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+ # ==============================================================================
15
+ """Factory class that creates NodeState instances."""
16
+
17
+
18
+ import threading
19
+ from typing import Optional
20
+
21
+ from .in_memory_nodestate import InMemoryNodeState
22
+ from .nodestate import NodeState
23
+
24
+
25
+ class NodeStateFactory:
26
+ """Factory class that creates NodeState instances."""
27
+
28
+ def __init__(self) -> None:
29
+ self.state_instance: Optional[NodeState] = None
30
+ self.lock = threading.RLock()
31
+
32
+ def state(self) -> NodeState:
33
+ """Return a State instance and create it, if necessary."""
34
+ # Lock access to NodeStateFactory to prevent returning different instances
35
+ with self.lock:
36
+ if self.state_instance is None:
37
+ self.state_instance = InMemoryNodeState()
38
+ return self.state_instance