Merge remote-tracking branch 'origin/fix-dp-optim-parallel' into fix-…

…dp-optim-parallel
wangzyphysics · Mar 30, 2024 · 23e36db · 23e36db
2 parents 89e6ba7 + fe02139
commit 23e36db
Show file tree

Hide file tree

Showing 6 changed files with 65 additions and 70 deletions.
diff --git a/tests/test_dpgen_loop.py b/tests/test_dpgen_loop.py
@@ -46,9 +46,46 @@
 except ModuleNotFoundError:
     # case of upload everything to argo, no context needed
     pass
+from context import (
+    default_host,
+    default_image,
+    skip_ut_with_dflow,
+    skip_ut_with_dflow_reason,
+    upload_python_packages,
+)
 from dflow.python import (
     FatalError,
 )
+from mocked_ops import (
+    MockedCollectData,
+    MockedCollectDataCheckOptParam,
+    MockedCollectDataFailed,
+    MockedCollectDataRestart,
+    MockedConfSelector,
+    MockedConstTrustLevelStageScheduler,
+    MockedExplorationReport,
+    MockedExplorationTaskGroup,
+    MockedExplorationTaskGroup1,
+    MockedExplorationTaskGroup2,
+    MockedPrepDPTrain,
+    MockedPrepVasp,
+    MockedRunDPTrain,
+    MockedRunDPTrainCheckOptParam,
+    MockedRunLmp,
+    MockedRunVasp,
+    MockedRunVaspFail1,
+    MockedRunVaspRestart,
+    MockedSelectConfs,
+    MockedStage,
+    MockedStage1,
+    MockedStage2,
+    make_mocked_init_data,
+    make_mocked_init_models,
+    mocked_incar_template,
+    mocked_numb_models,
+    mocked_numb_select,
+    mocked_template_script,
+)
 
 from dpgen2.constants import (
     fp_task_pattern,
@@ -98,44 +135,6 @@
 )
 from dpgen2.utils.step_config import normalize as normalize_step_dict
 
-from context import (
-    default_host,
-    default_image,
-    skip_ut_with_dflow,
-    skip_ut_with_dflow_reason,
-    upload_python_packages,
-)
-from mocked_ops import (
-    MockedCollectData,
-    MockedCollectDataCheckOptParam,
-    MockedCollectDataFailed,
-    MockedCollectDataRestart,
-    MockedConfSelector,
-    MockedConstTrustLevelStageScheduler,
-    MockedExplorationReport,
-    MockedExplorationTaskGroup,
-    MockedExplorationTaskGroup1,
-    MockedExplorationTaskGroup2,
-    MockedPrepDPTrain,
-    MockedPrepVasp,
-    MockedRunDPTrain,
-    MockedRunDPTrainCheckOptParam,
-    MockedRunLmp,
-    MockedRunVasp,
-    MockedRunVaspFail1,
-    MockedRunVaspRestart,
-    MockedSelectConfs,
-    MockedStage,
-    MockedStage1,
-    MockedStage2,
-    make_mocked_init_data,
-    make_mocked_init_models,
-    mocked_incar_template,
-    mocked_numb_models,
-    mocked_numb_select,
-    mocked_template_script,
-)
-
 default_config = normalize_step_dict(
     {
         "template_config": {

diff --git a/tests/test_prep_run_caly.py b/tests/test_prep_run_caly.py
@@ -58,6 +58,12 @@
     skip_ut_with_dflow_reason,
     upload_python_packages,
 )
+from mocked_ops import (
+    MockedCollRunCaly,
+    MockedRunCalyModelDevi,
+    MockedRunDPOptim,
+    mocked_numb_models,
+)
 
 from dpgen2.exploration.task import (
     BaseExplorationTaskGroup,
@@ -80,13 +86,6 @@
 )
 from dpgen2.utils.step_config import normalize as normalize_step_dict
 
-from mocked_ops import (
-    MockedCollRunCaly,
-    MockedRunCalyModelDevi,
-    MockedRunDPOptim,
-    mocked_numb_models,
-)
-
 prep_default_config = normalize_step_dict(
     {
         "template_config": {

diff --git a/tests/test_prep_run_dp_train.py b/tests/test_prep_run_dp_train.py
@@ -49,16 +49,6 @@
     skip_ut_with_dflow_reason,
     upload_python_packages,
 )
-
-from dpgen2.constants import (
-    train_task_pattern,
-)
-from dpgen2.superop.prep_run_dp_train import (
-    ModifyTrainScript,
-    PrepRunDPTrain,
-)
-from dpgen2.utils.step_config import normalize as normalize_step_dict
-
 from mocked_ops import (
     MockedModifyTrainScript,
     MockedPrepDPTrain,
@@ -70,6 +60,15 @@
     mocked_template_script,
 )
 
+from dpgen2.constants import (
+    train_task_pattern,
+)
+from dpgen2.superop.prep_run_dp_train import (
+    ModifyTrainScript,
+    PrepRunDPTrain,
+)
+from dpgen2.utils.step_config import normalize as normalize_step_dict
+
 default_config = normalize_step_dict(
     {
         "template_config": {

diff --git a/tests/test_prep_run_lmp.py b/tests/test_prep_run_lmp.py
@@ -51,6 +51,10 @@
     skip_ut_with_dflow_reason,
     upload_python_packages,
 )
+from mocked_ops import (
+    MockedRunLmp,
+    mocked_numb_models,
+)
 
 from dpgen2.constants import (
     lmp_conf_name,
@@ -76,11 +80,6 @@
 )
 from dpgen2.utils.step_config import normalize as normalize_step_dict
 
-from mocked_ops import (
-    MockedRunLmp,
-    mocked_numb_models,
-)
-
 default_config = normalize_step_dict(
     {
         "template_config": {

diff --git a/tests/test_prep_run_vasp.py b/tests/test_prep_run_vasp.py
@@ -50,6 +50,11 @@
     skip_ut_with_dflow_reason,
     upload_python_packages,
 )
+from mocked_ops import (
+    MockedPrepVasp,
+    MockedRunVasp,
+    mocked_incar_template,
+)
 
 from dpgen2.constants import (
     fp_task_pattern,
@@ -65,12 +70,6 @@
 )
 from dpgen2.utils.step_config import normalize as normalize_step_dict
 
-from mocked_ops import (
-    MockedPrepVasp,
-    MockedRunVasp,
-    mocked_incar_template,
-)
-
 default_config = normalize_step_dict(
     {
         "template_config": {

diff --git a/tests/test_select_confs.py b/tests/test_select_confs.py
@@ -45,16 +45,16 @@
 except ModuleNotFoundError:
     # case of upload everything to argo, no context needed
     pass
-from dpgen2.op.select_confs import (
-    SelectConfs,
-)
-
 from mocked_ops import (
     MockedConfSelector,
     MockedExplorationReport,
     MockedSelectConfs,
 )
 
+from dpgen2.op.select_confs import (
+    SelectConfs,
+)
+
 
 class TestMockedSelectConfs(unittest.TestCase):
     def setUp(self):