flwr 1.17.0__py3-none-any.whl → 1.19.0__py3-none-any.whl

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 (286) hide show
  1. flwr/__init__.py +1 -1
  2. flwr/app/__init__.py +15 -0
  3. flwr/app/error.py +68 -0
  4. flwr/app/metadata.py +223 -0
  5. flwr/cli/__init__.py +1 -1
  6. flwr/cli/app.py +21 -2
  7. flwr/cli/build.py +83 -58
  8. flwr/cli/cli_user_auth_interceptor.py +1 -1
  9. flwr/cli/config_utils.py +53 -17
  10. flwr/cli/example.py +1 -1
  11. flwr/cli/install.py +1 -1
  12. flwr/cli/log.py +4 -4
  13. flwr/cli/login/__init__.py +1 -1
  14. flwr/cli/login/login.py +15 -8
  15. flwr/cli/ls.py +16 -37
  16. flwr/cli/new/__init__.py +1 -1
  17. flwr/cli/new/new.py +4 -4
  18. flwr/cli/new/templates/__init__.py +1 -1
  19. flwr/cli/new/templates/app/__init__.py +1 -1
  20. flwr/cli/new/templates/app/code/__init__.py +1 -1
  21. flwr/cli/new/templates/app/code/client.baseline.py.tpl +1 -1
  22. flwr/cli/new/templates/app/code/flwr_tune/__init__.py +1 -1
  23. flwr/cli/new/templates/app/code/flwr_tune/client_app.py.tpl +4 -4
  24. flwr/cli/new/templates/app/code/model.baseline.py.tpl +1 -1
  25. flwr/cli/new/templates/app/code/server.baseline.py.tpl +2 -3
  26. flwr/cli/new/templates/app/code/task.sklearn.py.tpl +1 -1
  27. flwr/cli/new/templates/app/pyproject.baseline.toml.tpl +14 -17
  28. flwr/cli/new/templates/app/pyproject.flowertune.toml.tpl +4 -4
  29. flwr/cli/new/templates/app/pyproject.huggingface.toml.tpl +1 -1
  30. flwr/cli/new/templates/app/pyproject.jax.toml.tpl +1 -1
  31. flwr/cli/new/templates/app/pyproject.mlx.toml.tpl +1 -1
  32. flwr/cli/new/templates/app/pyproject.numpy.toml.tpl +1 -1
  33. flwr/cli/new/templates/app/pyproject.pytorch.toml.tpl +1 -1
  34. flwr/cli/new/templates/app/pyproject.sklearn.toml.tpl +1 -1
  35. flwr/cli/new/templates/app/pyproject.tensorflow.toml.tpl +1 -1
  36. flwr/cli/run/__init__.py +1 -1
  37. flwr/cli/run/run.py +11 -19
  38. flwr/cli/stop.py +3 -3
  39. flwr/cli/utils.py +42 -17
  40. flwr/client/__init__.py +3 -3
  41. flwr/client/client.py +1 -1
  42. flwr/client/client_app.py +140 -138
  43. flwr/client/clientapp/__init__.py +1 -8
  44. flwr/client/clientapp/utils.py +1 -1
  45. flwr/client/dpfedavg_numpy_client.py +1 -1
  46. flwr/client/grpc_adapter_client/__init__.py +1 -1
  47. flwr/client/grpc_adapter_client/connection.py +5 -5
  48. flwr/client/grpc_rere_client/__init__.py +1 -1
  49. flwr/client/grpc_rere_client/client_interceptor.py +1 -1
  50. flwr/client/grpc_rere_client/connection.py +131 -61
  51. flwr/client/grpc_rere_client/grpc_adapter.py +35 -7
  52. flwr/client/message_handler/__init__.py +1 -1
  53. flwr/client/message_handler/message_handler.py +2 -2
  54. flwr/client/mod/__init__.py +1 -1
  55. flwr/client/mod/centraldp_mods.py +1 -1
  56. flwr/client/mod/comms_mods.py +39 -20
  57. flwr/client/mod/localdp_mod.py +6 -6
  58. flwr/client/mod/secure_aggregation/__init__.py +1 -1
  59. flwr/client/mod/secure_aggregation/secagg_mod.py +1 -1
  60. flwr/client/mod/secure_aggregation/secaggplus_mod.py +1 -1
  61. flwr/client/mod/utils.py +1 -1
  62. flwr/client/numpy_client.py +1 -1
  63. flwr/client/rest_client/__init__.py +1 -1
  64. flwr/client/rest_client/connection.py +174 -68
  65. flwr/client/run_info_store.py +1 -1
  66. flwr/client/typing.py +1 -1
  67. flwr/clientapp/__init__.py +15 -0
  68. flwr/common/__init__.py +3 -3
  69. flwr/common/address.py +1 -1
  70. flwr/common/args.py +1 -1
  71. flwr/common/auth_plugin/__init__.py +3 -1
  72. flwr/common/auth_plugin/auth_plugin.py +30 -4
  73. flwr/common/config.py +1 -1
  74. flwr/common/constant.py +37 -8
  75. flwr/common/context.py +1 -1
  76. flwr/common/date.py +1 -1
  77. flwr/common/differential_privacy.py +1 -1
  78. flwr/common/differential_privacy_constants.py +1 -1
  79. flwr/common/dp.py +1 -1
  80. flwr/common/event_log_plugin/event_log_plugin.py +3 -3
  81. flwr/common/exit/exit.py +6 -6
  82. flwr/common/exit_handlers.py +31 -1
  83. flwr/common/grpc.py +1 -1
  84. flwr/common/heartbeat.py +165 -0
  85. flwr/common/inflatable.py +290 -0
  86. flwr/common/inflatable_grpc_utils.py +99 -0
  87. flwr/common/inflatable_rest_utils.py +99 -0
  88. flwr/common/inflatable_utils.py +341 -0
  89. flwr/common/logger.py +1 -1
  90. flwr/common/message.py +137 -252
  91. flwr/common/object_ref.py +1 -1
  92. flwr/common/parameter.py +1 -1
  93. flwr/common/pyproject.py +1 -1
  94. flwr/common/record/__init__.py +3 -2
  95. flwr/common/record/array.py +323 -0
  96. flwr/common/record/arrayrecord.py +121 -243
  97. flwr/common/record/configrecord.py +71 -16
  98. flwr/common/record/conversion_utils.py +2 -2
  99. flwr/common/record/metricrecord.py +71 -20
  100. flwr/common/record/recorddict.py +207 -90
  101. flwr/common/record/typeddict.py +1 -1
  102. flwr/common/recorddict_compat.py +2 -2
  103. flwr/common/retry_invoker.py +15 -11
  104. flwr/common/secure_aggregation/__init__.py +1 -1
  105. flwr/common/secure_aggregation/crypto/__init__.py +1 -1
  106. flwr/common/secure_aggregation/crypto/shamir.py +52 -30
  107. flwr/common/secure_aggregation/crypto/symmetric_encryption.py +1 -1
  108. flwr/common/secure_aggregation/ndarrays_arithmetic.py +1 -1
  109. flwr/common/secure_aggregation/quantization.py +1 -1
  110. flwr/common/secure_aggregation/secaggplus_constants.py +1 -1
  111. flwr/common/secure_aggregation/secaggplus_utils.py +1 -1
  112. flwr/common/serde.py +60 -184
  113. flwr/common/serde_utils.py +175 -0
  114. flwr/common/telemetry.py +2 -2
  115. flwr/common/typing.py +6 -4
  116. flwr/common/version.py +1 -1
  117. flwr/compat/__init__.py +15 -0
  118. flwr/compat/client/__init__.py +15 -0
  119. flwr/{client → compat/client}/app.py +71 -211
  120. flwr/{client → compat/client}/grpc_client/__init__.py +1 -1
  121. flwr/{client → compat/client}/grpc_client/connection.py +13 -13
  122. flwr/compat/common/__init__.py +15 -0
  123. flwr/compat/server/__init__.py +15 -0
  124. flwr/compat/server/app.py +174 -0
  125. flwr/compat/simulation/__init__.py +15 -0
  126. flwr/proto/__init__.py +1 -1
  127. flwr/proto/fleet_pb2.py +32 -27
  128. flwr/proto/fleet_pb2.pyi +49 -35
  129. flwr/proto/fleet_pb2_grpc.py +117 -13
  130. flwr/proto/fleet_pb2_grpc.pyi +47 -6
  131. flwr/proto/heartbeat_pb2.py +33 -0
  132. flwr/proto/heartbeat_pb2.pyi +66 -0
  133. flwr/proto/heartbeat_pb2_grpc.py +4 -0
  134. flwr/proto/heartbeat_pb2_grpc.pyi +4 -0
  135. flwr/proto/message_pb2.py +28 -11
  136. flwr/proto/message_pb2.pyi +125 -0
  137. flwr/proto/recorddict_pb2.py +16 -28
  138. flwr/proto/recorddict_pb2.pyi +46 -64
  139. flwr/proto/run_pb2.py +24 -32
  140. flwr/proto/run_pb2.pyi +4 -52
  141. flwr/proto/serverappio_pb2.py +32 -23
  142. flwr/proto/serverappio_pb2.pyi +45 -3
  143. flwr/proto/serverappio_pb2_grpc.py +138 -34
  144. flwr/proto/serverappio_pb2_grpc.pyi +54 -13
  145. flwr/proto/simulationio_pb2.py +12 -11
  146. flwr/proto/simulationio_pb2_grpc.py +35 -0
  147. flwr/proto/simulationio_pb2_grpc.pyi +14 -0
  148. flwr/server/__init__.py +2 -2
  149. flwr/server/app.py +69 -187
  150. flwr/server/client_manager.py +1 -1
  151. flwr/server/client_proxy.py +1 -1
  152. flwr/server/compat/__init__.py +1 -1
  153. flwr/server/compat/app.py +1 -1
  154. flwr/server/compat/app_utils.py +51 -29
  155. flwr/server/compat/legacy_context.py +1 -1
  156. flwr/server/criterion.py +1 -1
  157. flwr/server/fleet_event_log_interceptor.py +2 -2
  158. flwr/server/grid/grid.py +3 -3
  159. flwr/server/grid/grpc_grid.py +104 -34
  160. flwr/server/grid/inmemory_grid.py +5 -4
  161. flwr/server/history.py +1 -1
  162. flwr/server/run_serverapp.py +1 -1
  163. flwr/server/server.py +1 -1
  164. flwr/server/server_app.py +65 -58
  165. flwr/server/server_config.py +1 -1
  166. flwr/server/serverapp/__init__.py +1 -1
  167. flwr/server/serverapp/app.py +19 -1
  168. flwr/server/serverapp_components.py +1 -1
  169. flwr/server/strategy/__init__.py +1 -1
  170. flwr/server/strategy/aggregate.py +1 -1
  171. flwr/server/strategy/bulyan.py +2 -2
  172. flwr/server/strategy/dp_adaptive_clipping.py +17 -17
  173. flwr/server/strategy/dp_fixed_clipping.py +17 -17
  174. flwr/server/strategy/dpfedavg_adaptive.py +1 -1
  175. flwr/server/strategy/dpfedavg_fixed.py +1 -1
  176. flwr/server/strategy/fault_tolerant_fedavg.py +1 -1
  177. flwr/server/strategy/fedadagrad.py +1 -1
  178. flwr/server/strategy/fedadam.py +1 -1
  179. flwr/server/strategy/fedavg.py +1 -1
  180. flwr/server/strategy/fedavg_android.py +1 -1
  181. flwr/server/strategy/fedavgm.py +1 -1
  182. flwr/server/strategy/fedmedian.py +1 -1
  183. flwr/server/strategy/fedopt.py +1 -1
  184. flwr/server/strategy/fedprox.py +1 -1
  185. flwr/server/strategy/fedtrimmedavg.py +1 -1
  186. flwr/server/strategy/fedxgb_bagging.py +1 -1
  187. flwr/server/strategy/fedxgb_cyclic.py +1 -1
  188. flwr/server/strategy/fedxgb_nn_avg.py +3 -2
  189. flwr/server/strategy/fedyogi.py +1 -1
  190. flwr/server/strategy/krum.py +1 -1
  191. flwr/server/strategy/qfedavg.py +1 -1
  192. flwr/server/strategy/strategy.py +1 -1
  193. flwr/server/superlink/__init__.py +1 -1
  194. flwr/server/superlink/ffs/__init__.py +3 -1
  195. flwr/server/superlink/ffs/disk_ffs.py +1 -1
  196. flwr/server/superlink/ffs/ffs.py +1 -1
  197. flwr/server/superlink/ffs/ffs_factory.py +1 -1
  198. flwr/server/superlink/fleet/__init__.py +1 -1
  199. flwr/server/superlink/fleet/grpc_adapter/__init__.py +1 -1
  200. flwr/server/superlink/fleet/grpc_adapter/grpc_adapter_servicer.py +14 -4
  201. flwr/server/superlink/fleet/grpc_bidi/__init__.py +1 -1
  202. flwr/server/superlink/fleet/grpc_bidi/flower_service_servicer.py +1 -1
  203. flwr/server/superlink/fleet/grpc_bidi/grpc_bridge.py +1 -1
  204. flwr/server/superlink/fleet/grpc_bidi/grpc_client_proxy.py +1 -1
  205. flwr/server/superlink/fleet/grpc_bidi/grpc_server.py +13 -13
  206. flwr/server/superlink/fleet/grpc_rere/__init__.py +1 -1
  207. flwr/server/superlink/fleet/grpc_rere/fleet_servicer.py +102 -8
  208. flwr/server/superlink/fleet/grpc_rere/server_interceptor.py +1 -1
  209. flwr/server/superlink/fleet/message_handler/__init__.py +1 -1
  210. flwr/server/superlink/fleet/message_handler/message_handler.py +136 -19
  211. flwr/server/superlink/fleet/rest_rere/__init__.py +1 -1
  212. flwr/server/superlink/fleet/rest_rere/rest_api.py +73 -12
  213. flwr/server/superlink/fleet/vce/__init__.py +1 -1
  214. flwr/server/superlink/fleet/vce/backend/__init__.py +1 -1
  215. flwr/server/superlink/fleet/vce/backend/backend.py +1 -1
  216. flwr/server/superlink/fleet/vce/backend/raybackend.py +1 -1
  217. flwr/server/superlink/fleet/vce/vce_api.py +7 -4
  218. flwr/server/superlink/linkstate/__init__.py +1 -1
  219. flwr/server/superlink/linkstate/in_memory_linkstate.py +139 -44
  220. flwr/server/superlink/linkstate/linkstate.py +54 -21
  221. flwr/server/superlink/linkstate/linkstate_factory.py +1 -1
  222. flwr/server/superlink/linkstate/sqlite_linkstate.py +150 -56
  223. flwr/server/superlink/linkstate/utils.py +34 -30
  224. flwr/server/superlink/serverappio/serverappio_grpc.py +3 -0
  225. flwr/server/superlink/serverappio/serverappio_servicer.py +211 -57
  226. flwr/server/superlink/simulation/__init__.py +1 -1
  227. flwr/server/superlink/simulation/simulationio_grpc.py +1 -1
  228. flwr/server/superlink/simulation/simulationio_servicer.py +26 -2
  229. flwr/server/superlink/utils.py +45 -3
  230. flwr/server/typing.py +1 -1
  231. flwr/server/utils/__init__.py +1 -1
  232. flwr/server/utils/tensorboard.py +1 -1
  233. flwr/server/utils/validator.py +3 -3
  234. flwr/server/workflow/__init__.py +1 -1
  235. flwr/server/workflow/constant.py +1 -1
  236. flwr/server/workflow/default_workflows.py +1 -1
  237. flwr/server/workflow/secure_aggregation/__init__.py +1 -1
  238. flwr/server/workflow/secure_aggregation/secagg_workflow.py +1 -1
  239. flwr/server/workflow/secure_aggregation/secaggplus_workflow.py +1 -1
  240. flwr/serverapp/__init__.py +15 -0
  241. flwr/simulation/__init__.py +1 -1
  242. flwr/simulation/app.py +18 -1
  243. flwr/simulation/legacy_app.py +1 -1
  244. flwr/simulation/ray_transport/__init__.py +1 -1
  245. flwr/simulation/ray_transport/ray_actor.py +1 -1
  246. flwr/simulation/ray_transport/ray_client_proxy.py +1 -1
  247. flwr/simulation/ray_transport/utils.py +1 -1
  248. flwr/simulation/run_simulation.py +2 -2
  249. flwr/simulation/simulationio_connection.py +1 -1
  250. flwr/supercore/__init__.py +15 -0
  251. flwr/supercore/object_store/__init__.py +24 -0
  252. flwr/supercore/object_store/in_memory_object_store.py +229 -0
  253. flwr/supercore/object_store/object_store.py +192 -0
  254. flwr/supercore/object_store/object_store_factory.py +44 -0
  255. flwr/superexec/__init__.py +1 -1
  256. flwr/superexec/app.py +1 -1
  257. flwr/superexec/deployment.py +7 -3
  258. flwr/superexec/exec_event_log_interceptor.py +4 -4
  259. flwr/superexec/exec_grpc.py +8 -4
  260. flwr/superexec/exec_servicer.py +126 -24
  261. flwr/superexec/exec_user_auth_interceptor.py +38 -9
  262. flwr/superexec/executor.py +5 -1
  263. flwr/superexec/simulation.py +8 -2
  264. flwr/superlink/__init__.py +15 -0
  265. flwr/{client/supernode → supernode}/__init__.py +1 -8
  266. flwr/{client/nodestate/nodestate.py → supernode/cli/__init__.py} +8 -15
  267. flwr/{client/supernode/app.py → supernode/cli/flower_supernode.py} +4 -13
  268. flwr/supernode/cli/flwr_clientapp.py +81 -0
  269. flwr/{client → supernode}/nodestate/__init__.py +1 -1
  270. flwr/supernode/nodestate/in_memory_nodestate.py +190 -0
  271. flwr/supernode/nodestate/nodestate.py +212 -0
  272. flwr/{client → supernode}/nodestate/nodestate_factory.py +1 -1
  273. flwr/supernode/runtime/__init__.py +15 -0
  274. flwr/{client/clientapp/app.py → supernode/runtime/run_clientapp.py} +26 -57
  275. flwr/supernode/servicer/__init__.py +15 -0
  276. flwr/supernode/servicer/clientappio/__init__.py +24 -0
  277. flwr/{client/clientapp → supernode/servicer/clientappio}/clientappio_servicer.py +1 -1
  278. flwr/supernode/start_client_internal.py +491 -0
  279. {flwr-1.17.0.dist-info → flwr-1.19.0.dist-info}/METADATA +6 -5
  280. flwr-1.19.0.dist-info/RECORD +365 -0
  281. {flwr-1.17.0.dist-info → flwr-1.19.0.dist-info}/WHEEL +1 -1
  282. {flwr-1.17.0.dist-info → flwr-1.19.0.dist-info}/entry_points.txt +2 -2
  283. flwr/client/heartbeat.py +0 -74
  284. flwr/client/nodestate/in_memory_nodestate.py +0 -38
  285. flwr-1.17.0.dist-info/LICENSE +0 -202
  286. flwr-1.17.0.dist-info/RECORD +0 -333
