flwr 1.15.2__py3-none-any.whl → 1.17.0__py3-none-any.whl

flwr/server/utils/validator.py

@@ -16,93 +16,78 @@
 
 
 import time
-from typing import Union
 
+from flwr.common import Message
 from flwr.common.constant import SUPERLINK_NODE_ID
-from flwr.proto.task_pb2 import TaskIns, TaskRes  # pylint: disable=E0611
 
 
-# pylint: disable-next=too-many-branches,too-many-statements
-def validate_task_ins_or_res(tasks_ins_res: Union[TaskIns, TaskRes]) -> list[str]:
-    """Validate a TaskIns or TaskRes."""
+# pylint: disable-next=too-many-branches
+def validate_message(message: Message, is_reply_message: bool) -> list[str]:
+    """Validate a Message."""
     validation_errors = []
+    metadata = message.metadata
 
-    if tasks_ins_res.task_id != "":
-        validation_errors.append("non-empty `task_id`")
-
-    if not tasks_ins_res.HasField("task"):
-        validation_errors.append("`task` does not set field `task`")
+    if metadata.message_id != "":
+        validation_errors.append("non-empty `metadata.message_id`")
 
     # Created/delivered/TTL/Pushed
     if (
-        tasks_ins_res.task.created_at < 1711497600.0
-    ):  # unix timestamp of 27 March 2024 00h:00m:00s UTC
+        metadata.created_at < 1740700800.0
+    ):  # unix timestamp of 28 February 2025 00h:00m:00s UTC
         validation_errors.append(
-            "`created_at` must be a float that records the unix timestamp "
+            "`metadata.created_at` must be a float that records the unix timestamp "
             "in seconds when the message was created."
         )
-    if tasks_ins_res.task.delivered_at != "":
-        validation_errors.append("`delivered_at` must be an empty str")
-    if tasks_ins_res.task.ttl <= 0:
-        validation_errors.append("`ttl` must be higher than zero")
+    if metadata.delivered_at != "":
+        validation_errors.append("`metadata.delivered_at` must be an empty str")
+    if metadata.ttl <= 0:
+        validation_errors.append("`metadata.ttl` must be higher than zero")
 
     # Verify TTL and created_at time
     current_time = time.time()
-    if tasks_ins_res.task.created_at + tasks_ins_res.task.ttl <= current_time:
-        validation_errors.append("Task TTL has expired")
-
-    # TaskIns specific
-    if isinstance(tasks_ins_res, TaskIns):
-        # Task producer
-        if not tasks_ins_res.task.HasField("producer"):
-            validation_errors.append("`producer` does not set field `producer`")
-        if tasks_ins_res.task.producer.node_id != SUPERLINK_NODE_ID:
-            validation_errors.append(f"`producer.node_id` is not {SUPERLINK_NODE_ID}")
-
-        # Task consumer
-        if not tasks_ins_res.task.HasField("consumer"):
-            validation_errors.append("`consumer` does not set field `consumer`")
-        if tasks_ins_res.task.consumer.node_id == SUPERLINK_NODE_ID:
-            validation_errors.append("consumer MUST provide a valid `node_id`")
+    if metadata.created_at + metadata.ttl <= current_time:
+        validation_errors.append("Message TTL has expired")
 
-        # Content check
-        if tasks_ins_res.task.task_type == "":
-            validation_errors.append("`task_type` MUST be set")
-        if not (
-            tasks_ins_res.task.HasField("recordset")
-            ^ tasks_ins_res.task.HasField("error")
-        ):
-            validation_errors.append("Either `recordset` or `error` MUST be set")
+    # Source node is set and is not zero
+    if not metadata.src_node_id:
+        validation_errors.append("`metadata.src_node_id` is not set.")
 
-        # Ancestors
-        if len(tasks_ins_res.task.ancestry) != 0:
-            validation_errors.append("`ancestry` is not empty")
+    # Destination node is set and is not zero
+    if not metadata.dst_node_id:
+        validation_errors.append("`metadata.dst_node_id` is not set.")
 
