Spaces:
Running
Running
File size: 4,079 Bytes
76818cc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 | import os
from typing import Any, Dict, List, Tuple
import requests
ENV_URL = os.environ.get("ENV_URL", "http://localhost:8000")
_MAX_STEPS = {
"task_easy": 5,
"task_medium": 10,
"task_hard": 20,
"task_expert": 15,
}
def _extract_obs(payload: Dict[str, Any]) -> Dict[str, Any]:
if isinstance(payload, dict) and isinstance(payload.get("observation"), dict):
return payload["observation"]
return payload
def _extract_reward(payload: Dict[str, Any], obs: Dict[str, Any]) -> float:
if "reward" in payload and payload["reward"] is not None:
return float(payload["reward"] or 0.0)
return float(obs.get("reward") or 0.0)
def _extract_done(payload: Dict[str, Any], obs: Dict[str, Any]) -> bool:
if "done" in payload and payload["done"] is not None:
return bool(payload["done"])
return bool(obs.get("done", False))
def run_episode(task_id: str) -> Tuple[float, int]:
reset_resp = requests.post(
f"{ENV_URL}/reset", json={"task_id": task_id}, timeout=15
)
try:
reset_payload = reset_resp.json()
except Exception:
print(
f" ERROR: non-JSON reset response for {task_id}: {reset_resp.text[:200]}"
)
return 0.0, 0
obs = _extract_obs(reset_payload)
max_steps = int(_MAX_STEPS.get(task_id, 20))
steps_taken = 0
rewards: List[float] = []
done = bool(obs.get("done", False))
while not done and steps_taken < max_steps:
alerts = obs.get("alerts") or []
alert_id = ""
if isinstance(alerts, list) and alerts:
first = alerts[0]
if isinstance(first, dict):
alert_id = str(first.get("id") or "")
action_type = "fix"
if task_id == "task_easy":
kind = alert_id.split("-")[0] if alert_id else "disk"
if kind in ("cpu", "memory"):
action_type = "restart"
elif kind == "network":
action_type = "mitigate"
else:
action_type = "scale_up"
step_payload = requests.post(
f"{ENV_URL}/step",
json={
"action": {
"alert_id": alert_id,
"action_type": action_type,
"notes": "benchmark",
}
},
timeout=15,
)
try:
step_payload = step_payload.json()
except Exception:
# Skip this episode rather than crashing.
print(
f" ERROR: non-JSON step response for {task_id}: {step_payload.text[:200]}"
)
return 0.0, 0
obs = _extract_obs(step_payload)
reward = _extract_reward(step_payload, obs)
done = _extract_done(step_payload, obs)
rewards.append(reward)
steps_taken += 1
# Continue until the environment reports done (or we hit max_steps).
score = min(sum(rewards), 1.0) if rewards else 0.0
return score, steps_taken
def run_benchmark(n: int = 10) -> None:
tasks = ["task_easy", "task_medium", "task_hard", "task_expert"]
print(f"=== Benchmark Results (n={n} stub episodes each) ===")
overall_scores: List[float] = []
for task_id in tasks:
scores: List[float] = []
steps_list: List[int] = []
for _ in range(n):
score, steps_taken = run_episode(task_id)
scores.append(score)
steps_list.append(steps_taken)
overall_scores.append(score)
mean_score = sum(scores) / n if n else 0.0
min_score = min(scores) if scores else 0.0
max_score = max(scores) if scores else 0.0
mean_steps = sum(steps_list) / n if n else 0.0
print(
f"{task_id}: mean={mean_score:.3f} min={min_score:.3f} max={max_score:.3f} steps={mean_steps:.1f}"
)
overall_mean = sum(overall_scores) / len(overall_scores) if overall_scores else 0.0
print(f"=== Overall mean: {overall_mean:.3f} ===")
if __name__ == "__main__":
run_benchmark()
|