Repository: deepmodeling/dpgen2 (language: python)

Issue:
Bug: exceeded workflow templates limited 
The scheduler is a class with large amount of data. Passing scheduler as `dflow.Parameter` causes an error like 
```
workflow templates are limited to 128KB, this workflow is 172157 bytes
```

Fast fixing: encode the class in a file and pass the scheduler as a `dflow.Artifact`.



Discussion / hints:

Fix patch:
diff --git a/dpgen2/flow/dpgen_loop.py b/dpgen2/flow/dpgen_loop.py
index d8654737..2481f8ab 100644
--- a/dpgen2/flow/dpgen_loop.py
+++ b/dpgen2/flow/dpgen_loop.py
@@ -39,7 +39,7 @@ class SchedulerWrapper(OP):
     @classmethod
     def get_input_sign(cls):
         return OPIOSign({
-            "exploration_scheduler" : ExplorationScheduler,
+            "exploration_scheduler" : Artifact(Path),
             "exploration_report": ExplorationReport,
             "trajs": Artifact(List[Path]),
         })
@@ -47,7 +47,7 @@ def get_input_sign(cls):
     @classmethod
     def get_output_sign(cls):
         return OPIOSign({
-            "exploration_scheduler" : ExplorationScheduler,
+            "exploration_scheduler" : Artifact(Path),
             "converged" : bool,
             "lmp_task_grp" : Artifact(Path),
             "conf_selector" : ConfSelector,
@@ -58,20 +58,28 @@ def execute(
             self,
             ip : OPIO,
     ) -> OPIO:
-        scheduler = ip['exploration_scheduler']
+        scheduler_in = ip['exploration_scheduler']
         report = ip['exploration_report']
         trajs = ip['trajs']
+        lmp_task_grp_file = Path('lmp_task_grp.dat')
+        scheduler_file = Path('scheduler.dat')
+
+        with open(scheduler_in, 'rb') as fp:
+            scheduler = pickle.load(fp)
 
         conv, lmp_task_grp, selector = scheduler.plan_next_iteration(report, trajs)
 
-        with open('lmp_task_grp.dat', 'wb') as fp:
+        with open(lmp_task_grp_file, 'wb') as fp:
             pickle.dump(lmp_task_grp, fp)
 
+        with open(scheduler_file, 'wb') as fp:
+            pickle.dump(scheduler, fp)
+
         return OPIO({
-            "exploration_scheduler" : scheduler,
+            "exploration_scheduler" : scheduler_file,
             "converged" : conv,
             "conf_selector" : selector,
-            "lmp_task_grp" : Path('lmp_task_grp.dat'),
+            "lmp_task_grp" : lmp_task_grp_file,
         })
 
 
@@ -79,7 +87,7 @@ class MakeBlockId(OP):
     @classmethod
     def get_input_sign(cls):
         return OPIOSign({
-            "exploration_scheduler" : ExplorationScheduler,
+            "exploration_scheduler" : Artifact(Path),
         })
 
     @classmethod
@@ -93,8 +101,11 @@ def execute(
             self,
             ip : OPIO,
     ) -> OPIO:
-        scheduler = ip['exploration_scheduler']
+        scheduler_in = ip['exploration_scheduler']
         
+        with open(scheduler_in, 'rb') as fp:
+            scheduler = pickle.load(fp)
+
         stage = scheduler.get_stage()
         iteration = scheduler.get_iteration()
 
@@ -121,18 +132,18 @@ def __init__(
             "conf_selector" : InputParameter(),
             "fp_inputs" : InputParameter(),
             "fp_config" : InputParameter(),
-            "exploration_scheduler" : InputParameter(),
         }
         self._input_artifacts={
+            "exploration_scheduler" : InputArtifact(),
             "init_models" : InputArtifact(),
             "init_data" : InputArtifact(),
             "iter_data" : InputArtifact(),
             "lmp_task_grp" : InputArtifact(),
         }
         self._output_parameters={
-            "exploration_scheduler": OutputParameter(),
         }
         self._output_artifacts={
+            "exploration_scheduler": OutputArtifact(),
             "models": OutputArtifact(),
             "iter_data" : OutputArtifact(),
         }
@@ -213,17 +224,17 @@ def __init__(
             "lmp_config" : InputParameter(),
             "fp_inputs" : InputParameter(),
             "fp_config" : InputParameter(),
-            "exploration_scheduler" : InputParameter(),
         }
         self._input_artifacts={
+            "exploration_scheduler" : InputArtifact(),
             "init_models" : InputArtifact(),
             "init_data" : InputArtifact(),
             "iter_data" : InputArtifact(),
         }
         self._output_parameters={
-            "exploration_scheduler": OutputParameter(),
         }
         self._output_artifacts={
+            "exploration_scheduler": OutputArtifact(),
             "models": OutputArtifact(),
             "iter_data" : OutputArtifact(),
         }        
@@ -321,10 +332,10 @@ def _loop (
             python_packages = upload_python_package,
         ),
         parameters={
-            "exploration_scheduler": steps.inputs.parameters['exploration_scheduler'],
             "exploration_report": block_step.outputs.parameters['exploration_report'],
         },
         artifacts={
+            "exploration_scheduler": steps.inputs.artifacts['exploration_scheduler'],
             "trajs" : block_step.outputs.artifacts['trajs'],
         },
         key = step_keys['scheduler'],
@@ -339,9 +350,9 @@ def _loop (
             python_packages = upload_python_package,
         ),
         parameters={
-            "exploration_scheduler": scheduler_step.outputs.parameters['exploration_scheduler'],
         },
         artifacts={
+            "exploration_scheduler": scheduler_step.outputs.artifacts['exploration_scheduler'],
         },
         key = step_keys['id'],
     )
@@ -360,9 +371,9 @@ def _loop (
             "conf_selector" : scheduler_step.outputs.parameters["conf_selector"],
             "fp_inputs" : steps.inputs.parameters["fp_inputs"],
             "fp_config" : steps.inputs.parameters["fp_config"],
-            "exploration_scheduler" : scheduler_step.outputs.parameters["exploration_scheduler"],
         },
         artifacts={
+            "exploration_scheduler" : scheduler_step.outputs.artifacts["exploration_scheduler"],
             "lmp_task_grp" : scheduler_step.outputs.artifacts["lmp_task_grp"],
             "init_models" : block_step.outputs.artifacts['models'],
             "init_data" : steps.inputs.artifacts['init_data'],
@@ -372,11 +383,11 @@ def _loop (
     )
     steps.add(next_step)    
 
-    steps.outputs.parameters['exploration_scheduler'].value_from_expression = \
+    steps.outputs.artifacts['exploration_scheduler'].from_expression = \
         if_expression(
             _if = (scheduler_step.outputs.parameters['converged'] == True),
-            _then = scheduler_step.outputs.parameters['exploration_scheduler'],
-            _else = next_step.outputs.parameters['exploration_scheduler'],
+            _then = scheduler_step.outputs.artifacts['exploration_scheduler'],
+            _else = next_step.outputs.artifacts['exploration_scheduler'],
         )
     steps.outputs.artifacts['models'].from_expression = \
         if_expression(
@@ -411,10 +422,10 @@ def _dpgen(
             python_packages = upload_python_package,
         ),
         parameters={
-            "exploration_scheduler": steps.inputs.parameters['exploration_scheduler'],
             "exploration_report": None,
         },
         artifacts={
+            "exploration_scheduler": steps.inputs.artifacts['exploration_scheduler'],
             "trajs" : None,
         },
         key = step_keys['scheduler'],
@@ -429,9 +440,9 @@ def _dpgen(
             python_packages = upload_python_package,
         ),
         parameters={
-            "exploration_scheduler": scheduler_step.outputs.parameters['exploration_scheduler'],
         },
         artifacts={
+            "exploration_scheduler": scheduler_step.outputs.artifacts['exploration_scheduler'],
         },
         key = step_keys['id'],
     )
@@ -450,9 +461,9 @@ def _dpgen(
             "lmp_config" : steps.inputs.parameters['lmp_config'],
             "fp_inputs" : steps.inputs.parameters['fp_inputs'],
             "fp_config" : steps.inputs.parameters['fp_config'],
-            "exploration_scheduler" : scheduler_step.outputs.parameters['exploration_scheduler'],
         },
         artifacts={
+            "exploration_scheduler" : scheduler_step.outputs.artifacts['exploration_scheduler'],
             "lmp_task_grp" : scheduler_step.outputs.artifacts['lmp_task_grp'],
             "init_models": steps.inputs.artifacts["init_models"],
             "init_data": steps.inputs.artifacts["init_data"],
@@ -462,8 +473,8 @@ def _dpgen(
     )
     steps.add(loop_step)
 
-    steps.outputs.parameters["exploration_scheduler"].value_from_parameter = \
-        loop_step.outputs.parameters["exploration_scheduler"]
+    steps.outputs.artifacts["exploration_scheduler"]._from = \
+        loop_step.outputs.artifacts["exploration_scheduler"]
     steps.outputs.artifacts["models"]._from = \
         loop_step.outputs.artifacts["models"]
     steps.outputs.artifacts["iter_data"]._from = \


Test patch:
diff --git a/tests/test_dpgen_loop.py b/tests/test_dpgen_loop.py
index 773278cc..c8503442 100644
--- a/tests/test_dpgen_loop.py
+++ b/tests/test_dpgen_loop.py
@@ -1,4 +1,4 @@
-import os, textwrap
+import os, textwrap, pickle
 import numpy as np
 import unittest
 
@@ -203,7 +203,9 @@ def _setUp_data(self):
             max_numb_iter = 2,
         )
         self.scheduler.add_stage_scheduler(stage_scheduler)        
-        
+        with open('in_scheduler.dat', 'wb') as fp:
+            pickle.dump(self.scheduler, fp)
+        self.scheduler_artifact = upload_artifact('in_scheduler.dat')
         
     def setUp(self):
         self.name = 'dpgen'
@@ -219,7 +221,7 @@ def tearDown(self):
             name = Path(model_name_pattern % ii)
             if name.is_file():
                 os.remove(name)
-        for ii in [self.incar, self.potcar]:
+        for ii in [self.incar, self.potcar, Path('scheduler.dat'), Path('in_scheduler.dat')]:
             if ii.is_file():
                 os.remove(ii)
 
@@ -251,9 +253,9 @@ def test(self):
                 "lmp_config" : {},
                 'fp_inputs' : self.vasp_inputs,
                 "fp_config" : {},
-                "exploration_scheduler" : self.scheduler,
             },
             artifacts = {
+                "exploration_scheduler" : self.scheduler_artifact,
                 "init_models" : self.init_models,
                 "init_data" : self.init_data,
                 "iter_data" : self.iter_data,
@@ -270,7 +272,10 @@ def test(self):
         step = wf.query_step(name='dpgen-step')[0]
         self.assertEqual(step.phase, "Succeeded")        
         
-        scheduler = jsonpickle.decode(step.outputs.parameters['exploration_scheduler'].value)
+        # scheduler = jsonpickle.decode(step.outputs.parameters['exploration_scheduler'].value)
+        download_artifact(step.outputs.artifacts["exploration_scheduler"])
+        with open('scheduler.dat', 'rb') as fp:
+            scheduler = pickle.load(fp)
         download_artifact(step.outputs.artifacts["iter_data"], path = 'iter_data')
         download_artifact(step.outputs.artifacts["models"], path = Path('models')/self.name)
         self.assertEqual(scheduler.get_stage(), 2)
@@ -481,8 +486,10 @@ def _setUp_data(self):
             conv_accuracy = 0.7,
             max_numb_iter = 2,
         )
-        self.scheduler_0.add_stage_scheduler(stage_scheduler)        
-
+        self.scheduler_0.add_stage_scheduler(stage_scheduler)
+        with open('in_scheduler_0.dat', 'wb') as fp:
+            pickle.dump(self.scheduler_0, fp)
+        self.scheduler_0_artifact = upload_artifact('in_scheduler_0.dat')
 
         self.scheduler_1 = ExplorationScheduler()        
         self.trust_level = TrustLevel(0.1, 0.3)
@@ -502,6 +509,10 @@ def _setUp_data(self):
             max_numb_iter = 2,
         )
         self.scheduler_1.add_stage_scheduler(stage_scheduler)        
