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

flwr/server/superlink/driver/{driver_grpc.py → serverappio_grpc.py}

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Driver gRPC API."""
+"""ServerAppIo gRPC API."""
 
 from logging import INFO
 from typing import Optional
@@ -21,37 +21,40 @@ import grpc
 
 from flwr.common import GRPC_MAX_MESSAGE_LENGTH
 from flwr.common.logger import log
-from flwr.proto.driver_pb2_grpc import (  # pylint: disable=E0611
-    add_DriverServicer_to_server,
+from flwr.proto.serverappio_pb2_grpc import (  # pylint: disable=E0611
+    add_ServerAppIoServicer_to_server,
 )
 from flwr.server.superlink.ffs.ffs_factory import FfsFactory
-from flwr.server.superlink.state import StateFactory
+from flwr.server.superlink.linkstate import LinkStateFactory
 
 from ..fleet.grpc_bidi.grpc_server import generic_create_grpc_server
-from .driver_servicer import DriverServicer
+from .serverappio_servicer import ServerAppIoServicer
 
 
-def run_driver_api_grpc(
+def run_serverappio_api_grpc(
     address: str,
-    state_factory: StateFactory,
+    state_factory: LinkStateFactory,
     ffs_factory: FfsFactory,
     certificates: Optional[tuple[bytes, bytes, bytes]],
 ) -> grpc.Server:
-    """Run Driver API (gRPC, request-response)."""
-    # Create Driver API gRPC server
-    driver_servicer: grpc.Server = DriverServicer(
+    """Run ServerAppIo API (gRPC, request-response)."""
+    # Create ServerAppIo API gRPC server
+    serverappio_servicer: grpc.Server = ServerAppIoServicer(
         state_factory=state_factory,
         ffs_factory=ffs_factory,
     )
-    driver_add_servicer_to_server_fn = add_DriverServicer_to_server
-    driver_grpc_server = generic_create_grpc_server(
-        servicer_and_add_fn=(driver_servicer, driver_add_servicer_to_server_fn),
+    serverappio_add_servicer_to_server_fn = add_ServerAppIoServicer_to_server
+    serverappio_grpc_server = generic_create_grpc_server(
+        servicer_and_add_fn=(
+            serverappio_servicer,
+            serverappio_add_servicer_to_server_fn,
+        ),
         server_address=address,
         max_message_length=GRPC_MAX_MESSAGE_LENGTH,
         certificates=certificates,
     )
 
-    log(INFO, "Flower ECE: Starting Driver API (gRPC-rere) on %s", address)
-    driver_grpc_server.start()
+    log(INFO, "Flower ECE: Starting ServerAppIo API (gRPC-rere) on %s", address)
+    serverappio_grpc_server.start()
 
-    return driver_grpc_server
+    return serverappio_grpc_server

flwr/server/superlink/driver/{driver_servicer.py → serverappio_servicer.py}

@@ -12,62 +12,80 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Driver API servicer."""
+"""ServerAppIo API servicer."""
 
 
+import threading
 import time
-from logging import DEBUG
+from logging import DEBUG, INFO
 from typing import Optional
 from uuid import UUID
 
 import grpc
 
+from flwr.common import ConfigsRecord
+from flwr.common.constant import Status
 from flwr.common.logger import log
 from flwr.common.serde import (
+    context_from_proto,
+    context_to_proto,
     fab_from_proto,
     fab_to_proto,
+    run_status_from_proto,
+    run_to_proto,
     user_config_from_proto,
-    user_config_to_proto,
-)
-from flwr.common.typing import Fab
-from flwr.proto import driver_pb2_grpc  # pylint: disable=E0611
-from flwr.proto.driver_pb2 import (  # pylint: disable=E0611
-    GetNodesRequest,
-    GetNodesResponse,
-    PullTaskResRequest,
-    PullTaskResResponse,
-    PushTaskInsRequest,
-    PushTaskInsResponse,
 )
+from flwr.common.typing import Fab, RunStatus
+from flwr.proto import serverappio_pb2_grpc  # pylint: disable=E0611
 from flwr.proto.fab_pb2 import GetFabRequest, GetFabResponse  # pylint: disable=E0611
