ci: Make OpenMPI yield when idle for the MPI example notebooks

The data-initialization notebook's point-to-point routing (NBX exchange)
hung the 4-engine ipyparallel cluster under OpenMPI, while IntelMPI ran
fine. OpenMPI defaults to an aggressive busy-wait: with the engines
contending for cores, the probe/barrier loop spins at 100% and starves the
peers it waits on, so MPI never progresses. IntelMPI yields by default.

Set OMPI_MCA_mpi_yield_when_idle=1 in the examples-mpi job so OpenMPI yields
the CPU while waiting. Also revert the earlier Python-level sleep(0) in
nbx_exchange: sched_yield at the wrong layer (the spin is in OpenMPI's C
progress engine), which did not help.

Author: mloubout

.github/workflows/examples-mpi.yaml

@@ -44,6 +44,11 @@ jobs:
       DEVITO_ARCH: "gcc"
       CC: "gcc"
       CXX: "g++"
+      # Make OpenMPI yield the CPU while waiting instead of busy-spinning. With
+      # 4 ipyparallel engines contending for cores, the point-to-point routing
+      # in the data notebooks otherwise livelocks under OpenMPI (IntelMPI yields
+      # by default, so it is unaffected and ignores this OMPI_* setting).
+      OMPI_MCA_mpi_yield_when_idle: "1"
 
     steps:
     - name: Checkout devito
@@ -71,7 +76,7 @@ jobs:
         ipcluster start --profile=mpi --engines=mpi -n 4 --daemonize
         # A few seconds to ensure workers are ready
         sleep 10
-        pytest -v --nbval examples/mpi
+        pytest -vvv --nbval examples/mpi
         ipcluster stop --profile=mpi
 
     - name: Test seismic examples

devito/data/distributed/transport.py

@@ -8,8 +8,6 @@
 graph communicator without affecting the layers above.
 """
 
-import time
-
 import numpy as np
 
 from devito.mpi import MPI
@@ -85,19 +83,12 @@ def nbx_exchange(comm, sendbufs, dtype, tag=0):
             buf = np.empty(count, dtype=dtype)
             comm.Recv([buf, mpitype], source=src, tag=tag)
             recvd[src] = buf
-            # Drain any further ready messages before yielding
-            continue
-        if barrier is None:
+        elif barrier is None:
             if MPI.Request.Testall(sends):
                 # All my sends were matched -> announce I am done sending
                 barrier = comm.Ibarrier()
         elif barrier.Test():
             # Everyone is done sending and nothing is in flight
             break
-        # Nothing was ready this pass: yield the CPU so co-scheduled ranks can
-        # make progress. Without this the probe loop busy-waits at 100%, which
-        # deadlocks the exchange when ranks are oversubscribed (more ranks than
-        # cores, e.g. the 4-engine ipyparallel cluster on a 2-core CI runner).
-        time.sleep(0)
 
     return recvd