From 7fb3e3995db4cc117e2f5ba45c9fad4a687d0a40 Mon Sep 17 00:00:00 2001 From: Nan Date: Sun, 20 Sep 2026 01:53:52 +0000 Subject: [PATCH 1/3] metrics(async): report selected policy provenance --- docs/user-guide/fully-async.md | 27 +++++++- miles/rollout/filter_hub/common_filters.py | 65 +++++++++++++++++-- miles/rollout/fully_async_data_buffer.py | 64 +++++++++++++++++- tests/fast/rollout/test_filters.py | 25 +++++++ .../fast/rollout/test_fully_async_rollout.py | 38 ++++++++++- 5 files changed, 207 insertions(+), 12 deletions(-) diff --git a/docs/user-guide/fully-async.md b/docs/user-guide/fully-async.md index 375bffecb46..95ff56ff6a7 100644 --- a/docs/user-guide/fully-async.md +++ b/docs/user-guide/fully-async.md @@ -308,13 +308,34 @@ rollout/fully_async/queue_size rollout/fully_async/aborted_groups_filtered rollout/fully_async/stale_groups_filtered rollout/fully_async/avg_staleness, rollout/fully_async/max_staleness +rollout/fully_async/stale_filtered_avg_staleness, rollout/fully_async/stale_filtered_max_staleness +rollout/fully_async/avg_post_generation_staleness, rollout/fully_async/max_post_generation_staleness +rollout/fully_async/avg_generation_version_span, rollout/fully_async/max_generation_version_span +rollout/fully_async/token_weighted_staleness +rollout/fully_async/weight_version_sample_coverage rollout/fully_async/buffer_avg_staleness, rollout/fully_async/buffer_max_staleness rollout/dynamic_filter/drop_ ``` -The `avg_staleness` and `max_staleness` pair covers the groups training actually -consumed, while the `buffer_` pair covers the groups still sitting in the buffer when -the step drained it. +Staleness is measured in published rollout weight versions: with +`--update-weights-interval 1`, one version normally corresponds to one optimizer step; +otherwise it does not. The metrics separate three parts of policy provenance: + +- `avg_staleness` and `max_staleness` are the gap from the current version to the + oldest version in each selected group. The `stale_filtered_` pair reports the same + quantity only for groups rejected by `--max-weight-staleness`. +- `post_generation_staleness` is the gap from the current version to the newest version + in each selected group. `generation_version_span` is the newest minus oldest version + within the group. Together they distinguish aging after generation from updates that + landed while a trajectory was being generated. +- `token_weighted_staleness` averages `current_version - span.version`, weighting each + numeric version span by its number of generated tokens. It therefore describes the + typical version gap of stamped generation tokens rather than the oldest token only. + +`weight_version_sample_coverage` is the fraction of selected samples with at least one +numeric version span. Missing provenance is excluded from lag averages, not treated as +zero. The `buffer_` pair covers groups still sitting in the buffer when the step drained +it. A `queue_size` pinned at zero means rollout is the bottleneck, so scale rollout capacity or lower per-sample generation cost. A `queue_size` pinned at capacity means training is diff --git a/miles/rollout/filter_hub/common_filters.py b/miles/rollout/filter_hub/common_filters.py index 498c63c0102..0725b82bb37 100644 --- a/miles/rollout/filter_hub/common_filters.py +++ b/miles/rollout/filter_hub/common_filters.py @@ -1,4 +1,5 @@ from argparse import Namespace +from dataclasses import dataclass import torch @@ -8,6 +9,64 @@ Group = list[Sample | list[Sample]] +@dataclass(frozen=True) +class GroupWeightVersionStats: + sample_count: int + versioned_sample_count: int + versioned_token_count: int + token_version_sum: int + oldest_version: int | None + newest_version: int | None + + def oldest_lag(self, current_version: int | None) -> int | None: + if self.oldest_version is None or current_version is None: + return None + return current_version - self.oldest_version + + def newest_lag(self, current_version: int | None) -> int | None: + if self.newest_version is None or current_version is None: + return None + return current_version - self.newest_version + + def token_weighted_lag(self, current_version: int | None) -> float | None: + if current_version is None or self.versioned_token_count == 0: + return None + return current_version - self.token_version_sum / self.versioned_token_count + + +def group_weight_version_stats(group: Group) -> GroupWeightVersionStats: + sample_count = 0 + versioned_sample_count = 0 + versioned_token_count = 0 + token_version_sum = 0 + oldest_version = None + newest_version = None + + for sample in iter_samples(group): + sample_count += 1 + sample_has_version = False + for span in sample.all_weight_version_spans: + if not str(span.version).isdigit(): + continue + version = int(span.version) + token_count = span.abs_end - span.abs_start + versioned_token_count += token_count + token_version_sum += version * token_count + oldest_version = version if oldest_version is None else min(oldest_version, version) + newest_version = version if newest_version is None else max(newest_version, version) + sample_has_version = True + versioned_sample_count += int(sample_has_version) + + return GroupWeightVersionStats( + sample_count=sample_count, + versioned_sample_count=versioned_sample_count, + versioned_token_count=versioned_token_count, + token_version_sum=token_version_sum, + oldest_version=oldest_version, + newest_version=newest_version, + ) + + def apply_preput_filters(args: Namespace, dynamic_filter, samples: Group, **kwargs) -> FilterOutput: output = apply_aborted_filter(args, samples, **kwargs) if not output.keep: @@ -43,8 +102,4 @@ def apply_reward_nonzero_std_filter(args, samples: list[Sample | list[Sample]], def group_staleness(group: Group, current_version: int | None) -> int | None: - versions = [version for sample in iter_samples(group) if (version := sample.oldest_weight_version) is not None] - oldest = min(versions) if versions else None - if oldest is None or current_version is None: - return None - return current_version - oldest + return group_weight_version_stats(group).oldest_lag(current_version) diff --git a/miles/rollout/fully_async_data_buffer.py b/miles/rollout/fully_async_data_buffer.py index 58fd5cd8366..43ac0929955 100644 --- a/miles/rollout/fully_async_data_buffer.py +++ b/miles/rollout/fully_async_data_buffer.py @@ -15,7 +15,13 @@ from dataclasses import dataclass from miles.rollout.filter_hub.base_types import MetricGatherer, call_dynamic_filter -from miles.rollout.filter_hub.common_filters import apply_aborted_filter, apply_missing_reward_filter, group_staleness +from miles.rollout.filter_hub.common_filters import ( + GroupWeightVersionStats, + apply_aborted_filter, + apply_missing_reward_filter, + group_staleness, + group_weight_version_stats, +) from miles.utils.function_registry import load_function from miles.utils.types import Sample @@ -109,6 +115,13 @@ def __init__(self, input: DataBufferConstructorInput): self._metric_aborted_groups = 0 self._metric_stale_groups = 0 self._metric_consumed_staleness: list[int] = [] + self._metric_stale_group_staleness: list[int] = [] + self._metric_selected_newest_lag: list[int] = [] + self._metric_selected_version_span: list[int] = [] + self._metric_selected_token_lag_sum = 0.0 + self._metric_selected_versioned_tokens = 0 + self._metric_selected_samples = 0 + self._metric_selected_versioned_samples = 0 async def put(self, input: DataBufferInput) -> None: if not self._preput_filter(input): @@ -148,16 +161,37 @@ async def get(self, current_version: int | None = None, **_) -> DataBufferInput: entry = self._buffer.pop(0) self._cond.notify_all() # wake producers blocked on a full buffer - staleness = group_staleness(entry.group, current_version) + version_stats = group_weight_version_stats(entry.group) + staleness = version_stats.oldest_lag(current_version) if staleness is not None: - self._metric_consumed_staleness.append(staleness) if self._args.max_weight_staleness is not None and staleness > self._args.max_weight_staleness: logger.info(f"Filtered stale group ({staleness=} > max={self._args.max_weight_staleness})") self._metric_stale_groups += 1 + self._metric_stale_group_staleness.append(staleness) self._unused_handler_fn(entry.prompt_group) continue + self._metric_consumed_staleness.append(staleness) + self._record_selected_version_stats(version_stats, current_version) return entry + def _record_selected_version_stats( + self, + stats: GroupWeightVersionStats, + current_version: int | None, + ) -> None: + self._metric_selected_samples += stats.sample_count + self._metric_selected_versioned_samples += stats.versioned_sample_count + + newest_lag = stats.newest_lag(current_version) + if newest_lag is not None: + self._metric_selected_newest_lag.append(newest_lag) + self._metric_selected_version_span.append(stats.newest_version - stats.oldest_version) + + token_weighted_lag = stats.token_weighted_lag(current_version) + if token_weighted_lag is not None: + self._metric_selected_token_lag_sum += token_weighted_lag * stats.versioned_token_count + self._metric_selected_versioned_tokens += stats.versioned_token_count + def get_metrics(self) -> dict[str, float]: prefix = "rollout/fully_async/" metrics = { @@ -169,6 +203,23 @@ def get_metrics(self) -> dict[str, float]: if consumed := self._metric_consumed_staleness: metrics[f"{prefix}avg_staleness"] = sum(consumed) / len(consumed) metrics[f"{prefix}max_staleness"] = max(consumed) + if filtered := self._metric_stale_group_staleness: + metrics[f"{prefix}stale_filtered_avg_staleness"] = sum(filtered) / len(filtered) + metrics[f"{prefix}stale_filtered_max_staleness"] = max(filtered) + if newest_lag := self._metric_selected_newest_lag: + metrics[f"{prefix}avg_post_generation_staleness"] = sum(newest_lag) / len(newest_lag) + metrics[f"{prefix}max_post_generation_staleness"] = max(newest_lag) + if version_span := self._metric_selected_version_span: + metrics[f"{prefix}avg_generation_version_span"] = sum(version_span) / len(version_span) + metrics[f"{prefix}max_generation_version_span"] = max(version_span) + if self._metric_selected_versioned_tokens: + metrics[f"{prefix}token_weighted_staleness"] = ( + self._metric_selected_token_lag_sum / self._metric_selected_versioned_tokens + ) + if self._metric_selected_samples: + metrics[f"{prefix}weight_version_sample_coverage"] = ( + self._metric_selected_versioned_samples / self._metric_selected_samples + ) buffered = [ s for entry in self._buffer if (s := group_staleness(entry.group, self._current_version)) is not None ] @@ -178,5 +229,12 @@ def get_metrics(self) -> dict[str, float]: self._metric_gatherer = MetricGatherer() self._metric_consumed_staleness = [] + self._metric_stale_group_staleness = [] + self._metric_selected_newest_lag = [] + self._metric_selected_version_span = [] + self._metric_selected_token_lag_sum = 0.0 + self._metric_selected_versioned_tokens = 0 + self._metric_selected_samples = 0 + self._metric_selected_versioned_samples = 0 self._metric_aborted_groups = self._metric_stale_groups = 0 return metrics diff --git a/tests/fast/rollout/test_filters.py b/tests/fast/rollout/test_filters.py index dfc4066915e..7d26b072162 100644 --- a/tests/fast/rollout/test_filters.py +++ b/tests/fast/rollout/test_filters.py @@ -15,6 +15,7 @@ apply_preput_filters, apply_reward_nonzero_std_filter, group_staleness, + group_weight_version_stats, ) from miles.utils.function_registry import load_function from miles.utils.types import Sample, WeightVersionSpan, WeightVersionsPerCall @@ -248,3 +249,27 @@ def test_group_staleness_uses_oldest_version_across_nested_samples(): assert group_staleness(group, current_version=None) is None assert group_staleness([make_sample()], current_version=10) is None assert group_staleness([make_sample(weight_versions=("12",))], current_version=10) == -2 + + +def test_group_weight_version_stats_preserve_sample_and_token_populations(): + first = make_sample() + first.weight_versions = [ + WeightVersionsPerCall( + spans=[ + WeightVersionSpan(version="4", abs_start=2, abs_end=5), + WeightVersionSpan(version="6", abs_start=5, abs_end=6), + ] + ) + ] + second = make_sample(weight_versions=("not-numeric",)) + + stats = group_weight_version_stats([first, [second]]) + + assert stats.sample_count == 2 + assert stats.versioned_sample_count == 1 + assert stats.versioned_token_count == 4 + assert stats.oldest_lag(current_version=10) == 6 + assert stats.newest_lag(current_version=10) == 4 + assert stats.token_weighted_lag(current_version=10) == 5.5 + assert stats.oldest_lag(current_version=None) is None + assert stats.token_weighted_lag(current_version=None) is None diff --git a/tests/fast/rollout/test_fully_async_rollout.py b/tests/fast/rollout/test_fully_async_rollout.py index 0d3a15caba2..3118ba8faa5 100644 --- a/tests/fast/rollout/test_fully_async_rollout.py +++ b/tests/fast/rollout/test_fully_async_rollout.py @@ -256,7 +256,8 @@ def get_samples_with_fresh_versions(num_samples): assert data_source.recycled == [stale] assert output.metrics["rollout/fully_async/stale_groups_filtered"] == 1 - assert output.metrics["rollout/fully_async/max_staleness"] == 5 + assert output.metrics["rollout/fully_async/max_staleness"] == 0 + assert output.metrics["rollout/fully_async/stale_filtered_max_staleness"] == 5 async def test_stale_group_dropped_by_default(monkeypatch): @@ -587,6 +588,41 @@ async def test_buffer_staleness_metrics(): assert metrics["rollout/fully_async/buffer_max_staleness"] == 4 +async def test_buffer_separates_selected_and_stale_filtered_staleness(): + buffer, _ = make_buffer(max_groups=8, max_staleness=2) + await put_group(buffer, make_group(1, weight_versions=["2", "4"])) + await put_group(buffer, make_group(2, weight_versions=["8", "10"])) + + assert (await buffer.get(current_version=10)).group[0].group_index == 2 + metrics = buffer.get_metrics() + + assert metrics["rollout/fully_async/avg_staleness"] == 2 + assert metrics["rollout/fully_async/max_staleness"] == 2 + assert metrics["rollout/fully_async/stale_filtered_avg_staleness"] == 8 + assert metrics["rollout/fully_async/stale_filtered_max_staleness"] == 8 + assert metrics["rollout/fully_async/avg_post_generation_staleness"] == 0 + assert metrics["rollout/fully_async/avg_generation_version_span"] == 2 + assert metrics["rollout/fully_async/token_weighted_staleness"] == 1 + assert metrics["rollout/fully_async/weight_version_sample_coverage"] == 1 + + +async def test_buffer_reports_missing_weight_version_coverage_without_inventing_lag(): + buffer, _ = make_buffer(max_groups=8) + partial_group = make_group(1, weight_versions=["7"]) + partial_group[1].weight_versions = [] + await put_group(buffer, partial_group) + + await buffer.get(current_version=10) + metrics = buffer.get_metrics() + + assert metrics["rollout/fully_async/weight_version_sample_coverage"] == 0.5 + assert metrics["rollout/fully_async/token_weighted_staleness"] == 3 + + metrics = buffer.get_metrics() + assert "rollout/fully_async/weight_version_sample_coverage" not in metrics + assert "rollout/fully_async/token_weighted_staleness" not in metrics + + class RecordingBuffer(data_buffer.DefaultDataBuffer): constructed_with = None From b6bb79e9fb40101a8dd731f2ad1b39d48a6642c3 Mon Sep 17 00:00:00 2001 From: Nan Date: Sun, 20 Sep 2026 03:49:44 +0000 Subject: [PATCH 2/3] refactor(metrics): keep actionable async provenance --- docs/user-guide/fully-async.md | 5 ++--- miles/rollout/fully_async_data_buffer.py | 6 ------ tests/fast/rollout/test_fully_async_rollout.py | 5 +---- 3 files changed, 3 insertions(+), 13 deletions(-) diff --git a/docs/user-guide/fully-async.md b/docs/user-guide/fully-async.md index 95ff56ff6a7..7541d601b54 100644 --- a/docs/user-guide/fully-async.md +++ b/docs/user-guide/fully-async.md @@ -308,7 +308,6 @@ rollout/fully_async/queue_size rollout/fully_async/aborted_groups_filtered rollout/fully_async/stale_groups_filtered rollout/fully_async/avg_staleness, rollout/fully_async/max_staleness -rollout/fully_async/stale_filtered_avg_staleness, rollout/fully_async/stale_filtered_max_staleness rollout/fully_async/avg_post_generation_staleness, rollout/fully_async/max_post_generation_staleness rollout/fully_async/avg_generation_version_span, rollout/fully_async/max_generation_version_span rollout/fully_async/token_weighted_staleness @@ -322,8 +321,8 @@ Staleness is measured in published rollout weight versions: with otherwise it does not. The metrics separate three parts of policy provenance: - `avg_staleness` and `max_staleness` are the gap from the current version to the - oldest version in each selected group. The `stale_filtered_` pair reports the same - quantity only for groups rejected by `--max-weight-staleness`. + oldest version in each selected group. `stale_groups_filtered` counts groups that + exceeded `--max-weight-staleness`; rejected groups do not enter these averages. - `post_generation_staleness` is the gap from the current version to the newest version in each selected group. `generation_version_span` is the newest minus oldest version within the group. Together they distinguish aging after generation from updates that diff --git a/miles/rollout/fully_async_data_buffer.py b/miles/rollout/fully_async_data_buffer.py index 43ac0929955..17e1a4ff9f9 100644 --- a/miles/rollout/fully_async_data_buffer.py +++ b/miles/rollout/fully_async_data_buffer.py @@ -115,7 +115,6 @@ def __init__(self, input: DataBufferConstructorInput): self._metric_aborted_groups = 0 self._metric_stale_groups = 0 self._metric_consumed_staleness: list[int] = [] - self._metric_stale_group_staleness: list[int] = [] self._metric_selected_newest_lag: list[int] = [] self._metric_selected_version_span: list[int] = [] self._metric_selected_token_lag_sum = 0.0 @@ -167,7 +166,6 @@ async def get(self, current_version: int | None = None, **_) -> DataBufferInput: if self._args.max_weight_staleness is not None and staleness > self._args.max_weight_staleness: logger.info(f"Filtered stale group ({staleness=} > max={self._args.max_weight_staleness})") self._metric_stale_groups += 1 - self._metric_stale_group_staleness.append(staleness) self._unused_handler_fn(entry.prompt_group) continue self._metric_consumed_staleness.append(staleness) @@ -203,9 +201,6 @@ def get_metrics(self) -> dict[str, float]: if consumed := self._metric_consumed_staleness: metrics[f"{prefix}avg_staleness"] = sum(consumed) / len(consumed) metrics[f"{prefix}max_staleness"] = max(consumed) - if filtered := self._metric_stale_group_staleness: - metrics[f"{prefix}stale_filtered_avg_staleness"] = sum(filtered) / len(filtered) - metrics[f"{prefix}stale_filtered_max_staleness"] = max(filtered) if newest_lag := self._metric_selected_newest_lag: metrics[f"{prefix}avg_post_generation_staleness"] = sum(newest_lag) / len(newest_lag) metrics[f"{prefix}max_post_generation_staleness"] = max(newest_lag) @@ -229,7 +224,6 @@ def get_metrics(self) -> dict[str, float]: self._metric_gatherer = MetricGatherer() self._metric_consumed_staleness = [] - self._metric_stale_group_staleness = [] self._metric_selected_newest_lag = [] self._metric_selected_version_span = [] self._metric_selected_token_lag_sum = 0.0 diff --git a/tests/fast/rollout/test_fully_async_rollout.py b/tests/fast/rollout/test_fully_async_rollout.py index 3118ba8faa5..8f923fc8554 100644 --- a/tests/fast/rollout/test_fully_async_rollout.py +++ b/tests/fast/rollout/test_fully_async_rollout.py @@ -257,7 +257,6 @@ def get_samples_with_fresh_versions(num_samples): assert data_source.recycled == [stale] assert output.metrics["rollout/fully_async/stale_groups_filtered"] == 1 assert output.metrics["rollout/fully_async/max_staleness"] == 0 - assert output.metrics["rollout/fully_async/stale_filtered_max_staleness"] == 5 async def test_stale_group_dropped_by_default(monkeypatch): @@ -588,7 +587,7 @@ async def test_buffer_staleness_metrics(): assert metrics["rollout/fully_async/buffer_max_staleness"] == 4 -async def test_buffer_separates_selected_and_stale_filtered_staleness(): +async def test_buffer_reports_selected_policy_provenance(): buffer, _ = make_buffer(max_groups=8, max_staleness=2) await put_group(buffer, make_group(1, weight_versions=["2", "4"])) await put_group(buffer, make_group(2, weight_versions=["8", "10"])) @@ -598,8 +597,6 @@ async def test_buffer_separates_selected_and_stale_filtered_staleness(): assert metrics["rollout/fully_async/avg_staleness"] == 2 assert metrics["rollout/fully_async/max_staleness"] == 2 - assert metrics["rollout/fully_async/stale_filtered_avg_staleness"] == 8 - assert metrics["rollout/fully_async/stale_filtered_max_staleness"] == 8 assert metrics["rollout/fully_async/avg_post_generation_staleness"] == 0 assert metrics["rollout/fully_async/avg_generation_version_span"] == 2 assert metrics["rollout/fully_async/token_weighted_staleness"] == 1 From c51b92cd4af1b7ae68b23c1e52d9901f57501e8c Mon Sep 17 00:00:00 2001 From: Nan Date: Mon, 21 Sep 2026 21:13:24 +0000 Subject: [PATCH 3/3] fix(metrics): decouple generation span from trainer version --- miles/rollout/fully_async_data_buffer.py | 4 +++- tests/fast/rollout/test_fully_async_rollout.py | 12 ++++++++++++ 2 files changed, 15 insertions(+), 1 deletion(-) diff --git a/miles/rollout/fully_async_data_buffer.py b/miles/rollout/fully_async_data_buffer.py index 17e1a4ff9f9..aa1611beaef 100644 --- a/miles/rollout/fully_async_data_buffer.py +++ b/miles/rollout/fully_async_data_buffer.py @@ -180,10 +180,12 @@ def _record_selected_version_stats( self._metric_selected_samples += stats.sample_count self._metric_selected_versioned_samples += stats.versioned_sample_count + if stats.oldest_version is not None and stats.newest_version is not None: + self._metric_selected_version_span.append(stats.newest_version - stats.oldest_version) + newest_lag = stats.newest_lag(current_version) if newest_lag is not None: self._metric_selected_newest_lag.append(newest_lag) - self._metric_selected_version_span.append(stats.newest_version - stats.oldest_version) token_weighted_lag = stats.token_weighted_lag(current_version) if token_weighted_lag is not None: diff --git a/tests/fast/rollout/test_fully_async_rollout.py b/tests/fast/rollout/test_fully_async_rollout.py index 8f923fc8554..508d27f9ae3 100644 --- a/tests/fast/rollout/test_fully_async_rollout.py +++ b/tests/fast/rollout/test_fully_async_rollout.py @@ -603,6 +603,18 @@ async def test_buffer_reports_selected_policy_provenance(): assert metrics["rollout/fully_async/weight_version_sample_coverage"] == 1 +async def test_buffer_reports_generation_span_without_current_version(): + buffer, _ = make_buffer(max_groups=8) + await put_group(buffer, make_group(1, weight_versions=["2", "4"])) + + await buffer.get(current_version=None) + metrics = buffer.get_metrics() + + assert metrics["rollout/fully_async/avg_generation_version_span"] == 2 + assert "rollout/fully_async/avg_post_generation_staleness" not in metrics + assert "rollout/fully_async/token_weighted_staleness" not in metrics + + async def test_buffer_reports_missing_weight_version_coverage_without_inventing_lag(): buffer, _ = make_buffer(max_groups=8) partial_group = make_group(1, weight_versions=["7"])