flwr-nightly 1.13.0.dev20241021__py3-none-any.whl → 1.13.0.dev20241111__py3-none-any.whl

flwr/server/superlink/{state → linkstate}/utils.py

@@ -15,22 +15,32 @@
 """Utility functions for State."""
 
 
-import time
-from logging import ERROR
 from os import urandom
-from uuid import uuid4
 
-from flwr.common import log
-from flwr.common.constant import ErrorCode
-from flwr.proto.error_pb2 import Error  # pylint: disable=E0611
-from flwr.proto.node_pb2 import Node  # pylint: disable=E0611
-from flwr.proto.task_pb2 import Task, TaskIns, TaskRes  # pylint: disable=E0611
+from flwr.common import ConfigsRecord, Context, serde
+from flwr.common.constant import Status, SubStatus
+from flwr.common.typing import RunStatus
+from flwr.proto.message_pb2 import Context as ProtoContext  # pylint: disable=E0611
+
+# pylint: disable=E0611
+from flwr.proto.recordset_pb2 import ConfigsRecord as ProtoConfigsRecord
 
 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),
+}
+VALID_RUN_SUB_STATUSES = {
+    SubStatus.COMPLETED,
+    SubStatus.FAILED,
+    SubStatus.STOPPED,
+}
+
 
 def generate_rand_int_from_bytes(num_bytes: int) -> int:
     """Generate a random unsigned integer from `num_bytes` bytes."""
@@ -123,26 +133,67 @@ def convert_sint64_values_in_dict_to_uint64(
             data_dict[key] = convert_sint64_to_uint64(data_dict[key])
 
 
-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 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.
+    """
+    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 == ""

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.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 ..fleet.grpc_bidi.grpc_server import generic_create_grpc_server
+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,153 @@
+# 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_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,
+    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
+
+
+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()
+        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()
+
+        # 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 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)
+        )

flwr/simulation/__init__.py

@@ -17,12 +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`.
 
@@ -37,6 +39,8 @@ To install the necessary dependencies, install `flwr` with the `simulation` extr
 
 
 __all__ = [
+    "SimulationIoConnection",
     "run_simulation",
+    "run_simulation_process",
     "start_simulation",
 ]