add UT for fold_keys

Signed-off-by: zjgemi <liuxin_zijian@163.com>
deepmodeling · Mar 14, 2024 · a04edc2 · a04edc2
1 parent 4c8a827
commit a04edc2
Show file tree

Hide file tree

Showing 2 changed files with 116 additions and 5 deletions.
diff --git a/dpgen2/entrypoint/submit.py b/dpgen2/entrypoint/submit.py
@@ -743,22 +743,24 @@ def successful_step_keys(wf):
 def get_superop(key):
     if "prep-train" in key:
         return key.replace("prep-train", "prep-run-train")
-    elif "run-train" in key:
+    elif "run-train-" in key:
         return re.sub("run-train-[0-9]*", "prep-run-train", key)
     elif "prep-lmp" in key:
         return key.replace("prep-lmp", "prep-run-explore")
-    elif "run-lmp" in key:
+    elif "run-lmp-" in key:
         return re.sub("run-lmp-[0-9]*", "prep-run-explore", key)
     elif "prep-fp" in key:
         return key.replace("prep-fp", "prep-run-fp")
-    elif "run-fp" in key:
+    elif "run-fp-" in key:
         return re.sub("run-fp-[0-9]*", "prep-run-fp", key)
     elif "prep-caly-input" in key:
         return key.replace("prep-caly-input", "prep-run-explore")
-    elif "collect-run-calypso" in key:
+    elif "collect-run-calypso-" in key:
         return re.sub("collect-run-calypso-[0-9]*-[0-9]*", "prep-run-explore", key)
-    elif "prep-run-dp-optim" in key:
+    elif "prep-run-dp-optim-" in key:
         return re.sub("prep-run-dp-optim-[0-9]*-[0-9]*", "prep-run-explore", key)
+    elif "run-caly-model-devi" in key:
+        return key.replace("run-caly-model-devi", "prep-run-explore")
     return None
 
 

diff --git a/tests/entrypoint/test_submit.py b/tests/entrypoint/test_submit.py
@@ -19,6 +19,7 @@
 from dpgen2.entrypoint.submit import (
     copy_scheduler_plans,
     expand_idx,
+    fold_keys,
     print_list_steps,
     submit_concurrent_learning,
     update_reuse_step_scheduler,
@@ -956,3 +957,111 @@ def test(self):
 PRINT ARG=restraint.bias
 """
 )
+
+
+def test_fold_keys_lmp():
+    all_step_keys = [
+        'init--scheduler',
+        'init--id',
+        'iter-000000--prep-run-train',
+        'iter-000000--prep-train',
+        'iter-000000--run-train-0000',
+        'iter-000000--run-train-0001',
+        'iter-000000--run-train-0002',
+        'iter-000000--run-train-0003',
+        'iter-000000--prep-run-explore',
+        'iter-000000--prep-lmp',
+        'iter-000000--run-lmp-000000',
+        'iter-000000--run-lmp-000001',
+        'iter-000000--select-confs',
+        'iter-000000--prep-run-fp',
+        'iter-000000--prep-fp',
+        'iter-000000--run-fp-000000',
+        'iter-000000--run-fp-000001',
+        'iter-000000--run-fp-000002',
+    ]
+    folded_keys = fold_keys(all_step_keys)
+    assert folded_keys == {
+        "init--scheduler": ["init--scheduler"],
+        "init--id": ["init--id"],
+        "iter-000000--prep-run-train" : [
+            'iter-000000--prep-train',
+            'iter-000000--run-train-0000',
+            'iter-000000--run-train-0001',
+            'iter-000000--run-train-0002',
+            'iter-000000--run-train-0003',
+        ],
+        "iter-000000--prep-run-explore": [
+            'iter-000000--prep-lmp',
+            'iter-000000--run-lmp-000000',
+            'iter-000000--run-lmp-000001',
+        ],
+        "iter-000000--select-confs": ["iter-000000--select-confs"],
+        "iter-000000--prep-run-fp": [
+            'iter-000000--prep-fp',
+            'iter-000000--run-fp-000000',
+            'iter-000000--run-fp-000001',
+            'iter-000000--run-fp-000002',
+        ],
+    }
+
+
+def test_fold_keys_caly():
+    all_step_keys = [
+        'init--scheduler',
+        'init--id',
+        'iter-000000--prep-run-train',
+        'iter-000000--prep-train',
+        'iter-000000--run-train-0000',
+        'iter-000000--run-train-0001',
+        'iter-000000--run-train-0002',
+        'iter-000000--run-train-0003',
+        'iter-000000--prep-run-explore',
+        'iter-000000--prep-caly-input',
+        'iter-000000--prep-run-dp-optim-000000-0',
+        'iter-000000--prep-run-dp-optim-000000-1',
+        'iter-000000--prep-run-dp-optim-000001-0',
+        'iter-000000--prep-run-dp-optim-000001-1',
+        'iter-000000--collect-run-calypso-000000-0',
+        'iter-000000--collect-run-calypso-000000-1',
+        'iter-000000--collect-run-calypso-000001-0',
+        'iter-000000--collect-run-calypso-000001-1',
+        'iter-000000--run-caly-model-devi',
+        'iter-000000--select-confs',
+        'iter-000000--prep-run-fp',
+        'iter-000000--prep-fp',
+        'iter-000000--run-fp-000000',
+        'iter-000000--run-fp-000001',
+        'iter-000000--run-fp-000002',
+    ]
+    folded_keys = fold_keys(all_step_keys)
+    assert folded_keys == {
+        "init--scheduler": ["init--scheduler"],
+        "init--id": ["init--id"],
+        "iter-000000--prep-run-train" : [
+            'iter-000000--prep-train',
+            'iter-000000--run-train-0000',
+            'iter-000000--run-train-0001',
+            'iter-000000--run-train-0002',
+            'iter-000000--run-train-0003',
+        ],
+        "iter-000000--prep-run-explore": [
+            'iter-000000--prep-caly-input',
+            'iter-000000--prep-run-dp-optim-000000-0',
+            'iter-000000--prep-run-dp-optim-000000-1',
+            'iter-000000--prep-run-dp-optim-000001-0',
+            'iter-000000--prep-run-dp-optim-000001-1',
+            'iter-000000--collect-run-calypso-000000-0',
+            'iter-000000--collect-run-calypso-000000-1',
+            'iter-000000--collect-run-calypso-000001-0',
+            'iter-000000--collect-run-calypso-000001-1',
+            'iter-000000--run-caly-model-devi',
+        ],
+        "iter-000000--select-confs": ["iter-000000--select-confs"],
+        "iter-000000--prep-run-fp": [
+            'iter-000000--prep-fp',
+            'iter-000000--run-fp-000000',
+            'iter-000000--run-fp-000001',
+            'iter-000000--run-fp-000002',
+        ],
+    }