env.py

"""Module that sets tensorflow working environment and exports inportant constants."""

import logging
import os
import re
import platform
from configparser import ConfigParser
from importlib import reload
from pathlib import Path
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
from packaging.version import Version

import numpy as np

if TYPE_CHECKING:
    from types import ModuleType

# import tensorflow v1 compatability
try:
    import tensorflow.compat.v1 as tf

    tf.disable_v2_behavior()
except ImportError:
    import tensorflow as tf
try:
    import tensorflow.compat.v2 as tfv2
except ImportError:
    tfv2 = None

__all__ = [
    "GLOBAL_CONFIG",
    "GLOBAL_TF_FLOAT_PRECISION",
    "GLOBAL_NP_FLOAT_PRECISION",
    "GLOBAL_ENER_FLOAT_PRECISION",
    "global_float_prec",
    "global_cvt_2_tf_float",
    "global_cvt_2_ener_float",
    "MODEL_VERSION",
    "SHARED_LIB_MODULE",
    "default_tf_session_config",
    "reset_default_tf_session_config",
    "op_module",
    "op_grads_module",
    "TRANSFER_PATTERN",
    "FITTING_NET_PATTERN",
    "EMBEDDING_NET_PATTERN",
    "TYPE_EMBEDDING_PATTERN",
    "ATTENTION_LAYER_PATTERN",
    "TF_VERSION"
]

SHARED_LIB_MODULE = "op"

# Python library version
try:
    tf_py_version = tf.version.VERSION
except AttributeError:
    tf_py_version = tf.__version__

EMBEDDING_NET_PATTERN = str(
    r"filter_type_\d+/matrix_\d+_\d+|"
    r"filter_type_\d+/bias_\d+_\d+|"
    r"filter_type_\d+/idt_\d+_\d+|"
    r"filter_type_all/matrix_\d+|"
    r"filter_type_all/matrix_\d+_\d+|"
    r"filter_type_all/matrix_\d+_\d+_\d+|"
    r"filter_type_all/bias_\d+|"
    r"filter_type_all/bias_\d+_\d+|"
    r"filter_type_all/bias_\d+_\d+_\d+|"
    r"filter_type_all/idt_\d+|"
    r"filter_type_all/idt_\d+_\d+|"
)

FITTING_NET_PATTERN = str(
    r"layer_\d+/matrix|"
    r"layer_\d+_type_\d+/matrix|"
    r"layer_\d+/bias|"
    r"layer_\d+_type_\d+/bias|"
    r"layer_\d+/idt|"
    r"layer_\d+_type_\d+/idt|"
    r"final_layer/matrix|"
    r"final_layer_type_\d+/matrix|"
    r"final_layer/bias|"
    r"final_layer_type_\d+/bias|"
)

TYPE_EMBEDDING_PATTERN = str(
    r"type_embed_net+/matrix_\d+|"
    r"type_embed_net+/bias_\d+|"
    r"type_embed_net+/idt_\d+|"
)

ATTENTION_LAYER_PATTERN = str(
    r"attention_layer_\d+/c_query/matrix|"
    r"attention_layer_\d+/c_query/bias|"
    r"attention_layer_\d+/c_key/matrix|"
    r"attention_layer_\d+/c_key/bias|"
    r"attention_layer_\d+/c_value/matrix|"
    r"attention_layer_\d+/c_value/bias|"
    r"attention_layer_\d+/c_out/matrix|"
    r"attention_layer_\d+/c_out/bias|"
    r"attention_layer_\d+/layer_normalization/beta|"
    r"attention_layer_\d+/layer_normalization/gamma|"
    r"attention_layer_\d+/layer_normalization_\d+/beta|"
    r"attention_layer_\d+/layer_normalization_\d+/gamma|"
)

TRANSFER_PATTERN = \
    EMBEDDING_NET_PATTERN + \
    FITTING_NET_PATTERN + \
    TYPE_EMBEDDING_PATTERN + \
    str(
        r"descrpt_attr/t_avg|"
        r"descrpt_attr/t_std|"
        r"fitting_attr/t_fparam_avg|"
        r"fitting_attr/t_fparam_istd|"
        r"fitting_attr/t_aparam_avg|"
        r"fitting_attr/t_aparam_istd|"
        r"model_attr/t_tab_info|"
        r"model_attr/t_tab_data|"
)

def set_env_if_empty(key: str, value: str, verbose: bool = True):
    """Set environment variable only if it is empty.

    Parameters
    ----------
    key : str
        env variable name
    value : str
        env variable value
    verbose : bool, optional
        if True action will be logged, by default True
    """
    if os.environ.get(key) is None:
        os.environ[key] = value
        if verbose:
            logging.warn(
                f"Environment variable {key} is empty. Use the default value {value}"
            )


