#!/usr/bin/env python3
"""Verify every tensor of a safetensors model is finite (no NaN/Inf) and
sane. Usage: python3 weight_check.py /models/<MODEL_DIR>

Run inside the vllm container (it has torch + safetensors):
    docker cp diag/weight_check.py vllm:/tmp/ && \
    docker exec vllm python3 /tmp/weight_check.py /models/Qwen3.6-35B-A3B
"""
import glob
import os
import sys

import torch
from safetensors import safe_open


def main():
    model_dir = sys.argv[1] if len(sys.argv) > 1 else "/models"
    files = sorted(glob.glob(os.path.join(model_dir, "*.safetensors")))
    if not files:
        print(f"no safetensors found in {model_dir}")
        return 2
    total = bad = 0
    for f in files:
        with safe_open(f, framework="pt", device="cpu") as st:
            for k in st.keys():
                t = st.get_tensor(k)
                total += 1
                if torch.is_floating_point(t):
                    if not torch.isfinite(t).all():
                        print(f"BAD {os.path.basename(f)}::{k} "
                              f"nan={int(torch.isnan(t).sum())} inf={int(torch.isinf(t).sum())}",
                              flush=True)
                        bad += 1
                elif t.dtype != torch.bool and (t.abs() > 1e6).any():
                    print(f"ODD {os.path.basename(f)}::{k} max={t.abs().max().item()}",
                          flush=True)
    print(f"DONE shards={len(files)} tensors={total} bad={bad}")
    return 1 if bad else 0


if __name__ == "__main__":
    sys.exit(main())