@@ -16,38 +16,52 @@
16
16
 
17
17
 
18
18
  import threading
19
- from logging import DEBUG, INFO
19
+ from logging import DEBUG, ERROR, INFO
20
20
  from typing import Optional
21
- from uuid import UUID
22
21
 
23
22
  import grpc
24
23
 
25
- from flwr.common import ConfigRecord, Message
24
+ from flwr.common import Message
26
25
  from flwr.common.constant import SUPERLINK_NODE_ID, Status
26
+ from flwr.common.inflatable import (
27
+ UnexpectedObjectContentError,
28
+ get_all_nested_objects,
29
+ get_object_tree,
30
+ no_object_id_recompute,
31
+ )
27
32
  from flwr.common.logger import log
28
33
  from flwr.common.serde import (
29
34
  context_from_proto,
30
35
  context_to_proto,
31
- fab_from_proto,
32
36
  fab_to_proto,
33
37
  message_from_proto,
34
38
  message_to_proto,
35
39
  run_status_from_proto,
36
40
  run_status_to_proto,
37
41
  run_to_proto,
38
- user_config_from_proto,
39
42
  )
40
43
  from flwr.common.typing import Fab, RunStatus
41
44
  from flwr.proto import serverappio_pb2_grpc # pylint: disable=E0611
