job startup hook, to update env vars

Fix #178

sisyphus/global_constants.py

@@ -10,6 +10,7 @@
 JOB_FINISHED_MARKER = "finished"
 JOB_FINISHED_ARCHIVE = "finished.tar.gz"
 JOB_INFO = "info"
+JOB_STARTUP_HOOK_PY = "startup_hook.py"
 
 # engine path
 ENGINE_LOG = "log"

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"""
 
@@ -225,6 +224,7 @@ def _sis_init(self, args, kwargs, parsed_args):
             self._sis_environment = tools.EnvironmentModifier()
             self._sis_environment.keep(gs.DEFAULT_ENVIRONMENT_KEEP)
             self._sis_environment.set(gs.DEFAULT_ENVIRONMENT_SET)
+        self._sis_environ_updates = {}
 
         if gs.AUTO_SET_JOB_INIT_ATTRIBUTES:
             self.set_attrs(parsed_args)
@@ -286,6 +286,15 @@ def _sis_setup_directory(self, force=False):
                 if not os.path.isdir(link_name):
                     os.symlink(src=os.path.abspath(str(creator._sis_path())), dst=link_name, target_is_directory=True)
 
+        if self._sis_environ_updates:
+            startup_hook_py_file = self._sis_path(gs.JOB_STARTUP_HOOK_PY)
+            with open(startup_hook_py_file, "w") as f:
+                f.write("import os\n\n")
+                f.write("os.environ.update({\n")
+                for k, v in self._sis_environ_updates.items():
+                    f.write(f"    {k!r}: {v!r},\n")
+                f.write("})\n\n")
+
         # export the actual job
         with gzip.open(self._sis_path(gs.JOB_SAVE), "w") as f:
             pickle.dump(self, f)
@@ -1137,6 +1146,10 @@ def update_rqmt(self, task_name, rqmt):
         self._sis_task_rqmt_overwrite[task_name] = rqmt.copy(), False
         return self
 
+    def putenv(self, key: str, value: str):
+        """this environment var will be set at job startup"""
+        self._sis_environ_updates[key] = value
+
     def tasks(self) -> Iterator[Task]:
         """
         :return: yields Task's

sisyphus/worker.py

@@ -248,6 +248,13 @@ def worker_helper(args):
     if hasattr(task._job, "_sis_environment") and task._job._sis_environment:
         task._job._sis_environment.modify_environment()
 
+    startup_hook_py_file = args.jobdir + os.path.sep + gs.JOB_STARTUP_HOOK_PY
+    if os.path.exists(startup_hook_py_file):
+        source = open(startup_hook_py_file).read()
+        co = compile(source, startup_hook_py_file, "exec")
+        user_ns = {"__file__": startup_hook_py_file, "__name__": startup_hook_py_file, "task": task, "job": job}
+        eval(co, user_ns, user_ns)
+
     try:
         # run task
         task.run(task_id, resume_job, logging_thread=logging_thread)