+from flwr.proto.log_pb2 import (  # pylint: disable=E0611
+    PushLogsRequest,
+    PushLogsResponse,
+)
 from flwr.proto.node_pb2 import Node  # pylint: disable=E0611
 from flwr.proto.run_pb2 import (  # pylint: disable=E0611
     CreateRunRequest,
     CreateRunResponse,
     GetRunRequest,
     GetRunResponse,
-    Run,
+    UpdateRunStatusRequest,
+    UpdateRunStatusResponse,
+)
+from flwr.proto.serverappio_pb2 import (  # pylint: disable=E0611
+    GetNodesRequest,
+    GetNodesResponse,
+    PullServerAppInputsRequest,
+    PullServerAppInputsResponse,
+    PullTaskResRequest,
+    PullTaskResResponse,
+    PushServerAppOutputsRequest,
+    PushServerAppOutputsResponse,
+    PushTaskInsRequest,
+    PushTaskInsResponse,
 )
 from flwr.proto.task_pb2 import TaskRes  # pylint: disable=E0611
 from flwr.server.superlink.ffs.ffs import Ffs
 from flwr.server.superlink.ffs.ffs_factory import FfsFactory
-from flwr.server.superlink.state import State, StateFactory
+from flwr.server.superlink.linkstate import LinkState, LinkStateFactory
 from flwr.server.utils.validator import validate_task_ins_or_res
 
 
-class DriverServicer(driver_pb2_grpc.DriverServicer):
-    """Driver API servicer."""
+class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
+    """ServerAppIo API servicer."""
 
-    def __init__(self, state_factory: StateFactory, ffs_factory: FfsFactory) -> None:
+    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 GetNodes(
         self, request: GetNodesRequest, context: grpc.ServicerContext
     ) -> GetNodesResponse:
         """Get available nodes."""
-        log(DEBUG, "DriverServicer.GetNodes")
-        state: State = self.state_factory.state()
+        log(DEBUG, "ServerAppIoServicer.GetNodes")
+        state: LinkState = self.state_factory.state()
         all_ids: set[int] = state.get_nodes(request.run_id)
         nodes: list[Node] = [
             Node(node_id=node_id, anonymous=False) for node_id in all_ids
@@ -78,8 +96,8 @@ class DriverServicer(driver_pb2_grpc.DriverServicer):
         self, request: CreateRunRequest, context: grpc.ServicerContext
     ) -> CreateRunResponse:
         """Create run ID."""
-        log(DEBUG, "DriverServicer.CreateRun")
-        state: State = self.state_factory.state()
+        log(DEBUG, "ServerAppIoServicer.CreateRun")
+        state: LinkState = self.state_factory.state()
         if request.HasField("fab"):
             fab = fab_from_proto(request.fab)
             ffs: Ffs = self.ffs_factory.ffs()
@@ -95,6 +113,7 @@ class DriverServicer(driver_pb2_grpc.DriverServicer):
             request.fab_version,
             fab_hash,
             user_config_from_proto(request.override_config),
+            ConfigsRecord(),
         )
         return CreateRunResponse(run_id=run_id)
 
@@ -102,7 +121,7 @@ class DriverServicer(driver_pb2_grpc.DriverServicer):
         self, request: PushTaskInsRequest, context: grpc.ServicerContext
     ) -> PushTaskInsResponse:
         """Push a set of TaskIns."""
-        log(DEBUG, "DriverServicer.PushTaskIns")
+        log(DEBUG, "ServerAppIoServicer.PushTaskIns")
 
         # Set pushed_at (timestamp in seconds)
         pushed_at = time.time()
@@ -116,7 +135,7 @@ class DriverServicer(driver_pb2_grpc.DriverServicer):
             _raise_if(bool(validation_errors), ", ".join(validation_errors))
 
         # Init state
-        state: State = self.state_factory.state()
+        state: LinkState = self.state_factory.state()
 
         # Store each TaskIns
         task_ids: list[Optional[UUID]] = []
@@ -132,17 +151,20 @@ class DriverServicer(driver_pb2_grpc.DriverServicer):
         self, request: PullTaskResRequest, context: grpc.ServicerContext
     ) -> PullTaskResResponse:
         """Pull a set of TaskRes."""
-        log(DEBUG, "DriverServicer.PullTaskRes")
+        log(DEBUG, "ServerAppIoServicer.PullTaskRes")
 
         # Convert each task_id str to UUID
         task_ids: set[UUID] = {UUID(task_id) for task_id in request.task_ids}
 
         # Init state
-        state: State = self.state_factory.state()
+        state: LinkState = self.state_factory.state()
 
         # Register callback
         def on_rpc_done() -> None:
-            log(DEBUG, "DriverServicer.PullTaskRes callback: delete TaskIns/TaskRes")
+            log(
+                DEBUG,
+                "ServerAppIoServicer.PullTaskRes callback: delete TaskIns/TaskRes",
+            )
 
             if context.is_active():
                 return
@@ -164,10 +186,10 @@ class DriverServicer(driver_pb2_grpc.DriverServicer):
         self, request: GetRunRequest, context: grpc.ServicerContext
     ) -> GetRunResponse:
         """Get run information."""
-        log(DEBUG, "DriverServicer.GetRun")
+        log(DEBUG, "ServerAppIoServicer.GetRun")
 
         # Init state
-        state: State = self.state_factory.state()
+        state: LinkState = self.state_factory.state()
 
         # Retrieve run information
         run = state.get_run(request.run_id)
@@ -175,21 +197,13 @@ class DriverServicer(driver_pb2_grpc.DriverServicer):
         if run is None:
             return GetRunResponse()
 
-        return GetRunResponse(
-            run=Run(
-                run_id=run.run_id,
-                fab_id=run.fab_id,
-                fab_version=run.fab_version,
-                override_config=user_config_to_proto(run.override_config),
-                fab_hash=run.fab_hash,
-            )
-        )
+        return GetRunResponse(run=run_to_proto(run))
 
     def GetFab(
         self, request: GetFabRequest, context: grpc.ServicerContext
     ) -> GetFabResponse:
         """Get FAB from Ffs."""
-        log(DEBUG, "DriverServicer.GetFab")
+        log(DEBUG, "ServerAppIoServicer.GetFab")
 
         ffs: Ffs = self.ffs_factory.ffs()
         if result := ffs.get(request.hash_str):
@@ -198,6 +212,78 @@ class DriverServicer(driver_pb2_grpc.DriverServicer):
 
         raise ValueError(f"Found no FAB with hash: {request.hash_str}")
 
