scripts/benchmark/trainer-benchmark.py/main
def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--base-cmd",
        default=None,
        type=str,
        required=True,
        help="Base cmd",
    )
    parser.add_argument(
        "--variations",
        default=None,
        type=str,
        nargs="+",
        required=True,
        help="Multi-dimensional variations, example: '|--fp16|--bf16' '|--tf32'",
    )
    parser.add_argument(
        "--base-variation",
        default=None,
        type=str,
        help="Baseline variation to compare to. if None the minimal target value will be used to compare against",
    )
    parser.add_argument(
        "--target-metric-key",
        default=None,
        type=str,
        required=True,
        help="Target metric key in output_dir/all_results.json, e.g., train_samples_per_second",
    )
    parser.add_argument(
        "--report-metric-keys",
        default="",
        type=str,
        help="Report metric keys - other metric keys from output_dir/all_results.json to report, e.g., train_loss. Use a single argument e.g., 'train_loss train_samples",
    )
    parser.add_argument(
        "--repeat-times",
        default=1,
        type=int,
        help="How many times to re-run each variation - an average will be reported",
    )
    parser.add_argument(
        "--output_dir",
        default="output_benchmark",
        type=str,
        help="The output directory where all the benchmark reports will go to and additionally this directory will be used to override --output_dir in the script that is being benchmarked",
    )
    parser.add_argument(
        "--verbose",
        default=False,
        action="store_true",
        help="Whether to show the outputs of each run or just the benchmark progress",
    )
    args = parser.parse_args()

    output_dir = args.output_dir
    Path(output_dir).mkdir(exist_ok=True)
    base_cmd = get_base_command(args, output_dir)

    # split each dimension into its --foo variations
    dims = [list(map(str.strip, re.split(r"\|", x))) for x in args.variations]
    # build a cartesian product of dimensions and convert those back into cmd-line arg strings,
    # while stripping white space for inputs that were empty
    variations = list(map(str.strip, map(" ".join, itertools.product(*dims))))
    longest_variation_len = max(len(x) for x in variations)

    # split wanted keys
    report_metric_keys = args.report_metric_keys.split()

    # capture prints into a log file for convenience
    report_fn = f"benchmark-report-{datetime.datetime.now().strftime('%Y-%m-%d-%H-%M-%S')}.txt"
    print(f"\nNote: each run's output is also logged under {output_dir}/log.*.std*.txt")
    print(f"and this script's output is also piped into {report_fn}")

    sys.stdout = Tee(report_fn)

    print(f"\n*** Running {len(variations)} benchmarks:")
    print(f"Base command: {' '.join(base_cmd)}")

    variation_key = "variation"
    results = []
    for id, variation in enumerate(tqdm(variations, desc="Total completion: ", leave=False)):
        cmd = base_cmd + variation.split()
        results.append(
            process_run(
                id + 1,
                cmd,
                variation_key,
                variation,
                longest_variation_len,
                args.target_metric_key,
                report_metric_keys,
                args.repeat_times,
                output_dir,
                args.verbose,
            )
        )

    process_results(results, args.target_metric_key, report_metric_keys, args.base_variation, output_dir)