src/inferscale/prediction.py

from __future__ import annotations

from dataclasses import dataclass


PREDICTOR_SCOPES = {"global_ema", "per_tool_ema"}


@dataclass
class _Estimate:
    mean_s: float
    count: int = 0


class OnlineToolGapPredictor:
    """Tiny online tool-duration predictor used by adaptive KV tiering.

    The predictor deliberately stays transparent: it keeps an exponentially
    weighted moving average (EWMA) of observed tool-call durations. In
    ``per_tool_ema`` mode each tool family gets its own estimate, backed by a
    global EWMA until enough observations have arrived. No future tool duration
    is visible at decision time.
    """

    def __init__(
        self,
        *,
        initial_mean_s: float,
        alpha: float = 0.30,
        min_observations: int = 2,
        scope: str = "per_tool_ema",
    ) -> None:
        if scope not in PREDICTOR_SCOPES:
            raise ValueError(f"Unsupported predictor scope: {scope}")
        self.scope = scope
        self.alpha = min(max(float(alpha), 0.01), 1.0)
        self.min_observations = max(int(min_observations), 1)
        self.global_estimate = _Estimate(max(float(initial_mean_s), 0.0), 0)
        self.per_tool: dict[str, _Estimate] = {}

    @staticmethod
    def _update(estimate: _Estimate, value_s: float, alpha: float) -> None:
        value_s = max(float(value_s), 0.0)
        if estimate.count == 0:
            estimate.mean_s = value_s
        else:
            estimate.mean_s = alpha * value_s + (1.0 - alpha) * estimate.mean_s
        estimate.count += 1

    def predict(self, tool_kind: str) -> tuple[float, str, int]:
        if self.scope == "per_tool_ema":
            estimate = self.per_tool.get(tool_kind)
            if estimate is not None and estimate.count >= self.min_observations:
                return estimate.mean_s, "tool", estimate.count
        return self.global_estimate.mean_s, "global", self.global_estimate.count

    def observe(self, tool_kind: str, duration_s: float) -> None:
        self._update(self.global_estimate, duration_s, self.alpha)
        estimate = self.per_tool.setdefault(tool_kind, _Estimate(self.global_estimate.mean_s, 0))
        self._update(estimate, duration_s, self.alpha)

    def snapshot(self) -> dict:
        return {
            "scope": self.scope,
            "alpha": self.alpha,
            "min_observations": self.min_observations,
            "global": {"mean_s": self.global_estimate.mean_s, "count": self.global_estimate.count},
            "per_tool": {
                key: {"mean_s": value.mean_s, "count": value.count}
                for key, value in sorted(self.per_tool.items())
            },
        }