42
45
  from flwr.proto.fab_pb2 import GetFabRequest, GetFabResponse # pylint: disable=E0611
46
+ from flwr.proto.heartbeat_pb2 import ( # pylint: disable=E0611
47
+ SendAppHeartbeatRequest,
48
+ SendAppHeartbeatResponse,
49
+ )
43
50
  from flwr.proto.log_pb2 import ( # pylint: disable=E0611
44
51
  PushLogsRequest,
45
52
  PushLogsResponse,
46
53
  )
54
+ from flwr.proto.message_pb2 import ( # pylint: disable=E0611
55
+ ConfirmMessageReceivedRequest,
56
+ ConfirmMessageReceivedResponse,
57
+ ObjectIDs,
58
+ PullObjectRequest,
59
+ PullObjectResponse,
60
+ PushObjectRequest,
61
+ PushObjectResponse,
62
+ )
47
63
  from flwr.proto.node_pb2 import Node # pylint: disable=E0611
48
64
  from flwr.proto.run_pb2 import ( # pylint: disable=E0611
49
- CreateRunRequest,
50
- CreateRunResponse,
51
65
  GetRunRequest,
52
66
  GetRunResponse,
53
67
  GetRunStatusRequest,
@@ -72,16 +86,23 @@ from flwr.server.superlink.ffs.ffs_factory import FfsFactory
72
86
  from flwr.server.superlink.linkstate import LinkState, LinkStateFactory
73
87
  from flwr.server.superlink.utils import abort_if
74
88
  from flwr.server.utils.validator import validate_message
89
+ from flwr.supercore.object_store import NoObjectInStoreError, ObjectStoreFactory
90
+
91
+ from ..utils import store_mapping_and_register_objects
75
92
 
76
93
 
77
94
  class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
78
95
  """ServerAppIo API servicer."""
79
96
 
80
97
  def __init__(
81
- self, state_factory: LinkStateFactory, ffs_factory: FfsFactory
98
+ self,
99
+ state_factory: LinkStateFactory,
100
+ ffs_factory: FfsFactory,
101
+ objectstore_factory: ObjectStoreFactory,
82
102
  ) -> None:
83
103
  self.state_factory = state_factory
84
104
  self.ffs_factory = ffs_factory
105
+ self.objectstore_factory = objectstore_factory
85
106
  self.lock = threading.RLock()
86
107
 
87
108
  def GetNodes(
@@ -90,14 +111,16 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
90
111
  """Get available nodes."""
91
112
  log(DEBUG, "ServerAppIoServicer.GetNodes")
92
113
 
93
- # Init state
94
- state: LinkState = self.state_factory.state()
114
+ # Init state and store
115
+ state = self.state_factory.state()
116
+ store = self.objectstore_factory.store()
95
117
 
96
118
  # Abort if the run is not running
97
119
  abort_if(
98
120
  request.run_id,
99
121
  [Status.PENDING, Status.STARTING, Status.FINISHED],
100
122
  state,
123
+ store,
101
124
  context,
102
125
  )
103
126
 
@@ -105,46 +128,22 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
105
128
  nodes: list[Node] = [Node(node_id=node_id) for node_id in all_ids]
106
129
  return GetNodesResponse(nodes=nodes)
107
130
 
108
- def CreateRun(
109
- self, request: CreateRunRequest, context: grpc.ServicerContext
110
- ) -> CreateRunResponse:
111
- """Create run ID."""
112
- log(DEBUG, "ServerAppIoServicer.CreateRun")
113
- state: LinkState = self.state_factory.state()
114
- if request.HasField("fab"):
115
- fab = fab_from_proto(request.fab)
116
- ffs: Ffs = self.ffs_factory.ffs()
117
- fab_hash = ffs.put(fab.content, {})
118
- _raise_if(
119
- validation_error=fab_hash != fab.hash_str,
120
- request_name="CreateRun",
121
- detail=f"FAB ({fab.hash_str}) hash from request doesn't match contents",
122
- )
123
- else:
124
- fab_hash = ""
125
- run_id = state.create_run(
126
- request.fab_id,
127
- request.fab_version,
128
- fab_hash,
129
- user_config_from_proto(request.override_config),
130
- ConfigRecord(),
131
- )
132
- return CreateRunResponse(run_id=run_id)
133
-
134
131
  def PushMessages(
135
132
  self, request: PushInsMessagesRequest, context: grpc.ServicerContext
136
133
  ) -> PushInsMessagesResponse:
137
134
  """Push a set of Messages."""
138
135
  log(DEBUG, "ServerAppIoServicer.PushMessages")
139
136
 
140
- # Init state
141
- state: LinkState = self.state_factory.state()
137
+ # Init state and store
138
+ state = self.state_factory.state()
139
+ store = self.objectstore_factory.store()
142
140
 
143
141
  # Abort if the run is not running
144
142
  abort_if(
145
143
  request.run_id,
146
144
  [Status.PENDING, Status.STARTING, Status.FINISHED],
147
145
  state,
146
+ store,
148
147
  context,
149
148
  )
150
149
 
@@ -154,9 +153,8 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
154
153
  request_name="PushMessages",
155
154
  detail="`messages_list` must not be empty",
156
155
  )
157
- message_ids: list[Optional[UUID]] = []
158
- while request.messages_list:
159
- message_proto = request.messages_list.pop(0)
156
+ message_ids: list[Optional[str]] = []
157
+ for message_proto in request.messages_list:
160
158
  message = message_from_proto(message_proto=message_proto)
161
159
  validation_errors = validate_message(message, is_reply_message=False)
162
160
  _raise_if(
@@ -170,49 +168,70 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
170
168
  detail="`Message.metadata` has mismatched `run_id`",
171
169
  )
172
170
  # Store
173
- message_id: Optional[UUID] = state.store_message_ins(message=message)
171
+ message_id: Optional[str] = state.store_message_ins(message=message)
174
172
  message_ids.append(message_id)
175
173
 
174
+ # Store Message object to descendants mapping and preregister objects
175
+ objects_to_push = store_mapping_and_register_objects(store, request=request)
176
+
176
177
  return PushInsMessagesResponse(
177
178
  message_ids=[
178
179
  str(message_id) if message_id else "" for message_id in message_ids
179
- ]
180
+ ],
181
+ objects_to_push=objects_to_push,
180
182
  )
