92 lines
2.7 KiB
Python
92 lines
2.7 KiB
Python
#!/usr/bin/env python3
|
|
from __future__ import annotations
|
|
|
|
import importlib.util
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
|
|
HERE = Path(__file__).resolve().parent
|
|
|
|
|
|
def load_model():
|
|
spec = importlib.util.spec_from_file_location(
|
|
"active_intervention_model", HERE / "model.py"
|
|
)
|
|
module = importlib.util.module_from_spec(spec)
|
|
assert spec.loader is not None
|
|
sys.modules[spec.name] = module
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
def example(model, decision: str, action: str, pressure: float, target: float):
|
|
outcome = {
|
|
name: 0.5 for name in model.PREFIX_FEATURES
|
|
}
|
|
telemetry = {name: 0.0 for name in model.TELEMETRY_FEATURES}
|
|
telemetry["queue_waiting_mean"] = pressure
|
|
telemetry["batch_size_mean"] = pressure
|
|
return {
|
|
"decision_id": decision,
|
|
"source": {
|
|
"mns": 16,
|
|
"mbbt": 8192,
|
|
"offered_rate_per_gpu": 2.0,
|
|
"outcome": outcome,
|
|
"telemetry": telemetry,
|
|
},
|
|
"action": {
|
|
"id": action,
|
|
"target_mns": 64 if action == "mns" else 16,
|
|
"target_mbbt": 8192 if action == "mns" else 16384,
|
|
},
|
|
"target_normalized_goodput": target,
|
|
"target_delta_normalized_goodput": target - 0.5,
|
|
}
|
|
|
|
|
|
def main() -> None:
|
|
model = load_model()
|
|
examples = []
|
|
for index, pressure in enumerate((0.2, 0.5, 0.8), 1):
|
|
examples.extend(
|
|
(
|
|
example(model, f"d{index}", "mns", pressure, 0.5 + pressure / 2),
|
|
example(model, f"d{index}", "mbbt", pressure, 0.6 - pressure / 4),
|
|
)
|
|
)
|
|
fitted = model.fit_ridge(
|
|
examples, include_telemetry=True, regularization=1.0
|
|
)
|
|
encoded = fitted.to_json()
|
|
restored = model.RidgeModel.from_json(encoded)
|
|
names, values = model.feature_vector(examples[-2], include_telemetry=True)
|
|
assert tuple(names) == restored.feature_names
|
|
assert abs(fitted.predict(values) - restored.predict(values)) < 1e-12
|
|
ensemble = model.fit_jackknife_ensemble(
|
|
examples, include_telemetry=True, regularization=1.0
|
|
)
|
|
decision = model.select_action(
|
|
ensemble, examples[-2:], include_telemetry=True, minimum_margin=0.0
|
|
)
|
|
assert decision["selected_action"] == "mns"
|
|
assert all(-1.0 <= row["prediction"]["mean"] <= 1.0 for row in decision["candidates"])
|
|
noop = example(model, "noop", "noop", 0.8, 0.5)
|
|
noop["action"]["target_mbbt"] = 8192
|
|
prediction = model.ensemble_predict(
|
|
ensemble, noop, include_telemetry=True
|
|
)
|
|
assert prediction == {
|
|
"mean": 0.0,
|
|
"std": 0.0,
|
|
"min": 0.0,
|
|
"max": 0.0,
|
|
"distinct_n": 1,
|
|
}
|
|
print("active intervention model: PASS")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|