How to use typer in mpi environment properly?

[xiechao06]

First Check

• I added a very descriptive title here.
• I used the GitHub search to find a similar question and didn't find it.
• I searched the Typer documentation, with the integrated search.
• I already searched in Google "How to X in Typer" and didn't find any information.
• I already read and followed all the tutorial in the docs and didn't find an answer.
• I already checked if it is not related to Typer but to Click.

Commit to Help

• I commit to help with one of those options 👆

Example Code

import typer
from mpi4py import MPI

TEACHER_RANK = 0

world = MPI.COMM_WORLD
rank = world.rank


def main():
    if rank == TEACHER_RANK:
        # Only the teacher should handle input
        typer.run(teacher)
    else:
        student()

    MPI.Finalize()


def teacher(x: int, y: int):
    """teacher receive x and y, asks the students to calculate "x + y"

    Args:
        x (int): _description_
        y (int): _description_
    """
    # broadcast the question!
    world.bcast((x, y))
    size = world.size

    remaining = size - 1
    while remaining > 0:
        status = MPI.Status()
        answer = world.recv(None, source=MPI.ANY_SOURCE, status=status)
        judgement = "correct" if answer == x + y else "wrong"
        print(f"teacher: you are {judgement}! student {status.source}", flush=True)
        remaining -= 1


def student():
    # receive the broadcast
    x, y = world.bcast(None)
    answer = x + y
    print(f"student {rank}: answer is {answer}!", flush=True)
    world.send(answer, dest=TEACHER_RANK)


if __name__ == "__main__":
    main()

Description

• Message Passing Interface (MPI) is a standard for parallelism computation.
• In this snippet, a teacher receives 2 integers x and y from command line arguments and asks each student to calculate "x + y" in parallelism.
• This snippet is executed as:

$ mpiexec -n 4 python question.py

here we have 1 teacher and 3 students. (since 4 = 1+3)

• At the beginning mpiexec creates 4 processes and I ask only the teacher should handle cli arguments, otherwise when I execute:

$ mpiexec -n 4 python question.py --help

the help screen will be printed 4 times.

• The problem is if I execute the script as:

$ mpiexec -n 4 python question.py --help

Since "typer" will handle the "--help" option in the teacher process and exit the teacher process, the students just reside there and wait for the question from the teacher.

Similarly, if I execute the script with invalid arguments like

$ mpiexec  -n 4 python question.py 1 2 3

The teacher exits abnormally and the students just hang there.

• Suggestion: "typer" could provide a mechanism to intercept the "--help" option and invalid arguments handling, thus I could abort the students.

Operating System

macOS

Operating System Details

No response

Typer Version

0.9.0

Python Version

3.11.4

Additional Context

No response