flwr 1.23.0__py3-none-any.whl → 1.25.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 (339) hide show
  1. flwr/__init__.py +16 -5
  2. flwr/app/error.py +2 -2
  3. flwr/app/exception.py +3 -3
  4. flwr/cli/app.py +19 -0
  5. flwr/cli/{new/templates → app_cmd}/__init__.py +9 -1
  6. flwr/cli/app_cmd/publish.py +285 -0
  7. flwr/cli/app_cmd/review.py +262 -0
  8. flwr/cli/auth_plugin/auth_plugin.py +4 -5
  9. flwr/cli/auth_plugin/noop_auth_plugin.py +54 -11
  10. flwr/cli/auth_plugin/oidc_cli_plugin.py +32 -9
  11. flwr/cli/build.py +60 -18
  12. flwr/cli/cli_account_auth_interceptor.py +24 -7
  13. flwr/cli/config_utils.py +101 -13
  14. flwr/cli/{new/templates/app/code/flwr_tune → federation}/__init__.py +10 -1
  15. flwr/cli/federation/ls.py +140 -0
  16. flwr/cli/federation/show.py +318 -0
  17. flwr/cli/install.py +91 -13
  18. flwr/cli/log.py +52 -9
  19. flwr/cli/login/login.py +7 -4
  20. flwr/cli/ls.py +211 -130
  21. flwr/cli/new/new.py +123 -331
  22. flwr/cli/pull.py +10 -5
  23. flwr/cli/run/run.py +71 -29
  24. flwr/cli/run_utils.py +148 -0
  25. flwr/cli/stop.py +26 -8
  26. flwr/cli/supernode/ls.py +25 -12
  27. flwr/cli/supernode/register.py +9 -4
  28. flwr/cli/supernode/unregister.py +5 -3
  29. flwr/cli/utils.py +239 -16
  30. flwr/client/__init__.py +1 -1
  31. flwr/client/dpfedavg_numpy_client.py +4 -1
  32. flwr/client/grpc_adapter_client/connection.py +8 -9
  33. flwr/client/grpc_rere_client/connection.py +16 -14
  34. flwr/client/grpc_rere_client/grpc_adapter.py +6 -2
  35. flwr/client/grpc_rere_client/node_auth_client_interceptor.py +2 -1
  36. flwr/client/message_handler/message_handler.py +2 -2
  37. flwr/client/mod/secure_aggregation/secaggplus_mod.py +3 -3
  38. flwr/client/numpy_client.py +1 -1
  39. flwr/client/rest_client/connection.py +18 -18
  40. flwr/client/run_info_store.py +4 -5
  41. flwr/client/typing.py +1 -1
  42. flwr/clientapp/client_app.py +9 -10
  43. flwr/clientapp/mod/centraldp_mods.py +16 -17
  44. flwr/clientapp/mod/localdp_mod.py +8 -9
  45. flwr/clientapp/typing.py +1 -1
  46. flwr/clientapp/utils.py +3 -3
  47. flwr/common/address.py +1 -2
  48. flwr/common/args.py +3 -4
  49. flwr/common/config.py +13 -16
  50. flwr/common/constant.py +5 -2
  51. flwr/common/differential_privacy.py +3 -4
  52. flwr/common/event_log_plugin/event_log_plugin.py +3 -4
  53. flwr/common/exit/exit.py +15 -2
  54. flwr/common/exit/exit_code.py +19 -0
  55. flwr/common/exit/exit_handler.py +6 -2
  56. flwr/common/exit/signal_handler.py +5 -5
  57. flwr/common/grpc.py +6 -6
  58. flwr/common/inflatable_protobuf_utils.py +1 -1
  59. flwr/common/inflatable_utils.py +38 -21
  60. flwr/common/logger.py +19 -19
  61. flwr/common/message.py +4 -4
  62. flwr/common/object_ref.py +7 -7
  63. flwr/common/record/array.py +3 -3
  64. flwr/common/record/arrayrecord.py +18 -30
  65. flwr/common/record/configrecord.py +3 -3
  66. flwr/common/record/recorddict.py +5 -5
  67. flwr/common/record/typeddict.py +9 -2
  68. flwr/common/recorddict_compat.py +7 -10
  69. flwr/common/retry_invoker.py +20 -20
  70. flwr/common/secure_aggregation/ndarrays_arithmetic.py +3 -3
  71. flwr/common/serde.py +11 -4
  72. flwr/common/serde_utils.py +2 -2
  73. flwr/common/telemetry.py +9 -5
  74. flwr/common/typing.py +58 -37
  75. flwr/compat/client/app.py +38 -37
  76. flwr/compat/client/grpc_client/connection.py +11 -11
  77. flwr/compat/server/app.py +5 -6
  78. flwr/proto/appio_pb2.py +13 -3
  79. flwr/proto/appio_pb2.pyi +134 -65
  80. flwr/proto/appio_pb2_grpc.py +20 -0
  81. flwr/proto/appio_pb2_grpc.pyi +27 -0
  82. flwr/proto/clientappio_pb2.py +17 -7
  83. flwr/proto/clientappio_pb2.pyi +15 -0
  84. flwr/proto/clientappio_pb2_grpc.py +206 -40
  85. flwr/proto/clientappio_pb2_grpc.pyi +168 -53
  86. flwr/proto/control_pb2.py +71 -52
  87. flwr/proto/control_pb2.pyi +277 -111
  88. flwr/proto/control_pb2_grpc.py +249 -40
  89. flwr/proto/control_pb2_grpc.pyi +185 -52
  90. flwr/proto/error_pb2.py +13 -3
  91. flwr/proto/error_pb2.pyi +24 -6
  92. flwr/proto/error_pb2_grpc.py +20 -0
  93. flwr/proto/error_pb2_grpc.pyi +27 -0
  94. flwr/proto/fab_pb2.py +14 -4
  95. flwr/proto/fab_pb2.pyi +59 -31
  96. flwr/proto/fab_pb2_grpc.py +20 -0
  97. flwr/proto/fab_pb2_grpc.pyi +27 -0
  98. flwr/proto/federation_pb2.py +38 -0
  99. flwr/proto/federation_pb2.pyi +56 -0
  100. flwr/proto/federation_pb2_grpc.py +24 -0
  101. flwr/proto/federation_pb2_grpc.pyi +31 -0
  102. flwr/proto/fleet_pb2.py +24 -14
  103. flwr/proto/fleet_pb2.pyi +141 -61
  104. flwr/proto/fleet_pb2_grpc.py +189 -48
  105. flwr/proto/fleet_pb2_grpc.pyi +175 -61
  106. flwr/proto/grpcadapter_pb2.py +14 -4
  107. flwr/proto/grpcadapter_pb2.pyi +38 -16
  108. flwr/proto/grpcadapter_pb2_grpc.py +35 -4
  109. flwr/proto/grpcadapter_pb2_grpc.pyi +38 -7
  110. flwr/proto/heartbeat_pb2.py +17 -7
  111. flwr/proto/heartbeat_pb2.pyi +51 -22
  112. flwr/proto/heartbeat_pb2_grpc.py +20 -0
  113. flwr/proto/heartbeat_pb2_grpc.pyi +27 -0
  114. flwr/proto/log_pb2.py +13 -3
  115. flwr/proto/log_pb2.pyi +34 -11
  116. flwr/proto/log_pb2_grpc.py +20 -0
  117. flwr/proto/log_pb2_grpc.pyi +27 -0
  118. flwr/proto/message_pb2.py +15 -5
  119. flwr/proto/message_pb2.pyi +154 -86
  120. flwr/proto/message_pb2_grpc.py +20 -0
  121. flwr/proto/message_pb2_grpc.pyi +27 -0
  122. flwr/proto/node_pb2.py +15 -5
  123. flwr/proto/node_pb2.pyi +50 -25
  124. flwr/proto/node_pb2_grpc.py +20 -0
  125. flwr/proto/node_pb2_grpc.pyi +27 -0
  126. flwr/proto/recorddict_pb2.py +13 -3
  127. flwr/proto/recorddict_pb2.pyi +184 -107
  128. flwr/proto/recorddict_pb2_grpc.py +20 -0
  129. flwr/proto/recorddict_pb2_grpc.pyi +27 -0
  130. flwr/proto/run_pb2.py +40 -31
  131. flwr/proto/run_pb2.pyi +158 -84
  132. flwr/proto/run_pb2_grpc.py +20 -0
  133. flwr/proto/run_pb2_grpc.pyi +27 -0
  134. flwr/proto/serverappio_pb2.py +13 -3
  135. flwr/proto/serverappio_pb2.pyi +32 -8
  136. flwr/proto/serverappio_pb2_grpc.py +246 -65
  137. flwr/proto/serverappio_pb2_grpc.pyi +221 -85
  138. flwr/proto/simulationio_pb2.py +16 -8
  139. flwr/proto/simulationio_pb2.pyi +15 -0
  140. flwr/proto/simulationio_pb2_grpc.py +162 -41
  141. flwr/proto/simulationio_pb2_grpc.pyi +149 -55
  142. flwr/proto/transport_pb2.py +20 -10
  143. flwr/proto/transport_pb2.pyi +249 -160
  144. flwr/proto/transport_pb2_grpc.py +35 -4
  145. flwr/proto/transport_pb2_grpc.pyi +38 -8
  146. flwr/server/app.py +39 -17
  147. flwr/server/client_manager.py +4 -5
  148. flwr/server/client_proxy.py +10 -11
  149. flwr/server/compat/app.py +4 -5
  150. flwr/server/compat/app_utils.py +2 -1
  151. flwr/server/compat/grid_client_proxy.py +10 -12
  152. flwr/server/compat/legacy_context.py +3 -4
  153. flwr/server/fleet_event_log_interceptor.py +2 -1
  154. flwr/server/grid/grid.py +2 -3
  155. flwr/server/grid/grpc_grid.py +10 -8
  156. flwr/server/grid/inmemory_grid.py +4 -4
  157. flwr/server/run_serverapp.py +2 -3
  158. flwr/server/server.py +34 -39
  159. flwr/server/server_app.py +7 -8
  160. flwr/server/server_config.py +1 -2
  161. flwr/server/serverapp/app.py +34 -28
  162. flwr/server/serverapp_components.py +4 -5
  163. flwr/server/strategy/aggregate.py +9 -8
  164. flwr/server/strategy/bulyan.py +13 -11
  165. flwr/server/strategy/dp_adaptive_clipping.py +16 -20
  166. flwr/server/strategy/dp_fixed_clipping.py +12 -17
  167. flwr/server/strategy/dpfedavg_adaptive.py +3 -4
  168. flwr/server/strategy/dpfedavg_fixed.py +6 -10
  169. flwr/server/strategy/fault_tolerant_fedavg.py +14 -13
  170. flwr/server/strategy/fedadagrad.py +18 -14
  171. flwr/server/strategy/fedadam.py +16 -14
  172. flwr/server/strategy/fedavg.py +16 -17
  173. flwr/server/strategy/fedavg_android.py +15 -15
  174. flwr/server/strategy/fedavgm.py +21 -18
  175. flwr/server/strategy/fedmedian.py +2 -3
  176. flwr/server/strategy/fedopt.py +11 -10
  177. flwr/server/strategy/fedprox.py +10 -9
  178. flwr/server/strategy/fedtrimmedavg.py +12 -11
  179. flwr/server/strategy/fedxgb_bagging.py +13 -11
  180. flwr/server/strategy/fedxgb_cyclic.py +6 -6
  181. flwr/server/strategy/fedxgb_nn_avg.py +4 -4
  182. flwr/server/strategy/fedyogi.py +16 -14
  183. flwr/server/strategy/krum.py +12 -11
  184. flwr/server/strategy/qfedavg.py +16 -15
  185. flwr/server/strategy/strategy.py +6 -9
  186. flwr/server/superlink/fleet/grpc_adapter/grpc_adapter_servicer.py +2 -1
  187. flwr/server/superlink/fleet/grpc_bidi/flower_service_servicer.py +1 -2
  188. flwr/server/superlink/fleet/grpc_bidi/grpc_bridge.py +3 -4
  189. flwr/server/superlink/fleet/grpc_bidi/grpc_client_proxy.py +10 -12
  190. flwr/server/superlink/fleet/grpc_bidi/grpc_server.py +1 -3
  191. flwr/server/superlink/fleet/grpc_rere/fleet_servicer.py +4 -4
  192. flwr/server/superlink/fleet/grpc_rere/node_auth_server_interceptor.py +3 -2
  193. flwr/server/superlink/fleet/message_handler/message_handler.py +75 -30
  194. flwr/server/superlink/fleet/rest_rere/rest_api.py +2 -2
  195. flwr/server/superlink/fleet/vce/backend/backend.py +1 -1
  196. flwr/server/superlink/fleet/vce/backend/raybackend.py +5 -5
  197. flwr/server/superlink/fleet/vce/vce_api.py +15 -9
  198. flwr/server/superlink/linkstate/in_memory_linkstate.py +148 -149
  199. flwr/server/superlink/linkstate/linkstate.py +91 -43
  200. flwr/server/superlink/linkstate/linkstate_factory.py +22 -5
  201. flwr/server/superlink/linkstate/sqlite_linkstate.py +502 -436
  202. flwr/server/superlink/linkstate/utils.py +6 -6
  203. flwr/server/superlink/serverappio/serverappio_grpc.py +1 -2
  204. flwr/server/superlink/serverappio/serverappio_servicer.py +26 -21
  205. flwr/server/superlink/simulation/simulationio_grpc.py +1 -2
  206. flwr/server/superlink/simulation/simulationio_servicer.py +18 -13
  207. flwr/server/superlink/utils.py +4 -6
  208. flwr/server/typing.py +1 -1
  209. flwr/server/utils/tensorboard.py +15 -8
  210. flwr/server/workflow/default_workflows.py +5 -5
  211. flwr/server/workflow/secure_aggregation/secagg_workflow.py +2 -4
  212. flwr/server/workflow/secure_aggregation/secaggplus_workflow.py +8 -8
  213. flwr/serverapp/strategy/bulyan.py +16 -15
  214. flwr/serverapp/strategy/dp_adaptive_clipping.py +12 -11
  215. flwr/serverapp/strategy/dp_fixed_clipping.py +11 -14
  216. flwr/serverapp/strategy/fedadagrad.py +10 -11
  217. flwr/serverapp/strategy/fedadam.py +10 -11
  218. flwr/serverapp/strategy/fedavg.py +9 -10
  219. flwr/serverapp/strategy/fedavgm.py +17 -16
  220. flwr/serverapp/strategy/fedmedian.py +2 -2
  221. flwr/serverapp/strategy/fedopt.py +10 -11
  222. flwr/serverapp/strategy/fedprox.py +7 -8
  223. flwr/serverapp/strategy/fedtrimmedavg.py +9 -9
  224. flwr/serverapp/strategy/fedxgb_bagging.py +3 -3
  225. flwr/serverapp/strategy/fedxgb_cyclic.py +9 -9
  226. flwr/serverapp/strategy/fedyogi.py +9 -11
  227. flwr/serverapp/strategy/krum.py +7 -7
  228. flwr/serverapp/strategy/multikrum.py +9 -9
  229. flwr/serverapp/strategy/qfedavg.py +17 -16
  230. flwr/serverapp/strategy/strategy.py +6 -9
  231. flwr/serverapp/strategy/strategy_utils.py +7 -8
  232. flwr/simulation/app.py +46 -42
  233. flwr/simulation/legacy_app.py +12 -12
  234. flwr/simulation/ray_transport/ray_actor.py +10 -11
  235. flwr/simulation/ray_transport/ray_client_proxy.py +11 -12
  236. flwr/simulation/run_simulation.py +43 -43
  237. flwr/simulation/simulationio_connection.py +4 -4
  238. flwr/supercore/cli/flower_superexec.py +3 -4
  239. flwr/supercore/constant.py +34 -1
  240. flwr/supercore/corestate/corestate.py +24 -3
  241. flwr/supercore/corestate/in_memory_corestate.py +138 -0
  242. flwr/supercore/corestate/sqlite_corestate.py +157 -0
  243. flwr/supercore/ffs/disk_ffs.py +1 -2
  244. flwr/supercore/ffs/ffs.py +1 -2
  245. flwr/supercore/ffs/ffs_factory.py +1 -2
  246. flwr/{common → supercore}/heartbeat.py +20 -25
  247. flwr/supercore/object_store/in_memory_object_store.py +1 -2
  248. flwr/supercore/object_store/object_store.py +1 -2
  249. flwr/supercore/object_store/object_store_factory.py +1 -2
  250. flwr/supercore/object_store/sqlite_object_store.py +8 -7
  251. flwr/supercore/primitives/asymmetric.py +1 -1
  252. flwr/supercore/primitives/asymmetric_ed25519.py +11 -1
  253. flwr/supercore/sqlite_mixin.py +37 -34
  254. flwr/supercore/superexec/plugin/base_exec_plugin.py +1 -2
  255. flwr/supercore/superexec/plugin/exec_plugin.py +3 -3
  256. flwr/supercore/superexec/run_superexec.py +9 -13
  257. flwr/supercore/utils.py +190 -0
  258. flwr/superlink/artifact_provider/artifact_provider.py +1 -2
  259. flwr/superlink/auth_plugin/auth_plugin.py +6 -9
  260. flwr/superlink/auth_plugin/noop_auth_plugin.py +6 -9
  261. flwr/{cli/new/templates/app → superlink/federation}/__init__.py +10 -1
  262. flwr/superlink/federation/federation_manager.py +64 -0
  263. flwr/superlink/federation/noop_federation_manager.py +71 -0
  264. flwr/superlink/servicer/control/control_account_auth_interceptor.py +22 -13
  265. flwr/superlink/servicer/control/control_event_log_interceptor.py +7 -7
  266. flwr/superlink/servicer/control/control_grpc.py +7 -6
  267. flwr/superlink/servicer/control/control_license_interceptor.py +3 -3
  268. flwr/superlink/servicer/control/control_servicer.py +190 -23
  269. flwr/supernode/cli/flower_supernode.py +58 -3
  270. flwr/supernode/nodestate/in_memory_nodestate.py +121 -49
  271. flwr/supernode/nodestate/nodestate.py +52 -8
  272. flwr/supernode/nodestate/nodestate_factory.py +7 -4
  273. flwr/supernode/runtime/run_clientapp.py +41 -22
  274. flwr/supernode/servicer/clientappio/clientappio_servicer.py +46 -10
  275. flwr/supernode/start_client_internal.py +165 -46
  276. {flwr-1.23.0.dist-info → flwr-1.25.0.dist-info}/METADATA +9 -11
  277. flwr-1.25.0.dist-info/RECORD +393 -0
  278. flwr/cli/new/templates/app/.gitignore.tpl +0 -163
  279. flwr/cli/new/templates/app/LICENSE.tpl +0 -202
  280. flwr/cli/new/templates/app/README.baseline.md.tpl +0 -127
  281. flwr/cli/new/templates/app/README.flowertune.md.tpl +0 -68
  282. flwr/cli/new/templates/app/README.md.tpl +0 -37
  283. flwr/cli/new/templates/app/code/__init__.baseline.py.tpl +0 -1
  284. flwr/cli/new/templates/app/code/__init__.py +0 -15
  285. flwr/cli/new/templates/app/code/__init__.py.tpl +0 -1
  286. flwr/cli/new/templates/app/code/__init__.pytorch_legacy_api.py.tpl +0 -1
  287. flwr/cli/new/templates/app/code/client.baseline.py.tpl +0 -75
  288. flwr/cli/new/templates/app/code/client.huggingface.py.tpl +0 -93
  289. flwr/cli/new/templates/app/code/client.jax.py.tpl +0 -71
  290. flwr/cli/new/templates/app/code/client.mlx.py.tpl +0 -102
  291. flwr/cli/new/templates/app/code/client.numpy.py.tpl +0 -46
  292. flwr/cli/new/templates/app/code/client.pytorch.py.tpl +0 -80
  293. flwr/cli/new/templates/app/code/client.pytorch_legacy_api.py.tpl +0 -55
  294. flwr/cli/new/templates/app/code/client.sklearn.py.tpl +0 -108
  295. flwr/cli/new/templates/app/code/client.tensorflow.py.tpl +0 -82
  296. flwr/cli/new/templates/app/code/client.xgboost.py.tpl +0 -110
  297. flwr/cli/new/templates/app/code/dataset.baseline.py.tpl +0 -36
  298. flwr/cli/new/templates/app/code/flwr_tune/client_app.py.tpl +0 -92
  299. flwr/cli/new/templates/app/code/flwr_tune/dataset.py.tpl +0 -87
  300. flwr/cli/new/templates/app/code/flwr_tune/models.py.tpl +0 -56
  301. flwr/cli/new/templates/app/code/flwr_tune/server_app.py.tpl +0 -73
  302. flwr/cli/new/templates/app/code/flwr_tune/strategy.py.tpl +0 -78
  303. flwr/cli/new/templates/app/code/model.baseline.py.tpl +0 -66
  304. flwr/cli/new/templates/app/code/server.baseline.py.tpl +0 -43
  305. flwr/cli/new/templates/app/code/server.huggingface.py.tpl +0 -42
  306. flwr/cli/new/templates/app/code/server.jax.py.tpl +0 -39
  307. flwr/cli/new/templates/app/code/server.mlx.py.tpl +0 -41
  308. flwr/cli/new/templates/app/code/server.numpy.py.tpl +0 -38
  309. flwr/cli/new/templates/app/code/server.pytorch.py.tpl +0 -41
  310. flwr/cli/new/templates/app/code/server.pytorch_legacy_api.py.tpl +0 -31
  311. flwr/cli/new/templates/app/code/server.sklearn.py.tpl +0 -44
  312. flwr/cli/new/templates/app/code/server.tensorflow.py.tpl +0 -38
  313. flwr/cli/new/templates/app/code/server.xgboost.py.tpl +0 -56
  314. flwr/cli/new/templates/app/code/strategy.baseline.py.tpl +0 -1
  315. flwr/cli/new/templates/app/code/task.huggingface.py.tpl +0 -98
  316. flwr/cli/new/templates/app/code/task.jax.py.tpl +0 -57
  317. flwr/cli/new/templates/app/code/task.mlx.py.tpl +0 -102
  318. flwr/cli/new/templates/app/code/task.numpy.py.tpl +0 -7
  319. flwr/cli/new/templates/app/code/task.pytorch.py.tpl +0 -98
  320. flwr/cli/new/templates/app/code/task.pytorch_legacy_api.py.tpl +0 -111
  321. flwr/cli/new/templates/app/code/task.sklearn.py.tpl +0 -67
  322. flwr/cli/new/templates/app/code/task.tensorflow.py.tpl +0 -52
  323. flwr/cli/new/templates/app/code/task.xgboost.py.tpl +0 -67
  324. flwr/cli/new/templates/app/code/utils.baseline.py.tpl +0 -1
  325. flwr/cli/new/templates/app/pyproject.baseline.toml.tpl +0 -146
  326. flwr/cli/new/templates/app/pyproject.flowertune.toml.tpl +0 -80
  327. flwr/cli/new/templates/app/pyproject.huggingface.toml.tpl +0 -65
  328. flwr/cli/new/templates/app/pyproject.jax.toml.tpl +0 -52
  329. flwr/cli/new/templates/app/pyproject.mlx.toml.tpl +0 -56
  330. flwr/cli/new/templates/app/pyproject.numpy.toml.tpl +0 -49
  331. flwr/cli/new/templates/app/pyproject.pytorch.toml.tpl +0 -53
  332. flwr/cli/new/templates/app/pyproject.pytorch_legacy_api.toml.tpl +0 -53
  333. flwr/cli/new/templates/app/pyproject.sklearn.toml.tpl +0 -52
  334. flwr/cli/new/templates/app/pyproject.tensorflow.toml.tpl +0 -53
  335. flwr/cli/new/templates/app/pyproject.xgboost.toml.tpl +0 -61
  336. flwr/supercore/object_store/utils.py +0 -43
  337. flwr-1.23.0.dist-info/RECORD +0 -439
  338. {flwr-1.23.0.dist-info → flwr-1.25.0.dist-info}/WHEEL +0 -0
  339. {flwr-1.23.0.dist-info → flwr-1.25.0.dist-info}/entry_points.txt +0 -0
