flwr-nightly 1.9.0.dev20240420__py3-none-any.whl → 1.9.0.dev20240509__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.

Potentially problematic release.


This version of flwr-nightly might be problematic. Click here for more details.

Files changed (71) hide show
  1. flwr/cli/app.py +2 -0
  2. flwr/cli/build.py +151 -0
  3. flwr/cli/config_utils.py +18 -46
  4. flwr/cli/new/new.py +44 -18
  5. flwr/cli/new/templates/app/code/client.hf.py.tpl +55 -0
  6. flwr/cli/new/templates/app/code/client.mlx.py.tpl +70 -0
  7. flwr/cli/new/templates/app/code/client.pytorch.py.tpl +1 -1
  8. flwr/cli/new/templates/app/code/client.sklearn.py.tpl +94 -0
  9. flwr/cli/new/templates/app/code/client.tensorflow.py.tpl +15 -29
  10. flwr/cli/new/templates/app/code/server.hf.py.tpl +17 -0
  11. flwr/cli/new/templates/app/code/server.mlx.py.tpl +15 -0
  12. flwr/cli/new/templates/app/code/server.pytorch.py.tpl +1 -1
  13. flwr/cli/new/templates/app/code/server.sklearn.py.tpl +17 -0
  14. flwr/cli/new/templates/app/code/server.tensorflow.py.tpl +9 -1
  15. flwr/cli/new/templates/app/code/task.hf.py.tpl +87 -0
  16. flwr/cli/new/templates/app/code/task.mlx.py.tpl +89 -0
  17. flwr/cli/new/templates/app/code/task.tensorflow.py.tpl +29 -0
  18. flwr/cli/new/templates/app/pyproject.hf.toml.tpl +31 -0
  19. flwr/cli/new/templates/app/pyproject.mlx.toml.tpl +28 -0
  20. flwr/cli/new/templates/app/pyproject.numpy.toml.tpl +7 -4
  21. flwr/cli/new/templates/app/pyproject.pytorch.toml.tpl +7 -4
  22. flwr/cli/new/templates/app/pyproject.sklearn.toml.tpl +27 -0
  23. flwr/cli/new/templates/app/pyproject.tensorflow.toml.tpl +7 -4
  24. flwr/cli/run/run.py +1 -1
  25. flwr/cli/utils.py +18 -17
  26. flwr/client/__init__.py +1 -1
  27. flwr/client/app.py +17 -93
  28. flwr/client/grpc_client/connection.py +6 -1
  29. flwr/client/grpc_rere_client/client_interceptor.py +158 -0
  30. flwr/client/grpc_rere_client/connection.py +17 -2
  31. flwr/client/mod/centraldp_mods.py +4 -2
  32. flwr/client/mod/localdp_mod.py +9 -3
  33. flwr/client/rest_client/connection.py +5 -1
  34. flwr/client/supernode/__init__.py +2 -0
  35. flwr/client/supernode/app.py +181 -7
  36. flwr/common/grpc.py +5 -1
  37. flwr/common/logger.py +37 -4
  38. flwr/common/message.py +105 -86
  39. flwr/common/record/parametersrecord.py +0 -1
  40. flwr/common/record/recordset.py +17 -5
  41. flwr/common/secure_aggregation/crypto/symmetric_encryption.py +35 -1
  42. flwr/server/__init__.py +0 -2
  43. flwr/server/app.py +118 -2
  44. flwr/server/compat/app.py +5 -56
  45. flwr/server/compat/app_utils.py +1 -1
  46. flwr/server/compat/driver_client_proxy.py +27 -72
  47. flwr/server/driver/__init__.py +3 -0
  48. flwr/server/driver/driver.py +12 -242
  49. flwr/server/driver/grpc_driver.py +315 -0
  50. flwr/server/history.py +20 -20
  51. flwr/server/run_serverapp.py +18 -4
  52. flwr/server/server.py +2 -5
  53. flwr/server/strategy/dp_adaptive_clipping.py +5 -3
  54. flwr/server/strategy/dp_fixed_clipping.py +6 -3
  55. flwr/server/superlink/driver/driver_servicer.py +1 -1
  56. flwr/server/superlink/fleet/grpc_bidi/grpc_server.py +3 -1
  57. flwr/server/superlink/fleet/grpc_rere/server_interceptor.py +215 -0
  58. flwr/server/superlink/fleet/vce/backend/raybackend.py +9 -6
  59. flwr/server/superlink/fleet/vce/vce_api.py +1 -1
  60. flwr/server/superlink/state/in_memory_state.py +76 -8
  61. flwr/server/superlink/state/sqlite_state.py +116 -11
  62. flwr/server/superlink/state/state.py +35 -3
  63. flwr/simulation/__init__.py +2 -2
  64. flwr/simulation/app.py +16 -1
  65. flwr/simulation/run_simulation.py +14 -9
  66. {flwr_nightly-1.9.0.dev20240420.dist-info → flwr_nightly-1.9.0.dev20240509.dist-info}/METADATA +3 -2
  67. {flwr_nightly-1.9.0.dev20240420.dist-info → flwr_nightly-1.9.0.dev20240509.dist-info}/RECORD +70 -55
  68. {flwr_nightly-1.9.0.dev20240420.dist-info → flwr_nightly-1.9.0.dev20240509.dist-info}/entry_points.txt +1 -1
  69. flwr/server/driver/abc_driver.py +0 -140
  70. {flwr_nightly-1.9.0.dev20240420.dist-info → flwr_nightly-1.9.0.dev20240509.dist-info}/LICENSE +0 -0
  71. {flwr_nightly-1.9.0.dev20240420.dist-info → flwr_nightly-1.9.0.dev20240509.dist-info}/WHEEL +0 -0
