flwr 1.13.1__py3-none-any.whl → 1.14.0__py3-none-any.whl

flwr/server/superlink/linkstate/linkstate.py

@@ -139,8 +139,19 @@ class LinkState(abc.ABC):  # pylint: disable=R0904
         """
 
     @abc.abstractmethod
-    def delete_tasks(self, task_ids: set[UUID]) -> None:
-        """Delete all delivered TaskIns/TaskRes pairs."""
+    def delete_tasks(self, task_ins_ids: set[UUID]) -> None:
+        """Delete TaskIns/TaskRes pairs based on provided TaskIns IDs.
+
+        Parameters
+        ----------
+        task_ins_ids : set[UUID]
+            A set of TaskIns IDs. For each ID in the set, the corresponding
+            TaskIns and its associated TaskRes will be deleted.
+        """
+
+    @abc.abstractmethod
+    def get_task_ids_from_run_id(self, run_id: int) -> set[UUID]:
+        """Get all TaskIns IDs for the given run_id."""
 
     @abc.abstractmethod
     def create_node(
@@ -273,6 +284,10 @@ class LinkState(abc.ABC):  # pylint: disable=R0904
     def get_server_public_key(self) -> Optional[bytes]:
         """Retrieve `server_public_key` in urlsafe bytes."""
 
+    @abc.abstractmethod
+    def clear_supernode_auth_keys_and_credentials(self) -> None:
+        """Clear stored `node_public_keys` and credentials in the link state if any."""
+
     @abc.abstractmethod
     def store_node_public_keys(self, public_keys: set[bytes]) -> None:
         """Store a set of `node_public_keys` in the link state."""

flwr/server/superlink/linkstate/sqlite_linkstate.py

@@ -14,12 +14,12 @@
 # ==============================================================================
 """SQLite based implemenation of the link state."""
 
+
 # pylint: disable=too-many-lines
 
 import json
 import re
 import sqlite3
-import threading
 import time
 from collections.abc import Sequence
 from logging import DEBUG, ERROR, WARNING
@@ -183,7 +183,6 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         """
         self.database_path = database_path
         self.conn: Optional[sqlite3.Connection] = None
-        self.lock = threading.RLock()
 
     def initialize(self, log_queries: bool = False) -> list[tuple[str]]:
         """Create tables if they don't exist yet.
@@ -216,7 +215,6 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         cur.execute(SQL_CREATE_TABLE_PUBLIC_KEY)
         cur.execute(SQL_CREATE_INDEX_ONLINE_UNTIL)
         res = cur.execute("SELECT name FROM sqlite_schema;")
-
         return res.fetchall()
 
     def query(
@@ -569,9 +567,6 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         data: list[Any] = [delivered_at] + task_res_ids
         self.query(query, data)
 
-        # Cleanup
-        self._force_delete_tasks_by_ids(set(ret.keys()))
-
         return list(ret.values())
 
     def num_task_ins(self) -> int:
@@ -595,68 +590,50 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         result: dict[str, int] = rows[0]
         return result["num"]
 
-    def delete_tasks(self, task_ids: set[UUID]) -> None:
-        """Delete all delivered TaskIns/TaskRes pairs."""
-        ids = list(task_ids)
-        if len(ids) == 0:
-            return None
+    def delete_tasks(self, task_ins_ids: set[UUID]) -> None:
+        """Delete TaskIns/TaskRes pairs based on provided TaskIns IDs."""
+        if not task_ins_ids:
+            return
+        if self.conn is None:
+            raise AttributeError("LinkState not initialized")
 
-        placeholders = ",".join([f":id_{index}" for index in range(len(task_ids))])
-        data = {f"id_{index}": str(task_id) for index, task_id in enumerate(task_ids)}
+        placeholders = ",".join(["?"] * len(task_ins_ids))
+        data = tuple(str(task_id) for task_id in task_ins_ids)
 
-        # 1. Query: Delete task_ins which have a delivered task_res
+        # Delete task_ins
         query_1 = f"""
             DELETE FROM task_ins
-            WHERE delivered_at != ''
-            AND task_id IN (
-                SELECT ancestry
-                FROM task_res
-                WHERE ancestry IN ({placeholders})
-                AND delivered_at != ''
-            );
+            WHERE task_id IN ({placeholders});
         """
 
-        # 2. Query: Delete delivered task_res to be run after 1. Query
+        # Delete task_res
         query_2 = f"""
             DELETE FROM task_res
