Hierarchical Regional Aggregation Topologies

from __future__ import annotations

import math
from dataclasses import dataclass, field
from typing import Mapping, Sequence


class RegionTooSmall(RuntimeError):
    """A regional aggregate would expose too few clients to forward."""


@dataclass(frozen=True)
class Topology:
    """A two-tier aggregation plan and the properties that follow from it."""

    clients: int
    fan_in: int
    model_mb: float
    regional_rounds: int = 1
    min_clients_per_region: int = 8

    @property
    def regions(self) -> int:
        return max(1, math.ceil(self.clients / self.fan_in))

    def server_ingress_mb(self) -> float:
        """Traffic reaching the global server per global round."""
        return self.regions * self.model_mb

    def regional_ingress_mb(self) -> float:
        """Traffic reaching ONE regional aggregator per global round."""
        return self.fan_in * self.model_mb * self.regional_rounds

    def flat_ingress_mb(self) -> float:
        return self.clients * self.model_mb

    def saving(self) -> float:
        return 1.0 - self.server_ingress_mb() / self.flat_ingress_mb()


@dataclass
class RegionalAggregator:
    """One middle-tier node. Holds masked updates only if SecAgg reaches the region."""

    region: str
    topology: Topology
    _received: list[tuple[str, float]] = field(default_factory=list)

    def submit(self, client_id: str, masked_norm: float) -> None:
        self._received.append((client_id, masked_norm))

    def forward(self) -> tuple[float, int]:
        """Aggregate and forward, refusing to expose an under-populated region.

        Refusing is the correct default: a regional aggregate over a handful of
        clients is close to an individual update, and forwarding it upward would
        make the middle tier a disclosure point rather than a bandwidth optimisation.
        """
        if len(self._received) < self.topology.min_clients_per_region:
            raise RegionTooSmall(
                f"{self.region}: {len(self._received)} clients, floor is "
                f"{self.topology.min_clients_per_region} — hold or forward masked"
            )
        total = sum(v for _, v in self._received)
        count = len(self._received)
        self._received.clear()
        return total / count, count


def divergence_bound(regional_rounds: int, lr: float, grad_norm: float) -> float:
    """Worst-case drift between two regional models before a global merge.

    Regional specialisation is the point of the middle tier and its risk: after
    `regional_rounds` local steps two regions can be this far apart, and a naive
    average of two badly diverged models is worse than either.
    """
    return 2.0 * regional_rounds * lr * grad_norm

def _validate() -> None:
    flat = Topology(clients=1024, fan_in=1024, model_mb=12.4)
    tiered = Topology(clients=1024, fan_in=8, model_mb=12.4)

    # 1. The tier must actually reduce server ingress by roughly the fan-in.
    assert tiered.server_ingress_mb() < flat.flat_ingress_mb() / 7
    assert 0.85 < tiered.saving() < 0.95, tiered.saving()

    # 2. The saving at the server is paid for at the regional node.
    assert tiered.regional_ingress_mb() == 8 * 12.4

    # 3. More regional rounds increase regional traffic, not server traffic.
    chatty = Topology(clients=1024, fan_in=8, model_mb=12.4, regional_rounds=3)
    assert chatty.server_ingress_mb() == tiered.server_ingress_mb()
    assert chatty.regional_ingress_mb() == 3 * tiered.regional_ingress_mb()

    # 4. An under-populated region must refuse to forward.
    node = RegionalAggregator("islet", tiered)
    for i in range(4):
        node.submit(f"c{i}", 1.0)
    try:
        node.forward()
    except RegionTooSmall as exc:
        assert "floor is 8" in str(exc)
    else:
        raise AssertionError("a small region must refuse to forward")

    # 5. A populated region forwards a mean and its contributor count.
    for i in range(8):
        node.submit(f"d{i}", 2.0)
    mean, count = node.forward()
    assert abs(mean - 2.0) < 1e-9 and count == 12

    # 6. Divergence grows linearly in the number of regional rounds.
    assert abs(divergence_bound(3, 0.01, 1.0) / divergence_bound(1, 0.01, 1.0) - 3.0) < 1e-9

    print("hierarchical topology: all assertions passed")


_validate()