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

Code before:

    if 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"]

        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, "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")
        t0 = time()
        results: list[tuple[float, list[float]]] = thread_map(send_prompt, data, max_workers=parallel, chunksize=1)
    finally:
        if server is not None:
            server["process"].terminate()
            server["process"].wait()
        if session is not None:
            session.close()

    prompt_t = []
    token_t = []
    depth_sum: int = 0
    for pn, (t_submit, tat) in zip(prompt_n, results):
        prompt_t.append(tat[0] - t_submit)
        token_t += tat
        n_tokens: int = len(tat)
        depth_sum += n_tokens * pn
        depth_sum += n_tokens * (n_tokens + 1) // 2
    assert len(token_t) > 0

Review comment (suggestion):
```suggestion
                "n_predict": random.randint(n_predict_min, n_predict), "seed": (3 * (seed_offset + 1000 * i) + 2) if seed_offset >= 0 else -1})
```

Code after:

    if 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"]

        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, "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")
        t0 = time()
        results: list[tuple[float, list[float]]] = thread_map(send_prompt, data, max_workers=parallel, chunksize=1)
    finally:
        if server is not None:
            server["process"].terminate()
            server["process"].wait()
        if session is not None:
            session.close()

    prompt_t = []
    token_t = []
    depth_sum: int = 0
    for pn, (t_submit, tat) in zip(prompt_n, results):
        prompt_t.append(tat[0] - t_submit)
        token_t += tat
        n_tokens: int = len(tat)
        depth_sum += n_tokens * pn
        depth_sum += n_tokens * (n_tokens + 1) // 2
    assert len(token_t) > 0