@@ -1,4 +1,4 @@
1
- # Copyright 2022 Flower Labs GmbH. All Rights Reserved.
1
+ # Copyright 2024 Flower Labs GmbH. All Rights Reserved.
2
2
  #
3
3
  # Licensed under the Apache License, Version 2.0 (the "License");
4
4
  # you may not use this file except in compliance with the License.
@@ -12,180 +12,19 @@
12
12
  # See the License for the specific language governing permissions and
13
13
  # limitations under the License.
14
14
  # ==============================================================================
15
- """Flower driver service client."""
15
+ """Driver (abstract base class)."""
16
16
 
17
- import time
18
- import warnings
19
- from logging import DEBUG, ERROR, WARNING
20
- from typing import Iterable, List, Optional, Tuple
21
17
 
22
- import grpc
18
+ from abc import ABC, abstractmethod
19
+ from typing import Iterable, List, Optional
23
20
 
24
- from flwr.common import DEFAULT_TTL, EventType, Message, Metadata, RecordSet, event
25
- from flwr.common.grpc import create_channel
26
- from flwr.common.logger import log
27
- from flwr.common.serde import message_from_taskres, message_to_taskins
28
- from flwr.proto.driver_pb2 import ( # pylint: disable=E0611
29
- CreateRunRequest,
30
- CreateRunResponse,
31
- GetNodesRequest,
32
- GetNodesResponse,
33
- PullTaskResRequest,
34
- PullTaskResResponse,
35
- PushTaskInsRequest,
36
- PushTaskInsResponse,
37
- )
38
- from flwr.proto.driver_pb2_grpc import DriverStub # pylint: disable=E0611
39
- from flwr.proto.node_pb2 import Node # pylint: disable=E0611
40
- from flwr.proto.task_pb2 import TaskIns # pylint: disable=E0611
21
+ from flwr.common import Message, RecordSet
41
22
 
42
- DEFAULT_SERVER_ADDRESS_DRIVER = "[::]:9091"
43
23
 
