flwr-nightly 1.17.0.dev20250318__py3-none-any.whl → 1.17.0.dev20250320__py3-none-any.whl

flwr/server/grid/inmemory_grid.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Flower in-memory Driver."""
+"""Flower in-memory Grid."""
 
 
 import time
@@ -20,22 +20,23 @@ from collections.abc import Iterable
 from typing import Optional, cast
 from uuid import UUID
 
-from flwr.common import DEFAULT_TTL, Message, Metadata, RecordSet
+from flwr.common import Message, RecordDict
 from flwr.common.constant import SUPERLINK_NODE_ID
+from flwr.common.logger import warn_deprecated_feature
 from flwr.common.typing import Run
 from flwr.proto.node_pb2 import Node  # pylint: disable=E0611
 from flwr.server.superlink.linkstate import LinkStateFactory
 
-from .grid import Driver
+from .grid import Grid
 
 
-class InMemoryDriver(Driver):
-    """`InMemoryDriver` class provides an interface to the ServerAppIo API.
+class InMemoryGrid(Grid):
+    """`InMemoryGrid` class provides an interface to the ServerAppIo API.
 
     Parameters
     ----------
     state_factory : StateFactory
-        A StateFactory embedding a state that this driver can interface with.
+        A StateFactory embedding a state that this grid can interface with.
     pull_interval : float (default=0.1)
         Sleep duration between calls to `pull_messages`.
     """
@@ -53,10 +54,8 @@ class InMemoryDriver(Driver):
     def _check_message(self, message: Message) -> None:
         # Check if the message is valid
         if not (
-            message.metadata.run_id == cast(Run, self._run).run_id
-            and message.metadata.src_node_id == self.node.node_id
-            and message.metadata.message_id == ""
-            and message.metadata.reply_to_message == ""
+            message.metadata.message_id == ""
+            and message.metadata.reply_to_message_id == ""
             and message.metadata.ttl > 0
             and message.metadata.delivered_at == ""
         ):
