Fix attribute forwarding for tasks with dynamic dependencies

luigi/worker.py

@@ -37,6 +37,7 @@
 import signal
 import subprocess
 import sys
+import contextlib
 
 try:
  import Queue
@@ -135,16 +136,8 @@ def __init__(self, task, worker_id, result_queue, status_reporter,
  self.check_unfulfilled_deps = check_unfulfilled_deps
 
  def _run_get_new_deps(self):
- # forward some attributes before running
- for reporter_attr, task_attr in six.iteritems(self.forward_reporter_attributes):
- setattr(self.task, task_attr, getattr(self.status_reporter, reporter_attr))
-
  task_gen = self.task.run()
 
- # reset attributes again
- for reporter_attr, task_attr in six.iteritems(self.forward_reporter_attributes):
- setattr(self.task, task_attr, None)
-
  if not isinstance(task_gen, types.GeneratorType):
  return None
 
@@ -202,7 +195,8 @@ def run(self):
  expl = 'Task is an external data dependency ' \
  'and data does not exist (yet?).'
  else:
- new_deps = self._run_get_new_deps()
+ with self._forward_attributes():
+ new_deps = self._run_get_new_deps()
  status = DONE if not new_deps else PENDING
 
  if new_deps:
@@ -258,6 +252,20 @@ def terminate(self):
  except ImportError:
  return super(TaskProcess, self).terminate()
 
+ @contextlib.contextmanager
+ def _forward_attributes(self):
+ # forward configured attributes to the task
+ for reporter_attr, task_attr in six.iteritems(self.forward_reporter_attributes):
+ setattr(self.task, task_attr, getattr(self.status_reporter, reporter_attr))
+
+ try:
+ yield self
+
+ finally:
+ # reset attributes again
+ for reporter_attr, task_attr in six.iteritems(self.forward_reporter_attributes):
+ setattr(self.task, task_attr, None)
+
 
 # This code and the task_process_context config key currently feels a bit ad-hoc.
 # Discussion on generalizing it into a plugin system: https://github.com/spotify/luigi/issues/1897

test/task_forwarded_attributes_test.py

@@ -0,0 +1,76 @@
+# -*- coding: utf-8 -*-
+#
+# Copyright 2012-2015 Spotify AB
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+from helpers import LuigiTestCase, RunOnceTask
+
+import luigi
+import luigi.scheduler
+import luigi.worker
+
+
+forwarded_attributes = set(luigi.worker.TaskProcess.forward_reporter_attributes.values())
+
+
+class NonYieldingTask(RunOnceTask):
+
+ # need to accept messages in order for the "scheduler_message" attribute to be not None
+ accepts_messages = True
+
+ def gather_forwarded_attributes(self):
+ attrs = set()
+ for attr in forwarded_attributes:
+ if getattr(self, attr, None) is not None:
+ attrs.add(attr)
+ return attrs
+
+ def run(self):
+ self.attributes_while_running = self.gather_forwarded_attributes()
+
+ RunOnceTask.run(self)
+
+
+class YieldingTask(NonYieldingTask):
+
+ def run(self):
+ self.attributes_before_yield = self.gather_forwarded_attributes()
+
+ yield RunOnceTask()
+
+ self.attributes_after_yield = self.gather_forwarded_attributes()
+
+ RunOnceTask.run(self)
+
+
+class TaskForwardedAttributesTest(LuigiTestCase):
+
+ def run_task(self, task):
+ sch = luigi.scheduler.Scheduler()
+ with luigi.worker.Worker(scheduler=sch) as w:
+ w.add(task)
+ w.run()
+ return task
+
+ def test_non_yielding_task(self):
+ task = self.run_task(NonYieldingTask())
+
+ self.assertEqual(task.attributes_while_running, forwarded_attributes)
+
+ def test_yielding_task(self):
+ task = self.run_task(YieldingTask())
+
+ self.assertEqual(task.attributes_before_yield, forwarded_attributes)
+ self.assertEqual(task.attributes_after_yield, forwarded_attributes)