store worker wrapper in Job (#184)

sisyphus/job.py

@@ -210,6 +210,7 @@ def _sis_init(self, args, kwargs, parsed_args):
         self._sis_outputs = {}
         self._sis_keep_value = None
         self._sis_hold_job = False
+        self._sis_worker_wrapper = gs.worker_wrapper
 
         self._sis_blocks = set()
         self._sis_kwargs = parsed_args
@@ -314,6 +315,7 @@ def __getstate__(self):
             "current_block",
             "_sis_cleanable_cache",
             "_sis_cleaned_or_not_cleanable",
+            "_sis_worker_wrapper",
         ]:
             if key in d:
                 del d[key]

sisyphus/task.py

@@ -481,5 +481,9 @@ def get_worker_call(self, task_id=None):
         call += [gs.CMD_WORKER, os.path.relpath(self.path()), self.name()]
         if task_id is not None:
             call.append(str(task_id))
-        call = gs.worker_wrapper(getattr(self, "_job", None), self.name(), call)
+        if hasattr(self, "_job"):
+            call = self._job._sis_worker_wrapper(self._job, self.name(), call)
+        else:
+            logging.warning(f"Task {self.name()} run without an associated Job. Using global worker_wrapper.")
+            call = gs.worker_wrapper(None, self.name(), call)
         return call