Support some dflow.Workflow commands (#112)

- also provides UTs for command line argument parsing.

Co-authored-by: Han Wang <wang_han@iapcm.ac.cn>

dpgen2/entrypoint/main.py

@@ -31,14 +31,15 @@
     watch,
     default_watching_keys,
 )
+from .workflow import (
+    workflow_subcommands,
+    add_subparser_workflow_subcommand,
+    execute_workflow_subcommand,
+)
 from dpgen2 import (
     __version__
 )
 
-#####################################
-# logging
-logging.basicConfig(level=logging.INFO)
-
 
 def main_parser() -> argparse.ArgumentParser:
     """DPGEN2 commandline options argument parser.
@@ -182,8 +183,13 @@ def main_parser() -> argparse.ArgumentParser:
         help="if specified, download regardless whether check points exist."
     )
 
+    # workflow subcommands
+    for cmd in workflow_subcommands:
+        add_subparser_workflow_subcommand(subparsers, cmd)
+
     # --version
     parser.add_argument(
+        "-v",
         '--version', 
         action='version', 
         version='DPGEN v%s' % __version__,
@@ -211,6 +217,10 @@ def parse_args(args: Optional[List[str]] = None):
 
 
 def main():
+    #####################################
+    # logging
+    logging.basicConfig(level=logging.INFO)
+
     args = parse_args()
     dict_args = vars(args)
 
@@ -268,6 +278,11 @@ def main():
             prefix=args.prefix,
             chk_pnt=args.no_check_point,
         )
+    elif args.command in workflow_subcommands:
+        with open(args.CONFIG) as fp:
+            config = json.load(fp)
+        wfid = args.ID
+        execute_workflow_subcommand(args.command, wfid, config)
     elif args.command is None:
         pass
     else:

dpgen2/entrypoint/workflow.py

@@ -0,0 +1,48 @@
+import argparse, os, json, logging
+from dflow import (
+    Workflow,
+)
+from typing import (
+    Optional,
+)
+from dpgen2.entrypoint.args import (
+    normalize as normalize_args,
+)
+from dpgen2.entrypoint.common import (
+    global_config_workflow,
+)
+
+workflow_subcommands = ["terminate", 
+                        "stop",
+                        "suspend",
+                        "delete",
+                        "retry",
+                        "resume",
+                        ]
+
+def add_subparser_workflow_subcommand(
+        subparsers,
+        command: str
+):
+    parser_cmd = subparsers.add_parser(
+        command,
+        help=f"{command.capitalize()} a DPGEN2 workflow.",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parser_cmd.add_argument(
+        "CONFIG", help="the config file in json format."
+    )
+    parser_cmd.add_argument(
+        "ID", help="the ID of the workflow."
+    )
+
+def execute_workflow_subcommand(
+        command: str,
+        wfid: str,
+        wf_config: Optional[dict] = {},
+):
+    wf_config = normalize_args(wf_config)    
+    global_config_workflow(wf_config)
+    wf = Workflow(id=wfid)
+    getattr(wf, command)()
+

tests/entrypoint/test_argparse.py

@@ -0,0 +1,61 @@
+import unittest, json, shutil, os
+
+from dpgen2.entrypoint.main import (
+    main_parser,
+    parse_args,
+    workflow_subcommands,
+)
+
+class ParserTest(unittest.TestCase):
+    def setUp(self):
+        self.parser = main_parser()
+
+    def test_commands(self):
+        tested_commands = ['resubmit', 'status', 'download', 'watch']
+        tested_commands += workflow_subcommands
+
+        for cmd in tested_commands:
+            parsed = self.parser.parse_args([cmd, 'foo', 'bar'])
+            self.assertEqual(parsed.command, cmd)
+            self.assertEqual(parsed.CONFIG, 'foo')
+            self.assertEqual(parsed.ID, 'bar')
+
+        tested_commands = ['submit']
+        for cmd in tested_commands:
+            parsed = self.parser.parse_args([cmd, 'foo'])
+            self.assertEqual(parsed.command, cmd)
+            self.assertEqual(parsed.CONFIG, 'foo')
+
+    def test_watch(self):
+        parsed = self.parser.parse_args(
+            ['watch', 'foo', 'bar', 
+             '-k', 'foo', 'bar', 'tar', 
+             '-f', '10',
+             '-d',
+             '-p', 'myprefix',
+             ]
+        )
+        self.assertEqual(parsed.keys, ['foo', 'bar', 'tar'])
+        self.assertEqual(parsed.download, True)
+        self.assertEqual(parsed.frequency, 10)
+        self.assertEqual(parsed.prefix, 'myprefix')
+
+    def test_dld(self):
+        parsed = self.parser.parse_args(
+            ['download', 'foo', 'bar', 
+             '-k', 'foo', 'bar', 'tar', 
+             '-p', 'myprefix',
+             ]
+        )
+        self.assertEqual(parsed.keys, ['foo', 'bar', 'tar'])
+        self.assertEqual(parsed.prefix, 'myprefix')
+
+    def test_resubmit(self):
+        parsed = self.parser.parse_args(
+            ['resubmit', 'foo', 'bar', 
+             '-l',
+             '--reuse', '0', '10-20',
+             ]
+        )
+        self.assertEqual(parsed.list, True)
+        self.assertEqual(parsed.reuse, ['0', '10-20'])

tests/entrypoint/test_workflow.py

@@ -0,0 +1,115 @@
+import unittest, json, shutil, os
+import dflow
+from dflow import (
+    Workflow,
+)
+import mock, textwrap
+from dpgen2.entrypoint.workflow import (
+    execute_workflow_subcommand,
+)
+
+class ParserTest(unittest.TestCase):
+    @mock.patch('dflow.Workflow.terminate')
+    def test_terminate(self, mocked_f):
+        config = json.loads(foo_str)
+        execute_workflow_subcommand("terminate", "foo", config)
+        mocked_f.assert_called_with()
+
+    @mock.patch('dflow.Workflow.stop')
+    def test_stop(self, mocked_f):
+        config = json.loads(foo_str)
+        execute_workflow_subcommand("stop", "foo", config)
+        mocked_f.assert_called_with()
+
+    @mock.patch('dflow.Workflow.suspend')
+    def test_suspend(self, mocked_f):
+        config = json.loads(foo_str)
+        execute_workflow_subcommand("suspend", "foo", config)
+        mocked_f.assert_called_with()
+
+    @mock.patch('dflow.Workflow.delete')
+    def test_delete(self, mocked_f):
+        config = json.loads(foo_str)
+        execute_workflow_subcommand("delete", "foo", config)
+        mocked_f.assert_called_with()
+
+    @mock.patch('dflow.Workflow.retry')
+    def test_retry(self, mocked_f):
+        config = json.loads(foo_str)
+        execute_workflow_subcommand("retry", "foo", config)
+        mocked_f.assert_called_with()
+
+    @mock.patch('dflow.Workflow.resume')
+    def test_resume(self, mocked_f):
+        config = json.loads(foo_str)
+        execute_workflow_subcommand("resume", "foo", config)
+        mocked_f.assert_called_with()
+
+
+foo_str = textwrap.dedent("""
+{
+    "default_step_config" : {
+	"template_config" : {
+	    "image" : "dflow:1.1.4",
+	    "_comment" : "all"
+	},
+	"_comment" : "all"
+    },
+
+    "step_configs":{
+	"_comment" : "all"
+    },
+
+    "upload_python_packages" : "/path/to/dpgen2",
+
+    "inputs": {
+	"type_map":		["Al", "Mg"],
+	"mass_map":		[27, 24],
+	"init_data_prefix":	"",
+	"init_data_sys":	[
+	    "init/al.fcc.01x01x01/02.md/sys-0004/deepmd",
+	    "init/mg.fcc.01x01x01/02.md/sys-0004/deepmd"
+	],
+	"_comment" : "all"
+    },
+    "train":{
+	"type" :	"dp",
+	"numb_models" : 4,
+	"config" : {},
+	"template_script" : "dp_input_template",
+	"_comment" : "all"
+    },
+
+    "explore" : {
+	"type" : "lmp",
+	"config" : {
+	    "command": "lmp -var restart 0"
+	},
+	"max_numb_iter" :	5,
+	"conv_accuracy" :	0.9,
+	"fatal_at_max" :	false,
+	"f_trust_lo":		0.05,
+	"f_trust_hi":		0.50,
+	"configuration_prefix": null, 
+	"configuration":	[
+	],
+	"stages":	[
+	],
+	"_comment" : "all"
+    },
+    "fp" : {
+	"type" :	"vasp",
+	"run_config" : {
+	    "command": "source /opt/intel/oneapi/setvars.sh && mpirun -n 16 vasp_std"
+	},
+	"task_max":	2,
+	"inputs_config" : {
+	    "pp_files":	{"Al" : "vasp/POTCAR.Al", "Mg" : "vasp/POTCAR.Mg"},
+	    "incar":    "vasp/INCAR",
+	    "kspacing":	0.32,
+	    "kgamma":	true
+	},
+	"_comment" : "all"
+    }
+}
+""")