-            WHERE ancestry IN ({placeholders})
-            AND delivered_at != '';
+            WHERE ancestry IN ({placeholders});
         """
 
-        if self.conn is None:
-            raise AttributeError("LinkState not intitialized")
-
         with self.conn:
             self.conn.execute(query_1, data)
             self.conn.execute(query_2, data)
 
-        return None
-
-    def _force_delete_tasks_by_ids(self, task_ids: set[UUID]) -> None:
-        """Delete tasks based on a set of TaskIns IDs."""
-        if not task_ids:
-            return
+    def get_task_ids_from_run_id(self, run_id: int) -> set[UUID]:
+        """Get all TaskIns IDs for the given run_id."""
         if self.conn is None:
             raise AttributeError("LinkState not initialized")
 
-        placeholders = ",".join([f":id_{index}" for index in range(len(task_ids))])
-        data = {f"id_{index}": str(task_id) for index, task_id in enumerate(task_ids)}
-
-        # Delete task_ins
-        query_1 = f"""
-            DELETE FROM task_ins
-            WHERE task_id IN ({placeholders});
+        query = """
+            SELECT task_id
+            FROM task_ins
+            WHERE run_id = :run_id;
         """
 
-        # Delete task_res
-        query_2 = f"""
-            DELETE FROM task_res
-            WHERE ancestry IN ({placeholders});
-        """
+        sint64_run_id = convert_uint64_to_sint64(run_id)
+        data = {"run_id": sint64_run_id}
 
         with self.conn:
-            self.conn.execute(query_1, data)
-            self.conn.execute(query_2, data)
+            rows = self.conn.execute(query, data).fetchall()
+
+        return {UUID(row["task_id"]) for row in rows}
 
     def create_node(
         self, ping_interval: float, public_key: Optional[bytes] = None
@@ -784,8 +761,6 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
                 "federation_options, pending_at, starting_at, running_at, finished_at, "
                 "sub_status, details) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?);"
             )
