flwr-nightly 1.19.0.dev20250611__py3-none-any.whl → 1.19.0.dev20250613__py3-none-any.whl

flwr/supercore/object_store/in_memory_object_store.py

@@ -15,44 +15,95 @@
 """Flower in-memory ObjectStore implementation."""
 
 
+import threading
+from dataclasses import dataclass
 from typing import Optional
 
-from flwr.common.inflatable import get_object_id, is_valid_sha256_hash
+from flwr.common.inflatable import (
+    get_object_children_ids_from_object_content,
+    get_object_id,
+    is_valid_sha256_hash,
+    iterate_object_tree,
+)
 from flwr.common.inflatable_utils import validate_object_content
+from flwr.proto.message_pb2 import ObjectTree  # pylint: disable=E0611
 
 from .object_store import NoObjectInStoreError, ObjectStore
 
 
+@dataclass
+class ObjectEntry:
+    """Data class representing an object entry in the store."""
+
+    content: bytes
+    is_available: bool
+    ref_count: int  # Number of references (direct parents) to this object
+    runs: set[int]  # Set of run IDs that used this object
+
+
 class InMemoryObjectStore(ObjectStore):
     """In-memory implementation of the ObjectStore interface."""
 
     def __init__(self, verify: bool = True) -> None:
         self.verify = verify
-        self.store: dict[str, bytes] = {}
-        # Mapping the Object ID of a message to the list of children object IDs
-        self.msg_children_objects_mapping: dict[str, list[str]] = {}
-
-    def preregister(self, object_ids: list[str]) -> list[str]:
+        self.store: dict[str, ObjectEntry] = {}
+        self.lock_store = threading.RLock()
+        # Mapping the Object ID of a message to the list of descendant object IDs
+        self.msg_descendant_objects_mapping: dict[str, list[str]] = {}
+        self.lock_msg_mapping = threading.RLock()
+        # Mapping each run ID to a set of object IDs that are used in that run
+        self.run_objects_mapping: dict[int, set[str]] = {}
+
+    def preregister(self, run_id: int, object_tree: ObjectTree) -> list[str]:
         """Identify and preregister missing objects."""
         new_objects = []
-        for obj_id in object_ids:
+        if run_id not in self.run_objects_mapping:
+            self.run_objects_mapping[run_id] = set()
+
+        for tree_node in iterate_object_tree(object_tree):
+            obj_id = tree_node.object_id
             # Verify object ID format (must be a valid sha256 hash)
             if not is_valid_sha256_hash(obj_id):
                 raise ValueError(f"Invalid object ID format: {obj_id}")
-            if obj_id not in self.store:
-                self.store[obj_id] = b""
-                new_objects.append(obj_id)
+            with self.lock_store:
+                if obj_id not in self.store:
+                    self.store[obj_id] = ObjectEntry(
+                        content=b"",  # Initially empty content
+                        is_available=False,  # Initially not available
+                        ref_count=0,  # Reference count starts at 0
+                        runs={run_id},  # Start with the current run ID
+                    )
+
+                    # Increment the reference count for all its children
+                    # Post-order traversal ensures that children are registered
+                    # before parents
+                    for child_node in tree_node.children:
+                        child_id = child_node.object_id
+                        self.store[child_id].ref_count += 1
+
+                    # Add the object ID to the run's mapping
+                    self.run_objects_mapping[run_id].add(obj_id)
+
+                    # Add to the list of new objects
+                    new_objects.append(obj_id)
+                else:
+                    # Object is in store, retrieve it
+                    obj_entry = self.store[obj_id]
+
+                    # Add to the list of new objects if not available
+                    if not obj_entry.is_available:
+                        new_objects.append(obj_id)
+
+                    # If the object is already registered but not in this run,
+                    # add the run ID to its runs
+                    if obj_id not in self.run_objects_mapping[run_id]:
+                        obj_entry.runs.add(run_id)
+                        self.run_objects_mapping[run_id].add(obj_id)
 
         return new_objects
 
     def put(self, object_id: str, object_content: bytes) -> None:
         """Put an object into the store."""
-        # Only allow adding the object if it has been preregistered
-        if object_id not in self.store:
-            raise NoObjectInStoreError(
-                f"Object with ID '{object_id}' was not pre-registered."
-            )
-
         if self.verify:
             # Verify object_id and object_content match
             object_id_from_content = get_object_id(object_content)