44
- ERROR_MESSAGE_DRIVER_NOT_CONNECTED = """
45
- [Driver] Error: Not connected.
46
-
47
- Call `connect()` on the `GrpcDriverHelper` instance before calling any of the other
48
- `GrpcDriverHelper` methods.
49
- """
50
-
51
-
52
- class GrpcDriverHelper:
53
- """`GrpcDriverHelper` provides access to the gRPC Driver API/service."""
54
-
55
- def __init__(
56
- self,
57
- driver_service_address: str = DEFAULT_SERVER_ADDRESS_DRIVER,
58
- root_certificates: Optional[bytes] = None,
59
- ) -> None:
60
- self.driver_service_address = driver_service_address
61
- self.root_certificates = root_certificates
62
- self.channel: Optional[grpc.Channel] = None
63
- self.stub: Optional[DriverStub] = None
64
-
65
- def connect(self) -> None:
66
- """Connect to the Driver API."""
67
- event(EventType.DRIVER_CONNECT)
68
- if self.channel is not None or self.stub is not None:
69
- log(WARNING, "Already connected")
70
- return
71
- self.channel = create_channel(
72
- server_address=self.driver_service_address,
73
- insecure=(self.root_certificates is None),
74
- root_certificates=self.root_certificates,
75
- )
76
- self.stub = DriverStub(self.channel)
77
- log(DEBUG, "[Driver] Connected to %s", self.driver_service_address)
78
-
79
- def disconnect(self) -> None:
80
- """Disconnect from the Driver API."""
81
- event(EventType.DRIVER_DISCONNECT)
82
- if self.channel is None or self.stub is None:
83
- log(DEBUG, "Already disconnected")
84
- return
85
- channel = self.channel
86
- self.channel = None
87
- self.stub = None
88
- channel.close()
89
- log(DEBUG, "[Driver] Disconnected")
90
-
91
- def create_run(self, req: CreateRunRequest) -> CreateRunResponse:
92
- """Request for run ID."""
93
- # Check if channel is open
94
- if self.stub is None:
95
- log(ERROR, ERROR_MESSAGE_DRIVER_NOT_CONNECTED)
96
- raise ConnectionError("`GrpcDriverHelper` instance not connected")
97
-
98
- # Call Driver API
99
- res: CreateRunResponse = self.stub.CreateRun(request=req)
100
- return res
101
-
102
- def get_nodes(self, req: GetNodesRequest) -> GetNodesResponse:
103
- """Get client IDs."""
104
- # Check if channel is open
105
- if self.stub is None:
106
- log(ERROR, ERROR_MESSAGE_DRIVER_NOT_CONNECTED)
107
- raise ConnectionError("`GrpcDriverHelper` instance not connected")
108
-
109
- # Call gRPC Driver API
110
- res: GetNodesResponse = self.stub.GetNodes(request=req)
111
- return res
112
-
113
- def push_task_ins(self, req: PushTaskInsRequest) -> PushTaskInsResponse:
114
- """Schedule tasks."""
115
- # Check if channel is open
116
- if self.stub is None:
117
- log(ERROR, ERROR_MESSAGE_DRIVER_NOT_CONNECTED)
118
- raise ConnectionError("`GrpcDriverHelper` instance not connected")
119
-
120
- # Call gRPC Driver API
121
- res: PushTaskInsResponse = self.stub.PushTaskIns(request=req)
122
- return res
123
-
124
- def pull_task_res(self, req: PullTaskResRequest) -> PullTaskResResponse:
125
- """Get task results."""
126
- # Check if channel is open
127
- if self.stub is None:
128
- log(ERROR, ERROR_MESSAGE_DRIVER_NOT_CONNECTED)
129
- raise ConnectionError("`GrpcDriverHelper` instance not connected")
130
-
131
- # Call Driver API
132
- res: PullTaskResResponse = self.stub.PullTaskRes(request=req)
133
- return res
134
-
135
-
136
- class Driver:
137
- """`Driver` class provides an interface to the Driver API.
138
-
139
- Parameters
140
- ----------
141
- driver_service_address : Optional[str]
142
- The IPv4 or IPv6 address of the Driver API server.
143
- Defaults to `"[::]:9091"`.
144
- certificates : bytes (default: None)
145
- Tuple containing root certificate, server certificate, and private key
146
- to start a secure SSL-enabled server. The tuple is expected to have
147
- three bytes elements in the following order:
148
-
149
- * CA certificate.
150
- * server certificate.
151
- * server private key.
152
- """
153
-
154
- def __init__(
155
- self,
156
- driver_service_address: str = DEFAULT_SERVER_ADDRESS_DRIVER,
157
- root_certificates: Optional[bytes] = None,
158
- ) -> None:
159
- self.addr = driver_service_address
160
- self.root_certificates = root_certificates
161
- self.grpc_driver_helper: Optional[GrpcDriverHelper] = None
162
- self.run_id: Optional[int] = None
163
- self.node = Node(node_id=0, anonymous=True)
164
-
165
- def _get_grpc_driver_helper_and_run_id(self) -> Tuple[GrpcDriverHelper, int]:
166
- # Check if the GrpcDriverHelper is initialized
167
- if self.grpc_driver_helper is None or self.run_id is None:
168
- # Connect and create run
169
- self.grpc_driver_helper = GrpcDriverHelper(
170
- driver_service_address=self.addr,
171
- root_certificates=self.root_certificates,
172
- )
173
- self.grpc_driver_helper.connect()
174
- res = self.grpc_driver_helper.create_run(CreateRunRequest())
175
- self.run_id = res.run_id
176
- return self.grpc_driver_helper, self.run_id
177
-
178
- def _check_message(self, message: Message) -> None:
179
- # Check if the message is valid
180
- if not (
181
- message.metadata.run_id == self.run_id
182
- and message.metadata.src_node_id == self.node.node_id
183
- and message.metadata.message_id == ""
184
- and message.metadata.reply_to_message == ""
185
- and message.metadata.ttl > 0
186
- ):
187
- raise ValueError(f"Invalid message: {message}")
24
+ class Driver(ABC):
25
+ """Abstract base Driver class for the Driver API."""
188
26
 
