main.py

#!/usr/bin/env python3

import sys
from typing import Annotated

import typer

from . import __version__
from .runner import run_with_input
from .utils import get_mpi_rank

app = typer.Typer(
    help=(
        "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"
    ),
    add_completion=False,
)


def _version_callback(value: bool):
    if value:
        typer.echo(f"xcl-lens {__version__}")
        raise typer.Exit()


@app.command(
    context_settings={
        "allow_extra_args": True,
        "ignore_unknown_options": True,
        "help_option_names": ["-h", "--help"],
    }
)
def main(
    ctx: typer.Context,
    version: Annotated[
        bool,
        typer.Option(
            "--version",
            help="Show version and exit",
            is_flag=True,
            is_eager=True,
            callback=_version_callback,
        ),
    ] = False,
    summary: Annotated[
        bool,
        typer.Option("-s", "--summary", help="Print summary report only"),
    ] = False,
    verbose: Annotated[bool, typer.Option("-v", "--verbose", help="Print verbose reports")] = False,
    hosts: Annotated[str | None, typer.Option(help="Filter by hosts (comma-separated)")] = None,
    ranks: Annotated[str | None, typer.Option(help="Filter by ranks (comma-separated)")] = None,
):
    """RCCL Log Parser Wrapper."""
    rank = get_mpi_rank()
    command: list[str] = ctx.args

    try:
        hosts_list = [h.strip() for h in hosts.split(",")] if hosts else None
        ranks_list = [r.strip() for r in ranks.split(",")] if ranks else None
        exit_code = run_with_input(
            command=command,
            summary=summary,
            verbose=verbose,
            hosts=hosts_list,
            ranks=ranks_list,
            rank=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:
        raise typer.Exit(code=1)
    sys.exit(1)


if __name__ == "__main__":
    app()