code: def with_mpi(nproc=2, timeout=30, skip_if_no_mpi=True):
      def outer_thunk(fn):
          @wraps(fn)
          def thunk(*args, **kwargs):
              serialized_fn = base64.b64encode(
                  cloudpickle.dumps(lambda: fn(*args, **kwargs))
              )
              subprocess.check_call(
                  [
                      "mpiexec",
                      "-n",
                      str(nproc),
                      sys.executable,
                      "-m",
                      "baselines.common.tests.test_with_mpi",
                      serialized_fn,
                  ],
                  env=os.environ,
                  timeout=timeout,
              )
  
          if skip_if_no_mpi:
              return pytest.mark.skipif(
                  MPI is None, reason="MPI not present"
              )(thunk)
          else:
              return thunk
  
      return outer_thunk
  
file_path: baselines/baselines/common/tests/test_with_mpi.py
imports: import os
  import sys
  import subprocess
  import cloudpickle
  import base64
  import pytest
  from functools import wraps
label_desc: Use image node_id 1 for calling a global function with example usage: with_mpi(nproc, timeout, skip_if_no_mpi) and returns: outer_thunk, pytest, thunk