Fix ATOMesh KV-cache counters / 修复 ATOMesh KV-cache 计数

inferencex-e2e/infx/results/agentic/__init__.py

@@ -231,6 +231,7 @@ def build_result(
         framework=framework,
         records=records,
         server_logs=server_logs,
+        env=env,
     )
 
     if "total_tput_tps" in request_flat and num_gpus > 0:

inferencex-e2e/infx/results/agentic/backends/atom.py

@@ -2,12 +2,24 @@
 
 from __future__ import annotations
 
+import logging
+import math
+from collections.abc import Mapping
 from typing import Any
 
-from infx.results.agentic.common import gauge_stat, normalize_fraction, rate, sum_stat
+from infx.results.agentic.common import (
+    gauge_stat,
+    metric_series,
+    normalize_fraction,
+    rate,
+    series_stat,
+    sum_stat,
+)
 
 from .base import ServerMetricsBackend, counter_int
 
+logger = logging.getLogger(__name__)
+
 
 def _atom_names(suffix: str) -> list[str]:
     """Accept direct worker metrics and Atomesh's colon-normalized form."""
@@ -19,6 +31,98 @@ def _atom_counter_names(stem: str) -> list[str]:
     return [*_atom_names(stem), *_atom_names(f"{stem}_total")]
 
 
+def _counter_values(
+    metrics: dict[str, dict[str, Any]], stem: str, endpoint: str
+) -> dict[tuple, float] | None:
+    names = _atom_counter_names(f"prefix_cache_{stem}_tokens")
+    values = {}
+    for series in metric_series(metrics, names):
+        if series.get("endpoint_url") != endpoint:
+            continue
+        labels = series.get("labels")
+        labels = {} if labels is None else labels
+        value = series_stat(series, preferred_keys=("total",))
+        if (
+            not isinstance(labels, dict)
+            or any(not isinstance(k, str) or not isinstance(v, str) for k, v in labels.items())
+            or value is None
+            or not math.isfinite(value)
+            or value < 0
+        ):
+            return None
+        key = tuple(sorted(labels.items()))
+        if key in values:
+            logger.warning("ATOM %s cache counter duplicated at %s", stem, endpoint)
+            return None
+        values[key] = value
+    return values
+
+
+def _cache_counts(
+    metrics: dict[str, dict[str, Any]],
+    endpoint_roles: Mapping[str, str] | None,
+    disaggregated: bool | None,
+    cpu_offload_disabled: bool = False,
+) -> tuple[float | None, float | None, float | None]:
+    names = [
+        name
+        for stem in ("full", "cached", "offload")
+        for name in _atom_counter_names(f"prefix_cache_{stem}_tokens")
+    ]
+    endpoints = [s.get("endpoint_url") for s in metric_series(metrics, names)]
+    if any(not isinstance(endpoint, str) or not endpoint for endpoint in endpoints):
+        return None, None, None
+    observed = set(endpoints)
+    if endpoint_roles is None:
+        accounted = observed if disaggregated is False and len(observed) == 1 else set()
+    else:
+        # Decode counts the prompt again without serving its prefill cache reuse.
+        accounted = {url for url, role in endpoint_roles.items() if role in {"prefill", "unified"}}
+        if observed - endpoint_roles.keys():
+            accounted = set()
+    if not accounted:
+        logger.warning("ATOM cache endpoint roles missing or ambiguous")
+        return None, None, None
+    full_total = cached_total = offload_total = 0.0
+    cpu_complete = True
+    for endpoint in accounted:
+        full, cached, offload = (
+            _counter_values(metrics, stem, endpoint) for stem in ("full", "cached", "offload")
+        )
+        if (
+            not full
+            or cached is None
+            or full.keys() != cached.keys()
+            or any(cached[key] > full[key] for key in full)
+        ):
+            return None, None, None
+        if cpu_offload_disabled:
+            offload = (
+                dict.fromkeys(full, 0.0)
+                if offload is not None
+                and offload.keys() <= full.keys()
+                and not any(offload.values())
+                else None
+            )
+        full_total += sum(full.values())
+        cached_total += sum(cached.values())
+        if (
+            offload is None
+            or offload.keys() != full.keys()
+            or any(cached[key] + offload[key] > full[key] for key in full)
+        ):
+            cpu_complete = False
+        else:
+            offload_total += sum(offload.values())
+    if full_total <= 0 or not all(math.isfinite(v) for v in (full_total, cached_total)):
+        return None, None, None
+    return (
+        full_total,
+        cached_total,
+        offload_total if cpu_complete and math.isfinite(offload_total) else None,
+    )
+
+
 class AtomBackend(ServerMetricsBackend):
     name = "atom"
 
@@ -34,12 +138,32 @@ def populate(
         metrics: dict[str, dict[str, Any]],
         flat: dict[str, Any],
         nested: dict[str, Any],
+        *,
+        endpoint_roles: Mapping[str, str] | None = None,
+        disaggregated: bool | None = None,
+        cpu_offload_disabled: bool = False,
     ) -> None:
         """Fill token totals, prefix and LMCache hit rates and KV usage from ATOM ``metrics``."""
