flwr-nightly 1.16.0.dev20250306__py3-none-any.whl → 1.16.0.dev20250308__py3-none-any.whl

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

@@ -29,6 +29,7 @@ from typing import Callable, Optional
 from flwr.client.client_app import ClientApp, ClientAppException, LoadClientAppError
 from flwr.client.clientapp.utils import get_load_client_app_fn
 from flwr.client.run_info_store import DeprecatedRunInfoStore
+from flwr.common import Message
 from flwr.common.constant import (
     NUM_PARTITIONS_KEY,
     PARTITION_ID_KEY,
@@ -37,9 +38,7 @@ from flwr.common.constant import (
 )
 from flwr.common.logger import log
 from flwr.common.message import Error
-from flwr.common.serde import message_from_taskins, message_to_taskres
 from flwr.common.typing import Run
-from flwr.proto.task_pb2 import TaskIns, TaskRes  # pylint: disable=E0611
 from flwr.server.superlink.linkstate import LinkState, LinkStateFactory
 
 from .backend import Backend, error_messages_backends, supported_backends
@@ -87,33 +86,33 @@ def _register_node_info_stores(
 
 # pylint: disable=too-many-arguments,too-many-locals
 def worker(
-    taskins_queue: "Queue[TaskIns]",
-    taskres_queue: "Queue[TaskRes]",
+    messageins_queue: Queue[Message],
+    messageres_queue: Queue[Message],
     node_info_store: dict[int, DeprecatedRunInfoStore],
     backend: Backend,
     f_stop: threading.Event,
 ) -> None:
-    """Get TaskIns from queue and pass it to an actor in the pool to execute it."""
+    """Process messages from the queue, execute them, update context, and enqueue
+    replies."""
     while not f_stop.is_set():
         out_mssg = None
         try:
             # Fetch from queue with timeout. We use a timeout so
             # the stopping event can be evaluated even when the queue is empty.
-            task_ins: TaskIns = taskins_queue.get(timeout=1.0)
-            node_id = task_ins.task.consumer.node_id
+            message: Message = messageins_queue.get(timeout=1.0)
+            node_id = message.metadata.dst_node_id
 
             # Retrieve context
-            context = node_info_store[node_id].retrieve_context(run_id=task_ins.run_id)
-
-            # Convert TaskIns to Message
-            message = message_from_taskins(task_ins)
+            context = node_info_store[node_id].retrieve_context(
+                run_id=message.metadata.run_id
+            )
 
             # Let backend process message
             out_mssg, updated_context = backend.process_message(message, context)
 
             # Update Context
             node_info_store[node_id].update_context(
-                task_ins.run_id, context=updated_context
+                message.metadata.run_id, context=updated_context
             )
         except Empty:
             # An exception raised if queue.get times out
@@ -137,35 +136,33 @@ def worker(
 
         finally:
             if out_mssg:
-                # Convert to TaskRes
-                task_res = message_to_taskres(out_mssg)
-                # Store TaskRes in state
-                taskres_queue.put(task_res)
+                # Store reply Messages in state
+                messageres_queue.put(out_mssg)
 
 
-def add_taskins_to_queue(
+def add_messages_to_queue(
     state: LinkState,
-    queue: "Queue[TaskIns]",
+    queue: Queue[Message],
     nodes_mapping: NodeToPartitionMapping,
     f_stop: threading.Event,
 ) -> None:
-    """Put TaskIns in a queue from State."""
+    """Put Messages in the queue from the LinkState."""
     while not f_stop.is_set():
         for node_id in nodes_mapping.keys():
-            task_ins_list = state.get_task_ins(node_id=node_id, limit=1)
-            for task_ins in task_ins_list:
-                queue.put(task_ins)
+            message_ins_list = state.get_message_ins(node_id=node_id, limit=1)
+            for msg in message_ins_list:
+                queue.put(msg)
         sleep(0.1)
 
 
-def put_taskres_into_state(
-    state: LinkState, queue: "Queue[TaskRes]", f_stop: threading.Event
+def put_message_into_state(
+    state: LinkState, queue: Queue[Message], f_stop: threading.Event
 ) -> None:
-    """Put TaskRes into State from a queue."""
+    """Store reply Messages into the LinkState from the queue."""
     while not f_stop.is_set():
         try:
-            taskres = queue.get(timeout=1.0)
-            state.store_task_res(taskres)
+            message_reply = queue.get(timeout=1.0)
+            state.store_message_res(message_reply)
         except Empty:
             # queue is empty when timeout was triggered
             pass
@@ -181,8 +178,8 @@ def run_api(
     f_stop: threading.Event,
 ) -> None:
     """Run the VCE."""
-    taskins_queue: Queue[TaskIns] = Queue()
-    taskres_queue: Queue[TaskRes] = Queue()
+    messageins_queue: Queue[Message] = Queue()
+    messageres_queue: Queue[Message] = Queue()
 
     try:
 
@@ -196,10 +193,10 @@ def run_api(
         state = state_factory.state()
 
         extractor_th = threading.Thread(
-            target=add_taskins_to_queue,
+            target=add_messages_to_queue,
             args=(
                 state,
-                taskins_queue,
+                messageins_queue,
                 nodes_mapping,
                 f_stop,
             ),
@@ -207,10 +204,10 @@ def run_api(
         extractor_th.start()
 
         injector_th = threading.Thread(
-            target=put_taskres_into_state,
+            target=put_message_into_state,
             args=(
                 state,
-                taskres_queue,
+                messageres_queue,
                 f_stop,
             ),
         )
@@ -220,8 +217,8 @@ def run_api(
             _ = [
                 executor.submit(
                     worker,
-                    taskins_queue,
-                    taskres_queue,
+                    messageins_queue,
+                    messageres_queue,
                     node_info_stores,
                     backend,
                     f_stop,

flwr/server/superlink/linkstate/in_memory_linkstate.py

@@ -33,18 +33,15 @@ from flwr.common.constant import (
 )
 from flwr.common.record import ConfigsRecord
 from flwr.common.typing import Run, RunStatus, UserConfig
-from flwr.proto.task_pb2 import TaskIns, TaskRes  # pylint: disable=E0611
 from flwr.server.superlink.linkstate.linkstate import LinkState
-from flwr.server.utils import validate_message, validate_task_ins_or_res
+from flwr.server.utils import validate_message
 
 from .utils import (
     generate_rand_int_from_bytes,
     has_valid_sub_status,
     is_valid_transition,
     verify_found_message_replies,
-    verify_found_taskres,
     verify_message_ids,
-    verify_taskins_ids,
 )
 
 
@@ -71,9 +68,6 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         self.run_ids: dict[int, RunRecord] = {}
         self.contexts: dict[int, Context] = {}
         self.federation_options: dict[int, ConfigsRecord] = {}
-        self.task_ins_store: dict[UUID, TaskIns] = {}
-        self.task_res_store: dict[UUID, TaskRes] = {}
-        self.task_ins_id_to_task_res_id: dict[UUID, UUID] = {}
         self.message_ins_store: dict[UUID, Message] = {}
         self.message_res_store: dict[UUID, Message] = {}
         self.message_ins_id_to_message_res_id: dict[UUID, UUID] = {}
@@ -82,45 +76,6 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
 
         self.lock = threading.RLock()
 
-    def store_task_ins(self, task_ins: TaskIns) -> Optional[UUID]:
-        """Store one TaskIns."""
-        # Validate task
-        errors = validate_task_ins_or_res(task_ins)
-        if any(errors):
-            log(ERROR, errors)
-            return None
-        # Validate run_id
-        if task_ins.run_id not in self.run_ids:
-            log(ERROR, "Invalid run ID for TaskIns: %s", task_ins.run_id)
-            return None
-        # Validate source node ID
-        if task_ins.task.producer.node_id != SUPERLINK_NODE_ID:
-            log(
-                ERROR,
-                "Invalid source node ID for TaskIns: %s",
-                task_ins.task.producer.node_id,
-            )
-            return None
-        # Validate destination node ID
-        if task_ins.task.consumer.node_id not in self.node_ids:
-            log(
-                ERROR,
-                "Invalid destination node ID for TaskIns: %s",
-                task_ins.task.consumer.node_id,
-            )
-            return None
-
-        # Create task_id
-        task_id = uuid4()
-
-        # Store TaskIns
-        task_ins.task_id = str(task_id)
-        with self.lock:
-            self.task_ins_store[task_id] = task_ins
-
-        # Return the new task_id
-        return task_id
-
     def store_message_ins(self, message: Message) -> Optional[UUID]:
         """Store one Message."""
         # Validate message
@@ -161,33 +116,6 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         # Return the new message_id
         return message_id
 
-    def get_task_ins(self, node_id: int, limit: Optional[int]) -> list[TaskIns]:
-        """Get all TaskIns that have not been delivered yet."""
-        if limit is not None and limit < 1:
-            raise AssertionError("`limit` must be >= 1")
-
-        # Find TaskIns for node_id that were not delivered yet
-        task_ins_list: list[TaskIns] = []
-        current_time = time.time()
-        with self.lock:
-            for _, task_ins in self.task_ins_store.items():
-                if (
-                    task_ins.task.consumer.node_id == node_id
-                    and task_ins.task.delivered_at == ""
-                    and task_ins.task.created_at + task_ins.task.ttl > current_time
-                ):
-                    task_ins_list.append(task_ins)
-                if limit and len(task_ins_list) == limit:
-                    break
-
-        # Mark all of them as delivered
-        delivered_at = now().isoformat()
-        for task_ins in task_ins_list:
-            task_ins.task.delivered_at = delivered_at
-
-        # Return TaskIns
-        return task_ins_list
-
     def get_message_ins(self, node_id: int, limit: Optional[int]) -> list[Message]:
         """Get all Messages that have not been delivered yet."""
         if limit is not None and limit < 1:
@@ -216,78 +144,6 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         # Return list of messages
         return message_ins_list
 
-    # pylint: disable=R0911
-    def store_task_res(self, task_res: TaskRes) -> Optional[UUID]:
-        """Store one TaskRes."""
-        # Validate task
-        errors = validate_task_ins_or_res(task_res)
-        if any(errors):
-            log(ERROR, errors)
-            return None
-
-        with self.lock:
-            # Check if the TaskIns it is replying to exists and is valid
-            task_ins_id = task_res.task.ancestry[0]
-            task_ins = self.task_ins_store.get(UUID(task_ins_id))
-
-            # Ensure that the consumer_id of taskIns matches the producer_id of taskRes.
-            if (
-                task_ins
-                and task_res
-                and task_ins.task.consumer.node_id != task_res.task.producer.node_id
-            ):
-                return None
-
-            if task_ins is None:
-                log(ERROR, "TaskIns with task_id %s does not exist.", task_ins_id)
-                return None
-
-            if task_ins.task.created_at + task_ins.task.ttl <= time.time():
-                log(
-                    ERROR,
-                    "Failed to store TaskRes: TaskIns with task_id %s has expired.",
-                    task_ins_id,
-                )
-                return None
-
-            # Fail if the TaskRes TTL exceeds the
-            # expiration time of the TaskIns it replies to.
-            # Condition: TaskIns.created_at + TaskIns.ttl ≥
-            #            TaskRes.created_at + TaskRes.ttl
-            # A small tolerance is introduced to account
-            # for floating-point precision issues.
-            max_allowed_ttl = (
-                task_ins.task.created_at + task_ins.task.ttl - task_res.task.created_at
-            )
-            if task_res.task.ttl and (
-                task_res.task.ttl - max_allowed_ttl > MESSAGE_TTL_TOLERANCE
-            ):
-                log(
-                    WARNING,
-                    "Received TaskRes with TTL %.2f "
-                    "exceeding the allowed maximum TTL %.2f.",
-                    task_res.task.ttl,
-                    max_allowed_ttl,
-                )
-                return None
-
-        # Validate run_id
-        if task_res.run_id not in self.run_ids:
-            log(ERROR, "`run_id` is invalid")
-            return None
-
-        # Create task_id
-        task_id = uuid4()
-
-        # Store TaskRes
-        task_res.task_id = str(task_id)
-        with self.lock:
-            self.task_res_store[task_id] = task_res
-            self.task_ins_id_to_task_res_id[UUID(task_ins_id)] = task_id
-
-        # Return the new task_id
-        return task_id
-
     # pylint: disable=R0911
     def store_message_res(self, message: Message) -> Optional[UUID]:
         """Store one Message."""
@@ -369,43 +225,6 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         # Return the new message_id
         return message_id
 
-    def get_task_res(self, task_ids: set[UUID]) -> list[TaskRes]:
-        """Get TaskRes for the given TaskIns IDs."""
-        ret: dict[UUID, TaskRes] = {}
-
-        with self.lock:
-            current = time.time()
-
-            # Verify TaskIns IDs
-            ret = verify_taskins_ids(
-                inquired_taskins_ids=task_ids,
-                found_taskins_dict=self.task_ins_store,
-                current_time=current,
-            )
-
-            # Find all TaskRes
-            task_res_found: list[TaskRes] = []
-            for task_id in task_ids:
-                # If TaskRes exists and is not delivered, add it to the list
-                if task_res_id := self.task_ins_id_to_task_res_id.get(task_id):
-                    task_res = self.task_res_store[task_res_id]
-                    if task_res.task.delivered_at == "":
-                        task_res_found.append(task_res)
-            tmp_ret_dict = verify_found_taskres(
-                inquired_taskins_ids=task_ids,
-                found_taskins_dict=self.task_ins_store,
-                found_taskres_list=task_res_found,
-                current_time=current,
-            )
-            ret.update(tmp_ret_dict)
-
-            # Mark existing TaskRes to be returned as delivered
-            delivered_at = now().isoformat()
-            for task_res in task_res_found:
-                task_res.task.delivered_at = delivered_at
-
-        return list(ret.values())
-
     def get_message_res(self, message_ids: set[UUID]) -> list[Message]:
         """Get reply Messages for the given Message IDs."""
         ret: dict[UUID, Message] = {}
@@ -445,21 +264,6 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
 
         return list(ret.values())
 
-    def delete_tasks(self, task_ins_ids: set[UUID]) -> None:
-        """Delete TaskIns/TaskRes pairs based on provided TaskIns IDs."""
-        if not task_ins_ids:
-            return
-
-        with self.lock:
-            for task_id in task_ins_ids:
-                # Delete TaskIns
-                if task_id in self.task_ins_store:
-                    del self.task_ins_store[task_id]
-                # Delete TaskRes
-                if task_id in self.task_ins_id_to_task_res_id:
-                    task_res_id = self.task_ins_id_to_task_res_id.pop(task_id)
-                    del self.task_res_store[task_res_id]
-
     def delete_messages(self, message_ins_ids: set[UUID]) -> None:
         """Delete a Message and its reply based on provided Message IDs."""
         if not message_ins_ids:
@@ -477,16 +281,6 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
                     )
                     del self.message_res_store[message_res_id]
 
-    def get_task_ids_from_run_id(self, run_id: int) -> set[UUID]:
-        """Get all TaskIns IDs for the given run_id."""
-        task_id_list: set[UUID] = set()
-        with self.lock:
-            for task_id, task_ins in self.task_ins_store.items():
-                if task_ins.run_id == run_id:
-                    task_id_list.add(task_id)
-
-        return task_id_list
-
     def get_message_ids_from_run_id(self, run_id: int) -> set[UUID]:
         """Get all instruction Message IDs for the given run_id."""
         message_id_list: set[UUID] = set()
@@ -497,13 +291,6 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
 
         return message_id_list
 
-    def num_task_ins(self) -> int:
-        """Calculate the number of task_ins in store.
-
-        This includes delivered but not yet deleted task_ins.
-        """
-        return len(self.task_ins_store)
-
     def num_message_ins(self) -> int:
         """Calculate the number of instruction Messages in store.
 
@@ -511,13 +298,6 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         """
         return len(self.message_ins_store)
 
-    def num_task_res(self) -> int:
-        """Calculate the number of task_res in store.
-
-        This includes delivered but not yet deleted task_res.
-        """
-        return len(self.task_res_store)
-
     def num_message_res(self) -> int:
         """Calculate the number of reply Messages in store.
 

flwr/server/superlink/linkstate/linkstate.py

@@ -22,30 +22,11 @@ from uuid import UUID
 from flwr.common import Context, Message
 from flwr.common.record import ConfigsRecord
 from flwr.common.typing import Run, RunStatus, UserConfig
-from flwr.proto.task_pb2 import TaskIns, TaskRes  # pylint: disable=E0611
 
 
 class LinkState(abc.ABC):  # pylint: disable=R0904
     """Abstract LinkState."""
 
-    @abc.abstractmethod
-    def store_task_ins(self, task_ins: TaskIns) -> Optional[UUID]:
-        """Store one TaskIns.
-
-        Usually, the ServerAppIo API calls this to schedule instructions.
-
-        Stores the value of the `task_ins` in the link state and, if successful,
-        returns the `task_id` (UUID) of the `task_ins`. If, for any reason,
-        storing the `task_ins` fails, `None` is returned.
-
-        Constraints
-        -----------
-        `task_ins.task.consumer.node_id` MUST be set (not constant.DRIVER_NODE_ID)
-
-        If `task_ins.run_id` is invalid, then
-        storing the `task_ins` MUST fail.
-        """
-
     @abc.abstractmethod
     def store_message_ins(self, message: Message) -> Optional[UUID]:
         """Store one Message.
@@ -64,28 +45,6 @@ class LinkState(abc.ABC):  # pylint: disable=R0904
         storing the `message` MUST fail.
         """
 
-    @abc.abstractmethod
-    def get_task_ins(self, node_id: int, limit: Optional[int]) -> list[TaskIns]:
-        """Get TaskIns optionally filtered by node_id.
-
-        Usually, the Fleet API calls this for Nodes planning to work on one or more
-        TaskIns.
-
-        Constraints
-        -----------
-        Retrieve all TaskIns where
-
-            1. the `task_ins.task.consumer.node_id` equals `node_id` AND
-            2. the `task_ins.task.delivered_at` equals `""`.
-
-
-        If `delivered_at` MUST BE set (not `""`) otherwise the TaskIns MUST not be in
-        the result.
-
-        If `limit` is not `None`, return, at most, `limit` number of `task_ins`. If
-        `limit` is set, it has to be greater zero.
-        """
-
     @abc.abstractmethod
     def get_message_ins(self, node_id: int, limit: Optional[int]) -> list[Message]:
         """Get zero or more `Message` objects for the provided `node_id`.
@@ -101,24 +60,6 @@ class LinkState(abc.ABC):  # pylint: disable=R0904
         `limit` is set, it has to be greater zero.
         """
 
-    @abc.abstractmethod
-    def store_task_res(self, task_res: TaskRes) -> Optional[UUID]:
-        """Store one TaskRes.
-
-        Usually, the Fleet API calls this for Nodes returning results.
-
-        Stores the TaskRes and, if successful, returns the `task_id` (UUID) of
-        the `task_res`. If storing the `task_res` fails, `None` is returned.
-
-        Constraints
-        -----------
-
-        `task_res.task.consumer.node_id` MUST be set (not constant.DRIVER_NODE_ID)
-
-        If `task_res.run_id` is invalid, then
-        storing the `task_res` MUST fail.
-        """
-
     @abc.abstractmethod
     def store_message_res(self, message: Message) -> Optional[UUID]:
         """Store one Message.
@@ -136,31 +77,6 @@ class LinkState(abc.ABC):  # pylint: disable=R0904
         storing the `message` MUST fail.
         """
 
-    @abc.abstractmethod
-    def get_task_res(self, task_ids: set[UUID]) -> list[TaskRes]:
-        """Get TaskRes for the given TaskIns IDs.
-
-        This method is typically called by the ServerAppIo API to obtain
-        results (TaskRes) for previously scheduled instructions (TaskIns).
-        For each task_id provided, this method returns one of the following responses:
-
-        - An error TaskRes if the corresponding TaskIns does not exist or has expired.
-        - An error TaskRes if the corresponding TaskRes exists but has expired.
-        - The valid TaskRes if the TaskIns has a corresponding valid TaskRes.
-        - Nothing if the TaskIns is still valid and waiting for a TaskRes.
-
-        Parameters
-        ----------
-        task_ids : set[UUID]
-            A set of TaskIns IDs for which to retrieve results (TaskRes).
-
-        Returns
-        -------
-        list[TaskRes]
-            A list of TaskRes corresponding to the given task IDs. If no
-            TaskRes could be found for any of the task IDs, an empty list is returned.
-        """
-
     @abc.abstractmethod
     def get_message_res(self, message_ids: set[UUID]) -> list[Message]:
         """Get reply Messages for the given Message IDs.
@@ -188,39 +104,14 @@ class LinkState(abc.ABC):  # pylint: disable=R0904
             carrying an Error.
         """
 
-    @abc.abstractmethod
-    def num_task_ins(self) -> int:
-        """Calculate the number of task_ins in store.
-
-        This includes delivered but not yet deleted task_ins.
-        """
-
     @abc.abstractmethod
     def num_message_ins(self) -> int:
         """Calculate the number of Messages awaiting a reply."""
 
-    @abc.abstractmethod
-    def num_task_res(self) -> int:
-        """Calculate the number of task_res in store.
-
-        This includes delivered but not yet deleted task_res.
-        """
-
     @abc.abstractmethod
     def num_message_res(self) -> int:
         """Calculate the number of reply Messages in store."""
 
-    @abc.abstractmethod
-    def delete_tasks(self, task_ins_ids: set[UUID]) -> None:
-        """Delete TaskIns/TaskRes pairs based on provided TaskIns IDs.
-
-        Parameters
-        ----------
-        task_ins_ids : set[UUID]
-            A set of TaskIns IDs. For each ID in the set, the corresponding
-            TaskIns and its associated TaskRes will be deleted.
-        """
-
     @abc.abstractmethod
     def delete_messages(self, message_ins_ids: set[UUID]) -> None:
         """Delete a Message and its reply based on provided Message IDs.
@@ -232,10 +123,6 @@ class LinkState(abc.ABC):  # pylint: disable=R0904
             Message and its associated reply Message will be deleted.
         """
 
-    @abc.abstractmethod
-    def get_task_ids_from_run_id(self, run_id: int) -> set[UUID]:
-        """Get all TaskIns IDs for the given run_id."""
-
     @abc.abstractmethod
     def get_message_ids_from_run_id(self, run_id: int) -> set[UUID]:
         """Get all instruction Message IDs for the given run_id."""