ChilleD/agent_world_model_env
3
1"""2Stress test simulating large-scale agentic RL with AWM environments.3 4Simulates one RL step: 2000 environments reset in parallel, each runs a5multi-turn episode (random tool calls with LLM-like latency), then closes.6 7Phases per session:8 1. connect + reset — env startup9 2. list_tools — tool discovery10 3. N turns of tool calls — simulate multi-turn agent interaction11 (random tool, empty args, random "thinking" delay between turns)12 4. done + close — episode end13 14Usage:15 # Terminal 1: Start server16 PYTHONPATH=src:envs uv run uvicorn \17 envs.agent_world_model_env.server.app:app \18 --host 0.0.0.0 --port 889919 20 # Terminal 2: Run RL simulation (default 2000 envs)21 PYTHONPATH=src:envs uv run python \22 envs/agent_world_model_env/example_stress_test.py23 24 # Custom scale25 PYTHONPATH=src:envs uv run python \26 envs/agent_world_model_env/example_stress_test.py \27 --scale 500 --concurrency 100 --min-turns 1 --max-turns 328"""29 30import argparse31import asyncio32import json33import logging34import os35import random36import statistics37import sys38import time39from dataclasses import dataclass, field40 41import httpx42import psutil43from openenv.core.env_server.mcp_types import CallToolAction, ListToolsAction44 45from agent_world_model_env import AWMEnv46 47logging.basicConfig(48 level=logging.INFO,49 format="%(asctime)s [%(levelname)s] %(message)s",50 datefmt="%H:%M:%S",51)52log = logging.getLogger("rl_stress")53 54BASE_URL = "http://localhost:8899"55CLIENT_TIMEOUT: float = 600.056 57# Scenarios to cycle through58SCENARIOS = [59 "e_commerce_33",60 "inventory_management_7",61 "document_management_5",62 "billing_payments_3",63 "hris_employee_management_1",64]65 66# RL simulation defaults67MIN_TURNS = 368MAX_TURNS = 2069# Simulate LLM rollout time: uniform [min, max] seconds70LLM_THINK_MIN = 1.071LLM_THINK_MAX = 20.072 73 74# ---------------------------------------------------------------------------75# Data classes76# ---------------------------------------------------------------------------77@dataclass78class SessionResult:79 session_id: int80 scenario: str81 task_idx: int82 num_turns: int # planned turns83 turns_completed: int = 084 connect_s: float = 0.085 reset_s: float = 0.086 list_tools_s: float = 0.087 tool_call_latencies: list[float] = field(default_factory=list)88 done_s: float = 0.089 total_s: float = 0.090 success: bool = False91 error: str | None = None92 num_tools: int = 093 tools_discovered: list[str] = field(default_factory=list)94 95 96# ---------------------------------------------------------------------------97# System resource monitor98# ---------------------------------------------------------------------------99class ResourceMonitor:100 """Periodically samples CPU and memory in a background task."""101 102 def __init__(self, interval: float = 2.0):103 self._interval = interval104 self._samples: list[dict] = []105 self._task: asyncio.Task | None = None106 self._process = psutil.Process(os.getpid())107 self._server_pid: int | None = None108 109 def start(self, server_pid: int | None = None):110 self._server_pid = server_pid111 self._task = asyncio.create_task(self._loop())112 113 async def stop(self):114 if self._task:115 self._task.cancel()116 try:117 await self._task118 except asyncio.CancelledError:119 pass120 121 async def _loop(self):122 while True:123 sample = {124 "time": time.monotonic(),125 "system_cpu_pct": psutil.cpu_percent(interval=0),126 "system_mem_pct": psutil.virtual_memory().percent,127 "system_mem_used_gb": round(128 psutil.virtual_memory().used / (1024**3), 2129 ),130 "client_mem_mb": round(self._process.memory_info().rss / (1024**2), 1),131 }132 if self._server_pid:133 try:134 server_proc = psutil.Process(self._server_pid)135 children = server_proc.children(recursive=True)136 server_mem = server_proc.memory_info().rss137 for child in children:138 try:139 server_mem += child.memory_info().rss140 except (psutil.NoSuchProcess, psutil.AccessDenied):141 pass142 sample["server_tree_mem_mb"] = round(server_mem / (1024**2), 1)143 sample["server_children"] = len(children)144 except (psutil.NoSuchProcess, psutil.AccessDenied):145 pass146 self._samples.append(sample)147 await asyncio.sleep(self._interval)148 149 def summary(self) -> dict:150 if not self._samples:151 return {}152 cpu_vals = [s["system_cpu_pct"] for s in self._samples]153 mem_vals = [s["system_mem_used_gb"] for s in self._samples]154 client_mem = [s["client_mem_mb"] for s in self._samples]155 result = {156 "samples": len(self._samples),157 "cpu_pct": {158 "mean": round(statistics.mean(cpu_vals), 1),159 "max": round(max(cpu_vals), 1),160 },161 "system_mem_gb": {162 "min": round(min(mem_vals), 2),163 "max": round(max(mem_vals), 2),164 },165 "client_mem_mb": {166 "min": round(min(client_mem), 1),167 "max": round(max(client_mem), 1),168 },169 }170 server_mem = [171 s["server_tree_mem_mb"] for s in self._samples if "server_tree_mem_mb" in s172 ]173 if server_mem:174 result["server_tree_mem_mb"] = {175 "min": round(min(server_mem), 1),176 "max": round(max(server_mem), 1),177 }178 server_children = [179 s["server_children"] for s in self._samples if "server_children" in s180 ]181 if server_children:182 result["server_subprocess_peak"] = max(server_children)183 return result184 185 186# ---------------------------------------------------------------------------187# Helpers188# ---------------------------------------------------------------------------189def latency_stats(values: list[float]) -> dict:190 if not values:191 return {}192 s = sorted(values)193 return {194 "count": len(s),195 "min": round(min(s), 3),196 "p50": round(s[len(s) // 2], 3),197 "p90": round(s[int(len(s) * 0.9)], 3),198 "p99": round(s[int(len(s) * 0.99)], 3),199 "max": round(max(s), 3),200 "mean": round(statistics.mean(s), 3),201 }202 203 204async def check_server(url: str) -> int | None:205 """Check server is up, return server PID if available."""206 async with httpx.AsyncClient() as client:207 resp = await client.get(f"{url}/docs", timeout=10)208 resp.raise_for_status()209 # Try to find server PID by matching the port in the URL210 from urllib.parse import urlparse211 212 port = str(urlparse(url).port or "8899")213 for proc in psutil.process_iter(["pid", "cmdline"]):214 try:215 cmdline = " ".join(proc.info["cmdline"] or [])216 if "uvicorn" in cmdline and port in cmdline:217 return proc.info["pid"]218 except (psutil.NoSuchProcess, psutil.AccessDenied):219 pass220 return None221 222 223async def fetch_server_stats(url: str) -> dict | None:224 try:225 async with httpx.AsyncClient() as client:226 resp = await client.get(f"{url}/stats", timeout=5)227 return resp.json()228 except Exception:229 return None230 231 232# ---------------------------------------------------------------------------233# Single session: full RL episode234# ---------------------------------------------------------------------------235@dataclass236class ProgressCounters:237 """Shared counters updated by each session for live progress reporting."""238 239 done: int = 0240 ok: int = 0241 fail: int = 0242 resets_done: int = 0243 turns_done: int = 0244 245 246async def run_rl_episode(247 session_id: int,248 scenario: str,249 task_idx: int,250 num_turns: int,251 reset_semaphore: asyncio.Semaphore,252 interact_semaphore: asyncio.Semaphore,253 counters: ProgressCounters,254) -> SessionResult:255 """Simulate a full RL episode: reset -> list_tools -> N tool calls -> done."""256 r = SessionResult(257 session_id=session_id,258 scenario=scenario,259 task_idx=task_idx,260 num_turns=num_turns,261 )262 session_start = time.monotonic()263 phase = "init"264 env = AWMEnv(265 base_url=BASE_URL, message_timeout_s=CLIENT_TIMEOUT, connect_timeout_s=60.0266 )267 268 try:269 # -- Phase 1: connect + reset (rate-limited to avoid thundering herd) --270 async with reset_semaphore:271 phase = "connect"272 t0 = time.monotonic()273 await env.connect()274 r.connect_s = time.monotonic() - t0275 276 phase = "reset"277 t0 = time.monotonic()278 result = await env.reset(scenario=scenario, task_idx=task_idx)279 r.reset_s = time.monotonic() - t0280 281 if result.observation.reward_type not in ("reset_ok", "reset_warning"):282 r.error = f"reset failed: {result.observation.error}"283 counters.done += 1284 counters.fail += 1285 return r286 287 r.num_tools = result.observation.num_tools or 0288 counters.resets_done += 1289 290 # -- Phase 2: list_tools --291 phase = "list_tools"292 t0 = time.monotonic()293 result = await env.step(ListToolsAction())294 r.list_tools_s = time.monotonic() - t0295 296 # Collect tool names for random calling297 obs = result.observation298 if hasattr(obs, "tools") and obs.tools:299 r.tools_discovered = [300 t.get("name", t.get("tool_name", ""))301 for t in obs.tools302 if isinstance(t, dict)303 ]304 if not r.tools_discovered:305 r.tools_discovered = ["unknown_tool"]306 307 # -- Phase 3: multi-turn tool calling (simulate agent interaction) --308 async with interact_semaphore:309 for turn in range(num_turns):310 phase = f"turn_{turn}"311 312 # Simulate LLM thinking time (async sleep = non-blocking)313 think_time = random.uniform(LLM_THINK_MIN, LLM_THINK_MAX)314 await asyncio.sleep(think_time)315 316 # Pick a random tool and call with empty args (will fail, that's fine)317 tool_name = random.choice(r.tools_discovered)318 t0 = time.monotonic()319 try:320 result = await env.step(321 CallToolAction(tool_name=tool_name, arguments={})322 )323 except Exception:324 # Tool call failure is expected (no args), just measure latency325 pass326 r.tool_call_latencies.append(time.monotonic() - t0)327 r.turns_completed += 1328 counters.turns_done += 1329 330 # -- Phase 4: done + close --331 phase = "done"332 t0 = time.monotonic()333 await env.step(334 CallToolAction(tool_name="done", arguments={"keep_session": False})335 )336 r.done_s = time.monotonic() - t0337 338 r.success = True339 counters.ok += 1340 341 except Exception as e:342 r.error = f"[{phase}] {type(e).__name__}: {str(e)[:200]}"343 counters.fail += 1344 finally:345 r.total_s = time.monotonic() - session_start346 counters.done += 1347 try:348 await env.close()349 except Exception:350 pass351 352 return r353 354 355# ---------------------------------------------------------------------------356# Progress reporter357# ---------------------------------------------------------------------------358async def progress_reporter(359 counters: ProgressCounters,360 total: int,361 total_turns: int,362 monitor: ResourceMonitor,363 interval: float = 10.0,364):365 """Periodically log progress while the test runs."""366 start = time.monotonic()367 while True:368 await asyncio.sleep(interval)369 elapsed = time.monotonic() - start370 in_flight = total - counters.done371 372 stats = await fetch_server_stats(BASE_URL)373 server_sessions = stats.get("total_sessions", "?") if stats else "?"374 375 # Current resource snapshot376 samples = monitor._samples377 last = samples[-1] if samples else {}378 cpu = last.get("system_cpu_pct", "?")379 mem = last.get("system_mem_used_gb", "?")380 server_mem = last.get("server_tree_mem_mb", "?")381 children = last.get("server_children", "?")382 383 log.info(384 f"[{elapsed:.0f}s] episodes={counters.done}/{total} "385 f"ok={counters.ok} fail={counters.fail} "386 f"resets={counters.resets_done} "387 f"turns={counters.turns_done}/{total_turns} "388 f"in_flight={in_flight} | "389 f"server={server_sessions} subprocs={children} | "390 f"cpu={cpu}% mem={mem}GB server={server_mem}MB"391 )392 393 394# ---------------------------------------------------------------------------395# Main test396# ---------------------------------------------------------------------------397async def run_rl_step(398 scale: int,399 concurrency: int,400 min_turns: int,401 max_turns: int,402) -> tuple[list[SessionResult], dict]:403 """Run one RL step: launch `scale` episodes in parallel."""404 405 log.info("=" * 78)406 log.info(407 f"RL STEP SIMULATION: {scale} envs, concurrency={concurrency}, "408 f"turns={min_turns}-{max_turns}, timeout={CLIENT_TIMEOUT}s"409 )410 log.info("=" * 78)411 412 # Discover server PID for resource monitoring413 server_pid = await check_server(BASE_URL)414 log.info(f"Server reachable (pid={server_pid})")415 416 monitor = ResourceMonitor(interval=2.0)417 monitor.start(server_pid)418 419 # Two semaphores:420 # - reset_semaphore: limits concurrent resets (heavy: subprocess spawn)421 # - interact_semaphore: limits concurrent multi-turn interaction422 reset_semaphore = asyncio.Semaphore(concurrency)423 interact_semaphore = asyncio.Semaphore(scale) # no limit on interaction424 425 # Pre-assign turns per session426 turn_counts = [random.randint(min_turns, max_turns) for _ in range(scale)]427 total_planned_turns = sum(turn_counts)428 429 counters = ProgressCounters()430 431 # Launch progress reporter432 progress_task = asyncio.create_task(433 progress_reporter(counters, scale, total_planned_turns, monitor)434 )435 436 wall_start = time.monotonic()437 438 tasks = []439 for i in range(scale):440 scenario = SCENARIOS[i % len(SCENARIOS)]441 task_idx = i % 10442 tasks.append(443 run_rl_episode(444 session_id=i,445 scenario=scenario,446 task_idx=task_idx,447 num_turns=turn_counts[i],448 reset_semaphore=reset_semaphore,449 interact_semaphore=interact_semaphore,450 counters=counters,451 )452 )453 454 completed = await asyncio.gather(*tasks)455 wall_s = time.monotonic() - wall_start456 457 progress_task.cancel()458 try:459 await progress_task460 except asyncio.CancelledError:461 pass462 await monitor.stop()463 464 resource_summary = monitor.summary()465 466 # --------------- Report ---------------467 ok = [r for r in completed if r.success]468 failed = [r for r in completed if not r.success]469 total_turns = sum(r.turns_completed for r in completed)470 total_planned = sum(r.num_turns for r in completed)471 472 log.info("")473 log.info(f"{'=' * 78}")474 log.info(475 f"RESULTS: {len(ok)}/{scale} succeeded, {len(failed)} failed, wall={wall_s:.1f}s"476 )477 log.info(f"Total turns: {total_turns}/{total_planned} completed")478 log.info(f"{'=' * 78}")479 480 # Latency distributions481 for label, values in [482 ("connect", [r.connect_s for r in ok]),483 ("reset", [r.reset_s for r in ok]),484 ("list_tools", [r.list_tools_s for r in ok]),485 ("tool_call", [lat for r in ok for lat in r.tool_call_latencies]),486 ("done", [r.done_s for r in ok]),487 ("episode_total", [r.total_s for r in ok]),488 ]:489 stats = latency_stats(values)490 if stats:491 log.info(f" {label:>14s}: {json.dumps(stats)}")492 493 # Resource summary494 log.info("")495 log.info(f" {'RESOURCES':>14s}: {json.dumps(resource_summary)}")496 497 # Turn distribution498 if ok:499 turn_dist = [r.num_turns for r in ok]500 log.info(501 f" {'turns/episode':>14s}: min={min(turn_dist)} max={max(turn_dist)} "502 f"mean={statistics.mean(turn_dist):.1f}"503 )504 505 # Failures506 if failed:507 log.warning("")508 log.warning(f" {len(failed)} failures:")509 for r in failed[:20]:510 log.warning(511 f" session {r.session_id} ({r.scenario}/{r.task_idx}, "512 f"turns={r.turns_completed}/{r.num_turns}): {r.error}"513 )514 if len(failed) > 20:515 log.warning(f" ... and {len(failed) - 20} more")516 517 return list(completed), resource_summary518 519 520# ---------------------------------------------------------------------------521# CLI522# ---------------------------------------------------------------------------523def parse_args():524 p = argparse.ArgumentParser(525 description="AWM stress test — simulates large-scale agentic RL"526 )527 p.add_argument(528 "--scale",529 type=int,530 default=2000,531 help="Number of parallel environments per RL step (default: 2000)",532 )533 p.add_argument(534 "--concurrency",535 type=int,536 default=256,537 help="Max concurrent resets (default: 256)",538 )539 p.add_argument(540 "--min-turns",541 type=int,542 default=3,543 help="Min tool-call turns per episode (default: 3)",544 )545 p.add_argument(546 "--max-turns",547 type=int,548 default=20,549 help="Max tool-call turns per episode (default: 20)",550 )551 p.add_argument(552 "--think-min",553 type=float,554 default=1.0,555 help="Min LLM rollout time per turn in seconds (default: 1.0)",556 )557 p.add_argument(558 "--think-max",559 type=float,560 default=20.0,561 help="Max LLM rollout time per turn in seconds (default: 20.0)",562 )563 p.add_argument(564 "--url",565 default="http://localhost:8899",566 help="Server base URL (default: http://localhost:8899)",567 )568 p.add_argument(569 "--client-timeout",570 type=float,571 default=600.0,572 help="Client message timeout in seconds (default: 600)",573 )574 return p.parse_args()575 576 577async def main():578 args = parse_args()579 global BASE_URL, CLIENT_TIMEOUT, MIN_TURNS, MAX_TURNS, LLM_THINK_MIN, LLM_THINK_MAX580 BASE_URL = args.url581 CLIENT_TIMEOUT = args.client_timeout582 MIN_TURNS = args.min_turns583 MAX_TURNS = args.max_turns584 LLM_THINK_MIN = args.think_min585 LLM_THINK_MAX = args.think_max586 587 log.info(f"AWM RL Stress Test — server: {BASE_URL}")588 try:589 await check_server(BASE_URL)590 except Exception as e:591 log.error(f"Cannot reach server at {BASE_URL}: {e}")592 sys.exit(1)593 594 results, resources = await run_rl_step(595 args.scale, args.concurrency, args.min_turns, args.max_turns596 )597 598 ok = sum(1 for r in results if r.success)599 fail = len(results) - ok600 601 log.info("")602 log.info("=" * 78)603 log.info("FINAL SUMMARY")604 log.info("=" * 78)605 log.info(606 f" scale={args.scale} concurrency={args.concurrency} ok={ok} fail={fail}"607 )608 log.info(609 f" turns_range=[{args.min_turns},{args.max_turns}] "610 f"total_turns={sum(r.turns_completed for r in results)}"611 )612 if resources:613 log.info(f" resources={json.dumps(resources)}")614 615 if fail > 0:616 log.error("SOME EPISODES FAILED")617 sys.exit(1)618 else:619 log.info("ALL EPISODES PASSED")620 621 622if __name__ == "__main__":623 asyncio.run(main())624