Source code for schrodinger.utils.openmpi

"""
Provide functions for integration with OpenMPI.
"""
import enum
import os
import pathlib
from typing import Dict
from typing import Optional

from schrodinger.job import util

from .env import prepend_path


@enum.unique
class OpenMPIVersion(enum.Enum):
    VERSION_1_6_5 = enum.auto()
    LATEST = enum.auto()


@enum.unique
class EnvVar(str, enum.Enum):
    LD_LIBRARY_PATH = "LD_LIBRARY_PATH"
    OPAL_PREFIX = "OPAL_PREFIX"
    PATH = "PATH"
    SCHRODINGER_OPENMPI = "SCHRODINGER_OPENMPI"


def get_openmpi_prefix(version: OpenMPIVersion,
                       *,
                       env: Optional[Dict[str, str]] = None) -> pathlib.Path:
    """
    Return openmpi prefix (OPAL_PREFIX).

    :param version: version of openmpi to use
    :param env: Optional environment to start calculations from, defaults to
        os.environ

    :return: path to openmpi base directory
    """
    env_override = os.environ.get(EnvVar.SCHRODINGER_OPENMPI)
    if env_override:
        return pathlib.Path(env_override)
    if version == OpenMPIVersion.VERSION_1_6_5:
        return pathlib.Path(
            os.environ["SCHRODINGER"]) / "internal" / "lib" / "openmpi-1.6.5"
    elif version == OpenMPIVersion.LATEST:
        return pathlib.Path(util.hunt("mmshare",
                                      util.DirectoryType.LIB)) / "openmpi"
    raise RuntimeError(f"{version} not in {OpenMPIVersion}")


def get_openmpi_env(version=OpenMPIVersion.LATEST,
                    *,
                    env: Optional[Dict[str, str]] = None) -> Dict[str, str]:
    """
    :param version: version of openmpi to use
    :param env: Optional environment to start calculations from, defaults to
        os.environ

    :return: updated environment directionary
    """

    if not env:
        env = os.environ
    opal_prefix = get_openmpi_prefix(version, env=env)
    openmpi_env = dict(
        OPAL_PREFIX=str(opal_prefix),
        PATH=os.environ[EnvVar.PATH],
        LD_LIBRARY_PATH=os.environ.get(EnvVar.LD_LIBRARY_PATH),
        # OMPI_MCA_mca_base_component_show_load_errors hides component load
        # errors like:

        # mca_base_component_repository_open: unable to open mca_ras_lsf:
        # libbat.so: cannot open shared object file: No such file or directory
        # (ignored)
        OMPI_MCA_mca_base_component_show_load_errors="0")
    prepend_path(openmpi_env, EnvVar.PATH, str(opal_prefix / "bin"))
    prepend_path(openmpi_env, EnvVar.LD_LIBRARY_PATH, str(opal_prefix / "lib"))
    return openmpi_env