Spaces:
Running
Running
File size: 4,448 Bytes
ee44678 | 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 | # SPDX-License-Identifier: BSD-3-Clause
"""Measure how many episodes the environment itself can supply.
Answers the question a training or eval run actually needs: is the environment
the bottleneck, or the model? Runs an agentic-shaped episode — reset, three
looks, optionally a pin, a guess — across a range of worker counts, both as
threads and as processes, and reports episodes and steps per second.
Every worker gets its own environment, because an episode is stateful; sharing
one interleaves resets and silently produces nonsense. A perfect guess is
asserted to score zero distance, so an interference bug fails loudly instead of
reporting excellent throughput.
Usage:
python scripts/benchmark_throughput.py threads 32
REVEAL=0 PIN=0 python scripts/benchmark_throughput.py threads 64
python scripts/benchmark_throughput.py processes 32
Environment variables:
REVEAL=0 skip the guess reveal map, as a training run would
PIN=0 skip the pin, leaving view renders only
"""
import concurrent.futures as cf
import os
import pathlib
import statistics
import sys
import threading
import time
_LOCAL = threading.local()
_ROOT = pathlib.Path(__file__).resolve().parents[2]
sys.path.insert(0, str(_ROOT))
sys.path.insert(0, str(_ROOT.parent / "src"))
INDEX = str(_ROOT / "geoguesser_env" / "tasks" / "pano_v1.jsonl")
CACHE = str(_ROOT / "geoguesser_env" / "data" / "panos")
def episode(task_index: int) -> tuple[float, int]:
"""One agentic-shaped episode: reset, 3 looks, 1 pin, 1 guess."""
from geoguesser_env.models import GuessAction, LookAction, PinAction, to_wire
from geoguesser_env.server.geoguesser_environment import GeoGuesserEnvironment
from geoguesser_env.server.render import minimap
minimap.set_street_detail(False) # training config: never touch Overpass
# An episode is stateful, so every worker needs its own environment. Sharing
# one through a module global interleaves resets between threads and silently
# produces nonsense — which is exactly what the first version of this
# benchmark measured.
if not hasattr(_LOCAL, "env"):
_LOCAL.env = GeoGuesserEnvironment(
index_path=INDEX,
cache_dir=CACHE,
allow_fetch=False,
view_size=640,
reveal_map=os.environ.get("REVEAL", "1") == "1",
)
env = _LOCAL.env
started = time.perf_counter()
env.reset(task_index=task_index)
steps = 1
for heading in (0, 120, 240):
env.step(to_wire(LookAction(heading_deg=heading, fov_deg=90)))
steps += 1
if os.environ.get("PIN", "1") == "1":
env.step(to_wire(PinAction(lat=10.0, lon=10.0, span_deg=7.0)))
steps += 1
truth = env._task.truth
result = env.step(to_wire(GuessAction(lat=truth[0], lon=truth[1])))
steps += 1
# A perfect guess must score ~1 minus costs. If workers interfered, this
# fails instead of quietly reporting a great throughput number.
assert result.distance_km is not None and result.distance_km < 0.01, (
f"cross-worker interference: task {task_index} scored "
f"{result.distance_km} km on a perfect guess"
)
return time.perf_counter() - started, steps
def run(pool_cls, workers: int, episodes: int) -> dict:
tasks = [i % 100 for i in range(episodes)]
wall0 = time.perf_counter()
with pool_cls(max_workers=workers) as pool:
results = list(pool.map(episode, tasks))
wall = time.perf_counter() - wall0
latencies = [r[0] for r in results]
total_steps = sum(r[1] for r in results)
return {
"workers": workers,
"episodes": episodes,
"wall_s": wall,
"episodes_per_s": episodes / wall,
"steps_per_s": total_steps / wall,
"median_episode_s": statistics.median(latencies),
}
if __name__ == "__main__":
kind = sys.argv[1]
episodes = int(sys.argv[2]) if len(sys.argv) > 2 else 24
pool = cf.ThreadPoolExecutor if kind == "threads" else cf.ProcessPoolExecutor
print(f"{kind}, {os.cpu_count()} cpus, warm cache, street detail off")
print(f"{'workers':>7} {'eps/s':>7} {'steps/s':>8} {'wall':>7} {'median ep':>10}")
for workers in (1, 2, 4, 8, 16):
r = run(pool, workers, episodes)
print(
f"{r['workers']:>7} {r['episodes_per_s']:>7.2f} {r['steps_per_s']:>8.1f} "
f"{r['wall_s']:>6.1f}s {r['median_episode_s']:>9.2f}s"
)
|