flwr-nightly 1.8.0.dev20240315__py3-none-any.whl → 1.15.0.dev20250114__py3-none-any.whl

flwr/common/message.py

@@ -14,14 +14,20 @@
 # ==============================================================================
 """Message."""
 
+
 from __future__ import annotations
 
-from dataclasses import dataclass
+import time
+from logging import WARNING
+from typing import Optional, cast
 
+from .constant import MESSAGE_TTL_TOLERANCE
+from .logger import log
 from .record import RecordSet
 
+DEFAULT_TTL = 3600
+
 
-@dataclass
 class Metadata:  # pylint: disable=too-many-instance-attributes
     """A dataclass holding metadata associated with the current message.
 
@@ -40,28 +46,14 @@ class Metadata:  # pylint: disable=too-many-instance-attributes
     group_id : str
         An identifier for grouping messages. In some settings,
         this is used as the FL round.
-    ttl : str
-        Time-to-live for this message.
+    ttl : float
+        Time-to-live for this message in seconds.
     message_type : str
         A string that encodes the action to be executed on
         the receiving end.
-    partition_id : Optional[int]
-        An identifier that can be used when loading a particular
-        data partition for a ClientApp. Making use of this identifier
-        is more relevant when conducting simulations.
     """
 
-    _run_id: int
-    _message_id: str
-    _src_node_id: int
-    _dst_node_id: int
-    _reply_to_message: str
-    _group_id: str
-    _ttl: str
-    _message_type: str
-    _partition_id: int | None
-
-    def __init__(  # pylint: disable=too-many-arguments
+    def __init__(  # pylint: disable=too-many-arguments,too-many-positional-arguments
         self,
         run_id: int,
         message_id: str,
@@ -69,92 +61,103 @@ class Metadata:  # pylint: disable=too-many-instance-attributes
         dst_node_id: int,
         reply_to_message: str,
         group_id: str,
-        ttl: str,
+        ttl: float,
         message_type: str,
-        partition_id: int | None = None,
     ) -> None:
-        self._run_id = run_id
-        self._message_id = message_id
-        self._src_node_id = src_node_id
-        self._dst_node_id = dst_node_id
-        self._reply_to_message = reply_to_message
-        self._group_id = group_id
-        self._ttl = ttl
-        self._message_type = message_type
-        self._partition_id = partition_id
+        var_dict = {
+            "_run_id": run_id,
+            "_message_id": message_id,
+            "_src_node_id": src_node_id,
+            "_dst_node_id": dst_node_id,
+            "_reply_to_message": reply_to_message,
+            "_group_id": group_id,
+            "_ttl": ttl,
+            "_message_type": message_type,
+        }
+        self.__dict__.update(var_dict)
 
     @property
     def run_id(self) -> int:
         """An identifier for the current run."""
-        return self._run_id
+        return cast(int, self.__dict__["_run_id"])
 
     @property
     def message_id(self) -> str:
         """An identifier for the current message."""
-        return self._message_id
+        return cast(str, self.__dict__["_message_id"])
 
     @property
     def src_node_id(self) -> int:
         """An identifier for the node sending this message."""
-        return self._src_node_id
+        return cast(int, self.__dict__["_src_node_id"])
 
     @property
     def reply_to_message(self) -> str:
         """An identifier for the message this message replies to."""
-        return self._reply_to_message
+        return cast(str, self.__dict__["_reply_to_message"])
 
     @property
     def dst_node_id(self) -> int:
         """An identifier for the node receiving this message."""
-        return self._dst_node_id
+        return cast(int, self.__dict__["_dst_node_id"])
 
     @dst_node_id.setter
     def dst_node_id(self, value: int) -> None:
         """Set dst_node_id."""
-        self._dst_node_id = value
+        self.__dict__["_dst_node_id"] = value
 
     @property
     def group_id(self) -> str:
         """An identifier for grouping messages."""
-        return self._group_id
+        return cast(str, self.__dict__["_group_id"])
 
     @group_id.setter
     def group_id(self, value: str) -> None:
         """Set group_id."""
-        self._group_id = value
+        self.__dict__["_group_id"] = value
+
+    @property
+    def created_at(self) -> float:
+        """Unix timestamp when the message was created."""
+        return cast(float, self.__dict__["_created_at"])
+
+    @created_at.setter
+    def created_at(self, value: float) -> None:
+        """Set creation timestamp for this message."""
+        self.__dict__["_created_at"] = value
 
     @property
-    def ttl(self) -> str:
+    def ttl(self) -> float:
         """Time-to-live for this message."""
-        return self._ttl
+        return cast(float, self.__dict__["_ttl"])
 
     @ttl.setter
-    def ttl(self, value: str) -> None:
+    def ttl(self, value: float) -> None:
         """Set ttl."""
-        self._ttl = value
+        self.__dict__["_ttl"] = value
 
     @property
     def message_type(self) -> str:
         """A string that encodes the action to be executed on the receiving end."""
-        return self._message_type
+        return cast(str, self.__dict__["_message_type"])
 
     @message_type.setter
     def message_type(self, value: str) -> None:
         """Set message_type."""
-        self._message_type = value
+        self.__dict__["_message_type"] = value
 
-    @property
-    def partition_id(self) -> int | None:
-        """An identifier telling which data partition a ClientApp should use."""
-        return self._partition_id
+    def __repr__(self) -> str:
+        """Return a string representation of this instance."""
+        view = ", ".join([f"{k.lstrip('_')}={v!r}" for k, v in self.__dict__.items()])
+        return f"{self.__class__.__qualname__}({view})"
 
-    @partition_id.setter
-    def partition_id(self, value: int) -> None:
-        """Set patition_id."""
-        self._partition_id = value
+    def __eq__(self, other: object) -> bool:
+        """Compare two instances of the class."""
+        if not isinstance(other, self.__class__):
+            raise NotImplementedError
+        return self.__dict__ == other.__dict__
 
 
-@dataclass
 class Error:
     """A dataclass that stores information about an error that occurred.
 
@@ -166,25 +169,35 @@ class Error:
         A reason for why the error arose (e.g. an exception stack-trace)
     """
 
-    _code: int
-    _reason: str | None = None
-
     def __init__(self, code: int, reason: str | None = None) -> None:
-        self._code = code
-        self._reason = reason
+        var_dict = {
+            "_code": code,
+            "_reason": reason,
+        }
+        self.__dict__.update(var_dict)
 
     @property
     def code(self) -> int:
         """Error code."""
-        return self._code
+        return cast(int, self.__dict__["_code"])
 
     @property
     def reason(self) -> str | None:
         """Reason reported about the error."""
-        return self._reason
+        return cast(Optional[str], self.__dict__["_reason"])
+
+    def __repr__(self) -> str:
+        """Return a string representation of this instance."""
+        view = ", ".join([f"{k.lstrip('_')}={v!r}" for k, v in self.__dict__.items()])
+        return f"{self.__class__.__qualname__}({view})"
+
+    def __eq__(self, other: object) -> bool:
+        """Compare two instances of the class."""
+        if not isinstance(other, self.__class__):
+            raise NotImplementedError
+        return self.__dict__ == other.__dict__
 
 
-@dataclass
 class Message:
     """State of your application from the viewpoint of the entity using it.
 
@@ -200,105 +213,108 @@ class Message:
         when processing another message.
     """
 
-    _metadata: Metadata
-    _content: RecordSet | None = None
-    _error: Error | None = None
-
     def __init__(
         self,
         metadata: Metadata,
         content: RecordSet | None = None,
         error: Error | None = None,
     ) -> None:
-        self._metadata = metadata
-
         if not (content is None) ^ (error is None):
             raise ValueError("Either `content` or `error` must be set, but not both.")
 
-        self._content = content
-        self._error = error
+        metadata.created_at = time.time()  # Set the message creation timestamp
+        var_dict = {
+            "_metadata": metadata,
+            "_content": content,
+            "_error": error,
+        }
+        self.__dict__.update(var_dict)
 
     @property
     def metadata(self) -> Metadata:
         """A dataclass including information about the message to be executed."""
-        return self._metadata
+        return cast(Metadata, self.__dict__["_metadata"])
 
     @property
     def content(self) -> RecordSet:
         """The content of this message."""
-        if self._content is None:
+        if self.__dict__["_content"] is None:
             raise ValueError(
                 "Message content is None. Use <message>.has_content() "
                 "to check if a message has content."
             )
-        return self._content
+        return cast(RecordSet, self.__dict__["_content"])
 
     @content.setter
     def content(self, value: RecordSet) -> None:
         """Set content."""
-        if self._error is None:
-            self._content = value
+        if self.__dict__["_error"] is None:
+            self.__dict__["_content"] = value
         else:
             raise ValueError("A message with an error set cannot have content.")
 
     @property
     def error(self) -> Error:
         """Error captured by this message."""
-        if self._error is None:
+        if self.__dict__["_error"] is None:
             raise ValueError(
                 "Message error is None. Use <message>.has_error() "
                 "to check first if a message carries an error."
             )
-        return self._error
+        return cast(Error, self.__dict__["_error"])
 
     @error.setter
     def error(self, value: Error) -> None:
         """Set error."""
         if self.has_content():
             raise ValueError("A message with content set cannot carry an error.")
-        self._error = value
+        self.__dict__["_error"] = value
 
     def has_content(self) -> bool:
         """Return True if message has content, else False."""
-        return self._content is not None
+        return self.__dict__["_content"] is not None
 
     def has_error(self) -> bool:
         """Return True if message has an error, else False."""
-        return self._error is not None
-
-    def _create_reply_metadata(self, ttl: str) -> Metadata:
-        """Construct metadata for a reply message."""
-        return Metadata(
-            run_id=self.metadata.run_id,
-            message_id="",
-            src_node_id=self.metadata.dst_node_id,
-            dst_node_id=self.metadata.src_node_id,
-            reply_to_message=self.metadata.message_id,
-            group_id=self.metadata.group_id,
-            ttl=ttl,
-            message_type=self.metadata.message_type,
-            partition_id=self.metadata.partition_id,
-        )
+        return self.__dict__["_error"] is not None
 
-    def create_error_reply(
-        self,
-        error: Error,
-        ttl: str,
-    ) -> Message:
+    def create_error_reply(self, error: Error, ttl: float | None = None) -> Message:
         """Construct a reply message indicating an error happened.
 
         Parameters
         ----------
         error : Error
             The error that was encountered.
-        ttl : str
-            Time-to-live for this message.
+        ttl : Optional[float] (default: None)
+            Time-to-live for this message in seconds. If unset, it will be set based
+            on the remaining time for the received message before it expires. This
+            follows the equation:
+
+            ttl = msg.meta.ttl - (reply.meta.created_at - msg.meta.created_at)
+
+        Returns
+        -------
+        message : Message
+            A Message containing only the relevant error and metadata.
         """
+        # If no TTL passed, use default for message creation (will update after
+        # message creation)
+        ttl_ = DEFAULT_TTL if ttl is None else ttl
         # Create reply with error
-        message = Message(metadata=self._create_reply_metadata(ttl), error=error)
+        message = Message(metadata=_create_reply_metadata(self, ttl_), error=error)
+
+        if ttl is None:
+            # Set TTL equal to the remaining time for the received message to expire
+            ttl = self.metadata.ttl - (
+                message.metadata.created_at - self.metadata.created_at
+            )
+            message.metadata.ttl = ttl
+
+        self._limit_task_res_ttl(message)
+
         return message
 
-    def create_reply(self, content: RecordSet, ttl: str) -> Message:
+    def create_reply(self, content: RecordSet, ttl: float | None = None) -> Message:
         """Create a reply to this message with specified content and TTL.
 
         The method generates a new `Message` as a reply to this message.
@@ -309,15 +325,84 @@ class Message:
         ----------
         content : RecordSet
             The content for the reply message.
-        ttl : str
-            Time-to-live for this message.
+        ttl : Optional[float] (default: None)
+            Time-to-live for this message in seconds. If unset, it will be set based
+            on the remaining time for the received message before it expires. This
+            follows the equation:
+
+            ttl = msg.meta.ttl - (reply.meta.created_at - msg.meta.created_at)
 
         Returns
         -------
         Message
             A new `Message` instance representing the reply.
         """
-        return Message(
-            metadata=self._create_reply_metadata(ttl),
+        # If no TTL passed, use default for message creation (will update after
+        # message creation)
+        ttl_ = DEFAULT_TTL if ttl is None else ttl
+
+        message = Message(
+            metadata=_create_reply_metadata(self, ttl_),
             content=content,
         )
+
+        if ttl is None:
+            # Set TTL equal to the remaining time for the received message to expire
+            ttl = self.metadata.ttl - (
+                message.metadata.created_at - self.metadata.created_at
+            )
+            message.metadata.ttl = ttl
+
+        self._limit_task_res_ttl(message)
+
+        return message
+
+    def __repr__(self) -> str:
+        """Return a string representation of this instance."""
+        view = ", ".join(
+            [
+                f"{k.lstrip('_')}={v!r}"
+                for k, v in self.__dict__.items()
+                if v is not None
+            ]
+        )
+        return f"{self.__class__.__qualname__}({view})"
+
+    def _limit_task_res_ttl(self, message: Message) -> None:
+        """Limit the TaskRes TTL to not exceed the expiration time of the TaskIns it
+        replies to.
+
+        Parameters
+        ----------
+        message : Message
+            The message to which the TaskRes is replying.
+        """
+        # Calculate the maximum allowed TTL
+        max_allowed_ttl = (
+            self.metadata.created_at + self.metadata.ttl - message.metadata.created_at
+        )
+
+        if message.metadata.ttl - max_allowed_ttl > MESSAGE_TTL_TOLERANCE:
+            log(
+                WARNING,
+                "The reply TTL of %.2f seconds exceeded the "
+                "allowed maximum of %.2f seconds. "
+                "The TTL has been updated to the allowed maximum.",
+                message.metadata.ttl,
+                max_allowed_ttl,
+            )
+            message.metadata.ttl = max_allowed_ttl
+
+
+def _create_reply_metadata(msg: Message, ttl: float) -> Metadata:
+    """Construct metadata for a reply message."""
+    return Metadata(
+        run_id=msg.metadata.run_id,
+        message_id="",
+        src_node_id=msg.metadata.dst_node_id,
+        dst_node_id=msg.metadata.src_node_id,
+        reply_to_message=msg.metadata.message_id,
+        group_id=msg.metadata.group_id,
+        ttl=ttl,
+        message_type=msg.metadata.message_type,
+    )