@@ -15,18 +15,28 @@
15
15
  """In-memory NodeState implementation."""
16
16
 
17
17
 
18
- import secrets
19
18
  from collections.abc import Sequence
20
19
  from dataclasses import dataclass
21
- from threading import Lock
22
- from typing import Optional
23
-
24
- from flwr.common import Context, Message
25
- from flwr.common.constant import FLWR_APP_TOKEN_LENGTH
20
+ from threading import Lock, RLock
21
+
22
+ from flwr.common import Context, Error, Message, now
23
+ from flwr.common.constant import ErrorCode
24
+ from flwr.common.inflatable import (
25
+ get_all_nested_objects,
26
+ get_object_tree,
27
+ no_object_id_recompute,
28
+ )
26
29
  from flwr.common.typing import Run
30
+ from flwr.supercore.constant import MESSAGE_TIME_ENTRY_MAX_AGE_SECONDS
31
+ from flwr.supercore.corestate.in_memory_corestate import InMemoryCoreState
32
+ from flwr.supercore.object_store import ObjectStore
27
33
 
28
34
  from .nodestate import NodeState
29
35
 
36
+ CLIENT_APP_CRASHED_ERROR = Error(
37
+ ErrorCode.CLIENT_APP_CRASHED, "ClientApp stopped responding."
38
+ )
39
+
30
40
 