27
+ @abstractmethod
189
28
  def create_message( # pylint: disable=too-many-arguments
190
29
  self,
191
30
  content: RecordSet,
@@ -223,35 +62,12 @@ class Driver:
223
62
  message : Message
224
63
  A new `Message` instance with the specified content and metadata.
225
64
  """
226
- _, run_id = self._get_grpc_driver_helper_and_run_id()
227
- if ttl:
228
- warnings.warn(
229
- "A custom TTL was set, but note that the SuperLink does not enforce "
230
- "the TTL yet. The SuperLink will start enforcing the TTL in a future "
231
- "version of Flower.",
232
- stacklevel=2,
233
- )
234
-
235
- ttl_ = DEFAULT_TTL if ttl is None else ttl
236
- metadata = Metadata(
237
- run_id=run_id,
238
- message_id="", # Will be set by the server
239
- src_node_id=self.node.node_id,
240
- dst_node_id=dst_node_id,
241
- reply_to_message="",
242
- group_id=group_id,
243
- ttl=ttl_,
244
- message_type=message_type,
245
- )
246
- return Message(metadata=metadata, content=content)
247
65
 
66
+ @abstractmethod
248
67
  def get_node_ids(self) -> List[int]:
249
68
  """Get node IDs."""
250
- grpc_driver_helper, run_id = self._get_grpc_driver_helper_and_run_id()
251
- # Call GrpcDriverHelper method
252
- res = grpc_driver_helper.get_nodes(GetNodesRequest(run_id=run_id))
253
- return [node.node_id for node in res.nodes]
254
69
 
70
+ @abstractmethod
255
71
  def push_messages(self, messages: Iterable[Message]) -> Iterable[str]:
256
72
  """Push messages to specified node IDs.
257
73
 
@@ -269,22 +85,8 @@ class Driver:
269
85
  An iterable of IDs for the messages that were sent, which can be used
270
86
  to pull replies.
271
87
  """
272
- grpc_driver_helper, _ = self._get_grpc_driver_helper_and_run_id()
273
- # Construct TaskIns
274
- task_ins_list: List[TaskIns] = []
275
- for msg in messages:
276
- # Check message
277
- self._check_message(msg)
278
- # Convert Message to TaskIns
279
- taskins = message_to_taskins(msg)
280
- # Add to list
281
- task_ins_list.append(taskins)
282
- # Call GrpcDriverHelper method
283
- res = grpc_driver_helper.push_task_ins(
284
- PushTaskInsRequest(task_ins_list=task_ins_list)
285
- )
286
- return list(res.task_ids)
287
88
 
89
+ @abstractmethod
288
90
  def pull_messages(self, message_ids: Iterable[str]) -> Iterable[Message]:
289
91
  """Pull messages based on message IDs.
290
92
 
@@ -301,15 +103,8 @@ class Driver:
301
103
  messages : Iterable[Message]
302
104
  An iterable of messages received.
303
105
  """
304
- grpc_driver, _ = self._get_grpc_driver_helper_and_run_id()
305
- # Pull TaskRes
306
- res = grpc_driver.pull_task_res(
307
- PullTaskResRequest(node=self.node, task_ids=message_ids)
308
- )
309
- # Convert TaskRes to Message
310
- msgs = [message_from_taskres(taskres) for taskres in res.task_res_list]
311
- return msgs
312
106
 
107
+ @abstractmethod
313
108
  def send_and_receive(
314
109
  self,
315
110
  messages: Iterable[Message],
@@ -343,28 +138,3 @@ class Driver:
343
138
  replies for all sent messages. A message remains valid until its TTL,
344
139
  which is not affected by `timeout`.
345
140
  """
346
- # Push messages
347
- msg_ids = set(self.push_messages(messages))
348
-
349
- # Pull messages
350
- end_time = time.time() + (timeout if timeout is not None else 0.0)
351
- ret: List[Message] = []
352
- while timeout is None or time.time() < end_time:
353
- res_msgs = self.pull_messages(msg_ids)
354
- ret.extend(res_msgs)
355
- msg_ids.difference_update(
356
- {msg.metadata.reply_to_message for msg in res_msgs}
357
- )
358
- if len(msg_ids) == 0:
359
- break
360
- # Sleep
361
- time.sleep(3)
362
- return ret
363
-
364
- def close(self) -> None:
365
- """Disconnect from the SuperLink if connected."""
366
- # Check if GrpcDriverHelper is initialized
367
- if self.grpc_driver_helper is None:
368
- return
369
- # Disconnect
370
- self.grpc_driver_helper.disconnect()
@@ -0,0 +1,315 @@
1
+ # Copyright 2022 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 gRPC Driver."""
16
+
17
+ import time
18
+ import warnings
19
+ from logging import DEBUG, ERROR, WARNING
20
+ from typing import Iterable, List, Optional, Tuple
21
+
22
+ import grpc
23
+
24
+ from flwr.common import DEFAULT_TTL, EventType, Message, Metadata, RecordSet, event
25
+ from flwr.common.grpc import create_channel
26
+ from flwr.common.logger import log
27
+ from flwr.common.serde import message_from_taskres, message_to_taskins
28
+ from flwr.proto.driver_pb2 import ( # pylint: disable=E0611
29
+ CreateRunRequest,
30
+ CreateRunResponse,
31
+ GetNodesRequest,
32
+ GetNodesResponse,
33
+ PullTaskResRequest,
34
+ PullTaskResResponse,
35
+ PushTaskInsRequest,
36
+ PushTaskInsResponse,
37
+ )
38
+ from flwr.proto.driver_pb2_grpc import DriverStub # pylint: disable=E0611
39
+ from flwr.proto.node_pb2 import Node # pylint: disable=E0611
40
+ from flwr.proto.task_pb2 import TaskIns # pylint: disable=E0611
41
+
42
+ from .driver import Driver
43
+
44
+ DEFAULT_SERVER_ADDRESS_DRIVER = "[::]:9091"
45
+
46
+ ERROR_MESSAGE_DRIVER_NOT_CONNECTED = """
47
+ [Driver] Error: Not connected.
48
+
49
+ Call `connect()` on the `GrpcDriverHelper` instance before calling any of the other
50
+ `GrpcDriverHelper` methods.
51
+ """
52
+
53
+
54
+ class GrpcDriverHelper:
55
+ """`GrpcDriverHelper` provides access to the gRPC Driver API/service."""
56
+
57
+ def __init__(
58
+ self,
59
+ driver_service_address: str = DEFAULT_SERVER_ADDRESS_DRIVER,
60
+ root_certificates: Optional[bytes] = None,
61
+ ) -> None:
62
+ self.driver_service_address = driver_service_address
63
+ self.root_certificates = root_certificates
64
+ self.channel: Optional[grpc.Channel] = None
65
+ self.stub: Optional[DriverStub] = None
66
+
67
+ def connect(self) -> None:
68
+ """Connect to the Driver API."""
69
+ event(EventType.DRIVER_CONNECT)
70
+ if self.channel is not None or self.stub is not None:
71
+ log(WARNING, "Already connected")
72
+ return
73
+ self.channel = create_channel(
74
+ server_address=self.driver_service_address,
75
+ insecure=(self.root_certificates is None),
76
+ root_certificates=self.root_certificates,
77
+ )
78
+ self.stub = DriverStub(self.channel)
79
+ log(DEBUG, "[Driver] Connected to %s", self.driver_service_address)
80
+
81
+ def disconnect(self) -> None:
82
+ """Disconnect from the Driver API."""
83
+ event(EventType.DRIVER_DISCONNECT)
84
+ if self.channel is None or self.stub is None:
85
+ log(DEBUG, "Already disconnected")
86
+ return
87
+ channel = self.channel
88
+ self.channel = None
89
+ self.stub = None
90
+ channel.close()
91
+ log(DEBUG, "[Driver] Disconnected")
92
+
93
+ def create_run(self, req: CreateRunRequest) -> CreateRunResponse:
94
+ """Request for run ID."""
95
+ # Check if channel is open
96
+ if self.stub is None:
97
+ log(ERROR, ERROR_MESSAGE_DRIVER_NOT_CONNECTED)
98
+ raise ConnectionError("`GrpcDriverHelper` instance not connected")
99
+
100
+ # Call Driver API
101
+ res: CreateRunResponse = self.stub.CreateRun(request=req)
102
+ return res
103
+
104
+ def get_nodes(self, req: GetNodesRequest) -> GetNodesResponse:
105
+ """Get client IDs."""
106
+ # Check if channel is open
107
+ if self.stub is None:
108
+ log(ERROR, ERROR_MESSAGE_DRIVER_NOT_CONNECTED)
109
+ raise ConnectionError("`GrpcDriverHelper` instance not connected")
110
+
111
+ # Call gRPC Driver API
112
+ res: GetNodesResponse = self.stub.GetNodes(request=req)
113
+ return res
114
+
115
+ def push_task_ins(self, req: PushTaskInsRequest) -> PushTaskInsResponse:
116
+ """Schedule tasks."""
117
+ # Check if channel is open
118
+ if self.stub is None:
119
+ log(ERROR, ERROR_MESSAGE_DRIVER_NOT_CONNECTED)
120
+ raise ConnectionError("`GrpcDriverHelper` instance not connected")
121
+
122
+ # Call gRPC Driver API
123
+ res: PushTaskInsResponse = self.stub.PushTaskIns(request=req)
124
+ return res
125
+
126
+ def pull_task_res(self, req: PullTaskResRequest) -> PullTaskResResponse:
127
+ """Get task results."""
128
+ # Check if channel is open
129
+ if self.stub is None:
130
+ log(ERROR, ERROR_MESSAGE_DRIVER_NOT_CONNECTED)
131
+ raise ConnectionError("`GrpcDriverHelper` instance not connected")
132
+
133
+ # Call Driver API
134
+ res: PullTaskResResponse = self.stub.PullTaskRes(request=req)
135
+ return res
136
+
137
+
138
+ class GrpcDriver(Driver):
139
+ """`Driver` class provides an interface to the Driver API.
140
+
141
+ Parameters
142
+ ----------
143
+ driver_service_address : Optional[str]
144
+ The IPv4 or IPv6 address of the Driver API server.
145
+ Defaults to `"[::]:9091"`.
146
+ certificates : bytes (default: None)
147
+ Tuple containing root certificate, server certificate, and private key
148
+ to start a secure SSL-enabled server. The tuple is expected to have
149
+ three bytes elements in the following order:
150
+
151
+ * CA certificate.
152
+ * server certificate.
153
+ * server private key.
154
+ fab_id : str (default: None)
155
+ The identifier of the FAB used in the run.
156
+ fab_version : str (default: None)
157
+ The version of the FAB used in the run.
158
+ """
159
+
160
+ def __init__(
161
+ self,
162
+ driver_service_address: str = DEFAULT_SERVER_ADDRESS_DRIVER,
163
+ root_certificates: Optional[bytes] = None,
164
+ fab_id: Optional[str] = None,
165
+ fab_version: Optional[str] = None,
166
+ ) -> None:
167
+ self.addr = driver_service_address
168
+ self.root_certificates = root_certificates
169
+ self.driver_helper: Optional[GrpcDriverHelper] = None
170
+ self.run_id: Optional[int] = None
171
+ self.fab_id = fab_id if fab_id is not None else ""
172
+ self.fab_version = fab_version if fab_version is not None else ""
173
+ self.node = Node(node_id=0, anonymous=True)
174
+
175
+ def _get_grpc_driver_helper_and_run_id(self) -> Tuple[GrpcDriverHelper, int]:
176
+ # Check if the GrpcDriverHelper is initialized
177
+ if self.driver_helper is None or self.run_id is None:
178
+ # Connect and create run
179
+ self.driver_helper = GrpcDriverHelper(
180
+ driver_service_address=self.addr,
181
+ root_certificates=self.root_certificates,
182
+ )
183
+ self.driver_helper.connect()
184
+ req = CreateRunRequest(fab_id=self.fab_id, fab_version=self.fab_version)
185
+ res = self.driver_helper.create_run(req)
186
+ self.run_id = res.run_id
187
+ return self.driver_helper, self.run_id
188
+
189
+ def _check_message(self, message: Message) -> None:
190
+ # Check if the message is valid
191
+ if not (
192
+ message.metadata.run_id == self.run_id
193
+ and message.metadata.src_node_id == self.node.node_id
194
+ and message.metadata.message_id == ""
195
+ and message.metadata.reply_to_message == ""
196
+ and message.metadata.ttl > 0
197
+ ):
198
+ raise ValueError(f"Invalid message: {message}")
199
+
200
+ def create_message( # pylint: disable=too-many-arguments
201
+ self,
202
+ content: RecordSet,
203
+ message_type: str,
204
+ dst_node_id: int,
205
+ group_id: str,
206
+ ttl: Optional[float] = None,
207
+ ) -> Message:
208
+ """Create a new message with specified parameters.
209
+
210
+ This method constructs a new `Message` with given content and metadata.
211
+ The `run_id` and `src_node_id` will be set automatically.
212
+ """
213
+ _, run_id = self._get_grpc_driver_helper_and_run_id()
214
+ if ttl:
215
+ warnings.warn(
216
+ "A custom TTL was set, but note that the SuperLink does not enforce "
217
+ "the TTL yet. The SuperLink will start enforcing the TTL in a future "
218
+ "version of Flower.",
219
+ stacklevel=2,
220
+ )
221
+
222
+ ttl_ = DEFAULT_TTL if ttl is None else ttl
223
+ metadata = Metadata(
224
+ run_id=run_id,
225
+ message_id="", # Will be set by the server
226
+ src_node_id=self.node.node_id,
227
+ dst_node_id=dst_node_id,
228
+ reply_to_message="",
229
+ group_id=group_id,
230
+ ttl=ttl_,
231
+ message_type=message_type,
232
+ )
233
+ return Message(metadata=metadata, content=content)
234
+
235
+ def get_node_ids(self) -> List[int]:
236
+ """Get node IDs."""
237
+ grpc_driver_helper, run_id = self._get_grpc_driver_helper_and_run_id()
238
+ # Call GrpcDriverHelper method
239
+ res = grpc_driver_helper.get_nodes(GetNodesRequest(run_id=run_id))
240
+ return [node.node_id for node in res.nodes]
241
+
242
+ def push_messages(self, messages: Iterable[Message]) -> Iterable[str]:
243
+ """Push messages to specified node IDs.
244
+
245
+ This method takes an iterable of messages and sends each message
246
+ to the node specified in `dst_node_id`.
247
+ """
248
+ grpc_driver_helper, _ = self._get_grpc_driver_helper_and_run_id()
249
+ # Construct TaskIns
250
+ task_ins_list: List[TaskIns] = []
251
+ for msg in messages:
252
+ # Check message
253
+ self._check_message(msg)
254
+ # Convert Message to TaskIns
255
+ taskins = message_to_taskins(msg)
256
+ # Add to list
257
+ task_ins_list.append(taskins)
258
+ # Call GrpcDriverHelper method
259
+ res = grpc_driver_helper.push_task_ins(
260
+ PushTaskInsRequest(task_ins_list=task_ins_list)
261
+ )
262
+ return list(res.task_ids)
263
+
264
+ def pull_messages(self, message_ids: Iterable[str]) -> Iterable[Message]:
265
+ """Pull messages based on message IDs.
266
+
267
+ This method is used to collect messages from the SuperLink that correspond to a
268
+ set of given message IDs.
269
+ """
270
+ grpc_driver, _ = self._get_grpc_driver_helper_and_run_id()
271
+ # Pull TaskRes
272
+ res = grpc_driver.pull_task_res(
273
+ PullTaskResRequest(node=self.node, task_ids=message_ids)
274
+ )
275
+ # Convert TaskRes to Message
276
+ msgs = [message_from_taskres(taskres) for taskres in res.task_res_list]
277
+ return msgs
278
+
279
+ def send_and_receive(
280
+ self,
281
+ messages: Iterable[Message],
282
+ *,
283
+ timeout: Optional[float] = None,
284
+ ) -> Iterable[Message]:
285
+ """Push messages to specified node IDs and pull the reply messages.
286
+
287
+ This method sends a list of messages to their destination node IDs and then
288
+ waits for the replies. It continues to pull replies until either all replies are
289
+ received or the specified timeout duration is exceeded.
290
+ """
291
+ # Push messages
292
+ msg_ids = set(self.push_messages(messages))
293
+
294
+ # Pull messages
295
+ end_time = time.time() + (timeout if timeout is not None else 0.0)
296
+ ret: List[Message] = []
297
+ while timeout is None or time.time() < end_time:
298
+ res_msgs = self.pull_messages(msg_ids)
299
+ ret.extend(res_msgs)
300
+ msg_ids.difference_update(
301
+ {msg.metadata.reply_to_message for msg in res_msgs}
302
+ )
303
+ if len(msg_ids) == 0:
304
+ break
305
+ # Sleep
306
+ time.sleep(3)
307
+ return ret
308
+
309
+ def close(self) -> None:
310
+ """Disconnect from the SuperLink if connected."""
311
+ # Check if GrpcDriverHelper is initialized
312
+ if self.driver_helper is None:
313
+ return
314
+ # Disconnect
315
+ self.driver_helper.disconnect()
flwr/server/history.py CHANGED
@@ -91,32 +91,32 @@ class History:
91
91
  """
92
92
  rep = ""
93
93
  if self.losses_distributed:
94
- rep += "History (loss, distributed):\n" + pprint.pformat(
95
- reduce(
96
- lambda a, b: a + b,
97
- [
98
- f"\tround {server_round}: {loss}\n"
99
- for server_round, loss in self.losses_distributed
100
- ],
101
- )
94
+ rep += "History (loss, distributed):\n" + reduce(
95
+ lambda a, b: a + b,
96
+ [
97
+ f"\tround {server_round}: {loss}\n"
98
+ for server_round, loss in self.losses_distributed
99
+ ],
102
100
  )
103
101
  if self.losses_centralized:
104
- rep += "History (loss, centralized):\n" + pprint.pformat(
105
- reduce(
106
- lambda a, b: a + b,
107
- [
108
- f"\tround {server_round}: {loss}\n"
109
- for server_round, loss in self.losses_centralized
110
- ],
111
- )
102
+ rep += "History (loss, centralized):\n" + reduce(
103
+ lambda a, b: a + b,
104
+ [
105
+ f"\tround {server_round}: {loss}\n"
106
+ for server_round, loss in self.losses_centralized
107
+ ],
112
108
  )
113
109
  if self.metrics_distributed_fit:
114
- rep += "History (metrics, distributed, fit):\n" + pprint.pformat(
115
- self.metrics_distributed_fit
110
+ rep += (
111
+ "History (metrics, distributed, fit):\n"
112
+ + pprint.pformat(self.metrics_distributed_fit)
113
+ + "\n"
116
114
  )
117
115
  if self.metrics_distributed:
118
- rep += "History (metrics, distributed, evaluate):\n" + pprint.pformat(
119
- self.metrics_distributed
116
+ rep += (
117
+ "History (metrics, distributed, evaluate):\n"
118
+ + pprint.pformat(self.metrics_distributed)
119
+ + "\n"
120
120
  )
121
121
  if self.metrics_centralized:
122
122
  rep += "History (metrics, centralized):\n" + pprint.pformat(