w/ set_stat()/get_state()

src/py/flwr/client/client.py

@@ -141,6 +141,14 @@ def evaluate(self, ins: EvaluateIns) -> EvaluateRes:
             metrics={},
         )
 
+    def get_state(self) -> WorkloadState:
+        """Get the workload state from this client."""
+        return self.state
+
+    def set_state(self, state: WorkloadState) -> None:
+        """Apply a workload state to this client."""
+        self.state = state
+
     def to_client(self) -> Client:
         """Return client (itself)."""
         return self

src/py/flwr/client/message_handler/message_handler.py

@@ -101,7 +101,7 @@ def handle(
     if server_msg is None:
         # Instantiate the client
         client = client_fn("-1")
-        client.state = state
+        client.set_state(state)
         # Secure Aggregation
         if task_ins.task.HasField("sa") and isinstance(
             client, SecureAggregationHandler
@@ -118,7 +118,7 @@ def handle(
                     sa=SecureAggregation(named_values=serde.named_values_to_proto(res)),
                 ),
             )
-            return task_res, client.state
+            return task_res, client.get_state()
         raise NotImplementedError()
     client_msg, updated_state = handle_legacy_message(client_fn, state, server_msg)
     task_res = wrap_client_message_in_task_res(client_msg)
@@ -152,7 +152,7 @@ def handle_legacy_message(
 
     # Instantiate the client
     client = client_fn("-1")
-    client.state = state
+    client.set_state(state)
     # Execute task
     message = None
     if field == "get_properties_ins":
@@ -164,7 +164,7 @@ def handle_legacy_message(
     if field == "evaluate_ins":
         message = _evaluate(client, server_msg.evaluate_ins)
     if message:
-        return message, client.state
+        return message, client.get_state()
     raise UnknownServerMessage()