# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD 3-Clause license found in the
# LICENSE file in the root directory of this source tree.
#!/usr/bin/env python3
"""
Script which can be used to parse the log file generated by the torchtitan,
and calculate the training performance metrics (mdian tokens/second and peak memory usage).

Usage:
    python parse_torchtitan_logs.py --log-file <log_file_path>
"""

import os
import re
import statistics
from argparse import ArgumentParser, Namespace


def main(args: Namespace):
    print("\n=====================================================")
    print(" Calculating training performance metrics")
    print("=====================================================")

    log_pattern = re.compile(r"step: (\d+).*?memory: ([\d.]+)GiB.*?tps: ([\d,]+)")

    assert os.path.exists(args.log_file), f"{args.log_file} does not exist"

    with open(args.log_file, "r") as f:
        log_data = f.read()

    matches = re.findall(log_pattern, log_data)

    tokens_per_second = []
    max_memory_usage = 0.0
    for match in matches:
        step = int(match[0])
        memory_usage = float(match[1])
        tps = float(match[2].replace(",", ""))

        # update peak memory usage
        max_memory_usage = max(max_memory_usage, memory_usage)

        # collect tokens per second, excluding step 1 which has initialization overhead
        if step != 1:
            tokens_per_second.append(tps)

    # calculate median tokens per second
    median_tps = statistics.median(tokens_per_second) if tokens_per_second else 0

    print(f"Median Tokens/Second (excluding step 1): {median_tps}")
    print(f"Max Memory Usage: {max_memory_usage} GiB")


if __name__ == "__main__":
    argparser = ArgumentParser()
    argparser.add_argument(
        "--log-file", type=str, required=True, help="torchtitan log file"
    )
    args = argparser.parse_args()
    main(args)