Refine disabling of workers (#1839)

This commit's PR started small but grew quite big. But there's two major changes:

 * Bugfix: don't give tasks to disabled assistants

Previously assistants would keep getting tasks and never stop. Despite
that a user have clicked "disable worker" in the WebUI.

 * Invent a worker-state, pass it along in get_work and worker_list

This gives an opportunity for workers to know if the scheduler considers them disabled. It also allows the WebUI to show if a worker is disabled or not. Making the UI experience ever so more responsive. 

The old "greyed out worker" now means will mean "disabled". Previously it was equivalent to "has no pending tasks". But because we also associate "greyed out" with "can't click disable", the WebUI have
confused at least me for a while.

Author: Tarrasch

luigi/scheduler.py

@@ -69,6 +69,9 @@
     DISABLED: UPSTREAM_DISABLED,
 }
 
+WORKER_STATE_DISABLED = 'disabled'
+WORKER_STATE_ACTIVE = 'active'
+
 TASK_FAMILY_RE = re.compile(r'([^(_]+)[(_]')
 
 RPC_METHODS = {}
@@ -319,6 +322,17 @@ def is_trivial_worker(self, state):
     def assistant(self):
         return self.info.get('assistant', False)
 
+    @property
+    def enabled(self):
+        return not self.disabled
+
+    @property
+    def state(self):
+        if self.enabled:
+            return WORKER_STATE_ACTIVE
+        else:
+            return WORKER_STATE_DISABLED
+
     def __str__(self):
         return self.id
 
@@ -527,7 +541,7 @@ def get_active_workers(self, last_active_lt=None, last_get_work_gt=None):
         for worker in six.itervalues(self._active_workers):
             if last_active_lt is not None and worker.last_active >= last_active_lt:
                 continue
-            last_get_work = getattr(worker, 'last_get_work', None)
+            last_get_work = worker.last_get_work
             if last_get_work_gt is not None and (
                             last_get_work is None or last_get_work <= last_get_work_gt):
                 continue
@@ -554,10 +568,10 @@ def _remove_workers_from_tasks(self, workers, remove_stakeholders=True):
                 task.stakeholders.difference_update(workers)
             task.workers.difference_update(workers)
 
-    def disable_workers(self, workers):
-        self._remove_workers_from_tasks(workers, remove_stakeholders=False)
-        for worker in workers:
-            self.get_worker(worker).disabled = True
+    def disable_workers(self, worker_ids):
+        self._remove_workers_from_tasks(worker_ids, remove_stakeholders=False)
+        for worker_id in worker_ids:
+            self.get_worker(worker_id).disabled = True
 
 
 class Scheduler(object):
@@ -623,13 +637,12 @@ def _prune_tasks(self):
 
         self._state.inactivate_tasks(remove_tasks)
 
-    def update(self, worker_id, worker_reference=None, get_work=False):
-        """
-        Keep track of whenever the worker was last active.
-        """
+    def _update_worker(self, worker_id, worker_reference=None, get_work=False):
+        # Keep track of whenever the worker was last active.
+        # For convenience also return the worker object.
         worker = self._state.get_worker(worker_id)
         worker.update(worker_reference, get_work=get_work)
-        return not getattr(worker, 'disabled', False)
+        return worker
 
     def _update_priority(self, task, prio, worker):
         """
@@ -663,10 +676,10 @@ def add_task(self, task_id=None, status=PENDING, runnable=True,
         """
         assert worker is not None
         worker_id = worker
-        worker_enabled = self.update(worker_id)
+        worker = self._update_worker(worker_id)
         retry_policy = self._generate_retry_policy(retry_policy_dict)
 
