Repository: spotify/luigi

Issue:
Workers dying on errors that can't be unpickled
If I have a task that raises an error on run that can't be unpickled, it kills the whole worker immediately. Example code:

```
import luigi

class MyError(BaseException):
    def __init__(self, a, b):
        pass


class MyTask(luigi.Task):
    def run(self):
        raise MyError(1, 2)


if __name__ == '__main__':
    luigi.run()
```

If I run this with the latest master, the whole worker dies on the error with the following traceback:

```
Traceback (most recent call last):
  File "test.py", line 14, in <module>
    luigi.run()
  File "/Users/buck/workspace/luigi/luigi/interface.py", line 331, in run
    return _schedule_and_run(tasks, worker_scheduler_factory)
  File "/Users/buck/workspace/luigi/luigi/interface.py", line 184, in _schedule_and_run
    success &= w.run()
  File "/Users/buck/workspace/luigi/luigi/worker.py", line 790, in run
    self._handle_next_task()
  File "/Users/buck/workspace/luigi/luigi/worker.py", line 690, in _handle_next_task
    timeout=self._config.wait_interval))
  File "/usr/local/Cellar/python/2.7.10/Frameworks/Python.framework/Versions/2.7/lib/python2.7/multiprocessing/queues.py", line 135, in get
    res = self._recv()
TypeError: ('__init__() takes exactly 3 arguments (1 given)', <class '__main__.MyError'>, ())
```

whereas if I run with older code, it recovers from the failure and carries on. The problem seems to be that the exception is now put into the result queue and cannot be successfully unpickled when it's retrieved from the queue in the parent process. This is the same error you get if you pickle and try to unpickle the error.



Discussion:
Nice catch. This must have been introduced in #1205 

We should add a test for this and possibly revert right away.

It's weird that even that simple exception class can't be unpickled. Are we doing something wrong with how we pickle?

What??? Pickle doesn't support exceptions?

```
>>> import pickle
>>> class MyException(BaseException):
...    def __init__(self, a): self.a = a
... 
>>> a = pickle.dumps(MyException(42))
>>> a
"c__main__\nMyException\np0\n(tRp1\n(dp2\nS'a'\np3\nI42\nsb."
>>> pickle.loads(a)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/System/Library/Frameworks/Python.framework/Versions/2.7/lib/python2.7/pickle.py", line 1382, in loads
    return Unpickler(file).load()
  File "/System/Library/Frameworks/Python.framework/Versions/2.7/lib/python2.7/pickle.py", line 858, in load
    dispatch[key](self)
  File "/System/Library/Frameworks/Python.framework/Versions/2.7/lib/python2.7/pickle.py", line 1133, in load_reduce
    value = func(*args)
TypeError: __init__() takes exactly 2 arguments (1 given)
```

Seems like this is an old bug: http://bugs.python.org/issue1692335

Either way let's just revert #1205?

It was a fairly easy catch for me, as my workers don't last very long before something throws a custom exception. Reverting #1205 makes sense, as this is a pretty serious bug.


Resolving patch:
diff --git a/luigi/worker.py b/luigi/worker.py
--- a/luigi/worker.py
+++ b/luigi/worker.py
@@ -128,8 +128,6 @@ def run(self):
         expl = ''
         missing = []
         new_deps = []
-        exception = None
-
         try:
             # Verify that all the tasks are fulfilled!
             missing = [dep.task_id for dep in self.task.deps() if not dep.complete()]
@@ -169,19 +167,16 @@ def run(self):
             logger.exception("[pid %s] Worker %s failed    %s", os.getpid(), self.worker_id, self.task)
             self.task.trigger_event(Event.FAILURE, self.task, ex)
             subject = "Luigi: %s FAILED" % self.task
-            exception = ex
 
             raw_error_message = self.task.on_failure(ex)
             notification_error_message = notifications.wrap_traceback(raw_error_message)
             expl = json.dumps(raw_error_message)
-
             formatted_error_message = notifications.format_task_error(subject, self.task,
                                                                       formatted_exception=notification_error_message)
             notifications.send_error_email(subject, formatted_error_message, self.task.owner_email)
         finally:
-            # multiprocessing.Queue will pickle it
-            self.result_queue.put((
-                self.task.task_id, status, expl, missing, new_deps, exception))
+            self.result_queue.put(
+                (self.task.task_id, status, expl, missing, new_deps))
 
     def _recursive_terminate(self):
         import psutil
@@ -331,7 +326,7 @@ class Worker(object):
     * asks for stuff to do (pulls it in a loop and runs it)
     """
 
-    def __init__(self, scheduler=None, worker_id=None, worker_processes=1, assistant=False, raise_on_error=False, **kwargs):
+    def __init__(self, scheduler=None, worker_id=None, worker_processes=1, assistant=False, **kwargs):
         if scheduler is None:
             scheduler = CentralPlannerScheduler()
 
@@ -374,9 +369,6 @@ def __init__(self, scheduler=None, worker_id=None, worker_processes=1, assistant
         self._add_task_history = []
         self._get_work_response_history = []
 
-        # whether to raise when a task throws an exception
-        self._raise_on_error = raise_on_error
-
     def _add_task(self, *args, **kwargs):
         """
         Call ``self._scheduler.add_task``, but store the values too so we can
