Repository: deepmodeling/dpgen2 (language: python)

Issue:
Implement the op RunLmp
Implement the op `RunLmp`. This OP run a LAMMPS MD task prepared by `PrepLmp`, and outputs a MD trajectory and the model deviation of each configuration on the trajectory.

We have to implement the `execute` method of the class. What the op does is explained in the [docstr of the class](https://github.com/wanghan-iapcm/dpgen2/blob/ef08f44f7a39c81fe469885ffaa4cd8be9d4d0a3/dpgen2/op/run_lmp.py#L12-L18) and the interface of the `execute` method is provided in [the docstr of the method](https://github.com/wanghan-iapcm/dpgen2/blob/ef08f44f7a39c81fe469885ffaa4cd8be9d4d0a3/dpgen2/op/run_lmp.py#L44-L66).

One can also take the [implementation of `RunDPTrain`](https://github.com/wanghan-iapcm/dpgen2/blob/master/dpgen2/op/run_dp_train.py) as an example for the `RunLmp`


Discussion / hints:

Fix patch:
diff --git a/dpgen2/flow/dpgen_loop.py b/dpgen2/flow/dpgen_loop.py
index c3f71dbe..8d26346f 100644
--- a/dpgen2/flow/dpgen_loop.py
+++ b/dpgen2/flow/dpgen_loop.py
@@ -368,6 +368,9 @@ def _loop(
         executor=step_executor,
         **step_config,
     )
+    scheduler_step.template.outputs.parameters[
+        "exploration_scheduler"
+    ].global_name = "exploration_scheduler"
     steps.add(scheduler_step)
 
     id_step = Step(
@@ -467,6 +470,9 @@ def _dpgen(
         executor=step_executor,
         **step_config,
     )
+    scheduler_step.template.outputs.parameters[
+        "exploration_scheduler"
+    ].global_name = "exploration_scheduler"
     steps.add(scheduler_step)
 
     id_step = Step(
diff --git a/dpgen2/utils/dflow_query.py b/dpgen2/utils/dflow_query.py
index bb700a29..f3a136de 100644
--- a/dpgen2/utils/dflow_query.py
+++ b/dpgen2/utils/dflow_query.py
@@ -1,3 +1,4 @@
+import logging
 import re
 from typing import (
     Any,
@@ -50,20 +51,29 @@ def get_last_scheduler(
     """
     get the output Scheduler of the last successful iteration
     """
+    outputs = wf.query_global_outputs()
+    if (
+        outputs is not None
+        and hasattr(outputs, "parameters")
+        and "exploration_scheduler" in outputs.parameters
+    ):
+        return outputs.parameters["exploration_scheduler"].value
+
+    logging.warn("Exploration scheduler not found in the global outputs")
     scheduler_keys_ = []
     for ii in keys:
         if get_subkey(ii) == "scheduler":
             scheduler_keys_.append(ii)
-    wf_info = wf.query()
+    scheduler_steps = wf.query_step_by_key(scheduler_keys_)
     scheduler_keys = []
-    for ii in scheduler_keys_:
-        if wf_info.get_step(key=ii)[0]["phase"] == "Succeeded":
-            scheduler_keys.append(ii)
+    for step in scheduler_steps:
+        if step["phase"] == "Succeeded":
+            scheduler_keys.append(step.key)
     if len(scheduler_keys) == 0:
         return None
     else:
         skey = sorted(scheduler_keys)[-1]
-        step = wf_info.get_step(key=skey)[0]
+        step = [step for step in scheduler_steps if step.key == skey][0]
         return step.outputs.parameters["exploration_scheduler"].value
 
 


Test patch:
diff --git a/tests/utils/test_dflow_query.py b/tests/utils/test_dflow_query.py
index 34589769..c24a8b7a 100644
--- a/tests/utils/test_dflow_query.py
+++ b/tests/utils/test_dflow_query.py
@@ -122,7 +122,8 @@ def __init__(self):
 
 
 class MockedBar:
-    def __init__(self, xx):
+    def __init__(self, xx, kk):
+        self.key = kk
         self.outputs = MockedFoo()
         self.outputs.parameters["exploration_scheduler"].value = xx * 10
 
@@ -133,11 +134,11 @@ def __getitem__(self, key):
 
 def _get_step(key=None):
     if key == "init--scheduler":
-        return [MockedBar(2)]
+        return [MockedBar(2, key)]
     elif key == "iter-0--scheduler":
-        return [MockedBar(0)]
+        return [MockedBar(0, key)]
     elif key == "iter-1--scheduler":
-        return [MockedBar(1)]
+        return [MockedBar(1, key)]
     else:
         raise RuntimeError("unexpected key")
 
@@ -148,12 +149,29 @@ def get_step(self, key=None):
 
 
 class MockedWF:
+    def __init__(
+        self,
+        none_global=True,
+    ):
+        self.none_global = none_global
+
     def query_step(self, key=None):
         return _get_step(key)
 
     def query(self):
         return MockedWFInfo()
 
+    def query_global_outputs(self):
+        # mocked return None: non-global scheduler output
+        if self.none_global:
+            return None
+        else:
+            return MockedFoo()
+
+    def query_step_by_key(self, keys):
+        ret = [_get_step(kk)[0] for kk in keys]
+        return ret
+
 
 class TestDflowQuery(unittest.TestCase):
     def test_get_subkey(self):
@@ -180,6 +198,13 @@ def test_get_last_scheduler(self):
         )
         self.assertEqual(value, 10)
 
+    def test_get_last_scheduler(self):
+        value = get_last_scheduler(
+            MockedWF(none_global=False),
+            ["iter-1--scheduler", "foo", "bar", "iter-0--scheduler", "init--scheduler"],
+        )
+        self.assertEqual(value, 10)
+
     def test_get_all_schedulers(self):
         value = get_all_schedulers(
             MockedWF(),