diff --git a/IMPL.md b/IMPL.md index 443d9a48d..b73865d65 100644 --- a/IMPL.md +++ b/IMPL.md @@ -8,7 +8,58 @@ ROADMAP's "Recently shipped" and clear the relevant block below. ## Current focus -_No in-flight work._ The corpus and measurement apparatus moved to +**`num-examples` carried a batch count, not a sample count.** All four call sites in +`client_app.py` passed `len(trainloader)` / `len(testloader)`; `len()` on a DataLoader is +the number of **batches**. Confirmed against torch: 33 samples and 64 samples both report +2 at `batch_size=32`, as do 500 and 501 at 16. + +That key is not decorative. `flwr` 1.38's `FedAvg` takes `weighted_by_key="num-examples"` +by default, so the batch count was weighting both the adapter aggregate and the reported +loss/accuracy. Because batches are `ceil(n/32)`, the error is a quantisation that +systematically over-weights the smallest partitions — largest exactly under the Dirichlet +skew the testbed exists to study. Fixed to `len(loader.dataset)`. + +Found by re-reading a batch-vs-sample normalisation defect logged against the older +`fl-execution-framework-dev` testbed and checking whether the same shape existed here. It +did. The related finding there — that per-client local test shards are not a global test +set — also applies, and is now a ROADMAP v2 item rather than a silent caveat. + +`fl.round.ess` lands alongside: effective sample size over those same weights, the +generalisable half of the LQR-Fed weight diagnostic from that testbed. It is what makes +this class of bug visible rather than silent — a weighting that quietly concentrates on a +few clients shows up as ESS far below the client count. + +Not portable, and not ported: the LQR-Fed strategy itself (phalanx is FedAvg-only by +scope discipline), the SLSQP-to-closed-form solver, and that repo's CI and smoke +plumbing. + +**Review round.** Four things the first cut got wrong, all now locked by tests: + +- `_num_examples` indexed `msg.content["metrics"]` by literal record name. flwr addresses + the record by type, so a reply naming its MetricRecord anything else raised `KeyError` + *inside* `aggregate_train` — a telemetry read aborting the round it was only meant to + observe. Reproduced against a real `Message`, now reads `metric_records` by type and + returns None rather than raising. +- ESS as `1/Σ(wᵢ/Σw)²` is not float-exact: an even five-way split read + `4.999999999999999`, ten-way `9.999999999999996`. Kish's `(Σwᵢ)²/Σwᵢ²` over the raw + weights with `math.fsum` is exact for every n in 2..32. (Clamping with + `min(ess, n)` does not help — the error runs below n, not above.) +- The four changed `client_app` lines were executed by nothing: no test imports the + module, the CI job named `smoke-test` only runs `flwr build`, and `ty` sees `torch.**` + as `Any`. Extracted `_sample_count` and covered it in `tests/test_client.py`. +- `docs/architecture.md` and `docs/getting-started.md` enumerate the round span + attributes and metric names; both were missing `fl.ess` and, already, `fl.failures`. + +**Unverified here:** `ty check` and `pytest` need the app env, and `ray` publishes no +macOS x86_64 wheel, so `uv sync` cannot build on an Intel Mac. `ruff format --check` and +`ruff check` pass; the rest rides on CI. No end-to-end `flwr run` has exercised the +`client_app` change — `tests/test_client.py` covers the expression, not a live round. + +--- + +## Background + +The corpus and measurement apparatus moved to [`ajbarea/sphragis`](https://github.com/ajbarea/sphragis) on 2026-09-13; see ROADMAP's `corpus` section for why. Open roadmap items here are the v2 observability and v3 breadth lines. diff --git a/ROADMAP.md b/ROADMAP.md index 91b352832..62bb22300 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -64,6 +64,18 @@ use lands. spans land in a single trace, viewable end-to-end in Jaeger. The genuinely novel OTel↔FL piece (Flower ships no such bridge). `phalanx/telemetry.py` `traceparent_for` / `context_from_traceparent`. +- [x] **Aggregation-weight ESS — `fl.round.ess`.** `1/Σwᵢ²` over the `num-examples` + weights FedAvg actually aggregates by: equal to the client count when shares are + even, falling toward 1.0 as one client dominates. Under Dirichlet skew it reports how + much less than `clients` a round really averaged over, which participation counts + cannot show. Emitted as a round metric and an `fl.ess` span attribute. +- [ ] **Global test set / centralized evaluation.** Every client currently evaluates on + a 20% holdout carved from *its own* partition (`task.py` `load_data`), and the round + figure is a `num-examples`-weighted mean of those local shards. Under Dirichlet the + holdout inherits the partition's label skew, so a client scores well by predicting its + majority label — the aggregate accuracy is therefore not comparable across alphas, and + the "round-2 Dirichlet collapse" reading below rests on it. A shared held-out split + evaluated server-side would make the number mean one thing. - [ ] **Round wall-time + comm-cost metrics** — per-round duration histogram and bytes-on-the-wire (adapter payload size), alongside loss/accuracy/participation. - [ ] **Jaeger / OTel-Collector `compose` recipe** — one command to bring up a backend diff --git a/docs/architecture.md b/docs/architecture.md index 5edea61db..e0097d77b 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -51,8 +51,9 @@ so tests can re-initialise between cases. `init_telemetry` chooses an exporter: - otherwise telemetry is recorded but not exported. Server-side, each round emits an `fl.round` span (`fl.round`, `fl.loss`, -`fl.accuracy`, `fl.clients`) and the metrics `fl.round.loss` / `fl.round.accuracy` / -`fl.round.clients`. Client-side, each pass emits an `fl.client.train` or +`fl.accuracy`, `fl.clients`, `fl.ess`, `fl.failures`) and the metrics `fl.round.loss` / +`fl.round.accuracy` / `fl.round.clients` / `fl.round.ess` / `fl.round.failures`. +Client-side, each pass emits an `fl.client.train` or `fl.client.evaluate` span and `fl.client.examples` / `fl.client.loss` metrics. Because the simulation runs clients in separate Ray processes, client spans are diff --git a/docs/getting-started.md b/docs/getting-started.md index 8220d7708..669d8291e 100644 --- a/docs/getting-started.md +++ b/docs/getting-started.md @@ -66,8 +66,9 @@ OTEL_TRACES_EXPORTER=console make run # this is what `make trace` does ``` Each round produces an `fl.round` span (attributes: `fl.round`, `fl.loss`, -`fl.accuracy`, `fl.clients`) and FL metrics (`fl.round.loss`, `fl.round.accuracy`, -`fl.round.clients`); each participating client produces an `fl.client.train` or +`fl.accuracy`, `fl.clients`, `fl.ess`, `fl.failures`) and FL metrics (`fl.round.loss`, +`fl.round.accuracy`, `fl.round.clients`, `fl.round.ess`, `fl.round.failures`); each +participating client produces an `fl.client.train` or `fl.client.evaluate` span and `fl.client.*` metrics. ## Develop diff --git a/phalanx/client_app.py b/phalanx/client_app.py index a77bcb62c..db0fe9867 100644 --- a/phalanx/client_app.py +++ b/phalanx/client_app.py @@ -37,6 +37,16 @@ app = ClientApp() +def _sample_count(loader: Any) -> int: + """Rows behind a loader, which is what FedAvg must weight by. + + ``len(loader)`` counts batches, not rows: 33 rows and 64 rows both report 2 at + ``batch_size=32``. FedAvg takes ``weighted_by_key="num-examples"``, so a batch count + here quantises the adapter aggregate toward the smallest partitions. + """ + return len(loader.dataset) + + def _device() -> torch.device: return torch.device("cuda:0" if torch.cuda.is_available() else "cpu") @@ -86,12 +96,16 @@ def train(msg: Message, context: Context) -> Message: device = _device() model.to(device) loss = train_fn(model, trainloader, epochs=int(cfg["local-epochs"]), device=device) - record_client_metrics(partition_id=partition_id, num_examples=len(trainloader), loss=loss) + record_client_metrics( + partition_id=partition_id, num_examples=_sample_count(trainloader), loss=loss + ) content = RecordDict( { "arrays": ArrayRecord(get_adapter_state(model)), - "metrics": MetricRecord({"num-examples": len(trainloader), "train_loss": loss}), + "metrics": MetricRecord( + {"num-examples": _sample_count(trainloader), "train_loss": loss} + ), } ) return Message(content=content, reply_to=msg) @@ -124,12 +138,14 @@ def evaluate(msg: Message, context: Context) -> Message: device = _device() model.to(device) loss, accuracy = test_fn(model, testloader, device=device) - record_client_metrics(partition_id=partition_id, num_examples=len(testloader), loss=loss) + record_client_metrics( + partition_id=partition_id, num_examples=_sample_count(testloader), loss=loss + ) content = RecordDict( { "metrics": MetricRecord( - {"num-examples": len(testloader), "loss": loss, "accuracy": accuracy} + {"num-examples": _sample_count(testloader), "loss": loss, "accuracy": accuracy} ) } ) diff --git a/phalanx/server_app.py b/phalanx/server_app.py index 74d03cf7d..932df0b4a 100644 --- a/phalanx/server_app.py +++ b/phalanx/server_app.py @@ -3,12 +3,13 @@ ``ObservableFedAvg`` subclasses Flower's ``FedAvg`` and hooks the per-round entry points inside ``strategy.start()``: it counts participating clients in ``aggregate_train`` and, after ``aggregate_evaluate``, emits an ``fl.round`` span -plus aggregated loss/accuracy/participation metrics. Only the LoRA adapters are +plus aggregated loss/accuracy/participation/ESS metrics. Only the LoRA adapters are federated (the initial arrays come from the adapter state, not the full model). """ from __future__ import annotations +import math from collections.abc import Iterable from typing import Any @@ -40,12 +41,49 @@ def _round_summary(metrics: MetricRecord | None) -> tuple[float, float]: return loss, accuracy +def effective_sample_size(weights: Iterable[float]) -> float: + """Clients effectively contributing to the aggregate: Kish's ``(Σwᵢ)² / Σwᵢ²``. + + Equals the client count when every client carries the same weight, falls toward 1.0 + as one client's share dominates, and is NaN when nothing was aggregated. FedAvg + weights by ``num-examples``, so under a skewed partition ESS reports how much less + than ``clients`` the round actually averaged over. + + Measured over the **train** replies, matching ``fl.clients``. Train and evaluate + sample their clients independently, so ``fl.ess`` describes the aggregation that + produced the adapters, not the client set behind ``fl.loss`` / ``fl.accuracy``. + """ + w = [float(x) for x in weights] + total = math.fsum(w) + if not w or total <= 0: + return float("nan") + # Kish's form over the raw weights, not 1/Σ(wᵢ/Σw)²: dividing each term by the total + # first leaves an equal split reading 4.999999999999999 for five clients. + return total * total / math.fsum(x * x for x in w) + + +def _num_examples(msg: Message) -> float | None: + """The sample count a client reported, or None when the reply does not carry one. + + Addresses the record by type rather than by the literal name ``client_app`` happens + to use, the way flwr's own aggregation does — a telemetry read must not be the thing + that aborts a round. MetricRecord values are a broad numeric union, so the cast + reads through Any, as ``_round_summary`` does for loss/accuracy. + """ + record = next(iter(msg.content.metric_records.values()), None) + if record is None or "num-examples" not in record: + return None + metrics: Any = record + return float(metrics["num-examples"]) + + def observe_round( *, server_round: int, metrics: MetricRecord | None, clients: int, failures: int = 0, + ess: float = float("nan"), span: Any | None = None, ) -> None: """Decorate the round span with aggregated metrics + status, then end it. @@ -60,6 +98,7 @@ def observe_round( span.set_attribute("fl.loss", loss) span.set_attribute("fl.accuracy", accuracy) span.set_attribute("fl.clients", clients) + span.set_attribute("fl.ess", ess) span.set_attribute("fl.failures", failures) if failures: # Surface client/worker failures in the trace, not just the participation count. @@ -71,6 +110,7 @@ def observe_round( accuracy=accuracy, clients=clients, failures=failures, + ess=ess, ) span.end() @@ -81,6 +121,7 @@ class ObservableFedAvg(FedAvg): def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) self._round_clients: dict[int, int] = {} + self._round_ess: dict[int, float] = {} self._round_failures: dict[int, int] = {} self._round_spans: dict[int, Any] = {} @@ -108,6 +149,9 @@ def aggregate_train( replies = list(replies) self._round_clients[server_round] = sum(1 for m in replies if not m.has_error()) self._round_failures[server_round] = sum(1 for m in replies if m.has_error()) + # ESS over the same key FedAvg aggregates by, so it describes the actual weights. + counts = (_num_examples(m) for m in replies if not m.has_error()) + self._round_ess[server_round] = effective_sample_size(n for n in counts if n is not None) return super().aggregate_train(server_round, replies) def aggregate_evaluate( @@ -121,6 +165,7 @@ def aggregate_evaluate( metrics=metrics, clients=self._round_clients.pop(server_round, 0), failures=self._round_failures.pop(server_round, 0) + eval_failures, + ess=self._round_ess.pop(server_round, float("nan")), span=self._round_spans.pop(server_round, None), ) return metrics diff --git a/phalanx/telemetry.py b/phalanx/telemetry.py index 707e43d84..7bf9e4c25 100644 --- a/phalanx/telemetry.py +++ b/phalanx/telemetry.py @@ -92,6 +92,7 @@ def init_telemetry( "round_loss": _meter.create_gauge("fl.round.loss"), "round_accuracy": _meter.create_gauge("fl.round.accuracy"), "round_clients": _meter.create_gauge("fl.round.clients"), + "round_ess": _meter.create_gauge("fl.round.ess"), "round_failures": _meter.create_counter("fl.round.failures"), "client_examples": _meter.create_counter("fl.client.examples"), "client_loss": _meter.create_gauge("fl.client.loss"), @@ -174,14 +175,21 @@ def context_from_traceparent(traceparent: str) -> Any: def record_round_metrics( - *, rnd: int, loss: float, accuracy: float, clients: int, failures: int = 0 + *, + rnd: int, + loss: float, + accuracy: float, + clients: int, + failures: int = 0, + ess: float = float("nan"), ) -> None: - """Record aggregated server-round metrics (loss, accuracy, participation, failures).""" + """Record aggregated server-round metrics (loss, accuracy, participation, ESS, failures).""" _ensure_init() attrs = {"fl.round": rnd} _instruments["round_loss"].set(loss, attributes=attrs) _instruments["round_accuracy"].set(accuracy, attributes=attrs) _instruments["round_clients"].set(clients, attributes=attrs) + _instruments["round_ess"].set(ess, attributes=attrs) _instruments["round_failures"].add(failures, attributes=attrs) diff --git a/tests/test_client.py b/tests/test_client.py new file mode 100644 index 000000000..445d2d52d --- /dev/null +++ b/tests/test_client.py @@ -0,0 +1,35 @@ +"""The sample count a client reports as its FedAvg aggregation weight. + +``len(DataLoader)`` is a batch count; FedAvg's ``weighted_by_key="num-examples"`` needs +a row count. These lock the distinction, which no federated run in CI would surface. +""" + +from __future__ import annotations + +from typing import Any + +import torch +from torch.utils.data import DataLoader, TensorDataset + +from phalanx.client_app import _sample_count + + +def _loader(rows: int, batch_size: int = 32) -> DataLoader[Any]: + return DataLoader(TensorDataset(torch.zeros(rows)), batch_size=batch_size) + + +def test_sample_count_reports_rows_not_batches() -> None: + # The pairs that collide under len(loader): both are 2 batches, and 500/501 both 16. + assert _sample_count(_loader(33)) == 33 + assert _sample_count(_loader(64)) == 64 + assert _sample_count(_loader(500)) == 500 + assert _sample_count(_loader(501)) == 501 + + +def test_sample_count_is_independent_of_batch_size() -> None: + # The weight must describe the partition, not how it was chopped up. + assert {_sample_count(_loader(501, bs)) for bs in (1, 7, 32, 512, 1024)} == {501} + + +def test_sample_count_handles_an_empty_partition() -> None: + assert _sample_count(_loader(0)) == 0 diff --git a/tests/test_server.py b/tests/test_server.py index 9f383c93d..e789c4cfe 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -9,12 +9,12 @@ import math from typing import Any -from flwr.app import MetricRecord +from flwr.app import Message, MetricRecord, RecordDict from opentelemetry.sdk.metrics.export import InMemoryMetricReader from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter from opentelemetry.trace import StatusCode -from phalanx.server_app import observe_round +from phalanx.server_app import _num_examples, effective_sample_size, observe_round from phalanx.telemetry import init_telemetry @@ -29,6 +29,11 @@ def _setup() -> tuple[InMemorySpanExporter, InMemoryMetricReader]: return span_exporter, metric_reader +def _reply(content: RecordDict) -> Message: + """A client reply carrying `content` (the shape aggregate_train iterates).""" + return Message(content=content, dst_node_id=0, message_type="train") + + def _attrs(span: Any) -> dict[str, Any]: """A span's attributes as a plain dict (never None).""" return dict(span.attributes or {}) @@ -86,3 +91,64 @@ def test_observe_round_clean_round_is_not_error() -> None: observe_round(server_round=1, metrics=MetricRecord({"loss": 0.5, "accuracy": 0.6}), clients=2) span = next(s for s in span_exporter.get_finished_spans() if s.name == "fl.round") assert span.status.status_code != StatusCode.ERROR + + +def test_effective_sample_size_reports_weight_concentration() -> None: + # Uniform weights average over every client; a dominant client collapses ESS to ~1. + assert effective_sample_size([10, 10, 10, 10]) == 4.0 + assert effective_sample_size([1, 1]) == 2.0 + assert math.isclose(effective_sample_size([999_999, 1]), 1.0, abs_tol=1e-4) + # Scale-invariant: only the shares matter, not the absolute counts. + assert math.isclose(effective_sample_size([3, 1]), effective_sample_size([300, 100])) + # Between the extremes for a skewed but not degenerate split. + assert 1.0 < effective_sample_size([8, 1, 1]) < 3.0 + + +def test_effective_sample_size_is_exact_for_an_even_split() -> None: + # Normalising each term before squaring reads 4.999999999999999 at n=5 and + # 9.999999999999996 at n=10; Kish's form over the raw weights is exact. + for n in range(2, 33): + assert effective_sample_size([10] * n) == float(n), f"inexact at n={n}" + + +def test_effective_sample_size_never_exceeds_the_client_count() -> None: + for weights in ([1, 2, 3], [7, 7, 7, 1], [10] * 9, [5, 4], [1] * 17): + assert effective_sample_size(weights) <= len(weights) + 1e-12 + + +def test_effective_sample_size_is_nan_when_nothing_aggregated() -> None: + assert math.isnan(effective_sample_size([])) + assert math.isnan(effective_sample_size([0, 0])) + + +def test_num_examples_reads_the_record_by_type_not_by_name() -> None: + # client_app names its record "metrics"; nothing guarantees that, and flwr addresses + # it by type. A differently-named record must still yield the weight. + msg = _reply(RecordDict({"whatever-name": MetricRecord({"num-examples": 40.0})})) + assert _num_examples(msg) == 40.0 + + +def test_num_examples_is_none_when_the_reply_carries_no_count() -> None: + # Must not raise: a telemetry read cannot be what aborts a round. + assert _num_examples(_reply(RecordDict({"metrics": MetricRecord({"loss": 0.5})}))) is None + assert _num_examples(_reply(RecordDict({}))) is None + + +def test_observe_round_records_ess() -> None: + span_exporter, metric_reader = _setup() + observe_round( + server_round=1, + metrics=MetricRecord({"loss": 0.5, "accuracy": 0.6}), + clients=2, + ess=1.6, + ) + span = next(s for s in span_exporter.get_finished_spans() if s.name == "fl.round") + assert _attrs(span)["fl.ess"] == 1.6 + assert "fl.round.ess" in _metric_names(metric_reader) + + +def test_observe_round_ess_defaults_to_nan() -> None: + span_exporter, _ = _setup() + observe_round(server_round=1, metrics=None, clients=0) + span = next(s for s in span_exporter.get_finished_spans() if s.name == "fl.round") + assert math.isnan(_attrs(span)["fl.ess"])