-    # TaskRes specific
-    if isinstance(tasks_ins_res, TaskRes):
-        # Task producer
-        if not tasks_ins_res.task.HasField("producer"):
-            validation_errors.append("`producer` does not set field `producer`")
-        if tasks_ins_res.task.producer.node_id == SUPERLINK_NODE_ID:
-            validation_errors.append("producer MUST provide a valid `node_id`")
+    # Message type
+    if metadata.message_type == "":
+        validation_errors.append("`metadata.message_type` MUST be set")
 
-        # Task consumer
-        if not tasks_ins_res.task.HasField("consumer"):
-            validation_errors.append("`consumer` does not set field `consumer`")
-        if tasks_ins_res.task.consumer.node_id != SUPERLINK_NODE_ID:
-            validation_errors.append(f"consumer is not {SUPERLINK_NODE_ID}")
-
-        # Content check
-        if tasks_ins_res.task.task_type == "":
-            validation_errors.append("`task_type` MUST be set")
-        if not (
-            tasks_ins_res.task.HasField("recordset")
-            ^ tasks_ins_res.task.HasField("error")
-        ):
-            validation_errors.append("Either `recordset` or `error` MUST be set")
+    # Content
+    if not message.has_content() != message.has_error():
+        validation_errors.append(
+            "Either message `content` or `error` MUST be set (but not both)"
+        )
 
-        # Ancestors
-        if len(tasks_ins_res.task.ancestry) == 0:
-            validation_errors.append("`ancestry` is empty")
+    # Link respose to original message
+    if not is_reply_message:
+        if metadata.reply_to_message_id != "":
+            validation_errors.append("`metadata.reply_to_message_id` MUST not be set.")
+        if metadata.src_node_id != SUPERLINK_NODE_ID:
+            validation_errors.append(
+                f"`metadata.src_node_id` is not {SUPERLINK_NODE_ID} (SuperLink node ID)"
+            )
+        if metadata.dst_node_id == SUPERLINK_NODE_ID:
+            validation_errors.append(
+                f"`metadata.dst_node_id` is {SUPERLINK_NODE_ID} (SuperLink node ID)"
+            )
+    else:
+        if metadata.reply_to_message_id == "":
+            validation_errors.append("`metadata.reply_to_message_id` MUST be set.")
+        if metadata.src_node_id == SUPERLINK_NODE_ID:
+            validation_errors.append(
+                f"`metadata.src_node_id` is {SUPERLINK_NODE_ID} (SuperLink node ID)"
+            )
+        if metadata.dst_node_id != SUPERLINK_NODE_ID:
+            validation_errors.append(
+                f"`metadata.dst_node_id` is not {SUPERLINK_NODE_ID} (SuperLink node ID)"
+            )
 
     return validation_errors

flwr/server/workflow/default_workflows.py

@@ -20,15 +20,16 @@ import timeit
 from logging import INFO, WARN
 from typing import Optional, Union, cast
 
-import flwr.common.recordset_compat as compat
+import flwr.common.recorddict_compat as compat
 from flwr.common import (
+    ArrayRecord,
     Code,
-    ConfigsRecord,
+    ConfigRecord,
     Context,
     EvaluateRes,
     FitRes,
     GetParametersIns,
-    ParametersRecord,
+    Message,
     log,
 )
 from flwr.common.constant import MessageType, MessageTypeLegacy
@@ -36,7 +37,7 @@ from flwr.common.constant import MessageType, MessageTypeLegacy
 from ..client_proxy import ClientProxy
 from ..compat.app_utils import start_update_client_manager_thread
 from ..compat.legacy_context import LegacyContext
-from ..driver import Driver
+from ..grid import Grid
 from ..typing import Workflow
 from .constant import MAIN_CONFIGS_RECORD, MAIN_PARAMS_RECORD, Key
 
@@ -56,7 +57,7 @@ class DefaultWorkflow:
         self.fit_workflow: Workflow = fit_workflow
         self.evaluate_workflow: Workflow = evaluate_workflow
 
