code: def clear_mpi_env_vars():
      """Clears the MPI of environment variables.
  
      `from mpi4py import MPI` will call `MPI_Init` by default.
      If the child process has MPI environment variables, MPI will think that the child process
      is an MPI process just like the parent and do bad things such as hang.
  
      This context manager is a hacky way to clear those environment variables
      temporarily such as when we are starting multiprocessing Processes.
  
      Yields:
          Yields for the context manager
      """
      removed_environment = {}
      for k, v in list(os.environ.items()):
          for prefix in ["OMPI_", "PMI_"]:
              if k.startswith(prefix):
                  removed_environment[k] = v
                  del os.environ[k]
      try:
          yield
      finally:
          os.environ.update(removed_environment)
  
imports: import contextlib
  import os
label_desc: Use image node_id 1 for calling a global function with example usage: clear_mpi_env_vars() without return types