#!/usr/bin/env python3
"""Scan FP8-E4M3 weight-scale tensors in a safetensors checkpoint for NaN bytes.

ModelOpt 0.44 can occasionally emit literal E4M3 NaN encodings (0x7F / 0xFF) when the
float32->FP8 cast of a per-block weight scale rounds above the E4M3 max of 448.
A single NaN in any weight_scale collapses served output. Expect ZERO hits.

Usage: python scan_weight_scale_nan.py /path/to/checkpoint_dir_or_file [--all-fp8]
"""
import argparse
import glob
import json
import os
import struct
import sys


def iter_safetensors_headers(path):
    with open(path, "rb") as f:
        header_len = struct.unpack("<Q", f.read(8))[0]
        header = json.loads(f.read(header_len))
        data_start = 8 + header_len
        return header, data_start


def scan_file(path, all_fp8=False):
    header, data_start = iter_safetensors_headers(path)
    findings = []
    scanned = 0
    total_bytes = 0
    with open(path, "rb") as f:
        for name, meta in header.items():
            if name == "__metadata__":
                continue
            dtype = meta.get("dtype", "")
            if dtype not in ("F8_E4M3", "FP8_E4M3", "F8-E4M3"):
                continue
            if not all_fp8 and "weight_scale" not in name:
                continue
            begin, end = meta["data_offsets"]
            f.seek(data_start + begin)
            buf = f.read(end - begin)
            scanned += 1
            total_bytes += len(buf)
            n_nan = buf.count(0x7F) + buf.count(0xFF)
            if n_nan:
                findings.append((name, n_nan, len(buf)))
    return scanned, total_bytes, findings


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("target")
    ap.add_argument("--all-fp8", action="store_true",
                    help="scan every FP8 tensor, not just *weight_scale*")
    args = ap.parse_args()

    if os.path.isdir(args.target):
        files = sorted(glob.glob(os.path.join(args.target, "*.safetensors")))
    else:
        files = [args.target]
    if not files:
        print(f"no safetensors found under {args.target}", file=sys.stderr)
        sys.exit(2)

    any_nan = False
    for path in files:
        scanned, total_bytes, findings = scan_file(path, all_fp8=args.all_fp8)
        print(f"{os.path.basename(path)}: {scanned} FP8-E4M3 "
              f"{'tensors' if args.all_fp8 else 'weight_scale tensors'}, "
              f"{total_bytes/1e6:.1f} MB scanned")
        for name, n_nan, size in findings:
            any_nan = True
            print(f"  NAN  {name}: {n_nan} NaN byte(s) in {size} bytes")
    if any_nan:
        print("\nRESULT: FAIL — NaN-encoded scale bytes found (requantize on modelopt 0.45)")
        sys.exit(1)
    print("\nRESULT: PASS — zero E4M3 NaN bytes (0x7F/0xFF)")


if __name__ == "__main__":
    main()