561 lines
21 KiB
Python
561 lines
21 KiB
Python
# 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 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 harness import (
|
||
ExoClient,
|
||
ExoHttpError,
|
||
add_common_instance_args,
|
||
ensure_cuda_available,
|
||
instance_id_from_instance,
|
||
nodes_used_in_instance,
|
||
resolve_model_short_id,
|
||
run_planning_phase,
|
||
settle_and_fetch_placements,
|
||
validate_vllm_args,
|
||
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"],
|
||
)
|
||
)
|
||
|
||
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
|
||
|
||
# Default: use AutoTokenizer
|
||
return AutoTokenizer.from_pretrained(model_id, 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...")
|
||
|
||
|
||
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
|
||
) -> 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}],
|
||
"stream": False,
|
||
"max_tokens": tg,
|
||
}
|
||
|
||
t0 = time.perf_counter()
|
||
out = client.post_bench_chat_completions(payload)
|
||
elapsed = time.perf_counter() - t0
|
||
|
||
stats = out.get("generation_stats")
|
||
|
||
# Extract preview, handling None content (common for thinking models)
|
||
choices = out.get("choices") or [{}]
|
||
message = choices[0].get("message", {}) if choices else {}
|
||
content = message.get("content") or ""
|
||
preview = content[:200] if content else ""
|
||
|
||
return {
|
||
"elapsed_s": elapsed,
|
||
"output_text_preview": preview,
|
||
"stats": stats,
|
||
}, 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}]
|
||
ids = tokenizer.apply_chat_template(
|
||
messages, tokenize=True, add_generation_prompt=True
|
||
)
|
||
# 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.",
|
||
)
|
||
args = ap.parse_args()
|
||
validate_vllm_args(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
|
||
|
||
# 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)
|
||
if args.ensure_cuda:
|
||
ensure_cuda_available(client)
|
||
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
|
||
|
||
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")
|
||
|
||
all_rows: list[dict[str, Any]] = []
|
||
|
||
for preview in selected:
|
||
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}"
|
||
)
|
||
|
||
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)
|
||
|
||
try:
|
||
for i in range(args.warmup):
|
||
run_one_completion(
|
||
client, full_model_id, pp_list[0], tg_list[0], prompt_sizer
|
||
)
|
||
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]] = []
|
||
for r in range(args.repeat):
|
||
time.sleep(3)
|
||
|
||
if concurrency <= 1:
|
||
# Sequential: single request
|
||
try:
|
||
row, actual_pp_tokens = run_one_completion(
|
||
client, full_model_id, pp, tg, prompt_sizer
|
||
)
|
||
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
|
||
# Each thread gets its own ExoClient (separate HTTP connection)
|
||
batch_results: list[tuple[dict[str, Any], int]] = []
|
||
batch_errors = 0
|
||
|
||
def _run_concurrent(
|
||
idx: int, *, _pp: int = pp, _tg: int = tg
|
||
) -> tuple[dict[str, Any], int]:
|
||
c = ExoClient(
|
||
args.host, args.port, timeout_s=args.timeout
|
||
)
|
||
return run_one_completion(
|
||
c, full_model_id, _pp, _tg, prompt_sizer
|
||
)
|
||
|
||
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
|
||
|
||
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
|
||
]
|
||
agg_gen_tps = (
|
||
mean(valid_gen_tps) if valid_gen_tps else 0.0
|
||
)
|
||
gen_tps = agg_gen_tps / concurrency
|
||
logger.info(
|
||
f"[concurrent {concurrency}x] "
|
||
f"agg_gen_tps={agg_gen_tps:.2f} "
|
||
f"gen_tps={gen_tps:.2f} "
|
||
f"errors={batch_errors}"
|
||
)
|
||
|
||
if runs:
|
||
prompt_tps = mean(x["stats"]["prompt_tps"] for x in runs)
|
||
gen_tps = mean(
|
||
x["stats"]["generation_tps"] / x["concurrency"]
|
||
for x in runs
|
||
)
|
||
ptok = mean(x["stats"]["prompt_tokens"] for x in runs)
|
||
gtok = mean(x["stats"]["generation_tokens"] for x in runs)
|
||
peak = mean(
|
||
x["stats"]["peak_memory_usage"]["inBytes"] for x in runs
|
||
)
|
||
|
||
logger.info(
|
||
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)}\n"
|
||
)
|
||
time.sleep(2)
|
||
finally:
|
||
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)
|
||
|
||
if args.stdout:
|
||
json.dump(all_rows, sys.stdout, indent=2, ensure_ascii=False)
|
||
elif args.json_out:
|
||
with open(args.json_out, "w", encoding="utf-8") as f:
|
||
json.dump(all_rows, f, indent=2, ensure_ascii=False)
|
||
logger.debug(f"\nWrote results JSON: {args.json_out}")
|
||
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
raise SystemExit(main())
|