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

flwr/server/superlink/linkstate/utils.py

@@ -0,0 +1,399 @@
+# Copyright 2024 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.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Utility functions for State."""
+
+
+from logging import ERROR
+from os import urandom
+from typing import Optional, Union
+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.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
+
+# 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),
+    (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, exclude: Optional[list[int]] = None
+) -> int:
+    """Generate a random unsigned integer from `num_bytes` bytes.
+
+    If `exclude` is set, this function guarantees such number is not returned.
+    """
+    num = int.from_bytes(urandom(num_bytes), "little", signed=False)
+
+    if exclude:
+        while num in exclude:
+            num = int.from_bytes(urandom(num_bytes), "little", signed=False)
+    return num
+
+
+def convert_uint64_to_sint64(u: int) -> int:
+    """Convert a uint64 value to a sint64 value with the same bit sequence.
+
+    Parameters
+    ----------
+    u : int
+        The unsigned 64-bit integer to convert.
+
+    Returns
+    -------
+    int
+        The signed 64-bit integer equivalent.
+
+        The signed 64-bit integer will have the same bit pattern as the
+        unsigned 64-bit integer but may have a different decimal value.
+
+        For numbers within the range [0, `sint64` max value], the decimal
+        value remains the same. However, for numbers greater than the `sint64`
+        max value, the decimal value will differ due to the wraparound caused
+        by the sign bit.
+    """
+    if u >= (1 << 63):
+        return u - (1 << 64)
+    return u
+
+
+def convert_sint64_to_uint64(s: int) -> int:
+    """Convert a sint64 value to a uint64 value with the same bit sequence.
+
+    Parameters
+    ----------
+    s : int
+        The signed 64-bit integer to convert.
+
+    Returns
+    -------
+    int
+        The unsigned 64-bit integer equivalent.
+
+        The unsigned 64-bit integer will have the same bit pattern as the
+        signed 64-bit integer but may have a different decimal value.
+
+        For negative `sint64` values, the conversion adds 2^64 to the
+        signed value to obtain the equivalent `uint64` value. For non-negative
+        `sint64` values, the decimal value remains unchanged in the `uint64`
+        representation.
+    """
+    if s < 0:
+        return s + (1 << 64)
+    return s
+
+
+def convert_uint64_values_in_dict_to_sint64(
+    data_dict: dict[str, int], keys: list[str]
+) -> None:
+    """Convert uint64 values to sint64 in the given dictionary.
+
+    Parameters
+    ----------
+    data_dict : dict[str, int]
+        A dictionary where the values are integers to be converted.
+    keys : list[str]
+        A list of keys in the dictionary whose values need to be converted.
+    """
+    for key in keys:
+        if key in data_dict:
+            data_dict[key] = convert_uint64_to_sint64(data_dict[key])
+
+
+def convert_sint64_values_in_dict_to_uint64(
+    data_dict: dict[str, int], keys: list[str]
+) -> None:
+    """Convert sint64 values to uint64 in the given dictionary.
+
+    Parameters
+    ----------
+    data_dict : dict[str, int]
+        A dictionary where the values are integers to be converted.
+    keys : list[str]
+        A list of keys in the dictionary whose values need to be converted.
+    """
+    for key in keys:
+        if key in data_dict:
+            data_dict[key] = convert_sint64_to_uint64(data_dict[key])
+
+
+def context_to_bytes(context: Context) -> bytes:
+    """Serialize `Context` to bytes."""
+    return serde.context_to_proto(context).SerializeToString()
+
+
+def context_from_bytes(context_bytes: bytes) -> Context:
+    """Deserialize `Context` from bytes."""
+    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 configsrecord_from_bytes(configsrecord_bytes: bytes) -> ConfigsRecord:
+    """Deserialize `ConfigsRecord` from bytes."""
+    return serde.configs_record_from_proto(
+        ProtoConfigsRecord.FromString(configsrecord_bytes)
+    )
+
+
+def is_valid_transition(current_status: RunStatus, new_status: RunStatus) -> bool:
+    """Check if a transition between two run statuses is valid.
+
+    Parameters
+    ----------
+    current_status : RunStatus
+        The current status of the run.
+    new_status : RunStatus
+        The new status to transition to.
+
+    Returns
+    -------
+    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,
+    ) in VALID_RUN_STATUS_TRANSITIONS
+
+
+def has_valid_sub_status(status: RunStatus) -> bool:
+    """Check if the 'sub_status' field of the given status is valid.
+
+    Parameters
+    ----------
+    status : RunStatus
+        The status object to be checked.
+
+    Returns
+    -------
+    bool
+        True if the status object has a valid sub-status, False otherwise.
+
+    Notes
+    -----
+    Only an empty string (i.e., "") is considered a valid sub-status for
+    non-finished statuses. The sub-status of a finished status cannot be empty.
+    """
+    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=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,
+            ),
+        ),
+    )
+
+
+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=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,
+            ),
+        ),
+    )
+
+
+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/__init__.py

@@ -0,0 +1,15 @@
+# Copyright 2024 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.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Flower SimulationIo service."""

