Merge pull request #19 from amcadmus/api

api update from amcadmus

.github/workflows/build_cc.yml

@@ -0,0 +1,30 @@
+on:
+  push:
+  pull_request:
+name: Build C++
+jobs:
+  testpython:
+    name: Build C++
+    runs-on: ubuntu-20.04
+    strategy:  
+      matrix:
+        include:
+        - float_prec: high
+          variant: cpu
+        - float_prec: low
+          variant: cpu
+        - float_prec: high
+          variant: gpu
+        - float_prec: low
+          variant: gpu
+    steps:
+    - uses: actions/checkout@master
+      with:
+        submodules: true
+    - run: sudo apt update && sudo apt install nvidia-cuda-toolkit
+      if: matrix.variant == 'gpu'
+    - run: source/install/build_cc.sh
+      env:
+        FLOAT_PREC: ${{ matrix.float_prec }}
+        CC: gcc-7
+        CXX: g++-7

.github/workflows/lint_python.yml

@@ -17,9 +17,9 @@ jobs:
         python-version: ${{ matrix.python }}
     - name: Install requirements
       run: pip install -r requirements.txt
-    - uses: marian-code/python-lint-annotate@master
+    - uses: marian-code/python-lint-annotate@v2.5.0
       with:
-        python-root-list: "./deepmd/'*'.py ./deepmd/'*'/'*'.py ./source/train/'*'.py ./source/tests/'*'.py ./source/scripts/'*'.py ./source/op/'*'.py"
+        python-root-list: "./deepmd/*.py ./deepmd/*/*.py ./source/train/*.py ./source/tests/*.py ./source/op/*.py"
         use-black: true
         use-isort: true
         use-mypy: true

.github/workflows/test_cc.yml

@@ -0,0 +1,11 @@
+on:
+  push:
+  pull_request:
+name: Test C++
+jobs:
+  testpython:
+    name: Test C++
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@master
+    - run: source/install/test_cc.sh

.github/workflows/test_python.yml

@@ -5,7 +5,7 @@ name: Test Python
 jobs:
   testpython:
     name: Test Python
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-18.04
     strategy:  
       matrix:
         include:
@@ -59,15 +59,15 @@ jobs:
       with:
         path: ~/.cache/pip
         key:
-          ${{ runner.os }}-pip-${{ hashFiles('**/setup.py') }}
+          ${{ runner.os }}-pip-${{ hashFiles('**/setup.py') }}-${{ hashFiles('**/requirements.txt') }}
         restore-keys: |
           ${{ runner.os }}-pip-
     - run: |
         sudo apt update
         sudo apt install gcc-${{ matrix.gcc }} g++-${{ matrix.gcc }}
-    - run: pip install .[cpu,test]
+    - run: pip install -e .[cpu,test] codecov
       env:
         CC: gcc-${{ matrix.gcc }}
         CXX: g++-${{ matrix.gcc }}
         TENSORFLOW_VERSION: ${{ matrix.tf }}
-    - run: cd source/tests && python -m unittest
+    - run: pytest --cov=deepmd source/tests && codecov

.gitmodules

@@ -1,3 +1,3 @@
-[submodule "source/op/cuda/cub"]
-	path = source/op/cuda/cub
-	url = git://github.com/NVlabs/cub.git
+[submodule "source/lib/src/cuda/cub"]
+	path = source/lib/src/cuda/cub
+	url = git://github.com/NVIDIA/cub.git

README.md

@@ -44,8 +44,6 @@ The code is organized as follows:
 
 * `source/op`: tensorflow op implementation. working with library.
 
-* `source/scripts`: Python script for model freezing.
-
 * `source/train`: Python modules and scripts for training and testing.
 
 

deepmd/__init__.py

@@ -1,18 +1,11 @@
 """Root of the deepmd package, exposes all public classes and submodules."""
 
-from . import descriptor
-from . import fit
-from . import loss
-from . import utils
 import deepmd.utils.network as network
-from .infer.deep_eval import DeepEval
-from .infer.deep_pot import DeepPot
-from .infer.deep_dipole import DeepDipole
-from .infer.deep_polar import DeepPolar
-from .infer.deep_polar import DeepGlobalPolar
-from .infer.deep_wfc import DeepWFC
-from .infer.data_modifier import DipoleChargeModifier
+
+from . import cluster, descriptor, fit, loss, utils
 from .env import set_mkl
