flwr-nightly 1.14.0.dev20241204__py3-none-any.whl → 1.14.0.dev20241216__py3-none-any.whl

flwr/server/superlink/fleet/vce/backend/__init__.py

@@ -14,6 +14,7 @@
 # ==============================================================================
 """Simulation Engine Backends."""
 
+
 import importlib
 
 from .backend import Backend, BackendConfig

flwr/server/superlink/fleet/vce/backend/raybackend.py

@@ -14,6 +14,7 @@
 # ==============================================================================
 """Ray backend for the Fleet API using the Simulation Engine."""
 
+
 import sys
 from logging import DEBUG, ERROR
 from typing import Callable, Optional, Union

flwr/server/superlink/linkstate/in_memory_linkstate.py

@@ -265,41 +265,15 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
             for task_res in task_res_found:
                 task_res.task.delivered_at = delivered_at
 
-            # Cleanup
-            self._force_delete_tasks_by_ids(set(ret.keys()))
-
         return list(ret.values())
 
-    def delete_tasks(self, task_ids: set[UUID]) -> None:
-        """Delete all delivered TaskIns/TaskRes pairs."""
-        task_ins_to_be_deleted: set[UUID] = set()
-        task_res_to_be_deleted: set[UUID] = set()
-
-        with self.lock:
-            for task_ins_id in task_ids:
-                # Find the task_id of the matching task_res
-                for task_res_id, task_res in self.task_res_store.items():
-                    if UUID(task_res.task.ancestry[0]) != task_ins_id:
-                        continue
-                    if task_res.task.delivered_at == "":
-                        continue
-
-                    task_ins_to_be_deleted.add(task_ins_id)
-                    task_res_to_be_deleted.add(task_res_id)
-
-            for task_id in task_ins_to_be_deleted:
-                del self.task_ins_store[task_id]
-                del self.task_ins_id_to_task_res_id[task_id]
-            for task_id in task_res_to_be_deleted:
-                del self.task_res_store[task_id]
-
-    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:
+    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
 
         with self.lock:
-            for task_id in task_ids:
+            for task_id in task_ins_ids:
                 # Delete TaskIns
                 if task_id in self.task_ins_store:
                     del self.task_ins_store[task_id]
@@ -308,6 +282,16 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
                     task_res_id = self.task_ins_id_to_task_res_id.pop(task_id)
                     del self.task_res_store[task_res_id]
 
