Repository: spotify/luigi

Issue:
Local scheduler does not prune automatically
Take these script: a main task that depends on a task that only completes successfully after some time has passed since launch

``` python
import time
import luigi

class DepThatNeedsSomeTime(luigi.Task):
    ready_time = luigi.FloatParameter()

    def run(self):
        assert self.complete()

    def complete(self):
        return  self.ready_time <= time.time()

class MainTask(luigi.WrapperTask):
    launch_time = luigi.FloatParameter(default=time.time())

    def requires(self):
        return [DepThatNeedsSomeTime(self.launch_time + 10)]

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

run it with these configs to simplify testing/debugging:

```
python script.py MainTask --scheduler-retry-delay 1 --worker-keep-alive
```

If you run this with central scheduler, the dependency will fail initially, but after some it will be retried because the server calls `scheduler.prune` every minute:

``` python
    # prune work DAG every 60 seconds
    pruner = tornado.ioloop.PeriodicCallback(scheduler.prune, 60000)
```

and `scheduler.prune` causes rescheduling/retrying of the failed task (through triggering `SimpleTaskState.prune`)

So eventually all work will be done successfully.

However if you run this with the local scheduler (additional `--local-scheduler`), you end up in an endless loop because `scheduler.prune` is never called and the failed task never gets a chance on a retry.

(FYI: I lost quite some time figuring all this out when I was examining the retry behaviour of Luigi with the local scheduler, assuming it would work the same as the central scheduler)



Discussion:

Resolving patch:
diff --git a/luigi/interface.py b/luigi/interface.py
--- a/luigi/interface.py
+++ b/luigi/interface.py
@@ -109,7 +109,7 @@ class core(task.Config):
 class WorkerSchedulerFactory(object):
 
     def create_local_scheduler(self):
-        return scheduler.CentralPlannerScheduler()
+        return scheduler.CentralPlannerScheduler(prune_on_get_work=True)
 
     def create_remote_scheduler(self, host, port):
         return rpc.RemoteScheduler(host=host, port=port)
diff --git a/luigi/scheduler.py b/luigi/scheduler.py
--- a/luigi/scheduler.py
+++ b/luigi/scheduler.py
@@ -101,6 +101,8 @@ class scheduler(Config):
 
     visualization_graph = parameter.Parameter(default="svg", config_path=dict(section='scheduler', name='visualization-graph'))
 
+    prune_on_get_work = parameter.BoolParameter(default=False)
+
 
 def fix_time(x):
     # Backwards compatibility for a fix in Dec 2014. Prior to the fix, pickled state might store datetime objects
@@ -227,7 +229,7 @@ class Worker(object):
     def __init__(self, worker_id, last_active=None):
         self.id = worker_id
         self.reference = None  # reference to the worker in the real world. (Currently a dict containing just the host)
-        self.last_active = last_active  # seconds since epoch
+        self.last_active = last_active or time.time()  # seconds since epoch
         self.started = time.time()  # seconds since epoch
         self.tasks = set()  # task objects
         self.info = {}
@@ -710,6 +712,9 @@ def get_work(self, host=None, assistant=False, **kwargs):
         # TODO: remove tasks that can't be done, figure out if the worker has absolutely
         # nothing it can wait for
 
+        if self._config.prune_on_get_work:
+            self.prune()
+
         worker_id = kwargs['worker']
         # Return remaining tasks that have no FAILED descendents
         self.update(worker_id, {'host': host})


Test patch:
diff --git a/test/worker_external_task_test.py b/test/worker_external_task_test.py
--- a/test/worker_external_task_test.py
+++ b/test/worker_external_task_test.py
@@ -80,7 +80,8 @@ def setUp(self):
                                                  worker_disconnect_delay=3,
                                                  disable_persist=3,
                                                  disable_window=5,
-                                                 disable_failures=2)
+                                                 disable_failures=2,
+                                                 prune_on_get_work=True)
 
     def _assert_complete(self, tasks):
         for t in tasks:
@@ -127,15 +128,6 @@ def test_external_dependency_completes_later(self):
                                                            'retry-external-tasks',
                                                            False) is True
 
-        original_get_work = self.scheduler.get_work
-
-        def decorated_get_work(*args, **kwargs):
-            # need to call `prune()` to make the scheduler run the retry logic
-            self.scheduler.prune()
-            return original_get_work(*args, **kwargs)
-
-        self.scheduler.get_work = decorated_get_work
-
         tempdir = tempfile.mkdtemp(prefix='luigi-test-')
 
         with patch('random.randint', return_value=0.1):