flwr-nightly 1.16.0.dev20250226__tar.gz → 1.16.0.dev20250228__tar.gz

{flwr_nightly-1.16.0.dev20250226 → flwr_nightly-1.16.0.dev20250228}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: flwr-nightly
-Version: 1.16.0.dev20250226
+Version: 1.16.0.dev20250228
 Summary: Flower: A Friendly Federated AI Framework
 Home-page: https://flower.ai
 License: Apache-2.0

{flwr_nightly-1.16.0.dev20250226 → flwr_nightly-1.16.0.dev20250228}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "flwr-nightly"
-version = "1.16.0.dev20250226"
+version = "1.16.0.dev20250228"
 description = "Flower: A Friendly Federated AI Framework"
 license = "Apache-2.0"
 authors = ["The Flower Authors <hello@flower.ai>"]

{flwr_nightly-1.16.0.dev20250226 → flwr_nightly-1.16.0.dev20250228}/src/py/flwr/cli/log.py

@@ -38,6 +38,10 @@ from flwr.proto.exec_pb2_grpc import ExecStub
 from .utils import init_channel, try_obtain_cli_auth_plugin, unauthenticated_exc_handler
 
 
+class AllLogsRetrieved(BaseException):
+    """Raised when all logs are retrieved."""
+
+
 def start_stream(
     run_id: int, channel: grpc.Channel, refresh_period: int = CONN_REFRESH_PERIOD
 ) -> None:
@@ -56,10 +60,10 @@ def start_stream(
         # pylint: disable=E1101
         if e.code() == grpc.StatusCode.NOT_FOUND:
             logger(ERROR, "Invalid run_id `%s`, exiting", run_id)
-        if e.code() == grpc.StatusCode.CANCELLED:
-            pass
         else:
             raise e
+    except AllLogsRetrieved:
+        pass
     finally:
         channel.close()
 
@@ -94,6 +98,7 @@ def stream_logs(
         with unauthenticated_exc_handler():
             for res in stub.StreamLogs(req, timeout=duration):
                 print(res.log_output, end="")
+        raise AllLogsRetrieved()
     except grpc.RpcError as e:
         # pylint: disable=E1101
         if e.code() != grpc.StatusCode.DEADLINE_EXCEEDED:
@@ -108,27 +113,21 @@ def stream_logs(
 def print_logs(run_id: int, channel: grpc.Channel, timeout: int) -> None:
     """Print logs from the beginning of a run."""
     stub = ExecStub(channel)
-    req = StreamLogsRequest(run_id=run_id)
+    req = StreamLogsRequest(run_id=run_id, after_timestamp=0.0)
 
     try:
-        while True:
-            try:
-                with unauthenticated_exc_handler():
-                    # Enforce timeout for graceful exit
-                    for res in stub.StreamLogs(req, timeout=timeout):
-                        print(res.log_output)
-            except grpc.RpcError as e:
-                # pylint: disable=E1101
-                if e.code() == grpc.StatusCode.DEADLINE_EXCEEDED:
-                    break
-                if e.code() == grpc.StatusCode.NOT_FOUND:
-                    logger(ERROR, "Invalid run_id `%s`, exiting", run_id)
-                    break
-                if e.code() == grpc.StatusCode.CANCELLED:
-                    break
-                raise e
-    except KeyboardInterrupt:
-        logger(DEBUG, "Stream interrupted by user")
+        with unauthenticated_exc_handler():
+            # Enforce timeout for graceful exit
+            for res in stub.StreamLogs(req, timeout=timeout):
+                print(res.log_output)
+                break
+    except grpc.RpcError as e:
+        if e.code() == grpc.StatusCode.NOT_FOUND:  # pylint: disable=E1101
+            logger(ERROR, "Invalid run_id `%s`, exiting", run_id)
+        elif e.code() == grpc.StatusCode.DEADLINE_EXCEEDED:  # pylint: disable=E1101
+            pass
+        else:
+            raise e
     finally:
         channel.close()
         logger(DEBUG, "Channel closed")

{flwr_nightly-1.16.0.dev20250226 → flwr_nightly-1.16.0.dev20250228}/src/py/flwr/client/client_app.py

@@ -16,6 +16,8 @@
 
 
 import inspect
+from collections.abc import Iterator
+from contextlib import contextmanager
 from typing import Callable, Optional
 
 from flwr.client.client import Client
@@ -71,6 +73,11 @@ def _inspect_maybe_adapt_client_fn_signature(client_fn: ClientFnExt) -> ClientFn
     return client_fn
 
 
+@contextmanager
+def _empty_lifecycle(_: Context) -> Iterator[None]:
+    yield
+
+
 class ClientAppException(Exception):
     """Exception raised when an exception is raised while executing a ClientApp."""
 
@@ -135,29 +142,31 @@ class ClientApp:
         self._train: Optional[ClientAppCallable] = None
         self._evaluate: Optional[ClientAppCallable] = None
         self._query: Optional[ClientAppCallable] = None
+        self._lifecycle = _empty_lifecycle
 
     def __call__(self, message: Message, context: Context) -> Message:
         """Execute `ClientApp`."""
-        # Execute message using `client_fn`
-        if self._call:
-            return self._call(message, context)
-
-        # Execute message using a new
-        if message.metadata.message_type == MessageType.TRAIN:
-            if self._train:
-                return self._train(message, context)
-            raise ValueError("No `train` function registered")
-        if message.metadata.message_type == MessageType.EVALUATE:
-            if self._evaluate:
-                return self._evaluate(message, context)
-            raise ValueError("No `evaluate` function registered")
-        if message.metadata.message_type == MessageType.QUERY:
-            if self._query:
-                return self._query(message, context)
-            raise ValueError("No `query` function registered")
-
-        # Message type did not match one of the known message types abvoe
-        raise ValueError(f"Unknown message_type: {message.metadata.message_type}")
+        with self._lifecycle(context):
+            # Execute message using `client_fn`
+            if self._call:
+                return self._call(message, context)
+
+            # Execute message using a new
+            if message.metadata.message_type == MessageType.TRAIN:
+                if self._train:
+                    return self._train(message, context)
+                raise ValueError("No `train` function registered")
+            if message.metadata.message_type == MessageType.EVALUATE:
+                if self._evaluate:
+                    return self._evaluate(message, context)
+                raise ValueError("No `evaluate` function registered")
+            if message.metadata.message_type == MessageType.QUERY:
+                if self._query:
+                    return self._query(message, context)
+                raise ValueError("No `query` function registered")
+
+            # Message type did not match one of the known message types abvoe
+            raise ValueError(f"Unknown message_type: {message.metadata.message_type}")
 
     def train(
         self, mods: Optional[list[Mod]] = None
@@ -296,6 +305,66 @@ class ClientApp:
 
         return query_decorator
 
+    def lifecycle(
+        self,
+    ) -> Callable[
+        [Callable[[Context], Iterator[None]]], Callable[[Context], Iterator[None]]
+    ]:
+        """Return a decorator that registers the lifecycle fn with the client app.
+
+        The decorated function should accept a `Context` object and use `yield`
+        to define enter and exit behavior.
+
+        Examples
+        --------
+        >>> app = ClientApp()
+        >>>
+        >>> @app.lifecycle()
+        >>> def lifecycle(context: Context) -> None:
+        >>>     print("Initializing ClientApp")
+        >>>     yield
+        >>>     print("Cleaning up ClientApp")
+        """
+
+        def lifecycle_decorator(
+            lifecycle_fn: Callable[[Context], Iterator[None]]
+        ) -> Callable[[Context], Iterator[None]]:
+            """Register the lifecycle fn with the ServerApp object."""
+            warn_preview_feature("ClientApp-register-lifecycle-function")
+
+            @contextmanager
+            def decorated_lifecycle(context: Context) -> Iterator[None]:
+                # Execute the code before `yield` in lifecycle_fn
+                try:
+                    if not isinstance(it := lifecycle_fn(context), Iterator):
+                        raise StopIteration
+                    next(it)
+                except StopIteration:
+                    raise RuntimeError(
+                        "Lifecycle function should yield at least once."
+                    ) from None
+
+                try:
+                    # Enter the context
+                    yield
+                finally:
+                    try:
+                        # Execute the code after `yield` in lifecycle_fn
+                        next(it)
+                    except StopIteration:
+                        pass
+                    else:
+                        raise RuntimeError("Lifecycle function should only yield once.")
+
+            # Register provided function with the ClientApp object
+            # Ignore mypy error because of different argument names (`_` vs `context`)
+            self._lifecycle = decorated_lifecycle  # type: ignore
+
+            # Return provided function unmodified
+            return lifecycle_fn
+
+        return lifecycle_decorator
+
 
 class LoadClientAppError(Exception):
     """Error when trying to load `ClientApp`."""

{flwr_nightly-1.16.0.dev20250226 → flwr_nightly-1.16.0.dev20250228}/src/py/flwr/common/auth_plugin/auth_plugin.py

@@ -20,6 +20,7 @@ from collections.abc import Sequence
 from pathlib import Path
 from typing import Optional, Union
 
+from flwr.common.typing import UserInfo
 from flwr.proto.exec_pb2_grpc import ExecStub
 
 from ..typing import UserAuthCredentials, UserAuthLoginDetails
@@ -49,7 +50,7 @@ class ExecAuthPlugin(ABC):
     @abstractmethod
     def validate_tokens_in_metadata(
         self, metadata: Sequence[tuple[str, Union[str, bytes]]]
-    ) -> bool:
+    ) -> tuple[bool, Optional[UserInfo]]:
         """Validate authentication tokens in the provided metadata."""
 
     @abstractmethod

{flwr_nightly-1.16.0.dev20250226 → flwr_nightly-1.16.0.dev20250228}/src/py/flwr/common/constant.py

@@ -212,3 +212,19 @@ class AuthType:
     def __new__(cls) -> AuthType:
         """Prevent instantiation."""
         raise TypeError(f"{cls.__name__} cannot be instantiated.")
+
+
+class EventLogWriterType:
+    """Event log writer types."""
+
+    FALSE = "false"
+    STDOUT = "stdout"
+
+    def __new__(cls) -> EventLogWriterType:
+        """Prevent instantiation."""
+        raise TypeError(f"{cls.__name__} cannot be instantiated.")
+
+    @classmethod
+    def choices(cls) -> list[str]:
+        """Return a list of available log writer choices."""
+        return [cls.FALSE, cls.STDOUT]

flwr_nightly-1.16.0.dev20250228/src/py/flwr/common/event_log_plugin/__init__.py

@@ -0,0 +1,26 @@
+# 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.
+# ==============================================================================
+"""Event log plugin components."""
+
+
+from .event_log_plugin import EventLogRequest as EventLogRequest
+from .event_log_plugin import EventLogResponse as EventLogResponse
+from .event_log_plugin import EventLogWriterPlugin as EventLogWriterPlugin
+
+__all__ = [
+    "EventLogRequest",
+    "EventLogResponse",
+    "EventLogWriterPlugin",
+]

flwr_nightly-1.16.0.dev20250228/src/py/flwr/common/event_log_plugin/event_log_plugin.py

@@ -0,0 +1,87 @@
+# 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.
+# ==============================================================================
+"""Abstract class for Flower Event Log Writer Plugin."""
+
+
+from abc import ABC, abstractmethod
+from typing import Union
+
+import grpc
+
+from flwr.common.typing import LogEntry, UserInfo
+from flwr.proto.exec_pb2 import (  # pylint: disable=E0611
+    GetLoginDetailsRequest,
+    GetLoginDetailsResponse,
+    ListRunsRequest,
+    ListRunsResponse,
+    StartRunRequest,
+    StartRunResponse,
+    StopRunRequest,
+    StopRunResponse,
+    StreamLogsRequest,
+    StreamLogsResponse,
+)
+
+# Type variables for request and response messages
+EventLogRequest = Union[
+    StartRunRequest,
+    ListRunsRequest,
+    StreamLogsRequest,
+    StopRunRequest,
+    GetLoginDetailsRequest,
+]
+EventLogResponse = Union[
+    StartRunResponse,
+    ListRunsResponse,
+    StreamLogsResponse,
+    StopRunResponse,
+    GetLoginDetailsResponse,
+]
+
+
+class EventLogWriterPlugin(ABC):
+    """Abstract Flower Event Log Writer Plugin class for ExecServicer."""
+
+    @abstractmethod
+    def __init__(self) -> None:
+        """Abstract constructor."""
+
+    @abstractmethod
+    def compose_log_before_event(  # pylint: disable=too-many-arguments
+        self,
+        request: EventLogRequest,
+        context: grpc.ServicerContext,
+        user_info: UserInfo,
+        method_name: str,
+    ) -> LogEntry:
+        """Compose pre-event log entry from the provided request and context."""
+
+    @abstractmethod
+    def compose_log_after_event(  # pylint: disable=too-many-arguments,R0917
+        self,
+        request: EventLogRequest,
+        context: grpc.ServicerContext,
+        user_info: UserInfo,
+        method_name: str,
+        response: EventLogResponse,
+    ) -> LogEntry:
+        """Compose post-event log entry from the provided response and context."""
+
+    @abstractmethod
+    def write_log(
+        self,
+        log_entry: LogEntry,
+    ) -> None:
+        """Write the event log to the specified data sink."""

{flwr_nightly-1.16.0.dev20250226 → flwr_nightly-1.16.0.dev20250228}/src/py/flwr/common/object_ref.py

@@ -170,7 +170,6 @@ def load_app(  # pylint: disable= too-many-branches
                 module = importlib.import_module(module_str)
             else:
                 module = sys.modules[module_str]
-                _reload_modules(project_dir)
 
         except ModuleNotFoundError as err:
             raise error_type(
@@ -200,15 +199,6 @@ def _unload_modules(project_dir: Path) -> None:
             del sys.modules[name]
 
 
-def _reload_modules(project_dir: Path) -> None:
-    """Reload modules from the project directory."""
-    dir_str = str(project_dir.absolute())
-    for m in list(sys.modules.values()):
-        path: Optional[str] = getattr(m, "__file__", None)
-        if path is not None and path.startswith(dir_str):
-            importlib.reload(m)
-
-
 def _set_sys_path(directory: Optional[Union[str, Path]]) -> None:
     """Set the system path."""
     if directory is None:

{flwr_nightly-1.16.0.dev20250226 → flwr_nightly-1.16.0.dev20250228}/src/py/flwr/common/record/recordset.py

@@ -17,82 +17,79 @@
 
 from __future__ import annotations
 
-from dataclasses import dataclass
-from typing import cast
+from logging import WARN
+from textwrap import indent
+from typing import TypeVar, Union, cast
 
+from ..logger import log
 from .configsrecord import ConfigsRecord
 from .metricsrecord import MetricsRecord
 from .parametersrecord import ParametersRecord
 from .typeddict import TypedDict
 
+RecordType = Union[ParametersRecord, MetricsRecord, ConfigsRecord]
 
-@dataclass
-class RecordSetData:
-    """Inner data container for the RecordSet class."""
+T = TypeVar("T")
 
-    parameters_records: TypedDict[str, ParametersRecord]
-    metrics_records: TypedDict[str, MetricsRecord]
-    configs_records: TypedDict[str, ConfigsRecord]
 
-    def __init__(
-        self,
-        parameters_records: dict[str, ParametersRecord] | None = None,
-        metrics_records: dict[str, MetricsRecord] | None = None,
-        configs_records: dict[str, ConfigsRecord] | None = None,
-    ) -> None:
-        self.parameters_records = TypedDict[str, ParametersRecord](
-            self._check_fn_str, self._check_fn_params
-        )
-        self.metrics_records = TypedDict[str, MetricsRecord](
-            self._check_fn_str, self._check_fn_metrics
+def _check_key(key: str) -> None:
+    if not isinstance(key, str):
+        raise TypeError(
+            f"Expected `{str.__name__}`, but "
+            f"received `{type(key).__name__}` for the key."
         )
-        self.configs_records = TypedDict[str, ConfigsRecord](
-            self._check_fn_str, self._check_fn_configs
+
+
+def _check_value(value: RecordType) -> None:
+    if not isinstance(value, (ParametersRecord, MetricsRecord, ConfigsRecord)):
+        raise TypeError(
+            f"Expected `{ParametersRecord.__name__}`, `{MetricsRecord.__name__}`, "
+            f"or `{ConfigsRecord.__name__}` but received "
+            f"`{type(value).__name__}` for the value."
         )
-        if parameters_records is not None:
-            self.parameters_records.update(parameters_records)
-        if metrics_records is not None:
-            self.metrics_records.update(metrics_records)
-        if configs_records is not None:
-            self.configs_records.update(configs_records)
-
-    def _check_fn_str(self, key: str) -> None:
-        if not isinstance(key, str):
-            raise TypeError(
-                f"Expected `{str.__name__}`, but "
-                f"received `{type(key).__name__}` for the key."
-            )
 
-    def _check_fn_params(self, record: ParametersRecord) -> None:
-        if not isinstance(record, ParametersRecord):
-            raise TypeError(
-                f"Expected `{ParametersRecord.__name__}`, but "
-                f"received `{type(record).__name__}` for the value."
-            )
 
-    def _check_fn_metrics(self, record: MetricsRecord) -> None:
-        if not isinstance(record, MetricsRecord):
-            raise TypeError(
-                f"Expected `{MetricsRecord.__name__}`, but "
-                f"received `{type(record).__name__}` for the value."
-            )
+class _SyncedDict(TypedDict[str, T]):
+    """A synchronized dictionary that mirrors changes to an underlying RecordSet.
 
-    def _check_fn_configs(self, record: ConfigsRecord) -> None:
-        if not isinstance(record, ConfigsRecord):
+    This dictionary ensures that any modifications (set or delete operations)
+    are automatically reflected in the associated `RecordSet`. Only values of
+    the specified `allowed_type` are permitted.
+    """
+
+    def __init__(self, ref_recordset: RecordSet, allowed_type: type[T]) -> None:
+        if not issubclass(
+            allowed_type, (ParametersRecord, MetricsRecord, ConfigsRecord)
+        ):
+            raise TypeError(f"{allowed_type} is not a valid type.")
+        super().__init__(_check_key, self.check_value)
+        self.recordset = ref_recordset
+        self.allowed_type = allowed_type
+
+    def __setitem__(self, key: str, value: T) -> None:
+        super().__setitem__(key, value)
+        self.recordset[key] = cast(RecordType, value)
+
+    def __delitem__(self, key: str) -> None:
+        super().__delitem__(key)
+        del self.recordset[key]
+
+    def check_value(self, value: T) -> None:
+        """Check if value is of expected type."""
+        if not isinstance(value, self.allowed_type):
             raise TypeError(
-                f"Expected `{ConfigsRecord.__name__}`, but "
-                f"received `{type(record).__name__}` for the value."
+                f"Expected `{self.allowed_type.__name__}`, but "
+                f"received `{type(value).__name__}` for the value."
             )
 
 
-class RecordSet:
+class RecordSet(TypedDict[str, RecordType]):
     """RecordSet stores groups of parameters, metrics and configs.
 
-    A :code:`RecordSet` is the unified mechanism by which parameters,
-    metrics and configs can be either stored as part of a
-    `flwr.common.Context <flwr.common.Context.html>`_ in your apps
-    or communicated as part of a
-    `flwr.common.Message <flwr.common.Message.html>`_ between your apps.
+    A :class:`RecordSet` is the unified mechanism by which parameters,
+    metrics and configs can be either stored as part of a :class:`Context`
+    in your apps or communicated as part of a :class:`Message` between
+    your apps.
 
     Parameters
     ----------
@@ -127,12 +124,12 @@ class RecordSet:
     >>>  # We can create a ConfigsRecord
     >>>  c_record = ConfigsRecord({"lr": 0.1, "batch-size": 128})
     >>>  # Adding it to the record_set would look like this
-    >>>  my_recordset.configs_records["my_config"] = c_record
+    >>>  my_recordset["my_config"] = c_record
     >>>
     >>>  # We can create a MetricsRecord following a similar process
     >>>  m_record = MetricsRecord({"accuracy": 0.93, "losses": [0.23, 0.1]})
     >>>  # Adding it to the record_set would look like this
-    >>>  my_recordset.metrics_records["my_metrics"] = m_record
+    >>>  my_recordset["my_metrics"] = m_record
 
     Adding a :code:`ParametersRecord` follows the same steps as above but first,
     the array needs to be serialized and represented as a :code:`flwr.common.Array`.
@@ -151,7 +148,7 @@ class RecordSet:
     >>>  p_record = ParametersRecord({"my_array": arr})
     >>>
     >>>  # Adding it to the record_set would look like this
-    >>>  my_recordset.parameters_records["my_parameters"] = p_record
+    >>>  my_recordset["my_parameters"] = p_record
 
     For additional examples on how to construct each of the records types shown
     above, please refer to the documentation for :code:`ConfigsRecord`,
@@ -164,39 +161,57 @@ class RecordSet:
         metrics_records: dict[str, MetricsRecord] | None = None,
         configs_records: dict[str, ConfigsRecord] | None = None,
     ) -> None:
-        data = RecordSetData(
-            parameters_records=parameters_records,
-            metrics_records=metrics_records,
-            configs_records=configs_records,
-        )
-        self.__dict__["_data"] = data
+        super().__init__(_check_key, _check_value)
+        for key, p_record in (parameters_records or {}).items():
+            self[key] = p_record
+        for key, m_record in (metrics_records or {}).items():
+            self[key] = m_record
+        for key, c_record in (configs_records or {}).items():
+            self[key] = c_record
 
     @property
     def parameters_records(self) -> TypedDict[str, ParametersRecord]:
-        """Dictionary holding ParametersRecord instances."""
-        data = cast(RecordSetData, self.__dict__["_data"])
-        return data.parameters_records
+        """Dictionary holding only ParametersRecord instances."""
+        synced_dict = _SyncedDict[ParametersRecord](self, ParametersRecord)
+        for key, record in self.items():
+            if isinstance(record, ParametersRecord):
+                synced_dict[key] = record
+        return synced_dict
 
     @property
     def metrics_records(self) -> TypedDict[str, MetricsRecord]:
-        """Dictionary holding MetricsRecord instances."""
-        data = cast(RecordSetData, self.__dict__["_data"])
-        return data.metrics_records
+        """Dictionary holding only MetricsRecord instances."""
+        synced_dict = _SyncedDict[MetricsRecord](self, MetricsRecord)
+        for key, record in self.items():
+            if isinstance(record, MetricsRecord):
+                synced_dict[key] = record
+        return synced_dict
 
     @property
     def configs_records(self) -> TypedDict[str, ConfigsRecord]:
-        """Dictionary holding ConfigsRecord instances."""
-        data = cast(RecordSetData, self.__dict__["_data"])
-        return data.configs_records
+        """Dictionary holding only ConfigsRecord instances."""
+        synced_dict = _SyncedDict[ConfigsRecord](self, ConfigsRecord)
+        for key, record in self.items():
+            if isinstance(record, ConfigsRecord):
+                synced_dict[key] = record
+        return synced_dict
 
     def __repr__(self) -> str:
         """Return a string representation of this instance."""
         flds = ("parameters_records", "metrics_records", "configs_records")
-        view = ", ".join([f"{fld}={getattr(self, fld)!r}" for fld in flds])
-        return f"{self.__class__.__qualname__}({view})"
-
-    def __eq__(self, other: object) -> bool:
-        """Compare two instances of the class."""
-        if not isinstance(other, self.__class__):
-            raise NotImplementedError
-        return self.__dict__ == other.__dict__
+        fld_views = [f"{fld}={dict(getattr(self, fld))!r}" for fld in flds]
+        view = indent(",\n".join(fld_views), "  ")
+        return f"{self.__class__.__qualname__}(\n{view}\n)"
+
+    def __setitem__(self, key: str, value: RecordType) -> None:
+        """Set the given key to the given value after type checking."""
+        original_value = self.get(key, None)
+        super().__setitem__(key, value)
+        if original_value is not None and not isinstance(value, type(original_value)):
+            log(
+                WARN,
+                "Key '%s' was overwritten: record of type `%s` replaced with type `%s`",
+                key,
+                type(original_value).__name__,
+                type(value).__name__,
+            )

{flwr_nightly-1.16.0.dev20250226 → flwr_nightly-1.16.0.dev20250228}/src/py/flwr/common/secure_aggregation/quantization.py

@@ -25,7 +25,11 @@ from flwr.common.typing import NDArrayFloat, NDArrayInt
 def _stochastic_round(arr: NDArrayFloat) -> NDArrayInt:
     ret: NDArrayInt = np.ceil(arr).astype(np.int32)
     rand_arr = np.random.rand(*ret.shape)
-    ret[rand_arr < ret - arr] -= 1
+    if len(ret.shape) == 0:
+        if rand_arr < ret - arr:
+            ret -= 1
+    else:
+        ret[rand_arr < ret - arr] -= 1
     return ret