-    def __call__(self, driver: Driver, context: Context) -> None:
+    def __call__(self, grid: Grid, context: Context) -> None:
         """Execute the workflow."""
         if not isinstance(context, LegacyContext):
             raise TypeError(
@@ -64,19 +65,22 @@ class DefaultWorkflow:
             )
 
         # Start the thread updating nodes
-        thread, f_stop = start_update_client_manager_thread(
-            driver, context.client_manager
+        thread, f_stop, c_done = start_update_client_manager_thread(
+            grid, context.client_manager
         )
 
+        # Wait until the node registration done
+        c_done.wait()
+
         # Initialize parameters
         log(INFO, "[INIT]")
-        default_init_params_workflow(driver, context)
+        default_init_params_workflow(grid, context)
 
         # Run federated learning for num_rounds
         start_time = timeit.default_timer()
-        cfg = ConfigsRecord()
+        cfg = ConfigRecord()
         cfg[Key.START_TIME] = start_time
-        context.state.configs_records[MAIN_CONFIGS_RECORD] = cfg
+        context.state.config_records[MAIN_CONFIGS_RECORD] = cfg
 
         for current_round in range(1, context.config.num_rounds + 1):
             log(INFO, "")
@@ -84,13 +88,13 @@ class DefaultWorkflow:
             cfg[Key.CURRENT_ROUND] = current_round
 
             # Fit round
-            self.fit_workflow(driver, context)
+            self.fit_workflow(grid, context)
 
             # Centralized evaluation
-            default_centralized_evaluation_workflow(driver, context)
+            default_centralized_evaluation_workflow(grid, context)
 
             # Evaluate round
-            self.evaluate_workflow(driver, context)
+            self.evaluate_workflow(grid, context)
 
         # Bookkeeping and log results
         end_time = timeit.default_timer()
@@ -116,7 +120,7 @@ class DefaultWorkflow:
         thread.join()
 
 
-def default_init_params_workflow(driver: Driver, context: Context) -> None:
+def default_init_params_workflow(grid: Grid, context: Context) -> None:
     """Execute the default workflow for parameters initialization."""
     if not isinstance(context, LegacyContext):
         raise TypeError(f"Expect a LegacyContext, but get {type(context).__name__}.")
@@ -126,21 +130,19 @@ def default_init_params_workflow(driver: Driver, context: Context) -> None:
     )
     if parameters is not None:
         log(INFO, "Using initial global parameters provided by strategy")
-        paramsrecord = compat.parameters_to_parametersrecord(
-            parameters, keep_input=True
-        )
+        arr_record = compat.parameters_to_arrayrecord(parameters, keep_input=True)
     else:
         # Get initial parameters from one of the clients
         log(INFO, "Requesting initial parameters from one random client")
         random_client = context.client_manager.sample(1)[0]
         # Send GetParametersIns and get the response