31
41
  @dataclass
32
42
  class MessageEntry:
@@ -36,27 +46,37 @@ class MessageEntry:
36
46
  is_retrieved: bool = False
37
47
 
38
48
 
39
- class InMemoryNodeState(NodeState): # pylint: disable=too-many-instance-attributes
49
+ @dataclass
50
+ class TimeEntry:
51
+ """Data class to represent a time entry."""
52
+
53
+ starting_at: float
54
+ finished_at: float | None = None
55
+
56
+
57
+ class InMemoryNodeState(
58
+ NodeState, InMemoryCoreState
59
+ ): # pylint: disable=too-many-instance-attributes
40
60
  """In-memory NodeState implementation."""
41
61
 
42
- def __init__(self) -> None:
62
+ def __init__(self, object_store: ObjectStore) -> None:
63
+ super().__init__(object_store)
43
64
  # Store node_id
44
- self.node_id: Optional[int] = None
65
+ self.node_id: int | None = None
45
66
  # Store Object ID to MessageEntry mapping
46
67
  self.msg_store: dict[str, MessageEntry] = {}
47
- self.lock_msg_store = Lock()
68
+ self.lock_msg_store = RLock()
48
69
  # Store run ID to Run mapping
49
70
  self.run_store: dict[int, Run] = {}