+from .infer import DeepPotential
+from .infer.data_modifier import DipoleChargeModifier
 
 set_mkl()
 
@@ -26,12 +19,9 @@
     "fit",
     "loss",
     "utils",
+    "cluster",
     "network",
     "DeepEval",
-    "DeepPot",
-    "DeepDipole",
-    "DeepPolar",
-    "DeepGlobalPolar",
-    "DeepWFC",
+    "DeepPotential",
     "DipoleChargeModifier",
 ]

deepmd/cluster/__init__.py

@@ -0,0 +1,22 @@
+"""Module that reads node resources, auto detects if running local or on SLURM."""
+
+from .local import get_resource as get_local_res
+from .slurm import get_resource as get_slurm_res
+import os
+from typing import List, Tuple, Optional
+
+__all__ = ["get_resource"]
+
+
+def get_resource() -> Tuple[str, List[str], Optional[List[int]]]:
+    """Get local or slurm resources: nodename, nodelist, and gpus.
+
+    Returns
+    -------
+    Tuple[str, List[str], Optional[List[int]]]
+        nodename, nodelist, and gpus
+    """
+    if "SLURM_JOB_NODELIST" in os.environ:
+        return get_slurm_res()
+    else:
+        return get_local_res()

source/train/Local.py → deepmd/cluster/local.py

@@ -4,9 +4,11 @@
 import socket
 from typing import List, Tuple, Optional
 
+__all__ = ["get_resource"]
+
 
 def get_resource() -> Tuple[str, List[str], Optional[List[int]]]:
-    """Get loacl resources: nodename, nodelist, and gpus.
+    """Get local resources: nodename, nodelist, and gpus.
 
     Returns
     -------
@@ -16,10 +18,8 @@ def get_resource() -> Tuple[str, List[str], Optional[List[int]]]:
     nodename = socket.gethostname()
     nodelist = [nodename]
     gpus_env = os.getenv("CUDA_VISIBLE_DEVICES", None)
-    if gpus_env is None:
+    if not gpus_env:
         gpus = None
-    elif gpus_env == "":
-        gpus = []
     else:
         gpus = [int(gpu) for gpu in gpus_env.split(",")]
 

deepmd/cluster/slurm.py

@@ -0,0 +1,84 @@
+"""MOdule to get resources on SLURM cluster.
+
+References
+----------
+https://github.com/deepsense-ai/tensorflow_on_slurm ####
+"""
+
+import re
+import os
+from typing import List, Tuple, Optional, Iterable
+
+__all__ = ["get_resource"]
+
+
+def get_resource() -> Tuple[str, List[str], Optional[List[int]]]:
+    """Get SLURM resources: nodename, nodelist, and gpus.
+
+    Returns
+    -------
+    Tuple[str, List[str], Optional[List[int]]]
+        nodename, nodelist, and gpus
+
+    Raises
+    ------
+    RuntimeError
+        if number of nodes could not be retrieved
+    ValueError
+        list of nodes is not of the same length sa number of nodes
+    ValueError
+        if current nodename is not found in node list
+    """
+    nodelist = _expand_nodelist(os.environ["SLURM_JOB_NODELIST"])
+    nodename = os.environ["SLURMD_NODENAME"]
+    num_nodes_env = os.getenv("SLURM_JOB_NUM_NODES")
+    if num_nodes_env:
+        num_nodes = int(num_nodes_env)
+    else:
+        raise RuntimeError("Could not get SLURM number of nodes")
+
+    if len(nodelist) != num_nodes:
+        raise ValueError(
+            f"Number of slurm nodes {len(nodelist)} not equal to {num_nodes}"
+        )
+    if nodename not in nodelist:
+        raise ValueError(
+            f"Nodename({nodename}) not in nodelist({nodelist}). This should not happen!"
+        )
+    gpus_env = os.getenv("CUDA_VISIBLE_DEVICES")
+    if not gpus_env:
+        gpus = None
+    else:
+        gpus = [int(gpu) for gpu in gpus_env.split(",")]
+    return nodename, nodelist, gpus
+
+
+def _pad_zeros(iterable: Iterable, length: int):
+    return (str(t).rjust(length, "0") for t in iterable)
+
+
+def _expand_ids(ids: str) -> List[str]:
+    result = []
+    for _id in ids.split(","):
+        if "-" in _id:
+            str_end = _id.split("-")[1]
+            begin, end = [int(token) for token in _id.split("-")]
+            result.extend(_pad_zeros(range(begin, end + 1), len(str_end)))
+        else:
+            result.append(_id)
+    return result
+
+
+def _expand_nodelist(nodelist: str) -> List[str]:
+    result = []
+    interval_list = nodelist.split(",")
+    for interval in interval_list:
+        match = re.search(r"(.*)\[(.*)\]", interval)
+        if match:
+            prefix = match.group(1)
+            ids = match.group(2)
+            ids_list = _expand_ids(ids)
+            result.extend([f"{prefix}{_id}" for _id in ids_list])
+        else:
+            result.append(interval)
+    return result