@@ -62,41 +113,117 @@ class InMemoryObjectStore(ObjectStore):
             # Validate object content
             validate_object_content(content=object_content)
 
-        # Return if object is already present in the store
-        if self.store[object_id] != b"":
-            return
+        with self.lock_store:
+            # Only allow adding the object if it has been preregistered
+            if object_id not in self.store:
+                raise NoObjectInStoreError(
+                    f"Object with ID '{object_id}' was not pre-registered."
+                )
+
+            # Return if object is already present in the store
+            if self.store[object_id].is_available:
+                return
 
-        self.store[object_id] = object_content
+            # Update the object entry in the store
+            self.store[object_id].content = object_content
+            self.store[object_id].is_available = True
 
     def set_message_descendant_ids(
         self, msg_object_id: str, descendant_ids: list[str]
     ) -> None:
         """Store the mapping from a ``Message`` object ID to the object IDs of its
         descendants."""
-        self.msg_children_objects_mapping[msg_object_id] = descendant_ids
+        with self.lock_msg_mapping:
+            self.msg_descendant_objects_mapping[msg_object_id] = descendant_ids
 
     def get_message_descendant_ids(self, msg_object_id: str) -> list[str]:
         """Retrieve the object IDs of all descendants of a given Message."""
-        if msg_object_id not in self.msg_children_objects_mapping:
-            raise NoObjectInStoreError(
-                f"No message registered in Object Store with ID '{msg_object_id}'. "
-                "Mapping to descendants could not be found."
-            )
-        return self.msg_children_objects_mapping[msg_object_id]
+        with self.lock_msg_mapping:
+            if msg_object_id not in self.msg_descendant_objects_mapping:
+                raise NoObjectInStoreError(
+                    f"No message registered in Object Store with ID '{msg_object_id}'. "
+                    "Mapping to descendants could not be found."
+                )
+            return self.msg_descendant_objects_mapping[msg_object_id]
+
+    def delete_message_descendant_ids(self, msg_object_id: str) -> None:
+        """Delete the mapping from a ``Message`` object ID to its descendants."""
+        with self.lock_msg_mapping:
+            self.msg_descendant_objects_mapping.pop(msg_object_id, None)
 
     def get(self, object_id: str) -> Optional[bytes]:
         """Get an object from the store."""
-        return self.store.get(object_id)
+        with self.lock_store:
+            # Check if the object ID is pre-registered
+            if object_id not in self.store:
+                return None
+
+            # Return content (if not yet available, it will b"")
+            return self.store[object_id].content
 
     def delete(self, object_id: str) -> None:
-        """Delete an object from the store."""
-        if object_id in self.store:
-            del self.store[object_id]
+        """Delete an object and its unreferenced descendants from the store."""
+        with self.lock_store:
+            # If the object is not in the store, nothing to delete
+            if (object_entry := self.store.get(object_id)) is None:
+                return
+
+            # Delete the object if it has no references left
+            if object_entry.ref_count == 0:
+                del self.store[object_id]
+
+                # Remove the object from the run's mapping
+                for run_id in object_entry.runs:
+                    self.run_objects_mapping[run_id].discard(object_id)
+
+                # Decrease the reference count of its children
+                children_ids = get_object_children_ids_from_object_content(
+                    object_entry.content
+                )
+                for child_id in children_ids:
+                    self.store[child_id].ref_count -= 1
+
+                    # Recursively try to delete the child object
+                    self.delete(child_id)
+
+    def delete_objects_in_run(self, run_id: int) -> None:
+        """Delete all objects that were registered in a specific run."""
+        with self.lock_store:
+            if run_id not in self.run_objects_mapping:
+                return
+            for object_id in list(self.run_objects_mapping[run_id]):
+                # Check if the object is still in the store
+                if (object_entry := self.store.get(object_id)) is None:
+                    continue
+
+                # Remove the run ID from the object's runs
+                object_entry.runs.discard(run_id)
+
+                # Only message objects are allowed to have a `ref_count` of 0,
+                # and every message object must have a `ref_count` of 0
+                if object_entry.ref_count == 0:
+                    # Delete the message object and its unreferenced descendants
+                    self.delete(object_id)
+
+                    # Delete the message's descendants mapping
+                    self.delete_message_descendant_ids(object_id)
+
+            # Remove the run from the mapping
+            del self.run_objects_mapping[run_id]
 
     def clear(self) -> None:
         """Clear the store."""
