Add PMIx Horovod test

flx42 · flx42 · commit 863ccf39c8b1 · 2023-02-08T13:58:35.000-08:00
Signed-off-by: Felix Abecassis &lt;fabecassis@nvidia.com&gt;
diff --git a/tests/mpi.bats b/tests/mpi.bats
@@ -24,6 +24,11 @@ function teardown() {
     [ "$(sort -n <<< ${output})" == "$(seq 0 3)" ]
 }
 
+@test "PMIx Horovod" {
+    run_srun --mpi=pmix --ntasks=8 --container-image=nvcr.io/nvidia/tensorflow:23.01-tf2-py3 \
+             python -c "import os ; import horovod.tensorflow as hvd ; hvd.init(); assert(int(os.getenv('PMIX_RANK')) == hvd.rank())"
+}
+
 @test "PMI2 file descriptor" {
     run_srun --mpi=pmi2 --ntasks=4 --container-image=ubuntu:18.04 bash -c '[ -n "${PMI_FD}" ] && realpath /proc/self/fd/${PMI_FD} | grep -q socket'
 }

Original file line number	Diff line number	Diff line change
`@@ -24,6 +24,11 @@ function teardown() {`
`24`	`24`	`[ "$(sort -n <<< ${output})" == "$(seq 0 3)" ]`
`25`	`25`	`}`
`26`	`26`
	`27`	`+@test "PMIx Horovod" {`
	`28`	`+ run_srun --mpi=pmix --ntasks=8 --container-image=nvcr.io/nvidia/tensorflow:23.01-tf2-py3 \`
	`29`	`+ python -c "import os ; import horovod.tensorflow as hvd ; hvd.init(); assert(int(os.getenv('PMIX_RANK')) == hvd.rank())"`
	`30`	`+}`
	`31`	`+`
`27`	`32`	`@test "PMI2 file descriptor" {`
`28`	`33`	`run_srun --mpi=pmi2 --ntasks=4 --container-image=ubuntu:18.04 bash -c '[ -n "${PMI_FD}" ] && realpath /proc/self/fd/${PMI_FD} \| grep -q socket'`
`29`	`34`	`}`