Skip to content

Commit c0e6b4c

Browse files
wanghan-iapcmHan Wang
wanghan-iapcm
and
Han Wang
authored
Support some dflow.Workflow commands (#112)
- also provides UTs for command line argument parsing. Co-authored-by: Han Wang <[email protected]>
1 parent 96c08af commit c0e6b4c

File tree

4 files changed

+243
-4
lines changed

4 files changed

+243
-4
lines changed

‎dpgen2/entrypoint/main.py

Lines changed: 19 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -31,14 +31,15 @@
3131
watch,
3232
default_watching_keys,
3333
)
34+
from .workflow import (
35+
workflow_subcommands,
36+
add_subparser_workflow_subcommand,
37+
execute_workflow_subcommand,
38+
)
3439
from dpgen2 import (
3540
__version__
3641
)
3742

38-
#####################################
39-
# logging
40-
logging.basicConfig(level=logging.INFO)
41-
4243

4344
def main_parser() -> argparse.ArgumentParser:
4445
"""DPGEN2 commandline options argument parser.
@@ -182,8 +183,13 @@ def main_parser() -> argparse.ArgumentParser:
182183
help="if specified, download regardless whether check points exist."
183184
)
184185

186+
# workflow subcommands
187+
for cmd in workflow_subcommands:
188+
add_subparser_workflow_subcommand(subparsers, cmd)
189+
185190
# --version
186191
parser.add_argument(
192+
"-v",
187193
'--version',
188194
action='version',
189195
version='DPGEN v%s' % __version__,
@@ -211,6 +217,10 @@ def parse_args(args: Optional[List[str]] = None):
211217

212218

213219
def main():
220+
#####################################
221+
# logging
222+
logging.basicConfig(level=logging.INFO)
223+
214224
args = parse_args()
215225
dict_args = vars(args)
216226

@@ -268,6 +278,11 @@ def main():
268278
prefix=args.prefix,
269279
chk_pnt=args.no_check_point,
270280
)
281+
elif args.command in workflow_subcommands:
282+
with open(args.CONFIG) as fp:
283+
config = json.load(fp)
284+
wfid = args.ID
285+
execute_workflow_subcommand(args.command, wfid, config)
271286
elif args.command is None:
272287
pass
273288
else:

‎dpgen2/entrypoint/workflow.py

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,48 @@
1+
import argparse, os, json, logging
2+
from dflow import (
3+
Workflow,
4+
)
5+
from typing import (
6+
Optional,
7+
)
8+
from dpgen2.entrypoint.args import (
9+
normalize as normalize_args,
10+
)
11+
from dpgen2.entrypoint.common import (
12+
global_config_workflow,
13+
)
14+
15+
workflow_subcommands = ["terminate",
16+
"stop",
17+
"suspend",
18+
"delete",
19+
"retry",
20+
"resume",
21+
]
22+
23+
def add_subparser_workflow_subcommand(
24+
subparsers,
25+
command: str
26+
):
27+
parser_cmd = subparsers.add_parser(
28+
command,
29+
help=f"{command.capitalize()} a DPGEN2 workflow.",
30+
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
31+
)
32+
parser_cmd.add_argument(
33+
"CONFIG", help="the config file in json format."
34+
)
35+
parser_cmd.add_argument(
36+
"ID", help="the ID of the workflow."
37+
)
38+
39+
def execute_workflow_subcommand(
40+
command: str,
41+
wfid: str,
42+
wf_config: Optional[dict] = {},
43+
):
44+
wf_config = normalize_args(wf_config)
45+
global_config_workflow(wf_config)
46+
wf = Workflow(id=wfid)
47+
getattr(wf, command)()
48+

‎tests/entrypoint/test_argparse.py

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,61 @@
1+
import unittest, json, shutil, os
2+
3+
from dpgen2.entrypoint.main import (
4+
main_parser,
5+
parse_args,
6+
workflow_subcommands,
7+
)
8+
9+
class ParserTest(unittest.TestCase):
10+
def setUp(self):
11+
self.parser = main_parser()
12+
13+
def test_commands(self):
14+
tested_commands = ['resubmit', 'status', 'download', 'watch']
15+
tested_commands += workflow_subcommands
16+
17+
for cmd in tested_commands:
18+
parsed = self.parser.parse_args([cmd, 'foo', 'bar'])
19+
self.assertEqual(parsed.command, cmd)
20+
self.assertEqual(parsed.CONFIG, 'foo')
21+
self.assertEqual(parsed.ID, 'bar')
22+
23+
tested_commands = ['submit']
24+
for cmd in tested_commands:
25+
parsed = self.parser.parse_args([cmd, 'foo'])
26+
self.assertEqual(parsed.command, cmd)
27+
self.assertEqual(parsed.CONFIG, 'foo')
28+
29+
def test_watch(self):
30+
parsed = self.parser.parse_args(
31+
['watch', 'foo', 'bar',
32+
'-k', 'foo', 'bar', 'tar',
33+
'-f', '10',
34+
'-d',
35+
'-p', 'myprefix',
36+
]
37+
)
38+
self.assertEqual(parsed.keys, ['foo', 'bar', 'tar'])
39+
self.assertEqual(parsed.download, True)
40+
self.assertEqual(parsed.frequency, 10)
41+
self.assertEqual(parsed.prefix, 'myprefix')
42+
43+
def test_dld(self):
44+
parsed = self.parser.parse_args(
45+
['download', 'foo', 'bar',
46+
'-k', 'foo', 'bar', 'tar',
47+
'-p', 'myprefix',
48+
]
49+
)
50+
self.assertEqual(parsed.keys, ['foo', 'bar', 'tar'])
51+
self.assertEqual(parsed.prefix, 'myprefix')
52+
53+
def test_resubmit(self):
54+
parsed = self.parser.parse_args(
55+
['resubmit', 'foo', 'bar',
56+
'-l',
57+
'--reuse', '0', '10-20',
58+
]
59+
)
60+
self.assertEqual(parsed.list, True)
61+
self.assertEqual(parsed.reuse, ['0', '10-20'])

‎tests/entrypoint/test_workflow.py

Lines changed: 115 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,115 @@
1+
import unittest, json, shutil, os
2+
import dflow
3+
from dflow import (
4+
Workflow,
5+
)
6+
import mock, textwrap
7+
from dpgen2.entrypoint.workflow import (
8+
execute_workflow_subcommand,
9+
)
10+
11+
class ParserTest(unittest.TestCase):
12+
@mock.patch('dflow.Workflow.terminate')
13+
def test_terminate(self, mocked_f):
14+
config = json.loads(foo_str)
15+
execute_workflow_subcommand("terminate", "foo", config)
16+
mocked_f.assert_called_with()
17+
18+
@mock.patch('dflow.Workflow.stop')
19+
def test_stop(self, mocked_f):
20+
config = json.loads(foo_str)
21+
execute_workflow_subcommand("stop", "foo", config)
22+
mocked_f.assert_called_with()
23+
24+
@mock.patch('dflow.Workflow.suspend')
25+
def test_suspend(self, mocked_f):
26+
config = json.loads(foo_str)
27+
execute_workflow_subcommand("suspend", "foo", config)
28+
mocked_f.assert_called_with()
29+
30+
@mock.patch('dflow.Workflow.delete')
31+
def test_delete(self, mocked_f):
32+
config = json.loads(foo_str)
33+
execute_workflow_subcommand("delete", "foo", config)
34+
mocked_f.assert_called_with()
35+
36+
@mock.patch('dflow.Workflow.retry')
37+
def test_retry(self, mocked_f):
38+
config = json.loads(foo_str)
39+
execute_workflow_subcommand("retry", "foo", config)
40+
mocked_f.assert_called_with()
41+
42+
@mock.patch('dflow.Workflow.resume')
43+
def test_resume(self, mocked_f):
44+
config = json.loads(foo_str)
45+
execute_workflow_subcommand("resume", "foo", config)
46+
mocked_f.assert_called_with()
47+
48+
49+
foo_str = textwrap.dedent("""
50+
{
51+
"default_step_config" : {
52+
"template_config" : {
53+
"image" : "dflow:1.1.4",
54+
"_comment" : "all"
55+
},
56+
"_comment" : "all"
57+
},
58+
59+
"step_configs":{
60+
"_comment" : "all"
61+
},
62+
63+
"upload_python_packages" : "/path/to/dpgen2",
64+
65+
"inputs": {
66+
"type_map": ["Al", "Mg"],
67+
"mass_map": [27, 24],
68+
"init_data_prefix": "",
69+
"init_data_sys": [
70+
"init/al.fcc.01x01x01/02.md/sys-0004/deepmd",
71+
"init/mg.fcc.01x01x01/02.md/sys-0004/deepmd"
72+
],
73+
"_comment" : "all"
74+
},
75+
"train":{
76+
"type" : "dp",
77+
"numb_models" : 4,
78+
"config" : {},
79+
"template_script" : "dp_input_template",
80+
"_comment" : "all"
81+
},
82+
83+
"explore" : {
84+
"type" : "lmp",
85+
"config" : {
86+
"command": "lmp -var restart 0"
87+
},
88+
"max_numb_iter" : 5,
89+
"conv_accuracy" : 0.9,
90+
"fatal_at_max" : false,
91+
"f_trust_lo": 0.05,
92+
"f_trust_hi": 0.50,
93+
"configuration_prefix": null,
94+
"configuration": [
95+
],
96+
"stages": [
97+
],
98+
"_comment" : "all"
99+
},
100+
"fp" : {
101+
"type" : "vasp",
102+
"run_config" : {
103+
"command": "source /opt/intel/oneapi/setvars.sh && mpirun -n 16 vasp_std"
104+
},
105+
"task_max": 2,
106+
"inputs_config" : {
107+
"pp_files": {"Al" : "vasp/POTCAR.Al", "Mg" : "vasp/POTCAR.Mg"},
108+
"incar": "vasp/INCAR",
109+
"kspacing": 0.32,
110+
"kgamma": true
111+
},
112+
"_comment" : "all"
113+
}
114+
}
115+
""")

0 commit comments

Comments
 (0)