Repository: ggerganov/llama.cpp
File: scripts/server-bench.py (Python)

Code before:
        response = session.post(
            f"{server_address}/apply-template",
            json={"messages": [{"role": "user", "content": data["prompt"], "stream": True}]}
        )
        if response.status_code != 200:
            response_text = ""
            try:
                response_text = response.text
                response_text = ": {response_text}"
            except RuntimeError:
                pass
            raise RuntimeError(f"Server returned status code {response.status_code}{response_text}")
        prompt: str = json.loads(response.text)["prompt"]

        json_data: dict = {"prompt": prompt, "seed": data["seed"], "n_predict": data["n_predict"], "stream": True}
        response = session.post(f"{server_address}/completion", json=json_data, stream=True)

    lines = []
    token_arrival_times: list[float] = []
    for line in response.iter_lines(decode_unicode=False):
        if not line.startswith(b"data: "):
            continue
        lines.append(line)
        token_arrival_times.append(time())
    token_arrival_times = token_arrival_times[:-1]
    if "timings" in json.loads(lines[-2][6:]):
        token_arrival_times = token_arrival_times[:-1]

    if response.status_code != 200:
        response_text = ""
        try:
            response_text = response.text
            response_text = ": {response_text}"
        except RuntimeError:
            pass
        raise RuntimeError(f"Server returned status code {response.status_code}{response_text}")

    return (t_submit, token_arrival_times)


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"
    if not external_server and os.environ.get("LLAMA_ARG_N_GPU_LAYERS") is None:
        logger.info("LLAMA_ARG_N_GPU_LAYERS not explicitly set, using 999")
        os.environ["LLAMA_ARG_N_GPU_LAYERS"] = "999"
    if not external_server and os.environ.get("LLAMA_ARG_FLASH_ATTN") is None:

Review comment (suggestion):
```suggestion
    if len(lines) > 1 and "timings" in json.loads(lines[-2][6:]):
```

Code after:
            json={"messages": [{"role": "user", "content": data["prompt"], "stream": True}]}
        )
        response.raise_for_status()
        prompt: str = json.loads(response.text)["prompt"]

        json_data: dict = {"prompt": prompt, "seed": data["seed"], "n_predict": data["n_predict"], "stream": True}
        response = session.post(f"{server_address}/completion", json=json_data, stream=True)
    response.raise_for_status()

    lines = []
    token_arrival_times: list[float] = []
    for line in response.iter_lines(decode_unicode=False):
        if not line.startswith(b"data: "):
            continue
        lines.append(line)
        token_arrival_times.append(time())
    token_arrival_times = token_arrival_times[:-1]
    if len(lines) > 1 and "timings" in json.loads(lines[-2][6:]):
        token_arrival_times = token_arrival_times[:-1]

    return (t_submit, token_arrival_times)


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"
    if not external_server and os.environ.get("LLAMA_ARG_N_GPU_LAYERS") is None:
        logger.info("LLAMA_ARG_N_GPU_LAYERS not explicitly set, using 999")
        os.environ["LLAMA_ARG_N_GPU_LAYERS"] = "999"
    if not external_server and os.environ.get("LLAMA_ARG_FLASH_ATTN") is None:
        logger.info("LLAMA_ARG_FLASH_ATTN not explicitly set, using 'true'")
        os.environ["LLAMA_ARG_FLASH_ATTN"] = "true"

    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)