Function throughput_command

tensorrt_llm/bench/benchmark/throughput.py:292–487

Run a throughput test on a TRT-LLM engine.

(
    bench_env: BenchmarkEnvironment,
    **params,
)

)
@click.pass_obj
def throughput_command(
    bench_env: BenchmarkEnvironment,
    **params,
) -> None:
    """Run a throughput test on a TRT-LLM engine."""
    logger.info("Preparing to run throughput benchmark...")

    # Parameters from CLI
    image_data_format: str = params.get("image_data_format", "pt")
    data_device: str = params.get("data_device", "cpu")
    no_skip_tokenizer_init: bool = params.get("no_skip_tokenizer_init", False)

    # Get general CLI options using the centralized function
    options: GeneralExecSettings = get_general_cli_options(params, bench_env)
    tokenizer = initialize_tokenizer(options.checkpoint_path)

    # Extract throughput-specific options not handled by GeneralExecSettings
    max_batch_size = params.get("max_batch_size")
    max_num_tokens = params.get("max_num_tokens")
    enable_chunked_context: bool = params.get("enable_chunked_context")
    scheduler_policy: str = params.get("scheduler_policy")

    custom_module_dirs: list[Path] = params.pop("custom_module_dirs", [])
    for custom_module_dir in custom_module_dirs:
        try:
            import_custom_module_from_dir(custom_module_dir)
        except Exception as e:
            logger.error(
                f"Failed to import custom module from {custom_module_dir}: {e}")
            raise e

    # Runtime kwargs and option tracking.
    kwargs = {}

    # Dataset Loading and Preparation
    with open(options.dataset_path, "r") as dataset:
        metadata, requests = create_dataset_from_stream(
            tokenizer,
            dataset,
            num_requests=options.num_requests,
            model_dir=options.checkpoint_path,
            model_type=options.model_type,
            modality=options.modality,
            image_data_format=image_data_format,
            data_device=data_device,
            max_input_seq_len_for_multimodal=options.max_input_len)
        metadata.dataset_path = options.dataset_path
        params["target_input_len"] = params.get(
            "target_input_len") or metadata.avg_isl
        params["target_output_len"] = params.get(
            "target_output_len") or metadata.avg_osl

    if options.modality is None:
        # Log dataset info
        # NOTE: This table is only accurate for non-multimodal models.
        #       The accurate table for multimodal models will be logged after the benchmark is done.
        logger.info(metadata.get_summary_for_print())