Job, possibility to update env vars (#180)

* job startup hook, to update env vars

Fix #178

* set_env more consistent to other methods

* update env directly, instead of using generic startup hook code

* reuse EnvironmentModifier

* use set_vars_verbatim

* set_verbatim method instead

* fix old behavior for DEFAULT_ENVIRONMENT_SET

* set_env verbatim flag

sisyphus/job.py

@@ -65,7 +65,6 @@ def job_finished(path):
 
 
 class JobSingleton(type):
-
     """Meta class to ensure that every Job with the same hash value is
     only created once"""
 
@@ -220,9 +219,8 @@ def _sis_init(self, args, kwargs, parsed_args):
         self._sis_is_finished = False
         self._sis_setup_since_restart = False
 
-        self._sis_environment = None
-        if gs.CLEANUP_ENVIRONMENT:
-            self._sis_environment = tools.EnvironmentModifier()
+        self._sis_environment = tools.EnvironmentModifier(cleanup_env=gs.CLEANUP_ENVIRONMENT)
+        if gs.CLEANUP_ENVIRONMENT:  # for compat, only set those below if CLEANUP_ENVIRONMENT is enabled
             self._sis_environment.keep(gs.DEFAULT_ENVIRONMENT_KEEP)
             self._sis_environment.set(gs.DEFAULT_ENVIRONMENT_SET)
 
@@ -1137,6 +1135,19 @@ def update_rqmt(self, task_name, rqmt):
         self._sis_task_rqmt_overwrite[task_name] = rqmt.copy(), False
         return self
 
+    def set_env(self, key: str, value: str, *, verbatim: bool = True):
+        """
+        Set environment variable. This environment var will be set at job startup in the worker.
+
+        :param key: variable name
+        :param value:
+        :param verbatim: True: set it as-is; False: use string.Template(value).substitute(orig_env)
+        """
+        if verbatim:
+            self._sis_environment.set_verbatim(key, value)
+        else:
+            self._sis_environment.set(key, value)
+
     def tasks(self) -> Iterator[Task]:
         """
         :return: yields Task's

sisyphus/tools.py

@@ -132,7 +132,6 @@ def try_get(v):
 
 
 class execute_in_dir(object):
-
     """Object to be used by the with statement.
     All code after the with will be executed in the given directory,
     working directory will be changed back after with statement.
@@ -156,7 +155,6 @@ def __exit__(self, type, value, traceback):
 
 
 class cache_result(object):
-
     """decorated to cache the result of a function for x_seconds"""
 
     def __init__(self, cache_time=30, force_update=None, clear_cache=None):
@@ -481,42 +479,56 @@ class EnvironmentModifier:
     A class to cleanup the environment before a job starts
     """
 
-    def __init__(self):
+    def __init__(self, *, cleanup_env: bool = True):
+        self.cleanup_env = cleanup_env
         self.keep_vars = set()
         self.set_vars = {}
+        self.set_vars_verbatim = {}  # does not go through string.Template(v).substitute
 
     def keep(self, var):
-        if type(var) == str:
+        if isinstance(var, str):
             self.keep_vars.add(var)
         else:
             self.keep_vars.update(var)
 
     def set(self, var, value=None):
-        if type(var) == dict:
+        if isinstance(var, dict):
             self.set_vars.update(var)
         else:
             self.set_vars[var] = value
 
+    def set_verbatim(self, key: str, value: str):
+        self.set_vars_verbatim[key] = value
+
     def modify_environment(self):
         import os
         import string
 
         orig_env = dict(os.environ)
-        keys = list(os.environ.keys())
-        for k in keys:
-            if k not in self.keep_vars:
-                del os.environ[k]
+        if self.cleanup_env:
+            keys = list(os.environ.keys())
+            for k in keys:
+                if k not in self.keep_vars:
+                    del os.environ[k]
+
         for k, v in self.set_vars.items():
-            if type(v) == str:
+            if isinstance(v, str):
                 os.environ[k] = string.Template(v).substitute(orig_env)
             else:
                 os.environ[k] = str(v)
 
+        for k, v in self.set_vars_verbatim.items():
+            os.environ[k] = v
+
         for k, v in os.environ.items():
             logging.debug("environment var %s=%s" % (k, v))
 
     def __repr__(self):
-        return repr(self.keep_vars) + " " + repr(self.set_vars)
+        return (
+            f"cleanup_env={self.cleanup_env} "
+            + (f"keep={self.keep_vars!r} " if self.cleanup_env else "")
+            + f"set={self.set_vars!r}"
+        )
 
 
 class FinishedResultsCache:

sisyphus/worker.py

@@ -245,7 +245,7 @@ def worker_helper(args):
     gs.active_engine.init_worker(task)
 
     # cleanup environment
-    if hasattr(task._job, "_sis_environment") and task._job._sis_environment:
+    if getattr(task._job, "_sis_environment", None):
         task._job._sis_environment.modify_environment()
 
     try: