import torch
from vllm import LLM, SamplingParams
from vllm.utils import GiB_bytes
import nvtx

# Reference:
# https://docs.vllm.ai/en/latest/features/sleep_mode.html#sleep-mode
# https://github.com/vllm-project/vllm/blob/main/tests/basic_correctness/test_cumem.py

def main():
    torch.cuda.cudart().cudaProfilerStart()

    model = "Qwen/Qwen2.5-Coder-7B-Instruct"

    llm = LLM(model, enable_sleep_mode=True, tensor_parallel_size=2)

    free, total = torch.cuda.mem_get_info()
    used_bytes_baseline = total - free
    prompt = "How are you?"
    sampling_params = SamplingParams(temperature=0, max_tokens=10)

    with nvtx.annotate("generate", color="green"):
        output = llm.generate(prompt, sampling_params)

    with nvtx.annotate("sleep", color="red"):
        llm.sleep(level=2) # or level=1

    free_gpu_bytes_after_sleep, total = torch.cuda.mem_get_info()
    used_bytes = total - free_gpu_bytes_after_sleep - used_bytes_baseline
    assert used_bytes < 3 * GiB_bytes

    with nvtx.annotate("wake_up(weights)", color="brown"):
        llm.wake_up(tags=["weights"])

    with nvtx.annotate("reload_weights", color="purple"):
        llm.collective_rpc("reload_weights")
    free_gpu_bytes_wake_up_w, total = torch.cuda.mem_get_info()
    used_bytes = total - free_gpu_bytes_wake_up_w - used_bytes_baseline
    assert used_bytes < 4 * GiB_bytes

    with nvtx.annotate("wake_up(kv_cache)", color="brown"):
        llm.wake_up(tags=["kv_cache"])

    with nvtx.annotate("generate", color="green"):
        output2 = llm.generate(prompt, sampling_params)

    assert output[0].outputs[0].text == output2[0].outputs[0].text

    torch.cuda.cudart().cudaProfilerStop()


if __name__ == "__main__":
    main()