181
183
 
182
- def PullMessages(
184
+ def PullMessages( # pylint: disable=R0914
183
185
  self, request: PullResMessagesRequest, context: grpc.ServicerContext
184
186
  ) -> PullResMessagesResponse:
185
187
  """Pull a set of Messages."""
186
188
  log(DEBUG, "ServerAppIoServicer.PullMessages")
187
189
 
188
- # Init state
189
- state: LinkState = self.state_factory.state()
190
+ # Init state and store
191
+ state = self.state_factory.state()
192
+ store = self.objectstore_factory.store()
190
193
 
191
194
  # Abort if the run is not running
192
195
  abort_if(
193
196
  request.run_id,
194
197
  [Status.PENDING, Status.STARTING, Status.FINISHED],
195
198
  state,
199
+ store,
196
200
  context,
197
201
  )
198
202
 
199
- # Convert each message_id str to UUID
200
- message_ids: set[UUID] = {
201
- UUID(message_id) for message_id in request.message_ids
202
- }
203
-
204
203
  # Read from state
205
- messages_res: list[Message] = state.get_message_res(message_ids=message_ids)
204
+ messages_res: list[Message] = state.get_message_res(
205
+ message_ids=set(request.message_ids)
206
+ )
207
+
208
+ # Register messages generated by LinkState in the Store for consistency
209
+ for msg_res in messages_res:
210
+ if msg_res.metadata.src_node_id == SUPERLINK_NODE_ID:
211
+ with no_object_id_recompute():
212
+ all_objects = get_all_nested_objects(msg_res)
213
+ descendants = list(all_objects.keys())[:-1]
214
+ message_obj_id = msg_res.metadata.message_id
215
+ # Store mapping
216
+ store.set_message_descendant_ids(
217
+ msg_object_id=message_obj_id, descendant_ids=descendants
218
+ )
219
+ # Preregister
220
+ store.preregister(request.run_id, get_object_tree(msg_res))
221
+ # Store objects
222
+ for obj_id, obj in all_objects.items():
223
+ store.put(obj_id, obj.deflate())
206
224
 
207
225
  # Delete the instruction Messages and their replies if found
208
226
  message_ins_ids_to_delete = {
209
- UUID(msg_res.metadata.reply_to_message_id) for msg_res in messages_res
227
+ msg_res.metadata.reply_to_message_id for msg_res in messages_res
210
228
  }
211
229
 
212
230
  state.delete_messages(message_ins_ids=message_ins_ids_to_delete)
213
231
 
214
232
  # Convert Messages to proto
215
233
  messages_list = []
234
+ objects_to_pull: dict[str, ObjectIDs] = {}
216
235
  while messages_res:
217
236
  msg = messages_res.pop(0)
218
237
 
@@ -225,7 +244,19 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
225
244
  )
