From 3fb2a3916774276fbdfbec14d834c74cade2c393 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Tue, 22 Nov 2022 06:22:20 -0500 Subject: [PATCH] speed up tests with CLI Currently all tests use subprocess to call CLI. This might be slow, as it's expensive to import tensorflow (takes 2s on my computer). This patch adds a new argument to `deepmd.entrypoints.main.main` to allow arguments passed, so tests can call CLI without start a new subprocess. --- deepmd/common.py | 7 +++++ deepmd/entrypoints/main.py | 17 ++++++++++-- source/tests/common.py | 26 +++++++++++++++++++ source/tests/test_adjust_sel.py | 14 +++++----- source/tests/test_finetune_se_atten.py | 14 +++++----- source/tests/test_init_frz_model_se_a.py | 6 ++--- source/tests/test_init_frz_model_se_a_type.py | 6 ++--- source/tests/test_init_frz_model_se_atten.py | 6 ++--- source/tests/test_init_frz_model_se_r.py | 6 ++--- source/tests/test_mixed_prec_training.py | 4 +-- source/tests/test_model_compression_se_a.py | 8 +++--- ...ession_se_a_type_one_side_exclude_types.py | 8 +++--- source/tests/test_model_compression_se_r.py | 8 +++--- source/tests/test_model_compression_se_t.py | 8 +++--- source/tests/test_transfer.py | 4 +-- 15 files changed, 94 insertions(+), 48 deletions(-) diff --git a/deepmd/common.py b/deepmd/common.py index 5d23986b44..07232e8b47 100644 --- a/deepmd/common.py +++ b/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() diff --git a/deepmd/entrypoints/main.py b/deepmd/entrypoints/main.py index 2820f2ecfe..845f3940c7 100644 --- a/deepmd/entrypoints/main.py +++ b/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() diff --git a/source/tests/common.py b/source/tests/common.py index 38c5bc1aee..4b86ef477d 100644 --- a/source/tests/common.py +++ b/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 diff --git a/source/tests/test_adjust_sel.py b/source/tests/test_adjust_sel.py index d16cffe6ce..3f3494f040 100644 --- a/source/tests/test_adjust_sel.py +++ b/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 diff --git a/source/tests/test_finetune_se_atten.py b/source/tests/test_finetune_se_atten.py index 54eb663a9d..a3124c726c 100644 --- a/source/tests/test_finetune_se_atten.py +++ b/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") diff --git a/source/tests/test_init_frz_model_se_a.py b/source/tests/test_init_frz_model_se_a.py index cad93dbeca..1de2996524 100644 --- a/source/tests/test_init_frz_model_se_a.py +++ b/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") diff --git a/source/tests/test_init_frz_model_se_a_type.py b/source/tests/test_init_frz_model_se_a_type.py index 2e7a693e3f..7c66922d6b 100644 --- a/source/tests/test_init_frz_model_se_a_type.py +++ b/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") diff --git a/source/tests/test_init_frz_model_se_atten.py b/source/tests/test_init_frz_model_se_atten.py index 0decb022e9..d6b0b33260 100644 --- a/source/tests/test_init_frz_model_se_atten.py +++ b/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") diff --git a/source/tests/test_init_frz_model_se_r.py b/source/tests/test_init_frz_model_se_r.py index f74a0ea222..20d27b99da 100644 --- a/source/tests/test_init_frz_model_se_r.py +++ b/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") diff --git a/source/tests/test_mixed_prec_training.py b/source/tests/test_mixed_prec_training.py index 28a1d485f7..e51c64aaaf 100644 --- a/source/tests/test_mixed_prec_training.py +++ b/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): diff --git a/source/tests/test_model_compression_se_a.py b/source/tests/test_model_compression_se_a.py index 1c13a9d610..4db6e8a12f 100644 --- a/source/tests/test_model_compression_se_a.py +++ b/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 diff --git a/source/tests/test_model_compression_se_a_type_one_side_exclude_types.py b/source/tests/test_model_compression_se_a_type_one_side_exclude_types.py index b223742d2b..ab575283ab 100644 --- a/source/tests/test_model_compression_se_a_type_one_side_exclude_types.py +++ b/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 diff --git a/source/tests/test_model_compression_se_r.py b/source/tests/test_model_compression_se_r.py index 9b9218d175..9433c93017 100644 --- a/source/tests/test_model_compression_se_r.py +++ b/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 diff --git a/source/tests/test_model_compression_se_t.py b/source/tests/test_model_compression_se_t.py index 4c77a392c7..71e58fd212 100644 --- a/source/tests/test_model_compression_se_t.py +++ b/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 diff --git a/source/tests/test_transfer.py b/source/tests/test_transfer.py index 7b5884b2b4..7322be71fd 100644 --- a/source/tests/test_transfer.py +++ b/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)