50
71
  self.lock_run_store = Lock()
51
72
  # Store run ID to Context mapping
52
73
  self.ctx_store: dict[int, Context] = {}
53
74
  self.lock_ctx_store = Lock()
54
- # Store run ID to token mapping and token to run ID mapping
55
- self.token_store: dict[int, str] = {}
56
- self.token_to_run_id: dict[str, int] = {}
57
- self.lock_token_store = Lock()
75
+ # Store msg ID to TimeEntry mapping
76
+ self.time_store: dict[str, TimeEntry] = {}
77
+ self.lock_time_store = Lock()
58
78
 
59
- def set_node_id(self, node_id: Optional[int]) -> None:
79
+ def set_node_id(self, node_id: int | None) -> None:
60
80
  """Set the node ID."""
61
81
  self.node_id = node_id
62
82
 
@@ -66,8 +86,10 @@ class InMemoryNodeState(NodeState): # pylint: disable=too-many-instance-attribu
66
86
  raise ValueError("Node ID not set")
67
87
  return self.node_id
68
88
 
69
- def store_message(self, message: Message) -> Optional[str]:
89
+ def store_message(self, message: Message) -> str | None:
70
90
  """Store a message."""
91
+ # No need to check for expired tokens here
92
+ # The ClientAppIo servicer will first verify the token before storing messages
71
93
  with self.lock_msg_store:
72
94
  msg_id = message.metadata.message_id
73
95
  if msg_id == "" or msg_id in self.msg_store:
@@ -78,13 +100,14 @@ class InMemoryNodeState(NodeState): # pylint: disable=too-many-instance-attribu
78
100
  def get_messages(
79
101
  self,
80
102
  *,
81
- run_ids: Optional[Sequence[int]] = None,
82
- is_reply: Optional[bool] = None,
83
- limit: Optional[int] = None,
103
+ run_ids: Sequence[int] | None = None,
104
+ is_reply: bool | None = None,
105
+ limit: int | None = None,
84
106
  ) -> Sequence[Message]:
85
107
  """Retrieve messages based on the specified filters."""
86
- selected_messages: list[Message] = []
108
+ self._cleanup_expired_tokens()
87
109
 
110
+ selected_messages: list[Message] = []
88
111
  with self.lock_msg_store:
89
112
  # Iterate through all messages in the store
90
113
  for object_id in list(self.msg_store.keys()):
@@ -122,7 +145,7 @@ class InMemoryNodeState(NodeState): # pylint: disable=too-many-instance-attribu
122
145
  def delete_messages(
123
146
  self,
124
147
  *,
125
- message_ids: Optional[Sequence[str]] = None,
148
+ message_ids: Sequence[str] | None = None,
126
149
  ) -> None:
127
150
  """Delete messages based on the specified filters."""
128
151
  with self.lock_msg_store:
@@ -140,7 +163,7 @@ class InMemoryNodeState(NodeState): # pylint: disable=too-many-instance-attribu
140
163
  with self.lock_run_store:
141
164
  self.run_store[run.run_id] = run
142
165
 
143
- def get_run(self, run_id: int) -> Optional[Run]:
166
+ def get_run(self, run_id: int) -> Run | None:
144
167
  """Retrieve a run by its ID."""
145
168
  with self.lock_run_store:
146
169
  return self.run_store.get(run_id)
@@ -150,7 +173,7 @@ class InMemoryNodeState(NodeState): # pylint: disable=too-many-instance-attribu
150
173
  with self.lock_ctx_store:
151
174
  self.ctx_store[context.run_id] = context
152
175
 
153
- def get_context(self, run_id: int) -> Optional[Context]:
176
+ def get_context(self, run_id: int) -> Context | None:
154
177
  """Retrieve a context by its run ID."""
155
178
  with self.lock_ctx_store:
156
179
  return self.ctx_store.get(run_id)
@@ -171,29 +194,78 @@ class InMemoryNodeState(NodeState): # pylint: disable=too-many-instance-attribu
171
194
  ret -= set(self.token_store.keys())
172
195
  return list(ret)
173
196
 
