Fix download and watch for prep-run-explore (deepmodeling#215)

Signed-off-by: zjgemi <[email protected]>

docs/quickcli.md

@@ -29,11 +29,11 @@ The progress of a workflow can be watched on-the-fly
 ```bash
 $ dpgen2 watch input.json WFID
 INFO:root:steps iter-000000--prep-run-train----------------------- finished
-INFO:root:steps iter-000000--prep-run-lmp------------------------- finished
+INFO:root:steps iter-000000--prep-run-explore--------------------- finished
 INFO:root:steps iter-000000--prep-run-fp-------------------------- finished
 INFO:root:steps iter-000000--collect-data------------------------- finished
 INFO:root:steps iter-000001--prep-run-train----------------------- finished
-INFO:root:steps iter-000001--prep-run-lmp------------------------- finished
+INFO:root:steps iter-000001--prep-run-explore--------------------- finished
 ...
 ```
 The artifacts can be downloaded on-the-fly with `-d` flag. Note that the existing files are automatically skipped if one sets `dflow_config["archive_mode"] = None`.

dpgen2/entrypoint/main.py

@@ -163,9 +163,9 @@ def main_parser() -> argparse.ArgumentParser:
             $ dpgen2 download CONFIG ID
 
             3. donwload specified input/output artifacts of certain steps. For example
-            $ dpgen2 download CONFIG ID -i 0-8 8 9 -d prep-run-train/input/init_data prep-run-lmp/output/trajs
+            $ dpgen2 download CONFIG ID -i 0-8 8 9 -d prep-run-train/input/init_data prep-run-explore/output/trajs
 
-            The command will download the init_data of prep-run-train's input and trajs of the prep-run-lmp's output from iterations 0 to 9 (by -i 0-8 8 9).
+            The command will download the init_data of prep-run-train's input and trajs of the prep-run-explore's output from iterations 0 to 9 (by -i 0-8 8 9).
             The supported step and the names of input/output can be checked by the -l flag.
             """
             )
@@ -229,7 +229,7 @@ def main_parser() -> argparse.ArgumentParser:
         type=str,
         nargs="+",
         default=default_watching_keys,
-        help="the subkey to watch. For example, 'prep-run-train' 'prep-run-lmp'",
+        help="the subkey to watch. For example, 'prep-run-train' 'prep-run-explore'",
     )
     parser_watch.add_argument(
         "-f",

dpgen2/entrypoint/watch.py

@@ -24,7 +24,7 @@
 
 default_watching_keys = [
     "prep-run-train",
-    "prep-run-lmp",
+    "prep-run-explore",
     "prep-run-fp",
     "collect-data",
 ]

dpgen2/utils/download_dpgen2_artifacts.py

@@ -60,7 +60,7 @@ def add_output(
     .add_output("models")
     .add_output("logs")
     .add_output("lcurves"),
-    "prep-run-lmp": DownloadDefinition()
+    "prep-run-explore": DownloadDefinition()
     .add_output("logs")
     .add_output("trajs")
     .add_output("model_devis"),
@@ -268,7 +268,7 @@ def _get_all_iterations(step_keys):
     ret = []
     for kk in step_keys:
         ii = get_iteration(kk)
-        if ii != "init":
+        if ii.startswith("iter-"):
             ii = int(ii.split("-")[1])
             ret.append(ii)
     ret = sorted(list(set(ret)))
@@ -341,7 +341,7 @@ def _filter_def_by_availability(
                 name in op_download_setting[subkey].output_def.keys()
             )
         else:
-            raise RuntimeError("unknown io style {io}")
+            raise RuntimeError(f"unknown io style {io}")
         if not avail:
             logging.warning(f"cannot find download settings for {dd}")
         else:

tests/utils/test_dflow_query.py

@@ -81,7 +81,7 @@
     "iter-000000--run-lmp-000002",
     "iter-000000--run-lmp-000003",
     "iter-000000--run-lmp-000000",
-    "iter-000000--prep-run-lmp",
+    "iter-000000--prep-run-explore",
     "iter-000000--select-confs",
     "iter-000000--prep-fp",
     "iter-000000--run-fp-000001",
@@ -103,7 +103,7 @@
     "iter-000001--run-lmp-000005",
     "iter-000001--run-lmp-000002",
     "iter-000001--run-lmp-000004",
-    "iter-000001--prep-run-lmp",
+    "iter-000001--prep-run-explore",
     "iter-000001--select-confs",
     "iter-000001--prep-fp",
     "iter-000001--run-fp-000001",
@@ -248,7 +248,7 @@ def test_sort_slice_ops(self):
             "iter-000000--run-lmp-000003",
             "iter-000000--run-lmp-000004",
             "iter-000000--run-lmp-000005",
-            "iter-000000--prep-run-lmp",
+            "iter-000000--prep-run-explore",
             "iter-000000--select-confs",
             "iter-000000--prep-fp",
             "iter-000000--run-fp-000000",

tests/utils/test_dl_dpgen2_arti.py

@@ -103,21 +103,21 @@ def test_train_download(self, mocked_dl):
 
     @mock.patch("dpgen2.utils.download_dpgen2_artifacts.download_artifact")
     def test_lmp_download(self, mocked_dl):
-        download_dpgen2_artifacts(Mockedwf(), "iter-000001--prep-run-lmp", None)
+        download_dpgen2_artifacts(Mockedwf(), "iter-000001--prep-run-explore", None)
         expected = [
             mock.call(
                 "arti-logs",
-                path=Path("iter-000001/prep-run-lmp/outputs"),
+                path=Path("iter-000001/prep-run-explore/outputs"),
                 skip_exists=True,
             ),
             mock.call(
                 "arti-trajs",
-                path=Path("iter-000001/prep-run-lmp/outputs"),
+                path=Path("iter-000001/prep-run-explore/outputs"),
                 skip_exists=True,
             ),
             mock.call(
                 "arti-model_devis",
-                path=Path("iter-000001/prep-run-lmp/outputs"),
+                path=Path("iter-000001/prep-run-explore/outputs"),
                 skip_exists=True,
             ),
         ]
@@ -216,23 +216,23 @@ def test_empty_download(self, mocked_dl):
     @mock.patch("dpgen2.utils.download_dpgen2_artifacts.download_artifact")
     def test_update_finished_steps_exist_steps(self, mocked_dl):
         wf = Mockedwf()
-        wf.keys = ["iter-000000--prep-run-train", "iter-000001--prep-run-lmp"]
+        wf.keys = ["iter-000000--prep-run-train", "iter-000001--prep-run-explore"]
         finished_keys = update_finished_steps(wf, ["iter-000000--prep-run-train"], True)
         self.assertEqual(finished_keys, wf.keys)
         expected = [
             mock.call(
                 "arti-logs",
-                path=Path("iter-000001/prep-run-lmp/outputs"),
+                path=Path("iter-000001/prep-run-explore/outputs"),
                 skip_exists=True,
             ),
             mock.call(
                 "arti-trajs",
-                path=Path("iter-000001/prep-run-lmp/outputs"),
+                path=Path("iter-000001/prep-run-explore/outputs"),
                 skip_exists=True,
             ),
             mock.call(
                 "arti-model_devis",
-                path=Path("iter-000001/prep-run-lmp/outputs"),
+                path=Path("iter-000001/prep-run-explore/outputs"),
                 skip_exists=True,
             ),
         ]
@@ -243,7 +243,7 @@ def test_update_finished_steps_exist_steps(self, mocked_dl):
     @mock.patch("dpgen2.utils.download_dpgen2_artifacts.download_artifact")
     def test_update_finished_steps_none_steps(self, mocked_dl):
         wf = Mockedwf()
-        wf.keys = ["iter-000000--prep-run-train", "iter-000001--prep-run-lmp"]
+        wf.keys = ["iter-000000--prep-run-train", "iter-000001--prep-run-explore"]
         finished_keys = update_finished_steps(wf, None, True)
         self.assertEqual(finished_keys, wf.keys)
         expected = [
@@ -284,17 +284,17 @@ def test_update_finished_steps_none_steps(self, mocked_dl):
             ),
             mock.call(
                 "arti-logs",
-                path=Path("iter-000001/prep-run-lmp/outputs"),
+                path=Path("iter-000001/prep-run-explore/outputs"),
                 skip_exists=True,
             ),
             mock.call(
                 "arti-trajs",
-                path=Path("iter-000001/prep-run-lmp/outputs"),
+                path=Path("iter-000001/prep-run-explore/outputs"),
                 skip_exists=True,
             ),
             mock.call(
                 "arti-model_devis",
-                path=Path("iter-000001/prep-run-lmp/outputs"),
+                path=Path("iter-000001/prep-run-explore/outputs"),
                 skip_exists=True,
             ),
         ]

tests/utils/test_dl_dpgen2_arti_by_def.py

@@ -59,7 +59,7 @@ class Mockedwf:
     keys = [
         "iter-000000--prep-run-train",
         "iter-000001--prep-run-train",
-        "iter-000000--prep-run-lmp",
+        "iter-000000--prep-run-explore",
     ]
 
     def query_step_by_key(self, key):
@@ -117,16 +117,16 @@ def test_download(self, mocked_dl):
                 step_defs=[
                     "prep-run-train/input/init_models",
                     "prep-run-train/output/logs",
-                    "prep-run-lmp/input/foo",
-                    "prep-run-lmp/output/trajs",
+                    "prep-run-explore/input/foo",
+                    "prep-run-explore/output/trajs",
                 ],
                 prefix="foo",
                 chk_pnt=False,
             )
         self.assertEqual(len(log.output), 1)
         self.assertEqual(len(log.records), 1)
         self.assertIn(
-            "cannot find download settings for prep-run-lmp/input/foo",
+            "cannot find download settings for prep-run-explore/input/foo",
             log.output[0],
         )
         expected = [
@@ -142,7 +142,7 @@ def test_download(self, mocked_dl):
             ),
             mock.call(
                 "arti-trajs",
-                path=Path("foo/iter-000000/prep-run-lmp/output/trajs"),
+                path=Path("foo/iter-000000/prep-run-explore/output/trajs"),
                 skip_exists=True,
             ),
             mock.call(
@@ -197,16 +197,16 @@ def test_download_with_ckpt(self, mocked_dl):
                 step_defs=[
                     "prep-run-train/input/init_models",
                     "prep-run-train/output/logs",
-                    "prep-run-lmp/input/foo",
-                    "prep-run-lmp/output/trajs",
+                    "prep-run-explore/input/foo",
+                    "prep-run-explore/output/trajs",
                 ],
                 prefix="foo",
                 chk_pnt=True,
             )
         self.assertEqual(len(log.output), 1)
         self.assertEqual(len(log.records), 1)
         self.assertIn(
-            "cannot find download settings for prep-run-lmp/input/foo",
+            "cannot find download settings for prep-run-explore/input/foo",
             log.output[0],
         )
         expected = [
@@ -222,7 +222,7 @@ def test_download_with_ckpt(self, mocked_dl):
             ),
             mock.call(
                 "arti-trajs",
-                path=Path("foo/iter-000000/prep-run-lmp/output/trajs"),
+                path=Path("foo/iter-000000/prep-run-explore/output/trajs"),
                 skip_exists=True,
             ),
             mock.call(
@@ -246,16 +246,16 @@ def test_download_with_ckpt(self, mocked_dl):
             step_defs=[
                 "prep-run-train/input/init_models",
                 "prep-run-train/output/logs",
-                "prep-run-lmp/output/trajs",
-                "prep-run-lmp/output/model_devis",
+                "prep-run-explore/output/trajs",
+                "prep-run-explore/output/model_devis",
             ],
             prefix="foo",
             chk_pnt=True,
         )
         expected = [
             mock.call(
                 "arti-model_devis",
-                path=Path("foo/iter-000000/prep-run-lmp/output/model_devis"),
+                path=Path("foo/iter-000000/prep-run-explore/output/model_devis"),
                 skip_exists=True,
             ),
         ]