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

flwr/server/superlink/linkstate/utils.py

@@ -15,29 +15,26 @@
 """Utility functions for State."""
 
 
-from logging import ERROR
 from os import urandom
-from typing import Optional, Union
+from typing import Optional
 from uuid import UUID, uuid4
 
-from flwr.common import ConfigsRecord, Context, log, now, serde
-from flwr.common.constant import SUPERLINK_NODE_ID, ErrorCode, Status, SubStatus
+from flwr.common import ConfigRecord, Context, Error, Message, Metadata, now, serde
+from flwr.common.constant import (
+    SUPERLINK_NODE_ID,
+    ErrorCode,
+    MessageType,
+    Status,
+    SubStatus,
+)
+from flwr.common.message import make_message
 from flwr.common.typing import RunStatus
 
 # pylint: disable=E0611
-from flwr.proto.error_pb2 import Error
 from flwr.proto.message_pb2 import Context as ProtoContext
-from flwr.proto.node_pb2 import Node
-from flwr.proto.recordset_pb2 import ConfigsRecord as ProtoConfigsRecord
-from flwr.proto.task_pb2 import Task, TaskIns, TaskRes
+from flwr.proto.recorddict_pb2 import ConfigRecord as ProtoConfigRecord
 
 # pylint: enable=E0611
