from mpi4py import MPI

COMM = MPI.COMM_WORLD
RANK = COMM.Get_rank()
SIZE = COMM.Get_size()

import socket
hostname = socket.gethostname()

from time import sleep

sleep(10)  # 等待10秒，延长程序运行时间以便于squeue查看任务进度。

for i in range(SIZE):
    if RANK == i:
        print(f"I am rank#{RANK}, I am in node={hostname}")
    else:
        pass
    COMM.barrier()

for i in range(SIZE):
    for j in range(SIZE):
        if i != j:
            if RANK == i:
                COMM.send("Sending message ...", dest=j)
            elif RANK == j:
                msg = COMM.recv(source=i)
                print("mpi between", i, "and", j, "works.", flush=True)
            else:
                pass