+        self.scheduler_1.add_stage_scheduler(stage_scheduler)
+        with open('in_scheduler_1.dat', 'wb') as fp:
+            pickle.dump(self.scheduler_1, fp)
+        self.scheduler_1_artifact = upload_artifact('in_scheduler_1.dat')
         
         
     def setUp(self):
@@ -521,6 +532,15 @@ def tearDown(self):
         for ii in [self.incar, self.potcar]:
             if ii.is_file():
                 os.remove(ii)
+        for ii in [Path('scheduler_0.dat'), Path('in_scheduler_0.dat')]:
+            if ii.is_file():
+                os.remove(ii)
+        for ii in [Path('scheduler_1.dat'), Path('in_scheduler_1.dat')]:
+            if ii.is_file():
+                os.remove(ii)
+        for ii in [Path('scheduler.dat'), Path('scheduler_new.dat')]:
+            if ii.is_file():
+                os.remove(ii)
 
     def get_restart_step(
             self,
@@ -551,9 +571,9 @@ def test_update_artifact(self):
                 "lmp_config" : {},
                 'fp_inputs' : self.vasp_inputs,
                 "fp_config" : {},
-                "exploration_scheduler" : self.scheduler_0,
             },
             artifacts = {
+                "exploration_scheduler" : self.scheduler_0_artifact,
                 "init_models" : self.init_models,
                 "init_data" : self.init_data,
                 "iter_data" : self.iter_data,
@@ -623,9 +643,9 @@ def test_update_artifact(self):
                 "lmp_config" : {},
                 'fp_inputs' : self.vasp_inputs,
                 "fp_config" : {},
-                "exploration_scheduler" : self.scheduler_1,
             },
             artifacts = {
+                "exploration_scheduler" : self.scheduler_1_artifact,
                 "init_models" : self.init_models,
                 "init_data" : self.init_data,
                 "iter_data" : self.iter_data,
@@ -645,7 +665,10 @@ def test_update_artifact(self):
         self.assertEqual(step.phase, "Succeeded")
         download_artifact(step.outputs.artifacts["iter_data"], path = 'iter_data')
         download_artifact(step.outputs.artifacts["models"], path = Path('models')/self.name)
-        scheduler = jsonpickle.decode(step.outputs.parameters['exploration_scheduler'].value)
+        # scheduler = jsonpickle.decode(step.outputs.parameters['exploration_scheduler'].value)
+        download_artifact(step.outputs.artifacts['exploration_scheduler'])
+        with open('scheduler.dat', 'rb') as fp:
+            scheduler = pickle.load(fp)
         self.assertEqual(scheduler.get_stage(), 2)
         self.assertEqual(scheduler.get_iteration(), 1)
         
@@ -695,9 +718,9 @@ def test_update_artifact_scheduler(self):
                 "lmp_config" : {},
                 'fp_inputs' : self.vasp_inputs,
                 "fp_config" : {},
-                "exploration_scheduler" : self.scheduler_0,
             },
             artifacts = {
+                "exploration_scheduler" : self.scheduler_0_artifact,
                 "init_models" : self.init_models,
                 "init_data" : self.init_data,
                 "iter_data" : self.iter_data,
@@ -755,11 +778,18 @@ def test_update_artifact_scheduler(self):
         self.assertTrue(scheduler_idx is not None)
         step_scheduler = steps_0.pop(scheduler_idx)
         self.assertEqual(step_scheduler['phase'], 'Succeeded')
-        old_scheduler = jsonpickle.decode(step_scheduler.outputs.parameters['exploration_scheduler'].value)
+        # old_scheduler = jsonpickle.decode(step_scheduler.outputs.parameters['exploration_scheduler'].value)
+        download_artifact(step_scheduler.outputs.artifacts['exploration_scheduler'])
+        with open('scheduler.dat', 'rb') as fp:
+            old_scheduler = pickle.load(fp)
         self.assertEqual(old_scheduler.get_stage(), 0)
         # update a stage scheduler
         old_scheduler.stage_schedulers[1] = self.scheduler_1.stage_schedulers[1]        
-        step_scheduler.modify_output_parameter("exploration_scheduler", old_scheduler)
+        # step_scheduler.modify_output_parameter("exploration_scheduler", old_scheduler)
+        with open('scheduler_new.dat', 'wb') as fp:
+            pickle.dump(old_scheduler, fp)
+        old_scheduler_artifact = upload_artifact('scheduler_new.dat')
+        step_scheduler.modify_output_artifact("exploration_scheduler", old_scheduler_artifact)
         steps_0.append(step_scheduler)
 
         dpgen_step_1 = Step(
@@ -773,9 +803,9 @@ def test_update_artifact_scheduler(self):
                 "lmp_config" : {},
                 'fp_inputs' : self.vasp_inputs,
                 "fp_config" : {},
-                "exploration_scheduler" : self.scheduler_0,
             },
             artifacts = {
+                "exploration_scheduler" : self.scheduler_0_artifact,
                 "init_models" : self.init_models,
                 "init_data" : self.init_data,
                 "iter_data" : self.iter_data,
@@ -795,7 +825,10 @@ def test_update_artifact_scheduler(self):
         self.assertEqual(step.phase, "Succeeded")
         download_artifact(step.outputs.artifacts["iter_data"], path = 'iter_data')
         download_artifact(step.outputs.artifacts["models"], path = Path('models')/self.name)
-        scheduler = jsonpickle.decode(step.outputs.parameters['exploration_scheduler'].value)
+        # scheduler = jsonpickle.decode(step.outputs.parameters['exploration_scheduler'].value)
+        download_artifact(step.outputs.artifacts['exploration_scheduler'])
+        with open('scheduler.dat', 'rb') as fp:
+            scheduler = pickle.load(fp)
         self.assertEqual(scheduler.get_stage(), 2)
         self.assertEqual(scheduler.get_iteration(), 1)
         
@@ -844,9 +877,9 @@ def test_update_slice_item_input(self):
                 "lmp_config" : {},
                 'fp_inputs' : self.vasp_inputs,
                 "fp_config" : {},
-                "exploration_scheduler" : self.scheduler_0,
             },
             artifacts = {
+                "exploration_scheduler" : self.scheduler_0_artifact,
                 "init_models" : self.init_models,
                 "init_data" : self.init_data,
                 "iter_data" : self.iter_data,
@@ -912,9 +945,9 @@ def test_update_slice_item_input(self):
                 "lmp_config" : {},
                 'fp_inputs' : self.vasp_inputs,
                 "fp_config" : {},
-                "exploration_scheduler" : self.scheduler_0,
             },
             artifacts = {
+                "exploration_scheduler" : self.scheduler_0_artifact,
                 "init_models" : self.init_models,
                 "init_data" : self.init_data,
                 "iter_data" : self.iter_data,
@@ -934,7 +967,10 @@ def test_update_slice_item_input(self):
         self.assertEqual(step.phase, "Succeeded")
         download_artifact(step.outputs.artifacts["iter_data"], path = 'iter_data')
         download_artifact(step.outputs.artifacts["models"], path = Path('models')/self.name)
-        scheduler = jsonpickle.decode(step.outputs.parameters['exploration_scheduler'].value)
+        # scheduler = jsonpickle.decode(step.outputs.parameters['exploration_scheduler'].value)
+        download_artifact(step.outputs.artifacts['exploration_scheduler'])
+        with open('scheduler.dat', 'rb') as fp:
+            scheduler = pickle.load(fp)
         self.assertEqual(scheduler.get_stage(), 2)
         self.assertEqual(scheduler.get_iteration(), 1)
         
@@ -995,9 +1031,9 @@ def test_update_slice_item_output(self):
                 "lmp_config" : {},
                 'fp_inputs' : self.vasp_inputs,
                 "fp_config" : {},
-                "exploration_scheduler" : self.scheduler_0,
             },
             artifacts = {
+                "exploration_scheduler" : self.scheduler_0_artifact,
                 "init_models" : self.init_models,
                 "init_data" : self.init_data,
                 "iter_data" : self.iter_data,
@@ -1061,9 +1097,9 @@ def test_update_slice_item_output(self):
                 "lmp_config" : {},
                 'fp_inputs' : self.vasp_inputs,
                 "fp_config" : {},
-                "exploration_scheduler" : self.scheduler_0,
             },
             artifacts = {
+                "exploration_scheduler" : self.scheduler_0_artifact,
                 "init_models" : self.init_models,
                 "init_data" : self.init_data,
                 "iter_data" : self.iter_data,
@@ -1083,7 +1119,10 @@ def test_update_slice_item_output(self):
         self.assertEqual(step.phase, "Succeeded")
         download_artifact(step.outputs.artifacts["iter_data"], path = 'iter_data')
         download_artifact(step.outputs.artifacts["models"], path = Path('models')/self.name)
-        scheduler = jsonpickle.decode(step.outputs.parameters['exploration_scheduler'].value)
+        # scheduler = jsonpickle.decode(step.outputs.parameters['exploration_scheduler'].value)
+        download_artifact(step.outputs.artifacts['exploration_scheduler'])
+        with open('scheduler.dat', 'rb') as fp:
+            scheduler = pickle.load(fp)
         self.assertEqual(scheduler.get_stage(), 2)
         self.assertEqual(scheduler.get_iteration(), 1)
         
@@ -1146,9 +1185,9 @@ def test_update_slice_item_output_1(self):
                 "lmp_config" : {},
                 'fp_inputs' : self.vasp_inputs,
                 "fp_config" : {},
-                "exploration_scheduler" : self.scheduler_0,
             },
             artifacts = {
+                "exploration_scheduler" : self.scheduler_0_artifact,
                 "init_models" : self.init_models,
                 "init_data" : self.init_data,
                 "iter_data" : self.iter_data,
@@ -1211,9 +1250,9 @@ def test_update_slice_item_output_1(self):
                 "lmp_config" : {},
                 'fp_inputs' : self.vasp_inputs,
                 "fp_config" : {},
-                "exploration_scheduler" : self.scheduler_0,
             },
             artifacts = {
+                "exploration_scheduler" : self.scheduler_0_artifact,
                 "init_models" : self.init_models,
                 "init_data" : self.init_data,
                 "iter_data" : self.iter_data,
@@ -1233,7 +1272,10 @@ def test_update_slice_item_output_1(self):
         self.assertEqual(step.phase, "Succeeded")
         download_artifact(step.outputs.artifacts["iter_data"], path = 'iter_data')
         download_artifact(step.outputs.artifacts["models"], path = Path('models')/self.name)
-        scheduler = jsonpickle.decode(step.outputs.parameters['exploration_scheduler'].value)
+        # scheduler = jsonpickle.decode(step.outputs.parameters['exploration_scheduler'].value)
+        download_artifact(step.outputs.artifacts['exploration_scheduler'])
+        with open('scheduler.dat', 'rb') as fp:
+            scheduler = pickle.load(fp)
         self.assertEqual(scheduler.get_stage(), 2)
         self.assertEqual(scheduler.get_iteration(), 1)