import json
import subprocess
import time
from vllm import LLM, SamplingParams

def npu_mem():
    try:
        out = subprocess.check_output(['npu-smi', 'info'], text=True, stderr=subprocess.STDOUT)
        return [line for line in out.splitlines() if '910B' in line or 'NPU' in line or 'Memory-Usage' in line][:20]
    except Exception as e:
        return [f'npu-smi failed: {e}']

print('MEM_BEFORE', json.dumps(npu_mem(), ensure_ascii=False), flush=True)
start = time.perf_counter()
llm = LLM(
    model='/data/vllm-ascend-smoke/models/Qwen/Qwen2.5-0.5B-Instruct',
    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('LOAD_SECONDS', round(time.perf_counter() - start, 4), flush=True)
params = SamplingParams(temperature=0.0, max_tokens=8)
latencies = []
for i in range(100):
    t = time.perf_counter()
    out = llm.generate([f'Number {i}: say hello in five words.'], params, use_tqdm=False)
    latencies.append(time.perf_counter() - t)
    if not out or not out[0].outputs or len(out[0].outputs[0].token_ids) == 0:
        raise RuntimeError(f'empty output at {i}')
print('STABILITY_RESULT', json.dumps({
    'requests': 100,
    'min_s': round(min(latencies), 4),
    'max_s': round(max(latencies), 4),
    'avg_s': round(sum(latencies)/len(latencies), 4),
}, sort_keys=True), flush=True)
print('MEM_AFTER', json.dumps(npu_mem(), ensure_ascii=False), flush=True)
