Function benchmark

subprojects/llama.cpp/scripts/server-bench.py:147–271

(
        path_server: str, path_log: Optional[str], path_db: Optional[str], name: Optional[str], prompt_source: str, n_prompts: int,
        n_predict: int, n_predict_min: int, seed_offset: int)

def benchmark(
        path_server: str, path_log: Optional[str], path_db: Optional[str], name: Optional[str], prompt_source: str, n_prompts: int,
        n_predict: int, n_predict_min: int, seed_offset: int):
    external_server: bool = path_server.startswith("http://") or path_server.startswith("https://")
    if os.environ.get("LLAMA_ARG_N_PARALLEL") is None:
        logger.info("LLAMA_ARG_N_PARALLEL not explicitly set, using 32")
        os.environ["LLAMA_ARG_N_PARALLEL"] = "32"

    parallel: int = int(os.environ.get("LLAMA_ARG_N_PARALLEL")) # type: ignore
    prompts: Union[None, list[str], list[list[int]]] = get_prompts_text(prompt_source, n_prompts)
    synthetic_prompts: bool = prompts is None
    prompt_n = []

    if synthetic_prompts:
        prompt_source_split: list[str] = prompt_source.split("-")
        assert len(prompt_source_split) == 3
        assert prompt_source_split[0].lower() == "rng"
        prompt_length_min: int = int(prompt_source_split[1])
        prompt_length_max: int = int(prompt_source_split[2])
        logger.info("Generating random prompts...")
        prompt_n = get_prompt_lengths_rng(n_prompts, prompt_length_min, prompt_length_max, seed_offset)
        prompts = get_prompts_rng(prompt_n)
    else:
        n_predict_min = n_predict

    if not external_server and os.environ.get("LLAMA_ARG_CTX_SIZE") is None:
        context_per_slot: int = int(1.05 * (n_predict + (np.max(prompt_n) if synthetic_prompts else 2048)))
        context_total: int = context_per_slot * parallel
        os.environ["LLAMA_ARG_CTX_SIZE"] = str(context_total)
        logger.info(f"LLAMA_ARG_CTX_SIZE not explicitly set, using {context_total} ({context_per_slot} per slot).")

    server: Optional[dict] = None
    session = None
    try:
        server = get_server(path_server, path_log)
        server_address: str = server["address"]
        assert external_server == (server["process"] is None)

        adapter = requests.adapters.HTTPAdapter(pool_connections=parallel, pool_maxsize=parallel)  # type: ignore
        session = requests.Session()
        session.mount("http://", adapter)
        session.mount("https://", adapter)

        data: list[dict] = []

        for i, p in enumerate(prompts):
            if seed_offset >= 0:
                random.seed(3 * (seed_offset + 1000 * i) + 1)
            data.append({
                "session": session, "server_address": server_address, "external_server": external_server, "prompt": p,
                "synthetic_prompt": synthetic_prompts, "n_predict": random.randint(n_predict_min, n_predict),
                "seed": (3 * (seed_offset + 1000 * i) + 2) if seed_offset >= 0 else -1})

        if not synthetic_prompts:
            logger.info("Getting the prompt lengths...")
            prompt_n = [get_prompt_length(d) for d in data]

        logger.info("Starting the benchmark...\n")