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,404 @@
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 simulation app."""
16
+
17
+
18
+ import asyncio
19
+ import logging
20
+ import sys
21
+ import threading
22
+ import traceback
23
+ import warnings
24
+ from logging import ERROR, INFO
25
+ from typing import Any, Optional, Union
26
+
27
+ import ray
28
+ from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy
29
+
30
+ from flwr.client import ClientFnExt
31
+ from flwr.common import EventType, event
32
+ from flwr.common.constant import NODE_ID_NUM_BYTES, SUPERLINK_NODE_ID
33
+ from flwr.common.logger import (
34
+ log,
35
+ set_logger_propagation,
36
+ warn_deprecated_feature,
37
+ warn_unsupported_feature,
38
+ )
39
+ from flwr.server.client_manager import ClientManager
40
+ from flwr.server.history import History
41
+ from flwr.server.server import Server, init_defaults, run_fl
42
+ from flwr.server.server_config import ServerConfig
43
+ from flwr.server.strategy import Strategy
44
+ from flwr.server.superlink.linkstate.utils import generate_rand_int_from_bytes
45
+ from flwr.simulation.ray_transport.ray_actor import (
46
+ ClientAppActor,
47
+ VirtualClientEngineActor,
48
+ VirtualClientEngineActorPool,
49
+ pool_size_from_resources,
50
+ )
51
+ from flwr.simulation.ray_transport.ray_client_proxy import RayActorClientProxy
52
+
53
+ INVALID_ARGUMENTS_START_SIMULATION = """
54
+ INVALID ARGUMENTS ERROR
55
+
56
+ Invalid Arguments in method:
57
+
58
+ `start_simulation(
59
+ *,
60
+ client_fn: ClientFn,
61
+ num_clients: int,
62
+ clients_ids: Optional[List[str]] = None,
63
+ client_resources: Optional[Dict[str, float]] = None,
64
+ server: Optional[Server] = None,
65
+ config: ServerConfig = None,
66
+ strategy: Optional[Strategy] = None,
67
+ client_manager: Optional[ClientManager] = None,
68
+ ray_init_args: Optional[Dict[str, Any]] = None,
69
+ ) -> None:`
70
+
71
+ REASON:
72
+ Method requires:
73
+ - Either `num_clients`[int] or `clients_ids`[List[str]]
74
+ to be set exclusively.
75
+ OR
76
+ - `len(clients_ids)` == `num_clients`
77
+
78
+ """
79
+
80
+ NodeToPartitionMapping = dict[int, int]
81
+
82
+
83
+ def _create_node_id_to_partition_mapping(
84
+ num_clients: int,
85
+ ) -> NodeToPartitionMapping:
86
+ """Generate a node_id:partition_id mapping."""
87
+ nodes_mapping: NodeToPartitionMapping = {} # {node-id; partition-id}
88
+ for i in range(num_clients):
89
+ while True:
90
+ node_id = generate_rand_int_from_bytes(
91
+ NODE_ID_NUM_BYTES, exclude=[SUPERLINK_NODE_ID, 0]
92
+ )
93
+ if node_id not in nodes_mapping:
94
+ break
95
+ nodes_mapping[node_id] = i
96
+ return nodes_mapping
97
+
98
+
99
+ # pylint: disable=too-many-arguments,too-many-statements,too-many-branches
100
+ def start_simulation(
101
+ *,
102
+ client_fn: ClientFnExt,
103
+ num_clients: int,
104
+ clients_ids: Optional[list[str]] = None, # UNSUPPORTED, WILL BE REMOVED
105
+ client_resources: Optional[dict[str, float]] = None,
106
+ server: Optional[Server] = None,
107
+ config: Optional[ServerConfig] = None,
108
+ strategy: Optional[Strategy] = None,
109
+ client_manager: Optional[ClientManager] = None,
110
+ ray_init_args: Optional[dict[str, Any]] = None,
111
+ keep_initialised: Optional[bool] = False,
112
+ actor_type: type[VirtualClientEngineActor] = ClientAppActor,
113
+ actor_kwargs: Optional[dict[str, Any]] = None,
114
+ actor_scheduling: Union[str, NodeAffinitySchedulingStrategy] = "DEFAULT",
115
+ ) -> History:
116
+ """Start a Ray-based Flower simulation server.
117
+
118
+ Warning
119
+ -------
120
+ This function is deprecated since 1.13.0. Use :code: `flwr run` to start a Flower
121
+ simulation.
122
+
123
+ Parameters
124
+ ----------
125
+ client_fn : ClientFnExt
126
+ A function creating `Client` instances. The function must have the signature
127
+ `client_fn(context: Context). It should return
128
+ a single client instance of type `Client`. Note that the created client
129
+ instances are ephemeral and will often be destroyed after a single method
130
+ invocation. Since client instances are not long-lived, they should not attempt
131
+ to carry state over method invocations. Any state required by the instance
132
+ (model, dataset, hyperparameters, ...) should be (re-)created in either the
133
+ call to `client_fn` or the call to any of the client methods (e.g., load
134
+ evaluation data in the `evaluate` method itself).
135
+ num_clients : int
136
+ The total number of clients in this simulation.
137
+ clients_ids : Optional[List[str]]
138
+ UNSUPPORTED, WILL BE REMOVED. USE `num_clients` INSTEAD.
139
+ List `client_id`s for each client. This is only required if
140
+ `num_clients` is not set. Setting both `num_clients` and `clients_ids`
141
+ with `len(clients_ids)` not equal to `num_clients` generates an error.
142
+ Using this argument will raise an error.
143
+ client_resources : Optional[Dict[str, float]] (default: `{"num_cpus": 1, "num_gpus": 0.0}`)
144
+ CPU and GPU resources for a single client. Supported keys
145
+ are `num_cpus` and `num_gpus`. To understand the GPU utilization caused by
146
+ `num_gpus`, as well as using custom resources, please consult the Ray
147
+ documentation.
148
+ server : Optional[flwr.server.Server] (default: None).
149
+ An implementation of the abstract base class `flwr.server.Server`. If no
150
+ instance is provided, then `start_server` will create one.
151
+ config: ServerConfig (default: None).
152
+ Currently supported values are `num_rounds` (int, default: 1) and
153
+ `round_timeout` in seconds (float, default: None).
154
+ strategy : Optional[flwr.server.Strategy] (default: None)
155
+ An implementation of the abstract base class `flwr.server.Strategy`. If
156
+ no strategy is provided, then `start_server` will use
157
+ `flwr.server.strategy.FedAvg`.
158
+ client_manager : Optional[flwr.server.ClientManager] (default: None)
159
+ An implementation of the abstract base class `flwr.server.ClientManager`.
160
+ If no implementation is provided, then `start_simulation` will use
161
+ `flwr.server.client_manager.SimpleClientManager`.
162
+ ray_init_args : Optional[Dict[str, Any]] (default: None)
163
+ Optional dictionary containing arguments for the call to `ray.init`.
164
+ If ray_init_args is None (the default), Ray will be initialized with
165
+ the following default args:
166
+
167
+ { "ignore_reinit_error": True, "include_dashboard": False }
168
+
169
+ An empty dictionary can be used (ray_init_args={}) to prevent any
170
+ arguments from being passed to ray.init.
171
+ keep_initialised: Optional[bool] (default: False)
172
+ Set to True to prevent `ray.shutdown()` in case `ray.is_initialized()=True`.
173
+
174
+ actor_type: VirtualClientEngineActor (default: ClientAppActor)
175
+ Optionally specify the type of actor to use. The actor object, which
176
+ persists throughout the simulation, will be the process in charge of
177
+ executing a ClientApp wrapping input argument `client_fn`.
178
+
179
+ actor_kwargs: Optional[Dict[str, Any]] (default: None)
180
+ If you want to create your own Actor classes, you might need to pass
181
+ some input argument. You can use this dictionary for such purpose.
182
+
183
+ actor_scheduling: Optional[Union[str, NodeAffinitySchedulingStrategy]]
184
+ (default: "DEFAULT")
185
+ Optional string ("DEFAULT" or "SPREAD") for the VCE to choose in which
186
+ node the actor is placed. If you are an advanced user needed more control
187
+ you can use lower-level scheduling strategies to pin actors to specific
188
+ compute nodes (e.g. via NodeAffinitySchedulingStrategy). Please note this
189
+ is an advanced feature. For all details, please refer to the Ray documentation:
190
+ https://docs.ray.io/en/latest/ray-core/scheduling/index.html
191
+
192
+ Returns
193
+ -------
194
+ hist : flwr.server.history.History
195
+ Object containing metrics from training.
196
+ """ # noqa: E501
197
+ # pylint: disable-msg=too-many-locals
198
+ msg = (
199
+ "flwr.simulation.start_simulation() is deprecated."
200
+ "\n\tInstead, use the `flwr run` CLI command to start a local simulation "
201
+ "in your Flower app, as shown for example below:"
202
+ "\n\n\t\t$ flwr new # Create a new Flower app from a template"
203
+ "\n\n\t\t$ flwr run # Run the Flower app in Simulation Mode"
204
+ "\n\n\tUsing `start_simulation()` is deprecated."
205
+ )
206
+ warn_deprecated_feature(name=msg)
207
+
208
+ event(
209
+ EventType.START_SIMULATION_ENTER,
210
+ {"num_clients": len(clients_ids) if clients_ids is not None else num_clients},
211
+ )
212
+
213
+ if clients_ids is not None:
214
+ warn_unsupported_feature(
215
+ "Passing `clients_ids` to `start_simulation` is deprecated and not longer "
216
+ "used by `start_simulation`. Use `num_clients` exclusively instead."
217
+ )
218
+ log(ERROR, "`clients_ids` argument used.")
219
+ sys.exit()
220
+
221
+ # Set logger propagation
222
+ loop: Optional[asyncio.AbstractEventLoop] = None
223
+ try:
224
+ loop = asyncio.get_running_loop()
225
+ except RuntimeError:
226
+ loop = None
227
+ finally:
228
+ if loop and loop.is_running():
229
+ # Set logger propagation to False to prevent duplicated log output in Colab.
230
+ logger = logging.getLogger("flwr")
231
+ _ = set_logger_propagation(logger, False)
232
+
233
+ # Initialize server and server config
234
+ initialized_server, initialized_config = init_defaults(
235
+ server=server,
236
+ config=config,
237
+ strategy=strategy,
238
+ client_manager=client_manager,
239
+ )
240
+
241
+ log(
242
+ INFO,
243
+ "Starting Flower simulation, config: %s",
244
+ initialized_config,
245
+ )
246
+
247
+ # Create node-id to partition-id mapping
248
+ nodes_mapping = _create_node_id_to_partition_mapping(num_clients)
249
+
250
+ # Default arguments for Ray initialization
251
+ if not ray_init_args:
252
+ ray_init_args = {
253
+ "ignore_reinit_error": True,
254
+ "include_dashboard": False,
255
+ }
256
+
257
+ # Shut down Ray if it has already been initialized (unless asked not to)
258
+ if ray.is_initialized() and not keep_initialised:
259
+ ray.shutdown()
260
+
261
+ # Initialize Ray
262
+ ray.init(**ray_init_args)
263
+ cluster_resources = ray.cluster_resources()
264
+ log(
265
+ INFO,
266
+ "Flower VCE: Ray initialized with resources: %s",
267
+ cluster_resources,
268
+ )
269
+
270
+ log(
271
+ INFO,
272
+ "Optimize your simulation with Flower VCE: "
273
+ "https://flower.ai/docs/framework/how-to-run-simulations.html",
274
+ )
275
+
276
+ # Log the resources that a single client will be able to use
277
+ if client_resources is None:
278
+ log(
279
+ INFO,
280
+ "No `client_resources` specified. Using minimal resources for clients.",
281
+ )
282
+ client_resources = {"num_cpus": 1, "num_gpus": 0.0}
283
+
284
+ # Each client needs at the very least one CPU
285
+ if "num_cpus" not in client_resources:
286
+ warnings.warn(
287
+ "No `num_cpus` specified in `client_resources`. "
288
+ "Using `num_cpus=1` for each client.",
289
+ stacklevel=2,
290
+ )
291
+ client_resources["num_cpus"] = 1
292
+
293
+ log(
294
+ INFO,
295
+ "Flower VCE: Resources for each Virtual Client: %s",
296
+ client_resources,
297
+ )
298
+
299
+ actor_args = {} if actor_kwargs is None else actor_kwargs
300
+
301
+ # An actor factory. This is called N times to add N actors
302
+ # to the pool. If at some point the pool can accommodate more actors
303
+ # this will be called again.
304
+ def create_actor_fn() -> type[VirtualClientEngineActor]:
305
+ return actor_type.options( # type: ignore
306
+ **client_resources,
307
+ scheduling_strategy=actor_scheduling,
308
+ ).remote(**actor_args)
309
+
310
+ # Instantiate ActorPool
311
+ pool = VirtualClientEngineActorPool(
312
+ create_actor_fn=create_actor_fn,
313
+ client_resources=client_resources,
314
+ )
315
+
316
+ f_stop = threading.Event()
317
+
318
+ # Periodically, check if the cluster has grown (i.e. a new
319
+ # node has been added). If this happens, we likely want to grow
320
+ # the actor pool by adding more Actors to it.
321
+ def update_resources(f_stop: threading.Event) -> None:
322
+ """Periodically check if more actors can be added to the pool.
323
+
324
+ If so, extend the pool.
325
+ """
326
+ if not f_stop.is_set():
327
+ num_max_actors = pool_size_from_resources(client_resources)
328
+ if num_max_actors > pool.num_actors:
329
+ num_new = num_max_actors - pool.num_actors
330
+ log(
331
+ INFO, "The cluster expanded. Adding %s actors to the pool.", num_new
332
+ )
333
+ pool.add_actors_to_pool(num_actors=num_new)
334
+
335
+ threading.Timer(10, update_resources, [f_stop]).start()
336
+
337
+ update_resources(f_stop)
338
+
339
+ log(
340
+ INFO,
341
+ "Flower VCE: Creating %s with %s actors",
342
+ pool.__class__.__name__,
343
+ pool.num_actors,
344
+ )
345
+
346
+ # Register one RayClientProxy object for each client with the ClientManager
347
+ for node_id, partition_id in nodes_mapping.items():
348
+ client_proxy = RayActorClientProxy(
349
+ client_fn=client_fn,
350
+ node_id=node_id,
351
+ partition_id=partition_id,
352
+ num_partitions=num_clients,
353
+ actor_pool=pool,
354
+ )
355
+ initialized_server.client_manager().register(client=client_proxy)
356
+
357
+ hist = History()
358
+ # pylint: disable=broad-except
359
+ try:
360
+ # Start training
361
+ hist = run_fl(
362
+ server=initialized_server,
363
+ config=initialized_config,
364
+ )
365
+ except Exception as ex:
366
+ log(ERROR, ex)
367
+ log(ERROR, traceback.format_exc())
368
+ log(
369
+ ERROR,
370
+ "Your simulation crashed :(. This could be because of several reasons. "
371
+ "The most common are: "
372
+ "\n\t > Sometimes, issues in the simulation code itself can cause crashes. "
373
+ "It's always a good idea to double-check your code for any potential bugs "
374
+ "or inconsistencies that might be contributing to the problem. "
375
+ "For example: "
376
+ "\n\t\t - You might be using a class attribute in your clients that "
377
+ "hasn't been defined."
378
+ "\n\t\t - There could be an incorrect method call to a 3rd party library "
379
+ "(e.g., PyTorch)."
380
+ "\n\t\t - The return types of methods in your clients/strategies might be "
381
+ "incorrect."
382
+ "\n\t > Your system couldn't fit a single VirtualClient: try lowering "
383
+ "`client_resources`."
384
+ "\n\t > All the actors in your pool crashed. This could be because: "
385
+ "\n\t\t - You clients hit an out-of-memory (OOM) error and actors couldn't "
386
+ "recover from it. Try launching your simulation with more generous "
387
+ "`client_resources` setting (i.e. it seems %s is "
388
+ "not enough for your run). Use fewer concurrent actors. "
389
+ "\n\t\t - You were running a multi-node simulation and all worker nodes "
390
+ "disconnected. The head node might still be alive but cannot accommodate "
391
+ "any actor with resources: %s."
392
+ "\nTake a look at the Flower simulation examples for guidance "
393
+ "<https://flower.ai/docs/framework/how-to-run-simulations.html>.",
394
+ client_resources,
395
+ client_resources,
396
+ )
397
+ raise RuntimeError("Simulation crashed.") from ex
398
+
399
+ finally:
400
+ # Stop time monitoring resources in cluster
401
+ f_stop.set()
402
+ event(EventType.START_SIMULATION_LEAVE)
403
+
404
+ return hist
@@ -1,4 +1,4 @@
1
- # Copyright 2020 Flower Labs GmbH. All Rights Reserved.
1
+ # Copyright 2021 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.
@@ -14,33 +14,23 @@
14
14
  # ==============================================================================
15
15
  """Ray-based Flower Actor and ActorPool implementation."""
16
16
 
17
- import asyncio
17
+
18
18
  import threading
19
- import traceback
20
19
  from abc import ABC
21
20
  from logging import DEBUG, ERROR, WARNING
22
- from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Type, Union
21
+ from typing import Any, Callable, Optional, Union
23
22
 
24
23
  import ray
25
24
  from ray import ObjectRef
26
25
  from ray.util.actor_pool import ActorPool
27
26
 
28
- from flwr.client.client_app import ClientApp, LoadClientAppError
27
+ from flwr.client.client_app import ClientApp, ClientAppException, LoadClientAppError
29
28
  from flwr.common import Context, Message
30
29
  from flwr.common.logger import log
31
30
 
32
31
  ClientAppFn = Callable[[], ClientApp]
33
32
 
34
33
 
35
- class ClientException(Exception):
36
- """Raised when client side logic crashes with an exception."""
37
-
38
- def __init__(self, message: str):
39
- div = ">" * 7
40
- self.message = "\n" + div + "A ClientException occurred." + message
41
- super().__init__(self.message)
42
-
43
-
44
34
  class VirtualClientEngineActor(ABC):
45
35
  """Abstract base class for VirtualClientEngine Actors."""
46
36
 
@@ -55,7 +45,7 @@ class VirtualClientEngineActor(ABC):
55
45
  message: Message,
56
46
  cid: str,
57
47
  context: Context,
58
- ) -> Tuple[str, Message, Context]:
48
+ ) -> tuple[str, Message, Context]:
59
49
  """Run a client run."""
60
50
  # Pass message through ClientApp and return a message
61
51
  # return also cid which is needed to ensure results
@@ -71,17 +61,7 @@ class VirtualClientEngineActor(ABC):
71
61
  raise load_ex
72
62
 
73
63
  except Exception as ex:
74
- client_trace = traceback.format_exc()
75
- mssg = (
76
- "\n\tSomething went wrong when running your client run."
77
- "\n\tClient "
78
- + cid
79
- + " crashed when the "
80
- + self.__class__.__name__
81
- + " was running its run."
82
- "\n\tException triggered on the client side: " + client_trace,
83
- )
84
- raise ClientException(str(mssg)) from ex
64
+ raise ClientAppException(str(ex)) from ex
85
65
 
86
66
  return cid, out_message, context
87
67
 
@@ -102,7 +82,7 @@ class ClientAppActor(VirtualClientEngineActor):
102
82
  on_actor_init_fn()
103
83
 
104
84
 
105
- def pool_size_from_resources(client_resources: Dict[str, Union[int, float]]) -> int:
85
+ def pool_size_from_resources(client_resources: dict[str, Union[int, float]]) -> int:
106
86
  """Calculate number of Actors that fit in the cluster.
107
87
 
108
88
  For this we consider the resources available on each node and those required per
@@ -145,14 +125,14 @@ def pool_size_from_resources(client_resources: Dict[str, Union[int, float]]) ->
145
125
  WARNING,
146
126
  "The ActorPool is empty. The system (CPUs=%s, GPUs=%s) "
147
127
  "does not meet the criteria to host at least one client with resources:"
148
- " %s. Lowering the `client_resources` could help.",
128
+ " %s. Lowering these resources could help.",
149
129
  num_cpus,
150
130
  num_gpus,
151
131
  client_resources,
152
132
  )
153
133
  raise ValueError(
154
134
  "ActorPool is empty. Stopping Simulation. "
155
- "Check 'client_resources' passed to `start_simulation`"
135
+ "Check `num_cpus` and/or `num_gpus` passed to the simulation engine"
156
136
  )
157
137
 
158
138
  return total_num_actors
@@ -183,9 +163,9 @@ class VirtualClientEngineActorPool(ActorPool):
183
163
 
184
164
  def __init__(
185
165
  self,
186
- create_actor_fn: Callable[[], Type[VirtualClientEngineActor]],
187
- client_resources: Dict[str, Union[int, float]],
188
- actor_list: Optional[List[Type[VirtualClientEngineActor]]] = None,
166
+ create_actor_fn: Callable[[], type[VirtualClientEngineActor]],
167
+ client_resources: dict[str, Union[int, float]],
168
+ actor_list: Optional[list[type[VirtualClientEngineActor]]] = None,
189
169
  ):
190
170
  self.client_resources = client_resources
191
171
  self.create_actor_fn = create_actor_fn
@@ -204,10 +184,10 @@ class VirtualClientEngineActorPool(ActorPool):
204
184
 
205
185
  # A dict that maps cid to another dict containing: a reference to the remote job
206
186
  # and its status (i.e. whether it is ready or not)
207
- self._cid_to_future: Dict[
208
- str, Dict[str, Union[bool, Optional[ObjectRef[Any]]]]
187
+ self._cid_to_future: dict[
188
+ str, dict[str, Union[bool, Optional[ObjectRef[Any]]]]
209
189
  ] = {}
210
- self.actor_to_remove: Set[str] = set() # a set
190
+ self.actor_to_remove: set[str] = set() # a set
211
191
  self.num_actors = len(actors)
212
192
 
213
193
  self.lock = threading.RLock()
@@ -231,7 +211,7 @@ class VirtualClientEngineActorPool(ActorPool):
231
211
  self._idle_actors.extend(new_actors)
232
212
  self.num_actors += num_actors
233
213
 
234
- def submit(self, fn: Any, value: Tuple[ClientAppFn, Message, str, Context]) -> None:
214
+ def submit(self, fn: Any, value: tuple[ClientAppFn, Message, str, Context]) -> None:
235
215
  """Take an idle actor and assign it to run a client app and Message.
236
216
 
237
217
  Submit a job to an actor by first removing it from the list of idle actors, then
@@ -241,7 +221,7 @@ class VirtualClientEngineActorPool(ActorPool):
241
221
  actor = self._idle_actors.pop()
242
222
  if self._check_and_remove_actor_from_pool(actor):
243
223
  future = fn(actor, app_fn, mssg, cid, context)
244
- future_key = tuple(future) if isinstance(future, List) else future
224
+ future_key = tuple(future) if isinstance(future, list) else future
245
225
  self._future_to_actor[future_key] = (self._next_task_index, actor, cid)
246
226
  self._next_task_index += 1
247
227
 
@@ -249,7 +229,7 @@ class VirtualClientEngineActorPool(ActorPool):
249
229
  self._cid_to_future[cid]["future"] = future_key
250
230
 
251
231
  def submit_client_job(
252
- self, actor_fn: Any, job: Tuple[ClientAppFn, Message, str, Context]
232
+ self, actor_fn: Any, job: tuple[ClientAppFn, Message, str, Context]
253
233
  ) -> None:
254
234
  """Submit a job while tracking client ids."""
255
235
  _, _, cid, _ = job
@@ -289,7 +269,7 @@ class VirtualClientEngineActorPool(ActorPool):
289
269
 
290
270
  return self._cid_to_future[cid]["ready"] # type: ignore
291
271
 
292
- def _fetch_future_result(self, cid: str) -> Tuple[Message, Context]:
272
+ def _fetch_future_result(self, cid: str) -> tuple[Message, Context]:
293
273
  """Fetch result and updated context for a VirtualClient from Object Store.
294
274
 
295
275
  The job submitted by the ClientProxy interfacing with client with cid=cid is
@@ -403,7 +383,7 @@ class VirtualClientEngineActorPool(ActorPool):
403
383
 
404
384
  def get_client_result(
405
385
  self, cid: str, timeout: Optional[float]
406
- ) -> Tuple[Message, Context]:
386
+ ) -> tuple[Message, Context]:
407
387
  """Get result from VirtualClient with specific cid."""
408
388
  # Loop until all jobs submitted to the pool are completed. Break early
409
389
  # if the result for the ClientProxy calling this method is ready
@@ -419,27 +399,19 @@ class VirtualClientEngineActorPool(ActorPool):
419
399
  return self._fetch_future_result(cid)
420
400
 
421
401
 
422
- def init_ray(*args: Any, **kwargs: Any) -> None:
423
- """Intialises Ray if not already initialised."""
424
- if not ray.is_initialized():
425
- ray.init(*args, **kwargs)
426
-
427
-
428
402
  class BasicActorPool:
429
403
  """A basic actor pool."""
430
404
 
431
405
  def __init__(
432
406
  self,
433
- actor_type: Type[VirtualClientEngineActor],
434
- client_resources: Dict[str, Union[int, float]],
435
- actor_kwargs: Dict[str, Any],
407
+ actor_type: type[VirtualClientEngineActor],
408
+ client_resources: dict[str, Union[int, float]],
409
+ actor_kwargs: dict[str, Any],
436
410
  ):
437
411
  self.client_resources = client_resources
438
412
 
439
413
  # Queue of idle actors
440
- self.pool: "asyncio.Queue[Type[VirtualClientEngineActor]]" = asyncio.Queue(
441
- maxsize=1024
442
- )
414
+ self.pool: list[VirtualClientEngineActor] = []
443
415
  self.num_actors = 0
444
416
 
445
417
  # Resolve arguments to pass during actor init
@@ -453,38 +425,37 @@ class BasicActorPool:
453
425
  # Figure out how many actors can be created given the cluster resources
454
426
  # and the resources the user indicates each VirtualClient will need
455
427
  self.actors_capacity = pool_size_from_resources(client_resources)
456
- self._future_to_actor: Dict[Any, Type[VirtualClientEngineActor]] = {}
428
+ self._future_to_actor: dict[Any, VirtualClientEngineActor] = {}
457
429
 
458
430
  def is_actor_available(self) -> bool:
459
431
  """Return true if there is an idle actor."""
460
- return self.pool.qsize() > 0
432
+ return len(self.pool) > 0
461
433
 
462
- async def add_actors_to_pool(self, num_actors: int) -> None:
434
+ def add_actors_to_pool(self, num_actors: int) -> None:
463
435
  """Add actors to the pool.
464
436
 
465
437
  This method may be executed also if new resources are added to your Ray cluster
466
438
  (e.g. you add a new node).
467
439
  """
468
440
  for _ in range(num_actors):
469
- await self.pool.put(self.create_actor_fn()) # type: ignore
441
+ self.pool.append(self.create_actor_fn()) # type: ignore
470
442
  self.num_actors += num_actors
471
443
 
472
- async def terminate_all_actors(self) -> None:
444
+ def terminate_all_actors(self) -> None:
473
445
  """Terminate actors in pool."""
474
446
  num_terminated = 0
475
- while self.pool.qsize():
476
- actor = await self.pool.get()
447
+ for actor in self.pool:
477
448
  actor.terminate.remote() # type: ignore
478
449
  num_terminated += 1
479
450
 
480
451
  log(DEBUG, "Terminated %i actors", num_terminated)
481
452
 
482
- async def submit(
483
- self, actor_fn: Any, job: Tuple[ClientAppFn, Message, str, Context]
453
+ def submit(
454
+ self, actor_fn: Any, job: tuple[ClientAppFn, Message, str, Context]
484
455
  ) -> Any:
485
456
  """On idle actor, submit job and return future."""
486
457
  # Remove idle actor from pool
487
- actor = await self.pool.get()
458
+ actor = self.pool.pop()
488
459
  # Submit job to actor
489
460
  app_fn, mssg, cid, context = job
490
461
  future = actor_fn(actor, app_fn, mssg, cid, context)
@@ -493,14 +464,18 @@ class BasicActorPool:
493
464
  self._future_to_actor[future] = actor
494
465
  return future
495
466
 
496
- async def fetch_result_and_return_actor_to_pool(
467
+ def add_actor_back_to_pool(self, future: Any) -> None:
468
+ """Ad actor assigned to run future back into the pool."""
469
+ actor = self._future_to_actor.pop(future)
470
+ self.pool.append(actor)
471
+
472
+ def fetch_result_and_return_actor_to_pool(
497
473
  self, future: Any
498
- ) -> Tuple[Message, Context]:
474
+ ) -> tuple[Message, Context]:
499
475
  """Pull result given a future and add actor back to pool."""
500
- # Get actor that ran job
501
- actor = self._future_to_actor.pop(future)
502
- await self.pool.put(actor)
503
476
  # Retrieve result for object store
504
477
  # Instead of doing ray.get(future) we await it
505
- _, out_mssg, updated_context = await future
478
+ _, out_mssg, updated_context = ray.get(future)
479
+ # Get actor that ran job
480
+ self.add_actor_back_to_pool(future)
506
481
  return out_mssg, updated_context