def set_mkl():
    """Tuning MKL for the best performance.

    References
    ----------
    TF overview
    https://www.tensorflow.org/guide/performance/overview

    Fixing an issue in numpy built by MKL
    https://github.com/ContinuumIO/anaconda-issues/issues/11367
    https://github.com/numpy/numpy/issues/12374

    check whether the numpy is built by mkl, see
    https://github.com/numpy/numpy/issues/14751
    """
    if "mkl_rt" in np.__config__.get_info("blas_mkl_info").get("libraries", []):
        set_env_if_empty("KMP_BLOCKTIME", "0")
        set_env_if_empty(
            "KMP_AFFINITY", "granularity=fine,verbose,compact,1,0")
        reload(np)


def set_tf_default_nthreads():
    """Set TF internal number of threads to default=automatic selection.

    Notes
    -----
    `TF_INTRA_OP_PARALLELISM_THREADS` and `TF_INTER_OP_PARALLELISM_THREADS`
    control TF configuration of multithreading.
    """
    if "OMP_NUM_THREADS" not in os.environ or \
       "TF_INTRA_OP_PARALLELISM_THREADS" not in os.environ or \
       "TF_INTER_OP_PARALLELISM_THREADS" not in os.environ:
        logging.warning(
            "To get the best performance, it is recommended to adjust "
            "the number of threads by setting the environment variables "
            "OMP_NUM_THREADS, TF_INTRA_OP_PARALLELISM_THREADS, and "
            "TF_INTER_OP_PARALLELISM_THREADS.")
    set_env_if_empty("TF_INTRA_OP_PARALLELISM_THREADS", "0", verbose=False)
    set_env_if_empty("TF_INTER_OP_PARALLELISM_THREADS", "0", verbose=False)


def get_tf_default_nthreads() -> Tuple[int, int]:
    """Get TF paralellism settings.

    Returns
    -------
    Tuple[int, int]
        number of `TF_INTRA_OP_PARALLELISM_THREADS` and
        `TF_INTER_OP_PARALLELISM_THREADS`
    """
    return int(os.environ.get("TF_INTRA_OP_PARALLELISM_THREADS", "0")), int(
        os.environ.get("TF_INTER_OP_PARALLELISM_THREADS", "0")
    )


def get_tf_session_config() -> Any:
    """Configure tensorflow session.

    Returns
    -------
    Any
        session configure object
    """
    set_tf_default_nthreads()
    intra, inter = get_tf_default_nthreads()
    config = tf.ConfigProto(
        gpu_options=tf.GPUOptions(allow_growth=True),
        intra_op_parallelism_threads=intra, inter_op_parallelism_threads=inter
    )
    if Version(tf_py_version) >= Version('1.15') and int(os.environ.get("DP_AUTO_PARALLELIZATION", 0)):
        config.graph_options.rewrite_options.custom_optimizers.add().name = "dpparallel"
    return config


default_tf_session_config = get_tf_session_config()


def reset_default_tf_session_config(cpu_only: bool):
    """Limit tensorflow session to CPU or not.

    Parameters
    ----------
    cpu_only : bool
        If enabled, no GPU device is visible to the TensorFlow Session.
    """
    global default_tf_session_config
    if cpu_only:
        default_tf_session_config.device_count['GPU'] = 0
    else:
        if 'GPU' in default_tf_session_config.device_count:
            del default_tf_session_config.device_count['GPU']


