Speed up task_list when beyond limit

luigi/scheduler.py

@@ -474,6 +474,12 @@ def get_active_tasks(self):
     def get_active_tasks_by_status(self, *statuses):
         return itertools.chain.from_iterable(six.itervalues(self._status_tasks[status]) for status in statuses)
 
+    def get_active_task_count_for_status(self, status):
+        if status:
+            return len(self._status_tasks[status])
+        else:
+            return len(self._tasks)
+
     def get_batch_running_tasks(self, batch_id):
         assert batch_id is not None
         return [
@@ -1335,7 +1341,13 @@ def task_list(self, status='', upstream_status='', limit=True, search=None, max_
         """
         Query for a subset of tasks by status.
         """
+        if not search:
+            count_limit = max_shown_tasks or self._config.max_shown_tasks
+            pre_count = self._state.get_active_task_count_for_status(status)
+            if limit and pre_count > count_limit:
+                return {'num_tasks': -1 if upstream_status else pre_count}
         self.prune()
+
         result = {}
         upstream_status_table = {}  # used to memoize upstream status
         if search is None:

luigi/static/visualiser/js/visualiserApp.js

@@ -743,7 +743,7 @@ function visualiserApp(luigi) {
         var taskCount;
         /* Check for integers in tasks.  This indicates max-shown-tasks was exceeded */
         if (tasks.length === 1 && typeof(tasks[0]) === 'number') {
-            taskCount = tasks[0];
+            taskCount = tasks[0] === -1 ? 'unknown' : tasks[0];
             missingCategories[category] = {name: category, count: taskCount};
         }
         else {

test/scheduler_api_test.py

@@ -1767,6 +1767,21 @@ def test_task_list_filter_by_multiple_search_terms(self):
 
         self.search_pending('ClassA 2016-02-01 num', {expected})
 
+    def test_upstream_beyond_limit(self):
+        sch = Scheduler(max_shown_tasks=3)
+        for i in range(4):
+            sch.add_task(worker=WORKER, family='Test', params={'p': str(i)}, task_id='Test_%i' % i)
+        self.assertEqual({'num_tasks': -1}, sch.task_list('PENDING', 'FAILED'))
+        self.assertEqual({'num_tasks': 4}, sch.task_list('PENDING', ''))
+
+    def test_do_not_prune_on_beyond_limit_check(self):
+        sch = Scheduler(max_shown_tasks=3)
+        sch.prune = mock.Mock()
+        for i in range(4):
+            sch.add_task(worker=WORKER, family='Test', params={'p': str(i)}, task_id='Test_%i' % i)
+        self.assertEqual({'num_tasks': 4}, sch.task_list('PENDING', ''))
+        sch.prune.assert_not_called()
+
     def test_search_results_beyond_limit(self):
         sch = Scheduler(max_shown_tasks=3)
         for i in range(4):