-        self.store.clear()
+        with self.lock_store:
+            self.store.clear()
+            self.msg_descendant_objects_mapping.clear()
+            self.run_objects_mapping.clear()
 
     def __contains__(self, object_id: str) -> bool:
         """Check if an object_id is in the store."""
-        return object_id in self.store
+        with self.lock_store:
+            return object_id in self.store
+
+    def __len__(self) -> int:
+        """Get the number of objects in the store."""
+        with self.lock_store:
+            return len(self.store)

flwr/supercore/object_store/object_store.py

@@ -18,6 +18,8 @@
 import abc
 from typing import Optional
 
+from flwr.proto.message_pb2 import ObjectTree  # pylint: disable=E0611
+
 
 class NoObjectInStoreError(Exception):
     """Error when trying to access an element in the ObjectStore that does not exist."""
@@ -39,20 +41,23 @@ class ObjectStore(abc.ABC):
     """
 
     @abc.abstractmethod
-    def preregister(self, object_ids: list[str]) -> list[str]:
+    def preregister(self, run_id: int, object_tree: ObjectTree) -> list[str]:
         """Identify and preregister missing objects in the `ObjectStore`.
 
         Parameters
         ----------
-        object_ids : list[str]
-            A list of object IDs to check against the store. Any object ID not already
-            present will be preregistered.
+        run_id : int
+            The ID of the run for which to preregister objects.
+        object_tree : ObjectTree
+            The object tree containing the IDs of objects to preregister.
+            This tree should contain all objects that are expected to be
+            stored in the `ObjectStore`.
 
         Returns
         -------
         list[str]
-            A list of object IDs that were not present in the `ObjectStore` and have now
-            been preregistered.
+            A list of object IDs that were either not previously preregistered
+            in the `ObjectStore`, or were preregistered but are not yet available.
         """
 
     @abc.abstractmethod
@@ -84,12 +89,34 @@ class ObjectStore(abc.ABC):
 
     @abc.abstractmethod
     def delete(self, object_id: str) -> None:
-        """Delete an object from the store.
+        """Delete an object and its unreferenced descendants from the store.
+
+        This method attempts to recursively delete the specified object and its
+        descendants, if they are not referenced by any other object.
 
         Parameters
         ----------
         object_id : str
             The object_id under which the object is stored.
+
+        Notes
+        -----
+        The object of the given object_id will NOT be deleted if it is still referenced
+        by any other object in the store.
+        """
+
+    @abc.abstractmethod
+    def delete_objects_in_run(self, run_id: int) -> None:
+        """Delete all objects that were registered in a specific run.
+
+        Parameters
+        ----------
+        run_id : int
+            The ID of the run for which to delete objects.
+
+        Notes
+        -----
+        Objects that are still registered in other runs will NOT be deleted.
         """
 
     @abc.abstractmethod
@@ -129,6 +156,16 @@ class ObjectStore(abc.ABC):
             A list of object IDs of all descendant objects of the ``Message``.
         """
 
+    @abc.abstractmethod
+    def delete_message_descendant_ids(self, msg_object_id: str) -> None:
+        """Delete the mapping from a ``Message`` object ID to its descendants.
+
+        Parameters
+        ----------
+        msg_object_id : str
+            The object ID of the ``Message``.
+        """
+
     @abc.abstractmethod
     def __contains__(self, object_id: str) -> bool:
         """Check if an object_id is in the store.
