flwr-nightly 1.13.0.dev20241028__py3-none-any.whl → 1.13.0.dev20241030__py3-none-any.whl

flwr/proto/log_pb2_grpc.pyi

@@ -0,0 +1,4 @@
+"""
+@generated by mypy-protobuf.  Do not edit manually!
+isort:skip_file
+"""

flwr/server/app.py

@@ -382,9 +382,7 @@ def _flwr_serverapp_scheduler(
 
             log(
                 INFO,
-                "Launching `flwr-serverapp` subprocess with run-id %d. "
-                "Connects to SuperLink on %s",
-                pending_run_id,
+                "Launching `flwr-serverapp` subprocess. Connects to SuperLink on %s",
                 driver_api_address,
             )
             # Start ServerApp subprocess
@@ -392,8 +390,6 @@ def _flwr_serverapp_scheduler(
                 "flwr-serverapp",
                 "--superlink",
                 driver_api_address,
-                "--run-id",
-                str(pending_run_id),
             ]
             if ssl_ca_certfile:
                 command.append("--root-certificates")

flwr/server/driver/driver.py

@@ -26,6 +26,20 @@ from flwr.common.typing import Run
 class Driver(ABC):
     """Abstract base Driver class for the Driver API."""
 
+    @abstractmethod
+    def init_run(self, run_id: int) -> None:
+        """Request a run to the SuperLink with a given `run_id`.
+
+        If a Run with the specified `run_id` exists, a local Run
+        object will be created. It enables further functionality
+        in the driver, such as sending `Messages`.
+
+        Parameters
+        ----------
+        run_id : int
+            The `run_id` of the Run this Driver object operates in.
+        """
+
     @property
     @abstractmethod
     def run(self) -> Run:

flwr/server/driver/grpc_driver.py

@@ -60,8 +60,6 @@ class GrpcDriver(Driver):
 
     Parameters
     ----------
-    run_id : int
-        The identifier of the run.
     driver_service_address : str (default: "[::]:9091")
         The address (URL, IPv6, IPv4) of the SuperLink Driver API service.
     root_certificates : Optional[bytes] (default: None)
@@ -72,11 +70,9 @@ class GrpcDriver(Driver):
 
     def __init__(  # pylint: disable=too-many-arguments
         self,
-        run_id: int,
         driver_service_address: str = DRIVER_API_DEFAULT_ADDRESS,
         root_certificates: Optional[bytes] = None,
     ) -> None:
-        self._run_id = run_id
         self._addr = driver_service_address
         self._cert = root_certificates
         self._run: Optional[Run] = None
@@ -116,15 +112,17 @@ class GrpcDriver(Driver):
         channel.close()
         log(DEBUG, "[Driver] Disconnected")
 
-    def _init_run(self) -> None:
+    def init_run(self, run_id: int) -> None:
+        """Initialize the run."""
         # Check if is initialized
         if self._run is not None:
             return
+
         # Get the run info
-        req = GetRunRequest(run_id=self._run_id)
+        req = GetRunRequest(run_id=run_id)
         res: GetRunResponse = self._stub.GetRun(req)
         if not res.HasField("run"):
-            raise RuntimeError(f"Cannot find the run with ID: {self._run_id}")
+            raise RuntimeError(f"Cannot find the run with ID: {run_id}")
         self._run = Run(
             run_id=res.run.run_id,
             fab_id=res.run.fab_id,
@@ -136,7 +134,6 @@ class GrpcDriver(Driver):
     @property
     def run(self) -> Run:
         """Run information."""
-        self._init_run()
         return Run(**vars(self._run))
 
     @property
@@ -150,7 +147,7 @@ class GrpcDriver(Driver):
         # Check if the message is valid
         if not (
             # Assume self._run being initialized
-            message.metadata.run_id == self._run_id
+            message.metadata.run_id == cast(Run, self._run).run_id
             and message.metadata.src_node_id == self.node.node_id
             and message.metadata.message_id == ""
             and message.metadata.reply_to_message == ""
@@ -171,7 +168,6 @@ class GrpcDriver(Driver):
         This method constructs a new `Message` with given content and metadata.
         The `run_id` and `src_node_id` will be set automatically.
         """
-        self._init_run()
         if ttl:
             warnings.warn(
                 "A custom TTL was set, but note that the SuperLink does not enforce "
@@ -182,7 +178,7 @@ class GrpcDriver(Driver):
 
         ttl_ = DEFAULT_TTL if ttl is None else ttl
         metadata = Metadata(
-            run_id=self._run_id,
+            run_id=cast(Run, self._run).run_id,
             message_id="",  # Will be set by the server
             src_node_id=self.node.node_id,
             dst_node_id=dst_node_id,
@@ -195,10 +191,9 @@ class GrpcDriver(Driver):
 
     def get_node_ids(self) -> list[int]:
         """Get node IDs."""
-        self._init_run()
         # Call GrpcDriverStub method
         res: GetNodesResponse = self._stub.GetNodes(
-            GetNodesRequest(run_id=self._run_id)
+            GetNodesRequest(run_id=cast(Run, self._run).run_id)
         )
         return [node.node_id for node in res.nodes]
 
@@ -208,7 +203,6 @@ class GrpcDriver(Driver):
         This method takes an iterable of messages and sends each message
         to the node specified in `dst_node_id`.
         """
-        self._init_run()
         # Construct TaskIns
         task_ins_list: list[TaskIns] = []
         for msg in messages:
@@ -230,7 +224,6 @@ class GrpcDriver(Driver):
         This method is used to collect messages from the SuperLink that correspond to a
         set of given message IDs.
         """
-        self._init_run()
         # Pull TaskRes
         res: PullTaskResResponse = self._stub.PullTaskRes(
             PullTaskResRequest(node=self.node, task_ids=message_ids)

flwr/server/driver/inmemory_driver.py

@@ -35,8 +35,6 @@ class InMemoryDriver(Driver):
 
     Parameters
     ----------
-    run_id : int
-        The identifier of the run.
     state_factory : StateFactory
         A StateFactory embedding a state that this driver can interface with.
     pull_interval : float (default=0.1)
@@ -45,18 +43,15 @@ class InMemoryDriver(Driver):
 
     def __init__(
         self,
-        run_id: int,
         state_factory: LinkStateFactory,
         pull_interval: float = 0.1,
     ) -> None:
-        self._run_id = run_id
         self._run: Optional[Run] = None
         self.state = state_factory.state()
         self.pull_interval = pull_interval
         self.node = Node(node_id=0, anonymous=True)
 
     def _check_message(self, message: Message) -> None:
-        self._init_run()
         # Check if the message is valid
         if not (
             message.metadata.run_id == cast(Run, self._run).run_id
@@ -67,19 +62,18 @@ class InMemoryDriver(Driver):
         ):
             raise ValueError(f"Invalid message: {message}")
 
-    def _init_run(self) -> None:
+    def init_run(self, run_id: int) -> None:
         """Initialize the run."""
         if self._run is not None:
             return
-        run = self.state.get_run(self._run_id)
+        run = self.state.get_run(run_id)
         if run is None:
-            raise RuntimeError(f"Cannot find the run with ID: {self._run_id}")
+            raise RuntimeError(f"Cannot find the run with ID: {run_id}")
         self._run = run
 
     @property
     def run(self) -> Run:
         """Run ID."""
-        self._init_run()
         return Run(**vars(cast(Run, self._run)))
 
     def create_message(  # pylint: disable=too-many-arguments,R0917
@@ -95,7 +89,6 @@ class InMemoryDriver(Driver):
         This method constructs a new `Message` with given content and metadata.
         The `run_id` and `src_node_id` will be set automatically.
         """
-        self._init_run()
         if ttl:
             warnings.warn(
                 "A custom TTL was set, but note that the SuperLink does not enforce "
@@ -119,7 +112,6 @@ class InMemoryDriver(Driver):
 
     def get_node_ids(self) -> list[int]:
         """Get node IDs."""
-        self._init_run()
         return list(self.state.get_nodes(cast(Run, self._run).run_id))
 
     def push_messages(self, messages: Iterable[Message]) -> Iterable[str]:

flwr/server/run_serverapp.py

@@ -171,11 +171,11 @@ def run_server_app() -> None:
     if app_path is None:
         # User provided `--run-id`, but not `app_dir`
         driver = GrpcDriver(
-            run_id=args.run_id,
             driver_service_address=args.superlink,
             root_certificates=root_certificates,
         )
         flwr_dir = get_flwr_dir(args.flwr_dir)
+        driver.init_run(args.run_id)
         run_ = driver.run
         if not run_.fab_hash:
             raise ValueError("FAB hash not provided.")
@@ -193,7 +193,6 @@ def run_server_app() -> None:
         # User provided `app_dir`, but not `--run-id`
         # Create run if run_id is not provided
         driver = GrpcDriver(
-            run_id=0,  # Will be overwritten
             driver_service_address=args.superlink,
             root_certificates=root_certificates,
         )
@@ -204,8 +203,8 @@ def run_server_app() -> None:
         # Create run
         req = CreateRunRequest(fab_id=fab_id, fab_version=fab_version)
         res: CreateRunResponse = driver._stub.CreateRun(req)  # pylint: disable=W0212
-        # Overwrite driver._run_id
-        driver._run_id = res.run_id  # pylint: disable=W0212
+        # Fetch full `Run` using `run_id`
+        driver.init_run(res.run_id)  # pylint: disable=W0212
 
     # Obtain server app reference and the run config
     server_app_attr = config["tool"]["flwr"]["app"]["components"]["serverapp"]

flwr/server/serverapp/app.py

@@ -16,17 +16,53 @@
 
 import argparse
 import sys
-from logging import DEBUG, INFO, WARN
+from logging import DEBUG, ERROR, INFO, WARN
 from os.path import isfile
 from pathlib import Path
+from queue import Queue
+from time import sleep
 from typing import Optional
 
-from flwr.common.logger import log
+from flwr.cli.config_utils import get_fab_metadata
+from flwr.cli.install import install_from_fab
+from flwr.common.config import (
+    get_flwr_dir,
+    get_fused_config_from_dir,
+    get_project_config,
+    get_project_dir,
+)
+from flwr.common.constant import Status, SubStatus
+from flwr.common.logger import (
+    log,
+    mirror_output_to_queue,
+    restore_output,
+    start_log_uploader,
+    stop_log_uploader,
+)
+from flwr.common.serde import (
+    context_from_proto,
+    context_to_proto,
+    fab_from_proto,
+    run_from_proto,
+    run_status_to_proto,
+)
+from flwr.common.typing import RunStatus
+from flwr.proto.driver_pb2 import (  # pylint: disable=E0611
+    PullServerAppInputsRequest,
+    PullServerAppInputsResponse,
+    PushServerAppOutputsRequest,
+)
+from flwr.proto.run_pb2 import UpdateRunStatusRequest  # pylint: disable=E0611
 from flwr.server.driver.grpc_driver import GrpcDriver
+from flwr.server.run_serverapp import run as run_
 
 
 def flwr_serverapp() -> None:
     """Run process-isolated Flower ServerApp."""
+    # Capture stdout/stderr
+    log_queue: Queue[Optional[str]] = Queue()
+    mirror_output_to_queue(log_queue)
+
     log(INFO, "Starting Flower ServerApp")
 
     parser = argparse.ArgumentParser(
@@ -38,11 +74,10 @@ def flwr_serverapp() -> None:
         help="Address of SuperLink's DriverAPI",
     )
     parser.add_argument(
-        "--run-id",
-        type=int,
-        required=False,
-        help="Id of the Run this process should start. If not supplied, this "
-        "function will request a pending run to the LinkState.",
+        "--run-once",
+        action="store_true",
+        help="When set, this process will start a single ServerApp "
+        "for a pending Run. If no pending run the process will exit. ",
     )
     parser.add_argument(
         "--flwr-dir",
@@ -75,18 +110,20 @@ def flwr_serverapp() -> None:
 
     log(
         DEBUG,
-        "Staring isolated `ServerApp` connected to SuperLink DriverAPI at %s "
-        "for run-id %s",
+        "Staring isolated `ServerApp` connected to SuperLink DriverAPI at %s",
         args.superlink,
-        args.run_id,
     )
     run_serverapp(
         superlink=args.superlink,
-        run_id=args.run_id,
+        log_queue=log_queue,
+        run_once=args.run_once,
         flwr_dir_=args.flwr_dir,
         certificates=certificates,
     )
 
+    # Restore stdout/stderr
+    restore_output()
+
 
 def _try_obtain_certificates(
     args: argparse.Namespace,
@@ -121,21 +158,114 @@ def _try_obtain_certificates(
     return root_certificates
 
 
-def run_serverapp(  # pylint: disable=R0914
+def run_serverapp(  # pylint: disable=R0914, disable=W0212
     superlink: str,
-    run_id: Optional[int] = None,
+    log_queue: Queue[Optional[str]],
+    run_once: bool,
     flwr_dir_: Optional[str] = None,
     certificates: Optional[bytes] = None,
 ) -> None:
     """Run Flower ServerApp process."""
-    _ = GrpcDriver(
-        run_id=run_id if run_id else 0,
+    driver = GrpcDriver(
         driver_service_address=superlink,
         root_certificates=certificates,
     )
 
-    log(INFO, "%s", flwr_dir_)
+    # Resolve directory where FABs are installed
+    flwr_dir = get_flwr_dir(flwr_dir_)
+    log_uploader = None
+
+    while True:
+
+        try:
+            # Pull ServerAppInputs from LinkState
+            req = PullServerAppInputsRequest()
+            res: PullServerAppInputsResponse = driver._stub.PullServerAppInputs(req)
+            if not res.HasField("run"):
+                sleep(3)
+                run_status = None
+                continue
+
+            context = context_from_proto(res.context)
+            run = run_from_proto(res.run)
+            fab = fab_from_proto(res.fab)
+
+            driver.init_run(run.run_id)
+
+            # Start log uploader for this run
+            log_uploader = start_log_uploader(
+                log_queue=log_queue,
+                node_id=0,
+                run_id=run.run_id,
+                stub=driver._stub,
+            )
+
+            log(DEBUG, "ServerApp process starts FAB installation.")
+            install_from_fab(fab.content, flwr_dir=flwr_dir, skip_prompt=True)
+
+            fab_id, fab_version = get_fab_metadata(fab.content)
+
+            app_path = str(get_project_dir(fab_id, fab_version, fab.hash_str, flwr_dir))
+            config = get_project_config(app_path)
+
+            # Obtain server app reference and the run config
+            server_app_attr = config["tool"]["flwr"]["app"]["components"]["serverapp"]
+            server_app_run_config = get_fused_config_from_dir(
+                Path(app_path), run.override_config
+            )
+
+            # Update run_config in context
+            context.run_config = server_app_run_config
+
+            log(
+                DEBUG,
+                "Flower will load ServerApp `%s` in %s",
+                server_app_attr,
+                app_path,
+            )
+
+            # Change status to Running
+            run_status_proto = run_status_to_proto(RunStatus(Status.RUNNING, "", ""))
+            driver._stub.UpdateRunStatus(
+                UpdateRunStatusRequest(run_id=run.run_id, run_status=run_status_proto)
+            )
+
+            # Load and run the ServerApp with the Driver
+            updated_context = run_(
+                driver=driver,
+                server_app_dir=app_path,
+                server_app_attr=server_app_attr,
+                context=context,
+            )
+
+            # Send resulting context
+            context_proto = context_to_proto(updated_context)
+            out_req = PushServerAppOutputsRequest(
+                run_id=run.run_id, context=context_proto
+            )
+            _ = driver._stub.PushServerAppOutputs(out_req)
+
+            run_status = RunStatus(Status.FINISHED, SubStatus.COMPLETED, "")
+
+        except Exception as ex:  # pylint: disable=broad-exception-caught
+            exc_entity = "ServerApp"
+            log(ERROR, "%s raised an exception", exc_entity, exc_info=ex)
+            run_status = RunStatus(Status.FINISHED, SubStatus.FAILED, str(ex))
+
+        finally:
+            if run_status:
+                run_status_proto = run_status_to_proto(run_status)
+                driver._stub.UpdateRunStatus(
+                    UpdateRunStatusRequest(
+                        run_id=run.run_id, run_status=run_status_proto
+                    )
+                )
 
-    # Then, GetServerInputs
+            # Stop log uploader for this run
+            if log_uploader:
+                stop_log_uploader(log_queue, log_uploader)
+                log_uploader = None
 
-    # Then, run ServerApp
+        # Stop the loop if `flwr-serverapp` is expected to process a single run
+        if run_once:
+            break

flwr/server/superlink/driver/driver_servicer.py

@@ -31,6 +31,7 @@ from flwr.common.serde import (
     context_to_proto,
     fab_from_proto,
     fab_to_proto,
+    run_status_from_proto,
     run_to_proto,
     user_config_from_proto,
 )
@@ -49,12 +50,18 @@ from flwr.proto.driver_pb2 import (  # pylint: disable=E0611
     PushTaskInsResponse,
 )
 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,
+    UpdateRunStatusRequest,
+    UpdateRunStatusResponse,
 )
 from flwr.proto.task_pb2 import TaskRes  # pylint: disable=E0611
 from flwr.server.superlink.ffs.ffs import Ffs
@@ -212,11 +219,8 @@ class DriverServicer(driver_pb2_grpc.DriverServicer):
 
         # Lock access to LinkState, preventing obtaining the same pending run_id
         with self.lock:
-            # If run_id is provided, use it, otherwise use the pending run_id
-            if _has_field(request, "run_id"):
-                run_id: Optional[int] = request.run_id
-            else:
-                run_id = state.get_pending_run_id()
+            # 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()
@@ -228,14 +232,12 @@ class DriverServicer(driver_pb2_grpc.DriverServicer):
             if run and run.fab_hash:
                 if result := ffs.get(run.fab_hash):
                     fab = Fab(run.fab_hash, result[0])
-            if run and fab:
+            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) if serverapp_ctxt else None
-                        ),
+                        context=context_to_proto(serverapp_ctxt),
                         run=run_to_proto(run),
                         fab=fab_to_proto(fab),
                     )
@@ -253,6 +255,31 @@ class DriverServicer(driver_pb2_grpc.DriverServicer):
         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, "DriverServicer.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/linkstate/in_memory_linkstate.py

@@ -17,7 +17,8 @@
 
 import threading
 import time
-from dataclasses import dataclass
+from bisect import bisect_right
+from dataclasses import dataclass, field
 from logging import ERROR, WARNING
 from typing import Optional
 from uuid import UUID, uuid4
@@ -43,7 +44,7 @@ from .utils import (
 
 
 @dataclass
-class RunRecord:
+class RunRecord:  # pylint: disable=R0902
     """The record of a specific run, including its status and timestamps."""
 
     run: Run
@@ -52,6 +53,8 @@ class RunRecord:
     starting_at: str = ""
     running_at: str = ""
     finished_at: str = ""
+    logs: list[tuple[float, str]] = field(default_factory=list)
+    log_lock: threading.Lock = field(default_factory=threading.Lock)
 
 
 class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
@@ -511,3 +514,26 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         if run_id not in self.run_ids:
             raise ValueError(f"Run {run_id} not found")
         self.contexts[run_id] = context
+
+    def add_serverapp_log(self, run_id: int, log_message: str) -> None:
+        """Add a log entry to the serverapp logs for the specified `run_id`."""
+        if run_id not in self.run_ids:
+            raise ValueError(f"Run {run_id} not found")
+        run = self.run_ids[run_id]
+        with run.log_lock:
+            run.logs.append((now().timestamp(), log_message))
+
+    def get_serverapp_log(
+        self, run_id: int, after_timestamp: Optional[float]
+    ) -> tuple[str, float]:
+        """Get the serverapp logs for the specified `run_id`."""
+        if run_id not in self.run_ids:
+            raise ValueError(f"Run {run_id} not found")
+        run = self.run_ids[run_id]
+        if after_timestamp is None:
+            after_timestamp = 0.0
+        with run.log_lock:
+            # Find the index where the timestamp would be inserted
+            index = bisect_right(run.logs, (after_timestamp, ""))
+            latest_timestamp = run.logs[-1][0] if index < len(run.logs) else 0.0
+            return "".join(log for _, log in run.logs[index:]), latest_timestamp

flwr/server/superlink/linkstate/linkstate.py

@@ -299,3 +299,38 @@ class LinkState(abc.ABC):  # pylint: disable=R0904
         context : Context
             The context to be associated with the specified `run_id`.
         """
+
+    @abc.abstractmethod
+    def add_serverapp_log(self, run_id: int, log_message: str) -> None:
+        """Add a log entry to the ServerApp logs for the specified `run_id`.
+
+        Parameters
+        ----------
+        run_id : int
+            The identifier of the run for which to add a log entry.
+        log_message : str
+            The log entry to be added to the ServerApp logs.
+        """
+
+    @abc.abstractmethod
+    def get_serverapp_log(
+        self, run_id: int, after_timestamp: Optional[float]
+    ) -> tuple[str, float]:
+        """Get the ServerApp logs for the specified `run_id`.
+
+        Parameters
+        ----------
+        run_id : int
+            The identifier of the run for which to retrieve the ServerApp logs.
+
+        after_timestamp : Optional[float]
+            Retrieve logs after this timestamp. If set to `None`, retrieve all logs.
+
+        Returns
+        -------
+        tuple[str, float]
+            A tuple containing:
+            - The ServerApp logs associated with the specified `run_id`.
+            - The timestamp of the latest log entry in the returned logs.
+              Returns `0` if no logs are returned.
+        """