Source code for schrodinger.utils.env

import contextlib
import glob
import os
import pathlib
import subprocess
import sys
from typing import Dict
from typing import Optional
from typing import Union

from schrodinger.job import util

DYLD_LIBRARY_PATH = 'DYLD_LIBRARY_PATH'
ORIGINAL_LD_LIBRARY_PATH = "ORIGINAL_LD_LIBRARY_PATH"
LD_LIBRARY_PATH = 'LD_LIBRARY_PATH'
PATH = 'PATH'
PERL5LIB = 'PERL5LIB'

SCHRODINGER = os.getenv("SCHRODINGER")
INTERNAL_LIB_DIR = os.path.join(SCHRODINGER, "internal", "lib")
TF_PATH = os.path.join(os.path.abspath(os.path.dirname(__file__)), "..", "..",
                       "tensorflow", "python", "_pywrap_tensorflow_internal.so")


def append_path(env_dict: Dict[str, str], env_var: str, path: str):
    """
    Appends a path to env_var (defined as constants in schrodinger.utils.env)
    in the env_dict. This uses the traditional path seperator.
    We normalize the path argument.

    :param env_dict: Dictionary of environment variables.
    :param env_var: existing environment variable (example:
        schrodinger.utils.env.PATH)
    :param path: path to extend with

    """
    path = os.path.normpath(path)
    if env_var not in env_dict:
        env_dict[env_var] = path
    else:
        env_dict[env_var] += os.pathsep + path


def prepend_path(env_dict: Dict[str, str], env_var: str, path: str):
    """
    Appends a path to env_var (defined as constants in schrodinger.utils.env)
    We normalize the path argument.

    :param env_dict: Dictionary of environment variables.
    :param env_var: existing environment variable (example:
                    schrodinger.utils.env.PATH)
    :param path: path to extend with

    """
    path = os.path.normpath(path)
    if env_var not in env_dict:
        env_dict[env_var] = path
    else:
        env_dict[env_var] = path + os.pathsep + env_dict[env_var]


@contextlib.contextmanager
def swap_ld_library_path():
    """
    Restore LD_LIBRARY_PATH to the value that it was when we first ran a
    toplevel script. ORIGINAL_LD_LIBRARY_PATH is set by toplevel.sh
    """
    with EnvironmentContext(LD_LIBRARY_PATH,
                            os.environ.get(ORIGINAL_LD_LIBRARY_PATH)):
        yield


class EnvironmentContext:
    """
    Context manager to temporarily set an environment variable
    and then return it to its original value upon exit.
    """

    def __init__(self, key: str, value: str):
        self.key = key
        self.value = value

    def __enter__(self):
        self.orig_value = os.environ.get(self.key)
        if self.value is None:
            if self.key in os.environ:
                del os.environ[self.key]
        else:
            os.environ[self.key] = self.value

    def __exit__(self, *args):
        if self.orig_value is not None:
            os.environ[self.key] = self.orig_value
        elif self.key in os.environ:
            del os.environ[self.key]


class MultiEnvironmentContext(contextlib.ExitStack):
    """
    Context manager to temporarily set multiple environment variables
    """

    def __init__(self, env: dict):
        super().__init__()
        self._env = env

    def __enter__(self):
        for k, v in self._env.items():
            self.enter_context(EnvironmentContext(k, v))


def get_ld_library_path_for_system_ssh() -> Optional[str]:
    """
    In order to run ssh as a subprocess, sometimes we need to use system openssl
    libraries. This will return a modified LD_LIBRARY_PATH that does not have
    Schrodinger OpenSSL libraries in it.

    :return: LD_LIBRARY_PATH str or None if LD_LIBRARY_PATH was not originally
        set
    """
    # Do this on Linux only
    if sys.platform.startswith('linux') and LD_LIBRARY_PATH in os.environ:
        proc = subprocess.run(['ssh', '-V'], capture_output=True)
        if proc.returncode != 0:
            orig_path = os.environ[LD_LIBRARY_PATH]
            ssl_dir = os.path.normpath(
                os.path.join(SCHRODINGER, "internal", "lib", "ssl"))
            return orig_path.replace(ssl_dir, '')
    return os.environ.get(LD_LIBRARY_PATH)


@contextlib.contextmanager
def remove_ssl_ld_library_path():
    """
    Temporarily remove SSL library path from the LD_LIBRARY_PATH. This will
    force ssh, as called by MPI, to be linked against the system libcrypto
    rather than the one in the Schrodinger software (JOBCON-5027).
    But sometimes the user has ssh but not libcrypto and then their ssh
    has to use our libcrypto (DESMOND-7780).
    So test whether ssh fails or not by actually running it.  What a world!
    """
    with EnvironmentContext(LD_LIBRARY_PATH,
                            get_ld_library_path_for_system_ssh()):
        yield


@contextlib.contextmanager
def prepend_sys_path(path: Union[str, pathlib.Path]):
    """
    Prepend a specific path to sys.path. Prefer placing the module for import
    in an existing sys.path.

    :param path: directory name of path
    """
    sys.path.insert(0, str(path))
    try:
        yield
    finally:
        sys.path.pop(0)


def add_cuda_stubs_for_tensorflow(env_dict: Dict[str, str]):
    """
    Prepends the CUDA stubs library to LD_LIBRARY_PATH if libcuda.so is not
    linked to the tensorflow library.

    :param env_dict: dict of the environment
    """

    if not sys.platform.startswith("linux"):
        raise OSError("CUDA stubs library is only available for Linux")

    cuda_stubs = os.path.join(INTERNAL_LIB_DIR, "cuda-stubs")
    output = subprocess.check_output(["ldd", TF_PATH], universal_newlines=True)
    for line in output.splitlines():
        if all(x in line for x in ('libcuda.so', 'not found')):
            prepend_path(env_dict, LD_LIBRARY_PATH, cuda_stubs)
            return


def find_knime_java_path() -> str:
    """
    Find the path to the jre included with KNIME

    :raises RuntimeError: if KNIME is not installed or it cannot find java

    :return: the path to the java file in the jre
    """
    knime_dir = util.hunt('knime')
    if knime_dir == "":
        raise RuntimeError("KNIME is not installed")
    java_filename = "java"
    if sys.platform == 'win32':
        java_filename += '.exe'
    java_glob = glob.glob(
        f'{knime_dir}/knime/**/plugins/*/jre/**/'
        f'{java_filename}',
        recursive=True)
    if len(java_glob) != 1:
        raise RuntimeError(f"Could not find java path. Instead found "
                           f"{java_glob}")
    return java_glob[0]