speed up tests with CLI

deepmd/common.py

@@ -578,3 +578,10 @@ def wrapper(self, *args, **kwargs):
         else:
             return safe_cast_tensor(returned_tensor, self.precision, GLOBAL_TF_FLOAT_PRECISION)
     return wrapper
+
+
+def clear_session():
+    """Reset all state generated by DeePMD-kit."""
+    tf.reset_default_graph()
+    # TODO: remove this line when data_requirement is not a global variable
+    data_requirement.clear()

deepmd/entrypoints/main.py

@@ -7,6 +7,7 @@
 from typing import Dict, List, Optional
 
 from deepmd import __version__
+from deepmd.common import clear_session
 from deepmd.entrypoints import (
     compress,
     config,
@@ -548,15 +549,24 @@ def parse_args(args: Optional[List[str]] = None) -> argparse.Namespace:
     return parsed_args
 
 
-def main():
+def main(args: Optional[List[str]] = None):
     """DeePMD-Kit entry point.
 
+    Parameters
+    ----------
+    args: List[str], optional
+        list of command line arguments, used to avoid calling from the subprocess,
+        as it is quite slow to import tensorflow
+
     Raises
     ------
     RuntimeError
         if no command was input
     """
-    args = parse_args()
+    if args is not None:
+        clear_session()
+
+    args = parse_args(args=args)
 
     # do not set log handles for None, it is useless
     # log handles for train will be set separatelly
@@ -592,3 +602,6 @@ def main():
         pass
     else:
         raise RuntimeError(f"unknown command {args.command}")
+
+    if args is not None:
+        clear_session()

source/tests/common.py

@@ -6,6 +6,7 @@
 from deepmd.env import GLOBAL_NP_FLOAT_PRECISION
 from deepmd.common import j_loader as dp_j_loader
 from deepmd.utils import random as dp_random
+from deepmd.entrypoints.main import main
 
 if GLOBAL_NP_FLOAT_PRECISION == np.float32 :
     global_default_fv_hh = 1e-2
@@ -407,3 +408,28 @@ def strerch_box(old_coord, old_box, new_box):
     nbox = new_box.reshape(3,3)
     ncoord = ocoord @ np.linalg.inv(obox) @ nbox
     return ncoord.reshape(old_coord.shape)
+
+
+def run_dp(cmd: str) -> int:
+    """Run DP directly from the entry point instead of the subprocess.
+
+    It is quite slow to start DeePMD-kit with subprocess.
+
+    Parameters
+    ----------
+    cmd : str
+        The command to run.
+
+    Returns
+    -------
+    int
+        Always returns 0.
+    """
+    cmds = cmd.split()
+    if cmds[0] == 'dp':
+        cmds = cmds[1:]
+    else:
+        raise RuntimeError('The command is not dp')
+
+    main(cmds)
+    return 0

source/tests/test_adjust_sel.py

@@ -6,7 +6,7 @@
 from deepmd.infer import DeepPot
 from deepmd.env import MODEL_VERSION
 # from deepmd.entrypoints.compress import compress
-from common import j_loader, tests_path
+from common import j_loader, tests_path, run_dp
 
 from deepmd.env import GLOBAL_NP_FLOAT_PRECISION
 if GLOBAL_NP_FLOAT_PRECISION == np.float32 :
@@ -44,26 +44,26 @@ def _init_models():
     with open(INPUT, "w") as fp:
         json.dump(jdata, fp, indent=4)
 
-    ret = _subprocess_run("dp train " + INPUT + " --skip-neighbor-stat")
+    ret = run_dp("dp train " + INPUT + " --skip-neighbor-stat")
     np.testing.assert_equal(ret, 0, 'DP train failed!')
-    ret = _subprocess_run("dp freeze -o " + frozen_model)
+    ret = run_dp("dp freeze -o " + frozen_model)
     np.testing.assert_equal(ret, 0, 'DP freeze failed!')
 
     jdata["training"]["numb_steps"] = 0
     jdata["model"]["descriptor"]["sel"] = [2, 4] # equal to data
     with open(INPUT, "w") as fp:
         json.dump(jdata, fp, indent=4)
-    ret = _subprocess_run("dp train " + INPUT + " -f " + frozen_model + " --skip-neighbor-stat")
+    ret = run_dp("dp train " + INPUT + " -f " + frozen_model + " --skip-neighbor-stat")
     np.testing.assert_equal(ret, 0, 'DP model adjust sel failed!')
-    ret = _subprocess_run("dp freeze -o " + decreased_model)
+    ret = run_dp("dp freeze -o " + decreased_model)
     np.testing.assert_equal(ret, 0, 'DP freeze failed!')
 
     jdata["model"]["descriptor"]["sel"] = [300, 300] # equal to data
     with open(INPUT, "w") as fp:
         json.dump(jdata, fp, indent=4)
-    ret = _subprocess_run("dp train " + INPUT + " -f " + frozen_model + " --skip-neighbor-stat")
+    ret = run_dp("dp train " + INPUT + " -f " + frozen_model + " --skip-neighbor-stat")
     np.testing.assert_equal(ret, 0, 'DP model adjust sel failed!')
-    ret = _subprocess_run("dp freeze -o " + increased_model)
+    ret = run_dp("dp freeze -o " + increased_model)
     np.testing.assert_equal(ret, 0, 'DP freeze failed!')
     return INPUT, frozen_model, decreased_model, increased_model
 

source/tests/test_finetune_se_atten.py

@@ -2,7 +2,7 @@
 import numpy as np
 import unittest
 import subprocess as sp
-from common import j_loader, tests_path
+from common import j_loader, tests_path, run_dp
 from deepmd.utils.graph import get_tensor_by_name
 from deepmd.utils.argcheck import normalize
 from deepmd.utils.compat import update_deepmd_input
@@ -62,17 +62,17 @@ def _init_models():
     with open(INPUT_FINETUNE_MIX, "w") as fp:
         json.dump(jdata_finetune, fp, indent=4)
 
-    ret = _subprocess_run("dp train " + INPUT_PRE)
+    ret = run_dp("dp train " + INPUT_PRE)
     np.testing.assert_equal(ret, 0, 'DP train failed!')
-    ret = _subprocess_run("dp freeze -o " + pretrained_model)
+    ret = run_dp("dp freeze -o " + pretrained_model)
     np.testing.assert_equal(ret, 0, 'DP freeze failed!')
-    ret = _subprocess_run("dp train " + INPUT_FINETUNE + " -t " + pretrained_model)
+    ret = run_dp("dp train " + INPUT_FINETUNE + " -t " + pretrained_model)
     np.testing.assert_equal(ret, 0, 'DP finetune failed!')
-    ret = _subprocess_run("dp freeze -o " + finetuned_model)
+    ret = run_dp("dp freeze -o " + finetuned_model)
     np.testing.assert_equal(ret, 0, 'DP freeze failed!')
-    ret = _subprocess_run("dp train " + INPUT_FINETUNE_MIX + " -t " + pretrained_model)
+    ret = run_dp("dp train " + INPUT_FINETUNE_MIX + " -t " + pretrained_model)
     np.testing.assert_equal(ret, 0, 'DP finetune failed!')
-    ret = _subprocess_run("dp freeze -o " + finetuned_model_mixed_type)
+    ret = run_dp("dp freeze -o " + finetuned_model_mixed_type)
     np.testing.assert_equal(ret, 0, 'DP freeze failed!')
 
     jdata_pre = update_deepmd_input(jdata_pre, warning=True, dump="input_v2_compat.json")

source/tests/test_init_frz_model_se_a.py

@@ -2,7 +2,7 @@
 import numpy as np
 import unittest
 import subprocess as sp
-from common import j_loader, tests_path
+from common import j_loader, tests_path, run_dp
 from deepmd.train.trainer import DPTrainer
 from deepmd.train.run_options import RunOptions
 from deepmd.utils.argcheck import normalize
@@ -48,9 +48,9 @@ def _init_models():
     jdata["training"]["save_ckpt"] = ckpt
     with open(INPUT, "w") as fp:
         json.dump(jdata, fp, indent=4)
-    ret = _subprocess_run("dp train " + INPUT)
+    ret = run_dp("dp train " + INPUT)
     np.testing.assert_equal(ret, 0, 'DP train failed!')
-    ret = _subprocess_run("dp freeze -c " + str(tests_path) + " -o " + frozen_model)
+    ret = run_dp("dp freeze -c " + str(tests_path) + " -o " + frozen_model)
     np.testing.assert_equal(ret, 0, 'DP freeze failed!')
 
     jdata = update_deepmd_input(jdata, warning=True, dump="input_v2_compat.json")

source/tests/test_init_frz_model_se_a_type.py

@@ -2,7 +2,7 @@
 import numpy as np
 import unittest
 import subprocess as sp
-from common import j_loader, tests_path
+from common import j_loader, tests_path, run_dp
 from deepmd.train.trainer import DPTrainer
 from deepmd.train.run_options import RunOptions
 from deepmd.utils.argcheck import normalize
@@ -52,9 +52,9 @@ def _init_models():
     jdata['model']["type_embedding"] = type_embed
     with open(INPUT, "w") as fp:
         json.dump(jdata, fp, indent=4)
-    ret = _subprocess_run("dp train " + INPUT)
+    ret = run_dp("dp train " + INPUT)
     np.testing.assert_equal(ret, 0, 'DP train failed!')
-    ret = _subprocess_run("dp freeze -c " + str(tests_path) + " -o " + frozen_model)
+    ret = run_dp("dp freeze -c " + str(tests_path) + " -o " + frozen_model)
     np.testing.assert_equal(ret, 0, 'DP freeze failed!')
 
     jdata = update_deepmd_input(jdata, warning=True, dump="input_v2_compat.json")

source/tests/test_init_frz_model_se_atten.py

@@ -2,7 +2,7 @@
 import numpy as np
 import unittest
 import subprocess as sp
-from common import j_loader, tests_path
+from common import j_loader, tests_path, run_dp
 from deepmd.train.trainer import DPTrainer
 from deepmd.train.run_options import RunOptions
 from deepmd.utils.argcheck import normalize
@@ -51,9 +51,9 @@ def _init_models():
     jdata['model']["descriptor"]['sel'] = 120
     with open(INPUT, "w") as fp:
         json.dump(jdata, fp, indent=4)
-    ret = _subprocess_run("dp train " + INPUT)
+    ret = run_dp("dp train " + INPUT)
     np.testing.assert_equal(ret, 0, 'DP train failed!')
-    ret = _subprocess_run("dp freeze -c " + str(tests_path) + " -o " + frozen_model)
+    ret = run_dp("dp freeze -c " + str(tests_path) + " -o " + frozen_model)
     np.testing.assert_equal(ret, 0, 'DP freeze failed!')
 
     jdata = update_deepmd_input(jdata, warning=True, dump="input_v2_compat.json")

source/tests/test_init_frz_model_se_r.py

@@ -2,7 +2,7 @@
 import numpy as np
 import unittest
 import subprocess as sp
-from common import j_loader, tests_path
+from common import j_loader, tests_path, run_dp
 from deepmd.train.trainer import DPTrainer
 from deepmd.train.run_options import RunOptions
 from deepmd.utils.argcheck import normalize
@@ -56,9 +56,9 @@ def _init_models():
     jdata["training"]["save_ckpt"] = ckpt
     with open(INPUT, "w") as fp:
         json.dump(jdata, fp, indent=4)
-    ret = _subprocess_run("dp train " + INPUT)
+    ret = run_dp("dp train " + INPUT)
     np.testing.assert_equal(ret, 0, 'DP train failed!')
-    ret = _subprocess_run("dp freeze -c " + str(tests_path) + " -o " + frozen_model)
+    ret = run_dp("dp freeze -c " + str(tests_path) + " -o " + frozen_model)
     np.testing.assert_equal(ret, 0, 'DP freeze failed!')
 
     jdata = update_deepmd_input(jdata, warning=True, dump="input_v2_compat.json")

source/tests/test_mixed_prec_training.py

@@ -6,7 +6,7 @@
 
 from deepmd.infer import DeepPot
 # from deepmd.entrypoints.compress import compress
-from common import j_loader, tests_path
+from common import j_loader, tests_path, run_dp
 from deepmd.env import TF_VERSION
 
 
@@ -43,7 +43,7 @@ def test_training(self):
         _TF_VERSION = Version(TF_VERSION)
         # check the TF_VERSION, when TF < 1.12, mixed precision is not allowed 
         if _TF_VERSION >= Version('1.14.0'):
-            ret = _subprocess_run("dp train " + self.INPUT)
+            ret = run_dp("dp train " + self.INPUT)
             np.testing.assert_equal(ret, 0, 'DP train failed!')
 
     def tearDown(self):

source/tests/test_model_compression_se_a.py

@@ -6,7 +6,7 @@
 from deepmd.infer import DeepPot
 from deepmd.env import MODEL_VERSION
 # from deepmd.entrypoints.compress import compress
-from common import j_loader, tests_path
+from common import j_loader, tests_path, run_dp
 
 from deepmd.env import GLOBAL_NP_FLOAT_PRECISION
 if GLOBAL_NP_FLOAT_PRECISION == np.float32 :
@@ -41,11 +41,11 @@ def _init_models():
     with open(INPUT, "w") as fp:
         json.dump(jdata, fp, indent=4)
 
-    ret = _subprocess_run("dp train " + INPUT)
+    ret = run_dp("dp train " + INPUT)
     np.testing.assert_equal(ret, 0, 'DP train failed!')
-    ret = _subprocess_run("dp freeze -o " + frozen_model)
+    ret = run_dp("dp freeze -o " + frozen_model)
     np.testing.assert_equal(ret, 0, 'DP freeze failed!')
-    ret = _subprocess_run("dp compress " + " -i " + frozen_model + " -o " + compressed_model)
+    ret = run_dp("dp compress " + " -i " + frozen_model + " -o " + compressed_model)
     np.testing.assert_equal(ret, 0, 'DP model compression failed!')
     return INPUT, frozen_model, compressed_model
 

source/tests/test_model_compression_se_a_type_one_side_exclude_types.py

@@ -6,7 +6,7 @@
 from deepmd.infer import DeepPot
 from deepmd.env import MODEL_VERSION
 # from deepmd.entrypoints.compress import compress
-from common import j_loader, tests_path
+from common import j_loader, tests_path, run_dp
 
 from deepmd.env import GLOBAL_NP_FLOAT_PRECISION
 if GLOBAL_NP_FLOAT_PRECISION == np.float32 :
@@ -43,11 +43,11 @@ def _init_models():
     with open(INPUT, "w") as fp:
         json.dump(jdata, fp, indent=4)
 
-    ret = _subprocess_run("dp train " + INPUT)
+    ret = run_dp("dp train " + INPUT)
     np.testing.assert_equal(ret, 0, 'DP train failed!')
-    ret = _subprocess_run("dp freeze -o " + frozen_model)
+    ret = run_dp("dp freeze -o " + frozen_model)
     np.testing.assert_equal(ret, 0, 'DP freeze failed!')
-    ret = _subprocess_run("dp compress " + " -i " + frozen_model + " -o " + compressed_model)
+    ret = run_dp("dp compress " + " -i " + frozen_model + " -o " + compressed_model)
     np.testing.assert_equal(ret, 0, 'DP model compression failed!')
     return INPUT, frozen_model, compressed_model
 

source/tests/test_model_compression_se_r.py

@@ -6,7 +6,7 @@
 from deepmd.infer import DeepPot
 from deepmd.env import MODEL_VERSION
 # from deepmd.entrypoints.compress import compress
-from common import j_loader, tests_path
+from common import j_loader, tests_path, run_dp
 
 from deepmd.env import GLOBAL_NP_FLOAT_PRECISION
 if GLOBAL_NP_FLOAT_PRECISION == np.float32 :
@@ -49,11 +49,11 @@ def _init_models():
     with open(INPUT, "w") as fp:
         json.dump(jdata, fp, indent=4)
 
-    ret = _subprocess_run("dp train " + INPUT)
+    ret = run_dp("dp train " + INPUT)
     np.testing.assert_equal(ret, 0, 'DP train failed!')
-    ret = _subprocess_run("dp freeze -o " + frozen_model)
+    ret = run_dp("dp freeze -o " + frozen_model)
     np.testing.assert_equal(ret, 0, 'DP freeze failed!')
-    ret = _subprocess_run("dp compress " + " -i " + frozen_model + " -o " + compressed_model)
+    ret = run_dp("dp compress " + " -i " + frozen_model + " -o " + compressed_model)
     np.testing.assert_equal(ret, 0, 'DP model compression failed!')
     return INPUT, frozen_model, compressed_model
 

source/tests/test_model_compression_se_t.py

@@ -6,7 +6,7 @@
 from deepmd.infer import DeepPot
 from deepmd.env import MODEL_VERSION
 # from deepmd.entrypoints.compress import compress
-from common import j_loader, tests_path
+from common import j_loader, tests_path, run_dp
 
 from deepmd.env import GLOBAL_NP_FLOAT_PRECISION
 if GLOBAL_NP_FLOAT_PRECISION == np.float32 :
@@ -49,11 +49,11 @@ def _init_models():
     with open(INPUT, "w") as fp:
         json.dump(jdata, fp, indent=4)
 
-    ret = _subprocess_run("dp train " + INPUT)
+    ret = run_dp("dp train " + INPUT)
     np.testing.assert_equal(ret, 0, 'DP train failed!')
-    ret = _subprocess_run("dp freeze -o " + frozen_model)
+    ret = run_dp("dp freeze -o " + frozen_model)
     np.testing.assert_equal(ret, 0, 'DP freeze failed!')
-    ret = _subprocess_run("dp compress " + " -i " + frozen_model + " -o " + compressed_model)
+    ret = run_dp("dp compress " + " -i " + frozen_model + " -o " + compressed_model)
     np.testing.assert_equal(ret, 0, 'DP model compression failed!')
     return INPUT, frozen_model, compressed_model
 

source/tests/test_transfer.py

@@ -5,7 +5,7 @@
 
 from deepmd.env import tf
 from deepmd.infer import DeepPot
-from common import j_loader, tests_path
+from common import j_loader, tests_path, run_dp
 from infer.convert2pb import convert_pbtxt_to_pb
 from deepmd.entrypoints.transfer import load_graph, transform_graph
 
@@ -37,7 +37,7 @@ def setUpClass(self):
         self.new_model = str(tests_path / "dp-new.pb")
         convert_pbtxt_to_pb(str(tests_path / os.path.join("infer","deeppot.pbtxt")), self.old_model)
         convert_pbtxt_to_pb(str(tests_path / os.path.join("infer","deeppot-1.pbtxt")), self.raw_model)
-        ret = _subprocess_run("dp transfer -O " + self.old_model + " -r " + self.raw_model + " -o " + self.new_model)
+        ret = run_dp("dp transfer -O " + self.old_model + " -r " + self.raw_model + " -o " + self.new_model)
         np.testing.assert_equal(ret, 0, 'DP transfer failed!')
 
         self.dp = DeepPot(self.new_model)