@@ -143,3 +180,13 @@ class ObjectStore(abc.ABC):
         bool
             True if the object_id is in the store, False otherwise.
         """
+
+    @abc.abstractmethod
+    def __len__(self) -> int:
+        """Return the number of objects in the store.
+
+        Returns
+        -------
+        int
+            The number of objects currently stored.
+        """

flwr/superexec/deployment.py

@@ -132,6 +132,7 @@ class DeploymentEngine(Executor):
         self,
         fab: Fab,
         override_config: UserConfig,
+        flwr_aid: Optional[str],
     ) -> int:
         fab_hash = self.ffs.put(fab.content, {})
         if fab_hash != fab.hash_str:
@@ -141,7 +142,7 @@ class DeploymentEngine(Executor):
         fab_id, fab_version = get_fab_metadata(fab.content)
 
         run_id = self.linkstate.create_run(
-            fab_id, fab_version, fab_hash, override_config, ConfigRecord()
+            fab_id, fab_version, fab_hash, override_config, ConfigRecord(), flwr_aid
         )
         return run_id
 
@@ -161,6 +162,7 @@ class DeploymentEngine(Executor):
         fab_file: bytes,
         override_config: UserConfig,
         federation_options: ConfigRecord,
+        flwr_aid: Optional[str],
     ) -> Optional[int]:
         """Start run using the Flower Deployment Engine."""
         run_id = None
@@ -168,7 +170,9 @@ class DeploymentEngine(Executor):
 
             # Call SuperLink to create run
             run_id = self._create_run(
-                Fab(hashlib.sha256(fab_file).hexdigest(), fab_file), override_config
+                Fab(hashlib.sha256(fab_file).hexdigest(), fab_file),
+                override_config,
+                flwr_aid,
             )
 
             # Register context for the Run

flwr/superexec/exec_grpc.py

@@ -29,6 +29,7 @@ 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.linkstate import LinkStateFactory
+from flwr.supercore.object_store import ObjectStoreFactory
 from flwr.superexec.exec_event_log_interceptor import ExecEventLogInterceptor
 from flwr.superexec.exec_user_auth_interceptor import ExecUserAuthInterceptor
 
@@ -42,6 +43,7 @@ def run_exec_api_grpc(
     executor: Executor,
     state_factory: LinkStateFactory,
     ffs_factory: FfsFactory,
+    objectstore_factory: ObjectStoreFactory,
     certificates: Optional[tuple[bytes, bytes, bytes]],
     config: UserConfig,
     auth_plugin: Optional[ExecAuthPlugin] = None,
@@ -54,6 +56,7 @@ def run_exec_api_grpc(
     exec_servicer: grpc.Server = ExecServicer(
         linkstate_factory=state_factory,
         ffs_factory=ffs_factory,
+        objectstore_factory=objectstore_factory,
         executor=executor,
         auth_plugin=auth_plugin,
     )

flwr/superexec/exec_servicer.py

@@ -18,20 +18,25 @@
 import time
 from collections.abc import Generator
 from logging import ERROR, INFO
-from typing import Any, Optional
+from typing import Any, Optional, cast
 
 import grpc
 
 from flwr.common import now
 from flwr.common.auth_plugin import ExecAuthPlugin
-from flwr.common.constant import LOG_STREAM_INTERVAL, Status, SubStatus
+from flwr.common.constant import (
+    LOG_STREAM_INTERVAL,
+    RUN_ID_NOT_FOUND_MESSAGE,
+    Status,
+    SubStatus,
+)
 from flwr.common.logger import log
 from flwr.common.serde import (
     config_record_from_proto,
     run_to_proto,
     user_config_from_proto,
 )
-from flwr.common.typing import RunStatus
+from flwr.common.typing import Run, RunStatus
 from flwr.proto import exec_pb2_grpc  # pylint: disable=E0611
 from flwr.proto.exec_pb2 import (  # pylint: disable=E0611
     GetAuthTokensRequest,
@@ -49,22 +54,26 @@ from flwr.proto.exec_pb2 import (  # pylint: disable=E0611
 )
 from flwr.server.superlink.ffs.ffs_factory import FfsFactory
 from flwr.server.superlink.linkstate import LinkState, LinkStateFactory
+from flwr.supercore.object_store import ObjectStore, ObjectStoreFactory
 
+from .exec_user_auth_interceptor import shared_account_info
 from .executor import Executor
 
 
 class ExecServicer(exec_pb2_grpc.ExecServicer):
     """SuperExec API servicer."""
 
-    def __init__(
+    def __init__(  # pylint: disable=R0913, R0917
         self,
         linkstate_factory: LinkStateFactory,
         ffs_factory: FfsFactory,
+        objectstore_factory: ObjectStoreFactory,
         executor: Executor,
         auth_plugin: Optional[ExecAuthPlugin] = None,
     ) -> None:
         self.linkstate_factory = linkstate_factory
         self.ffs_factory = ffs_factory
+        self.objectstore_factory = objectstore_factory
         self.executor = executor
         self.executor.initialize(linkstate_factory, ffs_factory)
         self.auth_plugin = auth_plugin
@@ -75,10 +84,12 @@ class ExecServicer(exec_pb2_grpc.ExecServicer):
         """Create run ID."""
         log(INFO, "ExecServicer.StartRun")
 
+        flwr_aid = shared_account_info.get().flwr_aid if self.auth_plugin else None
         run_id = self.executor.start_run(
             request.fab.content,
             user_config_from_proto(request.override_config),
             config_record_from_proto(request.federation_options),
+            flwr_aid,
         )
 
         if run_id is None:
@@ -94,12 +105,20 @@ class ExecServicer(exec_pb2_grpc.ExecServicer):
         log(INFO, "ExecServicer.StreamLogs")
         state = self.linkstate_factory.state()
 
-        # Retrieve run ID
+        # Retrieve run ID and run
         run_id = request.run_id
+        run = state.get_run(run_id)
 
         # Exit if `run_id` not found
-        if not state.get_run(run_id):
-            context.abort(grpc.StatusCode.NOT_FOUND, "Run ID not found")
+        if not run:
+            context.abort(grpc.StatusCode.NOT_FOUND, RUN_ID_NOT_FOUND_MESSAGE)
+
+        # If user auth is enabled, check if `flwr_aid` matches the run's `flwr_aid`
+        if self.auth_plugin:
+            flwr_aid = shared_account_info.get().flwr_aid
+            _check_flwr_aid_in_run(
+                flwr_aid=flwr_aid, run=cast(Run, run), context=context
+            )
 
         after_timestamp = request.after_timestamp + 1e-6
         while context.is_active():
@@ -118,7 +137,10 @@ class ExecServicer(exec_pb2_grpc.ExecServicer):
             # is returned at this point and the server ends the stream.
             run_status = state.get_run_status({run_id})[run_id]
             if run_status.status == Status.FINISHED:
-                log(INFO, "All logs for run ID `%s` returned", request.run_id)
+                log(INFO, "All logs for run ID `%s` returned", run_id)
+
+                # Delete objects of the run from the object store
+                self.objectstore_factory.store().delete_objects_in_run(run_id)
                 break
 
             time.sleep(LOG_STREAM_INTERVAL)  # Sleep briefly to avoid busy waiting
@@ -130,11 +152,44 @@ class ExecServicer(exec_pb2_grpc.ExecServicer):
         log(INFO, "ExecServicer.List")
         state = self.linkstate_factory.state()
 
-        # Handle `flwr ls --runs`
+        # Build a set of run IDs for `flwr ls --runs`
         if not request.HasField("run_id"):
-            return _create_list_runs_response(state.get_run_ids(), state)
-        # Handle `flwr ls --run-id <run_id>`
-        return _create_list_runs_response({request.run_id}, state)
+            if self.auth_plugin:
+                # If no `run_id` is specified and user auth is enabled,
+                # return run IDs for the authenticated user
+                flwr_aid = shared_account_info.get().flwr_aid
+                if flwr_aid is None:
+                    context.abort(
+                        grpc.StatusCode.PERMISSION_DENIED,
+                        "️⛔️ User authentication is enabled, but `flwr_aid` is None",
+                    )
+                run_ids = state.get_run_ids(flwr_aid=flwr_aid)
+            else:
+                # If no `run_id` is specified and no user auth is enabled,
+                # return all run IDs
+                run_ids = state.get_run_ids(None)
+        # Build a set of run IDs for `flwr ls --run-id <run_id>`
+        else:
+            # Retrieve run ID and run
+            run_id = request.run_id
+            run = state.get_run(run_id)
+
+            # Exit if `run_id` not found
+            if not run:
+                context.abort(grpc.StatusCode.NOT_FOUND, RUN_ID_NOT_FOUND_MESSAGE)
+
+            # If user auth is enabled, check if `flwr_aid` matches the run's `flwr_aid`
+            if self.auth_plugin:
+                flwr_aid = shared_account_info.get().flwr_aid
+                _check_flwr_aid_in_run(
+                    flwr_aid=flwr_aid, run=cast(Run, run), context=context
+                )
+
+            run_ids = {run_id}
+
+        # Init the object store
+        store = self.objectstore_factory.store()
+        return _create_list_runs_response(run_ids, state, store)
 
     def StopRun(
         self, request: StopRunRequest, context: grpc.ServicerContext
@@ -143,30 +198,42 @@ class ExecServicer(exec_pb2_grpc.ExecServicer):
         log(INFO, "ExecServicer.StopRun")
         state = self.linkstate_factory.state()
 
+        # Retrieve run ID and run
+        run_id = request.run_id
+        run = state.get_run(run_id)
+
         # 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"
+        if not run:
+            context.abort(grpc.StatusCode.NOT_FOUND, RUN_ID_NOT_FOUND_MESSAGE)
+
+        # If user auth is enabled, check if `flwr_aid` matches the run's `flwr_aid`
+        if self.auth_plugin:
+            flwr_aid = shared_account_info.get().flwr_aid
+            _check_flwr_aid_in_run(
+                flwr_aid=flwr_aid, run=cast(Run, run), context=context
             )
 
-        run_status = state.get_run_status({request.run_id})[request.run_id]
+        run_status = state.get_run_status({run_id})[run_id]
         if run_status.status == Status.FINISHED:
             context.abort(
                 grpc.StatusCode.FAILED_PRECONDITION,
-                f"Run ID {request.run_id} is already finished",
+                f"Run ID {run_id} is already finished",
             )
 
         update_success = state.update_run_status(
-            run_id=request.run_id,
+            run_id=run_id,
             new_status=RunStatus(Status.FINISHED, SubStatus.STOPPED, ""),
         )
 
         if update_success:
-            message_ids: set[str] = state.get_message_ids_from_run_id(request.run_id)
+            message_ids: set[str] = state.get_message_ids_from_run_id(run_id)
 
             # Delete Messages and their replies for the `run_id`
             state.delete_messages(message_ids)
 
+            # Delete objects of the run from the object store
+            self.objectstore_factory.store().delete_objects_in_run(run_id)
+
         return StopRunResponse(success=update_success)
 
     def GetLoginDetails(
@@ -221,10 +288,46 @@ class ExecServicer(exec_pb2_grpc.ExecServicer):
         )
 
 
-def _create_list_runs_response(run_ids: set[int], state: LinkState) -> ListRunsResponse:
+def _create_list_runs_response(
+    run_ids: set[int], state: LinkState, store: ObjectStore
+) -> ListRunsResponse:
     """Create response for `flwr ls --runs` and `flwr ls --run-id <run_id>`."""
-    run_dict = {run_id: state.get_run(run_id) for run_id in run_ids}
+    run_dict = {run_id: run for run_id in run_ids if (run := state.get_run(run_id))}
+
+    # Delete objects of finished runs from the object store
+    for run_id, run in run_dict.items():
+        if run.status.status == Status.FINISHED:
+            store.delete_objects_in_run(run_id)
+
     return ListRunsResponse(
-        run_dict={run_id: run_to_proto(run) for run_id, run in run_dict.items() if run},
+        run_dict={run_id: run_to_proto(run) for run_id, run in run_dict.items()},
         now=now().isoformat(),
     )
+
+
+def _check_flwr_aid_in_run(
+    flwr_aid: Optional[str], run: Run, context: grpc.ServicerContext
+) -> None:
+    """Guard clause to check if `flwr_aid` matches the run's `flwr_aid`."""
+    # `flwr_aid` must not be None. Abort if it is None.
+    if flwr_aid is None:
+        context.abort(
+            grpc.StatusCode.PERMISSION_DENIED,
+            "️⛔️ User authentication is enabled, but `flwr_aid` is None",
+        )
+
+    # `run.flwr_aid` must not be an empty string. Abort if it is empty.
+    run_flwr_aid = run.flwr_aid
+    if not run_flwr_aid:
+        context.abort(
+            grpc.StatusCode.PERMISSION_DENIED,
+            "⛔️ User authentication is enabled, but the run is not associated "
+            "with a `flwr_aid`.",
+        )
+
+    # Exit if `flwr_aid` does not match the run's `flwr_aid`
+    if run_flwr_aid != flwr_aid:
+        context.abort(
+            grpc.StatusCode.PERMISSION_DENIED,
+            "⛔️ Run ID does not belong to the user",
+        )

flwr/superexec/executor.py

@@ -74,6 +74,7 @@ class Executor(ABC):
         fab_file: bytes,
         override_config: UserConfig,
         federation_options: ConfigRecord,
+        flwr_aid: Optional[str],
     ) -> Optional[int]:
         """Start a run using the given Flower FAB ID and version.
 
@@ -88,6 +89,9 @@ class Executor(ABC):
             The config overrides dict sent by the user (using `flwr run`).
         federation_options: ConfigRecord
             The federation options sent by the user (using `flwr run`).
+        flwr_aid : Optional[str]
+            The Flower Account ID of the user starting the run, if authentication is
+            enabled.
 
         Returns
         -------