flwr/server/superlink/simulation/simulationio_grpc.py

@@ -0,0 +1,65 @@
+# Copyright 2024 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.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""SimulationIo gRPC API."""
+
+
+from logging import INFO
+from typing import Optional
+
+import grpc
+
+from flwr.common import GRPC_MAX_MESSAGE_LENGTH
+from flwr.common.grpc import generic_create_grpc_server
+from flwr.common.logger import log
+from flwr.proto.simulationio_pb2_grpc import (  # pylint: disable=E0611
+    add_SimulationIoServicer_to_server,
+)
+from flwr.server.superlink.ffs.ffs_factory import FfsFactory
+from flwr.server.superlink.linkstate import LinkStateFactory
+
+from .simulationio_servicer import SimulationIoServicer
+
+
+def run_simulationio_api_grpc(
+    address: str,
+    state_factory: LinkStateFactory,
+    ffs_factory: FfsFactory,
+    certificates: Optional[tuple[bytes, bytes, bytes]],
+) -> grpc.Server:
+    """Run SimulationIo API (gRPC, request-response)."""
+    # Create SimulationIo API gRPC server
+    simulationio_servicer: grpc.Server = SimulationIoServicer(
+        state_factory=state_factory,
+        ffs_factory=ffs_factory,
+    )
+    simulationio_add_servicer_to_server_fn = add_SimulationIoServicer_to_server
+    simulationio_grpc_server = generic_create_grpc_server(
+        servicer_and_add_fn=(
+            simulationio_servicer,
+            simulationio_add_servicer_to_server_fn,
+        ),
+        server_address=address,
+        max_message_length=GRPC_MAX_MESSAGE_LENGTH,
+        certificates=certificates,
+    )
+
+    log(
+        INFO,
+        "Flower Simulation Engine: Starting SimulationIo API on %s",
+        address,
+    )
+    simulationio_grpc_server.start()
+
+    return simulationio_grpc_server

flwr/server/superlink/simulation/simulationio_servicer.py

