288 lines
9.1 KiB
Python
288 lines
9.1 KiB
Python
#!/usr/bin/env python3
|
|
"""Small-data action-response model for the active intervention pilot.
|
|
|
|
The model predicts the paired normalized SLO-goodput treatment effect from a
|
|
source measurement and a full MNS/MBBT action. Telemetry features are direct,
|
|
continuous engine measurements; there is no diagnosis-to-action rule here.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import math
|
|
from dataclasses import dataclass
|
|
from typing import Any, Iterable, Mapping, Sequence
|
|
|
|
import numpy as np
|
|
|
|
|
|
PREFIX_FEATURES = (
|
|
"normalized_slo_goodput",
|
|
"admitted_fraction",
|
|
"completed_over_admitted",
|
|
"completed_pass_rate",
|
|
"completed_fail_fraction_of_total",
|
|
"outstanding_over_admitted",
|
|
"ttft_max_over_slo_max",
|
|
"ttft_mean_over_slo_max",
|
|
"tpot_max_over_slo",
|
|
"tpot_mean_over_slo",
|
|
"admitted_input_tokens_mean_over_limit",
|
|
)
|
|
|
|
TELEMETRY_FEATURES = (
|
|
"scheduler_steps_per_s",
|
|
"batch_size_mean",
|
|
"batch_size_cv",
|
|
"batch_tokens_mean",
|
|
"batch_tokens_cv",
|
|
"decode_batch_size_mean",
|
|
"decode_batch_size_cv",
|
|
"prefill_token_fraction",
|
|
"queue_waiting_mean",
|
|
"queue_running_mean",
|
|
"preemptions_per_step",
|
|
"kv_usage_mean",
|
|
"kv_usage_max",
|
|
"kv_usage_end_minus_start",
|
|
"graph_none_share",
|
|
"graph_full_share",
|
|
"graph_padding_fraction",
|
|
)
|
|
|
|
|
|
def _finite(value: Any, name: str) -> float:
|
|
result = float(value)
|
|
if not math.isfinite(result):
|
|
raise ValueError(f"{name} must be finite")
|
|
return result
|
|
|
|
|
|
def feature_vector(
|
|
example: Mapping[str, Any], *, include_telemetry: bool
|
|
) -> tuple[list[str], np.ndarray]:
|
|
source = example["source"]
|
|
action = example["action"]
|
|
source_log_mns = math.log2(_finite(source["mns"], "source MNS"))
|
|
source_log_mbbt = math.log2(_finite(source["mbbt"], "source MBBT"))
|
|
target_log_mns = math.log2(_finite(action["target_mns"], "target MNS"))
|
|
target_log_mbbt = math.log2(_finite(action["target_mbbt"], "target MBBT"))
|
|
delta_mns = target_log_mns - source_log_mns
|
|
delta_mbbt = target_log_mbbt - source_log_mbbt
|
|
names = [
|
|
"source_log2_mns",
|
|
"source_log2_mbbt",
|
|
"target_log2_mns",
|
|
"target_log2_mbbt",
|
|
"delta_log2_mns",
|
|
"delta_log2_mbbt",
|
|
"delta_product",
|
|
"offered_rate_per_gpu",
|
|
]
|
|
values = [
|
|
source_log_mns,
|
|
source_log_mbbt,
|
|
target_log_mns,
|
|
target_log_mbbt,
|
|
delta_mns,
|
|
delta_mbbt,
|
|
delta_mns * delta_mbbt,
|
|
_finite(source["offered_rate_per_gpu"], "offered rate"),
|
|
]
|
|
for name in PREFIX_FEATURES:
|
|
names.append(f"outcome.{name}")
|
|
values.append(_finite(source["outcome"][name], name))
|
|
if include_telemetry:
|
|
for name in TELEMETRY_FEATURES:
|
|
value = _finite(source["telemetry"][name], name)
|
|
names.extend(
|
|
(
|
|
f"telemetry.{name}",
|
|
f"telemetry.{name}*delta_mns",
|
|
f"telemetry.{name}*delta_mbbt",
|
|
)
|
|
)
|
|
values.extend((value, value * delta_mns, value * delta_mbbt))
|
|
vector = np.asarray(values, dtype=np.float64)
|
|
if not np.all(np.isfinite(vector)):
|
|
raise ValueError("feature vector contains a non-finite value")
|
|
return names, vector
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class RidgeModel:
|
|
feature_names: tuple[str, ...]
|
|
mean: np.ndarray
|
|
scale: np.ndarray
|
|
weights: np.ndarray
|
|
intercept: float
|
|
regularization: float
|
|
|
|
def predict(self, values: np.ndarray) -> float:
|
|
if values.shape != self.mean.shape:
|
|
raise ValueError("ridge prediction feature shape mismatch")
|
|
normalized = (values - self.mean) / self.scale
|
|
return float(self.intercept + normalized @ self.weights)
|
|
|
|
def to_json(self) -> dict[str, Any]:
|
|
return {
|
|
"feature_names": list(self.feature_names),
|
|
"mean": self.mean.tolist(),
|
|
"scale": self.scale.tolist(),
|
|
"weights": self.weights.tolist(),
|
|
"intercept": self.intercept,
|
|
"regularization": self.regularization,
|
|
}
|
|
|
|
@classmethod
|
|
def from_json(cls, payload: Mapping[str, Any]) -> "RidgeModel":
|
|
return cls(
|
|
feature_names=tuple(str(value) for value in payload["feature_names"]),
|
|
mean=np.asarray(payload["mean"], dtype=np.float64),
|
|
scale=np.asarray(payload["scale"], dtype=np.float64),
|
|
weights=np.asarray(payload["weights"], dtype=np.float64),
|
|
intercept=float(payload["intercept"]),
|
|
regularization=float(payload["regularization"]),
|
|
)
|
|
|
|
|
|
def fit_ridge(
|
|
examples: Sequence[Mapping[str, Any]],
|
|
*,
|
|
include_telemetry: bool,
|
|
regularization: float,
|
|
) -> RidgeModel:
|
|
if not examples:
|
|
raise ValueError("ridge fit requires examples")
|
|
if regularization <= 0:
|
|
raise ValueError("ridge regularization must be positive")
|
|
encoded = [
|
|
feature_vector(example, include_telemetry=include_telemetry)
|
|
for example in examples
|
|
]
|
|
names = encoded[0][0]
|
|
if any(item[0] != names for item in encoded):
|
|
raise ValueError("feature names changed across examples")
|
|
x = np.stack([item[1] for item in encoded])
|
|
y = np.asarray(
|
|
[
|
|
_finite(example["target_delta_normalized_goodput"], "target effect")
|
|
for example in examples
|
|
],
|
|
dtype=np.float64,
|
|
)
|
|
mean = x.mean(axis=0)
|
|
scale = x.std(axis=0)
|
|
scale[scale < 1e-12] = 1.0
|
|
normalized = (x - mean) / scale
|
|
intercept = float(y.mean())
|
|
centered = y - intercept
|
|
system = normalized.T @ normalized + regularization * np.eye(x.shape[1])
|
|
weights = np.linalg.solve(system, normalized.T @ centered)
|
|
return RidgeModel(
|
|
feature_names=tuple(names),
|
|
mean=mean,
|
|
scale=scale,
|
|
weights=weights,
|
|
intercept=intercept,
|
|
regularization=regularization,
|
|
)
|
|
|
|
|
|
def fit_jackknife_ensemble(
|
|
examples: Sequence[Mapping[str, Any]],
|
|
*,
|
|
include_telemetry: bool,
|
|
regularization: float,
|
|
group_key: str = "decision_id",
|
|
) -> list[RidgeModel]:
|
|
groups = sorted({str(example[group_key]) for example in examples})
|
|
if len(groups) < 3:
|
|
raise ValueError("jackknife ensemble requires at least three groups")
|
|
models = []
|
|
for held_out in groups:
|
|
training = [
|
|
example for example in examples if str(example[group_key]) != held_out
|
|
]
|
|
models.append(
|
|
fit_ridge(
|
|
training,
|
|
include_telemetry=include_telemetry,
|
|
regularization=regularization,
|
|
)
|
|
)
|
|
return models
|
|
|
|
|
|
def ensemble_predict(
|
|
models: Sequence[RidgeModel],
|
|
example: Mapping[str, Any],
|
|
*,
|
|
include_telemetry: bool,
|
|
) -> dict[str, float]:
|
|
if not models:
|
|
raise ValueError("ensemble prediction requires models")
|
|
source = example["source"]
|
|
action = example["action"]
|
|
if (
|
|
int(action["target_mns"]) == int(source["mns"])
|
|
and int(action["target_mbbt"]) == int(source["mbbt"])
|
|
):
|
|
return {"mean": 0.0, "std": 0.0, "min": 0.0, "max": 0.0, "distinct_n": 1}
|
|
names, values = feature_vector(example, include_telemetry=include_telemetry)
|
|
if any(model.feature_names != tuple(names) for model in models):
|
|
raise ValueError("ensemble feature schema mismatch")
|
|
raw = np.asarray([model.predict(values) for model in models], dtype=np.float64)
|
|
clipped = np.clip(raw, -1.0, 1.0)
|
|
return {
|
|
"mean": float(clipped.mean()),
|
|
"std": float(clipped.std(ddof=0)),
|
|
"min": float(clipped.min()),
|
|
"max": float(clipped.max()),
|
|
"distinct_n": len(set(float(value) for value in clipped)),
|
|
}
|
|
|
|
|
|
def select_action(
|
|
models: Sequence[RidgeModel],
|
|
candidates: Sequence[Mapping[str, Any]],
|
|
*,
|
|
include_telemetry: bool,
|
|
confidence_z: float = 1.0,
|
|
minimum_margin: float = 0.02,
|
|
) -> dict[str, Any]:
|
|
if len(candidates) < 2:
|
|
raise ValueError("action selection requires at least two candidates")
|
|
rows = []
|
|
for example in candidates:
|
|
prediction = ensemble_predict(
|
|
models, example, include_telemetry=include_telemetry
|
|
)
|
|
rows.append(
|
|
{
|
|
"action_id": str(example["action"]["id"]),
|
|
"prediction": prediction,
|
|
"lower": prediction["mean"] - confidence_z * prediction["std"],
|
|
"upper": prediction["mean"] + confidence_z * prediction["std"],
|
|
}
|
|
)
|
|
rows.sort(key=lambda row: (-row["prediction"]["mean"], row["action_id"]))
|
|
best, second = rows[:2]
|
|
margin = float(best["prediction"]["mean"] - second["prediction"]["mean"])
|
|
confident = bool(
|
|
margin >= minimum_margin and best["lower"] > second["upper"]
|
|
)
|
|
return {
|
|
"selected_action": best["action_id"],
|
|
"confident": confident,
|
|
"predicted_margin": margin,
|
|
"candidates": rows,
|
|
}
|
|
|
|
|
|
def models_to_json(models: Iterable[RidgeModel]) -> list[dict[str, Any]]:
|
|
return [model.to_json() for model in models]
|
|
|
|
|
|
def models_from_json(payload: Iterable[Mapping[str, Any]]) -> list[RidgeModel]:
|
|
return [RidgeModel.from_json(item) for item in payload]
|