-
-NODE_UNAVAILABLE_ERROR_REASON = (
-    "Error: Node Unavailable - The destination node is currently unavailable. "
-    "It exceeds the time limit specified in its last ping."
-)
-
 VALID_RUN_STATUS_TRANSITIONS = {
     (Status.PENDING, Status.STARTING),
     (Status.STARTING, Status.RUNNING),
@@ -58,6 +55,10 @@ MESSAGE_UNAVAILABLE_ERROR_REASON = (
 REPLY_MESSAGE_UNAVAILABLE_ERROR_REASON = (
     "Error: Reply Message Unavailable - The reply message has expired."
 )
+NODE_UNAVAILABLE_ERROR_REASON = (
+    "Error: Node Unavailable - The destination node is currently unavailable. "
+    "It exceeds twice the time limit specified in its last ping."
+)
 
 
 def generate_rand_int_from_bytes(
@@ -171,15 +172,15 @@ def context_from_bytes(context_bytes: bytes) -> Context:
     return serde.context_from_proto(ProtoContext.FromString(context_bytes))
 
 
-def configsrecord_to_bytes(configs_record: ConfigsRecord) -> bytes:
-    """Serialize a `ConfigsRecord` to bytes."""
-    return serde.configs_record_to_proto(configs_record).SerializeToString()
+def configrecord_to_bytes(config_record: ConfigRecord) -> bytes:
+    """Serialize a `ConfigRecord` to bytes."""
+    return serde.config_record_to_proto(config_record).SerializeToString()
 
 
-def configsrecord_from_bytes(configsrecord_bytes: bytes) -> ConfigsRecord:
-    """Deserialize `ConfigsRecord` from bytes."""
-    return serde.configs_record_from_proto(
-        ProtoConfigsRecord.FromString(configsrecord_bytes)
+def configrecord_from_bytes(configrecord_bytes: bytes) -> ConfigRecord:
+    """Deserialize `ConfigRecord` from bytes."""
+    return serde.config_record_from_proto(
+        ProtoConfigRecord.FromString(configrecord_bytes)
     )
 
 
@@ -235,165 +236,202 @@ def has_valid_sub_status(status: RunStatus) -> bool:
     return status.sub_status == ""
 
 
-def create_taskres_for_unavailable_taskins(taskins_id: Union[str, UUID]) -> TaskRes:
-    """Generate a TaskRes with a TaskIns unavailable error.
-
-    Parameters
-    ----------
-    taskins_id : Union[str, UUID]
-        The ID of the unavailable TaskIns.
-
-    Returns
-    -------
-    TaskRes
-        A TaskRes with an error code MESSAGE_UNAVAILABLE to indicate that the
-        inquired TaskIns ID cannot be found (due to non-existence or expiration).
-    """
+def create_message_error_unavailable_res_message(
+    ins_metadata: Metadata, error_type: str
+) -> Message:
+    """Generate an error Message that the SuperLink returns carrying the specified
+    error."""
     current_time = now().timestamp()
-    return TaskRes(
-        task_id=str(uuid4()),
-        group_id="",  # Unknown group ID
-        run_id=0,  # Unknown run ID
-        task=Task(
-            # This function is only called by SuperLink, and thus it's the producer.
-            producer=Node(node_id=SUPERLINK_NODE_ID),
-            consumer=Node(node_id=SUPERLINK_NODE_ID),
-            created_at=current_time,
-            ttl=0,
-            ancestry=[str(taskins_id)],
-            task_type="",  # Unknown message type
-            error=Error(
-                code=ErrorCode.MESSAGE_UNAVAILABLE,
-                reason=MESSAGE_UNAVAILABLE_ERROR_REASON,
+    ttl = max(ins_metadata.ttl - (current_time - ins_metadata.created_at), 0)
+    metadata = Metadata(
+        run_id=ins_metadata.run_id,
+        message_id=str(uuid4()),
+        src_node_id=SUPERLINK_NODE_ID,
+        dst_node_id=SUPERLINK_NODE_ID,
+        reply_to_message_id=ins_metadata.message_id,
+        group_id=ins_metadata.group_id,
+        message_type=ins_metadata.message_type,
+        created_at=current_time,
+        ttl=ttl,
+    )
+
+    return make_message(
+        metadata=metadata,
+        error=Error(
+            code=(
+                ErrorCode.REPLY_MESSAGE_UNAVAILABLE
+                if error_type == "msg_unavail"
+                else ErrorCode.NODE_UNAVAILABLE
+            ),
+            reason=(
+                REPLY_MESSAGE_UNAVAILABLE_ERROR_REASON
+                if error_type == "msg_unavail"
+                else NODE_UNAVAILABLE_ERROR_REASON
             ),
         ),
     )
 
 
-def create_taskres_for_unavailable_taskres(ref_taskins: TaskIns) -> TaskRes:
-    """Generate a TaskRes with a reply message unavailable error from a TaskIns.
-
-    Parameters
-    ----------
-    ref_taskins : TaskIns
-        The reference TaskIns object.
+def create_message_error_unavailable_ins_message(reply_to_message_id: UUID) -> Message:
+    """Error to indicate that the enquired Message had expired before reply arrived or
+    that it isn't found."""
+    metadata = Metadata(
+        run_id=0,  # Unknown
+        message_id=str(uuid4()),
+        src_node_id=SUPERLINK_NODE_ID,
+        dst_node_id=SUPERLINK_NODE_ID,
+        reply_to_message_id=str(reply_to_message_id),
+        group_id="",  # Unknown
+        message_type=MessageType.SYSTEM,
+        created_at=now().timestamp(),
+        ttl=0,
+    )
 
-    Returns
-    -------
-    TaskRes
-        The generated TaskRes with an error code REPLY_MESSAGE_UNAVAILABLE_ERROR_REASON,
-        indicating that the original TaskRes has expired.
-    """
-    current_time = now().timestamp()
-    ttl = ref_taskins.task.ttl - (current_time - ref_taskins.task.created_at)
-    if ttl < 0:
-        log(ERROR, "Creating TaskRes for TaskIns that exceeds its TTL.")
-        ttl = 0
-    return TaskRes(
-        task_id=str(uuid4()),
-        group_id=ref_taskins.group_id,
-        run_id=ref_taskins.run_id,
-        task=Task(
-            # This function is only called by SuperLink, and thus it's the producer.
-            producer=Node(node_id=SUPERLINK_NODE_ID),
-            consumer=Node(node_id=SUPERLINK_NODE_ID),
-            created_at=current_time,
-            ttl=ttl,
-            ancestry=[ref_taskins.task_id],
-            task_type=ref_taskins.task.task_type,
-            error=Error(
-                code=ErrorCode.REPLY_MESSAGE_UNAVAILABLE,
-                reason=REPLY_MESSAGE_UNAVAILABLE_ERROR_REASON,
-            ),
+    return make_message(
+        metadata=metadata,
+        error=Error(
+            code=ErrorCode.MESSAGE_UNAVAILABLE,
+            reason=MESSAGE_UNAVAILABLE_ERROR_REASON,
         ),
     )
 
 
-def has_expired(task_ins_or_res: Union[TaskIns, TaskRes], current_time: float) -> bool:
-    """Check if the TaskIns/TaskRes has expired."""
-    return task_ins_or_res.task.ttl + task_ins_or_res.task.created_at < current_time
+def message_ttl_has_expired(message_metadata: Metadata, current_time: float) -> bool:
+    """Check if the Message has expired."""
+    return message_metadata.ttl + message_metadata.created_at < current_time
 
 
-def verify_taskins_ids(
-    inquired_taskins_ids: set[UUID],
-    found_taskins_dict: dict[UUID, TaskIns],
+def verify_message_ids(
+    inquired_message_ids: set[UUID],
+    found_message_ins_dict: dict[UUID, Message],
     current_time: Optional[float] = None,
     update_set: bool = True,
-) -> dict[UUID, TaskRes]:
-    """Verify found TaskIns and generate error TaskRes for invalid ones.
+) -> dict[UUID, Message]:
+    """Verify found Messages and generate error Messages for invalid ones.
 
     Parameters
     ----------
-    inquired_taskins_ids : set[UUID]
-        Set of TaskIns IDs for which to generate error TaskRes if invalid.
-    found_taskins_dict : dict[UUID, TaskIns]
-        Dictionary containing all found TaskIns indexed by their IDs.
+    inquired_message_ids : set[UUID]
+        Set of Message IDs for which to generate error Message if invalid.
+    found_message_ins_dict : dict[UUID, Message]
+        Dictionary containing all found Message indexed by their IDs.
     current_time : Optional[float] (default: None)
         The current time to check for expiration. If set to `None`, the current time
         will automatically be set to the current timestamp using `now().timestamp()`.
     update_set : bool (default: True)
-        If True, the `inquired_taskins_ids` will be updated to remove invalid ones,
+        If True, the `inquired_message_ids` will be updated to remove invalid ones,
         by default True.
 
     Returns
     -------
-    dict[UUID, TaskRes]
-        A dictionary of error TaskRes indexed by the corresponding TaskIns ID.
+    dict[UUID, Message]
+        A dictionary of error Message indexed by the corresponding ID of the message
+        they are a reply of.
     """
     ret_dict = {}
     current = current_time if current_time else now().timestamp()
-    for taskins_id in list(inquired_taskins_ids):
-        # Generate error TaskRes if the task_ins doesn't exist or has expired
-        taskins = found_taskins_dict.get(taskins_id)
-        if taskins is None or has_expired(taskins, current):
+    for message_id in list(inquired_message_ids):
+        # Generate error message if the inquired message doesn't exist or has expired
+        message_ins = found_message_ins_dict.get(message_id)
+        if message_ins is None or message_ttl_has_expired(
+            message_ins.metadata, current
+        ):
             if update_set:
-                inquired_taskins_ids.remove(taskins_id)
-            taskres = create_taskres_for_unavailable_taskins(taskins_id)
-            ret_dict[taskins_id] = taskres
+                inquired_message_ids.remove(message_id)
+            message_res = create_message_error_unavailable_ins_message(message_id)
+            ret_dict[message_id] = message_res
     return ret_dict
 
 
-def verify_found_taskres(
-    inquired_taskins_ids: set[UUID],
-    found_taskins_dict: dict[UUID, TaskIns],
-    found_taskres_list: list[TaskRes],
+def verify_found_message_replies(
+    inquired_message_ids: set[UUID],
+    found_message_ins_dict: dict[UUID, Message],
+    found_message_res_list: list[Message],
     current_time: Optional[float] = None,
     update_set: bool = True,
-) -> dict[UUID, TaskRes]:
-    """Verify found TaskRes and generate error TaskRes for invalid ones.
+) -> dict[UUID, Message]:
+    """Verify found Message replies and generate error Message for invalid ones.
 
     Parameters
     ----------
-    inquired_taskins_ids : set[UUID]
-        Set of TaskIns IDs for which to generate error TaskRes if invalid.
-    found_taskins_dict : dict[UUID, TaskIns]
-        Dictionary containing all found TaskIns indexed by their IDs.
-    found_taskres_list : dict[TaskIns, TaskRes]
-        List of found TaskRes to be verified.
+    inquired_message_ids : set[UUID]
+        Set of Message IDs for which to generate error Message if invalid.
+    found_message_ins_dict : dict[UUID, Message]
+        Dictionary containing all found instruction Messages indexed by their IDs.
+    found_message_res_list : dict[Message, Message]
+        List of found Message to be verified.
     current_time : Optional[float] (default: None)
         The current time to check for expiration. If set to `None`, the current time
         will automatically be set to the current timestamp using `now().timestamp()`.
     update_set : bool (default: True)
-        If True, the `inquired_taskins_ids` will be updated to remove ones
-        that have a TaskRes, by default True.
+        If True, the `inquired_message_ids` will be updated to remove ones
+        that have a reply Message, by default True.
 
     Returns
     -------
-    dict[UUID, TaskRes]
-        A dictionary of TaskRes indexed by the corresponding TaskIns ID.
+    dict[UUID, Message]
+        A dictionary of Message indexed by the corresponding Message ID.
     """
-    ret_dict: dict[UUID, TaskRes] = {}
+    ret_dict: dict[UUID, Message] = {}
     current = current_time if current_time else now().timestamp()
-    for taskres in found_taskres_list:
-        taskins_id = UUID(taskres.task.ancestry[0])
+    for message_res in found_message_res_list:
+        message_ins_id = UUID(message_res.metadata.reply_to_message_id)
         if update_set:
-            inquired_taskins_ids.remove(taskins_id)
-        # Check if the TaskRes has expired
-        if has_expired(taskres, current):
-            # No need to insert the error TaskRes
-            taskres = create_taskres_for_unavailable_taskres(
-                found_taskins_dict[taskins_id]
+            inquired_message_ids.remove(message_ins_id)
+        # Check if the reply Message has expired
+        if message_ttl_has_expired(message_res.metadata, current):
+            # No need to insert the error Message
+            message_res = create_message_error_unavailable_res_message(
+                found_message_ins_dict[message_ins_id].metadata, "msg_unavail"
+            )
+        ret_dict[message_ins_id] = message_res
+    return ret_dict
+
+
+def check_node_availability_for_in_message(
+    inquired_in_message_ids: set[UUID],
+    found_in_message_dict: dict[UUID, Message],
+    node_id_to_online_until: dict[int, float],
+    current_time: Optional[float] = None,
+    update_set: bool = True,
+) -> dict[UUID, Message]:
+    """Check node availability for given Message and generate error reply Message if
+    unavailable. A Message error indicating node unavailability will be generated for
+    each given Message whose destination node is offline or non-existent.
+
+    Parameters
+    ----------
+    inquired_in_message_ids : set[UUID]
+        Set of Message IDs for which to check destination node availability.
+    found_in_message_dict : dict[UUID, Message]
+        Dictionary containing all found Message indexed by their IDs.
+    node_id_to_online_until : dict[int, float]
+        Dictionary mapping node IDs to their online-until timestamps.
+    current_time : Optional[float] (default: None)
+        The current time to check for expiration. If set to `None`, the current time
+        will automatically be set to the current timestamp using `now().timestamp()`.
+    update_set : bool (default: True)
+        If True, the `inquired_in_message_ids` will be updated to remove invalid ones,
+        by default True.
+
+    Returns
+    -------
+    dict[UUID, Message]
+        A dictionary of error Message indexed by the corresponding Message ID.
+    """
+    ret_dict = {}
+    current = current_time if current_time else now().timestamp()
+    for in_message_id in list(inquired_in_message_ids):
+        in_message = found_in_message_dict[in_message_id]
+        node_id = in_message.metadata.dst_node_id
+        online_until = node_id_to_online_until.get(node_id)
+        # Generate a reply message containing an error reply
+        # if the node is offline or doesn't exist.
+        if online_until is None or online_until < current:
+            if update_set:
+                inquired_in_message_ids.remove(in_message_id)
+            reply_message = create_message_error_unavailable_res_message(
+                in_message.metadata, "node_unavail"
             )
-            taskres.task.delivered_at = now().isoformat()
-        ret_dict[taskins_id] = taskres
+            ret_dict[in_message_id] = reply_message
     return ret_dict

flwr/server/superlink/{driver → serverappio}/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2024 Flower Labs GmbH. All Rights Reserved.
+# Copyright 2025 Flower Labs GmbH. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

flwr/server/superlink/{driver → serverappio}/serverappio_grpc.py

@@ -1,4 +1,4 @@
-# Copyright 2024 Flower Labs GmbH. All Rights Reserved.
+# Copyright 2025 Flower Labs GmbH. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

flwr/server/superlink/{driver → serverappio}/serverappio_servicer.py

@@ -1,4 +1,4 @@
-# Copyright 2024 Flower Labs GmbH. All Rights Reserved.
+# Copyright 2025 Flower Labs GmbH. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -22,8 +22,8 @@ from uuid import UUID
 
 import grpc
 
-from flwr.common import ConfigsRecord
-from flwr.common.constant import Status
+from flwr.common import ConfigRecord, Message
+from flwr.common.constant import SUPERLINK_NODE_ID, Status
 from flwr.common.logger import log
 from flwr.common.serde import (
     context_from_proto,
@@ -31,9 +31,7 @@ from flwr.common.serde import (
     fab_from_proto,
     fab_to_proto,
     message_from_proto,
-    message_from_taskres,
     message_to_proto,
-    message_to_taskins,
     run_status_from_proto,
     run_status_to_proto,
     run_to_proto,
@@ -69,12 +67,11 @@ from flwr.proto.serverappio_pb2 import (  # pylint: disable=E0611
     PushServerAppOutputsRequest,
     PushServerAppOutputsResponse,
 )
-from flwr.proto.task_pb2 import TaskRes  # pylint: disable=E0611
 from flwr.server.superlink.ffs.ffs import Ffs
 from flwr.server.superlink.ffs.ffs_factory import FfsFactory
 from flwr.server.superlink.linkstate import LinkState, LinkStateFactory
 from flwr.server.superlink.utils import abort_if
-from flwr.server.utils.validator import validate_task_ins_or_res
+from flwr.server.utils.validator import validate_message
 
 
 class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
@@ -130,7 +127,7 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
             request.fab_version,
             fab_hash,
             user_config_from_proto(request.override_config),
-            ConfigsRecord(),
+            ConfigRecord(),
         )
         return CreateRunResponse(run_id=run_id)
 
@@ -161,20 +158,19 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
         while request.messages_list:
             message_proto = request.messages_list.pop(0)
             message = message_from_proto(message_proto=message_proto)
-            task_ins = message_to_taskins(message=message)
-            validation_errors = validate_task_ins_or_res(task_ins)
+            validation_errors = validate_message(message, is_reply_message=False)
             _raise_if(
                 validation_error=bool(validation_errors),
                 request_name="PushMessages",
                 detail=", ".join(validation_errors),
             )
             _raise_if(
-                validation_error=request.run_id != task_ins.run_id,
+                validation_error=request.run_id != message.metadata.run_id,
                 request_name="PushMessages",
-                detail="`task_ins` has mismatched `run_id`",
+                detail="`Message.metadata` has mismatched `run_id`",
             )
             # Store
-            message_id: Optional[UUID] = state.store_task_ins(task_ins=task_ins)
+            message_id: Optional[UUID] = state.store_message_ins(message=message)
             message_ids.append(message_id)
 
         return PushInsMessagesResponse(
@@ -200,32 +196,34 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
             context,
         )
 
-        # Convert each task_id str to UUID
+        # Convert each message_id str to UUID
         message_ids: set[UUID] = {
             UUID(message_id) for message_id in request.message_ids
         }
 
         # Read from state
-        task_res_list: list[TaskRes] = state.get_task_res(task_ids=message_ids)
+        messages_res: list[Message] = state.get_message_res(message_ids=message_ids)
 
-        # Delete the TaskIns/TaskRes pairs if TaskRes is found
-        task_ins_ids_to_delete = {
-            UUID(task_res.task.ancestry[0]) for task_res in task_res_list
+        # Delete the instruction Messages and their replies if found
+        message_ins_ids_to_delete = {
+            UUID(msg_res.metadata.reply_to_message_id) for msg_res in messages_res
         }
 
-        state.delete_tasks(task_ins_ids=task_ins_ids_to_delete)
+        state.delete_messages(message_ins_ids=message_ins_ids_to_delete)
 
-        # Convert to Messages
+        # Convert Messages to proto
         messages_list = []
-        while task_res_list:
-            task_res = task_res_list.pop(0)
-            _raise_if(
-                validation_error=request.run_id != task_res.run_id,
-                request_name="PullMessages",
-                detail="`task_res` has mismatched `run_id`",
-            )
-            message = message_from_taskres(taskres=task_res)
-            messages_list.append(message_to_proto(message))
+        while messages_res:
+            msg = messages_res.pop(0)
+
+            # Skip `run_id` check for SuperLink generated replies
+            if msg.metadata.src_node_id != SUPERLINK_NODE_ID:
+                _raise_if(
+                    validation_error=request.run_id != msg.metadata.run_id,
+                    request_name="PullMessages",
+                    detail="`message.metadata` has mismatched `run_id`",
+                )
+            messages_list.append(message_to_proto(msg))
 
         return PullResMessagesResponse(messages_list=messages_list)
 

flwr/server/superlink/simulation/simulationio_servicer.py

@@ -24,7 +24,7 @@ from grpc import ServicerContext
 from flwr.common.constant import Status
 from flwr.common.logger import log
 from flwr.common.serde import (
-    configs_record_to_proto,
+    config_record_to_proto,
     context_from_proto,
     context_to_proto,
     fab_to_proto,
@@ -182,5 +182,5 @@ class SimulationIoServicer(simulationio_pb2_grpc.SimulationIoServicer):
             )
             return GetFederationOptionsResponse()
         return GetFederationOptionsResponse(
-            federation_options=configs_record_to_proto(federation_options)
+            federation_options=config_record_to_proto(federation_options)
         )

flwr/server/typing.py

@@ -19,9 +19,9 @@ from typing import Callable
 
 from flwr.common import Context
 
-from .driver import Driver
+from .grid import Grid
 from .serverapp_components import ServerAppComponents
 
-ServerAppCallable = Callable[[Driver, Context], None]
-Workflow = Callable[[Driver, Context], None]
+ServerAppCallable = Callable[[Grid, Context], None]
+Workflow = Callable[[Grid, Context], None]
 ServerFn = Callable[[Context], ServerAppComponents]

flwr/server/utils/__init__.py

@@ -16,9 +16,9 @@
 
 
 from .tensorboard import tensorboard as tensorboard
-from .validator import validate_task_ins_or_res as validate_task_ins_or_res
+from .validator import validate_message as validate_message
 
 __all__ = [
     "tensorboard",
-    "validate_task_ins_or_res",
+    "validate_message",
 ]