support environmental variables in op template (#63)

* support environmental variables in op template * fix issue in ut Co-authored-by: Han Wang <wang_han@iapcm.ac.cn>
deepmodeling · Aug 5, 2022 · c43bb57 · c43bb57
1 parent 1b8041e
commit c43bb57
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 0 deletions.
diff --git a/dpgen2/utils/step_config.py b/dpgen2/utils/step_config.py
@@ -37,11 +37,13 @@ def template_conf_args():
     doc_timeout = 'The time limit of the OP. Unit is second.'
     doc_retry_on_transient_error = 'Retry the step if a TransientError is raised.'
     doc_timeout_as_transient_error = 'Treat the timeout as TransientError.'
+    doc_envs = 'The environmental variables.'
     return [
         Argument("image", str, optional=True, default=default_image, doc=doc_image),
         Argument("timeout", int, optional=True, default=None, doc=doc_timeout),
         Argument("retry_on_transient_error", bool, optional=True, default=None, doc=doc_retry_on_transient_error),
         Argument("timeout_as_transient_error", bool, optional=True, default=False, doc=doc_timeout_as_transient_error),
+        Argument("envs", dict, optional=True, default=None, doc=doc_envs),
     ]
 
 def step_conf_args():

diff --git a/tests/utils/test_step_config.py b/tests/utils/test_step_config.py
@@ -31,6 +31,7 @@ def test_success(self):
                 "timeout" : None,
                 "retry_on_transient_error" : None,
                 "timeout_as_transient_error" : False,
+                "envs" : None,
             },
             "continue_on_failed" : False,
             "continue_on_num_success" : None,
@@ -56,6 +57,7 @@ def test_empty(self):
                 "timeout" : None,
                 "retry_on_transient_error" : None,
                 "timeout_as_transient_error" : False,
+                "envs" : None,
             },
             "continue_on_failed" : False,
             "continue_on_num_success" : None,