Skills › Scientific Agent Skills › pufferlib › benchmark_vectorization.py
benchmark_vectorization.py
244 lines · skills/pufferlib/scripts/benchmark_vectorization.py at 5f2bc03
1#!/usr/bin/env python3
2"""Bounded synthetic vectorization benchmark with no PufferLib import."""
4from __future__ import annotations
6import argparse
7import multiprocessing as mp
8import os
9import platform
10import statistics
11import time
12from typing import Any
14try:
15 from ._common import UserInputError, bounded_int, emit_json
16 from .env_template import SyntheticGymEnv
17except ImportError: # Direct script execution.
18 from _common import UserInputError, bounded_int, emit_json
19 from env_template import SyntheticGymEnv
22def _run_partition(payload: tuple[int, int, int, int, int]) -> dict[str, Any]:
23 """Run an allowlisted synthetic partition in one process."""
24 start_index, num_envs, steps_per_env, seed, max_steps = payload
25 agent_steps = 0
26 checksum = 0.0
27 resets = 0
28 for offset in range(num_envs):
29 env_index = start_index + offset
30 env = SyntheticGymEnv(max_steps=max_steps)
31 env.reset(seed=seed + env_index)
32 env_resets = 0
33 for step_index in range(steps_per_env):
34 action = (seed + 17 * env_index + step_index) % env.action_space.n
35 observation, reward, terminated, truncated, _ = env.step(action)
36 checksum += reward + observation[0] * 1e-6
37 agent_steps += 1
38 if terminated or truncated:
39 env_resets += 1
40 resets += 1
41 env.reset(seed=seed + env_index + env_resets * 1_000_003)
42 env.close()
43 return {"agent_steps": agent_steps, "checksum": checksum, "resets": resets}
46def _partitions(
47 *,
48 num_envs: int,
49 workers: int,
50 steps_per_env: int,
51 seed: int,
52 max_steps: int,
53) -> list[tuple[int, int, int, int, int]]:
54 parts: list[tuple[int, int, int, int, int]] = []
55 base, remainder = divmod(num_envs, workers)
56 start = 0
57 for worker_index in range(workers):
58 count = base + int(worker_index < remainder)
59 if count:
60 parts.append((start, count, steps_per_env, seed, max_steps))
61 start += count
62 return parts
65def _summarize(samples: list[float]) -> dict[str, float]:
66 ordered = sorted(samples)
68 def percentile(fraction: float) -> float:
69 index = round((len(ordered) - 1) * fraction)
70 return ordered[index]
72 return {
73 "maximum": max(samples),
74 "median": statistics.median(samples),
75 "minimum": min(samples),
76 "p10": percentile(0.10),
77 "p90": percentile(0.90),
78 }
81def benchmark(
82 *,
83 backend: str,
84 num_envs: int,
85 workers: int,
86 steps_per_env: int,
87 repeats: int,
88 warmup_steps: int,
89 seed: int,
90 max_steps: int,
91 start_method: str,
92) -> dict[str, Any]:
93 """Measure a fixed synthetic workload; never claim upstream PufferLib SPS."""
94 workers = min(workers, num_envs)
95 payloads = _partitions(
96 num_envs=num_envs,
97 workers=workers,
98 steps_per_env=steps_per_env,
99 seed=seed,
100 max_steps=max_steps,
101 )
102 warmup_payloads = _partitions(
103 num_envs=num_envs,
104 workers=workers,
105 steps_per_env=warmup_steps,
106 seed=seed,
107 max_steps=max_steps,
108 )
110 elapsed_samples: list[float] = []
111 throughput_samples: list[float] = []
112 checksums: list[float] = []
113 observed_steps: list[int] = []
115 pool: Any = None
116 try:
117 if backend == "multiprocessing":
118 context = mp.get_context(start_method)
119 pool = context.Pool(processes=workers)
120 if warmup_steps:
121 pool.map(_run_partition, warmup_payloads)
122 elif warmup_steps:
123 for payload in warmup_payloads:
124 _run_partition(payload)
126 for repeat_index in range(repeats):
127 adjusted = [
128 (start, count, steps, run_seed + repeat_index, limit)
129 for start, count, steps, run_seed, limit in payloads
130 ]
131 started = time.perf_counter()
132 if backend == "multiprocessing":
133 results = pool.map(_run_partition, adjusted)
134 else:
135 results = [_run_partition(payload) for payload in adjusted]
136 elapsed = time.perf_counter() - started
137 agent_steps = sum(int(item["agent_steps"]) for item in results)
138 checksum = sum(float(item["checksum"]) for item in results)
139 elapsed_samples.append(elapsed)
140 throughput_samples.append(agent_steps / elapsed)
141 observed_steps.append(agent_steps)
142 checksums.append(checksum)
143 finally:
144 if pool is not None:
145 pool.close()
146 pool.join()
148 return {
149 "backend": backend,
150 "benchmark": "bundled-synthetic-harness",
151 "checksums": checksums,
152 "elapsed_seconds": _summarize(elapsed_samples),
153 "environment_construction_in_timed_region": True,
154 "hardware": {
155 "logical_cpus": os.cpu_count(),
156 "platform": platform.platform(),
157 "python": platform.python_version(),
158 },
159 "network_used": False,
160 "parameters": {
161 "max_steps": max_steps,
162 "num_envs": num_envs,
163 "repeats": repeats,
164 "seed": seed,
165 "start_method": start_method if backend == "multiprocessing" else None,
166 "steps_per_env": steps_per_env,
167 "warmup_steps": warmup_steps,
168 "workers": workers,
169 },
170 "steps_per_second": _summarize(throughput_samples),
171 "total_agent_steps_per_repeat": observed_steps,
172 "warning": (
173 "This measures the bundled synthetic harness, not PufferLib, an Ocean "
174 "environment, training throughput, or cross-machine performance."
175 ),
176 }
179def build_parser() -> argparse.ArgumentParser:
180 safe_methods = [method for method in ("spawn", "forkserver") if method in mp.get_all_start_methods()]
181 parser = argparse.ArgumentParser(
182 description=(
183 "Benchmark only the built-in synthetic environment. Defaults are CPU-only, "
184 "network-free, and intentionally small."
185 )
186 )
187 parser.add_argument(
188 "--backend", choices=["serial", "multiprocessing"], default="serial"
189 )
190 parser.add_argument("--envs", type=int, default=4, help="1..128")
191 parser.add_argument("--workers", type=int, default=2, help="1..32")
192 parser.add_argument("--steps-per-env", type=int, default=1_000, help="1..100000")
193 parser.add_argument("--repeats", type=int, default=3, help="1..7")
194 parser.add_argument("--warmup-steps", type=int, default=32, help="0..1000")
195 parser.add_argument("--max-steps", type=int, default=32, help="1..10000")
196 parser.add_argument("--seed", type=int, default=42)
197 parser.add_argument(
198 "--start-method",
199 choices=safe_methods,
200 default="spawn" if "spawn" in safe_methods else safe_methods[0],
201 help="fork is intentionally unavailable",
202 )
203 parser.add_argument("--compact", action="store_true")
204 return parser
207def main(argv: list[str] | None = None) -> int:
208 parser = build_parser()
209 args = parser.parse_args(argv)
210 try:
211 num_envs = bounded_int(args.envs, name="envs", minimum=1, maximum=128)
212 workers = bounded_int(args.workers, name="workers", minimum=1, maximum=32)
213 steps = bounded_int(
214 args.steps_per_env,
215 name="steps_per_env",
216 minimum=1,
217 maximum=100_000,
218 )
219 repeats = bounded_int(args.repeats, name="repeats", minimum=1, maximum=7)
220 warmup = bounded_int(
221 args.warmup_steps, name="warmup_steps", minimum=0, maximum=1_000
222 )
223 max_steps = bounded_int(
224 args.max_steps, name="max_steps", minimum=1, maximum=10_000
225 )
226 report = benchmark(
227 backend=args.backend,
228 num_envs=num_envs,
229 workers=workers,
230 steps_per_env=steps,
231 repeats=repeats,
232 warmup_steps=warmup,
233 seed=args.seed,
234 max_steps=max_steps,
235 start_method=args.start_method,
236 )
237 except (UserInputError, ValueError, RuntimeError) as exc:
238 parser.error(str(exc))
239 emit_json(report, pretty=not args.compact)
240 return 0
243if __name__ == "__main__":
244 raise SystemExit(main())