def _clear_mpi_env_vars()

in gym3/subproc.py [0:0]

• 12 lines of code
• 4 McCabe index (conditional complexity)

def _clear_mpi_env_vars():
    """
    from mpi4py import MPI will call MPI_Init by default.  If we spawn a child process that also calls MPI_Init and
    has MPI environment variables defined, MPI will think that the child process is an MPI process just like the
    parent and do bad things such as hang or crash.

    This context manager is a hacky way to clear those environment variables temporarily such as when we are starting
    multiprocessing Processes.
    """
    with _clear_lock:
        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)