@@ -0,0 +1,186 @@
+# Copyright 2024 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.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""SimulationIo API servicer."""
+
+
+import threading
+from logging import DEBUG, INFO
+
+import grpc
+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,
+    run_status_from_proto,
+    run_status_to_proto,
+    run_to_proto,
+)
+from flwr.common.typing import Fab, RunStatus
+from flwr.proto import simulationio_pb2_grpc
+from flwr.proto.log_pb2 import (  # pylint: disable=E0611
+    PushLogsRequest,
+    PushLogsResponse,
+)
+from flwr.proto.run_pb2 import (  # pylint: disable=E0611
+    GetFederationOptionsRequest,
+    GetFederationOptionsResponse,
+    GetRunStatusRequest,
+    GetRunStatusResponse,
+    UpdateRunStatusRequest,
+    UpdateRunStatusResponse,
+)
+from flwr.proto.simulationio_pb2 import (  # pylint: disable=E0611
+    PullSimulationInputsRequest,
+    PullSimulationInputsResponse,
+    PushSimulationOutputsRequest,
+    PushSimulationOutputsResponse,
+)
+from flwr.server.superlink.ffs.ffs_factory import FfsFactory
+from flwr.server.superlink.linkstate import LinkStateFactory
+from flwr.server.superlink.utils import abort_if
+
+
+class SimulationIoServicer(simulationio_pb2_grpc.SimulationIoServicer):
+    """SimulationIo API servicer."""
+
+    def __init__(
+        self, state_factory: LinkStateFactory, ffs_factory: FfsFactory
+    ) -> None:
+        self.state_factory = state_factory
+        self.ffs_factory = ffs_factory
+        self.lock = threading.RLock()
+
+    def PullSimulationInputs(
+        self, request: PullSimulationInputsRequest, context: ServicerContext
+    ) -> PullSimulationInputsResponse:
+        """Pull SimultionIo process inputs."""
+        log(DEBUG, "SimultionIoServicer.SimultionIoInputs")
+        # Init access to LinkState and Ffs
+        state = self.state_factory.state()
+        ffs = self.ffs_factory.ffs()
+
+        # Lock access to LinkState, preventing obtaining the same pending run_id
+        with self.lock:
+            # Attempt getting the run_id of a pending run
+            run_id = state.get_pending_run_id()
+            # If there's no pending run, return an empty response
+            if run_id is None:
+                return PullSimulationInputsResponse()
+
+            # Retrieve Context, Run and Fab for the run_id
+            serverapp_ctxt = state.get_serverapp_context(run_id)
+            run = state.get_run(run_id)
+            fab = None
+            if run and run.fab_hash:
+                if result := ffs.get(run.fab_hash):
+                    fab = Fab(run.fab_hash, result[0])
+            if run and fab and serverapp_ctxt:
+                # Update run status to STARTING
+                if state.update_run_status(run_id, RunStatus(Status.STARTING, "", "")):
+                    log(INFO, "Starting run %d", run_id)
+                    return PullSimulationInputsResponse(
+                        context=context_to_proto(serverapp_ctxt),
+                        run=run_to_proto(run),
+                        fab=fab_to_proto(fab),
+                    )
+
+        # Raise an exception if the Run or Fab is not found,
+        # or if the status cannot be updated to STARTING
+        raise RuntimeError(f"Failed to start run {run_id}")
+
+    def PushSimulationOutputs(
+        self, request: PushSimulationOutputsRequest, context: ServicerContext
+    ) -> PushSimulationOutputsResponse:
+        """Push Simulation process outputs."""
+        log(DEBUG, "SimultionIoServicer.PushSimulationOutputs")
+        state = self.state_factory.state()
+
+        # Abort if the run is not running
+        abort_if(
+            request.run_id,
+            [Status.PENDING, Status.STARTING, Status.FINISHED],
+            state,
+            context,
+        )
+
+        state.set_serverapp_context(request.run_id, context_from_proto(request.context))
+        return PushSimulationOutputsResponse()
+
+    def UpdateRunStatus(
+        self, request: UpdateRunStatusRequest, context: grpc.ServicerContext
+    ) -> UpdateRunStatusResponse:
+        """Update the status of a run."""
+        log(DEBUG, "SimultionIoServicer.UpdateRunStatus")
+        state = self.state_factory.state()
+
+        # Abort if the run is finished
+        abort_if(request.run_id, [Status.FINISHED], state, context)
+
+        # Update the run status
+        state.update_run_status(
+            run_id=request.run_id, new_status=run_status_from_proto(request.run_status)
+        )
+        return UpdateRunStatusResponse()
+
+    def GetRunStatus(
+        self, request: GetRunStatusRequest, context: ServicerContext
+    ) -> GetRunStatusResponse:
+        """Get status of requested runs."""
+        log(DEBUG, "SimultionIoServicer.GetRunStatus")
+        state = self.state_factory.state()
+
+        statuses = state.get_run_status(set(request.run_ids))
+
+        return GetRunStatusResponse(
+            run_status_dict={
+                run_id: run_status_to_proto(status)
+                for run_id, status in statuses.items()
+            }
+        )
+
+    def PushLogs(
+        self, request: PushLogsRequest, context: grpc.ServicerContext
+    ) -> PushLogsResponse:
+        """Push logs."""
+        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)
+        )