226
245
  messages_list.append(message_to_proto(msg))
227
246
 
228
- return PullResMessagesResponse(messages_list=messages_list)
247
+ try:
248
+ msg_object_id = msg.metadata.message_id
249
+ descendants = store.get_message_descendant_ids(msg_object_id)
250
+ # Add mapping of message object ID to its descendants
251
+ objects_to_pull[msg_object_id] = ObjectIDs(object_ids=descendants)
252
+ except NoObjectInStoreError as e:
253
+ log(ERROR, e.message)
254
+ # Delete message ins from state
255
+ state.delete_messages(message_ins_ids={msg_object_id})
256
+
257
+ return PullResMessagesResponse(
258
+ messages_list=messages_list, objects_to_pull=objects_to_pull
259
+ )
229
260
 
230
261
  def GetRun(
231
262
  self, request: GetRunRequest, context: grpc.ServicerContext
@@ -303,14 +334,16 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
303
334
  """Push ServerApp process outputs."""
304
335
  log(DEBUG, "ServerAppIoServicer.PushServerAppOutputs")
305
336
 
306
- # Init state
337
+ # Init state and store
307
338
  state = self.state_factory.state()
339
+ store = self.objectstore_factory.store()
308
340
 
309
341
  # Abort if the run is not running
310
342
  abort_if(
311
343
  request.run_id,
312
344
  [Status.PENDING, Status.STARTING, Status.FINISHED],
313
345
  state,
346
+ store,
314
347
  context,
315
348
  )
316
349
 
@@ -323,16 +356,23 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
323
356
  """Update the status of a run."""
324
357
  log(DEBUG, "ServerAppIoServicer.UpdateRunStatus")
325
358
 
326
- # Init state
359
+ # Init state and store
327
360
  state = self.state_factory.state()
361
+ store = self.objectstore_factory.store()
328
362
 
329
363
  # Abort if the run is finished
330
- abort_if(request.run_id, [Status.FINISHED], state, context)
364
+ abort_if(request.run_id, [Status.FINISHED], state, store, context)
331
365
 
332
366
  # Update the run status
333
367
  state.update_run_status(
334
368
  run_id=request.run_id, new_status=run_status_from_proto(request.run_status)
335
369
  )
370
+
371
+ # If the run is finished, delete the run from ObjectStore
372
+ if request.run_status.status == Status.FINISHED:
373
+ # Delete all objects related to the run
374
+ store.delete_objects_in_run(request.run_id)
375
+
336
376
  return UpdateRunStatusResponse()
337
377
 
338
378
  def PushLogs(
@@ -362,6 +402,120 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
362
402
  }
363
403
  return GetRunStatusResponse(run_status_dict=run_status_dict)
364
404
 
405
+ def SendAppHeartbeat(
406
+ self, request: SendAppHeartbeatRequest, context: grpc.ServicerContext
407
+ ) -> SendAppHeartbeatResponse:
408
+ """Handle a heartbeat from the ServerApp."""
409
+ log(DEBUG, "ServerAppIoServicer.SendAppHeartbeat")
410
+
411
+ # Init state
412
+ state = self.state_factory.state()
413
+
414
+ # Acknowledge the heartbeat
415
+ # The app heartbeat can only be acknowledged if the run is in
416
+ # starting or running status.
417
+ success = state.acknowledge_app_heartbeat(
418
+ run_id=request.run_id,
419
+ heartbeat_interval=request.heartbeat_interval,
420
+ )
421
+
422
+ return SendAppHeartbeatResponse(success=success)
423
+
424
+ def PushObject(
425
+ self, request: PushObjectRequest, context: grpc.ServicerContext
426
+ ) -> PushObjectResponse:
427
+ """Push an object to the ObjectStore."""
428
+ log(DEBUG, "ServerAppIoServicer.PushObject")
429
+
430
+ # Init state and store
431
+ state = self.state_factory.state()
432
+ store = self.objectstore_factory.store()
433
+
434
+ # Abort if the run is not running
435
+ abort_if(
436
+ request.run_id,
437
+ [Status.PENDING, Status.STARTING, Status.FINISHED],
438
+ state,
439
+ store,
440
+ context,
441
+ )
442
+
443
+ if request.node.node_id != SUPERLINK_NODE_ID:
444
+ # Cancel insertion in ObjectStore
445
+ context.abort(grpc.StatusCode.FAILED_PRECONDITION, "Unexpected node ID.")
446
+
447
+ # Insert in store
448
+ stored = False
449
+ try:
450
+ store.put(request.object_id, request.object_content)
451
+ stored = True
452
+ except (NoObjectInStoreError, ValueError) as e:
453
+ log(ERROR, str(e))
454
+ except UnexpectedObjectContentError as e:
455
+ # Object content is not valid
456
+ context.abort(grpc.StatusCode.FAILED_PRECONDITION, str(e))
457
+
458
+ return PushObjectResponse(stored=stored)
459
+
460
+ def PullObject(
461
+ self, request: PullObjectRequest, context: grpc.ServicerContext
462
+ ) -> PullObjectResponse:
463
+ """Pull an object from the ObjectStore."""
464
+ log(DEBUG, "ServerAppIoServicer.PullObject")
465
+
466
+ # Init state and store
467
+ state = self.state_factory.state()
468
+ store = self.objectstore_factory.store()
469
+
470
+ # Abort if the run is not running
471
+ abort_if(
472
+ request.run_id,
473
+ [Status.PENDING, Status.STARTING, Status.FINISHED],
474
+ state,
475
+ store,
476
+ context,
477
+ )
478
+
479
+ if request.node.node_id != SUPERLINK_NODE_ID:
480
+ # Cancel insertion in ObjectStore
481
+ context.abort(grpc.StatusCode.FAILED_PRECONDITION, "Unexpected node ID.")
482
+
483
+ # Fetch from store
484
+ content = store.get(request.object_id)
485
+ if content is not None:
486
+ object_available = content != b""
487
+ return PullObjectResponse(
488
+ object_found=True,
489
+ object_available=object_available,
490
+ object_content=content,
491
+ )
492
+ return PullObjectResponse(object_found=False, object_available=False)
493
+
494
+ def ConfirmMessageReceived(
495
+ self, request: ConfirmMessageReceivedRequest, context: grpc.ServicerContext
496
+ ) -> ConfirmMessageReceivedResponse:
497
+ """Confirm message received."""
498
+ log(DEBUG, "ServerAppIoServicer.ConfirmMessageReceived")
499
+
500
+ # Init state and store
501
+ state = self.state_factory.state()
502
+ store = self.objectstore_factory.store()
503
+
504
+ # Abort if the run is not running
505
+ abort_if(
506
+ request.run_id,
507
+ [Status.PENDING, Status.STARTING, Status.FINISHED],
508
+ state,
509
+ store,
510
+ context,
511
+ )
512
+
513
+ # Delete the message object
514
+ store.delete(request.message_object_id)
515
+ store.delete_message_descendant_ids(request.message_object_id)
516
+
517
+ return ConfirmMessageReceivedResponse()
518
+
365
519
 
366
520
  def _raise_if(validation_error: bool, request_name: str, detail: str) -> None:
367
521
  """Raise a `ValueError` with a detailed message if a validation error occurs."""
@@ -1,4 +1,4 @@
1
- # Copyright 2024 Flower Labs GmbH. All Rights Reserved.
1
+ # Copyright 2025 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 2024 Flower Labs GmbH. All Rights Reserved.
1
+ # Copyright 2025 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 2024 Flower Labs GmbH. All Rights Reserved.
1
+ # Copyright 2025 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.
@@ -34,6 +34,10 @@ from flwr.common.serde import (
34
34
  )
35
35
  from flwr.common.typing import Fab, RunStatus
36
36
  from flwr.proto import simulationio_pb2_grpc
37
+ from flwr.proto.heartbeat_pb2 import ( # pylint: disable=E0611
38
+ SendAppHeartbeatRequest,
39
+ SendAppHeartbeatResponse,
40
+ )
37
41
  from flwr.proto.log_pb2 import ( # pylint: disable=E0611
38
42
  PushLogsRequest,
39
43
  PushLogsResponse,
@@ -117,6 +121,7 @@ class SimulationIoServicer(simulationio_pb2_grpc.SimulationIoServicer):
117
121
  request.run_id,
118
122
  [Status.PENDING, Status.STARTING, Status.FINISHED],
119
123
  state,
124
+ None,
120
125
  context,
121
126
  )
122
127
 
@@ -131,7 +136,7 @@ class SimulationIoServicer(simulationio_pb2_grpc.SimulationIoServicer):
131
136
  state = self.state_factory.state()
132
137
 
133
138
  # Abort if the run is finished
134
- abort_if(request.run_id, [Status.FINISHED], state, context)
139
+ abort_if(request.run_id, [Status.FINISHED], state, None, context)
135
140
 
136
141
  # Update the run status
137
142
  state.update_run_status(
@@ -184,3 +189,22 @@ class SimulationIoServicer(simulationio_pb2_grpc.SimulationIoServicer):
184
189
  return GetFederationOptionsResponse(
185
190
  federation_options=config_record_to_proto(federation_options)
186
191
  )
192
+
193
+ def SendAppHeartbeat(
194
+ self, request: SendAppHeartbeatRequest, context: grpc.ServicerContext
195
+ ) -> SendAppHeartbeatResponse:
196
+ """Handle a heartbeat from the ServerApp in simulation."""
197
+ log(DEBUG, "SimultionIoServicer.SendAppHeartbeat")
198
+
199
+ # Init state
200
+ state = self.state_factory.state()
201
+
202
+ # Acknowledge the heartbeat
203
+ # The app heartbeat can only be acknowledged if the run is in
204
+ # starting or running status.
205
+ success = state.acknowledge_app_heartbeat(
206
+ run_id=request.run_id,
207
+ heartbeat_interval=request.heartbeat_interval,
208
+ )
209
+
210
+ return SendAppHeartbeatResponse(success=success)
@@ -1,4 +1,4 @@
1
- # Copyright 2024 Flower Labs GmbH. All Rights Reserved.
1
+ # Copyright 2025 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,18 @@
15
15
  """SuperLink utilities."""
16
16
 
17
17
 
18
- from typing import Union
18
+ from typing import Optional, Union
19
19
 
20
20
  import grpc
21
21
 
22
22
  from flwr.common.constant import Status, SubStatus
23
+ from flwr.common.inflatable import iterate_object_tree
23
24
  from flwr.common.typing import RunStatus
25
+ from flwr.proto.fleet_pb2 import PushMessagesRequest # pylint: disable=E0611
26
+ from flwr.proto.message_pb2 import ObjectIDs # pylint: disable=E0611
27
+ from flwr.proto.serverappio_pb2 import PushInsMessagesRequest # pylint: disable=E0611
24
28
  from flwr.server.superlink.linkstate import LinkState
29
+ from flwr.supercore.object_store import ObjectStore
25
30
 
26
31
  _STATUS_TO_MSG = {
27
32
  Status.PENDING: "Run is pending.",
@@ -35,6 +40,7 @@ def check_abort(
35
40
  run_id: int,
36
41
  abort_status_list: list[str],
37
42
  state: LinkState,
43
+ store: Optional[ObjectStore] = None,
38
44
  ) -> Union[str, None]:
39
45
  """Check if the status of the provided `run_id` is in `abort_status_list`."""
40
46
  run_status: RunStatus = state.get_run_status({run_id})[run_id]
@@ -45,6 +51,10 @@ def check_abort(
45
51
  msg += " Stopped by user."
46
52
  return msg
47
53
 
54
+ # Clear the objects of the run from the store if the run is finished
55
+ if store and run_status.status == Status.FINISHED:
56
+ store.delete_objects_in_run(run_id)
57
+
48
58
  return None
49
59
 
50
60
 
@@ -58,8 +68,40 @@ def abort_if(
58
68
  run_id: int,
59
69
  abort_status_list: list[str],
60
70
  state: LinkState,
71
+ store: Optional[ObjectStore],
61
72
  context: grpc.ServicerContext,
62
73
  ) -> None:
63
74
  """Abort context if status of the provided `run_id` is in `abort_status_list`."""
64
- msg = check_abort(run_id, abort_status_list, state)
75
+ msg = check_abort(run_id, abort_status_list, state, store)
65
76
  abort_grpc_context(msg, context)
77
+
78
+
79
+ def store_mapping_and_register_objects(
80
+ store: ObjectStore, request: Union[PushInsMessagesRequest, PushMessagesRequest]
81
+ ) -> dict[str, ObjectIDs]:
82
+ """Store Message object to descendants mapping and preregister objects."""
83
+ if not request.messages_list:
84
+ return {}
85
+
86
+ objects_to_push: dict[str, ObjectIDs] = {}
87
+
88
+ # Get run_id from the first message in the list
89
+ # All messages of a request should in the same run
90
+ run_id = request.messages_list[0].metadata.run_id
91
+
92
+ for object_tree in request.message_object_trees:
93
+ all_object_ids = [obj.object_id for obj in iterate_object_tree(object_tree)]
94
+ msg_object_id, descendant_ids = all_object_ids[-1], all_object_ids[:-1]
95
+ # Store mapping
96
+ store.set_message_descendant_ids(
97
+ msg_object_id=msg_object_id, descendant_ids=descendant_ids
98
+ )
99
+
100
+ # Preregister
101
+ object_ids_just_registered = store.preregister(run_id, object_tree)
102
+ # Keep track of objects that need to be pushed
103
+ objects_to_push[msg_object_id] = ObjectIDs(
104
+ object_ids=object_ids_just_registered
105
+ )
106
+
107
+ return objects_to_push
flwr/server/typing.py CHANGED
@@ -1,4 +1,4 @@
1
- # Copyright 2024 Flower Labs GmbH. All Rights Reserved.
1
+ # Copyright 2025 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 2021 Flower Labs GmbH. All Rights Reserved.
1
+ # Copyright 2025 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 2021 Flower Labs GmbH. All Rights Reserved.
1
+ # Copyright 2025 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 2023 Flower Labs GmbH. All Rights Reserved.
1
+ # Copyright 2025 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.
@@ -27,8 +27,8 @@ def validate_message(message: Message, is_reply_message: bool) -> list[str]:
27
27
  validation_errors = []
28
28
  metadata = message.metadata
29
29
 
30
- if metadata.message_id != "":
31
- validation_errors.append("non-empty `metadata.message_id`")
30
+ if metadata.message_id == "":
31
+ validation_errors.append("empty `metadata.message_id`")
32
32
 
33
33
  # Created/delivered/TTL/Pushed
34
34
  if (
@@ -1,4 +1,4 @@
1
- # Copyright 2024 Flower Labs GmbH. All Rights Reserved.
1
+ # Copyright 2025 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 2024 Flower Labs GmbH. All Rights Reserved.
1
+ # Copyright 2025 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 2024 Flower Labs GmbH. All Rights Reserved.
1
+ # Copyright 2025 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 2024 Flower Labs GmbH. All Rights Reserved.
1
+ # Copyright 2025 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 2024 Flower Labs GmbH. All Rights Reserved.
1
+ # Copyright 2025 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 2024 Flower Labs GmbH. All Rights Reserved.
1
+ # Copyright 2025 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.