flwr 1.20.0__py3-none-any.whl → 1.22.0__py3-none-any.whl

flwr/server/superlink/serverappio/serverappio_servicer.py

@@ -43,6 +43,8 @@ from flwr.common.serde import (
 from flwr.common.typing import Fab, RunStatus
 from flwr.proto import serverappio_pb2_grpc  # pylint: disable=E0611
 from flwr.proto.appio_pb2 import (  # pylint: disable=E0611
+    ListAppsToLaunchRequest,
+    ListAppsToLaunchResponse,
     PullAppInputsRequest,
     PullAppInputsResponse,
     PullAppMessagesRequest,
@@ -51,6 +53,8 @@ from flwr.proto.appio_pb2 import (  # pylint: disable=E0611
     PushAppMessagesResponse,
     PushAppOutputsRequest,
     PushAppOutputsResponse,
+    RequestTokenRequest,
+    RequestTokenResponse,
 )
 from flwr.proto.fab_pb2 import GetFabRequest, GetFabResponse  # pylint: disable=E0611
 from flwr.proto.heartbeat_pb2 import (  # pylint: disable=E0611
@@ -104,6 +108,42 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
         self.objectstore_factory = objectstore_factory
         self.lock = threading.RLock()
 
+    def ListAppsToLaunch(
+        self,
+        request: ListAppsToLaunchRequest,
+        context: grpc.ServicerContext,
+    ) -> ListAppsToLaunchResponse:
+        """Get run IDs with pending messages."""
+        log(DEBUG, "ServerAppIoServicer.ListAppsToLaunch")
+
+        # Initialize state connection
+        state = self.state_factory.state()
+
+        # Get IDs of runs in pending status
+        run_ids = state.get_run_ids(flwr_aid=None)
+        pending_run_ids = []
+        for run_id, status in state.get_run_status(run_ids).items():
+            if status.status == Status.PENDING:
+                pending_run_ids.append(run_id)
+
+        # Return run IDs
+        return ListAppsToLaunchResponse(run_ids=pending_run_ids)
+
+    def RequestToken(
+        self, request: RequestTokenRequest, context: grpc.ServicerContext
+    ) -> RequestTokenResponse:
+        """Request token."""
+        log(DEBUG, "ServerAppIoServicer.RequestToken")
+
+        # Initialize state connection
+        state = self.state_factory.state()
+
+        # Attempt to create a token for the provided run ID
+        token = state.create_token(request.run_id)
+
+        # Return the token
+        return RequestTokenResponse(token=token or "")
+
     def GetNodes(
         self, request: GetNodesRequest, context: grpc.ServicerContext
     ) -> GetNodesResponse:
@@ -289,14 +329,11 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
         # Init access to LinkState
         state = self.state_factory.state()
 
+        # Validate the token
+        run_id = self._verify_token(request.token, context)
+
         # Lock access to LinkState, preventing obtaining the same pending run_id
         with self.lock:
-            # Attempt getting the run_id of a pending run
-            run_id = state.get_pending_run_id()
-            # If there's no pending run, return an empty response
-            if run_id is None:
-                return PullAppInputsResponse()
-
             # Init access to Ffs
             ffs = self.ffs_factory.ffs()
 
@@ -327,6 +364,9 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
         """Push ServerApp process outputs."""
         log(DEBUG, "ServerAppIoServicer.PushAppOutputs")
 
+        # Validate the token
+        run_id = self._verify_token(request.token, context)
+
         # Init state and store
         state = self.state_factory.state()
         store = self.objectstore_factory.store()
@@ -341,6 +381,9 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
         )
 
         state.set_serverapp_context(request.run_id, context_from_proto(request.context))
+
+        # Remove the token
+        state.delete_token(run_id)
         return PushAppOutputsResponse()
 
     def UpdateRunStatus(
@@ -508,6 +551,18 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
 
         return ConfirmMessageReceivedResponse()
 
+    def _verify_token(self, token: str, context: grpc.ServicerContext) -> int:
+        """Verify the token and return the associated run ID."""
+        state = self.state_factory.state()
+        run_id = state.get_run_id_by_token(token)
+        if run_id is None or not state.verify_token(run_id, token):
+            context.abort(
+                grpc.StatusCode.PERMISSION_DENIED,
+                "Invalid token.",
+            )
+            raise RuntimeError("This line should never be reached.")
+        return run_id
+
 
 def _raise_if(validation_error: bool, request_name: str, detail: str) -> None:
     """Raise a `ValueError` with a detailed message if a validation error occurs."""

flwr/server/superlink/simulation/simulationio_servicer.py

@@ -34,6 +34,16 @@ from flwr.common.serde import (
 )
 from flwr.common.typing import Fab, RunStatus
 from flwr.proto import simulationio_pb2_grpc
+from flwr.proto.appio_pb2 import (  # pylint: disable=E0611
+    ListAppsToLaunchRequest,
+    ListAppsToLaunchResponse,
+    PullAppInputsRequest,
+    PullAppInputsResponse,
+    PushAppOutputsRequest,
+    PushAppOutputsResponse,
+    RequestTokenRequest,
+    RequestTokenResponse,
+)
 from flwr.proto.heartbeat_pb2 import (  # pylint: disable=E0611
     SendAppHeartbeatRequest,
     SendAppHeartbeatResponse,
@@ -45,17 +55,13 @@ from flwr.proto.log_pb2 import (  # pylint: disable=E0611
 from flwr.proto.run_pb2 import (  # pylint: disable=E0611
     GetFederationOptionsRequest,
     GetFederationOptionsResponse,
+    GetRunRequest,
+    GetRunResponse,
     GetRunStatusRequest,
     GetRunStatusResponse,
     UpdateRunStatusRequest,
     UpdateRunStatusResponse,
 )
-from flwr.proto.simulationio_pb2 import (  # pylint: disable=E0611
-    PullSimulationInputsRequest,
-    PullSimulationInputsResponse,
-    PushSimulationOutputsRequest,
-    PushSimulationOutputsResponse,
-)
 from flwr.server.superlink.linkstate import LinkStateFactory
 from flwr.server.superlink.utils import abort_if
 from flwr.supercore.ffs import FfsFactory
@@ -71,23 +77,73 @@ class SimulationIoServicer(simulationio_pb2_grpc.SimulationIoServicer):
         self.ffs_factory = ffs_factory
         self.lock = threading.RLock()
 
-    def PullSimulationInputs(
-        self, request: PullSimulationInputsRequest, context: ServicerContext
-    ) -> PullSimulationInputsResponse:
+    def ListAppsToLaunch(
+        self,
+        request: ListAppsToLaunchRequest,
+        context: grpc.ServicerContext,
+    ) -> ListAppsToLaunchResponse:
+        """Get run IDs with pending messages."""
+        log(DEBUG, "SimulationIoServicer.ListAppsToLaunch")
+
+        # Initialize state connection
+        state = self.state_factory.state()
+
+        # Get IDs of runs in pending status
+        run_ids = state.get_run_ids(flwr_aid=None)
+        pending_run_ids = []
+        for run_id, status in state.get_run_status(run_ids).items():
+            if status.status == Status.PENDING:
+                pending_run_ids.append(run_id)
+
+        # Return run IDs
+        return ListAppsToLaunchResponse(run_ids=pending_run_ids)
+
+    def RequestToken(
+        self, request: RequestTokenRequest, context: grpc.ServicerContext
+    ) -> RequestTokenResponse:
+        """Request token."""
+        log(DEBUG, "SimulationIoServicer.RequestToken")
+
+        # Initialize state connection
+        state = self.state_factory.state()
+
+        # Attempt to create a token for the provided run ID
+        token = state.create_token(request.run_id)
+
+        # Return the token
+        return RequestTokenResponse(token=token or "")
+
+    def GetRun(
+        self, request: GetRunRequest, context: grpc.ServicerContext
+    ) -> GetRunResponse:
+        """Get run information."""
+        log(DEBUG, "SimulationIoServicer.GetRun")
+
+        # Init state
+        state = self.state_factory.state()
+
+        # Retrieve run information
+        run = state.get_run(request.run_id)
+
+        if run is None:
+            return GetRunResponse()
+
+        return GetRunResponse(run=run_to_proto(run))
+
+    def PullAppInputs(
+        self, request: PullAppInputsRequest, context: ServicerContext
+    ) -> PullAppInputsResponse:
         """Pull SimultionIo process inputs."""
         log(DEBUG, "SimultionIoServicer.SimultionIoInputs")
         # Init access to LinkState and Ffs
         state = self.state_factory.state()
         ffs = self.ffs_factory.ffs()
 
+        # Validate the token
+        run_id = self._verify_token(request.token, context)
+
         # Lock access to LinkState, preventing obtaining the same pending run_id
         with self.lock:
-            # Attempt getting the run_id of a pending run
-            run_id = state.get_pending_run_id()
-            # If there's no pending run, return an empty response
-            if run_id is None:
-                return PullSimulationInputsResponse()
-
             # Retrieve Context, Run and Fab for the run_id
             serverapp_ctxt = state.get_serverapp_context(run_id)
             run = state.get_run(run_id)
@@ -99,7 +155,7 @@ class SimulationIoServicer(simulationio_pb2_grpc.SimulationIoServicer):
                 # Update run status to STARTING
                 if state.update_run_status(run_id, RunStatus(Status.STARTING, "", "")):
                     log(INFO, "Starting run %d", run_id)
-                    return PullSimulationInputsResponse(
+                    return PullAppInputsResponse(
                         context=context_to_proto(serverapp_ctxt),
                         run=run_to_proto(run),
                         fab=fab_to_proto(fab),
@@ -109,11 +165,16 @@ class SimulationIoServicer(simulationio_pb2_grpc.SimulationIoServicer):
         # or if the status cannot be updated to STARTING
         raise RuntimeError(f"Failed to start run {run_id}")
 
-    def PushSimulationOutputs(
-        self, request: PushSimulationOutputsRequest, context: ServicerContext
-    ) -> PushSimulationOutputsResponse:
+    def PushAppOutputs(
+        self, request: PushAppOutputsRequest, context: ServicerContext
+    ) -> PushAppOutputsResponse:
         """Push Simulation process outputs."""
-        log(DEBUG, "SimultionIoServicer.PushSimulationOutputs")
+        log(DEBUG, "SimultionIoServicer.PushAppOutputs")
+
+        # Validate the token
+        run_id = self._verify_token(request.token, context)
+
+        # Init access to LinkState
         state = self.state_factory.state()
 
         # Abort if the run is not running
@@ -126,7 +187,10 @@ class SimulationIoServicer(simulationio_pb2_grpc.SimulationIoServicer):
         )
 
         state.set_serverapp_context(request.run_id, context_from_proto(request.context))
-        return PushSimulationOutputsResponse()
+
+        # Remove the token
+        state.delete_token(run_id)
+        return PushAppOutputsResponse()
 
     def UpdateRunStatus(
         self, request: UpdateRunStatusRequest, context: grpc.ServicerContext
@@ -208,3 +272,15 @@ class SimulationIoServicer(simulationio_pb2_grpc.SimulationIoServicer):
         )
 
         return SendAppHeartbeatResponse(success=success)
+
+    def _verify_token(self, token: str, context: grpc.ServicerContext) -> int:
+        """Verify the token and return the associated run ID."""
+        state = self.state_factory.state()
+        run_id = state.get_run_id_by_token(token)
+        if run_id is None or not state.verify_token(run_id, token):
+            context.abort(
+                grpc.StatusCode.PERMISSION_DENIED,
+                "Invalid token.",
+            )
+            raise RuntimeError("This line should never be reached.")
+        return run_id

flwr/serverapp/__init__.py

@@ -13,3 +13,15 @@
 # limitations under the License.
 # ==============================================================================
 """Public Flower ServerApp APIs."""
+
+
+from flwr.server.grid import Grid
+from flwr.server.server_app import ServerApp
+
+from . import strategy
+
+__all__ = [
+    "Grid",
+    "ServerApp",
+    "strategy",
+]

flwr/serverapp/exception.py

@@ -0,0 +1,38 @@
+# 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 ServerApp exceptions."""
+
+
+from flwr.app.exception import AppExitException
+from flwr.common.exit import ExitCode
+
+
+class InconsistentMessageReplies(AppExitException):
+    """Exception triggered when replies are inconsistent and therefore aggregation must
+    be skipped."""
+
+    exit_code = ExitCode.SERVERAPP_STRATEGY_PRECONDITION_UNMET
+
+    def __init__(self, reason: str):
+        super().__init__(reason)
+
+
+class AggregationError(AppExitException):
+    """Exception triggered when aggregation fails."""
+
+    exit_code = ExitCode.SERVERAPP_STRATEGY_AGGREGATION_ERROR
+
+    def __init__(self, reason: str):
+        super().__init__(reason)

flwr/serverapp/strategy/__init__.py

@@ -0,0 +1,64 @@
+# 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.
+# ==============================================================================
+"""ServerApp strategies."""
+
+
+from .bulyan import Bulyan
+from .dp_adaptive_clipping import (
+    DifferentialPrivacyClientSideAdaptiveClipping,
+    DifferentialPrivacyServerSideAdaptiveClipping,
+)
+from .dp_fixed_clipping import (
+    DifferentialPrivacyClientSideFixedClipping,
+    DifferentialPrivacyServerSideFixedClipping,
+)
+from .fedadagrad import FedAdagrad
+from .fedadam import FedAdam
+from .fedavg import FedAvg
+from .fedavgm import FedAvgM
+from .fedmedian import FedMedian
+from .fedprox import FedProx
+from .fedtrimmedavg import FedTrimmedAvg
+from .fedxgb_bagging import FedXgbBagging
+from .fedxgb_cyclic import FedXgbCyclic
+from .fedyogi import FedYogi
+from .krum import Krum
+from .multikrum import MultiKrum
+from .qfedavg import QFedAvg
+from .result import Result
+from .strategy import Strategy
+
+__all__ = [
+    "Bulyan",
+    "DifferentialPrivacyClientSideAdaptiveClipping",
+    "DifferentialPrivacyClientSideFixedClipping",
+    "DifferentialPrivacyServerSideAdaptiveClipping",
+    "DifferentialPrivacyServerSideFixedClipping",
+    "FedAdagrad",
+    "FedAdam",
+    "FedAvg",
+    "FedAvgM",
+    "FedMedian",
+    "FedProx",
+    "FedTrimmedAvg",
+    "FedXgbBagging",
+    "FedXgbCyclic",
+    "FedYogi",
+    "Krum",
+    "MultiKrum",
+    "QFedAvg",
+    "Result",
+    "Strategy",
+]

flwr/serverapp/strategy/bulyan.py

@@ -0,0 +1,238 @@
+# 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.
+# ==============================================================================
+"""Bulyan [El Mhamdi et al., 2018] strategy.
+
+Paper: arxiv.org/abs/1802.07927
+"""
+
+
+from collections import OrderedDict
+from collections.abc import Iterable
+from logging import INFO, WARN
+from typing import Callable, Optional, cast
+
+import numpy as np
+
+from flwr.common import (
+    Array,
+    ArrayRecord,
+    Message,
+    MetricRecord,
+    NDArrays,
+    RecordDict,
+    log,
+)
+
+from .fedavg import FedAvg
+from .multikrum import select_multikrum
+
+
+# pylint: disable=too-many-instance-attributes
+class Bulyan(FedAvg):
+    """Bulyan strategy.
+
+    Implementation based on https://arxiv.org/abs/1802.07927.
+
+    Parameters
+    ----------
+    fraction_train : float (default: 1.0)
+        Fraction of nodes used during training. In case `min_train_nodes`
+        is larger than `fraction_train * total_connected_nodes`, `min_train_nodes`
+        will still be sampled.
+    fraction_evaluate : float (default: 1.0)
+        Fraction of nodes used during validation. In case `min_evaluate_nodes`
+        is larger than `fraction_evaluate * total_connected_nodes`,
+        `min_evaluate_nodes` will still be sampled.
+    min_train_nodes : int (default: 2)
+        Minimum number of nodes used during training.
+    min_evaluate_nodes : int (default: 2)
+        Minimum number of nodes used during validation.
+    min_available_nodes : int (default: 2)
+        Minimum number of total nodes in the system.
+    num_malicious_nodes : int (default: 0)
+        Number of malicious nodes in the system.
+    weighted_by_key : str (default: "num-examples")
+        The key within each MetricRecord whose value is used as the weight when
+        computing weighted averages for MetricRecords.
+    arrayrecord_key : str (default: "arrays")
+        Key used to store the ArrayRecord when constructing Messages.
+    configrecord_key : str (default: "config")
+        Key used to store the ConfigRecord when constructing Messages.
+    train_metrics_aggr_fn : Optional[callable] (default: None)
+        Function with signature (list[RecordDict], str) -> MetricRecord,
+        used to aggregate MetricRecords from training round replies.
+        If `None`, defaults to `aggregate_metricrecords`, which performs a weighted
+        average using the provided weight factor key.
+    evaluate_metrics_aggr_fn : Optional[callable] (default: None)
+        Function with signature (list[RecordDict], str) -> MetricRecord,
+        used to aggregate MetricRecords from training round replies.
+        If `None`, defaults to `aggregate_metricrecords`, which performs a weighted
+        average using the provided weight factor key.
+    selection_rule : Optional[Callable] (default: None)
+        Function with signature (list[RecordDict], int, int) -> list[RecordDict].
+        The inputs are:
+        - a list of contents from reply messages,
+        - the assumed number of malicious nodes (`num_malicious_nodes`),
+        - the number of nodes to select (`num_nodes_to_select`).
+
+        The function should implement a Byzantine-resilient selection rule that
+        serves as the first step of Bulyan. If None, defaults to `select_multikrum`,
+        which selects nodes according to the Multi-Krum algorithm.
+    """
+
+    # pylint: disable=too-many-arguments,too-many-positional-arguments
+    def __init__(
+        self,
+        fraction_train: float = 1.0,
+        fraction_evaluate: float = 1.0,
+        min_train_nodes: int = 2,
+        min_evaluate_nodes: int = 2,
+        min_available_nodes: int = 2,
+        num_malicious_nodes: int = 0,
+        weighted_by_key: str = "num-examples",
+        arrayrecord_key: str = "arrays",
+        configrecord_key: str = "config",
+        train_metrics_aggr_fn: Optional[
+            Callable[[list[RecordDict], str], MetricRecord]
+        ] = None,
+        evaluate_metrics_aggr_fn: Optional[
+            Callable[[list[RecordDict], str], MetricRecord]
+        ] = None,
+        selection_rule: Optional[
+            Callable[[list[RecordDict], int, int], list[RecordDict]]
+        ] = None,
+    ) -> None:
+        super().__init__(
+            fraction_train=fraction_train,
+            fraction_evaluate=fraction_evaluate,
+            min_train_nodes=min_train_nodes,
+            min_evaluate_nodes=min_evaluate_nodes,
+            min_available_nodes=min_available_nodes,
+            weighted_by_key=weighted_by_key,
+            arrayrecord_key=arrayrecord_key,
+            configrecord_key=configrecord_key,
+            train_metrics_aggr_fn=train_metrics_aggr_fn,
+            evaluate_metrics_aggr_fn=evaluate_metrics_aggr_fn,
+        )
+        self.num_malicious_nodes = num_malicious_nodes
+        self.selection_rule = selection_rule or select_multikrum
+
+    def summary(self) -> None:
+        """Log summary configuration of the strategy."""
+        log(INFO, "\t├──> Bulyan settings:")
+        log(INFO, "\t│\t├── Number of malicious nodes: %d", self.num_malicious_nodes)
+        log(INFO, "\t│\t└── Selection rule: %s", self.selection_rule.__name__)
+        super().summary()
+
+    def aggregate_train(
+        self,
+        server_round: int,
+        replies: Iterable[Message],
+    ) -> tuple[Optional[ArrayRecord], Optional[MetricRecord]]:
+        """Aggregate ArrayRecords and MetricRecords in the received Messages."""
+        valid_replies, _ = self._check_and_log_replies(replies, is_train=True)
+
+        # Check if sufficient replies have been received
+        if len(valid_replies) < 4 * self.num_malicious_nodes + 3:
+            log(
+                WARN,
+                "Insufficient replies, skipping Bulyan aggregation: "
+                "Required at least %d (4*num_malicious_nodes + 3), but received %d.",
+                4 * self.num_malicious_nodes + 3,
+                len(valid_replies),
+            )
+            return None, None
+
+        reply_contents = [msg.content for msg in valid_replies]
+
+        # Compute theta and beta
+        theta = len(valid_replies) - 2 * self.num_malicious_nodes
+        beta = theta - 2 * self.num_malicious_nodes
+
+        # Byzantine-resilient selection rule
+        selected_contents = self.selection_rule(
+            reply_contents, self.num_malicious_nodes, theta
+        )
+
+        # Convert each ArrayRecord to a list of NDArray for easier computation
+        key = list(selected_contents[0].array_records.keys())[0]
+        array_keys = list(selected_contents[0][key].keys())
+        selected_ndarrays = [
+            cast(ArrayRecord, ctnt[key]).to_numpy_ndarrays(keep_input=False)
+            for ctnt in selected_contents
+        ]
+
+        # Compute median
+        median_ndarrays = [np.median(arr, axis=0) for arr in zip(*selected_ndarrays)]
+
+        # Aggregate the beta closest weights element-wise
+        aggregated_ndarrays = aggregate_n_closest_weights(
+            median_ndarrays, selected_ndarrays, beta
+        )
+
+        # Convert to ArrayRecord
+        arrays = ArrayRecord(
+            OrderedDict(zip(array_keys, map(Array, aggregated_ndarrays)))
+        )
+
+        # Aggregate MetricRecords
+        metrics = self.train_metrics_aggr_fn(
+            selected_contents,
+            self.weighted_by_key,
+        )
+        return arrays, metrics
+
+
+def aggregate_n_closest_weights(
+    ref_weights: NDArrays, weights_list: list[NDArrays], beta: int
+) -> NDArrays:
+    """Compute the element-wise mean of the `beta` closest weight arrays.
+
+    For each element (i-th coordinate), the output is the average of the
+    `beta` weight arrays that are closest to the reference weights.
+
+    Parameters
+    ----------
+    ref_weights : NDArrays
+        Reference weights used to compute distances.
+    weights_list : list[NDArrays]
+        List of weight arrays (e.g., from selected nodes).
+    beta : int
+        Number of closest weight arrays to include in the averaging.
+
+    Returns
+    -------
+    aggregated_weights : NDArrays
+        Element-wise average of the `beta` closest weight arrays to the
+        reference weights.
+    """
+    aggregated_weights = []
+    for layer_id, ref_layer in enumerate(ref_weights):
+        # Shape: (n_models, *layer_shape)
+        layer_stack = np.stack([weights[layer_id] for weights in weights_list])
+
+        # Compute absolute differences: shape (n_models, *layer_shape)
+        diffs = np.abs(layer_stack - ref_layer)
+
+        # Find indices of `beta` smallest per coordinate
+        idx = np.argpartition(diffs, beta - 1, axis=0)[:beta]
+
+        # Gather the closest weights
+        closest = np.take_along_axis(layer_stack, idx, axis=0)
+
+        # Average them
+        aggregated_weights.append(np.mean(closest, axis=0))
+
+    return aggregated_weights