174
- def create_token(self, run_id: int) -> Optional[str]:
175
- """Create a token for the given run ID."""
176
- token = secrets.token_hex(FLWR_APP_TOKEN_LENGTH) # Generate a random token
177
- with self.lock_token_store:
178
- if run_id in self.token_store:
179
- return None # Token already created for this run ID
180
- self.token_store[run_id] = token
181
- self.token_to_run_id[token] = run_id
182
- return token
183
-
184
- def verify_token(self, run_id: int, token: str) -> bool:
185
- """Verify a token for the given run ID."""
186
- with self.lock_token_store:
187
- return self.token_store.get(run_id) == token
188
-
189
- def delete_token(self, run_id: int) -> None:
190
- """Delete the token for the given run ID."""
191
- with self.lock_token_store:
192
- token = self.token_store.pop(run_id, None)
193
- if token is not None:
194
- self.token_to_run_id.pop(token, None)
195
-
196
- def get_run_id_by_token(self, token: str) -> Optional[int]:
197
- """Get the run ID associated with a given token."""
198
- with self.lock_token_store:
199
- return self.token_to_run_id.get(token)
197
+ def _on_tokens_expired(self, expired_records: list[tuple[int, float]]) -> None:
198
+ """Insert error replies for messages associated with expired tokens."""
199
+ with self.lock_msg_store:
200
+ # Find all retrieved messages associated with expired run IDs
201
+ expired_run_ids = {run_id for run_id, _ in expired_records}
202
+ messages_to_reply: list[Message] = []
203
+ for entry in self.msg_store.values():
204
+ msg = entry.message
205
+ if msg.metadata.run_id in expired_run_ids and entry.is_retrieved:
206
+ messages_to_reply.append(msg)
207
+
208
+ # Create and store error replies for each message
209
+ for msg in messages_to_reply:
210
+ error_reply = Message(CLIENT_APP_CRASHED_ERROR, reply_to=msg)
211
+
212
+ # Insert objects of the error reply into the object store
213
+ with no_object_id_recompute():
214
+ # pylint: disable-next=W0212
215
+ error_reply.metadata._message_id = error_reply.object_id # type: ignore
216
+ object_tree = get_object_tree(error_reply)
217
+ self.object_store.preregister(msg.metadata.run_id, object_tree)
218
+ for obj_id, obj in get_all_nested_objects(error_reply).items():
219
+ self.object_store.put(obj_id, obj.deflate())
220
+
221
+ # Store the error reply message
222
+ self.store_message(error_reply)
223
+
224
+ def record_message_processing_start(self, message_id: str) -> None:
225
+ """Record the start time of message processing based on the message ID."""
226
+ with self.lock_time_store:
227
+ self.time_store[message_id] = TimeEntry(starting_at=now().timestamp())
228
+
229
+ def record_message_processing_end(self, message_id: str) -> None:
230
+ """Record the end time of message processing based on the message ID."""
231
+ with self.lock_time_store:
232
+ if message_id not in self.time_store:
233
+ raise ValueError(
234
+ f"Cannot record end time: Message ID {message_id} not found."
235
+ )
236
+ entry = self.time_store[message_id]
237
+ entry.finished_at = now().timestamp()
238
+
239
+ def get_message_processing_duration(self, message_id: str) -> float:
240
+ """Get the message processing duration based on the message ID."""
241
+ # Cleanup old message processing times
242
+ self._cleanup_old_message_times()
243
+ with self.lock_time_store:
244
+ if message_id not in self.time_store:
245
+ raise ValueError(f"Message ID {message_id} not found.")
246
+
247
+ entry = self.time_store[message_id]
248
+ if entry.starting_at is None or entry.finished_at is None:
249
+ raise ValueError(
250
+ f"Start time or end time for message ID {message_id} is missing."
251
+ )
252
+
253
+ duration = entry.finished_at - entry.starting_at
254
+ return duration
255
+
256
+ def _cleanup_old_message_times(self) -> None:
257
+ """Remove time entries older than MESSAGE_TIME_ENTRY_MAX_AGE_SECONDS."""
258
+ with self.lock_time_store:
259
+ cutoff = now().timestamp() - MESSAGE_TIME_ENTRY_MAX_AGE_SECONDS
260
+ # Find message IDs for entries that have a finishing_at time
261
+ # before the cutoff, and those that don't exist in msg_store
262
+ to_delete = [
263
+ msg_id
264
+ for msg_id, entry in self.time_store.items()
265
+ if (entry.finished_at and entry.finished_at < cutoff)
266
+ or msg_id not in self.msg_store
267
+ ]
268
+
269
+ # Delete the identified entries
270
+ for msg_id in to_delete:
271
+ del self.time_store[msg_id]
@@ -17,7 +17,6 @@
17
17
 
18
18
  from abc import abstractmethod
19
19
  from collections.abc import Sequence
20
- from typing import Optional
21
20
 
22
21
  from flwr.common import Context, Message
23
22
  from flwr.common.typing import Run
@@ -36,7 +35,7 @@ class NodeState(CoreState):
36
35
  """Get the node ID."""
37
36
 
38
37
  @abstractmethod
39
- def store_message(self, message: Message) -> Optional[str]:
38
+ def store_message(self, message: Message) -> str | None:
40
39
  """Store a message.
41
40
 
42
41
  Parameters
@@ -54,9 +53,9 @@ class NodeState(CoreState):
54
53
  def get_messages(
55
54
  self,
56
55
  *,
57
- run_ids: Optional[Sequence[int]] = None,
58
- is_reply: Optional[bool] = None,
59
- limit: Optional[int] = None,
56
+ run_ids: Sequence[int] | None = None,
57
+ is_reply: bool | None = None,
58
+ limit: int | None = None,
60
59
  ) -> Sequence[Message]:
61
60
  """Retrieve messages based on the specified filters.
62
61
 
@@ -89,7 +88,7 @@ class NodeState(CoreState):
89
88
  def delete_messages(
90
89
  self,
91
90
  *,
92
- message_ids: Optional[Sequence[str]] = None,
91
+ message_ids: Sequence[str] | None = None,
93
92
  ) -> None:
94
93
  """Delete messages based on the specified filters.
95
94
 
@@ -118,7 +117,7 @@ class NodeState(CoreState):
118
117
  """
119
118
 
120
119
  @abstractmethod
121
- def get_run(self, run_id: int) -> Optional[Run]:
120
+ def get_run(self, run_id: int) -> Run | None:
122
121
  """Retrieve a run by its ID.
123
122
 
124
123
  Parameters
@@ -143,7 +142,7 @@ class NodeState(CoreState):
143
142
  """
144
143
 
145
144
  @abstractmethod
146
- def get_context(self, run_id: int) -> Optional[Context]:
145
+ def get_context(self, run_id: int) -> Context | None:
147
146
  """Retrieve a context by its run ID.
148
147
 
149
148
  Parameters
@@ -169,3 +168,48 @@ class NodeState(CoreState):
169
168
  Sequence[int]
170
169
  Sequence of run IDs with pending messages.
171
170
  """
171
+
172
+ @abstractmethod
173
+ def record_message_processing_start(self, message_id: str) -> None:
174
+ """Record the start time of message processing based on the message ID.
175
+
176
+ Parameters
177
+ ----------
178
+ message_id : str
179
+ The ID of the message associated with the start time.
180
+ """
181
+
182
+ @abstractmethod
183
+ def record_message_processing_end(self, message_id: str) -> None:
184
+ """Record the end time of message processing based on the message ID.
185
+
186
+ Parameters
187
+ ----------
188
+ message_id : str
189
+ The ID of the message associated with the end time.
190
+
191
+ Raises
192
+ ------
193
+ ValueError
194
+ If the message ID is not found.
195
+ """
196
+
197
+ @abstractmethod
198
+ def get_message_processing_duration(self, message_id: str) -> float:
199
+ """Get the message processing duration based on the message ID.
200
+
201
+ Parameters
202
+ ----------
203
+ message_id : str
204
+ The ID of the message.
205
+
206
+ Returns
207
+ -------
208
+ float
209
+ The processing duration in seconds.
210
+
211
+ Raises
212
+ ------
213
+ ValueError
214
+ If the message ID is not found, or if start/end times are missing.
215
+ """
@@ -16,7 +16,8 @@
16
16
 
17
17
 
18
18
  import threading
19
- from typing import Optional
19
+
20
+ from flwr.supercore.object_store import ObjectStoreFactory
20
21
 
21
22
  from .in_memory_nodestate import InMemoryNodeState
22
23
  from .nodestate import NodeState
@@ -25,8 +26,9 @@ from .nodestate import NodeState
25
26
  class NodeStateFactory:
26
27
  """Factory class that creates NodeState instances."""
27
28
 
28
- def __init__(self) -> None:
29
- self.state_instance: Optional[NodeState] = None
29
+ def __init__(self, objectstore_factory: ObjectStoreFactory) -> None:
30
+ self.objectstore_factory = objectstore_factory
31
+ self.state_instance: NodeState | None = None
30
32
  self.lock = threading.RLock()
