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.
- flwr/cli/app.py +16 -2
- flwr/cli/build.py +181 -0
- flwr/cli/cli_user_auth_interceptor.py +90 -0
- flwr/cli/config_utils.py +343 -0
- flwr/cli/example.py +4 -1
- flwr/cli/install.py +253 -0
- flwr/cli/log.py +182 -0
- flwr/{server/superlink/state → cli/login}/__init__.py +4 -10
- flwr/cli/login/login.py +88 -0
- flwr/cli/ls.py +327 -0
- flwr/cli/new/__init__.py +1 -0
- flwr/cli/new/new.py +210 -66
- flwr/cli/new/templates/app/.gitignore.tpl +163 -0
- flwr/cli/new/templates/app/LICENSE.tpl +202 -0
- flwr/cli/new/templates/app/README.baseline.md.tpl +127 -0
- flwr/cli/new/templates/app/README.flowertune.md.tpl +66 -0
- flwr/cli/new/templates/app/README.md.tpl +16 -32
- flwr/cli/new/templates/app/code/__init__.baseline.py.tpl +1 -0
- flwr/cli/new/templates/app/code/__init__.py.tpl +1 -1
- flwr/cli/new/templates/app/code/client.baseline.py.tpl +58 -0
- flwr/cli/new/templates/app/code/client.huggingface.py.tpl +55 -0
- flwr/cli/new/templates/app/code/client.jax.py.tpl +50 -0
- flwr/cli/new/templates/app/code/client.mlx.py.tpl +73 -0
- flwr/cli/new/templates/app/code/client.numpy.py.tpl +7 -7
- flwr/cli/new/templates/app/code/client.pytorch.py.tpl +30 -21
- flwr/cli/new/templates/app/code/client.sklearn.py.tpl +63 -0
- flwr/cli/new/templates/app/code/client.tensorflow.py.tpl +57 -1
- flwr/cli/new/templates/app/code/dataset.baseline.py.tpl +36 -0
- flwr/cli/new/templates/app/code/flwr_tune/__init__.py +15 -0
- flwr/cli/new/templates/app/code/flwr_tune/client_app.py.tpl +126 -0
- flwr/cli/new/templates/app/code/flwr_tune/dataset.py.tpl +87 -0
- flwr/cli/new/templates/app/code/flwr_tune/models.py.tpl +78 -0
- flwr/cli/new/templates/app/code/flwr_tune/server_app.py.tpl +94 -0
- flwr/cli/new/templates/app/code/flwr_tune/strategy.py.tpl +83 -0
- flwr/cli/new/templates/app/code/model.baseline.py.tpl +80 -0
- flwr/cli/new/templates/app/code/server.baseline.py.tpl +46 -0
- flwr/cli/new/templates/app/code/server.huggingface.py.tpl +38 -0
- flwr/cli/new/templates/app/code/server.jax.py.tpl +26 -0
- flwr/cli/new/templates/app/code/server.mlx.py.tpl +31 -0
- flwr/cli/new/templates/app/code/server.numpy.py.tpl +22 -9
- flwr/cli/new/templates/app/code/server.pytorch.py.tpl +21 -18
- flwr/cli/new/templates/app/code/server.sklearn.py.tpl +36 -0
- flwr/cli/new/templates/app/code/server.tensorflow.py.tpl +29 -1
- flwr/cli/new/templates/app/code/strategy.baseline.py.tpl +1 -0
- flwr/cli/new/templates/app/code/task.huggingface.py.tpl +102 -0
- flwr/cli/new/templates/app/code/task.jax.py.tpl +57 -0
- flwr/cli/new/templates/app/code/task.mlx.py.tpl +102 -0
- flwr/cli/new/templates/app/code/task.numpy.py.tpl +7 -0
- flwr/cli/new/templates/app/code/task.pytorch.py.tpl +29 -24
- flwr/cli/new/templates/app/code/task.sklearn.py.tpl +67 -0
- flwr/cli/new/templates/app/code/task.tensorflow.py.tpl +53 -0
- flwr/cli/new/templates/app/code/utils.baseline.py.tpl +1 -0
- flwr/cli/new/templates/app/pyproject.baseline.toml.tpl +138 -0
- flwr/cli/new/templates/app/pyproject.flowertune.toml.tpl +68 -0
- flwr/cli/new/templates/app/pyproject.huggingface.toml.tpl +46 -0
- flwr/cli/new/templates/app/pyproject.jax.toml.tpl +35 -0
- flwr/cli/new/templates/app/pyproject.mlx.toml.tpl +39 -0
- flwr/cli/new/templates/app/pyproject.numpy.toml.tpl +25 -12
- flwr/cli/new/templates/app/pyproject.pytorch.toml.tpl +29 -14
- flwr/cli/new/templates/app/pyproject.sklearn.toml.tpl +35 -0
- flwr/cli/new/templates/app/pyproject.tensorflow.toml.tpl +29 -14
- flwr/cli/run/__init__.py +1 -0
- flwr/cli/run/run.py +212 -34
- flwr/cli/stop.py +130 -0
- flwr/cli/utils.py +240 -5
- flwr/client/__init__.py +3 -2
- flwr/client/app.py +432 -255
- flwr/client/client.py +1 -11
- flwr/client/client_app.py +74 -13
- flwr/client/clientapp/__init__.py +22 -0
- flwr/client/clientapp/app.py +259 -0
- flwr/client/clientapp/clientappio_servicer.py +244 -0
- flwr/client/clientapp/utils.py +115 -0
- flwr/client/dpfedavg_numpy_client.py +7 -8
- flwr/client/grpc_adapter_client/__init__.py +15 -0
- flwr/client/grpc_adapter_client/connection.py +98 -0
- flwr/client/grpc_client/connection.py +21 -7
- flwr/client/grpc_rere_client/__init__.py +1 -1
- flwr/client/grpc_rere_client/client_interceptor.py +176 -0
- flwr/client/grpc_rere_client/connection.py +163 -56
- flwr/client/grpc_rere_client/grpc_adapter.py +167 -0
- flwr/client/heartbeat.py +74 -0
- flwr/client/message_handler/__init__.py +1 -1
- flwr/client/message_handler/message_handler.py +10 -11
- flwr/client/mod/__init__.py +5 -5
- flwr/client/mod/centraldp_mods.py +4 -2
- flwr/client/mod/comms_mods.py +5 -4
- flwr/client/mod/localdp_mod.py +10 -5
- flwr/client/mod/secure_aggregation/__init__.py +1 -1
- flwr/client/mod/secure_aggregation/secaggplus_mod.py +26 -26
- flwr/client/mod/utils.py +2 -4
- flwr/client/nodestate/__init__.py +26 -0
- flwr/client/nodestate/in_memory_nodestate.py +38 -0
- flwr/client/nodestate/nodestate.py +31 -0
- flwr/client/nodestate/nodestate_factory.py +38 -0
- flwr/client/numpy_client.py +8 -31
- flwr/client/rest_client/__init__.py +1 -1
- flwr/client/rest_client/connection.py +199 -176
- flwr/client/run_info_store.py +112 -0
- flwr/client/supernode/__init__.py +24 -0
- flwr/client/supernode/app.py +321 -0
- flwr/client/typing.py +1 -0
- flwr/common/__init__.py +17 -11
- flwr/common/address.py +47 -3
- flwr/common/args.py +153 -0
- flwr/common/auth_plugin/__init__.py +24 -0
- flwr/common/auth_plugin/auth_plugin.py +121 -0
- flwr/common/config.py +243 -0
- flwr/common/constant.py +135 -1
- flwr/common/context.py +32 -2
- flwr/common/date.py +22 -4
- flwr/common/differential_privacy.py +2 -2
- flwr/common/dp.py +2 -4
- flwr/common/exit_handlers.py +3 -3
- flwr/common/grpc.py +164 -5
- flwr/common/logger.py +230 -12
- flwr/common/message.py +191 -106
- flwr/common/object_ref.py +179 -44
- flwr/common/pyproject.py +1 -0
- flwr/common/record/__init__.py +2 -1
- flwr/common/record/configsrecord.py +58 -18
- flwr/common/record/metricsrecord.py +57 -17
- flwr/common/record/parametersrecord.py +88 -20
- flwr/common/record/recordset.py +153 -30
- flwr/common/record/typeddict.py +30 -55
- flwr/common/recordset_compat.py +31 -12
- flwr/common/retry_invoker.py +123 -30
- flwr/common/secure_aggregation/__init__.py +1 -1
- flwr/common/secure_aggregation/crypto/__init__.py +1 -1
- flwr/common/secure_aggregation/crypto/shamir.py +11 -11
- flwr/common/secure_aggregation/crypto/symmetric_encryption.py +68 -4
- flwr/common/secure_aggregation/ndarrays_arithmetic.py +17 -17
- flwr/common/secure_aggregation/quantization.py +8 -8
- flwr/common/secure_aggregation/secaggplus_constants.py +1 -1
- flwr/common/secure_aggregation/secaggplus_utils.py +10 -12
- flwr/common/serde.py +304 -23
- flwr/common/telemetry.py +65 -29
- flwr/common/typing.py +120 -19
- flwr/common/version.py +17 -3
- flwr/proto/clientappio_pb2.py +45 -0
- flwr/proto/clientappio_pb2.pyi +132 -0
- flwr/proto/clientappio_pb2_grpc.py +135 -0
- flwr/proto/clientappio_pb2_grpc.pyi +53 -0
- flwr/proto/exec_pb2.py +62 -0
- flwr/proto/exec_pb2.pyi +212 -0
- flwr/proto/exec_pb2_grpc.py +237 -0
- flwr/proto/exec_pb2_grpc.pyi +93 -0
- flwr/proto/fab_pb2.py +31 -0
- flwr/proto/fab_pb2.pyi +65 -0
- flwr/proto/fab_pb2_grpc.py +4 -0
- flwr/proto/fab_pb2_grpc.pyi +4 -0
- flwr/proto/fleet_pb2.py +42 -23
- flwr/proto/fleet_pb2.pyi +123 -1
- flwr/proto/fleet_pb2_grpc.py +170 -0
- flwr/proto/fleet_pb2_grpc.pyi +61 -0
- flwr/proto/grpcadapter_pb2.py +32 -0
- flwr/proto/grpcadapter_pb2.pyi +43 -0
- flwr/proto/grpcadapter_pb2_grpc.py +66 -0
- flwr/proto/grpcadapter_pb2_grpc.pyi +24 -0
- flwr/proto/log_pb2.py +29 -0
- flwr/proto/log_pb2.pyi +39 -0
- flwr/proto/log_pb2_grpc.py +4 -0
- flwr/proto/log_pb2_grpc.pyi +4 -0
- flwr/proto/message_pb2.py +41 -0
- flwr/proto/message_pb2.pyi +128 -0
- flwr/proto/message_pb2_grpc.py +4 -0
- flwr/proto/message_pb2_grpc.pyi +4 -0
- flwr/proto/node_pb2.py +2 -2
- flwr/proto/node_pb2.pyi +1 -4
- flwr/proto/recordset_pb2.py +35 -33
- flwr/proto/recordset_pb2.pyi +40 -14
- flwr/proto/run_pb2.py +64 -0
- flwr/proto/run_pb2.pyi +268 -0
- flwr/proto/run_pb2_grpc.py +4 -0
- flwr/proto/run_pb2_grpc.pyi +4 -0
- flwr/proto/serverappio_pb2.py +52 -0
- flwr/proto/{driver_pb2.pyi → serverappio_pb2.pyi} +62 -20
- flwr/proto/serverappio_pb2_grpc.py +410 -0
- flwr/proto/serverappio_pb2_grpc.pyi +160 -0
- flwr/proto/simulationio_pb2.py +38 -0
- flwr/proto/simulationio_pb2.pyi +65 -0
- flwr/proto/simulationio_pb2_grpc.py +239 -0
- flwr/proto/simulationio_pb2_grpc.pyi +94 -0
- flwr/proto/task_pb2.py +7 -8
- flwr/proto/task_pb2.pyi +8 -5
- flwr/proto/transport_pb2.py +8 -8
- flwr/proto/transport_pb2.pyi +9 -6
- flwr/server/__init__.py +2 -10
- flwr/server/app.py +579 -402
- flwr/server/client_manager.py +8 -6
- flwr/server/compat/app.py +6 -62
- flwr/server/compat/app_utils.py +14 -9
- flwr/server/compat/driver_client_proxy.py +25 -59
- flwr/server/compat/legacy_context.py +5 -4
- flwr/server/driver/__init__.py +2 -0
- flwr/server/driver/driver.py +36 -131
- flwr/server/driver/grpc_driver.py +220 -81
- flwr/server/driver/inmemory_driver.py +183 -0
- flwr/server/history.py +28 -29
- flwr/server/run_serverapp.py +15 -126
- flwr/server/server.py +50 -44
- flwr/server/server_app.py +59 -10
- flwr/server/serverapp/__init__.py +22 -0
- flwr/server/serverapp/app.py +256 -0
- flwr/server/serverapp_components.py +52 -0
- flwr/server/strategy/__init__.py +2 -2
- flwr/server/strategy/aggregate.py +37 -23
- flwr/server/strategy/bulyan.py +9 -9
- flwr/server/strategy/dp_adaptive_clipping.py +25 -25
- flwr/server/strategy/dp_fixed_clipping.py +23 -22
- flwr/server/strategy/dpfedavg_adaptive.py +8 -8
- flwr/server/strategy/dpfedavg_fixed.py +13 -12
- flwr/server/strategy/fault_tolerant_fedavg.py +11 -11
- flwr/server/strategy/fedadagrad.py +9 -9
- flwr/server/strategy/fedadam.py +20 -10
- flwr/server/strategy/fedavg.py +16 -16
- flwr/server/strategy/fedavg_android.py +17 -17
- flwr/server/strategy/fedavgm.py +9 -9
- flwr/server/strategy/fedmedian.py +5 -5
- flwr/server/strategy/fedopt.py +6 -6
- flwr/server/strategy/fedprox.py +7 -7
- flwr/server/strategy/fedtrimmedavg.py +8 -8
- flwr/server/strategy/fedxgb_bagging.py +12 -12
- flwr/server/strategy/fedxgb_cyclic.py +10 -10
- flwr/server/strategy/fedxgb_nn_avg.py +6 -6
- flwr/server/strategy/fedyogi.py +9 -9
- flwr/server/strategy/krum.py +9 -9
- flwr/server/strategy/qfedavg.py +16 -16
- flwr/server/strategy/strategy.py +10 -10
- flwr/server/superlink/driver/__init__.py +2 -2
- flwr/server/superlink/driver/serverappio_grpc.py +61 -0
- flwr/server/superlink/driver/serverappio_servicer.py +361 -0
- flwr/server/superlink/ffs/__init__.py +24 -0
- flwr/server/superlink/ffs/disk_ffs.py +108 -0
- flwr/server/superlink/ffs/ffs.py +79 -0
- flwr/server/superlink/ffs/ffs_factory.py +47 -0
- flwr/server/superlink/fleet/__init__.py +1 -1
- flwr/server/superlink/fleet/grpc_adapter/__init__.py +15 -0
- flwr/server/superlink/fleet/grpc_adapter/grpc_adapter_servicer.py +162 -0
- flwr/server/superlink/fleet/grpc_bidi/__init__.py +1 -1
- flwr/server/superlink/fleet/grpc_bidi/flower_service_servicer.py +4 -2
- flwr/server/superlink/fleet/grpc_bidi/grpc_bridge.py +3 -2
- flwr/server/superlink/fleet/grpc_bidi/grpc_client_proxy.py +1 -1
- flwr/server/superlink/fleet/grpc_bidi/grpc_server.py +5 -154
- flwr/server/superlink/fleet/grpc_rere/__init__.py +1 -1
- flwr/server/superlink/fleet/grpc_rere/fleet_servicer.py +120 -13
- flwr/server/superlink/fleet/grpc_rere/server_interceptor.py +228 -0
- flwr/server/superlink/fleet/message_handler/__init__.py +1 -1
- flwr/server/superlink/fleet/message_handler/message_handler.py +156 -13
- flwr/server/superlink/fleet/rest_rere/__init__.py +1 -1
- flwr/server/superlink/fleet/rest_rere/rest_api.py +119 -81
- flwr/server/superlink/fleet/vce/__init__.py +1 -0
- flwr/server/superlink/fleet/vce/backend/__init__.py +4 -4
- flwr/server/superlink/fleet/vce/backend/backend.py +8 -9
- flwr/server/superlink/fleet/vce/backend/raybackend.py +87 -68
- flwr/server/superlink/fleet/vce/vce_api.py +208 -146
- flwr/server/superlink/linkstate/__init__.py +28 -0
- flwr/server/superlink/linkstate/in_memory_linkstate.py +569 -0
- flwr/server/superlink/linkstate/linkstate.py +376 -0
- flwr/server/superlink/{state/state_factory.py → linkstate/linkstate_factory.py} +19 -10
- flwr/server/superlink/linkstate/sqlite_linkstate.py +1196 -0
- flwr/server/superlink/linkstate/utils.py +399 -0
- flwr/server/superlink/simulation/__init__.py +15 -0
- flwr/server/superlink/simulation/simulationio_grpc.py +65 -0
- flwr/server/superlink/simulation/simulationio_servicer.py +186 -0
- flwr/server/superlink/utils.py +65 -0
- flwr/server/typing.py +2 -0
- flwr/server/utils/__init__.py +1 -1
- flwr/server/utils/tensorboard.py +5 -5
- flwr/server/utils/validator.py +40 -45
- flwr/server/workflow/default_workflows.py +70 -26
- flwr/server/workflow/secure_aggregation/secagg_workflow.py +1 -0
- flwr/server/workflow/secure_aggregation/secaggplus_workflow.py +40 -27
- flwr/simulation/__init__.py +12 -5
- flwr/simulation/app.py +247 -315
- flwr/simulation/legacy_app.py +404 -0
- flwr/simulation/ray_transport/__init__.py +1 -1
- flwr/simulation/ray_transport/ray_actor.py +42 -67
- flwr/simulation/ray_transport/ray_client_proxy.py +37 -17
- flwr/simulation/ray_transport/utils.py +1 -0
- flwr/simulation/run_simulation.py +306 -163
- flwr/simulation/simulationio_connection.py +89 -0
- flwr/superexec/__init__.py +15 -0
- flwr/superexec/app.py +59 -0
- flwr/superexec/deployment.py +188 -0
- flwr/superexec/exec_grpc.py +80 -0
- flwr/superexec/exec_servicer.py +231 -0
- flwr/superexec/exec_user_auth_interceptor.py +101 -0
- flwr/superexec/executor.py +96 -0
- flwr/superexec/simulation.py +124 -0
- {flwr_nightly-1.8.0.dev20240315.dist-info → flwr_nightly-1.15.0.dev20250115.dist-info}/METADATA +33 -26
- flwr_nightly-1.15.0.dev20250115.dist-info/RECORD +328 -0
- flwr_nightly-1.15.0.dev20250115.dist-info/entry_points.txt +12 -0
- flwr/cli/flower_toml.py +0 -140
- flwr/cli/new/templates/app/flower.toml.tpl +0 -13
- flwr/cli/new/templates/app/requirements.numpy.txt.tpl +0 -2
- flwr/cli/new/templates/app/requirements.pytorch.txt.tpl +0 -4
- flwr/cli/new/templates/app/requirements.tensorflow.txt.tpl +0 -4
- flwr/client/node_state.py +0 -48
- flwr/client/node_state_tests.py +0 -65
- flwr/proto/driver_pb2.py +0 -44
- flwr/proto/driver_pb2_grpc.py +0 -169
- flwr/proto/driver_pb2_grpc.pyi +0 -66
- flwr/server/superlink/driver/driver_grpc.py +0 -54
- flwr/server/superlink/driver/driver_servicer.py +0 -129
- flwr/server/superlink/state/in_memory_state.py +0 -230
- flwr/server/superlink/state/sqlite_state.py +0 -630
- flwr/server/superlink/state/state.py +0 -154
- flwr_nightly-1.8.0.dev20240315.dist-info/RECORD +0 -211
- flwr_nightly-1.8.0.dev20240315.dist-info/entry_points.txt +0 -9
- {flwr_nightly-1.8.0.dev20240315.dist-info → flwr_nightly-1.15.0.dev20250115.dist-info}/LICENSE +0 -0
- {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
|
@@ -14,33 +14,23 @@
|
|
14
14
|
# ==============================================================================
|
15
15
|
"""Ray-based Flower Actor and ActorPool implementation."""
|
16
16
|
|
17
|
-
|
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,
|
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
|
-
) ->
|
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
|
-
|
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:
|
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
|
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
|
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[[],
|
187
|
-
client_resources:
|
188
|
-
actor_list: Optional[
|
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:
|
208
|
-
str,
|
187
|
+
self._cid_to_future: dict[
|
188
|
+
str, dict[str, Union[bool, Optional[ObjectRef[Any]]]]
|
209
189
|
] = {}
|
210
|
-
self.actor_to_remove:
|
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:
|
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,
|
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:
|
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) ->
|
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
|
-
) ->
|
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:
|
434
|
-
client_resources:
|
435
|
-
actor_kwargs:
|
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:
|
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:
|
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
|
432
|
+
return len(self.pool) > 0
|
461
433
|
|
462
|
-
|
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
|
-
|
441
|
+
self.pool.append(self.create_actor_fn()) # type: ignore
|
470
442
|
self.num_actors += num_actors
|
471
443
|
|
472
|
-
|
444
|
+
def terminate_all_actors(self) -> None:
|
473
445
|
"""Terminate actors in pool."""
|
474
446
|
num_terminated = 0
|
475
|
-
|
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
|
-
|
483
|
-
self, actor_fn: Any, job:
|
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 =
|
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
|
-
|
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
|
-
) ->
|
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 =
|
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
|