bench.py

import multiprocessing as mp
import os
import random
import sys
import time
from typing import Any, AsyncIterable, Callable, Dict, Iterable, List

import backoff
import grpc
import pandas as pd
from tqdm import tqdm

from llm_perf.utils.logger import logger
from llm_perf.utils.reporter import ReportType

sys.path.append(os.path.dirname(os.path.dirname(os.path.realpath(__file__))))

from llm_perf.server import server_pb2, server_pb2_grpc
from llm_perf.server.pb import deserialize_value, serialize_value


@backoff.on_exception(backoff.expo, Exception, factor=0.1, max_value=1, max_tries=3)
def gen_stream_request(
    stub: server_pb2_grpc.InferenceStub,
    index: int,
    prompt: str,
    min_new_tokens: int,
    max_new_tokens: int,
    top_p: float,
    top_k: int,
    get_input_logits: int,
) -> Iterable[server_pb2.InferenceResponse]:
    req = server_pb2.InferenceRequest(
        req_id=str(index) + "_" + str(int(time.time())),
        inputs={
            "input_messages": serialize_value(prompt),
            "min_new_tokens": serialize_value(min_new_tokens),
            "max_new_tokens": serialize_value(max_new_tokens),
            "top_p": serialize_value(top_p),
            "top_k": serialize_value(top_k),
            "get_input_logits": serialize_value(get_input_logits),
        },
    )
    for res in stub.StreamingInference(req, wait_for_ready=False):
        yield res


def format_question(line) -> str:
    question = line["question"]
    for choice in ["A", "B", "C", "D"]:
        question += f"\n{choice}. {line[f'{choice}']}"
    question += "\n答案："
    return question


def bench_accuracy(stub, workload: Dict[str, Any], result_queue: mp.Queue):
    dataset: pd.DataFrame = pd.read_csv(workload["dataset"])
    result_queue.put("@start")
    for row_index, row in tqdm(dataset.iterrows(), total=len(dataset)):
        question = format_question(row)
        output_messages = ""
        perplexity_list: List[float] = []
        for res in gen_stream_request(
            stub,
            index=0,
            prompt=question,
            min_new_tokens=workload["min_new_tokens"],
            max_new_tokens=workload["max_new_tokens"],
            top_p=0,
            top_k=1,  # use greedy search for accuracy bench
            get_input_logits=1,
        ):
            res = {k: deserialize_value(v) for k, v in res.outputs.items()}
            if res["choice"]["message"] != "":
                output_messages = res["choice"]["message"]
            perplexity = res["choice"]["perplexity"]
            logits_dump = res["choice"]["logits_dump"]
            if not logits_dump:
                perplexity_list.append(perplexity)

        print("*"*150)
        logger.info("question: ")
        print(question)
        print("*"*150)
        logger.info("answer: ")
        print(output_messages)
        print("*"*150)

        result = {
            "output_message": output_messages,
            "perplexity": perplexity_list, 
            "logits_dump": logits_dump
        }
        result_queue.put(result)


def bench_performance(
    stub,
    index: int,
    workload: Dict[str, Any],
    input_tokens: int,
    result_queue: mp.Queue,
):
    result_queue.put("@start")

    accum_time = 0
    perf_time: int = workload["perf_time"] * int(1e9)

    while accum_time < perf_time:
        # make fake prompt, actual input_ids len may exceed input_tokens
        prompt = "我" * input_tokens
    
        st = time.perf_counter_ns()
        first_token_latency = 0

        min_new_tokens = workload["output_tokens"]
        max_new_tokens = workload["output_tokens"]

        output_messages: str = ""
        wait_time = []
        model_time = []
        post_process_time = []

        for res in gen_stream_request(
            stub,
            index=index,
            prompt=prompt,
            min_new_tokens=min_new_tokens,
            max_new_tokens=max_new_tokens,
            top_p=0,
            top_k=1,
            get_input_logits=0,
        ):
            res = {k: deserialize_value(v) for k, v in res.outputs.items()}
            if res["choice"]["message"] != "":
                output_messages = res["choice"]["message"]
            wait_time.append(res["choice"]["wait_time"])
            model_time.append(res["choice"]["model_time"])
            post_process_time.append(res["choice"]["post_process_time"])
            if first_token_latency == 0:
                first_token_latency = (time.perf_counter_ns() - st) / 1e6

        use_time = time.perf_counter_ns() - st
        accum_time += use_time

        # record context and decode len
        prompt_tokens = res["usage"]["prompt_tokens"]
        completion_tokens = res["usage"]["completion_tokens"]

        # seperate context and decode latency
        if completion_tokens > 1:
            per_token_latency = (use_time - first_token_latency) / (completion_tokens - 1) / 1e6
        else:
            per_token_latency = first_token_latency / 1e6

        context_wait_time = wait_time[0]
        context_model_time = model_time[0]
        context_postprocess_time = post_process_time[0]

        decode_wait_time = sum(wait_time[1:]) / len(wait_time[1:])
        decode_model_time = sum(model_time[1:]) / len(model_time[1:])
        decode_postprocess_time = sum(post_process_time[1:]) / len(post_process_time[1:])

        result = {
            "prompt_tokens": prompt_tokens,
            "completion_tokens": completion_tokens,
            "output_message": output_messages,
            "first_token_latency": first_token_latency,
            "per_token_latency": per_token_latency,

            "context_wait_time": context_wait_time, 
            "context_model_time": context_model_time, 
            "context_postprocess_time": context_postprocess_time, 

            "decode_wait_time": decode_wait_time, 
            "decode_model_time": decode_model_time, 
            "decode_postprocess_time": decode_postprocess_time, 
        }

        logger.debug(f"bench_{index} prompt response: {result}")
        result_queue.put(result)

def benchmark(
    index: int,
    start_wait: int, 
    workload: Dict[str, Any],
    report_type: ReportType,
    input_tokens: int,
    result_queue: mp.Queue,
    host, port
):
    with grpc.insecure_channel(f"{host}:{port}") as channel:
        stub = server_pb2_grpc.InferenceStub(channel)
        logger.debug(f"{report_type.name} bench_{index} start")
        
        # wait for start_wait seconds
        time.sleep(1 * start_wait)
        
        try:
            if report_type == ReportType.ACCURACY:
                bench_accuracy(stub, workload, result_queue)
            elif report_type == ReportType.PERFORMANCE:
                bench_performance(stub, index, workload, input_tokens, result_queue)
        except Exception as e:
            logger.error(f"{report_type.name} bench_{index} error: {e}")
            raise e

    logger.debug(f"{report_type.name} bench_{index} finish")
    result_queue.put(None)