code/scripts/generate.py · viai957/align at main

Add files using upload-large-folder tool

bcc8d35 verified 16 days ago

13.1 kB

#!/usr/bin/env python
"""Generate completions with the align engine (or the bucketed fallback).
    uv run python scripts/generate.py \
        --checkpoint /apps/development/rlhf/nano-mistral-o1/checkpoints/sft \
        --prompts data/eval/math500.json --out .tmp/gen/math500.jsonl
Sharding across the 4 A10Gs is done by the CALLER, one process per device --
the engine is single-replica and knows nothing about other GPUs (no
collectives, no NCCL). The idiom, which respects the one-heavy-job rule:
    for i in 0 1 2 3; do
      RANK=$i CUDA_VISIBLE_DEVICES=$i uv run python scripts/generate.py \
        --shard-index $i --num-shards 4 --out .tmp/gen/shard$i.jsonl &
    done; wait
`acquire_heavy_job_lock` is rank-aware, so only RANK=0 takes the exclusive
flock; shards 1.. do NOT take it, they VERIFY it (`verify_shard_lock`): a lock
held by anything other than another generate.py (their own shard 0) refuses
the launch, so a live training stage cannot be overlapped by three of the four
shards. Every shard also runs the host preflight. Never launch this against a
live training stage: that overlap is what killed the predecessor's first SFT
run. Completions carry rid "<shard>:<local i>"; the global prompt index of
local i is shard_index + i * num_shards (stride sharding).
Every output path is asserted onto /apps (`/` has ~13GB and /tmp is a noexec
tmpfs). The launch banner prints the checkpoint fingerprint, the attention
backend that ACTUALLY ran, the engine budget, the sampling params and the
per-gate status; a missing mandatory key raises rather than printing a partial
banner, and gate status is "PENDING" until recorded -- never assumed passed.
"""

from __future__ import annotations

import argparse
import json
import logging
import os
import sys
import time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from align.utils.guard import apply_resource_guards  # noqa: E402

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
logger = logging.getLogger("generate")

PROJECT_ROOT = "/apps/development/rlhf/nano-mistral-o1"


def validate_shards(shard_index: int, num_shards: int) -> None:
    """An out-of-range shard silently yields an EMPTY prompt list (and a "successful" empty run)."""
    if num_shards < 1 or not (0 <= shard_index < num_shards):
        raise ValueError(f"bad shard spec: --shard-index {shard_index} must be in [0, --num-shards {num_shards})")


def _proc_cmdline(pid: int) -> str:
    try:
        with open(f"/proc/{pid}/cmdline", "rb") as f:
            return f.read().replace(b"\0", b" ").decode(errors="replace")
    except OSError:
        return ""


def verify_shard_lock(rank: int, lock_inode: int | None = None, locks_text: str | None = None, cmdline_of=_proc_cmdline) -> None:
    """Shards 1.. never take the heavy-job flock (`acquire_heavy_job_lock` returns True for RANK != 0),
    so without this a generation launched while a trainer holds it would be refused on shard 0 only,
    while shards 1-3 each loaded a 7B onto GPUs 1-3. Read-only: scans /proc/locks for the holder of the
    lock file (taking the flock here would race shard 0's own non-blocking acquire). The only holder
    tolerated is another generate.py -- this launch's own shard 0."""
    if rank == 0:
        return
    from align.exit_codes import PreflightRefused
    from align.utils.guard import LOCK_PATH

    if lock_inode is None:
        try:
            lock_inode = os.stat(LOCK_PATH).st_ino
        except FileNotFoundError:
            return  # no lock file yet, so no holder
    if locks_text is None:
        with open("/proc/locks") as f:
            locks_text = f.read()
    for line in locks_text.splitlines():
        parts = line.split()
        if len(parts) < 6 or parts[1] != "FLOCK" or parts[5].split(":")[-1] != str(lock_inode):
            continue
        holder = cmdline_of(int(parts[4]))
        if "generate.py" not in holder:
            raise PreflightRefused(
                f"another heavy align job holds the lock (pid {parts[4]}: {holder.strip()[:120]!r}); this shard "
                f"(RANK={rank}) must not overlap it. Shard 0 refuses for the same reason."
            )


def read_prompts(path: str) -> list[str]:
    """Accept .json (list of str, or list of dicts with a prompt/problem/text
    key), .jsonl (same per line), or .txt (one prompt per line). Refuses an
    empty file loudly: an eval that silently generates nothing looks like a
    fast eval."""

    def _field(obj):
        if isinstance(obj, str):
            return obj
        for k in ("prompt", "problem", "question", "text"):
            if k in obj:
                return obj[k]
        raise ValueError(
            f"prompt record has none of prompt/problem/question/text: {sorted(obj)[:8]}"
        )

    if path.endswith(".jsonl"):
        with open(path) as f:
            rows = [_field(json.loads(line)) for line in f if line.strip()]
    elif path.endswith(".json"):
        with open(path) as f:
            rows = [_field(o) for o in json.load(f)]
    else:
        with open(path) as f:
            rows = [line.rstrip("\n") for line in f if line.strip()]
    if not rows:
        raise ValueError(f"{path} contains no prompts")
    return rows


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument(
        "--checkpoint", required=True, help="HF-layout checkpoint directory (safetensors)"
    )
    ap.add_argument("--prompts", required=True)
    ap.add_argument("--out", required=True, help="completions JSONL; under /apps")
    ap.add_argument("--tokenizer", default=None, help="defaults to --checkpoint")
    ap.add_argument("--n-slots", type=int, default=32)
    ap.add_argument("--max-seq-len", type=int, default=1024)
    ap.add_argument("--max-new-tokens", type=int, default=300)
    ap.add_argument("--temperature", type=float, default=0.0)
    ap.add_argument("--top-p", type=float, default=1.0)
    ap.add_argument(
        "--stop", action="append", default=[], help="repeatable; matched on DECODED TEXT"
    )
    ap.add_argument("--stop-policy", default="trim", choices=("trim", "keep"))
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument(
        "--attn", default=None, help="force an attention backend; unavailable => hard error"
    )
    ap.add_argument(
        "--bucketed",
        action="store_true",
        help="use the static-batching fallback instead of the "
        "engine (correctness reference; slower)",
    )
    ap.add_argument(
        "--benchmark",
        action="store_true",
        help="run BOTH and write the engine-vs-static-batching speedup receipt beside --out",
    )
    ap.add_argument(
        "--hash-weights",
        action="store_true",
        help="full content hash of the checkpoint (reads 14.5GB)",
    )
    ap.add_argument("--limit", type=int, default=None)
    ap.add_argument("--shard-index", type=int, default=0)
    ap.add_argument("--num-shards", type=int, default=1)
    ap.add_argument("--strict-host-check", default=True, action=argparse.BooleanOptionalAction)
    ap.add_argument(
        "--heavy-lock",
        default=True,
        action=argparse.BooleanOptionalAction,
        help="take the exclusive one-heavy-job flock. Keep it ON "
        "for any real checkpoint. --no-heavy-lock exists for "
        "the tiny-model CLI smoke test, which loads ~200KB "
        "and is not a heavy job; using it to squeeze a 7B "
        "run past a live training stage is the exact overlap "
        "that killed the predecessor's first SFT run.",
    )
    args = ap.parse_args()
    try:
        validate_shards(args.shard_index, args.num_shards)
    except ValueError as e:
        ap.error(str(e))

    apply_resource_guards(
        cpu_threads=8,
        heavy=args.heavy_lock,
        label="generate",
        strict_host_check=args.strict_host_check,
    )
    if args.heavy_lock:
        verify_shard_lock(int(os.environ.get("RANK", "0")))

    import torch

    from align.inference.bucketed import bucketed_generate
    from align.inference.engine import (
        Engine,
        GenerationRequest,
        InferenceConfig,
        benchmark_engine_vs_bucketed,
    )
    from align.inference.receipts import (
        checkpoint_fingerprint,
        engine_banner,
        load_gate_status,
        sha256_file,
        write_completions_jsonl,
    )
    from align.model import checkpoint as ckpt
    from align.model.mistral import NmoMistralForCausalLM
    from align.registry import render_banner
    from align.utils.guard import assert_paths_off_root

    assert_paths_off_root([os.path.dirname(os.path.abspath(args.out))])

    prompts = read_prompts(args.prompts)
    if args.limit:
        prompts = prompts[: args.limit]
    # Shard by stride, not by contiguous block: prompt files are often sorted
    # by length or difficulty, and contiguous blocks would give one replica
    # all the long ones and skew both the wall-clock and the benchmark.
    if args.num_shards > 1:
        prompts = prompts[args.shard_index :: args.num_shards]
        logger.info("shard %d/%d: %d prompts", args.shard_index, args.num_shards, len(prompts))

    from align.eval.golden import resolve_model_config
    from align.tokenizer import load_tokenizer

    tok = load_tokenizer(args.tokenizer or args.checkpoint)
    state = ckpt.load_hf_safetensors(args.checkpoint)
    # `read_config` (align.model.hf_config) states every architecture field
    # explicitly and REFUSES a config.json that omits one: transformers'
    # MistralConfig used to resolve the missing ones to its own defaults
    # (sliding_window=4096, rms_norm_eps=1e-6 -- neither of them v0.3), which
    # silently accepted a checkpoint that was a different model.
    model_cfg = resolve_model_config(args.checkpoint)
    model = NmoMistralForCausalLM.from_state_dict(
        model_cfg,
        ckpt.hf_to_nmo(state),
        attn_implementation=args.attn,
        device="cuda" if torch.cuda.is_available() else "cpu",
    ).eval()
    if torch.cuda.is_available():
        model = model.to("cuda", dtype=torch.bfloat16)

    cfg = InferenceConfig(
        n_slots=args.n_slots,
        max_seq_len=args.max_seq_len,
        max_new_tokens=args.max_new_tokens,
        temperature=args.temperature,
        top_p=args.top_p,
        stop=tuple(args.stop),
        stop_string_policy=args.stop_policy,
        seed=args.seed,
        eos_token_id=tok.eos_token_id,
        pad_token_id=tok.pad_token_id if tok.pad_token_id is not None else tok.eos_token_id,
    )
    engine = Engine(model, cfg, tokenizer=tok)

    requests = [
        GenerationRequest(
            rid=f"{args.shard_index}:{i}", prompt_ids=tok(p, add_special_tokens=True)["input_ids"]
        )
        for i, p in enumerate(prompts)
    ]

    gate_status = load_gate_status(args.checkpoint)
    banner = engine_banner(
        engine=engine,
        checkpoint=checkpoint_fingerprint(args.checkpoint, args.hash_weights),
        gate_status=gate_status,
        prompts_sha256=sha256_file(args.prompts),
        extra={
            "n_prompts": len(requests),
            "path_bucketed": args.bucketed,
            "shard": f"{args.shard_index}/{args.num_shards}",
            "global_prompt_index": f"{args.shard_index} + i * {args.num_shards} (i = local index in rid)",
        },
    )
    print(render_banner(banner, title="align generate"))
    if not banner["engine_generated_numbers_usable"] and not args.bucketed:
        logger.warning(
            "the 3-part engine gate has NOT fully passed for this checkpoint "
            "(%s). These completions are fine for inspection; section-10 must "
            "use --bucketed until gate (c) is recorded.",
            gate_status,
        )

    t0 = time.perf_counter()
    if args.bucketed:
        results = bucketed_generate(model, requests, cfg, tokenizer=tok)
    else:
        results = engine.generate(requests)
    wall = time.perf_counter() - t0
    n_tokens = sum(r.n_generated_tokens for r in results)
    receipt = dict(banner)
    receipt["throughput"] = {
        "wall_s": round(wall, 3),
        "generated_tokens": n_tokens,
        "tokens_per_s": round(n_tokens / wall, 2) if wall > 0 else float("nan"),
    }
    write_completions_jsonl(args.out, results, receipt, tokenizer=tok)
    logger.info(
        "%d completions, %d tokens, %.1fs (%.1f tok/s)",
        len(results),
        n_tokens,
        wall,
        n_tokens / wall if wall else float("nan"),
    )

    if args.benchmark:
        bench_path = os.path.splitext(args.out)[0] + ".benchmark.json"
        rec = benchmark_engine_vs_bucketed(model, requests, cfg, tokenizer=tok, out_path=bench_path)
        logger.info(
            "engine vs static batching: measured speedup %.2fx "
            "(engine %.1f tok/s, bucketed %.1f tok/s)",
            rec["speedup"],
            rec["engine_tokens_per_s"],
            rec["bucketed_tokens_per_s"],
        )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())

	#!/usr/bin/env python
	"""Generate completions with the align engine (or the bucketed fallback).

	uv run python scripts/generate.py \
	--checkpoint /apps/development/rlhf/nano-mistral-o1/checkpoints/sft \
	--prompts data/eval/math500.json --out .tmp/gen/math500.jsonl

	Sharding across the 4 A10Gs is done by the CALLER, one process per device --
	the engine is single-replica and knows nothing about other GPUs (no
	collectives, no NCCL). The idiom, which respects the one-heavy-job rule:

	for i in 0 1 2 3; do
	RANK=$i CUDA_VISIBLE_DEVICES=$i uv run python scripts/generate.py \
	--shard-index $i --num-shards 4 --out .tmp/gen/shard$i.jsonl &
	done; wait

	`acquire_heavy_job_lock` is rank-aware, so only RANK=0 takes the exclusive
	flock; shards 1.. do NOT take it, they VERIFY it (`verify_shard_lock`): a lock
	held by anything other than another generate.py (their own shard 0) refuses
	the launch, so a live training stage cannot be overlapped by three of the four
	shards. Every shard also runs the host preflight. Never launch this against a
	live training stage: that overlap is what killed the predecessor's first SFT
	run. Completions carry rid "<shard>:<local i>"; the global prompt index of
	local i is shard_index + i * num_shards (stride sharding).

	Every output path is asserted onto /apps (`/` has ~13GB and /tmp is a noexec
	tmpfs). The launch banner prints the checkpoint fingerprint, the attention
	backend that ACTUALLY ran, the engine budget, the sampling params and the
	per-gate status; a missing mandatory key raises rather than printing a partial
	banner, and gate status is "PENDING" until recorded -- never assumed passed.
	"""

	from __future__ import annotations

	import argparse
	import json
	import logging
	import os
	import sys
	import time
	from pathlib import Path

	sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

	from align.utils.guard import apply_resource_guards # noqa: E402

	logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
	logger = logging.getLogger("generate")

	PROJECT_ROOT = "/apps/development/rlhf/nano-mistral-o1"


	def validate_shards(shard_index: int, num_shards: int) -> None:
	"""An out-of-range shard silently yields an EMPTY prompt list (and a "successful" empty run)."""
	if num_shards < 1 or not (0 <= shard_index < num_shards):
	raise ValueError(f"bad shard spec: --shard-index {shard_index} must be in [0, --num-shards {num_shards})")


	def _proc_cmdline(pid: int) -> str:
	try:
	with open(f"/proc/{pid}/cmdline", "rb") as f:
	return f.read().replace(b"\0", b" ").decode(errors="replace")
	except OSError:
	return ""


	def verify_shard_lock(rank: int, lock_inode: int \| None = None, locks_text: str \| None = None, cmdline_of=_proc_cmdline) -> None:
	"""Shards 1.. never take the heavy-job flock (`acquire_heavy_job_lock` returns True for RANK != 0),
	so without this a generation launched while a trainer holds it would be refused on shard 0 only,
	while shards 1-3 each loaded a 7B onto GPUs 1-3. Read-only: scans /proc/locks for the holder of the
	lock file (taking the flock here would race shard 0's own non-blocking acquire). The only holder
	tolerated is another generate.py -- this launch's own shard 0."""
	if rank == 0:
	return
	from align.exit_codes import PreflightRefused
	from align.utils.guard import LOCK_PATH

	if lock_inode is None:
	try:
	lock_inode = os.stat(LOCK_PATH).st_ino
	except FileNotFoundError:
	return # no lock file yet, so no holder
	if locks_text is None:
	with open("/proc/locks") as f:
	locks_text = f.read()
	for line in locks_text.splitlines():
	parts = line.split()
	if len(parts) < 6 or parts[1] != "FLOCK" or parts[5].split(":")[-1] != str(lock_inode):
	continue
	holder = cmdline_of(int(parts[4]))
	if "generate.py" not in holder:
	raise PreflightRefused(
	f"another heavy align job holds the lock (pid {parts[4]}: {holder.strip()[:120]!r}); this shard "
	f"(RANK={rank}) must not overlap it. Shard 0 refuses for the same reason."
	)


	def read_prompts(path: str) -> list[str]:
	"""Accept .json (list of str, or list of dicts with a prompt/problem/text
	key), .jsonl (same per line), or .txt (one prompt per line). Refuses an
	empty file loudly: an eval that silently generates nothing looks like a
	fast eval."""

	def _field(obj):
	if isinstance(obj, str):
	return obj
	for k in ("prompt", "problem", "question", "text"):
	if k in obj:
	return obj[k]
	raise ValueError(
	f"prompt record has none of prompt/problem/question/text: {sorted(obj)[:8]}"
	)

	if path.endswith(".jsonl"):
	with open(path) as f:
	rows = [_field(json.loads(line)) for line in f if line.strip()]
	elif path.endswith(".json"):
	with open(path) as f:
	rows = [_field(o) for o in json.load(f)]
	else:
	with open(path) as f:
	rows = [line.rstrip("\n") for line in f if line.strip()]
	if not rows:
	raise ValueError(f"{path} contains no prompts")
	return rows


	def main() -> int:
	ap = argparse.ArgumentParser()
	ap.add_argument(
	"--checkpoint", required=True, help="HF-layout checkpoint directory (safetensors)"
	)
	ap.add_argument("--prompts", required=True)
	ap.add_argument("--out", required=True, help="completions JSONL; under /apps")
	ap.add_argument("--tokenizer", default=None, help="defaults to --checkpoint")
	ap.add_argument("--n-slots", type=int, default=32)
	ap.add_argument("--max-seq-len", type=int, default=1024)
	ap.add_argument("--max-new-tokens", type=int, default=300)
	ap.add_argument("--temperature", type=float, default=0.0)
	ap.add_argument("--top-p", type=float, default=1.0)
	ap.add_argument(
	"--stop", action="append", default=[], help="repeatable; matched on DECODED TEXT"
	)
	ap.add_argument("--stop-policy", default="trim", choices=("trim", "keep"))
	ap.add_argument("--seed", type=int, default=0)
	ap.add_argument(
	"--attn", default=None, help="force an attention backend; unavailable => hard error"
	)
	ap.add_argument(
	"--bucketed",
	action="store_true",
	help="use the static-batching fallback instead of the "
	"engine (correctness reference; slower)",
	)
	ap.add_argument(
	"--benchmark",
	action="store_true",
	help="run BOTH and write the engine-vs-static-batching speedup receipt beside --out",
	)
	ap.add_argument(
	"--hash-weights",
	action="store_true",
	help="full content hash of the checkpoint (reads 14.5GB)",
	)
	ap.add_argument("--limit", type=int, default=None)
	ap.add_argument("--shard-index", type=int, default=0)
	ap.add_argument("--num-shards", type=int, default=1)
	ap.add_argument("--strict-host-check", default=True, action=argparse.BooleanOptionalAction)
	ap.add_argument(
	"--heavy-lock",
	default=True,
	action=argparse.BooleanOptionalAction,
	help="take the exclusive one-heavy-job flock. Keep it ON "
	"for any real checkpoint. --no-heavy-lock exists for "
	"the tiny-model CLI smoke test, which loads ~200KB "
	"and is not a heavy job; using it to squeeze a 7B "
	"run past a live training stage is the exact overlap "
	"that killed the predecessor's first SFT run.",
	)
	args = ap.parse_args()
	try:
	validate_shards(args.shard_index, args.num_shards)
	except ValueError as e:
	ap.error(str(e))

	apply_resource_guards(
	cpu_threads=8,
	heavy=args.heavy_lock,
	label="generate",
	strict_host_check=args.strict_host_check,
	)
	if args.heavy_lock:
	verify_shard_lock(int(os.environ.get("RANK", "0")))

	import torch

	from align.inference.bucketed import bucketed_generate
	from align.inference.engine import (
	Engine,
	GenerationRequest,
	InferenceConfig,
	benchmark_engine_vs_bucketed,
	)
	from align.inference.receipts import (
	checkpoint_fingerprint,
	engine_banner,
	load_gate_status,
	sha256_file,
	write_completions_jsonl,
	)
	from align.model import checkpoint as ckpt
	from align.model.mistral import NmoMistralForCausalLM
	from align.registry import render_banner
	from align.utils.guard import assert_paths_off_root

	assert_paths_off_root([os.path.dirname(os.path.abspath(args.out))])

	prompts = read_prompts(args.prompts)
	if args.limit:
	prompts = prompts[: args.limit]
	# Shard by stride, not by contiguous block: prompt files are often sorted
	# by length or difficulty, and contiguous blocks would give one replica
	# all the long ones and skew both the wall-clock and the benchmark.
	if args.num_shards > 1:
	prompts = prompts[args.shard_index :: args.num_shards]
	logger.info("shard %d/%d: %d prompts", args.shard_index, args.num_shards, len(prompts))

	from align.eval.golden import resolve_model_config
	from align.tokenizer import load_tokenizer

	tok = load_tokenizer(args.tokenizer or args.checkpoint)
	state = ckpt.load_hf_safetensors(args.checkpoint)
	# `read_config` (align.model.hf_config) states every architecture field
	# explicitly and REFUSES a config.json that omits one: transformers'
	# MistralConfig used to resolve the missing ones to its own defaults
	# (sliding_window=4096, rms_norm_eps=1e-6 -- neither of them v0.3), which
	# silently accepted a checkpoint that was a different model.
	model_cfg = resolve_model_config(args.checkpoint)
	model = NmoMistralForCausalLM.from_state_dict(
	model_cfg,
	ckpt.hf_to_nmo(state),
	attn_implementation=args.attn,
	device="cuda" if torch.cuda.is_available() else "cpu",
	).eval()
	if torch.cuda.is_available():
	model = model.to("cuda", dtype=torch.bfloat16)

	cfg = InferenceConfig(
	n_slots=args.n_slots,
	max_seq_len=args.max_seq_len,
	max_new_tokens=args.max_new_tokens,
	temperature=args.temperature,
	top_p=args.top_p,
	stop=tuple(args.stop),
	stop_string_policy=args.stop_policy,
	seed=args.seed,
	eos_token_id=tok.eos_token_id,
	pad_token_id=tok.pad_token_id if tok.pad_token_id is not None else tok.eos_token_id,
	)
	engine = Engine(model, cfg, tokenizer=tok)

	requests = [
	GenerationRequest(
	rid=f"{args.shard_index}:{i}", prompt_ids=tok(p, add_special_tokens=True)["input_ids"]
	)
	for i, p in enumerate(prompts)
	]

	gate_status = load_gate_status(args.checkpoint)
	banner = engine_banner(
	engine=engine,
	checkpoint=checkpoint_fingerprint(args.checkpoint, args.hash_weights),
	gate_status=gate_status,
	prompts_sha256=sha256_file(args.prompts),
	extra={
	"n_prompts": len(requests),
	"path_bucketed": args.bucketed,
	"shard": f"{args.shard_index}/{args.num_shards}",
	"global_prompt_index": f"{args.shard_index} + i * {args.num_shards} (i = local index in rid)",
	},
	)
	print(render_banner(banner, title="align generate"))
	if not banner["engine_generated_numbers_usable"] and not args.bucketed:
	logger.warning(
	"the 3-part engine gate has NOT fully passed for this checkpoint "
	"(%s). These completions are fine for inspection; section-10 must "
	"use --bucketed until gate (c) is recorded.",
	gate_status,
	)

	t0 = time.perf_counter()
	if args.bucketed:
	results = bucketed_generate(model, requests, cfg, tokenizer=tok)
	else:
	results = engine.generate(requests)
	wall = time.perf_counter() - t0
	n_tokens = sum(r.n_generated_tokens for r in results)
	receipt = dict(banner)
	receipt["throughput"] = {
	"wall_s": round(wall, 3),
	"generated_tokens": n_tokens,
	"tokens_per_s": round(n_tokens / wall, 2) if wall > 0 else float("nan"),
	}
	write_completions_jsonl(args.out, results, receipt, tokenizer=tok)
	logger.info(
	"%d completions, %d tokens, %.1fs (%.1f tok/s)",
	len(results),
	n_tokens,
	wall,
	n_tokens / wall if wall else float("nan"),
	)

	if args.benchmark:
	bench_path = os.path.splitext(args.out)[0] + ".benchmark.json"
	rec = benchmark_engine_vs_bucketed(model, requests, cfg, tokenizer=tok, out_path=bench_path)
	logger.info(
	"engine vs static batching: measured speedup %.2fx "
	"(engine %.1f tok/s, bucketed %.1f tok/s)",
	rec["speedup"],
	rec["engine_tokens_per_s"],
	rec["bucketed_tokens_per_s"],
	)
	return 0


	if __name__ == "__main__":
	raise SystemExit(main())