main.py

#!/usr/bin/env python3

import argparse
import sys

from .runner import run_with_input
from .utils import get_mpi_rank


def create_parser():
    parser = argparse.ArgumentParser(
        description="RCCL Log Parser Wrapper\n\n"
        "Usage modes:\n"
        "  1. Pipe input:    cat log.txt | xcl-lens\n"
        "  2. Read files:    xcl-lens log1.txt log2.txt\n"
        "  3. Wrap command:  xcl-lens ./all_reduce_perf",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument("--summary", action="store_true", help="Print summary report only")
    parser.add_argument("-v", "--verbose", action="store_true", help="Print verbose reports")
    parser.add_argument(
        "command", nargs=argparse.REMAINDER, help="Executable to run, or log files to read"
    )
    return parser


def main():
    rank = get_mpi_rank()

    parser = create_parser()
    args = parser.parse_args()

    try:
        exit_code = run_with_input(args, rank)
        if exit_code is not None:
            sys.exit(exit_code)
    except KeyboardInterrupt:
        sys.exit(130)

    # If we got here, no command was provided and stdin is a tty
    if rank == 0:
        parser.print_help()
    sys.exit(1)


if __name__ == "__main__":
    main()