flwr-nightly 1.13.0.dev20241106__py3-none-any.whl → 1.13.0.dev20241117__py3-none-any.whl

flwr/server/superlink/linkstate/utils.py

@@ -15,21 +15,23 @@
 """Utility functions for State."""
 
 
-import time
 from logging import ERROR
 from os import urandom
-from uuid import uuid4
+from typing import Optional, Union
+from uuid import UUID, uuid4
 
-from flwr.common import ConfigsRecord, Context, log, serde
+from flwr.common import ConfigsRecord, Context, log, now, serde
 from flwr.common.constant import ErrorCode, Status, SubStatus
 from flwr.common.typing import RunStatus
-from flwr.proto.error_pb2 import Error  # pylint: disable=E0611
-from flwr.proto.message_pb2 import Context as ProtoContext  # pylint: disable=E0611
-from flwr.proto.node_pb2 import Node  # pylint: disable=E0611
 
 # 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  # pylint: disable=E0611
+from flwr.proto.task_pb2 import Task, TaskIns, TaskRes
+
+# pylint: enable=E0611
 
 NODE_UNAVAILABLE_ERROR_REASON = (
     "Error: Node Unavailable - The destination node is currently unavailable. "
@@ -40,12 +42,22 @@ VALID_RUN_STATUS_TRANSITIONS = {
     (Status.PENDING, Status.STARTING),
     (Status.STARTING, Status.RUNNING),
     (Status.RUNNING, Status.FINISHED),
+    # Any non-FINISHED status can transition to FINISHED
+    (Status.PENDING, Status.FINISHED),
+    (Status.STARTING, Status.FINISHED),
 }
 VALID_RUN_SUB_STATUSES = {
     SubStatus.COMPLETED,
     SubStatus.FAILED,
     SubStatus.STOPPED,
 }
+MESSAGE_UNAVAILABLE_ERROR_REASON = (
+    "Error: Message Unavailable - The requested message could not be found in the "
+    "database. It may have expired due to its TTL or never existed."
+)
+REPLY_MESSAGE_UNAVAILABLE_ERROR_REASON = (
+    "Error: Reply Message Unavailable - The reply message has expired."
+)
 
 
 def generate_rand_int_from_bytes(num_bytes: int) -> int:
@@ -161,31 +173,6 @@ def configsrecord_from_bytes(configsrecord_bytes: bytes) -> ConfigsRecord:
     )
 
 
-def make_node_unavailable_taskres(ref_taskins: TaskIns) -> TaskRes:
-    """Generate a TaskRes with a node unavailable error from a TaskIns."""
-    current_time = time.time()
-    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(
-            producer=Node(node_id=ref_taskins.task.consumer.node_id, anonymous=False),
-            consumer=Node(node_id=ref_taskins.task.producer.node_id, anonymous=False),
-            created_at=current_time,
-            ttl=ttl,
-            ancestry=[ref_taskins.task_id],
-            task_type=ref_taskins.task.task_type,
-            error=Error(
-                code=ErrorCode.NODE_UNAVAILABLE, reason=NODE_UNAVAILABLE_ERROR_REASON
-            ),
-        ),
-    )
-
-
 def is_valid_transition(current_status: RunStatus, new_status: RunStatus) -> bool:
     """Check if a transition between two run statuses is valid.
 
@@ -201,6 +188,14 @@ def is_valid_transition(current_status: RunStatus, new_status: RunStatus) -> boo
     bool
         True if the transition is valid, False otherwise.
     """
+    # Transition to FINISHED from a non-RUNNING status is only allowed
+    # if the sub-status is not COMPLETED
+    if (
+        current_status.status in [Status.PENDING, Status.STARTING]
+        and new_status.status == Status.FINISHED
+    ):
+        return new_status.sub_status != SubStatus.COMPLETED
+
     return (
         current_status.status,
         new_status.status,
@@ -228,3 +223,167 @@ def has_valid_sub_status(status: RunStatus) -> bool:
     if status.status == Status.FINISHED:
         return status.sub_status in VALID_RUN_SUB_STATUSES
     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).