-        if worker_enabled:
+        if worker.enabled:
             _default_task = self._make_task(
                 task_id=task_id, status=PENDING, deps=deps, resources=resources,
                 priority=priority, family=family, module=module, params=params,
@@ -676,7 +689,7 @@ def add_task(self, task_id=None, status=PENDING, runnable=True,
 
         task = self._state.get_task(task_id, setdefault=_default_task)
 
-        if task is None or (task.status != RUNNING and not worker_enabled):
+        if task is None or (task.status != RUNNING and not worker.enabled):
             return
 
         # for setting priority, we'll sometimes create tasks with unset family and params
@@ -728,7 +741,7 @@ def add_task(self, task_id=None, status=PENDING, runnable=True,
         if resources is not None:
             task.resources = resources
 
-        if worker_enabled and not assistant:
+        if worker.enabled and not assistant:
             task.stakeholders.add(worker_id)
 
             # Task dependencies might not exist yet. Let's create dummy tasks for them for now.
@@ -743,7 +756,7 @@ def add_task(self, task_id=None, status=PENDING, runnable=True,
         # before we know their retry_policy, we always set it here
         task.retry_policy = retry_policy
 
-        if runnable and status != FAILED and worker_enabled:
+        if runnable and status != FAILED and worker.enabled:
             task.workers.add(worker_id)
             self._state.get_worker(worker_id).tasks.add(task)
             task.runnable = runnable
@@ -837,8 +850,19 @@ def get_work(self, host=None, assistant=False, current_tasks=None, worker=None,
 
         assert worker is not None
         worker_id = worker
-        # Return remaining tasks that have no FAILED descendants
-        self.update(worker_id, {'host': host}, get_work=True)
+        worker = self._update_worker(
+            worker_id,
+            worker_reference={'host': host},
+            get_work=True)
+        if not worker.enabled:
+            reply = {'n_pending_tasks': 0,
+                     'running_tasks': [],
+                     'task_id': None,
+                     'n_unique_pending': 0,
+                     'worker_state': worker.state,
+                     }
+            return reply
+
         if assistant:
             self.add_worker(worker_id, [('assistant', assistant)])
 
@@ -942,7 +966,9 @@ def get_work(self, host=None, assistant=False, current_tasks=None, worker=None,
         reply = {'n_pending_tasks': locally_pending_tasks,
                  'running_tasks': running_tasks,
                  'task_id': None,
-                 'n_unique_pending': n_unique_pending}
+                 'n_unique_pending': n_unique_pending,
+                 'worker_state': worker.state,
+                 }
 
         if len(batched_tasks) > 1:
             batch_string = '|'.join(task.id for task in batched_tasks)
@@ -976,7 +1002,7 @@ def get_work(self, host=None, assistant=False, current_tasks=None, worker=None,
     @rpc_method(attempts=1)
     def ping(self, **kwargs):
         worker_id = kwargs['worker']
-        self.update(worker_id)
+        self._update_worker(worker_id)
 
     def _upstream_status(self, task_id, upstream_status_table):
         if task_id in upstream_status_table:
@@ -1142,7 +1168,7 @@ def filter_func(t):
                 return all(term in t.pretty_id for term in terms)
         for task in filter(filter_func, self._state.get_active_tasks(status)):
             if task.status != PENDING or not upstream_status or upstream_status == self._upstream_status(task.id, upstream_status_table):
-                serialized = self._serialize_task(task.id, False)
+                serialized = self._serialize_task(task.id, include_deps=False)
                 result[task.id] = serialized
         if limit and len(result) > (max_shown_tasks or self._config.max_shown_tasks):
             return {'num_tasks': len(result)}
@@ -1162,7 +1188,8 @@ def worker_list(self, include_running=True, **kwargs):
             dict(
                 name=worker.id,
                 last_active=worker.last_active,
-                started=getattr(worker, 'started', None),
+                started=worker.started,
+                state=worker.state,
                 first_task_display_name=self._first_task_display_name(worker),
                 **worker.info
             ) for worker in self._state.get_active_workers()]
@@ -1173,7 +1200,7 @@ def worker_list(self, include_running=True, **kwargs):
             num_uniques = collections.defaultdict(int)
             for task in self._state.get_pending_tasks():
                 if task.status == RUNNING and task.worker_running:
-                    running[task.worker_running][task.id] = self._serialize_task(task.id, False)
+                    running[task.worker_running][task.id] = self._serialize_task(task.id, include_deps=False)
                 elif task.status == PENDING:
                     for worker in task.workers:
                         num_pending[worker] += 1
@@ -1204,7 +1231,7 @@ def resource_list(self):
             for task in self._state.get_running_tasks():
                 if task.status == RUNNING and task.resources:
                     for resource, amount in six.iteritems(task.resources):
-                        consumers[resource][task.id] = self._serialize_task(task.id, False)
+                        consumers[resource][task.id] = self._serialize_task(task.id, include_deps=False)
             for resource in resources:
                 tasks = consumers[resource['name']]
                 resource['num_consumer'] = len(tasks)
@@ -1235,7 +1262,7 @@ def task_search(self, task_str, **kwargs):
         result = collections.defaultdict(dict)
         for task in self._state.get_active_tasks():
             if task.id.find(task_str) != -1:
-                serialized = self._serialize_task(task.id, False)
+                serialized = self._serialize_task(task.id, include_deps=False)
                 result[task.status][task.id] = serialized
         return result
 

luigi/static/visualiser/index.html

@@ -274,6 +274,9 @@ <h4 class="modal-title" id="disiableWorkerLabel">Disable worker?</h4>
                   </div>
                   <div class="modal-body">
                     Are you sure you want to disable this worker?
+                    <p>
+                        A disabled worker will finish its existing tasks but not start new ones.
+                    </p>
                   </div>
                   <div class="modal-footer">
                     <button type="button" class="btn btn-default" data-dismiss="modal">Cancel</button>
@@ -283,22 +286,22 @@ <h4 class="modal-title" id="disiableWorkerLabel">Disable worker?</h4>
               </div>
             </div>
             {{#workers}}
-            {{#num_pending}}
-            <div class="box">
-            {{/num_pending}}
-            {{^num_pending}}
+            {{#is_disabled}}
             <div class="box box-solid box-default">
-            {{/num_pending}}
+            {{/is_disabled}}
+            {{^is_disabled}}
+            <div class="box">
+            {{/is_disabled}}
                 <div class="box-header with-border">
                     <h3 class="box-title">{{name}}</h3>
                     <div class="box-tools pull-right">
-                      {{#num_pending}}
+                      {{^is_disabled}}
                       <div class="button-tooltip" data-toggle="tooltip" title="Disable Worker">
                         <button type="button" class="btn btn-danger btn-disable-worker" data-toggle="modal" data-target="#disableWorkerModal" data-worker="{{name}}">
                           <i class="fa fa-fire-extinguisher"></i>
                         </button>
                       </div>
-                      {{/num_pending}}
+                      {{/is_disabled}}
                     </div>
                 </div>
                 <div class="box-body">
@@ -308,6 +311,9 @@ <h3 class="box-title">{{name}}</h3>
                     Running: {{num_running}}<br>
                     Pending: {{num_pending}}<br>
                     Unique Pending: {{num_uniques}}<br>
+                    {{#is_disabled}}
+                    This worker is <b>disabled</b>. It will not start new tasks.<br>
+                    {{/is_disabled}}
 
                     {{#num_running}}
                     <hr>

luigi/static/visualiser/js/visualiserApp.js

@@ -240,6 +240,7 @@ function visualiserApp(luigi) {
         worker.tasks.sort(function(task1, task2) { return task1.timeRunning - task2.timeRunning; });
         worker.start_time = new Date(worker.started * 1000).toLocaleString();
         worker.active = new Date(worker.last_active * 1000).toLocaleString();
+        worker.is_disabled = worker.state === 'disabled';
         return worker;
     }
 

luigi/worker.py

@@ -54,6 +54,7 @@
 from luigi.event import Event
 from luigi.task_register import load_task
 from luigi.scheduler import DISABLED, DONE, FAILED, PENDING, UNKNOWN, Scheduler, RetryPolicy
+from luigi.scheduler import WORKER_STATE_ACTIVE, WORKER_STATE_DISABLED
 from luigi.target import Target
 from luigi.task import Task, flatten, getpaths, Config
 from luigi.task_register import TaskClassException
@@ -730,7 +731,7 @@ def _get_work_task_id(self, get_work_response):
 
     def _get_work(self):
         if self._stop_requesting_work:
-            return None, 0, 0, 0
+            return None, 0, 0, 0, WORKER_STATE_DISABLED
         logger.debug("Asking scheduler for work...")
         r = self._scheduler.get_work(
             worker=self._id,
@@ -741,6 +742,9 @@ def _get_work(self):
         n_pending_tasks = r['n_pending_tasks']
         running_tasks = r['running_tasks']
         n_unique_pending = r['n_unique_pending']
+        # TODO: For a tiny amount of time (a month?) we'll keep forwards compatibility
+        # That is you can user a newer client than server (Sep 2016)
+        worker_state = r.get('worker_state', WORKER_STATE_ACTIVE)  # state according to server!
         task_id = self._get_work_task_id(r)
 
         self._get_work_response_history.append({
@@ -773,7 +777,7 @@ def _get_work(self):
                 self._scheduled_tasks.get(batch_id) for batch_id in r['batch_task_ids']])
             self._batch_running_tasks[task_id] = batch_tasks
 
-        return task_id, running_tasks, n_pending_tasks, n_unique_pending
+        return task_id, running_tasks, n_pending_tasks, n_unique_pending, worker_state
 
     def _run_task(self, task_id):
         task = self._scheduled_tasks[task_id]
@@ -927,8 +931,15 @@ def handle_interrupt(self, signum, _):
         Stops the assistant from asking for more work on SIGUSR1
         """
         if signum == signal.SIGUSR1:
-            self._config.keep_alive = False
-            self._stop_requesting_work = True
+            self._start_phasing_out()
+
+    def _start_phasing_out(self):
+        """
+        Go into a mode where we dont ask for more work and quit once existing
+        tasks are done.
+        """
+        self._config.keep_alive = False
+        self._stop_requesting_work = True
 
     def run(self):
         """
@@ -946,7 +957,10 @@ def run(self):
                 logger.debug('%d running tasks, waiting for next task to finish', len(self._running_tasks))
                 self._handle_next_task()
 
-            task_id, running_tasks, n_pending_tasks, n_unique_pending = self._get_work()
+            task_id, running_tasks, n_pending_tasks, n_unique_pending, worker_state = self._get_work()
+
+            if worker_state == WORKER_STATE_DISABLED:
+                self._start_phasing_out()
 
             if task_id is None:
                 if not self._stop_requesting_work:

test/scheduler_api_test.py

@@ -1305,6 +1305,15 @@ def test_disable_worker_can_fail_task(self):
     def test_disable_worker_stays_disabled_on_new_deps(self):
         self._test_disable_worker_helper(new_status='PENDING', new_deps=['B', 'C'])
 
+    def test_disable_worker_assistant_gets_no_task(self):
+        self.setTime(0)
+        self.sch.add_task(worker=WORKER, task_id='A')
+        self.sch.add_worker('assistant', [('assistant', True)])
+        self.sch.ping(worker='assistant')
+        self.sch.disable_worker('assistant')
+        self.assertIsNone(self.sch.get_work(worker='assistant', assistant=True)['task_id'])
+        self.assertIsNotNone(self.sch.get_work(worker=WORKER)['task_id'])
+
     def test_prune_worker(self):
         self.setTime(1)
         self.sch.add_worker(worker=WORKER, info={})

test/scheduler_visualisation_test.py

@@ -532,6 +532,7 @@ def complete(self):
         self.assertEqual(0, worker['num_pending'])
         self.assertEqual(0, worker['num_uniques'])
         self.assertEqual(0, worker['num_running'])
+        self.assertEqual('active', worker['state'])
         self.assertEqual(1, worker['workers'])
 
     def test_worker_list_pending_uniques(self):
@@ -583,6 +584,17 @@ class X(luigi.Task):
         self.assertEqual(1, worker['num_pending'])
         self.assertEqual(1, worker['num_uniques'])
 
+    def test_worker_list_disabled_worker(self):
+        class X(luigi.Task):
+            pass
 
-if __name__ == '__main__':
-    unittest.main()
+        with luigi.worker.Worker(worker_id='w', scheduler=self.scheduler) as w:
+            w.add(X())  #
+            workers = self._remote().worker_list()
+            self.assertEqual(1, len(workers))
+            self.assertEqual('active', workers[0]['state'])
+            self.scheduler.disable_worker('w')
+            workers = self._remote().worker_list()
+            self.assertEqual(1, len(workers))
+            self.assertEqual(1, len(workers))
+            self.assertEqual('disabled', workers[0]['state'])