"""Compare graphs generated by vLLM-Spyre vs AFTU

Run `python -m pytest tests/aftu/test_compare_graphs.py`.
"""

import os
import sys
import tempfile

import pytest
from graph_compare_utils import (
    collect_graph_files,
    compare_graphs,
    get_model_path,
    run_inference_py_and_get_graphs,
)
from spyre_util import ModelInfo, environ_checkpoint, patch_environment
from vllm import LLM


@pytest.mark.spyre
@pytest.mark.chunked_prefill
def test_compare_graphs_chunked_prefill(
    model: ModelInfo,
    max_num_seqs: int,
    max_model_len: int,
    monkeypatch: pytest.MonkeyPatch,
):
    """Test if the graphs generated by vllm-spyre matches with the ones by
    AFTU with chunked prefill enabled"""

    model_path = get_model_path(model)

    chunk_size = 128

    # fmt: off
    inference_py_args = [
        sys.executable, "-m", "aiu_fms_testing_utils.scripts.inference",
        "--architecture", "hf_pretrained",
        "--model_path", model_path,
        "--tokenizer", model_path,
        "--unfuse_weights",
        "--device_type", "aiu",
        "--compile",
        "--cast_bf16_to_fp16",
        "--compile_dynamic",
        "--min_pad_length", "64",
        "--max_new_tokens", "5",
        "--batch_size", str(max_num_seqs),
        "--compile_dynamic_sendnn",
        "--prefill_chunk_size", str(chunk_size),
    ]
    # fmt: on

    if not model.is_quantized:
        inference_py_args += ["--default_dtype", "fp16"]
        inference_py_args += ["--attention_type", "paged"]
    else:
        inference_py_args += ["--attention_type", "paged_fp8"]

    extra_env = {
        "VLLM_DT_MAX_CONTEXT_LEN": str(max_model_len),
        "VLLM_DT_MAX_BATCH_SIZE": str(max_num_seqs),
        "VLLM_DT_MAX_BATCH_TKV_LIMIT": str(1024 * 128),
        "VLLM_DT_CHUNK_LEN": str(chunk_size),
    }
    aftu_graphs = run_inference_py_and_get_graphs(inference_py_args, extra_env)

    assert len(aftu_graphs) > 0

    # VLLM

    monkeypatch.setenv("DEE_DUMP_GRAPHS", "vllm_cb_cp")
    # Disable cache to produce the graphs
    monkeypatch.setenv("TORCH_SENDNN_CACHE_ENABLE", "0")

    # NB: On pytorch 2.10 currently the eager and sendnn_compile_only backends do not work for
    # quantized models.
    backend = "sendnn" if model.is_quantized else "sendnn_compile_only"
    patch_environment(
        backend=backend,
        monkeypatch=monkeypatch,
    )

    # NB: We need AFTU and vllm-spyre to have the same number of paged
    # attention blocks for the graph comparison.
    #
    # vllm-spyre adds 1 block to the full batch count for the padding block and
    # rounds up to a multiple of the batch size (see get_total_spyre_blocks in
    # spyre_model_runner.py).
    #
    # Adjusting the context length for vLLM by the block size counters the
    # padding block adjustment for alignment with AFTU
    adjusted_max_model_len = max_model_len - 64
    original_cwd = os.getcwd()
    try:
        # Change to temp dir to set the test environment clean
        with tempfile.TemporaryDirectory() as tmpdir:
            os.chdir(tmpdir)

            # We only need to load the model
            with environ_checkpoint():
                LLM(
                    model=model.name,
                    revision=model.revision,
                    max_model_len=adjusted_max_model_len,
                    tensor_parallel_size=1,
                    max_num_batched_tokens=chunk_size,
                    max_num_seqs=max_num_seqs,
                )

            vllm_graphs = collect_graph_files(tmpdir)
    finally:
        # Restore in case of exception
        os.chdir(original_cwd)

    assert compare_graphs(vllm_graphs, aftu_graphs)