-            if fab_hash:
-                fab_id, fab_version = "", ""
             override_config_json = json.dumps(override_config)
             data = [
                 sint64_run_id,
@@ -843,6 +818,12 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
             public_key = None
         return public_key
 
+    def clear_supernode_auth_keys_and_credentials(self) -> None:
+        """Clear stored `node_public_keys` and credentials in the link state if any."""
+        queries = ["DELETE FROM public_key;", "DELETE FROM credential;"]
+        for query in queries:
+            self.query(query)
+
     def store_node_public_keys(self, public_keys: set[bytes]) -> None:
         """Store a set of `node_public_keys` in the link state."""
         query = "INSERT INTO public_key (public_key) VALUES (?)"

flwr/server/superlink/simulation/simulationio_servicer.py

@@ -14,6 +14,7 @@
 # ==============================================================================
 """SimulationIo API servicer."""
 
+
 import threading
 from logging import DEBUG, INFO
 
@@ -28,6 +29,7 @@ from flwr.common.serde import (
     context_to_proto,
     fab_to_proto,
     run_status_from_proto,
+    run_status_to_proto,
     run_to_proto,
 )
 from flwr.common.typing import Fab, RunStatus
@@ -39,6 +41,8 @@ from flwr.proto.log_pb2 import (  # pylint: disable=E0611
 from flwr.proto.run_pb2 import (  # pylint: disable=E0611
     GetFederationOptionsRequest,
     GetFederationOptionsResponse,
+    GetRunStatusRequest,
+    GetRunStatusResponse,
     UpdateRunStatusRequest,
     UpdateRunStatusResponse,
 )
@@ -50,6 +54,7 @@ from flwr.proto.simulationio_pb2 import (  # pylint: disable=E0611
 )
 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):
@@ -106,6 +111,15 @@ class SimulationIoServicer(simulationio_pb2_grpc.SimulationIoServicer):
         """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()
 
@@ -116,12 +130,31 @@ class SimulationIoServicer(simulationio_pb2_grpc.SimulationIoServicer):
         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:

flwr/server/superlink/utils.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.
+# ==============================================================================
+"""SuperLink utilities."""
+
+
+from typing import Union
+
+import grpc
+
+from flwr.common.constant import Status, SubStatus
+from flwr.common.typing import RunStatus
+from flwr.server.superlink.linkstate import LinkState
+
+_STATUS_TO_MSG = {
+    Status.PENDING: "Run is pending.",
+    Status.STARTING: "Run is starting.",
+    Status.RUNNING: "Run is running.",
+    Status.FINISHED: "Run is finished.",
+}
+
+
+def check_abort(
+    run_id: int,
+    abort_status_list: list[str],
+    state: LinkState,
+) -> Union[str, None]:
+    """Check if the status of the provided `run_id` is in `abort_status_list`."""
+    run_status: RunStatus = state.get_run_status({run_id})[run_id]
+
+    if run_status.status in abort_status_list:
+        msg = _STATUS_TO_MSG[run_status.status]
+        if run_status.sub_status == SubStatus.STOPPED:
+            msg += " Stopped by user."
+        return msg
+
+    return None
+
+
+def abort_grpc_context(msg: Union[str, None], context: grpc.ServicerContext) -> None:
+    """Abort context with statuscode PERMISSION_DENIED if `msg` is not None."""
+    if msg is not None:
+        context.abort(grpc.StatusCode.PERMISSION_DENIED, msg)
+
+
+def abort_if(
+    run_id: int,
+    abort_status_list: list[str],
+    state: LinkState,
+    context: grpc.ServicerContext,
+) -> None:
+    """Abort context if status of the provided `run_id` is in `abort_status_list`."""
+    msg = check_abort(run_id, abort_status_list, state)
+    abort_grpc_context(msg, context)

flwr/simulation/app.py

@@ -14,6 +14,7 @@
 # ==============================================================================
 """Flower Simulation process."""
 
+
 import argparse
 import sys
 from logging import DEBUG, ERROR, INFO
@@ -23,7 +24,8 @@ from typing import Optional
 
 from flwr.cli.config_utils import get_fab_metadata
 from flwr.cli.install import install_from_fab
-from flwr.common import EventType
+from flwr.cli.utils import get_sha256_hash
+from flwr.common import EventType, event
 from flwr.common.args import add_args_flwr_app_common
 from flwr.common.config import (
     get_flwr_dir,
@@ -47,6 +49,7 @@ from flwr.common.logger import (
 from flwr.common.serde import (
     configs_record_from_proto,
     context_from_proto,
+    context_to_proto,
     fab_from_proto,
     run_from_proto,
     run_status_to_proto,
@@ -200,8 +203,17 @@ def run_simulation_process(  # pylint: disable=R0914, disable=W0212, disable=R09
             verbose: bool = fed_opt.get("verbose", False)
             enable_tf_gpu_growth: bool = fed_opt.get("enable_tf_gpu_growth", False)
 
+            event(
+                EventType.FLWR_SIMULATION_RUN_ENTER,
+                event_details={
+                    "backend": "ray",
+                    "num-supernodes": num_supernodes,
+                    "run-id-hash": get_sha256_hash(run.run_id),
+                },
+            )
+
             # Launch the simulation
-            _run_simulation(
+            updated_context = _run_simulation(
                 server_app_attr=server_app_attr,
                 client_app_attr=client_app_attr,
                 num_supernodes=num_supernodes,
@@ -212,11 +224,11 @@ def run_simulation_process(  # pylint: disable=R0914, disable=W0212, disable=R09
                 verbose_logging=verbose,
                 server_app_run_config=fused_config,
                 is_app=True,
-                exit_event=EventType.CLI_FLOWER_SIMULATION_LEAVE,
+                exit_event=EventType.FLWR_SIMULATION_RUN_LEAVE,
             )
 
             # Send resulting context
-            context_proto = None  # context_to_proto(updated_context)
+            context_proto = context_to_proto(updated_context)
             out_req = PushSimulationOutputsRequest(
                 run_id=run.run_id, context=context_proto
             )

flwr/simulation/ray_transport/ray_actor.py

@@ -14,6 +14,7 @@
 # ==============================================================================
 """Ray-based Flower Actor and ActorPool implementation."""
 
+
 import threading
 from abc import ABC
 from logging import DEBUG, ERROR, WARNING

flwr/simulation/ray_transport/utils.py

@@ -14,6 +14,7 @@
 # ==============================================================================
 """Utilities for Actors in the Virtual Client Engine."""
 
+
 import traceback
 import warnings
 from logging import ERROR

flwr/simulation/run_simulation.py

@@ -14,6 +14,7 @@
 # ==============================================================================
 """Flower Simulation."""
 
+
 import argparse
 import asyncio
 import json
@@ -23,10 +24,11 @@ import threading
 import traceback
 from logging import DEBUG, ERROR, INFO, WARNING
 from pathlib import Path
-from time import sleep
+from queue import Empty, Queue
 from typing import Any, Optional
 
 from flwr.cli.config_utils import load_and_validate
+from flwr.cli.utils import get_sha256_hash
 from flwr.client import ClientApp
 from flwr.common import Context, EventType, RecordSet, event, log, now
 from flwr.common.config import get_fused_config_from_dir, parse_config_args
@@ -126,7 +128,7 @@ def run_simulation_from_cli() -> None:
     run = Run.create_empty(run_id)
     run.override_config = override_config
 
-    _run_simulation(
+    _ = _run_simulation(
         server_app_attr=server_app_attr,
         client_app_attr=client_app_attr,
         num_supernodes=args.num_supernodes,
@@ -135,7 +137,6 @@ def run_simulation_from_cli() -> None:
         app_dir=args.app,
         run=run,
         enable_tf_gpu_growth=args.enable_tf_gpu_growth,
-        delay_start=args.delay_start,
         verbose_logging=args.verbose,
         server_app_run_config=fused_config,
         is_app=True,
@@ -207,7 +208,7 @@ def run_simulation(
             "\n\tflwr.simulation.run_simulationt(...)",
         )
 
-    _run_simulation(
+    _ = _run_simulation(
         num_supernodes=num_supernodes,
         client_app=client_app,
         server_app=server_app,
@@ -230,6 +231,7 @@ def run_serverapp_th(
     has_exception: threading.Event,
     enable_tf_gpu_growth: bool,
     run_id: int,
+    ctx_queue: "Queue[Context]",
 ) -> threading.Thread:
     """Run SeverApp in a thread."""
 
@@ -242,6 +244,7 @@ def run_serverapp_th(
         _server_app_run_config: UserConfig,
         _server_app_attr: Optional[str],
         _server_app: Optional[ServerApp],
+        _ctx_queue: "Queue[Context]",
     ) -> None:
         """Run SeverApp, after check if GPU memory growth has to be set.
 
@@ -262,13 +265,14 @@ def run_serverapp_th(
             )
 
             # Run ServerApp
-            _run(
+            updated_context = _run(
                 driver=_driver,
                 context=context,
                 server_app_dir=_server_app_dir,
                 server_app_attr=_server_app_attr,
                 loaded_server_app=_server_app,
             )
+            _ctx_queue.put(updated_context)
         except Exception as ex:  # pylint: disable=broad-exception-caught
             log(ERROR, "ServerApp thread raised an exception: %s", ex)
             log(ERROR, traceback.format_exc())
@@ -292,6 +296,7 @@ def run_serverapp_th(
             server_app_run_config,
             server_app_attr,
             server_app,
+            ctx_queue,
         ),
     )
     serverapp_th.start()
@@ -308,14 +313,13 @@ def _main_loop(
     enable_tf_gpu_growth: bool,
     run: Run,
     exit_event: EventType,
-    delay_start: int,
     flwr_dir: Optional[str] = None,
     client_app: Optional[ClientApp] = None,
     client_app_attr: Optional[str] = None,
     server_app: Optional[ServerApp] = None,
     server_app_attr: Optional[str] = None,
     server_app_run_config: Optional[UserConfig] = None,
-) -> None:
+) -> Context:
     """Start ServerApp on a separate thread, then launch Simulation Engine."""
     # Initialize StateFactory
     state_factory = LinkStateFactory(":flwr-in-memory-state:")
