add HDF5Datasets format of trajs to exploration scheduler

Signed-off-by: zjgemi <[email protected]>
deepmodeling · Sep 6, 2024 · ee022f7 · ee022f7
1 parent 687b9c5
commit ee022f7
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 6 deletions.
diff --git a/dpgen2/exploration/scheduler/convergence_check_stage_scheduler.py b/dpgen2/exploration/scheduler/convergence_check_stage_scheduler.py
@@ -5,11 +5,15 @@
     List,
     Optional,
     Tuple,
+    Union,
 )
 
 from dflow.python import (
     FatalError,
 )
+from dflow.python.opio import (
+    HDF5Dataset,
+)
 
 from dpgen2.exploration.report import (
     ExplorationReport,
@@ -67,7 +71,7 @@ def reached_max_iteration(self):
     def plan_next_iteration(
         self,
         report: Optional[ExplorationReport] = None,
-        trajs: Optional[List[Path]] = None,
+        trajs: Optional[Union[List[Path], List[HDF5Dataset]]] = None,
     ) -> Tuple[bool, Optional[BaseExplorationTaskGroup], Optional[ConfSelector]]:
         if self.complete():
             raise FatalError("Cannot plan because the stage has completed.")

diff --git a/dpgen2/exploration/scheduler/scheduler.py b/dpgen2/exploration/scheduler/scheduler.py
@@ -5,12 +5,16 @@
     List,
     Optional,
     Tuple,
+    Union,
 )
 
 import numpy as np
 from dflow.python import (
     FatalError,
 )
+from dflow.python.opio import (
+    HDF5Dataset,
+)
 
 from dpgen2.exploration.report import (
     ExplorationReport,
@@ -110,7 +114,7 @@ def force_stage_complete(self):
     def plan_next_iteration(
         self,
         report: Optional[ExplorationReport] = None,
-        trajs: Optional[List[Path]] = None,
+        trajs: Optional[Union[List[Path], List[HDF5Dataset]]] = None,
     ) -> Tuple[bool, Optional[ExplorationTaskGroup], Optional[ConfSelector]]:
         """
         Make the plan for the next DPGEN iteration.
@@ -119,7 +123,7 @@ def plan_next_iteration(
         ----------
         report : ExplorationReport
             The exploration report of this iteration.
-        trajs : List[Path]
+        trajs : Union[List[Path], List[HDF5Dataset]]
             A list of configurations generated during the exploration. May be used to generate new configurations for the next iteration.
 
         Returns

diff --git a/dpgen2/exploration/scheduler/stage_scheduler.py b/dpgen2/exploration/scheduler/stage_scheduler.py
@@ -8,8 +8,12 @@
 from typing import (
     List,
     Tuple,
+    Union,
 )
 
+from dflow.python.opio import (
+    HDF5Dataset,
+)
 from dpgen2.exploration.report import (
     ExplorationReport,
 )
@@ -87,7 +91,7 @@ def get_reports(self) -> List[ExplorationReport]:
     def plan_next_iteration(
         self,
         report: ExplorationReport,
-        trajs: List[Path],
+        trajs: Union[List[Path], List[HDF5Dataset]],
     ) -> Tuple[bool, ExplorationTaskGroup, ConfSelector]:
         """
         Make the plan for the next iteration of the stage.
@@ -100,7 +104,7 @@ def plan_next_iteration(
             The historical exploration report of the stage. If this is the first iteration of the stage, this list is empty.
         report : ExplorationReport
             The exploration report of this iteration.
-        confs : List[Path]
+        trajs: Union[List[Path], List[HDF5Dataset]]
             A list of configurations generated during the exploration. May be used to generate new configurations for the next iteration.
 
         Returns

diff --git a/dpgen2/flow/dpgen_loop.py b/dpgen2/flow/dpgen_loop.py
@@ -9,6 +9,7 @@
 from typing import (
     List,
     Optional,
+    Union,
 )
 
 import jsonpickle
@@ -35,6 +36,7 @@
     OPIO,
     Artifact,
     BigParameter,
+    HDF5Datasets,
     OPIOSign,
     PythonOPTemplate,
     Slices,
@@ -91,7 +93,7 @@ def get_input_sign(cls):
             {
                 "exploration_scheduler": BigParameter(ExplorationScheduler),
                 "exploration_report": BigParameter(ExplorationReport),
-                "trajs": Artifact(List[Path]),
+                "trajs": Artifact(Union[List[Path], HDF5Datasets]),
             }
         )