import mpi4py
mpi4py.rc.recv_mprobe = False

from mpi4py import MPI  # for libE communicator
import sys
import subprocess

task_nprocs = sys.argv[1]

size = MPI.COMM_WORLD.Get_size()
rank = MPI.COMM_WORLD.Get_rank()
runline = "mpirun -np ".split()
runline.append(task_nprocs)
runline.append('python')
runline.append('helloworld.py')

if rank == 0:
    print("Total sub-task procs: {}".format(size*int(task_nprocs)))
    print("Total procs (parent + sub-tasks): {}".format(size*(int(task_nprocs)+1)))

# print("Rank {}: {}".format(rank, " ".join(runline)))
output = 'task_' + str(rank) + '.out'
p = subprocess.Popen(runline, stdout=open(output, 'w'), shell=False)
p.wait()
