flwr 1.18.0__py3-none-any.whl → 1.19.0__py3-none-any.whl

flwr/supercore/object_store/in_memory_object_store.py

@@ -0,0 +1,229 @@
+# Copyright 2025 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.
+# ==============================================================================
+"""Flower in-memory ObjectStore implementation."""
+
+
+import threading
+from dataclasses import dataclass
+from typing import Optional
+
+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, 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 = []
+        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}")
+            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."""
+        if self.verify:
+            # Verify object_id and object_content match
+            object_id_from_content = get_object_id(object_content)
+            if object_id != object_id_from_content:
+                raise ValueError(f"Object ID {object_id} does not match content hash")
+
+            # Validate object content
+            validate_object_content(content=object_content)
+
+        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
+
+            # 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."""
+        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."""
+        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."""
+        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 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."""
+        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."""
+        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

@@ -0,0 +1,192 @@
+# Copyright 2025 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.
+# ==============================================================================
+"""Flower abstract ObjectStore definition."""
+
+
+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."""
+
+    def __init__(self, message: str):
+        super().__init__(message)
+        self.message = message
+
+    def __str__(self) -> str:
+        """Return formatted exception message string."""
+        return f"NoObjectInStoreError: {self.message}"
+
+
+class ObjectStore(abc.ABC):
+    """Abstract base class for `ObjectStore` implementations.
+
+    This class defines the interface for an object store that can store, retrieve, and
+    delete objects identified by object IDs.
+    """
+
+    @abc.abstractmethod
+    def preregister(self, run_id: int, object_tree: ObjectTree) -> list[str]:
+        """Identify and preregister missing objects in the `ObjectStore`.
+
+        Parameters
+        ----------
+        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 either not previously preregistered
+            in the `ObjectStore`, or were preregistered but are not yet available.
+        """
+
+    @abc.abstractmethod
+    def put(self, object_id: str, object_content: bytes) -> None:
+        """Put an object into the store.
+
+        Parameters
+        ----------
+        object_id : str
+            The object_id under which to store the object. Must be preregistered.
+        object_content : bytes
+            The deflated object to store.
+        """
+
+    @abc.abstractmethod
+    def get(self, object_id: str) -> Optional[bytes]:
+        """Get an object from the store.
+
+        Parameters
+        ----------
+        object_id : str
+            The object_id under which the object is stored.
+
+        Returns
+        -------
+        bytes
+            The object stored under the given object_id.
+        """
+
+    @abc.abstractmethod
+    def delete(self, object_id: str) -> None:
+        """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
+    def clear(self) -> None:
+        """Clear the store.
+
+        This method should remove all objects from the store.
+        """
+
+    @abc.abstractmethod
+    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.
+
+        Parameters
+        ----------
+        msg_object_id : str
+            The object ID of the ``Message``.
+        descendant_ids : list[str]
+            A list of object IDs representing all descendant objects of the ``Message``.
+        """
+
+    @abc.abstractmethod
+    def get_message_descendant_ids(self, msg_object_id: str) -> list[str]:
+        """Retrieve the object IDs of all descendants of a given ``Message``.
+
+        Parameters
+        ----------
+        msg_object_id : str
+            The object ID of the ``Message``.
+
+        Returns
+        -------
+        list[str]
+            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.
+
+        Parameters
+        ----------
+        object_id : str
+            The object_id to check.
+
+        Returns
+        -------
+        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/supercore/object_store/object_store_factory.py

@@ -0,0 +1,44 @@
+# Copyright 2025 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.
+# ==============================================================================
+"""Factory class that creates ObjectStore instances."""
+
+
+from logging import DEBUG
+from typing import Optional
+
+from flwr.common.logger import log
+
+from .in_memory_object_store import InMemoryObjectStore
+from .object_store import ObjectStore
+
+
+class ObjectStoreFactory:
+    """Factory class that creates ObjectStore instances."""
+
+    def __init__(self) -> None:
+        self.store_instance: Optional[ObjectStore] = None
+
+    def store(self) -> ObjectStore:
+        """Return an ObjectStore instance and create it, if necessary.
+
+        Returns
+        -------
+        ObjectStore
+            An ObjectStore instance for storing objects by object_id.
+        """
+        if self.store_instance is None:
+            self.store_instance = InMemoryObjectStore()
+        log(DEBUG, "Using InMemoryObjectStore")
+        return self.store_instance

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_event_log_interceptor.py

@@ -24,7 +24,7 @@ from google.protobuf.message import Message as GrpcMessage
 from flwr.common.event_log_plugin.event_log_plugin import EventLogWriterPlugin
 from flwr.common.typing import LogEntry
 
-from .exec_user_auth_interceptor import shared_user_info
+from .exec_user_auth_interceptor import shared_account_info
 
 
 class ExecEventLogInterceptor(grpc.ServerInterceptor):  # type: ignore
@@ -62,7 +62,7 @@ class ExecEventLogInterceptor(grpc.ServerInterceptor):  # type: ignore
             log_entry = self.log_plugin.compose_log_before_event(
                 request=request,
                 context=context,
-                user_info=shared_user_info.get(),
+                account_info=shared_account_info.get(),
                 method_name=method_name,
             )
             self.log_plugin.write_log(log_entry)
@@ -81,7 +81,7 @@ class ExecEventLogInterceptor(grpc.ServerInterceptor):  # type: ignore
                     log_entry = self.log_plugin.compose_log_after_event(
                         request=request,
                         context=context,
-                        user_info=shared_user_info.get(),
+                        account_info=shared_account_info.get(),
                         method_name=method_name,
                         response=unary_response or error,
                     )
@@ -111,7 +111,7 @@ class ExecEventLogInterceptor(grpc.ServerInterceptor):  # type: ignore
                         log_entry = self.log_plugin.compose_log_after_event(
                             request=request,
                             context=context,
-                            user_info=shared_user_info.get(),
+                            account_info=shared_account_info.get(),
                             method_name=method_name,
                             response=stream_response or error,
                         )

flwr/superexec/exec_grpc.py

@@ -21,7 +21,7 @@ from typing import Optional
 import grpc
 
 from flwr.common import GRPC_MAX_MESSAGE_LENGTH
-from flwr.common.auth_plugin import ExecAuthPlugin
+from flwr.common.auth_plugin import ExecAuthPlugin, ExecAuthzPlugin
 from flwr.common.event_log_plugin import EventLogWriterPlugin
 from flwr.common.grpc import generic_create_grpc_server
 from flwr.common.logger import log
@@ -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,9 +43,11 @@ 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,
+    authz_plugin: Optional[ExecAuthzPlugin] = None,
     event_log_plugin: Optional[EventLogWriterPlugin] = None,
 ) -> grpc.Server:
     """Run Exec API (gRPC, request-response)."""
@@ -53,12 +56,13 @@ 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,
     )
     interceptors: list[grpc.ServerInterceptor] = []
-    if auth_plugin is not None:
-        interceptors.append(ExecUserAuthInterceptor(auth_plugin))
+    if auth_plugin is not None and authz_plugin is not None:
+        interceptors.append(ExecUserAuthInterceptor(auth_plugin, authz_plugin))
     # Event log interceptor must be added after user auth interceptor
     if event_log_plugin is not None:
         interceptors.append(ExecEventLogInterceptor(event_log_plugin))