Spaces:
Sleeping
Sleeping
File size: 7,008 Bytes
ed41ca2 | 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 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 | """
stress_test.py — Run 1000 simulations across all tasks with varying seeds.
Collects stats: success rate, avg score, avg reward, avg steps, failures.
"""
import sys, os
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from env_v2 import RepairEnvV2
from models_v2 import Action, ActionType, ComponentTarget
from tasks_v2 import TASKS
import random
import time
from collections import defaultdict
# ---------------------------------------------------------------------------
# Reuse the same inspect-then-fix agent from inference.py
# ---------------------------------------------------------------------------
def inspect_then_fix_agent(obs, task_config: dict, step: int, rng: random.Random) -> Action:
root_cause = task_config["root_cause"]
if isinstance(root_cause, list):
root_cause = root_cause[0]
if step == 0:
inspect_map = {
"database": ActionType.inspect_database,
"api": ActionType.inspect_api,
"cache": ActionType.inspect_cache,
"queue": ActionType.inspect_queue,
}
return Action(
action_type=inspect_map.get(root_cause, ActionType.inspect_database),
target=ComponentTarget(root_cause),
)
worst_comp = max(obs.metrics.items(), key=lambda kv: kv[1].error_rate)[0]
repair_map = {
"database": ActionType.repair_database,
"api": ActionType.restart_service,
"cache": ActionType.clear_cache,
"queue": ActionType.restart_service,
}
return Action(
action_type=repair_map.get(worst_comp, ActionType.restart_service),
target=ComponentTarget(worst_comp),
)
def random_agent(obs, rng: random.Random) -> Action:
"""Completely random agent for chaos testing."""
return Action(
action_type=rng.choice(list(ActionType)),
target=rng.choice(list(ComponentTarget)),
)
# ---------------------------------------------------------------------------
# Run N simulations
# ---------------------------------------------------------------------------
def run_stress_test(n_total: int = 1000):
task_ids = list(TASKS.keys())
n_per_task = n_total // len(task_ids)
# Stats per task per agent
stats = {}
for task_id in task_ids:
stats[task_id] = {
"baseline": defaultdict(list),
"random": defaultdict(list),
}
errors = []
total_runs = 0
start = time.time()
print(f"\n{'='*60}")
print(f" Stress Test: {n_total} simulations")
print(f" Tasks: {task_ids}")
print(f" {n_per_task} runs per task x 2 agents = {n_per_task * 2 * len(task_ids)} total episodes")
print(f"{'='*60}\n")
for task_id in task_ids:
config = TASKS[task_id]
print(f" Running task '{task_id}' ({n_per_task} seeds x 2 agents)...", end="", flush=True)
for seed in range(n_per_task):
for agent_name in ["baseline", "random"]:
try:
env = RepairEnvV2()
obs = env.reset(task_config=config, task_id=task_id)
done = False
step = 0
rng = random.Random(seed)
while not done:
if agent_name == "baseline":
action = inspect_then_fix_agent(obs, config, step, rng)
else:
action = random_agent(obs, rng)
obs, reward, done, info = env.step(action)
step += 1
result = env.get_episode_result()
s = stats[task_id][agent_name]
s["scores"].append(result.score)
s["rewards"].append(result.total_reward)
s["steps"].append(result.steps_taken)
s["success"].append(int(result.success))
s["health"].append(result.final_health_ratio)
total_runs += 1
except Exception as e:
errors.append({"task": task_id, "seed": seed, "agent": agent_name, "error": str(e)})
print(f" done.")
elapsed = time.time() - start
# ---------------------------------------------------------------------------
# Print results
# ---------------------------------------------------------------------------
print(f"\n{'='*60}")
print(f" RESULTS ({total_runs} episodes in {elapsed:.1f}s)")
print(f"{'='*60}\n")
for task_id in task_ids:
config = TASKS[task_id]
print(f" Task: {task_id} [{config.get('difficulty','?')}] (optimal={config['optimal_steps']})")
print(f" {'-'*55}")
for agent_name in ["baseline", "random"]:
s = stats[task_id][agent_name]
n = len(s["scores"])
if n == 0:
print(f" {agent_name:<10}: NO DATA")
continue
avg_score = sum(s["scores"]) / n
avg_reward = sum(s["rewards"]) / n
avg_steps = sum(s["steps"]) / n
success_pct = sum(s["success"]) / n * 100
avg_health = sum(s["health"]) / n
min_score = min(s["scores"])
max_score = max(s["scores"])
print(f" {agent_name:<10}: n={n} success={success_pct:5.1f}% "
f"score={avg_score:.3f} [{min_score:.2f}-{max_score:.2f}] "
f"reward={avg_reward:+.2f} steps={avg_steps:.1f} health={avg_health:.2%}")
print()
# Summary table
print(f" {'Task':<12} {'Agent':<10} {'Success%':>9} {'AvgScore':>9} {'AvgReward':>10} {'AvgSteps':>9}")
print(f" {'-'*58}")
for task_id in task_ids:
for agent_name in ["baseline", "random"]:
s = stats[task_id][agent_name]
n = len(s["scores"])
if n == 0: continue
print(f" {task_id:<12} {agent_name:<10} "
f"{sum(s['success'])/n*100:>8.1f}% "
f"{sum(s['scores'])/n:>9.4f} "
f"{sum(s['rewards'])/n:>10.3f} "
f"{sum(s['steps'])/n:>9.1f}")
print(f" {'-'*58}")
if errors:
print(f"\n ERRORS ({len(errors)} total):")
for e in errors[:10]:
print(f" [{e['task']} seed={e['seed']} {e['agent']}] {e['error']}")
if len(errors) > 10:
print(f" ... and {len(errors) - 10} more")
else:
print(f"\n No errors! All {total_runs} episodes completed cleanly.")
print(f"\n{'='*60}")
print(f" Stress test complete. {total_runs}/{n_per_task * 2 * len(task_ids)} episodes ran successfully.")
print(f"{'='*60}\n")
return stats, errors
if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser()
parser.add_argument("--n", type=int, default=1000, help="Total simulations (split across tasks)")
args = parser.parse_args()
run_stress_test(n_total=args.n)
|