+    def PullServerAppInputs(
+        self, request: PullServerAppInputsRequest, context: grpc.ServicerContext
+    ) -> PullServerAppInputsResponse:
+        """Pull ServerApp process inputs."""
+        log(DEBUG, "ServerAppIoServicer.PullServerAppInputs")
+        # 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 PullServerAppInputsResponse()
+
+            # 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 PullServerAppInputsResponse(
+                        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 PushServerAppOutputs(
+        self, request: PushServerAppOutputsRequest, context: grpc.ServicerContext
+    ) -> PushServerAppOutputsResponse:
+        """Push ServerApp process outputs."""
+        log(DEBUG, "ServerAppIoServicer.PushServerAppOutputs")
+        state = self.state_factory.state()
+        state.set_serverapp_context(request.run_id, context_from_proto(request.context))
+        return PushServerAppOutputsResponse()
+
+    def UpdateRunStatus(
+        self, request: UpdateRunStatusRequest, context: grpc.ServicerContext
+    ) -> UpdateRunStatusResponse:
+        """Update the status of a run."""
+        log(DEBUG, "ControlServicer.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, "ServerAppIoServicer.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 _raise_if(validation_error: bool, detail: str) -> None:
     if validation_error:

flwr/server/superlink/fleet/grpc_adapter/grpc_adapter_servicer.py

@@ -48,7 +48,7 @@ from flwr.proto.grpcadapter_pb2 import MessageContainer  # pylint: disable=E0611
 from flwr.proto.run_pb2 import GetRunRequest, GetRunResponse  # pylint: disable=E0611
 from flwr.server.superlink.ffs.ffs_factory import FfsFactory
 from flwr.server.superlink.fleet.message_handler import message_handler
-from flwr.server.superlink.state import StateFactory
+from flwr.server.superlink.linkstate import LinkStateFactory
 
 T = TypeVar("T", bound=GrpcMessage)
 
@@ -77,7 +77,9 @@ def _handle(
 class GrpcAdapterServicer(grpcadapter_pb2_grpc.GrpcAdapterServicer):
     """Fleet API via GrpcAdapter servicer."""
 
-    def __init__(self, state_factory: StateFactory, ffs_factory: FfsFactory) -> None:
+    def __init__(
+        self, state_factory: LinkStateFactory, ffs_factory: FfsFactory
+    ) -> None:
         self.state_factory = state_factory
         self.ffs_factory = ffs_factory
 

flwr/server/superlink/fleet/grpc_bidi/grpc_server.py

@@ -30,7 +30,7 @@ from flwr.proto.transport_pb2_grpc import (  # pylint: disable=E0611
     add_FlowerServiceServicer_to_server,
 )
 from flwr.server.client_manager import ClientManager
-from flwr.server.superlink.driver.driver_servicer import DriverServicer
+from flwr.server.superlink.driver.serverappio_servicer import ServerAppIoServicer
 from flwr.server.superlink.fleet.grpc_adapter.grpc_adapter_servicer import (
     GrpcAdapterServicer,
 )
@@ -161,7 +161,7 @@ def generic_create_grpc_server(  # pylint: disable=too-many-arguments,R0917
         tuple[FleetServicer, AddServicerToServerFn],
         tuple[GrpcAdapterServicer, AddServicerToServerFn],
         tuple[FlowerServiceServicer, AddServicerToServerFn],
-        tuple[DriverServicer, AddServicerToServerFn],
+        tuple[ServerAppIoServicer, AddServicerToServerFn],
     ],
     server_address: str,
     max_concurrent_workers: int = 1000,

flwr/server/superlink/fleet/grpc_rere/fleet_servicer.py

@@ -37,13 +37,15 @@ from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
 from flwr.proto.run_pb2 import GetRunRequest, GetRunResponse  # pylint: disable=E0611
 from flwr.server.superlink.ffs.ffs_factory import FfsFactory
 from flwr.server.superlink.fleet.message_handler import message_handler
-from flwr.server.superlink.state import StateFactory
+from flwr.server.superlink.linkstate import LinkStateFactory
 
 
 class FleetServicer(fleet_pb2_grpc.FleetServicer):
     """Fleet API servicer."""
 
-    def __init__(self, state_factory: StateFactory, ffs_factory: FfsFactory) -> None:
+    def __init__(
+        self, state_factory: LinkStateFactory, ffs_factory: FfsFactory
+    ) -> None:
         self.state_factory = state_factory
         self.ffs_factory = ffs_factory
 

flwr/server/superlink/fleet/grpc_rere/server_interceptor.py

@@ -45,7 +45,7 @@ from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
 )
 from flwr.proto.node_pb2 import Node  # pylint: disable=E0611
 from flwr.proto.run_pb2 import GetRunRequest, GetRunResponse  # pylint: disable=E0611
-from flwr.server.superlink.state import State
+from flwr.server.superlink.linkstate import LinkState
 
 _PUBLIC_KEY_HEADER = "public-key"
 _AUTH_TOKEN_HEADER = "auth-token"
@@ -84,7 +84,7 @@ def _get_value_from_tuples(
 class AuthenticateServerInterceptor(grpc.ServerInterceptor):  # type: ignore
     """Server interceptor for node authentication."""
 
-    def __init__(self, state: State):
+    def __init__(self, state: LinkState):
         self.state = state
 
         self.node_public_keys = state.get_node_public_keys()

flwr/server/superlink/fleet/message_handler/message_handler.py

@@ -43,12 +43,12 @@ from flwr.proto.run_pb2 import (  # pylint: disable=E0611
 )
 from flwr.proto.task_pb2 import TaskIns, TaskRes  # pylint: disable=E0611
 from flwr.server.superlink.ffs.ffs import Ffs
-from flwr.server.superlink.state import State
+from flwr.server.superlink.linkstate import LinkState
 
 
 def create_node(
     request: CreateNodeRequest,  # pylint: disable=unused-argument
-    state: State,
+    state: LinkState,
 ) -> CreateNodeResponse:
     """."""
     # Create node
@@ -56,7 +56,7 @@ def create_node(
     return CreateNodeResponse(node=Node(node_id=node_id, anonymous=False))
 
 
-def delete_node(request: DeleteNodeRequest, state: State) -> DeleteNodeResponse:
+def delete_node(request: DeleteNodeRequest, state: LinkState) -> DeleteNodeResponse:
     """."""
     # Validate node_id
     if request.node.anonymous or request.node.node_id == 0:
@@ -69,14 +69,14 @@ def delete_node(request: DeleteNodeRequest, state: State) -> DeleteNodeResponse:
 
 def ping(
     request: PingRequest,  # pylint: disable=unused-argument
-    state: State,  # pylint: disable=unused-argument
+    state: LinkState,  # pylint: disable=unused-argument
 ) -> PingResponse:
     """."""
     res = state.acknowledge_ping(request.node.node_id, request.ping_interval)
     return PingResponse(success=res)
 
 
-def pull_task_ins(request: PullTaskInsRequest, state: State) -> PullTaskInsResponse:
+def pull_task_ins(request: PullTaskInsRequest, state: LinkState) -> PullTaskInsResponse:
     """Pull TaskIns handler."""
     # Get node_id if client node is not anonymous
     node = request.node  # pylint: disable=no-member
@@ -92,7 +92,7 @@ def pull_task_ins(request: PullTaskInsRequest, state: State) -> PullTaskInsRespo
     return response
 
 
-def push_task_res(request: PushTaskResRequest, state: State) -> PushTaskResResponse:
+def push_task_res(request: PushTaskResRequest, state: LinkState) -> PushTaskResResponse:
     """Push TaskRes handler."""
     # pylint: disable=no-member
     task_res: TaskRes = request.task_res_list[0]
@@ -113,7 +113,7 @@ def push_task_res(request: PushTaskResRequest, state: State) -> PushTaskResRespo
 
 
 def get_run(
-    request: GetRunRequest, state: State  # pylint: disable=W0613
+    request: GetRunRequest, state: LinkState  # pylint: disable=W0613
 ) -> GetRunResponse:
     """Get run information."""
     run = state.get_run(request.run_id)

flwr/server/superlink/fleet/rest_rere/rest_api.py

@@ -40,7 +40,7 @@ from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
 from flwr.proto.run_pb2 import GetRunRequest, GetRunResponse  # pylint: disable=E0611
 from flwr.server.superlink.ffs.ffs import Ffs
 from flwr.server.superlink.fleet.message_handler import message_handler
-from flwr.server.superlink.state import State
+from flwr.server.superlink.linkstate import LinkState
 
 try:
     from starlette.applications import Starlette
@@ -90,7 +90,7 @@ def rest_request_response(
 async def create_node(request: CreateNodeRequest) -> CreateNodeResponse:
     """Create Node."""
     # Get state from app
-    state: State = app.state.STATE_FACTORY.state()
+    state: LinkState = app.state.STATE_FACTORY.state()
 
     # Handle message
     return message_handler.create_node(request=request, state=state)
@@ -100,7 +100,7 @@ async def create_node(request: CreateNodeRequest) -> CreateNodeResponse:
 async def delete_node(request: DeleteNodeRequest) -> DeleteNodeResponse:
     """Delete Node Id."""
     # Get state from app
-    state: State = app.state.STATE_FACTORY.state()
+    state: LinkState = app.state.STATE_FACTORY.state()
 
     # Handle message
     return message_handler.delete_node(request=request, state=state)
@@ -110,7 +110,7 @@ async def delete_node(request: DeleteNodeRequest) -> DeleteNodeResponse:
 async def pull_task_ins(request: PullTaskInsRequest) -> PullTaskInsResponse:
     """Pull TaskIns."""
     # Get state from app
-    state: State = app.state.STATE_FACTORY.state()
+    state: LinkState = app.state.STATE_FACTORY.state()
 
     # Handle message
     return message_handler.pull_task_ins(request=request, state=state)
@@ -121,7 +121,7 @@ async def pull_task_ins(request: PullTaskInsRequest) -> PullTaskInsResponse:
 async def push_task_res(request: PushTaskResRequest) -> PushTaskResResponse:
     """Push TaskRes."""
     # Get state from app
-    state: State = app.state.STATE_FACTORY.state()
+    state: LinkState = app.state.STATE_FACTORY.state()
 
     # Handle message
     return message_handler.push_task_res(request=request, state=state)
@@ -131,7 +131,7 @@ async def push_task_res(request: PushTaskResRequest) -> PushTaskResResponse:
 async def ping(request: PingRequest) -> PingResponse:
     """Ping."""
     # Get state from app
-    state: State = app.state.STATE_FACTORY.state()
+    state: LinkState = app.state.STATE_FACTORY.state()
 
     # Handle message
     return message_handler.ping(request=request, state=state)
@@ -141,7 +141,7 @@ async def ping(request: PingRequest) -> PingResponse:
 async def get_run(request: GetRunRequest) -> GetRunResponse:
     """GetRun."""
     # Get state from app
-    state: State = app.state.STATE_FACTORY.state()
+    state: LinkState = app.state.STATE_FACTORY.state()
 
     # Handle message
     return message_handler.get_run(request=request, state=state)