import os
import time

from vllm import LLM, SamplingParams

model = "/data/vllm-ascend-smoke/models/Qwen/Qwen2.5-0.5B-Instruct"
print("SMOKE_MODEL", model, flush=True)
print(
    "SMOKE_ENV",
    {k: os.environ.get(k) for k in ["VLLM_TARGET_DEVICE", "ASCEND_RT_VISIBLE_DEVICES", "PYTHONPATH"]},
    flush=True,
)
start = time.time()
llm = LLM(
    model=model,
    trust_remote_code=True,
    dtype="float16",
    max_model_len=256,
    max_num_seqs=1,
    max_num_batched_tokens=256,
    tensor_parallel_size=1,
    enforce_eager=True,
    gpu_memory_utilization=0.6,
    offload_group_size=4,
    offload_num_in_group=1,
)
print("SMOKE_LOAD_SECONDS", round(time.time() - start, 2), flush=True)
outputs = llm.generate(
    ["Hello, briefly introduce vLLM Ascend."],
    SamplingParams(temperature=0.0, max_tokens=16),
)
for out in outputs:
    print("SMOKE_PROMPT", out.prompt, flush=True)
    print("SMOKE_OUTPUT", out.outputs[0].text, flush=True)
print("SMOKE_OK_OFFLOAD", flush=True)
