Add action-conditioned intervention feasibility model
This commit is contained in:
287
runs/active-intervention-v0/model.py
Normal file
287
runs/active-intervention-v0/model.py
Normal file
@@ -0,0 +1,287 @@
|
||||
#!/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]
|
||||
Reference in New Issue
Block a user