-        content = compat.getparametersins_to_recordset(GetParametersIns({}))
-        messages = driver.send_and_receive(
+        content = compat.getparametersins_to_recorddict(GetParametersIns({}))
+        messages = grid.send_and_receive(
             [
-                driver.create_message(
+                Message(
                     content=content,
-                    message_type=MessageTypeLegacy.GET_PARAMETERS,
                     dst_node_id=random_client.node_id,
+                    message_type=MessageTypeLegacy.GET_PARAMETERS,
                     group_id="0",
                 )
             ]
@@ -149,26 +151,26 @@ def default_init_params_workflow(driver: Driver, context: Context) -> None:
 
         if (
             msg.has_content()
-            and compat._extract_status_from_recordset(  # pylint: disable=W0212
+            and compat._extract_status_from_recorddict(  # pylint: disable=W0212
                 "getparametersres", msg.content
             ).code
             == Code.OK
         ):
             log(INFO, "Received initial parameters from one random client")
-            paramsrecord = next(iter(msg.content.parameters_records.values()))
+            arr_record = next(iter(msg.content.array_records.values()))
         else:
             log(
                 WARN,
                 "Failed to receive initial parameters from the client."
                 " Empty initial parameters will be used.",
             )
-            paramsrecord = ParametersRecord()
+            arr_record = ArrayRecord()
 
-    context.state.parameters_records[MAIN_PARAMS_RECORD] = paramsrecord
+    context.state.array_records[MAIN_PARAMS_RECORD] = arr_record
 
     # Evaluate initial parameters
     log(INFO, "Starting evaluation of initial global parameters")
-    parameters = compat.parametersrecord_to_parameters(paramsrecord, keep_input=True)
+    parameters = compat.arrayrecord_to_parameters(arr_record, keep_input=True)
     res = context.strategy.evaluate(0, parameters=parameters)
     if res is not None:
         log(
@@ -183,19 +185,19 @@ def default_init_params_workflow(driver: Driver, context: Context) -> None:
         log(INFO, "Evaluation returned no results (`None`)")
 
 
-def default_centralized_evaluation_workflow(_: Driver, context: Context) -> None:
+def default_centralized_evaluation_workflow(_: Grid, context: Context) -> None:
     """Execute the default workflow for centralized evaluation."""
     if not isinstance(context, LegacyContext):
         raise TypeError(f"Expect a LegacyContext, but get {type(context).__name__}.")
 
     # Retrieve current_round and start_time from the context
-    cfg = context.state.configs_records[MAIN_CONFIGS_RECORD]
+    cfg = context.state.config_records[MAIN_CONFIGS_RECORD]
     current_round = cast(int, cfg[Key.CURRENT_ROUND])
     start_time = cast(float, cfg[Key.START_TIME])
 
     # Centralized evaluation
-    parameters = compat.parametersrecord_to_parameters(
-        record=context.state.parameters_records[MAIN_PARAMS_RECORD],
+    parameters = compat.arrayrecord_to_parameters(
+        record=context.state.array_records[MAIN_PARAMS_RECORD],
         keep_input=True,
     )
     res_cen = context.strategy.evaluate(current_round, parameters=parameters)
@@ -215,20 +217,16 @@ def default_centralized_evaluation_workflow(_: Driver, context: Context) -> None
         )
 
 
-def default_fit_workflow(  # pylint: disable=R0914
-    driver: Driver, context: Context
-) -> None:
+def default_fit_workflow(grid: Grid, context: Context) -> None:  # pylint: disable=R0914
     """Execute the default workflow for a single fit round."""
     if not isinstance(context, LegacyContext):
         raise TypeError(f"Expect a LegacyContext, but get {type(context).__name__}.")
 
     # Get current_round and parameters
-    cfg = context.state.configs_records[MAIN_CONFIGS_RECORD]
+    cfg = context.state.config_records[MAIN_CONFIGS_RECORD]
     current_round = cast(int, cfg[Key.CURRENT_ROUND])
-    parametersrecord = context.state.parameters_records[MAIN_PARAMS_RECORD]
-    parameters = compat.parametersrecord_to_parameters(
-        parametersrecord, keep_input=True
-    )
+    arr_record = context.state.array_records[MAIN_PARAMS_RECORD]
+    parameters = compat.arrayrecord_to_parameters(arr_record, keep_input=True)
 
     # Get clients and their respective instructions from strategy
     client_instructions = context.strategy.configure_fit(
@@ -252,10 +250,10 @@ def default_fit_workflow(  # pylint: disable=R0914
 
     # Build out messages
     out_messages = [
-        driver.create_message(
-            content=compat.fitins_to_recordset(fitins, True),
-            message_type=MessageType.TRAIN,
+        Message(
+            content=compat.fitins_to_recorddict(fitins, True),
             dst_node_id=proxy.node_id,
+            message_type=MessageType.TRAIN,
             group_id=str(current_round),
         )
         for proxy, fitins in client_instructions
@@ -263,7 +261,7 @@ def default_fit_workflow(  # pylint: disable=R0914
 
     # Send instructions to clients and
     # collect `fit` results from all clients participating in this round
-    messages = list(driver.send_and_receive(out_messages))
+    messages = list(grid.send_and_receive(out_messages))
     del out_messages
     num_failures = len([msg for msg in messages if msg.has_error()])
 
@@ -281,7 +279,7 @@ def default_fit_workflow(  # pylint: disable=R0914
     for msg in messages:
         if msg.has_content():
             proxy = node_id_to_proxy[msg.metadata.src_node_id]
-            fitres = compat.recordset_to_fitres(msg.content, False)
+            fitres = compat.recorddict_to_fitres(msg.content, False)
             if fitres.status.code == Code.OK:
                 results.append((proxy, fitres))
             else:
@@ -294,28 +292,24 @@ def default_fit_workflow(  # pylint: disable=R0914
 
     # Update the parameters and write history
     if parameters_aggregated:
-        paramsrecord = compat.parameters_to_parametersrecord(
-            parameters_aggregated, True
-        )
-        context.state.parameters_records[MAIN_PARAMS_RECORD] = paramsrecord
+        arr_record = compat.parameters_to_arrayrecord(parameters_aggregated, True)
+        context.state.array_records[MAIN_PARAMS_RECORD] = arr_record
         context.history.add_metrics_distributed_fit(
             server_round=current_round, metrics=metrics_aggregated
         )
 
 
 # pylint: disable-next=R0914
-def default_evaluate_workflow(driver: Driver, context: Context) -> None:
+def default_evaluate_workflow(grid: Grid, context: Context) -> None:
     """Execute the default workflow for a single evaluate round."""
     if not isinstance(context, LegacyContext):
         raise TypeError(f"Expect a LegacyContext, but get {type(context).__name__}.")
 
     # Get current_round and parameters
-    cfg = context.state.configs_records[MAIN_CONFIGS_RECORD]
+    cfg = context.state.config_records[MAIN_CONFIGS_RECORD]
     current_round = cast(int, cfg[Key.CURRENT_ROUND])
-    parametersrecord = context.state.parameters_records[MAIN_PARAMS_RECORD]
-    parameters = compat.parametersrecord_to_parameters(
-        parametersrecord, keep_input=True
-    )
+    arr_record = context.state.array_records[MAIN_PARAMS_RECORD]
+    parameters = compat.arrayrecord_to_parameters(arr_record, keep_input=True)
 
     # Get clients and their respective instructions from strategy
     client_instructions = context.strategy.configure_evaluate(
@@ -338,10 +332,10 @@ def default_evaluate_workflow(driver: Driver, context: Context) -> None:
 
     # Build out messages
     out_messages = [
-        driver.create_message(
-            content=compat.evaluateins_to_recordset(evalins, True),
-            message_type=MessageType.EVALUATE,
+        Message(
+            content=compat.evaluateins_to_recorddict(evalins, True),
             dst_node_id=proxy.node_id,
+            message_type=MessageType.EVALUATE,
             group_id=str(current_round),
         )
         for proxy, evalins in client_instructions
@@ -349,7 +343,7 @@ def default_evaluate_workflow(driver: Driver, context: Context) -> None:
 
     # Send instructions to clients and
     # collect `evaluate` results from all clients participating in this round
-    messages = list(driver.send_and_receive(out_messages))
+    messages = list(grid.send_and_receive(out_messages))
     del out_messages
     num_failures = len([msg for msg in messages if msg.has_error()])
 
@@ -367,7 +361,7 @@ def default_evaluate_workflow(driver: Driver, context: Context) -> None:
     for msg in messages:
         if msg.has_content():
             proxy = node_id_to_proxy[msg.metadata.src_node_id]
-            evalres = compat.recordset_to_evaluateres(msg.content)
+            evalres = compat.recorddict_to_evaluateres(msg.content)
             if evalres.status.code == Code.OK:
                 results.append((proxy, evalres))
             else: