nemo_automodel.components.speculative.bench_common

Engine-agnostic workload machinery shared by the drafter benchmarks.

bench_sglang and bench_vllm drive the same OpenAI-style chat-completions workload and report the same throughput/speedup numbers; they differ only in how each engine exposes its acceptance statistics (SGLang’s /server_info vs vLLM’s Prometheus /metrics). This module holds the shared layer: the prompt loader, the timed HTTP workload runner, and the throughput / speedup / CLI-bounds helpers.

Module Contents

Classes

Name	Description
`WorkloadResult`	Aggregate timing for one workload pass against a server.

Functions

Name	Description
`_chat_completion`	POST one chat completion and return its `completion_tokens` (0 on no usage).
`_extract_prompt_text`	Normalize a raw (non-chat) dataset field into a single prompt string.
`_load_prompts`	Load up to `--num-prompts` chat prompts.
`_normalize_server_url`	Return the server root URL without a trailing slash or `/v1` suffix.
`_output_throughput`	Output tokens per wall-clock second, or `None` if nothing was timed.
`_report_summary`	Print a benchmark summary as JSON, optionally also write it to `output_json`.
`_run_workload`	Send every prompt through `<server>/v1/chat/completions` and time the pass.
`_speedup`	Return `spec / baseline` output throughput, or `None` if not computable.
`_validate_workload_args`	Reject invalid values of the CLI flags every benchmark shares.

Data

logger

API

class nemo_automodel.components.speculative.bench_common.WorkloadResult(
    wall_clock_s: float,
    output_tokens: int,
    completed: int,
    failed: int
)

Dataclass

Aggregate timing for one workload pass against a server.

completed

int

failed

int

output_tokens

int

wall_clock_s

float

nemo_automodel.components.speculative.bench_common._chat_completion(
    session,
    url: str,
    payload: dict[str, typing.Any],
    timeout_s: float,
    max_retries: int
) -> int

async

POST one chat completion and return its completion_tokens (0 on no usage).

nemo_automodel.components.speculative.bench_common._extract_prompt_text(
    value: typing.Any
) -> str | None

Normalize a raw (non-chat) dataset field into a single prompt string.

Accepts a plain string, or a non-empty list of strings (the first entry is used — e.g. MT-Bench’s two-turn prompt column, reduced to its first turn). Returns None for anything else, or a blank/whitespace-only string.

nemo_automodel.components.speculative.bench_common._load_prompts(
    args: argparse.Namespace
) -> list[list[dict[str, typing.Any]]]

Load up to --num-prompts chat prompts.

Reads from --prompt-column when set (a raw single-turn text field, wrapped into a fresh user message — for datasets like GSM8K/HumanEval that are not chat-messages-shaped), else from --messages-column (an existing OpenAI-messages list, trailing assistant turn dropped). prompt_column is an optional attribute — callers that only ever use messages_column (the two single-dataset benchmarks’ own argparse.Namespace) need not set it.

nemo_automodel.components.speculative.bench_common._normalize_server_url(
    url: str
) -> str

Return the server root URL without a trailing slash or /v1 suffix.

Chat completions live at <root>/v1/chat/completions and the engine’s stats endpoint at its own root-relative path; accept either http://host:port or the OpenAI-style http://host:port/v1 so the flag is forgiving.

nemo_automodel.components.speculative.bench_common._output_throughput(
    result: nemo_automodel.components.speculative.bench_common.WorkloadResult
) -> float | None

Output tokens per wall-clock second, or None if nothing was timed.

nemo_automodel.components.speculative.bench_common._report_summary(
    summary: dict[str, typing.Any] | None,
    output_json: str | None
) -> int

Print a benchmark summary as JSON, optionally also write it to output_json.

Shared tail of bench_sglang/bench_vllm’s _run: returns exit code 1 with nothing printed when summary is None (no usable prompts were loaded), else 0.

nemo_automodel.components.speculative.bench_common._run_workload(
    server: str,
    prompts: list[list[dict[str, typing.Any]]],
    gen_cfg: nemo_automodel.components.speculative.regenerate.GenerationConfig,
    concurrency: int,
    timeout_s: float,
    max_retries: int
) -> nemo_automodel.components.speculative.bench_common.WorkloadResult

async

Send every prompt through <server>/v1/chat/completions and time the pass.

nemo_automodel.components.speculative.bench_common._speedup(
    spec_throughput: float | None,
    baseline_throughput: float | None
) -> float | None

Return spec / baseline output throughput, or None if not computable.

nemo_automodel.components.speculative.bench_common._validate_workload_args(
    args: argparse.Namespace
) -> None

Reject invalid values of the CLI flags every benchmark shares.

nemo_automodel.components.speculative.bench_common.logger = logging.getLogger(__name__)