#!/usr/bin/env python3
from __future__ import annotations

import argparse

from torch.distributed.checkpoint.metadata import Metadata
from torch.distributed.checkpoint.filesystem import FileSystemReader


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("checkpoint_dir")
    parser.add_argument("--limit", type=int, default=240)
    args = parser.parse_args()

    reader = FileSystemReader(args.checkpoint_dir)
    metadata: Metadata = reader.read_metadata()
    count = 0
    skip_prefixes = ("optimizer.", "opt_param_scheduler", "rng_state", "iteration")
    for key in sorted(metadata.state_dict_metadata):
        if key.startswith(skip_prefixes):
            continue
        meta = metadata.state_dict_metadata[key]
        size = getattr(meta, "size", None)
        print(f"{key}\t{tuple(size) if size is not None else '?'}")
        count += 1
        if args.limit and count >= args.limit:
            break
    print(f"printed={count}")


if __name__ == "__main__":
    main()