+    def get_task_ids_from_run_id(self, run_id: int) -> set[UUID]:
+        """Get all TaskIns IDs for the given run_id."""
+        task_id_list: set[UUID] = set()
+        with self.lock:
+            for task_id, task_ins in self.task_ins_store.items():
+                if task_ins.run_id == run_id:
+                    task_id_list.add(task_id)
+
+        return task_id_list
+
     def num_task_ins(self) -> int:
         """Calculate the number of task_ins in store.
 

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(

flwr/server/superlink/linkstate/sqlite_linkstate.py

@@ -14,6 +14,7 @@
 # ==============================================================================
 """SQLite based implemenation of the link state."""
 
+
 # pylint: disable=too-many-lines
 
 import json
@@ -566,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:
@@ -592,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

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,
 )
@@ -122,6 +126,22 @@ class SimulationIoServicer(simulationio_pb2_grpc.SimulationIoServicer):
         )
         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

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,7 +24,6 @@ import threading
 import traceback
 from logging import DEBUG, ERROR, INFO, WARNING
 from pathlib import Path
-from time import sleep
 from typing import Any, Optional
 
 from flwr.cli.config_utils import load_and_validate
@@ -135,7 +135,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,
@@ -308,7 +307,6 @@ 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,
@@ -353,9 +351,6 @@ def _main_loop(
             run_id=run.run_id,
         )
 
-        # 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,
@@ -404,7 +399,6 @@ 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:
@@ -459,7 +453,6 @@ def _run_simulation(
         enable_tf_gpu_growth,
         run,
         exit_event,
-        delay_start,
         flwr_dir,
         client_app,
         client_app_attr,
@@ -537,13 +530,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

flwr/superexec/exec_grpc.py

@@ -14,18 +14,22 @@
 # ==============================================================================
 """SuperExec gRPC API."""
 
+
+from collections.abc import Sequence
 from logging import INFO
 from typing import Optional
 
 import grpc
 
 from flwr.common import GRPC_MAX_MESSAGE_LENGTH
+from flwr.common.auth_plugin import ExecAuthPlugin
 from flwr.common.logger import log
 from flwr.common.typing import UserConfig
 from flwr.proto.exec_pb2_grpc import add_ExecServicer_to_server
 from flwr.server.superlink.ffs.ffs_factory import FfsFactory
 from flwr.server.superlink.fleet.grpc_bidi.grpc_server import generic_create_grpc_server
 from flwr.server.superlink.linkstate import LinkStateFactory
+from flwr.superexec.exec_user_auth_interceptor import ExecUserAuthInterceptor
 
 from .exec_servicer import ExecServicer
 from .executor import Executor
@@ -39,6 +43,7 @@ def run_exec_api_grpc(
     ffs_factory: FfsFactory,
     certificates: Optional[tuple[bytes, bytes, bytes]],
     config: UserConfig,
+    auth_plugin: Optional[ExecAuthPlugin] = None,
 ) -> grpc.Server:
     """Run Exec API (gRPC, request-response)."""
     executor.set_config(config)
@@ -47,16 +52,29 @@ def run_exec_api_grpc(
         linkstate_factory=state_factory,
         ffs_factory=ffs_factory,
         executor=executor,
+        auth_plugin=auth_plugin,
     )
+    interceptors: Optional[Sequence[grpc.ServerInterceptor]] = None
+    if auth_plugin is not None:
+        interceptors = [ExecUserAuthInterceptor(auth_plugin)]
     exec_add_servicer_to_server_fn = add_ExecServicer_to_server
     exec_grpc_server = generic_create_grpc_server(
         servicer_and_add_fn=(exec_servicer, exec_add_servicer_to_server_fn),
         server_address=address,
         max_message_length=GRPC_MAX_MESSAGE_LENGTH,
         certificates=certificates,
+        interceptors=interceptors,
     )
 
-    log(INFO, "Flower Deployment Engine: Starting Exec API on %s", address)
+    if auth_plugin is None:
+        log(INFO, "Flower Deployment Engine: Starting Exec API on %s", address)
+    else:
+        log(
+            INFO,
+            "Flower Deployment Engine: Starting Exec API with user "
+            "authentication on %s",
+            address,
+        )
     exec_grpc_server.start()
 
     return exec_grpc_server

flwr/superexec/exec_servicer.py

@@ -18,24 +18,33 @@
 import time
 from collections.abc import Generator
 from logging import ERROR, INFO
-from typing import Any
+from typing import Any, Optional
+from uuid import UUID
 
 import grpc
 
 from flwr.common import now
-from flwr.common.constant import LOG_STREAM_INTERVAL, Status
+from flwr.common.auth_plugin import ExecAuthPlugin
+from flwr.common.constant import LOG_STREAM_INTERVAL, Status, SubStatus
 from flwr.common.logger import log
 from flwr.common.serde import (
     configs_record_from_proto,
     run_to_proto,
     user_config_from_proto,
 )
+from flwr.common.typing import RunStatus
 from flwr.proto import exec_pb2_grpc  # pylint: disable=E0611
 from flwr.proto.exec_pb2 import (  # pylint: disable=E0611
+    GetAuthTokensRequest,
+    GetAuthTokensResponse,
+    GetLoginDetailsRequest,
+    GetLoginDetailsResponse,
     ListRunsRequest,
     ListRunsResponse,
     StartRunRequest,
     StartRunResponse,
+    StopRunRequest,
+    StopRunResponse,
     StreamLogsRequest,
     StreamLogsResponse,
 )
@@ -53,11 +62,13 @@ class ExecServicer(exec_pb2_grpc.ExecServicer):
         linkstate_factory: LinkStateFactory,
         ffs_factory: FfsFactory,
         executor: Executor,
+        auth_plugin: Optional[ExecAuthPlugin] = None,
     ) -> None:
         self.linkstate_factory = linkstate_factory
         self.ffs_factory = ffs_factory
         self.executor = executor
         self.executor.initialize(linkstate_factory, ffs_factory)
+        self.auth_plugin = auth_plugin
 
     def StartRun(
         self, request: StartRunRequest, context: grpc.ServicerContext
@@ -126,6 +137,69 @@ class ExecServicer(exec_pb2_grpc.ExecServicer):
         # Handle `flwr ls --run-id <run_id>`
         return _create_list_runs_response({request.run_id}, state)
 
+    def StopRun(
+        self, request: StopRunRequest, context: grpc.ServicerContext
+    ) -> StopRunResponse:
+        """Stop a given run ID."""
+        log(INFO, "ExecServicer.StopRun")
+        state = self.linkstate_factory.state()
+
+        # Exit if `run_id` not found
+        if not state.get_run(request.run_id):
+            context.abort(
+                grpc.StatusCode.NOT_FOUND, f"Run ID {request.run_id} not found"
+            )
+
+        run_status = state.get_run_status({request.run_id})[request.run_id]
+        if run_status.status == Status.FINISHED:
+            context.abort(
+                grpc.StatusCode.FAILED_PRECONDITION,
+                f"Run ID {request.run_id} is already finished",
+            )
+
+        update_success = state.update_run_status(
+            run_id=request.run_id,
+            new_status=RunStatus(Status.FINISHED, SubStatus.STOPPED, ""),
+        )
+
+        if update_success:
+            task_ids: set[UUID] = state.get_task_ids_from_run_id(request.run_id)
+
+            # Delete TaskIns and TaskRes for the `run_id`
+            state.delete_tasks(task_ids)
+
+        return StopRunResponse(success=update_success)
+
+    def GetLoginDetails(
+        self, request: GetLoginDetailsRequest, context: grpc.ServicerContext
+    ) -> GetLoginDetailsResponse:
+        """Start login."""
+        log(INFO, "ExecServicer.GetLoginDetails")
+        if self.auth_plugin is None:
+            context.abort(
+                grpc.StatusCode.UNIMPLEMENTED,
+                "ExecServicer initialized without user authentication",
+            )
+            raise grpc.RpcError()  # This line is unreachable
+        return GetLoginDetailsResponse(
+            login_details=self.auth_plugin.get_login_details()
+        )
+
+    def GetAuthTokens(
+        self, request: GetAuthTokensRequest, context: grpc.ServicerContext
+    ) -> GetAuthTokensResponse:
+        """Get auth token."""
+        log(INFO, "ExecServicer.GetAuthTokens")
+        if self.auth_plugin is None:
+            context.abort(
+                grpc.StatusCode.UNIMPLEMENTED,
+                "ExecServicer initialized without user authentication",
+            )
+            raise grpc.RpcError()  # This line is unreachable
+        return GetAuthTokensResponse(
+            auth_tokens=self.auth_plugin.get_auth_tokens(dict(request.auth_details))
+        )
+
 
 def _create_list_runs_response(run_ids: set[int], state: LinkState) -> ListRunsResponse:
     """Create response for `flwr ls --runs` and `flwr ls --run-id <run_id>`."""