def get_module(module_name: str) -> "ModuleType":
    """Load force module.

    Returns
    -------
    ModuleType
        loaded force module

    Raises
    ------
    FileNotFoundError
        if module is not found in directory
    """
    if platform.system() == "Windows":
        ext = ".dll"
        prefix = ""
    #elif platform.system() == "Darwin":
    #    ext = ".dylib"
    else:
        ext = ".so"
        prefix = "lib"

    module_file = (
        (Path(__file__).parent / SHARED_LIB_MODULE / (prefix + module_name))
        .with_suffix(ext)
        .resolve()
    )

    if not module_file.is_file():
        raise FileNotFoundError(f"module {module_name} does not exist")
    else:
        try:
            module = tf.load_op_library(str(module_file))
        except tf.errors.NotFoundError as e:
            # check CXX11_ABI_FLAG is compatiblity
            # see https://gcc.gnu.org/onlinedocs/libstdc++/manual/using_dual_abi.html
            # ABI should be the same
            if 'CXX11_ABI_FLAG' in tf.__dict__:
                tf_cxx11_abi_flag = tf.CXX11_ABI_FLAG
            else:
                tf_cxx11_abi_flag = tf.sysconfig.CXX11_ABI_FLAG
            if TF_CXX11_ABI_FLAG != tf_cxx11_abi_flag:
                raise RuntimeError(
                    "This deepmd-kit package was compiled with "
                    "CXX11_ABI_FLAG=%d, but TensorFlow runtime was compiled "
                    "with CXX11_ABI_FLAG=%d. These two library ABIs are "
                    "incompatible and thus an error is raised when loading %s. "
                    "You need to rebuild deepmd-kit against this TensorFlow "
                    "runtime." % (
                        TF_CXX11_ABI_FLAG,
                        tf_cxx11_abi_flag,
                        module_name,
                    )) from e

            # different versions may cause incompatibility
            # see #406, #447, #557, #774, and #796 for example
            # throw a message if versions are different
            if TF_VERSION != tf_py_version:
                raise RuntimeError(
                    "The version of TensorFlow used to compile this "
                    "deepmd-kit package is %s, but the version of TensorFlow "
                    "runtime you are using is %s. These two versions are "
                    "incompatible and thus an error is raised when loading %s. "
                    "You need to install TensorFlow %s, or rebuild deepmd-kit "
                    "against TensorFlow %s.\nIf you are using a wheel from "
                    "pypi, you may consider to install deepmd-kit execuating "
                    "`pip install deepmd-kit --no-binary deepmd-kit` "
                    "instead." % (
                        TF_VERSION,
                        tf_py_version,
                        module_name,
                        TF_VERSION,
                        tf_py_version,
                    )) from e
            error_message = (
                "This deepmd-kit package is inconsitent with TensorFlow "
                "Runtime, thus an error is raised when loading %s. "
                "You need to rebuild deepmd-kit against this TensorFlow "
                "runtime." % (
                    module_name,
                )
            )
            if TF_CXX11_ABI_FLAG == 1:
                # #1791
                error_message += (
                    "\nWARNING: devtoolset on RHEL6 and RHEL7 does not support _GLIBCXX_USE_CXX11_ABI=1. "
                    "See https://bugzilla.redhat.com/show_bug.cgi?id=1546704"
                )
            raise RuntimeError(error_message) from e
        return module


def _get_package_constants(
    config_file: Path = Path(__file__).parent / "pkg_config/run_config.ini",
) -> Dict[str, str]:
    """Read package constants set at compile time by CMake to dictionary.

    Parameters
    ----------
    config_file : str, optional
        path to CONFIG file, by default "pkg_config/run_config.ini"

    Returns
    -------
    Dict[str, str]
        dictionary with package constants
    """
    config = ConfigParser()
    config.read(config_file)
    return dict(config.items("CONFIG"))


GLOBAL_CONFIG = _get_package_constants()
MODEL_VERSION = GLOBAL_CONFIG["model_version"]
TF_VERSION = GLOBAL_CONFIG["tf_version"]
TF_CXX11_ABI_FLAG = int(GLOBAL_CONFIG["tf_cxx11_abi_flag"])

op_module = get_module("op_abi")
op_grads_module = get_module("op_grads")

# FLOAT_PREC
dp_float_prec = os.environ.get("DP_INTERFACE_PREC", "high").lower()
if dp_float_prec in ("high", ""):
    # default is high
    GLOBAL_TF_FLOAT_PRECISION = tf.float64
    GLOBAL_NP_FLOAT_PRECISION = np.float64
    GLOBAL_ENER_FLOAT_PRECISION = np.float64
    global_float_prec = "double"
elif dp_float_prec == "low":
    GLOBAL_TF_FLOAT_PRECISION = tf.float32
    GLOBAL_NP_FLOAT_PRECISION = np.float32
    GLOBAL_ENER_FLOAT_PRECISION = np.float64
    global_float_prec = "float"
else:
    raise RuntimeError(
        "Unsupported float precision option: %s. Supported: high,"
        "low. Please set precision with environmental variable "
        "DP_INTERFACE_PREC." % dp_float_prec
    )


def global_cvt_2_tf_float(xx: tf.Tensor) -> tf.Tensor:
    """Cast tensor to globally set TF precision.

    Parameters
    ----------
    xx : tf.Tensor
        input tensor

    Returns
    -------
    tf.Tensor
        output tensor cast to `GLOBAL_TF_FLOAT_PRECISION`
    """
    return tf.cast(xx, GLOBAL_TF_FLOAT_PRECISION)


def global_cvt_2_ener_float(xx: tf.Tensor) -> tf.Tensor:
    """Cast tensor to globally set energy precision.

    Parameters
    ----------
    xx : tf.Tensor
        input tensor

    Returns
    -------
    tf.Tensor
        output tensor cast to `GLOBAL_ENER_FLOAT_PRECISION`
    """
    return tf.cast(xx, GLOBAL_ENER_FLOAT_PRECISION)