deepmd/common.py

@@ -20,7 +20,7 @@
 import yaml
 
 from deepmd.env import op_module, tf
-from deepmd.RunOptions import global_tf_float_precision, global_np_float_precision
+from deepmd.run_options import GLOBAL_TF_FLOAT_PRECISION, GLOBAL_NP_FLOAT_PRECISION
 
 if TYPE_CHECKING:
     _DICT_VAL = TypeVar("_DICT_VAL")
@@ -34,7 +34,7 @@
 
 # define constants
 PRECISION_DICT = {
-    "default": global_tf_float_precision,
+    "default": GLOBAL_TF_FLOAT_PRECISION,
     "float16": tf.float16,
     "float32": tf.float32,
     "float64": tf.float64,
@@ -427,8 +427,10 @@ def expand_sys_str(root_dir: Union[str, Path]) -> List[str]:
     List[str]
         list of string pointing to system directories
     """
-    matches = [d for d in Path(root_dir).rglob("*") if (d / "type.raw").is_file()]
-    return [str(m.relative_to(root_dir)) for m in matches]
+    matches = [str(d) for d in Path(root_dir).rglob("*") if (d / "type.raw").is_file()]
+    if (Path(root_dir) / "type.raw").is_file():
+        matches += [root_dir]
+    return matches
 
 
 def docstring_parameter(*sub: Tuple[str, ...]):
@@ -452,14 +454,32 @@ def dec(obj: "_OBJ") -> "_OBJ":
 
     return dec
 
-def get_np_precision(precision):
+
+def get_np_precision(precision: "_PRECISION") -> np.dtype:
+    """Get numpy precision constant from string.
+
+    Parameters
+    ----------
+    precision : _PRECISION
+        string name of numpy constant or default
+
+    Returns
+    -------
+    np.dtype
+        numpy presicion constant
+
+    Raises
+    ------
+    RuntimeError
+        if string is invalid
+    """
     if precision == "default":
-        return  global_np_float_precision
+        return GLOBAL_NP_FLOAT_PRECISION
     elif precision == "float16":
         return np.float16
     elif precision == "float32":
         return np.float32
     elif precision == "float64":
         return np.float64
     else:
-        raise RuntimeError("%d is not a valid precision" % precision)
+        raise RuntimeError(f"{precision} is not a valid precision")

deepmd/descriptor/hybrid.py

@@ -4,8 +4,8 @@
 from deepmd.env import tf
 from deepmd.common import ClassArg
 from deepmd.env import op_module
-from deepmd.RunOptions import global_tf_float_precision
-from deepmd.RunOptions import global_np_float_precision
+from deepmd.run_options import GLOBAL_TF_FLOAT_PRECISION
+from deepmd.run_options import GLOBAL_NP_FLOAT_PRECISION
 # from deepmd.descriptor import DescrptLocFrame
 # from deepmd.descriptor import DescrptSeA
 # from deepmd.descriptor import DescrptSeAT
@@ -164,7 +164,7 @@ def build (self,
         with tf.variable_scope('descrpt_attr' + suffix, reuse = reuse) :
             t_rcut = tf.constant(self.get_rcut(), 
                                  name = 'rcut', 
-                                 dtype = global_tf_float_precision)
+                                 dtype = GLOBAL_TF_FLOAT_PRECISION)
             t_ntypes = tf.constant(self.get_ntypes(), 
                                    name = 'ntypes', 
                                    dtype = tf.int32)