@@ -669,8 +661,7 @@ def _purge_children(self):
                 continue
 
             logger.info(error_msg)
-            self._task_result_queue.put(
-                (task_id, FAILED, error_msg, [], [], None))
+            self._task_result_queue.put((task_id, FAILED, error_msg, [], []))
 
     def _handle_next_task(self):
         """
@@ -685,7 +676,7 @@ def _handle_next_task(self):
             self._purge_children()  # Deal with subprocess failures
 
             try:
-                task_id, status, expl, missing, new_requirements, exception = (
+                task_id, status, expl, missing, new_requirements = (
                     self._task_result_queue.get(
                         timeout=self._config.wait_interval))
             except Queue.Empty:
@@ -720,9 +711,6 @@ def _handle_next_task(self):
                 continue
             self._running_tasks.pop(task_id)
 
-            if status == FAILED and exception and self._raise_on_error:
-                raise exception
-
             # re-add task to reschedule missing dependencies
             if missing:
                 reschedule = True


Test patch:
diff --git a/test/worker_task_test.py b/test/worker_task_test.py
--- a/test/worker_task_test.py
+++ b/test/worker_task_test.py
@@ -48,12 +48,9 @@ def on_success(self):
 
 
 class FailTask(luigi.Task):
-    def __init__(self):
-        super(FailTask, self).__init__()
-        self.ex = BaseException("Uh oh.")
 
     def run(self):
-        raise self.ex
+        raise BaseException("Uh oh.")
 
     def on_failure(self, exception):
         return "test failure expl"
@@ -88,7 +85,7 @@ def test_update_result_queue_on_success(self):
 
         with mock.patch.object(result_queue, 'put') as mock_put:
             task_process.run()
-            mock_put.assert_called_once_with((task.task_id, DONE, json.dumps("test success expl"), [], None, None))
+            mock_put.assert_called_once_with((task.task_id, DONE, json.dumps("test success expl"), [], None))
 
     def test_update_result_queue_on_failure(self):
         task = FailTask()
@@ -97,7 +94,7 @@ def test_update_result_queue_on_failure(self):
 
         with mock.patch.object(result_queue, 'put') as mock_put:
             task_process.run()
-            mock_put.assert_called_once_with((task.task_id, FAILED, json.dumps("test failure expl"), [], [], task.ex))
+            mock_put.assert_called_once_with((task.task_id, FAILED, json.dumps("test failure expl"), [], []))
 
     def test_cleanup_children_on_terminate(self):
         """
diff --git a/test/worker_test.py b/test/worker_test.py
--- a/test/worker_test.py
+++ b/test/worker_test.py
@@ -111,7 +111,6 @@ def setUp(self):
         # InstanceCache.disable()
         self.sch = CentralPlannerScheduler(retry_delay=100, remove_delay=1000, worker_disconnect_delay=10)
         self.w = Worker(scheduler=self.sch, worker_id='X')
-        self.w_raise = Worker(scheduler=self.sch, worker_id='X_raise', raise_on_error=True)
         self.w2 = Worker(scheduler=self.sch, worker_id='Y')
         self.time = time.time
 
@@ -193,11 +192,15 @@ def complete(self):
         self.assertFalse(b.has_run)
 
     def test_fail(self):
+        class CustomException(BaseException):
+            def __init__(self, msg):
+                self.msg = msg
+
         class A(Task):
 
             def run(self):
                 self.has_run = True
-                raise Exception()
+                raise CustomException('bad things')
 
             def complete(self):
                 return self.has_run
@@ -226,41 +229,6 @@ def complete(self):
         self.assertTrue(a.has_run)
         self.assertFalse(b.has_run)
 
-    def test_fail_raised(self):
-        class A(Task):
-
-            def run(self):
-                self.has_run = True
-                raise BaseException()
-
-            def complete(self):
-                return self.has_run
-
-        a = A()
-
-        class B(Task):
-
-            def requires(self):
-                return a
-
-            def run(self):
-                self.has_run = True
-
-            def complete(self):
-                return self.has_run
-
-        b = B()
-
-        a.has_run = False
-        b.has_run = False
-
-        self.assertTrue(self.w_raise.add(b))
-        with self.assertRaises(BaseException):
-            self.w_raise.run()
-
-        self.assertTrue(a.has_run)
-        self.assertFalse(b.has_run)
-
     def test_unknown_dep(self):
         # see central_planner_test.CentralPlannerTest.test_remove_dep
         class A(ExternalTask):