@@ -325,6 +329,13 @@ def _main_loop(
     server_app_thread_has_exception = threading.Event()
     serverapp_th = None
     success = True
+    updated_context = Context(
+        run_id=run.run_id,
+        node_id=0,
+        node_config=UserConfig(),
+        state=RecordSet(),
+        run_config=UserConfig(),
+    )
     try:
         # Register run
         log(DEBUG, "Pre-registering run with id %s", run.run_id)
@@ -339,6 +350,7 @@ def _main_loop(
         # Initialize Driver
         driver = InMemoryDriver(state_factory=state_factory)
         driver.set_run(run_id=run.run_id)
+        output_context_queue: "Queue[Context]" = Queue()
 
         # Get and run ServerApp thread
         serverapp_th = run_serverapp_th(
@@ -351,11 +363,9 @@ def _main_loop(
             has_exception=server_app_thread_has_exception,
             enable_tf_gpu_growth=enable_tf_gpu_growth,
             run_id=run.run_id,
+            ctx_queue=output_context_queue,
         )
 
-        # Buffer time so the `ServerApp` in separate thread is ready
-        log(DEBUG, "Buffer time delay: %ds", delay_start)
-        sleep(delay_start)
         # Start Simulation Engine
         vce.start_vce(
             num_supernodes=num_supernodes,
@@ -371,6 +381,11 @@ def _main_loop(
             flwr_dir=flwr_dir,
         )
 
+        updated_context = output_context_queue.get(timeout=3)
+
+    except Empty:
+        log(DEBUG, "Queue timeout. No context received.")
+
     except Exception as ex:
         log(ERROR, "An exception occurred !! %s", ex)
         log(ERROR, traceback.format_exc())
@@ -380,13 +395,20 @@ def _main_loop(
     finally:
         # Trigger stop event
         f_stop.set()
-        event(exit_event, event_details={"success": success})
+        event(
+            exit_event,
+            event_details={
+                "run-id-hash": get_sha256_hash(run.run_id),
+                "success": success,
+            },
+        )
         if serverapp_th:
             serverapp_th.join()
             if server_app_thread_has_exception.is_set():
                 raise RuntimeError("Exception in ServerApp thread")
 
     log(DEBUG, "Stopping Simulation Engine now.")
+    return updated_context
 
 
 # pylint: disable=too-many-arguments,too-many-locals,too-many-positional-arguments
@@ -404,10 +426,9 @@ def _run_simulation(
     flwr_dir: Optional[str] = None,
     run: Optional[Run] = None,
     enable_tf_gpu_growth: bool = False,
-    delay_start: int = 5,
     verbose_logging: bool = False,
     is_app: bool = False,
-) -> None:
+) -> Context:
     """Launch the Simulation Engine."""
     if backend_config is None:
         backend_config = {}
@@ -459,7 +480,6 @@ def _run_simulation(
         enable_tf_gpu_growth,
         run,
         exit_event,
-        delay_start,
         flwr_dir,
         client_app,
         client_app_attr,
@@ -487,7 +507,8 @@ def _run_simulation(
             # Set logger propagation to False to prevent duplicated log output in Colab.
             logger = set_logger_propagation(logger, False)
 
-        _main_loop(*args)
+        updated_context = _main_loop(*args)
+    return updated_context
 
 
 def _parse_args_run_simulation() -> argparse.ArgumentParser:
@@ -537,13 +558,6 @@ def _parse_args_run_simulation() -> argparse.ArgumentParser:
         "Read more about how `tf.config.experimental.set_memory_growth()` works in "
         "the TensorFlow documentation: https://www.tensorflow.org/api/stable.",
     )
-    parser.add_argument(
-        "--delay-start",
-        type=int,
-        default=3,
-        help="Buffer time (in seconds) to delay the start the simulation engine after "
-        "the `ServerApp`, which runs in a separate thread, has been launched.",
-    )
     parser.add_argument(
         "--verbose",
         action="store_true",

flwr/simulation/simulationio_connection.py

@@ -23,6 +23,7 @@ import grpc
 from flwr.common.constant import SIMULATIONIO_API_DEFAULT_CLIENT_ADDRESS
 from flwr.common.grpc import create_channel
 from flwr.common.logger import log
+from flwr.common.retry_invoker import _make_simple_grpc_retry_invoker, _wrap_stub
 from flwr.proto.simulationio_pb2_grpc import SimulationIoStub  # pylint: disable=E0611
 
 
@@ -48,6 +49,7 @@ class SimulationIoConnection:
         self._cert = root_certificates
         self._grpc_stub: Optional[SimulationIoStub] = None
         self._channel: Optional[grpc.Channel] = None
+        self._retry_invoker = _make_simple_grpc_retry_invoker()
 
     @property
     def _is_connected(self) -> bool:
@@ -72,6 +74,7 @@ class SimulationIoConnection:
             root_certificates=self._cert,
         )
         self._grpc_stub = SimulationIoStub(self._channel)
+        _wrap_stub(self._grpc_stub, self._retry_invoker)
         log(DEBUG, "[SimulationIO] Connected to %s", self._addr)
 
     def _disconnect(self) -> None:

flwr/superexec/app.py

@@ -14,6 +14,7 @@
 # ==============================================================================
 """Flower SuperExec app."""
 
+
 import argparse
 import sys
 from logging import INFO

flwr/superexec/deployment.py

@@ -14,6 +14,7 @@
 # ==============================================================================
 """Deployment engine executor."""
 
+
 import hashlib
 from logging import ERROR, INFO
 from pathlib import Path