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

Code before:


def get_server(path_server: str, path_log: Optional[str]) -> dict:
    if os.environ.get("LLAMA_ARG_HOST") is None:
        logger.info("LLAMA_ARG_HOST not explicitly set, using 127.0.0.1")
        os.environ["LLAMA_ARG_HOST"] = "127.0.0.1"
    if os.environ.get("LLAMA_ARG_PORT") is None:
        logger.info("LLAMA_ARG_PORT not explicitly set, using 8080")
        os.environ["LLAMA_ARG_PORT"] = "8080"
    hostname: Optional[str] = os.environ.get("LLAMA_ARG_HOST")
    port: Optional[str] = os.environ.get("LLAMA_ARG_PORT")
    assert hostname is not None
    assert port is not None
    address: str = f"http://{hostname}:{port}"
    logger.info(f"Starting the llama.cpp server under {address}...")

    fout = open(path_log.format(port=port), "w") if path_log is not None else subprocess.DEVNULL
    process = subprocess.Popen([path_server], stdout=fout, stderr=subprocess.STDOUT)

    n_failures: int = 0
    while True:
        try:
            sleep(1.0)
            exit_code = process.poll()
            if exit_code is not None:
                raise RuntimeError(f"llama.cpp server exited unexpectedly with exit code {exit_code}, see {path_log.format(port=port)}")
            response = requests.get(f"{address}/health")
            if response.status_code == 200:
                break
        except requests.ConnectionError:
            n_failures += 1
            if n_failures >= 10:
                raise RuntimeError("llama.cpp server is not healthy after 10 seconds")

    return {"process": process, "address": address, "fout": fout}


def get_prompt_length(data: dict) -> int:
    session = data["session"]
    server_address: str = data["server_address"]

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

Review comment (suggestion):
```suggestion
                raise RuntimeError(f"llama.cpp server exited unexpectedly with exit code {exit_code}{path_log and f', see {path_log.format(port=port)}' or ''}")
```

Code after:


def get_server(path_server: str, path_log: Optional[str]) -> dict:
    if os.environ.get("LLAMA_ARG_HOST") is None:
        logger.info("LLAMA_ARG_HOST not explicitly set, using 127.0.0.1")
        os.environ["LLAMA_ARG_HOST"] = "127.0.0.1"
    if os.environ.get("LLAMA_ARG_PORT") is None:
        logger.info("LLAMA_ARG_PORT not explicitly set, using 8080")
        os.environ["LLAMA_ARG_PORT"] = "8080"
    hostname: Optional[str] = os.environ.get("LLAMA_ARG_HOST")
    port: Optional[str] = os.environ.get("LLAMA_ARG_PORT")
    assert hostname is not None
    assert port is not None
    address: str = f"http://{hostname}:{port}"
    logger.info(f"Starting the llama.cpp server under {address}...")

    fout = open(path_log.format(port=port), "w") if path_log is not None else subprocess.DEVNULL
    process = subprocess.Popen([path_server], stdout=fout, stderr=subprocess.STDOUT)

    n_failures: int = 0
    while True:
        try:
            sleep(1.0)
            exit_code = process.poll()
            if exit_code is not None:
                raise RuntimeError(f"llama.cpp server exited unexpectedly with exit code {exit_code}{path_log and f', see {path_log.format(port=port)}' or ''}")
            response = requests.get(f"{address}/health")
            if response.status_code == 200:
                break
        except requests.ConnectionError:
            n_failures += 1
            if n_failures >= 10:
                raise RuntimeError("llama.cpp server is not healthy after 10 seconds")

    return {"process": process, "address": address, "fout": fout}


def get_prompt_length(data: dict) -> int:
    session = data["session"]
    server_address: str = data["server_address"]

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