vllm_llama_test

import asyncio
import random
import aiohttp
import time
import matplotlib.pyplot as plt
import json
import os

SERVER_HOST = "127.0.0.0"
SERVER_PORT = "8080"
SERVER_URL = f"http://{SERVER_HOST}:{SERVER_PORT}/v1/completions"

# Имя модели, указанное в параметре --served-model-name при запуске vLLM / llama.cpp
MODEL_NAME = "*"

# Базовые строки для построения промптов
BASE_DOG = "The quick brown fox jumps over the lazy dog. "
BASE_CAT = "The quick brown fox jumps over the lazy cat. "

async def send_vllm_request(session, prompt, slot_id):
    """
    Отправляет потоковый запрос в vLLM / llama.cpp с поддержкой сохранения кеша.
    """
    payload = {
        "model": MODEL_NAME,
        "prompt": prompt,
        "max_tokens": 64,
        "temperature": 1.0,
        "stream": True,
        #"id_slot": slot_id,          # Иногда с этим ллама хуже работает
        "cache_prompt": True,
        "stream_options": {"include_usage": True}
    }

    start_time = time.perf_counter()
    ttft = None
    generated_text = ""
    prompt_tokens = 0
    completion_tokens = 0

    async with session.post(SERVER_URL, json=payload, timeout=900) as resp:
        if resp.status != 200:
            err_text = await resp.text()
            raise RuntimeError(f"HTTP {resp.status}: {err_text}")

        async for line in resp.content:
            line = line.decode('utf-8').strip()
            if not line or line == "data: [DONE]":
                continue

            if line.startswith("data: "):
                data_str = line[6:]
                try:
                    data = json.loads(data_str)
                except json.JSONDecodeError:
                    continue

                # Фиксируем TTFT по первому пришедшему чанку с текстом
                if ttft is None and "choices" in data and len(data["choices"]) > 0:
                    ttft = time.perf_counter() - start_time

                if data.get("choices") and data["choices"][0].get("text"):
                    generated_text += data["choices"][0]["text"]

                if data.get("usage"):
                    prompt_tokens = data["usage"].get("prompt_tokens", 0)
                    completion_tokens = data["usage"].get("completion_tokens", 0)

    total_time = time.perf_counter() - start_time

    if completion_tokens == 0 and generated_text:
        completion_tokens = max(1, len(generated_text) // 4)
    elif completion_tokens == 0:
        completion_tokens = 1

    if ttft is not None:
        tg_time = max(0.001, total_time - ttft)
    else:
        ttft = total_time
        tg_time = 0.001

    return ttft, tg_time, prompt_tokens, completion_tokens, generated_text

async def run_experiment(num_streams, max_context, chunk_text):
    print(f"\n=== СТАРТ ТЕСТА: {num_streams} ПОТОК(ОВ) | ЛИМИТ {max_context} ТОКЕНОВ ===")

    # Инициализируем структуру. Промпт растет строго линейно.
    streams_state = [{"prompt": chunk_text, "ctx_size": 0} for _ in range(num_streams)]

    history = {
        "context_sizes": [],
        "pp_speeds_wall": [],
        "tg_speeds_wall": []
    }

    step = 1
    timeout = aiohttp.ClientTimeout(total=2400)

    async with aiohttp.ClientSession(timeout=timeout) as session:
        while True:
            if step > 1 and streams_state[0]["ctx_size"] >= max_context:
                break

            try:
                # Передаем индекс `i` в качестве slot_id. Он стабилен между итерациями (0, 1, 2...)
                responses = await asyncio.gather(*(
                    send_vllm_request(session, s["prompt"], i) for i, s in enumerate(streams_state)
                ))
            except Exception as e:
                print(f"[Ошибка запроса]: {e}")
                break

            pp_ms_list = []
            tg_ms_list = []
            total_pp_tokens = 0
            total_tg_tokens = 0
            ctx_sizes = []

            for idx, res in enumerate(responses):
                ttft, tg_time, p_tokens, c_tokens, content = res

                pp_ms_list.append(ttft * 1000)
                tg_ms_list.append(tg_time * 1000)

                # Вычисляем дельту новых токенов
                new_tokens = p_tokens - streams_state[idx]["ctx_size"]
                if new_tokens <= 0:
                    new_tokens = p_tokens

                total_pp_tokens += new_tokens
                total_tg_tokens += c_tokens

                streams_state[idx]["ctx_size"] = p_tokens

                # Важно: склеиваем без лишних пробелов в стыках, чтобы не ломать токенизацию границ строк в llama.cpp
                streams_state[idx]["prompt"] += content + chunk_text
                ctx_sizes.append(p_tokens)

            avg_pp_ms = sum(pp_ms_list) / len(pp_ms_list) if pp_ms_list else 0.001
            avg_tg_ms = sum(tg_ms_list) / len(tg_ms_list) if tg_ms_list else 0.001

            pp_speed_wall = total_pp_tokens / (avg_pp_ms / 1000)
            tg_speed_wall = total_tg_tokens / (avg_tg_ms / 1000)

            cache_str = f"{min(ctx_sizes)}..{max(ctx_sizes)}"
            pp_ms_str = f"[{', '.join(f'{x:.1f}' for x in pp_ms_list)}]"
            tg_ms_str = f"[{', '.join(f'{x:.1f}' for x in tg_ms_list)}]"

            print(f"Шаг {step:<2}, кэш {cache_str:<13} | "
                  f"pp(TTFT) - {pp_ms_str:<25}ms [≈{pp_speed_wall:<6.1f} t/s] | "
                  f"tg(TPOT) - {tg_ms_str:<25}ms [≈{tg_speed_wall:<6.1f} t/s]")

            history["context_sizes"].append(max(ctx_sizes))
            history["pp_speeds_wall"].append(pp_speed_wall)
            history["tg_speeds_wall"].append(tg_speed_wall)

            step += 1

    return history

def save_and_plot(all_experiments_data, filename_suffix):
    dir_name = os.path.dirname(filename_suffix)
    if dir_name:
        os.makedirs(dir_name, exist_ok=True)

    json_path = f"{filename_suffix}.json"
    img_path = f"{filename_suffix}.png"

    with open(json_path, "w") as f:
        json.dump(all_experiments_data, f, indent=4)

    colors = {1:'blue', 2:'green', 4:'orange', 8:'red', 16:'purple'}
    fig, axs = plt.subplots(2, 2, figsize=(20, 14))

    ax_left_pp, ax_right_pp = axs[0, 0], axs[0, 1]
    ax_left_tg, ax_right_tg = axs[1, 0], axs[1, 1]

    for num_streams, history in all_experiments_data.items():
        num_streams = int(num_streams)
        color = colors.get(num_streams, 'black')

        ctx_single = history["context_sizes"]
        ctx_total = [c * num_streams for c in ctx_single]

        pp_total = history["pp_speeds_wall"]
        pp_ind = [p / num_streams for p in pp_total]

        tg_total = history["tg_speeds_wall"]
        tg_ind = [t / num_streams for t in tg_total]

        ax_left_pp.plot(ctx_single, pp_total, linestyle='-', marker='x', color=color, label=f'{num_streams} Поток(ов) - Полный PP')
        ax_left_pp.plot(ctx_single, pp_ind, linestyle=':', marker='x', color=color, alpha=0.8, label=f'{num_streams} Поток(ов) - Индив. PP')

        ax_left_tg.plot(ctx_single, tg_total, linestyle='-', marker='x', color=color, label=f'{num_streams} Поток(ов) - Полный TG')
        ax_left_tg.plot(ctx_single, tg_ind, linestyle=':', marker='x', color=color, alpha=0.8, label=f'{num_streams} Поток(ов) - Индив. TG')

        ax_right_pp.plot(ctx_total, pp_total, linestyle='-', marker='x', color=color, label=f'{num_streams} Поток(ов) - Полный PP')
        ax_right_pp.plot(ctx_total, pp_ind, linestyle=':', marker='x', color=color, alpha=0.8, label=f'{num_streams} Поток(ов) - Индив. PP')

        ax_right_tg.plot(ctx_total, tg_total, linestyle='-', marker='x', color=color, label=f'{num_streams} Поток(ов) - Полный TG')
        ax_right_tg.plot(ctx_total, tg_ind, linestyle=':', marker='x', color=color, alpha=0.8, label=f'{num_streams} Поток(ов) - Индив. TG')

    for ax in [ax_left_pp, ax_left_tg, ax_right_pp, ax_right_tg]:
        ax.grid(True, linestyle=':', alpha=0.6)
        ax.legend(loc='upper right', fontsize=8)

    ax_left_pp.set_title('PP Скорость (TTFT) | Ось X: Кэш потока', fontweight='bold')
    ax_left_tg.set_title('TG Скорость (TPOT) | Ось X: Кэш потока', fontweight='bold')
    ax_right_pp.set_title('PP Скорость (TTFT) | Ось X: Сумма кэшей', fontweight='bold')
    ax_right_tg.set_title('TG Скорость (TPOT) | Ось X: Сумма кэшей', fontweight='bold')

    fig.suptitle(f"Комплексный профиль производительности\nБаза: {os.path.basename(filename_suffix)}", fontsize=14, fontweight='bold', y=0.98)
    fig.tight_layout(rect=[0, 0, 1, 0.95])
    fig.savefig(img_path, dpi=150)

    print(f"\n[Успех] Данные сохранены в: {json_path}")
    print(f"[Успех] График сохранен в: {img_path}")

async def main():
    experiments_config = [
        {"streams": 1, "max_context": 120000},
        {"streams": 2, "max_context": 80000},
        {"streams": 4, "max_context": 50000},
        {"streams": 8, "max_context": 25000},
        {"streams": 16, "max_context": 16000}
    ]

    all_results = {}

    for idx, config in enumerate(experiments_config):
        streams = config["streams"]
        max_ctx = config["max_context"]

        reps = 105 if streams == 1 else (51 if streams == 2 else 26)
        reps *= 4

        base_text = BASE_DOG if idx % 2 == 0 else BASE_CAT
        chunk_text = str(random.randint(0, 9)) + base_text * reps

        await asyncio.sleep(2)

        history = await run_experiment(streams, max_ctx, chunk_text)
        all_results[streams] = history

    save_and_plot(all_results, filename_suffix="result_custom_test/vllm_benchmark_results")

if __name__ == "__main__":
    asyncio.run(main())