+    """
+    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=0, anonymous=False),
+            consumer=Node(node_id=0, anonymous=False),
+            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,
+            ),
+        ),
+    )
+
+
+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.
+
+    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=0, anonymous=False),
+            consumer=Node(node_id=0, anonymous=False),
+            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,
+            ),
+        ),
+    )
+
+
+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 verify_taskins_ids(
+    inquired_taskins_ids: set[UUID],
+    found_taskins_dict: dict[UUID, TaskIns],
+    current_time: Optional[float] = None,
+    update_set: bool = True,
+) -> dict[UUID, TaskRes]:
+    """Verify found TaskIns and generate error TaskRes 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.
+    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,
+        by default True.
+
+    Returns
+    -------
+    dict[UUID, TaskRes]
+        A dictionary of error TaskRes indexed by the corresponding TaskIns ID.
+    """
+    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):
+            if update_set:
+                inquired_taskins_ids.remove(taskins_id)
+            taskres = create_taskres_for_unavailable_taskins(taskins_id)
+            ret_dict[taskins_id] = taskres
+    return ret_dict
+
+
+def verify_found_taskres(
+    inquired_taskins_ids: set[UUID],
+    found_taskins_dict: dict[UUID, TaskIns],
+    found_taskres_list: list[TaskRes],
+    current_time: Optional[float] = None,
+    update_set: bool = True,
+) -> dict[UUID, TaskRes]:
+    """Verify found TaskRes and generate error TaskRes 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.
+    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.
+
+    Returns
+    -------
+    dict[UUID, TaskRes]
+        A dictionary of TaskRes indexed by the corresponding TaskIns ID.
+    """
+    ret_dict: dict[UUID, TaskRes] = {}
+    current = current_time if current_time else now().timestamp()
+    for taskres in found_taskres_list:
+        taskins_id = UUID(taskres.task.ancestry[0])
+        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]
+            )
+            taskres.task.delivered_at = now().isoformat()
+        ret_dict[taskins_id] = taskres
+    return ret_dict

flwr/server/superlink/simulation/simulationio_servicer.py

@@ -23,6 +23,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,
     context_from_proto,
     context_to_proto,
     fab_to_proto,
@@ -36,6 +37,8 @@ from flwr.proto.log_pb2 import (  # pylint: disable=E0611
     PushLogsResponse,
 )
 from flwr.proto.run_pb2 import (  # pylint: disable=E0611
+    GetFederationOptionsRequest,
+    GetFederationOptionsResponse,
     UpdateRunStatusRequest,
     UpdateRunStatusResponse,
 )
@@ -123,10 +126,28 @@ class SimulationIoServicer(simulationio_pb2_grpc.SimulationIoServicer):
         self, request: PushLogsRequest, context: grpc.ServicerContext
     ) -> PushLogsResponse:
         """Push logs."""
-        log(DEBUG, "ServerAppIoServicer.PushLogs")
+        log(DEBUG, "SimultionIoServicer.PushLogs")
         state = self.state_factory.state()
 
         # Add logs to LinkState
         merged_logs = "".join(request.logs)
         state.add_serverapp_log(request.run_id, merged_logs)
         return PushLogsResponse()
+
+    def GetFederationOptions(
+        self, request: GetFederationOptionsRequest, context: ServicerContext
+    ) -> GetFederationOptionsResponse:
+        """Get Federation Options associated with a run."""
+        log(DEBUG, "SimultionIoServicer.GetFederationOptions")
+        state = self.state_factory.state()
+
+        federation_options = state.get_federation_options(request.run_id)
+        if federation_options is None:
+            context.abort(
+                grpc.StatusCode.FAILED_PRECONDITION,
+                "Expected federation options to be set, but none available.",
+            )
+            return GetFederationOptionsResponse()
+        return GetFederationOptionsResponse(
+            federation_options=configs_record_to_proto(federation_options)
+        )

flwr/simulation/__init__.py

@@ -17,13 +17,14 @@
 
 import importlib
 
+from flwr.simulation.app import run_simulation_process
 from flwr.simulation.run_simulation import run_simulation
 from flwr.simulation.simulationio_connection import SimulationIoConnection
 
 is_ray_installed = importlib.util.find_spec("ray") is not None
 
 if is_ray_installed:
-    from flwr.simulation.app import start_simulation
+    from flwr.simulation.legacy_app import start_simulation
 else:
     RAY_IMPORT_ERROR: str = """Unable to import module `ray`.
 
@@ -40,5 +41,6 @@ To install the necessary dependencies, install `flwr` with the `simulation` extr
 __all__ = [
     "SimulationIoConnection",
     "run_simulation",
+    "run_simulation_process",
     "start_simulation",
 ]