#!/usr/bin/env python3
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.

import argparse
import subprocess
import sys
import os
import json
from collections import namedtuple

SCRIPT_DIR = os.path.realpath(os.path.dirname(__file__))

def parse_args():
  parser = argparse.ArgumentParser(description="Runs BERT performance tests.")
  parser.add_argument("--binary_dir", required=True,
                      help="Path to the ORT binary directory.")
  parser.add_argument("--training_data_root", required=True,
                      help="Path to the training data root directory.")
  parser.add_argument("--model_root", required=True,
                      help="Path to the model root directory.")
  parser.add_argument("--gpu_sku", choices=['V100_16G', 'MI100_32G'], default='V100_16G', required=False,
                      help="GPU model (e.g. V100_16G, MI100_32G).")
  return parser.parse_args()

# using the same params from "GitHub Master Merge Schedule" in OneNotes
def main():
    args = parse_args()

    Config = namedtuple('Config', ['use_mixed_precision', 'max_seq_length', 'batch_size', 'max_predictions_per_seq', 'expected_perf'])
    configs = {}
    configs['V100_16G'] = [
        Config(True, 128, 76, 20, -1.0),
        Config(True, 512, 11, 80, -1.0),
        Config(False, 128, 39, 20, -1.0),
        Config(False, 512, 6, 80, -1.0)
    ]

    configs['MI100_32G'] = [
        Config(True, 128, 128, 20, 240),
    ]

    # run BERT training
    for c in configs[args.gpu_sku]:
        model = 'bert-large-uncased_L_24_H_1024_A_16_V_30528_S_512_Dp_0.1_optimized_layer_norm_opset12'
        precision_prefix = ('fp16' if c.use_mixed_precision else 'fp32')
        print("######## testing name - " + ('fp16-' if c.use_mixed_precision else 'fp32-') + str(c.max_seq_length) + " ##############")
        cmds = [
            os.path.join(args.binary_dir, "onnxruntime_training_bert"),
            "--model_name", os.path.join(
                args.model_root, "nv/bert-large/{}".format(model)),
            "--train_data_dir", os.path.join(
                args.training_data_root, str(c.max_seq_length), "books_wiki_en_corpus/train"),
            "--test_data_dir", os.path.join(
                args.training_data_root, str(c.max_seq_length), "books_wiki_en_corpus/test"),
            "--train_batch_size", str(c.batch_size),
            "--mode", "train",
            "--max_seq_length", str(c.max_seq_length),
            "--num_train_steps", "640",
            "--display_loss_steps", "5",
            "--optimizer", "Lamb",
            "--learning_rate", "3e-3",
            "--warmup_ratio", "0.2843",
            "--warmup_mode", "Poly",
            "--gradient_accumulation_steps", "1",
            "--max_predictions_per_seq", str(c.max_predictions_per_seq),
            "--lambda", "0",
            "--use_nccl",
            "--perf_output_dir", os.path.join(SCRIPT_DIR, "results"), 
        ]

        if c.use_mixed_precision: 
            cmds.append("--use_mixed_precision"),
            cmds.append("--allreduce_in_fp16"),

        subprocess.run(cmds).check_returncode()
        if c.expected_perf > 0.0:
            json_filename = 'onnxruntime_perf_metrics_{}.onnx_bert_{}_{}_Lamb.json'.format(model, precision_prefix, c.max_seq_length)
            with open(os.path.join(SCRIPT_DIR, 'results', json_filename)) as json_file:
                results = json.load(json_file)
                assert(results['EndToEndThroughput'] > 0.98*c.expected_perf)
    
    return 0

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