Assign MPI cpus-per-proc individually

I have a code that runs on 2 MPI process, each of them using multiple threads with OpenMP. My computer (Ubuntu) has 4 CPU (cores ?). Thus I made a small script to give 2 CPU per process such that each of them can run in parallel:

# Environment variables

export CPU_PER_PROC=2
export OMP_NUM_THREADS=${CPU_PER_PROC}
export OPTION="-map-by node:PE=${CPU_PER_PROC}"

mpiexec ${OPTION} -n 2 python3 main.py

But now I would like to give, for instance, 3 CPU to the process (rank) 0 and only 1 CPU to the rank 1 such that:

Process 0 : use 3 CPU with X (certainly 3 for better performances) threads.
Process 1 : use 1 CPU with Y (certainly 1) threads.

Solution

All MPI process inherit from the parent environment by default which contains the OMP_NUM_THREADS variable. To tweak the value regarding the rank, one solution is to change the environment variable in the child processes. In Python, you can do that using for example:

import os
os.environ['OMP_NUM_THREADS'] = str(3)

This line must be execute before the OpenMP initialization (that is before the first execution of a parallel directive). You can tweak the right-end-side regarding the MPI rank.

An alternative solution is to use the omp_set_num_threads function call but this is not so simple from a Python code. Another solution is to use a num_thread clause in a parallel section, but again, this is generally an issue from a Python code. A last alternative is to calla shell script that changes the environment before calling the Python script but this script can hardly get the MPI rank (AFAIK there are shell variable to get this but they are dependent of the MPI implementation and so not portable).