+        prompt_names = _atom_counter_names("prompt_tokens")
+        prompt_endpoints = {s.get("endpoint_url") for s in metric_series(metrics, prompt_names)}
+        if endpoint_roles is None:
+            accounted = (
+                prompt_endpoints if disaggregated is False and len(prompt_endpoints) == 1 else set()
+            )
+        else:
+            accounted = {
+                url for url, role in endpoint_roles.items() if role in {"prefill", "unified"}
+            }
+            if prompt_endpoints - endpoint_roles.keys():
+                accounted = set()
+        accounted = {url for url in accounted if isinstance(url, str) and url}
+        if accounted - prompt_endpoints:
+            accounted = set()
         prompt_total = sum_stat(
             metrics,
-            _atom_counter_names("prompt_tokens"),
-            preferred_keys=("total", "sum", "max", "avg"),
+            prompt_names,
+            preferred_keys=("total",),
+            series_filter=lambda s: s.get("endpoint_url") in accounted,
         )
         generation_total = sum_stat(
             metrics,
@@ -49,24 +173,17 @@ def populate(
         flat["total_prompt_tokens"] = counter_int(prompt_total)
         flat["total_generation_tokens"] = counter_int(generation_total)
 
-        cached_tokens = sum_stat(
-            metrics,
-            _atom_counter_names("prefix_cache_cached_tokens"),
-            preferred_keys=("total", "sum", "max", "avg"),
-        )
-        full_tokens = sum_stat(
-            metrics,
-            _atom_counter_names("prefix_cache_full_tokens"),
-            preferred_keys=("total", "sum", "max", "avg"),
+        full_tokens, cached_tokens, external_tokens = _cache_counts(
+            metrics, endpoint_roles, disaggregated, cpu_offload_disabled
         )
         cache_hit_rate = rate(cached_tokens, full_tokens)
-        external_tokens = sum_stat(
-            metrics,
-            _atom_counter_names("lmcache_loaded_tokens"),
-            preferred_keys=("total", "sum", "max", "avg"),
-        )
+        # Offload is CPU reuse beyond the GPU prefix; loaded tokens measure transfer volume.
         external_hit_rate = rate(external_tokens, full_tokens)
-        overall_hit_rate = cache_hit_rate
+        overall_hit_rate = (
+            rate(cached_tokens + external_tokens, full_tokens)
+            if cached_tokens is not None and external_tokens is not None
+            else None
+        )
 
         flat["server_gpu_cache_hit_rate"] = cache_hit_rate
         flat["server_cpu_cache_hit_rate"] = external_hit_rate
@@ -90,7 +207,9 @@ def populate(
                 "prefix_cache_hits": cached_tokens,
                 "prefix_cache_queries": full_tokens,
                 "external_prefix_cache_hits": external_tokens,
-                "external_prefix_cache_queries": full_tokens,
+                "external_prefix_cache_queries": (
+                    full_tokens if external_tokens is not None else None
+                ),
             }
         )
         nested["kv_cache"]["gpu_usage_pct"] = flat["gpu_kv_cache_usage_pct"]

inferencex-e2e/infx/results/agentic/server_metrics.py

@@ -2,10 +2,11 @@
 
 from __future__ import annotations
 
-from collections.abc import Iterable
+from collections.abc import Iterable, Mapping
 from typing import Any
 
 from .backends import detect_backend
+from .backends.atom import AtomBackend
 from .backends.base import (
     SERVER_CACHE_FLAT_FIELDS,
     apply_profile_totals,
@@ -14,12 +15,36 @@
 from .common import index_server_metrics
 
 
+def endpoint_roles_from_env(env: Mapping[str, str]) -> dict[str, str] | None:
+    if env.get("SRTCTL_FRONTEND_TYPE") == "dynamo":
+        return {}
+    if env.get("AIPERF_SERVER_METRICS_URLS"):
+        return None
+    groups = (
+        [("unified", env["SRT_AGG_ENDPOINTS"])]
+        if env.get("SRT_AGG_ENDPOINTS")
+        else [
+            ("prefill", env.get("SRT_PREFILL_ENDPOINTS", "")),
+            ("decode", env.get("SRT_DECODE_ENDPOINTS", "")),
+        ]
+    )
+    roles: dict[str, str] = {}
+    for role, endpoints in groups:
+        for endpoint in endpoints.split(","):
+            if endpoint:
+                url = f"http://{endpoint}/metrics"
+                if roles.setdefault(url, role) != role:
+                    return {}
+    return roles or None
+
+
 def compute_server_metrics(
     server_metrics: dict[str, Any],
     *,
     framework: str,
     records: list[dict[str, Any]],
     server_logs: Iterable[str | None] = (),
+    env: Mapping[str, str] | None = None,
 ) -> tuple[dict[str, Any], dict[str, Any], list[str]]:
     """Return flat and nested server metrics plus warnings from aiperf's ``server_metrics`` export.
 
@@ -50,7 +75,22 @@ def compute_server_metrics(
             f"framework={framework.lower()!r}, metric_names=[{preview}]"
         )
     nested["adapter"] = adapter.name
-    adapter.populate(metrics, flat, nested)
+    if isinstance(adapter, AtomBackend):
+        source_env = env or {}
+        adapter.populate(
+            metrics,
+            flat,
+            nested,
+            endpoint_roles=endpoint_roles_from_env(source_env),
+            disaggregated={"true": True, "false": False}.get(source_env.get("DISAGG", "").lower()),
+            cpu_offload_disabled=(source_env.get("KV_OFFLOADING", "").strip().lower() == "none")
+            and not (
+                source_env.get("KV_OFFLOAD_BACKEND")
+                or source_env.get("KV_OFFLOAD_BACKEND_METADATA")
+            ),
+        )
+    else:
+        adapter.populate(metrics, flat, nested)
 
     if framework.lower() == "dynamo-sglang" and adapter.name == "sglang":
         nested["kv_cache"]["gpu_total_tokens"] = None