Using DeePMD, zero num_intra_threads and num_inter_threads?

[nbidler]

Hello. I am trying to make sure that I am using all the resources I have access to when training the neural network.
Note: I am submitting jobs via SLURM to a computing cluster. I am trying to determine if there is a conflict with the scheduler, if I'm not setting environment variables, or if there's some other disconnect.

For example, this sets the environment variables and (as far as I know) results in these environment variables being present in whatever environment the network is being trained in.

#SBATCH --nodes=1
#SBATCH --ntasks=4

#SBATCH --ntasks-per-node=4
#SBATCH --cpus-per-task=2

OMP_NUM_THREADS=8
TF_INTRA_OP_PARALLELISM_THREADS=8
TF_INTER_OP_PARALLELISM_THREADS=4

#SBATCH --export=ALL

However, the output says that no intra/inter-threads exist.

DEEPMD INFO     _____               _____   __  __  _____           _     _  _
DEEPMD INFO    |  __ \             |  __ \ |  \/  ||  __ \         | |   (_)| |
DEEPMD INFO    | |  | |  ___   ___ | |__) || \  / || |  | | ______ | | __ _ | |_
DEEPMD INFO    | |  | | / _ \ / _ \|  ___/ | |\/| || |  | ||______|| |/ /| || __|
DEEPMD INFO    | |__| ||  __/|  __/| |     | |  | || |__| |        |   < | || |_
DEEPMD INFO    |_____/  \___| \___||_|     |_|  |_||_____/         |_|\_\|_| \__|
DEEPMD INFO    Please read and cite:
DEEPMD INFO    Wang, Zhang, Han and E, Comput.Phys.Comm. 228, 178-184 (2018)
DEEPMD INFO    installed to:         /project/_skbuild/linux-x86_64-3.10/cmake-install
DEEPMD INFO    source :              v2.1.1
DEEPMD INFO    source brach:         HEAD
DEEPMD INFO    source commit:        c4f0cec
DEEPMD INFO    source commit at:     2022-04-16 11:11:16 +0800
DEEPMD INFO    build float prec:     double
DEEPMD INFO    build with tf inc:    /opt/python/cp310-cp310/lib/python3.10/site-packages/tensorflow/include
DEEPMD INFO    build with tf lib:
DEEPMD INFO    ---Summary of the training---------------------------------------
DEEPMD INFO    distributed
DEEPMD INFO    world size:           4
DEEPMD INFO    my rank:              0
DEEPMD INFO    node list:            ['gpu-0-2']
DEEPMD INFO    running on:           gpu-0-2
DEEPMD INFO    computing device:     cpu:0
DEEPMD INFO    CUDA_VISIBLE_DEVICES: unset
DEEPMD INFO    Count of visible GPU: 0
DEEPMD INFO    num_intra_threads:    0
DEEPMD INFO    num_inter_threads:    0
DEEPMD INFO    -----------------------------------------------------------------

[wanghan-iapcm]

export OMP_NUM_THREADS=8
export TF_INTRA_OP_PARALLELISM_THREADS=8
export TF_INTER_OP_PARALLELISM_THREADS=4

[nbidler]

I've tried it with and without the "export" tag, the "#SBATCH --export=ALL" should export all local environment variables to the environment in which the code will execute.

[njzjz]

According to the sbatch documentation,

sbatch will stop processing further #SBATCH directives once the first non-comment non-whitespace line has been reached in the script.

So this line has never been processed.

[nbidler]

I have finally, FINALLY, found out how to package the environment variables so they are accepted. I am posting the solution here so others can see it, for the future

#SBATCH --mem-per-cpu=4040M
#SBATCH --time=20:00:00
#SBATCH --export=ALL

module purge

export OMP_NUM_THREADS=8
export TF_INTRA_OP_PARALLELISM_THREADS=8
export TF_INTER_OP_PARALLELISM_THREADS=4

It really was as simple as "put it below the last SBATCH command."