← back to Exo
bench/exo_bench.py
937 lines
# type: ignore
#!/usr/bin/env python3
"""Tool-calling eval for exo's OpenAI-compatible API.
Tests whether models correctly:
- Trigger tool calls when appropriate
- Return valid JSON arguments matching function schemas
- Handle multi-turn tool use (call -> result -> final answer)
- Avoid calling tools when unnecessary
Start exo with a model first, then run:
uv run python tool_call_eval.py --model <model-id>
uv run python tool_call_eval.py --model <model-id> --host 10.0.0.5 --port 52415
uv run python tool_call_eval.py --model <model-id> --repeat 3
uv run python tool_call_eval.py --model <model-id> --scenarios weather_simple calculator_multi_turn
"""
from __future__ import annotations
import argparse
import contextlib
import itertools
import json
import sys
import threading
import time
from collections.abc import Callable
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path
from statistics import mean
from typing import Any
from exo_tools.client import ExoClient, ExoHttpError
from exo_tools.harness import (
add_common_instance_args,
capture_cluster_snapshot,
find_existing_instance,
instance_id_from_instance,
node_ids_from_instance,
nodes_used_in_instance,
resolve_model_short_id,
run_planning_phase,
settle_and_fetch_placements,
wait_for_instance_gone,
wait_for_instance_ready,
)
from loguru import logger
from transformers import AutoTokenizer
# Monkey-patch for transformers 5.x compatibility
# Kimi's tokenization_kimi.py imports bytes_to_unicode from the old location
# which was moved in transformers 5.0.0rc2
try:
import transformers.models.gpt2.tokenization_gpt2 as gpt2_tokenization
from transformers.convert_slow_tokenizer import bytes_to_unicode
if not hasattr(gpt2_tokenization, "bytes_to_unicode"):
gpt2_tokenization.bytes_to_unicode = bytes_to_unicode # type: ignore[attr-defined]
except ImportError:
pass # transformers < 5.0 or bytes_to_unicode not available
def load_tokenizer_for_bench(model_id: str) -> Any:
"""
Load tokenizer for benchmarking, with special handling for Kimi models.
Kimi uses a custom TikTokenTokenizer that transformers 5.x can't load via AutoTokenizer.
This function replicates the logic from utils_mlx.py for bench compatibility.
"""
model_id_lower = model_id.lower()
if "kimi-k2" in model_id_lower:
import importlib.util
import types
from huggingface_hub import snapshot_download
# Download/get the model path
model_path = Path(
snapshot_download(
model_id,
allow_patterns=["*.json", "*.py", "*.tiktoken", "*.model", "*.jinja"],
)
)
sys.path.insert(0, str(model_path))
# Load tool_declaration_ts first (tokenization_kimi imports it with relative import)
tool_decl_path = model_path / "tool_declaration_ts.py"
if tool_decl_path.exists():
spec = importlib.util.spec_from_file_location(
"tool_declaration_ts", tool_decl_path
)
if spec and spec.loader:
tool_decl_module = importlib.util.module_from_spec(spec)
sys.modules["tool_declaration_ts"] = tool_decl_module
spec.loader.exec_module(tool_decl_module)
# Load tokenization_kimi with patched source (convert relative to absolute import)
tok_path = model_path / "tokenization_kimi.py"
source = tok_path.read_text()
source = source.replace("from .tool_declaration_ts", "from tool_declaration_ts")
spec = importlib.util.spec_from_file_location("tokenization_kimi", tok_path)
if spec:
tok_module = types.ModuleType("tokenization_kimi")
tok_module.__file__ = str(tok_path)
sys.modules["tokenization_kimi"] = tok_module
exec(compile(source, tok_path, "exec"), tok_module.__dict__) # noqa: S102
TikTokenTokenizer = tok_module.TikTokenTokenizer # noqa: N806
else:
from tokenization_kimi import TikTokenTokenizer # type: ignore[import-not-found] # noqa: I001
hf_tokenizer: Any = TikTokenTokenizer.from_pretrained(model_path)
# Patch encode to use internal tiktoken model directly
# transformers 5.x has a bug in the encode->pad path for slow tokenizers
def _patched_encode(text: str, **kwargs: object) -> list[int]:
# Pass allowed_special="all" to handle special tokens like <|im_user|>
return list(hf_tokenizer.model.encode(text, allowed_special="all"))
hf_tokenizer.encode = _patched_encode
return hf_tokenizer
# TODO: Change back to using only transformers
try:
return AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
except (AttributeError, ValueError):
from huggingface_hub import snapshot_download
from transformers import PretrainedConfig
model_path = Path(
snapshot_download(
model_id,
allow_patterns=[
"*.json",
"*.py",
"tokenizer.model",
"*.tiktoken",
"tiktoken.model",
"*.txt",
"*.jsonl",
"*.jinja",
],
)
)
stub_kwargs: dict[str, Any] = {}
config_file = model_path / "config.json"
if config_file.exists():
with open(config_file) as f:
raw = json.load(f)
for key in (
"model_type",
"max_position_embeddings",
"vocab_size",
"bos_token_id",
"eos_token_id",
"pad_token_id",
):
if key in raw:
stub_kwargs[key] = raw[key]
return AutoTokenizer.from_pretrained(
str(model_path),
config=PretrainedConfig(**stub_kwargs),
trust_remote_code=True,
)
def format_peak_memory(b: float) -> str:
for unit in ["B", "KB", "MB", "GB", "TB"]:
if b < 1024.0:
return f"{b:.2f}{unit}"
b /= 1024.0
raise ValueError("You're using petabytes of memory. Something went wrong...")
_SAMPLER_METRICS = ("gpuUsage", "temp", "sysPower", "pcpuUsage", "ecpuUsage")
class SystemMetricsSampler:
def __init__(self, client: ExoClient, node_ids: list[str], interval_s: float = 1.0):
self._client = client
self._node_ids = node_ids
self._interval_s = interval_s
self._samples: dict[str, list[tuple[float, dict[str, float]]]] = {
nid: [] for nid in node_ids
}
self._stop = threading.Event()
self._thread: threading.Thread | None = None
def start(self) -> None:
self._stop.clear()
self._thread = threading.Thread(target=self._poll_loop, daemon=True)
self._thread.start()
def stop(self) -> None:
self._stop.set()
if self._thread:
self._thread.join(timeout=5)
def _poll_loop(self) -> None:
while not self._stop.is_set():
t = time.monotonic()
for nid in self._node_ids:
try:
data = self._client.get_node_system(nid)
if data:
self._samples[nid].append(
(t, {k: data.get(k, 0.0) for k in _SAMPLER_METRICS})
)
except Exception:
pass
self._stop.wait(self._interval_s)
def energy_between(self, t0: float, t1: float) -> float:
total_joules = 0.0
for _nid, samples in self._samples.items():
window = [(t, s["sysPower"]) for t, s in samples if t0 <= t <= t1]
if len(window) >= 2:
for i in range(1, len(window)):
dt = window[i][0] - window[i - 1][0]
avg_power = (window[i][1] + window[i - 1][1]) / 2
total_joules += avg_power * dt
elif len(window) == 1:
total_joules += window[0][1] * (t1 - t0)
return total_joules
def summarize(self) -> dict[str, dict[str, dict[str, float]]]:
result: dict[str, dict[str, dict[str, float]]] = {}
for nid, samples in self._samples.items():
if not samples:
continue
metrics: dict[str, dict[str, float]] = {}
for key in _SAMPLER_METRICS:
values = [s[key] for t, s in samples]
metrics[key] = {
"min": round(min(values), 2),
"max": round(max(values), 2),
"mean": round(mean(values), 2),
"samples": len(values),
}
result[nid] = metrics
return result
def print_summary(self, placement_label: str) -> None:
summary = self.summarize()
if not summary:
return
logger.info(f"--- System Metrics ({placement_label}) ---")
for nid, metrics in summary.items():
gpu = metrics.get("gpuUsage", {})
temp = metrics.get("temp", {})
power = metrics.get("sysPower", {})
logger.info(
f" {nid}: "
f"GPU {gpu.get('mean', 0) * 100:.0f}% avg ({gpu.get('min', 0) * 100:.0f}–{gpu.get('max', 0) * 100:.0f}%) | "
f"{temp.get('mean', 0):.1f}°C avg | "
f"{power.get('mean', 0):.1f}W avg"
)
def parse_int_list(values: list[str]) -> list[int]:
items: list[int] = []
for v in values:
for part in v.split(","):
part = part.strip()
if part:
items.append(int(part))
return items
def run_one_completion(
client: ExoClient,
model_id: str,
pp_hint: int,
tg: int,
prompt_sizer: PromptSizer,
*,
use_prefix_cache: bool = False,
stream: bool = False,
) -> tuple[dict[str, Any], int]:
content, pp_tokens = prompt_sizer.build(pp_hint)
payload: dict[str, Any] = {
"model": model_id,
"messages": [{"role": "user", "content": content}],
"max_tokens": tg,
"logprobs": False,
"use_prefix_cache": use_prefix_cache,
}
if not stream:
payload["stream"] = False
t0 = time.perf_counter()
out = client.post_bench_chat_completions(payload)
elapsed = time.perf_counter() - t0
stats = out.get("generation_stats")
power_usage = out.get("power_usage")
choices = out.get("choices") or [{}]
message = choices[0].get("message", {}) if choices else {}
content = message.get("content") or ""
preview = content[:200] if content else ""
else:
tokens = 0
first_token_time = None
t0 = time.perf_counter()
text_parts: list[str] = []
stats = None
for raw_line in client.stream_bench_chat_completions(payload):
line = raw_line.strip()
if line.startswith(": generation_stats "):
with contextlib.suppress(json.JSONDecodeError):
stats = json.loads(line[len(": generation_stats ") :])
continue
if not line.startswith("data: "):
continue
data = line[6:]
if data == "[DONE]":
break
try:
chunk = json.loads(data)
delta = chunk.get("choices", [{}])[0].get("delta", {})
if delta.get("content"):
if first_token_time is None:
first_token_time = time.perf_counter()
tokens += 1
text_parts.append(delta["content"])
except json.JSONDecodeError:
pass
elapsed = time.perf_counter() - t0
preview = "".join(text_parts)[:200]
power_usage = None
if not stats:
ttft = (first_token_time - t0) if first_token_time else elapsed
gen_time = elapsed - ttft if tokens > 1 else elapsed
gen_tps = (tokens - 1) / gen_time if tokens > 1 and gen_time > 0 else 0.0
prompt_tps = pp_tokens / ttft if ttft > 0 else 0.0
stats = {
"prompt_tokens": pp_tokens,
"generation_tokens": tokens,
"prompt_tps": round(prompt_tps, 2),
"generation_tps": round(gen_tps, 2),
"peak_memory_usage": {"inBytes": 0},
}
return {
"elapsed_s": elapsed,
"output_text_preview": preview,
"stats": stats,
"power_usage": power_usage,
}, pp_tokens
class PromptSizer:
def __init__(self, tokenizer: Any, atom: str = "a "):
self.tokenizer = tokenizer
self.atom = atom
self.count_fn = PromptSizer._make_counter(tokenizer)
self.base_tokens = self.count_fn("")
@staticmethod
def _make_counter(tokenizer: Any) -> Callable[[str], int]:
def count_fn(user_content: str) -> int:
messages = [{"role": "user", "content": user_content}]
try:
ids = tokenizer.apply_chat_template(
messages, tokenize=True, add_generation_prompt=True
)
except ValueError:
# Models without a Jinja chat template (e.g. DeepSeek V4 which
# ships its own Python encoder). Use the exo-side V4 encoder.
from exo.worker.engines.mlx.deepseek_v4_encoding import (
encode_messages as encode_v4,
)
prompt = encode_v4(messages, thinking_mode="thinking")
ids = tokenizer.encode(prompt, add_special_tokens=False)
# Fix for transformers 5.x
if hasattr(ids, "input_ids"):
ids = ids.input_ids
return int(len(ids))
return count_fn
def build(self, target_prompt_tokens: int) -> tuple[str, int]:
target = int(target_prompt_tokens)
if target < self.base_tokens:
raise RuntimeError(
f"Target ({target}) is smaller than template overhead ({self.base_tokens})."
)
# Estimate tokens per atom using a sample
sample_count = 100
sample_content = self.atom * sample_count
sample_tokens = self.count_fn(sample_content) - self.base_tokens
tokens_per_atom = sample_tokens / sample_count
# Estimate starting point
needed_tokens = target - self.base_tokens
estimated_atoms = int(needed_tokens / tokens_per_atom)
# Binary search to find exact atom count
low, high = 0, estimated_atoms * 2 + 100
while low < high:
mid = (low + high) // 2
tok = self.count_fn(self.atom * mid)
if tok < target:
low = mid + 1
else:
high = mid
content = self.atom * low
tok = self.count_fn(content)
logger.info(f"{tok=}")
if tok != target:
raise RuntimeError(
f"Overshot: got {tok} tokens (target {target}). "
f"Pick a different atom (try ' a' or '\\n' or '0 ')."
)
return content, tok
def main() -> int:
ap = argparse.ArgumentParser(
prog="exo-bench",
description="Benchmark exo model throughput across placement previews.",
)
add_common_instance_args(ap)
ap.add_argument(
"--pp",
nargs="+",
required=True,
help="Prompt-size hints (ints). Accepts commas.",
)
ap.add_argument(
"--tg",
nargs="+",
required=True,
help="Generation lengths (ints). Accepts commas.",
)
ap.add_argument(
"--repeat", type=int, default=1, help="Repetitions per (pp,tg) pair."
)
ap.add_argument(
"--concurrency",
nargs="+",
default=["1"],
help="Concurrency levels (ints). Accepts commas. E.g. --concurrency 1,2,4,8. Default 1.",
)
ap.add_argument(
"--warmup",
type=int,
default=0,
help="Warmup runs per placement (uses first pp/tg).",
)
ap.add_argument(
"--json-out",
default="bench/results.json",
help="Write raw per-run results JSON to this path.",
)
ap.add_argument("--stdout", action="store_true", help="Write results to stdout")
ap.add_argument(
"--dry-run", action="store_true", help="List selected placements and exit."
)
ap.add_argument(
"--all-combinations",
action="store_true",
help="Force all pp×tg combinations (cartesian product) even when lists have equal length.",
)
ap.add_argument(
"--stream",
action="store_true",
help="Use /bench/chat/completions with streaming SSE response (bench=True still applies: no EOS detection, no KV cache).",
)
ap.add_argument(
"--no-system-metrics",
action="store_true",
help="Disable GPU utilization, temperature, and power collection during inference.",
)
ap.add_argument(
"--metrics-interval",
type=float,
default=1.0,
help="System metrics polling interval in seconds (default: 1.0).",
)
ap.add_argument(
"--use-prefix-cache",
action="store_true",
help="Enable KV prefix cache during bench (default: disabled for cold-cache measurements).",
)
args = ap.parse_args()
pp_list = parse_int_list(args.pp)
tg_list = parse_int_list(args.tg)
if not pp_list or not tg_list:
logger.error("pp and tg lists must be non-empty")
return 2
if args.repeat <= 0:
logger.error("--repeat must be >= 1")
return 2
concurrency_list = parse_int_list(args.concurrency)
if not concurrency_list or any(c <= 0 for c in concurrency_list):
logger.error("--concurrency values must be >= 1")
return 2
if args.use_prefix_cache:
logger.warning(
"--use-prefix-cache: prompt TPS will be approximate. See METHODOLOGY.md for details."
)
if pp_list != sorted(pp_list):
logger.warning(
"--pp values are not in ascending order: prompt TPS will be less accurate. Use ascending --pp for best results."
)
# Log pairing mode
use_combinations = args.all_combinations or len(pp_list) != len(tg_list)
if use_combinations:
logger.info(
f"pp/tg mode: combinations (product) - {len(pp_list) * len(tg_list)} pairs"
)
else:
logger.info(f"pp/tg mode: tandem (zip) - {len(pp_list)} pairs")
client = ExoClient(args.host, args.port, timeout_s=args.timeout)
short_id, full_model_id = resolve_model_short_id(
client, args.model, force_download=args.force_download
)
tokenizer = load_tokenizer_for_bench(full_model_id)
if tokenizer is None:
raise RuntimeError("[exo-bench] tokenizer load failed")
try:
prompt_sizer = PromptSizer(tokenizer)
logger.debug(f"[exo-bench] loaded tokenizer: {full_model_id} for prompt sizer")
except Exception:
logger.error("[exo-bench] tokenizer usable but prompt sizing failed")
raise
# Optionally reuse a running instance for this model
reused_instance_id: str | None = None
if args.reuse_instance:
existing = find_existing_instance(client, full_model_id)
if existing:
reused_instance_id = existing
logger.info(f"Reusing existing instance {reused_instance_id}")
else:
logger.warning(
"--reuse-instance: no existing instance found, creating a new one"
)
if reused_instance_id is not None:
# Use the existing instance directly — skip placement iteration
selected = []
download_duration_s = None
else:
selected = settle_and_fetch_placements(
client, full_model_id, args, settle_timeout=args.settle_timeout
)
if not selected:
logger.error("No valid placements matched your filters.")
return 1
selected.sort(
key=lambda p: (
str(p.get("instance_meta", "")),
str(p.get("sharding", "")),
nodes_used_in_instance(p["instance"]),
),
reverse=True,
)
logger.debug(f"exo-bench model: short_id={short_id} full_id={full_model_id}")
logger.info(f"placements: {len(selected)}")
for p in selected:
logger.info(
f" - {p['sharding']} / {p['instance_meta']} / nodes={nodes_used_in_instance(p['instance'])}"
)
if args.dry_run:
return 0
settle_deadline = (
time.monotonic() + args.settle_timeout if args.settle_timeout > 0 else None
)
logger.info("Planning phase: checking downloads...")
download_duration_s = run_planning_phase(
client,
full_model_id,
selected[0],
args.danger_delete_downloads,
args.timeout,
settle_deadline,
)
if download_duration_s is not None:
logger.info(f"Download: {download_duration_s:.1f}s (freshly downloaded)")
else:
logger.info("Download: model already cached")
cluster_snapshot = capture_cluster_snapshot(client)
all_rows: list[dict[str, Any]] = []
all_system_metrics: dict[str, dict[str, dict[str, float]]] = {}
# If reusing an existing instance, run a single benchmark pass against it
if reused_instance_id is not None:
selected = [None]
for preview in selected:
created_instance = False
if preview is not None:
instance = preview["instance"]
instance_id = instance_id_from_instance(instance)
sharding = str(preview["sharding"])
instance_meta = str(preview["instance_meta"])
n_nodes = nodes_used_in_instance(instance)
logger.info("=" * 80)
logger.info(
f"PLACEMENT: {sharding} / {instance_meta} / nodes={n_nodes} / instance_id={instance_id}"
)
# Delete any existing instances to free resources before placing
try:
state = client.request_json("GET", "/state")
for old_id in list(state.get("instances", {}).keys()):
logger.info(f"Deleting stale instance {old_id}")
with contextlib.suppress(ExoHttpError):
client.request_json("DELETE", f"/instance/{old_id}")
if state.get("instances"):
time.sleep(2)
except Exception as e:
logger.warning(f"Failed to clean up stale instances: {e}")
client.request_json("POST", "/instance", body={"instance": instance})
try:
wait_for_instance_ready(client, instance_id)
except (RuntimeError, TimeoutError) as e:
logger.error(f"Failed to initialize placement: {e}")
with contextlib.suppress(ExoHttpError):
client.request_json("DELETE", f"/instance/{instance_id}")
continue
time.sleep(1)
created_instance = True
else:
instance_id = reused_instance_id
sharding = "reused"
instance_meta = "reused"
n_nodes = 0
logger.info("=" * 80)
logger.info(f"Using existing instance {instance_id}")
sampler: SystemMetricsSampler | None = None
if not args.no_system_metrics and preview is not None:
nids = node_ids_from_instance(instance)
sampler = SystemMetricsSampler(
ExoClient(args.host, args.port, timeout_s=30),
nids,
interval_s=args.metrics_interval,
)
sampler.start()
def _do_one(c: ExoClient, pp: int, tg: int) -> tuple[dict[str, Any], int]:
return run_one_completion(
c,
full_model_id,
pp,
tg,
prompt_sizer,
use_prefix_cache=args.use_prefix_cache,
stream=args.stream,
)
try:
for i in range(args.warmup):
_do_one(client, pp_list[0], tg_list[0])
logger.debug(f" warmup {i + 1}/{args.warmup} done")
# If pp and tg lists have same length, run in tandem (zip)
# Otherwise (or if --all-combinations), run all combinations (cartesian product)
if use_combinations:
pp_tg_pairs = list(itertools.product(pp_list, tg_list))
else:
pp_tg_pairs = list(zip(pp_list, tg_list, strict=True))
for pp, tg in pp_tg_pairs:
for concurrency in concurrency_list:
logger.info(f"--- pp={pp} tg={tg} concurrency={concurrency} ---")
runs: list[dict[str, Any]] = []
inference_windows: list[tuple[float, float]] = []
for r in range(args.repeat):
time.sleep(3)
if concurrency <= 1:
# Sequential: single request
try:
inf_t0 = time.monotonic()
row, actual_pp_tokens = _do_one(client, pp, tg)
inference_windows.append((inf_t0, time.monotonic()))
except Exception as e:
logger.error(e)
continue
row.update(
{
"model_short_id": short_id,
"model_id": full_model_id,
"placement_sharding": sharding,
"placement_instance_meta": instance_meta,
"placement_nodes": n_nodes,
"instance_id": instance_id,
"pp_tokens": actual_pp_tokens,
"tg": tg,
"repeat_index": r,
"concurrency": 1,
**(
{"download_duration_s": download_duration_s}
if download_duration_s is not None
else {}
),
}
)
runs.append(row)
all_rows.append(row)
else:
# Concurrent: fire N requests in parallel
# Pre-build prompt once, barrier ensures simultaneous dispatch
content, actual_pp = prompt_sizer.build(pp)
pre_built_payload: dict[str, Any] = {
"model": full_model_id,
"messages": [{"role": "user", "content": content}],
"stream": False,
"max_tokens": tg,
"logprobs": False,
"use_prefix_cache": args.use_prefix_cache,
}
barrier = threading.Barrier(concurrency)
batch_start = threading.Event()
batch_t0: float = 0.0
batch_results: list[tuple[dict[str, Any], int]] = []
batch_errors = 0
def _run_concurrent(
idx: int,
_barrier: threading.Barrier = barrier,
_batch_start: threading.Event = batch_start,
_payload: dict[str, Any] = pre_built_payload,
_actual_pp: int = actual_pp,
) -> tuple[dict[str, Any], int]:
nonlocal batch_t0
c = ExoClient(
args.host, args.port, timeout_s=args.timeout
)
if _barrier.wait() == 0:
batch_t0 = time.perf_counter()
_batch_start.set()
else:
_batch_start.wait()
t0 = batch_t0
out = c.post_bench_chat_completions(_payload)
elapsed = time.perf_counter() - t0
stats = out.get("generation_stats")
power_usage = out.get("power_usage")
choices = out.get("choices") or [{}]
message = (
choices[0].get("message", {}) if choices else {}
)
text = message.get("content") or ""
return {
"elapsed_s": elapsed,
"output_text_preview": text[:200],
"stats": stats,
"power_usage": power_usage,
}, _actual_pp
inf_t0 = time.monotonic()
with ThreadPoolExecutor(max_workers=concurrency) as pool:
futures = {
pool.submit(_run_concurrent, i): i
for i in range(concurrency)
}
for fut in as_completed(futures):
try:
batch_results.append(fut.result())
except Exception as e:
logger.error(f"Concurrent request failed: {e}")
batch_errors += 1
batch_wall_s = (
max(x["elapsed_s"] for x, _ in batch_results)
if batch_results
else time.perf_counter() - batch_t0
)
inference_windows.append((inf_t0, time.monotonic()))
for idx, (row, actual_pp_tokens) in enumerate(
batch_results
):
row.update(
{
"model_short_id": short_id,
"model_id": full_model_id,
"placement_sharding": sharding,
"placement_instance_meta": instance_meta,
"placement_nodes": n_nodes,
"instance_id": instance_id,
"pp_tokens": actual_pp_tokens,
"tg": tg,
"repeat_index": r,
"concurrency": concurrency,
"concurrent_index": idx,
**(
{"download_duration_s": download_duration_s}
if download_duration_s is not None
else {}
),
}
)
runs.append(row)
all_rows.append(row)
if batch_results:
valid_gen_tps = [
x["stats"]["generation_tps"]
for x, _ in batch_results
if x["stats"]["generation_tps"] > 0
]
per_req_tps = (
max(valid_gen_tps) if valid_gen_tps else 0.0
)
agg_gen_tps = per_req_tps * concurrency
logger.info(
f"[concurrent {concurrency}x] "
f"agg_gen_tps={agg_gen_tps:.2f} "
f"per_req_tps={per_req_tps:.2f} "
f"wall_s={batch_wall_s:.2f} "
f"errors={batch_errors}"
)
if runs:
prompt_tps = mean(x["stats"]["prompt_tps"] for x in runs)
valid_gen = [
x["stats"]["generation_tps"]
for x in runs
if x["stats"]["generation_tps"] > 0
]
per_req_tps = max(valid_gen) if valid_gen else 0.0
gen_tps = per_req_tps * concurrency
ptok = mean(x["stats"]["prompt_tokens"] for x in runs)
gtok = mean(x["stats"]["generation_tokens"] for x in runs)
def _peak_bytes(s: dict[str, Any]) -> float:
pm = s["peak_memory_usage"]
return pm.get("inBytes") or pm.get("in_bytes", 0)
peak = mean(_peak_bytes(x["stats"]) for x in runs)
summary = (
f"prompt_tps={prompt_tps:.2f} gen_tps={gen_tps:.2f} "
f"prompt_tokens={ptok} gen_tokens={gtok} "
f"peak_memory={format_peak_memory(peak)}"
)
if sampler and inference_windows:
joules = sum(
sampler.energy_between(t0, t1)
for t0, t1 in inference_windows
)
inf_seconds = sum(t1 - t0 for t0, t1 in inference_windows)
avg_watts = joules / inf_seconds if inf_seconds > 0 else 0
summary += f" energy={joules:.1f}J ({avg_watts:.1f}W avg over {inf_seconds:.1f}s inference)"
# mean() not sum() across concurrent runs: each
# request's PowerSampler observes the same shared
# cluster state, so they all report the same figure.
prefill_energies = [
(x.get("power_usage") or {}).get("prefill_energy_joules")
for x in runs
]
gen_energies = [
(x.get("power_usage") or {}).get("generation_energy_joules")
for x in runs
]
prefill_vals = [e for e in prefill_energies if e is not None]
gen_vals = [e for e in gen_energies if e is not None]
if prefill_vals and gen_vals:
avg_pref = mean(prefill_vals)
avg_gen = mean(gen_vals)
summary += (
f" prefill_energy={avg_pref:.1f}J "
f"gen_energy={avg_gen:.1f}J"
)
logger.info(f"{summary}\n")
time.sleep(2)
finally:
if sampler:
sampler.stop()
placement_label = f"{sharding}/{instance_meta}/{n_nodes} nodes"
sampler.print_summary(placement_label)
placement_metrics = sampler.summarize()
if placement_metrics:
all_system_metrics.update(placement_metrics)
if created_instance and instance_id is not None:
try:
client.request_json("DELETE", f"/instance/{instance_id}")
except ExoHttpError as e:
if e.status != 404:
raise
wait_for_instance_gone(client, instance_id)
logger.debug(f"Deleted instance {instance_id}")
time.sleep(5)
output: dict[str, Any] = {"runs": all_rows}
if cluster_snapshot:
output["cluster"] = cluster_snapshot
if all_system_metrics:
output["system_metrics"] = all_system_metrics
if args.stdout:
json.dump(output, sys.stdout, indent=2, ensure_ascii=False)
elif args.json_out:
with open(args.json_out, "w", encoding="utf-8") as f:
json.dump(output, f, indent=2, ensure_ascii=False)
logger.debug(f"\nWrote results JSON: {args.json_out}")
return 0
if __name__ == "__main__":
raise SystemExit(main())