31
33
 
32
34
  def state(self) -> NodeState:
@@ -34,5 +36,6 @@ class NodeStateFactory:
34
36
  # Lock access to NodeStateFactory to prevent returning different instances
35
37
  with self.lock:
36
38
  if self.state_instance is None:
37
- self.state_instance = InMemoryNodeState()
39
+ object_store = self.objectstore_factory.store()
40
+ self.state_instance = InMemoryNodeState(object_store)
38
41
  return self.state_instance
@@ -15,9 +15,7 @@
15
15
  """Flower ClientApp process."""
16
16
 
17
17
 
18
- import gc
19
18
  from logging import DEBUG, ERROR, INFO
20
- from typing import Optional
21
19
 
22
20
  import grpc
23
21
 
@@ -28,6 +26,7 @@ from flwr.clientapp.utils import get_load_client_app_fn
28
26
  from flwr.common import Context, Message
29
27
  from flwr.common.config import get_flwr_dir
30
28
  from flwr.common.constant import ErrorCode
29
+ from flwr.common.exit import ExitCode, flwr_exit, register_signal_handlers
31
30
  from flwr.common.grpc import create_channel, on_channel_state_change
32
31
  from flwr.common.inflatable import (
33
32
  get_all_nested_objects,
@@ -50,6 +49,7 @@ from flwr.common.serde import (
50
49
  message_to_proto,
51
50
  run_from_proto,
52
51
  )
52
+ from flwr.common.telemetry import EventType, event
53
53
  from flwr.common.typing import Fab, Run
54
54
  from flwr.proto.appio_pb2 import ( # pylint: disable=E0611
55
55
  PullAppInputsRequest,
@@ -63,27 +63,41 @@ from flwr.proto.appio_pb2 import ( # pylint: disable=E0611
63
63
  from flwr.proto.clientappio_pb2_grpc import ClientAppIoStub
64
64
  from flwr.proto.node_pb2 import Node # pylint: disable=E0611
65
65
  from flwr.supercore.app_utils import start_parent_process_monitor
66
+ from flwr.supercore.heartbeat import HeartbeatSender, make_app_heartbeat_fn_grpc
66
67
  from flwr.supercore.utils import mask_string
67
68
 
68
69
 
69
70
  def run_clientapp( # pylint: disable=R0913, R0914, R0917
70
71
  clientappio_api_address: str,
71
72
  token: str,
72
- flwr_dir: Optional[str] = None,
73
- certificates: Optional[bytes] = None,
74
- parent_pid: Optional[int] = None,
73
+ flwr_dir: str | None = None,
74
+ certificates: bytes | None = None,
75
+ parent_pid: int | None = None,
75
76
  ) -> None:
76
77
  """Run Flower ClientApp process."""
77
78
  # Monitor the main process in case of SIGKILL
78
79
  if parent_pid is not None:
79
80
  start_parent_process_monitor(parent_pid)
80
81
 
82
+ event(EventType.FLWR_CLIENTAPP_RUN_ENTER)
83
+
81
84
  channel = create_channel(
82
85
  server_address=clientappio_api_address,
83
86
  insecure=(certificates is None),
84
87
  root_certificates=certificates,
85
88
  )
86
89
  channel.subscribe(on_channel_state_change)
90
+ heartbeat_sender = None
91
+
92
+ def on_exit() -> None:
93
+ if heartbeat_sender is not None and heartbeat_sender.is_running:
94
+ heartbeat_sender.stop()
95
+ channel.close()
96
+
97
+ register_signal_handlers(
98
+ event_type=EventType.FLWR_CLIENTAPP_RUN_LEAVE,
99
+ exit_handlers=[on_exit],
100
+ )
87
101
 
88
102
  # Resolve directory where FABs are installed
89
103
  flwr_dir_ = get_flwr_dir(flwr_dir)
@@ -91,22 +105,27 @@ def run_clientapp( # pylint: disable=R0913, R0914, R0917
91
105
  stub = ClientAppIoStub(channel)
92
106
  _wrap_stub(stub, _make_simple_grpc_retry_invoker())
93
107
 
108
+ # Start app heartbeat
109
+ heartbeat_sender = HeartbeatSender(make_app_heartbeat_fn_grpc(stub, token))
110
+ heartbeat_sender.start()
111
+
94
112
  # Pull Message, Context, Run and (optional) FAB from SuperNode
95
113
  message, context, run, fab = pull_clientappinputs(stub=stub, token=token)
96
114
 
97
- # Install FAB, if provided
98
- if fab:
99
- log(DEBUG, "[flwr-clientapp] Start FAB installation.")
100
- install_from_fab(fab.content, flwr_dir=flwr_dir_, skip_prompt=True)
115
+ try:
101
116
 
102
- load_client_app_fn = get_load_client_app_fn(
103
- default_app_ref="",
104
- app_path=None,
105
- multi_app=True,
106
- flwr_dir=str(flwr_dir_),
107
- )
117
+ # Install FAB, if provided
118
+ if fab:
119
+ log(DEBUG, "[flwr-clientapp] Start FAB installation.")
120
+ install_from_fab(fab.content, flwr_dir=flwr_dir_, skip_prompt=True)
121
+
122
+ load_client_app_fn = get_load_client_app_fn(
123
+ default_app_ref="",
124
+ app_path=None,
125
+ multi_app=True,
126
+ flwr_dir=str(flwr_dir_),
127
+ )
108
128
 
109
- try:
110
129
  # Load ClientApp
111
130
  log(DEBUG, "[flwr-clientapp] Start `ClientApp` Loading.")
112
131
  client_app: ClientApp = load_client_app_fn(
@@ -137,18 +156,18 @@ def run_clientapp( # pylint: disable=R0913, R0914, R0917
137
156
  stub=stub, token=token, message=reply_message, context=context
138
157
  )
139
158
 
140
- del client_app, message, context, run, fab, reply_message
141
- gc.collect()
142
-
143
159
  except grpc.RpcError as e:
144
160
  log(ERROR, "GRPC error occurred: %s", str(e))
145
- finally:
146
- channel.close()
161
+
162
+ flwr_exit(
163
+ code=ExitCode.SUCCESS,
164
+ event_type=EventType.FLWR_CLIENTAPP_RUN_LEAVE,
165
+ )
147
166
 
148
167
 
149
168
  def pull_clientappinputs(
150
169
  stub: ClientAppIoStub, token: str
151
- ) -> tuple[Message, Context, Run, Optional[Fab]]:
170
+ ) -> tuple[Message, Context, Run, Fab | None]:
152
171
  """Pull ClientAppInputs from SuperNode."""
153
172
  masked_token = mask_string(token)
154
173
  log(INFO, "[flwr-clientapp] Pull `ClientAppInputs` for token %s", masked_token)
@@ -35,7 +35,7 @@ from flwr.common.typing import Fab, Run
35
35
 
36
36
  # pylint: disable=E0611
37
37
  from flwr.proto import clientappio_pb2_grpc
38
- from flwr.proto.appio_pb2 import ( # pylint: disable=E0401
38
+ from flwr.proto.appio_pb2 import (
39
39
  ListAppsToLaunchRequest,
40
40
  ListAppsToLaunchResponse,
41
41
  PullAppInputsRequest,
@@ -49,6 +49,7 @@ from flwr.proto.appio_pb2 import ( # pylint: disable=E0401
49
49
  RequestTokenRequest,
50
50
  RequestTokenResponse,
51
51
  )
52
+ from flwr.proto.heartbeat_pb2 import SendAppHeartbeatRequest, SendAppHeartbeatResponse
52
53
  from flwr.proto.message_pb2 import (
53
54
  ConfirmMessageReceivedRequest,
54
55
  ConfirmMessageReceivedResponse,
@@ -57,12 +58,11 @@ from flwr.proto.message_pb2 import (
57
58
  PushObjectRequest,
58
59
  PushObjectResponse,
59
60
  )
60
- from flwr.proto.run_pb2 import GetRunRequest, GetRunResponse # pylint: disable=E0611
61
+ from flwr.proto.run_pb2 import GetRunRequest, GetRunResponse
61
62
 
62
63
  # pylint: disable=E0601
63
64
  from flwr.supercore.ffs import FfsFactory
64
65
  from flwr.supercore.object_store import NoObjectInStoreError, ObjectStoreFactory
65
- from flwr.supercore.object_store.utils import store_mapping_and_register_objects
66
66
  from flwr.supernode.nodestate import NodeStateFactory
67
67
 
68
68
 
@@ -151,7 +151,24 @@ class ClientAppIoServicer(clientappio_pb2_grpc.ClientAppIoServicer):
151
151
  # Retrieve context, run and fab for this run
152
152
  context = cast(Context, state.get_context(run_id))
153
153
  run = cast(Run, state.get_run(run_id))
154
- fab = Fab(run.fab_hash, ffs.get(run.fab_hash)[0], ffs.get(run.fab_hash)[1]) # type: ignore
154
+
155
+ # Retrieve FAB from FFS
156
+ if result := ffs.get(run.fab_hash):
157
+ content, verifications = result
158
+ log(
159
+ DEBUG,
160
+ "Retrieved FAB: hash=%s, content_len=%d, verifications=%s",
161
+ run.fab_hash,
162
+ len(content),
163
+ verifications,
164
+ )
165
+ fab = Fab(run.fab_hash, content, verifications)
166
+ else:
167
+ context.abort(
168
+ grpc.StatusCode.NOT_FOUND,
169
+ f"FAB with hash {run.fab_hash} not found in FFS.",
170
+ )
171
+ raise RuntimeError("This line should never be reached.")
155
172
 
156
173
  return PullAppInputsResponse(
157
174
  context=context_to_proto(context),
@@ -206,6 +223,9 @@ class ClientAppIoServicer(clientappio_pb2_grpc.ClientAppIoServicer):
206
223
  # Retrieve message for this run
207
224
  message = state.get_messages(run_ids=[run_id], is_reply=False)[0]
208
225
 
226
+ # Record message processing start time
227
+ state.record_message_processing_start(message_id=message.metadata.message_id)
228
+
209
229
  # Retrieve the object tree for the message
210
230
  object_tree = store.get_object_tree(message.metadata.message_id)
211
231
 
@@ -231,19 +251,35 @@ class ClientAppIoServicer(clientappio_pb2_grpc.ClientAppIoServicer):
231
251
  )
232
252
  raise RuntimeError("This line should never be reached.")
233
253
 
254
+ # Store Message object to descendants mapping and preregister objects
255
+ objects_to_push: set[str] = set()
256
+ for object_tree in request.message_object_trees:
257
+ objects_to_push |= set(store.preregister(run_id, object_tree))
234
258
  # Save the message to the state
235
259
  state.store_message(message_from_proto(request.messages_list[0]))
260
+ # Record message processing end time
261
+ state.record_message_processing_end(
262
+ message_id=request.messages_list[0].metadata.reply_to_message_id
263
+ )
264
+ return PushAppMessagesResponse(objects_to_push=objects_to_push)
236
265
 
237
- # Store Message object to descendants mapping and preregister objects
238
- objects_to_push = store_mapping_and_register_objects(store, request=request)
266
+ def SendAppHeartbeat(
267
+ self, request: SendAppHeartbeatRequest, context: grpc.ServicerContext
268
+ ) -> SendAppHeartbeatResponse:
269
+ """Handle a heartbeat from an app process."""
270
+ log(DEBUG, "ClientAppIoServicer.SendAppHeartbeat")
271
+ # Initialize state
272
+ state = self.state_factory.state()
239
273
 
240
- return PushAppMessagesResponse(objects_to_push=objects_to_push)
274
+ # Acknowledge the heartbeat
275
+ success = state.acknowledge_app_heartbeat(request.token)
276
+ return SendAppHeartbeatResponse(success=success)
241
277
 
242
278
  def PushObject(
243
279
  self, request: PushObjectRequest, context: grpc.ServicerContext
244
280
  ) -> PushObjectResponse:
245
281
  """Push an object to the ObjectStore."""
246
- log(DEBUG, "ServerAppIoServicer.PushObject")
282
+ log(DEBUG, "ClientAppIoServicer.PushObject")
247
283
 
248
284
  # Init state and store
249
285
  store = self.objectstore_factory.store()
@@ -265,7 +301,7 @@ class ClientAppIoServicer(clientappio_pb2_grpc.ClientAppIoServicer):
265
301
  self, request: PullObjectRequest, context: grpc.ServicerContext
266
302
  ) -> PullObjectResponse:
267
303
  """Pull an object from the ObjectStore."""
268
- log(DEBUG, "ServerAppIoServicer.PullObject")
304
+ log(DEBUG, "ClientAppIoServicer.PullObject")
269
305
 
270
306
  # Init state and store
271
307
  store = self.objectstore_factory.store()
@@ -285,7 +321,7 @@ class ClientAppIoServicer(clientappio_pb2_grpc.ClientAppIoServicer):
285
321
  self, request: ConfirmMessageReceivedRequest, context: grpc.ServicerContext
286
322
  ) -> ConfirmMessageReceivedResponse:
287
323
  """Confirm message received."""
288
- log(DEBUG, "ServerAppIoServicer.ConfirmMessageReceived")
324
+ log(DEBUG, "ClientAppIoServicer.ConfirmMessageReceived")
289
325
 
290
326
  # Init state and store
291
327
  store = self.objectstore_factory.store()