@@ -76,7 +75,7 @@ class InMemoryDriver(Driver):
 
     def create_message(  # pylint: disable=too-many-arguments,R0917
         self,
-        content: RecordSet,
+        content: RecordDict,
         message_type: str,
         dst_node_id: int,
         group_id: str,
@@ -87,19 +86,11 @@ class InMemoryDriver(Driver):
         This method constructs a new `Message` with given content and metadata.
         The `run_id` and `src_node_id` will be set automatically.
         """
-        ttl_ = DEFAULT_TTL if ttl is None else ttl
-
-        metadata = Metadata(
-            run_id=cast(Run, self._run).run_id,
-            message_id="",  # Will be set by the server
-            src_node_id=self.node.node_id,
-            dst_node_id=dst_node_id,
-            reply_to_message="",
-            group_id=group_id,
-            ttl=ttl_,
-            message_type=message_type,
+        warn_deprecated_feature(
+            "`Driver.create_message` / `Grid.create_message` is deprecated."
+            "Use `Message` constructor instead."
         )
-        return Message(metadata=metadata, content=content)
+        return Message(content, dst_node_id, message_type, ttl=ttl, group_id=group_id)
 
     def get_node_ids(self) -> Iterable[int]:
         """Get node IDs."""
@@ -113,6 +104,9 @@ class InMemoryDriver(Driver):
         """
         msg_ids: list[str] = []
         for msg in messages:
+            # Populate metadata
+            msg.metadata.__dict__["_run_id"] = cast(Run, self._run).run_id
+            msg.metadata.__dict__["_src_node_id"] = self.node.node_id
             # Check message
             self._check_message(msg)
             # Store in state
@@ -133,7 +127,7 @@ class InMemoryDriver(Driver):
         message_res_list = self.state.get_message_res(message_ids=msg_ids)
         # Get IDs of Messages these replies are for
         message_ins_ids_to_delete = {
-            UUID(msg_res.metadata.reply_to_message) for msg_res in message_res_list
+            UUID(msg_res.metadata.reply_to_message_id) for msg_res in message_res_list
         }
         # Delete
         self.state.delete_messages(message_ins_ids=message_ins_ids_to_delete)
@@ -162,7 +156,7 @@ class InMemoryDriver(Driver):
             res_msgs = self.pull_messages(msg_ids)
             ret.extend(res_msgs)
             msg_ids.difference_update(
-                {msg.metadata.reply_to_message for msg in res_msgs}
+                {msg.metadata.reply_to_message_id for msg in res_msgs}
             )
             if len(msg_ids) == 0:
                 break

flwr/server/run_serverapp.py

@@ -22,18 +22,18 @@ from flwr.common import Context
 from flwr.common.logger import log
 from flwr.common.object_ref import load_app
 
-from .grid import Driver
+from .grid import Grid
 from .server_app import LoadServerAppError, ServerApp
 
 
 def run(
-    driver: Driver,
+    grid: Grid,
     context: Context,
     server_app_dir: str,
     server_app_attr: Optional[str] = None,
     loaded_server_app: Optional[ServerApp] = None,
 ) -> Context:
-    """Run ServerApp with a given Driver."""
+    """Run ServerApp with a given Grid."""
     if not (server_app_attr is None) ^ (loaded_server_app is None):
         raise ValueError(
             "Either `server_app_attr` or `loaded_server_app` should be set "
@@ -59,7 +59,7 @@ def run(
     server_app = _load()
 
     # Call ServerApp
-    server_app(driver=driver, context=context)
+    server_app(grid=grid, context=context)
 
     log(DEBUG, "ServerApp finished running.")
     return context

flwr/server/server_app.py

@@ -15,6 +15,7 @@
 """Flower ServerApp."""
 
 
+import inspect
 from collections.abc import Iterator
 from contextlib import contextmanager
 from typing import Callable, Optional
@@ -24,8 +25,8 @@ from flwr.common.logger import warn_deprecated_feature_with_example
 from flwr.server.strategy import Strategy
 
 from .client_manager import ClientManager
-from .compat import start_driver
-from .grid import Driver
+from .compat import start_grid
+from .grid import Driver, Grid
 from .server import Server
 from .server_config import ServerConfig
 from .typing import ServerAppCallable, ServerFn
@@ -43,6 +44,21 @@ SERVER_FN_USAGE_EXAMPLE = """
         app = ServerApp(server_fn=server_fn)
 """
 
+GRID_USAGE_EXAMPLE = """
+                app = ServerApp()
+
+                @app.main()
+                def main(grid: Grid, context: Context) -> None:
+                    # Your existing ServerApp code ...
+"""
+
+DRIVER_DEPRECATION_MSG = """
+            The `Driver` class is deprecated, it will be removed in a future release.
+"""
+DRIVER_EXAMPLE_MSG = """
+            Instead, use `Grid` in the signature of your `ServerApp`. For example:
+"""
+
 
 @contextmanager
 def _empty_lifespan(_: Context) -> Iterator[None]:
@@ -54,7 +70,7 @@ class ServerApp:  # pylint: disable=too-many-instance-attributes
 
     Examples
     --------
-    Use the `ServerApp` with an existing `Strategy`:
+    Use the ``ServerApp`` with an existing ``Strategy``:
 
     >>> def server_fn(context: Context):
     >>>     server_config = ServerConfig(num_rounds=3)
@@ -66,12 +82,12 @@ class ServerApp:  # pylint: disable=too-many-instance-attributes
     >>>
     >>> app = ServerApp(server_fn=server_fn)
 
-    Use the `ServerApp` with a custom main function:
+    Use the ``ServerApp`` with a custom main function:
 
     >>> app = ServerApp()
     >>>
     >>> @app.main()
-    >>> def main(driver: Driver, context: Context) -> None:
+    >>> def main(grid: Grid, context: Context) -> None:
     >>>    print("ServerApp running")
     """
 
@@ -111,7 +127,7 @@ class ServerApp:  # pylint: disable=too-many-instance-attributes
         self._main: Optional[ServerAppCallable] = None
         self._lifespan = _empty_lifespan
 
-    def __call__(self, driver: Driver, context: Context) -> None:
+    def __call__(self, grid: Grid, context: Context) -> None:
         """Execute `ServerApp`."""
         with self._lifespan(context):
             # Compatibility mode
@@ -123,17 +139,17 @@ class ServerApp:  # pylint: disable=too-many-instance-attributes
                     self._config = components.config
                     self._strategy = components.strategy
                     self._client_manager = components.client_manager
-                start_driver(
+                start_grid(
                     server=self._server,
                     config=self._config,
                     strategy=self._strategy,
                     client_manager=self._client_manager,
-                    driver=driver,
+                    grid=grid,
                 )
                 return
 
             # New execution mode
-            self._main(driver, context)
+            self._main(grid, context)
 
     def main(self) -> Callable[[ServerAppCallable], ServerAppCallable]:
         """Return a decorator that registers the main fn with the server app.
@@ -143,7 +159,7 @@ class ServerApp:  # pylint: disable=too-many-instance-attributes
         >>> app = ServerApp()
         >>>
         >>> @app.main()
-        >>> def main(driver: Driver, context: Context) -> None:
+        >>> def main(grid: Grid, context: Context) -> None:
         >>>    print("ServerApp running")
         """
 
@@ -168,11 +184,21 @@ class ServerApp:  # pylint: disable=too-many-instance-attributes
                     >>> app = ServerApp()
                     >>>
                     >>> @app.main()
-                    >>> def main(driver: Driver, context: Context) -> None:
+                    >>> def main(grid: Grid, context: Context) -> None:
                     >>>    print("ServerApp running")
                     """,
                 )
 
+            sig = inspect.signature(main_fn)
+            param = list(sig.parameters.values())[0]
+            # Check if parameter name or the annotation should be updated
+            if param.name == "driver" or param.annotation is Driver:
+                warn_deprecated_feature_with_example(
+                    deprecation_message=DRIVER_DEPRECATION_MSG,
+                    example_message=DRIVER_EXAMPLE_MSG,
+                    code_example=GRID_USAGE_EXAMPLE,
+                )
+
             # Register provided function with the ServerApp object
             self._main = main_fn
 

flwr/server/serverapp/app.py

@@ -60,7 +60,7 @@ from flwr.proto.serverappio_pb2 import (  # pylint: disable=E0611
     PullServerAppInputsResponse,
     PushServerAppOutputsRequest,
 )
-from flwr.server.grid.grpc_grid import GrpcDriver
+from flwr.server.grid.grpc_grid import GrpcGrid
 from flwr.server.run_serverapp import run as run_
 
 
@@ -106,7 +106,7 @@ def run_serverapp(  # pylint: disable=R0914, disable=W0212, disable=R0915
     certificates: Optional[bytes] = None,
 ) -> None:
     """Run Flower ServerApp process."""
-    driver = GrpcDriver(
+    grid = GrpcGrid(
         serverappio_service_address=serverappio_api_address,
         root_certificates=certificates,
     )
@@ -123,7 +123,7 @@ def run_serverapp(  # pylint: disable=R0914, disable=W0212, disable=R0915
             # Pull ServerAppInputs from LinkState
             req = PullServerAppInputsRequest()
             log(DEBUG, "[flwr-serverapp] Pull ServerAppInputs")
-            res: PullServerAppInputsResponse = driver._stub.PullServerAppInputs(req)
+            res: PullServerAppInputsResponse = grid._stub.PullServerAppInputs(req)
             if not res.HasField("run"):
                 sleep(3)
                 run_status = None
@@ -135,14 +135,14 @@ def run_serverapp(  # pylint: disable=R0914, disable=W0212, disable=R0915
 
             hash_run_id = get_sha256_hash(run.run_id)
 
-            driver.set_run(run.run_id)
+            grid.set_run(run.run_id)
 
             # Start log uploader for this run
             log_uploader = start_log_uploader(
                 log_queue=log_queue,
                 node_id=0,
                 run_id=run.run_id,
-                stub=driver._stub,
+                stub=grid._stub,
             )
 
             log(DEBUG, "[flwr-serverapp] Start FAB installation.")
@@ -173,7 +173,7 @@ def run_serverapp(  # pylint: disable=R0914, disable=W0212, disable=R0915
 
             # Change status to Running
             run_status_proto = run_status_to_proto(RunStatus(Status.RUNNING, "", ""))
-            driver._stub.UpdateRunStatus(
+            grid._stub.UpdateRunStatus(
                 UpdateRunStatusRequest(run_id=run.run_id, run_status=run_status_proto)
             )
 
@@ -182,9 +182,9 @@ def run_serverapp(  # pylint: disable=R0914, disable=W0212, disable=R0915
                 event_details={"run-id-hash": hash_run_id},
             )
 
-            # Load and run the ServerApp with the Driver
+            # Load and run the ServerApp with the Grid
             updated_context = run_(
-                driver=driver,
+                grid=grid,
                 server_app_dir=app_path,
                 server_app_attr=server_app_attr,
                 context=context,
@@ -196,7 +196,7 @@ def run_serverapp(  # pylint: disable=R0914, disable=W0212, disable=R0915
             out_req = PushServerAppOutputsRequest(
                 run_id=run.run_id, context=context_proto
             )
-            _ = driver._stub.PushServerAppOutputs(out_req)
+            _ = grid._stub.PushServerAppOutputs(out_req)
 
             run_status = RunStatus(Status.FINISHED, SubStatus.COMPLETED, "")
         except RunNotRunningException:
@@ -221,7 +221,7 @@ def run_serverapp(  # pylint: disable=R0914, disable=W0212, disable=R0915
             # Update run status
             if run_status:
                 run_status_proto = run_status_to_proto(run_status)
-                driver._stub.UpdateRunStatus(
+                grid._stub.UpdateRunStatus(
                     UpdateRunStatusRequest(
                         run_id=run.run_id, run_status=run_status_proto
                     )

flwr/server/superlink/fleet/vce/vce_api.py

@@ -130,9 +130,7 @@ def worker(
                 e_code = ErrorCode.UNKNOWN
 
             reason = str(type(ex)) + ":<'" + str(ex) + "'>"
-            out_mssg = message.create_error_reply(
-                error=Error(code=e_code, reason=reason)
-            )
+            out_mssg = Message(Error(code=e_code, reason=reason), reply_to=message)
 
         finally:
             if out_mssg:

flwr/server/superlink/linkstate/in_memory_linkstate.py

@@ -27,6 +27,7 @@ from flwr.common import Context, Message, log, now
 from flwr.common.constant import (
     MESSAGE_TTL_TOLERANCE,
     NODE_ID_NUM_BYTES,
+    PING_PATIENCE,
     RUN_ID_NUM_BYTES,
     SUPERLINK_NODE_ID,
     Status,
@@ -37,6 +38,7 @@ from flwr.server.superlink.linkstate.linkstate import LinkState
 from flwr.server.utils import validate_message
 
 from .utils import (
+    check_node_availability_for_in_message,
     generate_rand_int_from_bytes,
     has_valid_sub_status,
     is_valid_transition,
@@ -156,7 +158,7 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         res_metadata = message.metadata
         with self.lock:
             # Check if the Message it is replying to exists and is valid
-            msg_ins_id = res_metadata.reply_to_message
+            msg_ins_id = res_metadata.reply_to_message_id
             msg_ins = self.message_ins_store.get(UUID(msg_ins_id))
 
             # Ensure that dst_node_id of original Message matches the src_node_id of
@@ -232,13 +234,28 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         with self.lock:
             current = time.time()
 
-            # Verify Messge IDs
+            # Verify Message IDs
             ret = verify_message_ids(
                 inquired_message_ids=message_ids,
                 found_message_ins_dict=self.message_ins_store,
                 current_time=current,
             )
 
+            # Check node availability
+            dst_node_ids = {
+                self.message_ins_store[message_id].metadata.dst_node_id
+                for message_id in message_ids
+            }
+            tmp_ret_dict = check_node_availability_for_in_message(
+                inquired_in_message_ids=message_ids,
+                found_in_message_dict=self.message_ins_store,
+                node_id_to_online_until={
+                    node_id: self.node_ids[node_id][0] for node_id in dst_node_ids
+                },
+                current_time=current,
+            )
+            ret.update(tmp_ret_dict)
+
             # Find all reply Messages
             message_res_found: list[Message] = []
             for message_id in message_ids:
@@ -317,6 +334,7 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
                 log(ERROR, "Unexpected node registration failure.")
                 return 0
 
+            # Mark the node online util time.time() + ping_interval
             self.node_ids[node_id] = (time.time() + ping_interval, ping_interval)
             return node_id
 
@@ -519,10 +537,17 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
             return self.federation_options[run_id]
 
     def acknowledge_ping(self, node_id: int, ping_interval: float) -> bool:
-        """Acknowledge a ping received from a node, serving as a heartbeat."""
+        """Acknowledge a ping received from a node, serving as a heartbeat.
+
+        It allows for one missed ping (in a PING_PATIENCE * ping_interval) before
+        marking the node as offline, where PING_PATIENCE = 2 in default.
+        """
         with self.lock:
             if node_id in self.node_ids:
-                self.node_ids[node_id] = (time.time() + ping_interval, ping_interval)
+                self.node_ids[node_id] = (
+                    time.time() + PING_PATIENCE * ping_interval,
+                    ping_interval,
+                )
                 return True
         return False
 

flwr/server/superlink/linkstate/sqlite_linkstate.py

@@ -30,28 +30,31 @@ from flwr.common import Context, Message, Metadata, log, now
 from flwr.common.constant import (
     MESSAGE_TTL_TOLERANCE,
     NODE_ID_NUM_BYTES,
+    PING_PATIENCE,
     RUN_ID_NUM_BYTES,
     SUPERLINK_NODE_ID,
     Status,
 )
+from flwr.common.message import make_message
 from flwr.common.record import ConfigsRecord
 from flwr.common.serde import (
     error_from_proto,
     error_to_proto,
-    recordset_from_proto,
-    recordset_to_proto,
+    recorddict_from_proto,
+    recorddict_to_proto,
 )
 from flwr.common.typing import Run, RunStatus, UserConfig
 
 # pylint: disable=E0611
 from flwr.proto.error_pb2 import Error as ProtoError
-from flwr.proto.recordset_pb2 import RecordSet as ProtoRecordSet
+from flwr.proto.recorddict_pb2 import RecordDict as ProtoRecordDict
 
 # pylint: enable=E0611
 from flwr.server.utils.validator import validate_message
 
 from .linkstate import LinkState
 from .utils import (
+    check_node_availability_for_in_message,
     configsrecord_from_bytes,
     configsrecord_to_bytes,
     context_from_bytes,
@@ -129,7 +132,7 @@ CREATE TABLE IF NOT EXISTS message_ins(
     run_id                  INTEGER,
     src_node_id             INTEGER,
     dst_node_id             INTEGER,
-    reply_to_message        TEXT,
+    reply_to_message_id     TEXT,
     created_at              REAL,
     delivered_at            TEXT,
     ttl                     REAL,
@@ -148,7 +151,7 @@ CREATE TABLE IF NOT EXISTS message_res(
     run_id                  INTEGER,
     src_node_id             INTEGER,
     dst_node_id             INTEGER,
-    reply_to_message        TEXT,
+    reply_to_message_id     TEXT,
     created_at              REAL,
     delivered_at            TEXT,
     ttl                     REAL,
@@ -371,7 +374,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
             return None
 
         res_metadata = message.metadata
-        msg_ins_id = res_metadata.reply_to_message
+        msg_ins_id = res_metadata.reply_to_message_id
         msg_ins = self.get_valid_message_ins(msg_ins_id)
         if msg_ins is None:
             log(
@@ -442,6 +445,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
 
     def get_message_res(self, message_ids: set[UUID]) -> list[Message]:
         """Get reply Messages for the given Message IDs."""
+        # pylint: disable-msg=too-many-locals
         ret: dict[UUID, Message] = {}
 
         # Verify Message IDs
@@ -465,11 +469,34 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
             current_time=current,
         )
 
+        # Check node availability
+        dst_node_ids: set[int] = set()
+        for message_id in message_ids:
+            in_message = found_message_ins_dict[message_id]
+            sint_node_id = convert_uint64_to_sint64(in_message.metadata.dst_node_id)
+            dst_node_ids.add(sint_node_id)
+        query = f"""
+                    SELECT node_id, online_until
+                    FROM node
+                    WHERE node_id IN ({",".join(["?"] * len(dst_node_ids))});
+                """
+        rows = self.query(query, tuple(dst_node_ids))
+        tmp_ret_dict = check_node_availability_for_in_message(
+            inquired_in_message_ids=message_ids,
+            found_in_message_dict=found_message_ins_dict,
+            node_id_to_online_until={
+                convert_sint64_to_uint64(row["node_id"]): row["online_until"]
+                for row in rows
+            },
+            current_time=current,
+        )
+        ret.update(tmp_ret_dict)
+
         # Find all reply Messages
         query = f"""
             SELECT *
             FROM message_res
-            WHERE reply_to_message IN ({",".join(["?"] * len(message_ids))})
+            WHERE reply_to_message_id IN ({",".join(["?"] * len(message_ids))})
             AND delivered_at = "";
         """
         rows = self.query(query, tuple(str(message_id) for message_id in message_ids))
@@ -542,7 +569,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         # Delete reply Message
         query_2 = f"""
             DELETE FROM message_res
-            WHERE reply_to_message IN ({placeholders});
+            WHERE reply_to_message_id IN ({placeholders});
         """
 
         with self.conn:
@@ -584,6 +611,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
             "VALUES (?, ?, ?, ?)"
         )
 
+        # Mark the node online util time.time() + ping_interval
         try:
             self.query(
                 query,
@@ -899,7 +927,11 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         return configsrecord_from_bytes(row["federation_options"])
 
     def acknowledge_ping(self, node_id: int, ping_interval: float) -> bool:
-        """Acknowledge a ping received from a node, serving as a heartbeat."""
+        """Acknowledge a ping received from a node, serving as a heartbeat.
+
+        It allows for one missed ping (in a PING_PATIENCE * ping_interval) before
+        marking the node as offline, where PING_PATIENCE = 2 in default.
+        """
         sint64_node_id = convert_uint64_to_sint64(node_id)
 
         # Check if the node exists in the `node` table
@@ -909,7 +941,14 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
 
         # Update `online_until` and `ping_interval` for the given `node_id`
         query = "UPDATE node SET online_until = ?, ping_interval = ? WHERE node_id = ?"
-        self.query(query, (time.time() + ping_interval, ping_interval, sint64_node_id))
+        self.query(
+            query,
+            (
+                time.time() + PING_PATIENCE * ping_interval,
+                ping_interval,
+                sint64_node_id,
+            ),
+        )
         return True
 
     def get_serverapp_context(self, run_id: int) -> Optional[Context]:
@@ -1026,7 +1065,7 @@ def message_to_dict(message: Message) -> dict[str, Any]:
         "run_id": message.metadata.run_id,
         "src_node_id": message.metadata.src_node_id,
         "dst_node_id": message.metadata.dst_node_id,
-        "reply_to_message": message.metadata.reply_to_message,
+        "reply_to_message_id": message.metadata.reply_to_message_id,
         "created_at": message.metadata.created_at,
         "delivered_at": message.metadata.delivered_at,
         "ttl": message.metadata.ttl,
@@ -1036,7 +1075,7 @@ def message_to_dict(message: Message) -> dict[str, Any]:
     }
 
     if message.has_content():
-        result["content"] = recordset_to_proto(message.content).SerializeToString()
+        result["content"] = recorddict_to_proto(message.content).SerializeToString()
     else:
         result["error"] = error_to_proto(message.error).SerializeToString()
 
@@ -1047,20 +1086,15 @@ def dict_to_message(message_dict: dict[str, Any]) -> Message:
     """Transform dict to Message."""
     content, error = None, None
     if (b_content := message_dict.pop("content")) is not None:
-        content = recordset_from_proto(ProtoRecordSet.FromString(b_content))
+        content = recorddict_from_proto(ProtoRecordDict.FromString(b_content))
     if (b_error := message_dict.pop("error")) is not None:
         error = error_from_proto(ProtoError.FromString(b_error))
 
     # Metadata constructor doesn't allow passing created_at. We set it later
     metadata = Metadata(
-        **{
-            k: v
-            for k, v in message_dict.items()
-            if k not in ["created_at", "delivered_at"]
-        }
+        **{k: v for k, v in message_dict.items() if k not in ["delivered_at"]}
     )
-    msg = Message(metadata=metadata, content=content, error=error)
-    msg.metadata.__dict__["_created_at"] = message_dict["created_at"]
+    msg = make_message(metadata=metadata, content=content, error=error)
     msg.metadata.delivered_at = message_dict["delivered_at"]
     return msg