Repository: dask/distributed

Issue:
No tasks on scheduler yet total_occupancy > 0 
I observed an example yesterday of the scheduler having no remaining tasks, yet total_occupancy was > 0. This mattered, because I was using Adaptive, and it meant there was a continuous cycle of spin up and start a new worker, only for it to be immediately closed.


Relevant log snippet below:
```
2017-11-09 20:01:13  distributed.deploy.adaptive (scheduler): CPU limit exceeded [353 occupancy / 0 cores]
2017-11-09 20:01:13  distributed.deploy.adaptive (scheduler): Scaling up to 1 workers
2017-11-09 20:01:17  distributed.comm.tcp (worker-group-1-7kl0): Setting TCP keepalive: nprobes=10, idle=10, interval=2
2017-11-09 20:01:17  distributed.comm.tcp (worker-group-1-7kl0): Setting TCP user timeout: 30000 ms
2017-11-09 20:01:17  distributed.comm.tcp (worker-group-1-7kl0): TLS connection with 'tls://scheduler:8786': protocol=TLSv1.2, cipher=ECDHE-RSA-AES256-GCM-SHA384, bits=256
2017-11-09 20:01:17  distributed.scheduler (scheduler): Register tls://10.154.0.5:32819
2017-11-09 20:01:17  distributed.scheduler (scheduler): Starting worker compute stream, tls://10.154.0.5:32819
2017-11-09 20:01:18  distributed.scheduler (scheduler): Closing worker tls://10.154.0.5:32819
2017-11-09 20:01:18  distributed.core (worker-group-1-7kl0): Message from 'tls://10.154.0.4:33068': {'reply': False, 'op': 'close'}
2017-11-09 20:01:18  distributed.scheduler (scheduler): Remove worker tls://10.154.0.5:32819
2017-11-09 20:01:18  distributed.scheduler (scheduler): Lost all workers
2017-11-09 20:01:18  distributed.scheduler (scheduler): Removed worker tls://10.154.0.5:32819
2017-11-09 20:01:18  distributed.deploy.adaptive (scheduler): Retiring workers ['tls://10.154.0.5:32819']
```

I suspect this is because in the scheduler, `transition_processing_released` (unlike in `transition_processing_memory`), there is no check to see if there is any outstanding processing, following which occupancy is set to zero. But before proposing a PR, I thought I would see if someone with a better understanding of the scheduler concurred.



Discussion:
This is on distributed master @dd269f4
> I suspect this is because in the scheduler, transition_processing_released (unlike in transition_processing_memory), there is no check to see if there is any outstanding processing, following which occupancy is set to zero. But before proposing a PR, I thought I would see if someone with a better understanding of the scheduler concurred.

That sounds like a fair assessment.  If you feel like submitting a PR that would be welcome.

Resolving patch:
diff --git a/distributed/scheduler.py b/distributed/scheduler.py
--- a/distributed/scheduler.py
+++ b/distributed/scheduler.py
@@ -651,6 +651,7 @@ def add_worker(self, comm=None, address=None, keys=(), ncores=None,
                 self.worker_bytes[address] = 0
                 self.processing[address] = dict()
                 self.occupancy[address] = 0
+                # Do not need to adjust self.total_occupancy as self.occupancy[address] cannot exist before this.
                 self.check_idle_saturated(address)
 
             # for key in keys:  # TODO
@@ -1145,7 +1146,8 @@ def validate_state(self, allow_overlap=False):
                 set(self.has_what) ==
                 set(self.processing) ==
                 set(self.worker_info) ==
-                set(self.worker_comms)):
+                set(self.worker_comms) ==
+                set(self.occupancy)):
             raise ValueError("Workers not the same in all collections")
 
         a = self.worker_bytes
@@ -1163,11 +1165,13 @@ def validate_state(self, allow_overlap=False):
 
         assert all(self.who_has.values())
 
-        # for worker, occ in self.occupancy.items():
+        for worker, occ in self.occupancy.items():
         #     for d in self.extensions['stealing'].in_flight.values():
         #         if worker in (d['thief'], d['victim']):
         #             continue
-        #     assert abs(sum(self.processing[worker].values()) - occ) < 1e-8
+            assert abs(sum(self.processing[worker].values()) - occ) < 1e-8
+
+        assert abs(sum(self.occupancy.values()) - self.total_occupancy) < 1e-8
 
     ###################
     # Manage Messages #
@@ -1425,6 +1429,7 @@ def handle_long_running(self, key=None, worker=None, compute_duration=None):
 
         worker = self.rprocessing[key]
         self.occupancy[actual_worker] -= self.processing[actual_worker][key]
+        self.total_occupancy -= self.processing[actual_worker][key]
         self.processing[actual_worker][key] = 0
 
     @gen.coroutine
@@ -2762,8 +2767,12 @@ def transition_processing_released(self, key):
             w = self.rprocessing.pop(key)
             if w in self.workers:
                 duration = self.processing[w].pop(key)
-                self.occupancy[w] -= duration
-                self.total_occupancy -= duration
+                if not self.processing[w]:
+                    self.total_occupancy -= self.occupancy[w]
+                    self.occupancy[w] = 0
+                else:
+                    self.total_occupancy -= duration
+                    self.occupancy[w] -= duration
                 self.check_idle_saturated(w)
                 self.release_resources(key, w)
                 self.worker_send(w, {'op': 'release-task', 'key': key})
@@ -2835,8 +2844,12 @@ def transition_processing_erred(self, key, cause=None, exception=None,
             w = self.rprocessing.pop(key)
             if w in self.processing:
                 duration = self.processing[w].pop(key)
-                self.occupancy[w] -= duration
-                self.total_occupancy -= duration
+                if not self.processing[w]:
+                    self.total_occupancy -= self.occupancy[w]
+                    self.occupancy[w] = 0
+                else:
+                    self.total_occupancy -= duration
+                    self.occupancy[w] -= duration
                 self.check_idle_saturated(w)
                 self.release_resources(key, w)
 
@@ -3146,8 +3159,9 @@ def reschedule(self, key=None, worker=None):
     # Assigning Tasks to Workers #
     ##############################
 
-    def check_idle_saturated(self, worker):
-        occ = self.occupancy[worker]
+    def check_idle_saturated(self, worker, occ=None):
+        if occ is None:
+            occ = self.occupancy[worker]
         nc = self.ncores[worker]
         p = len(self.processing[worker])
 
diff --git a/distributed/stealing.py b/distributed/stealing.py
--- a/distributed/stealing.py
+++ b/distributed/stealing.py
@@ -49,6 +49,7 @@ def __init__(self, scheduler):
         self.scheduler.events['stealing'] = deque(maxlen=100000)
         self.count = 0
         self.in_flight = dict()
+        self.in_flight_occupancy = defaultdict(lambda: 0)
 
         self.scheduler.worker_handlers['steal-response'] = self.move_task_confirm
 
@@ -165,15 +166,11 @@ def move_task_request(self, key, victim, thief):
                          thief, self.scheduler.occupancy[thief])
 
             victim_duration = self.scheduler.processing[victim][key]
-            self.scheduler.occupancy[victim] -= victim_duration
-            self.scheduler.total_occupancy -= victim_duration
 
             thief_duration = self.scheduler.task_duration.get(key_split(key), 0.5)
             thief_duration += sum(self.scheduler.nbytes[key] for key in
                                   self.scheduler.dependencies[key] -
                                   self.scheduler.has_what[thief]) / BANDWIDTH
-            self.scheduler.occupancy[thief] += thief_duration
-            self.scheduler.total_occupancy += thief_duration
 
             self.scheduler.worker_comms[victim].send({'op': 'steal-request',
                                                       'key': key})
@@ -182,6 +179,9 @@ def move_task_request(self, key, victim, thief):
                                    'thief': thief,
                                    'victim_duration': victim_duration,
                                    'thief_duration': thief_duration}
+
+            self.in_flight_occupancy[victim] -= victim_duration
+            self.in_flight_occupancy[thief] += thief_duration
         except CommClosedError:
             logger.info("Worker comm closed while stealing: %s", victim)
         except Exception as e:
@@ -202,10 +202,21 @@ def move_task_confirm(self, key=None, worker=None, state=None):
             logger.debug("Confirm move %s, %s -> %s.  State: %s", key, victim,
                          thief, state)
 
+            self.in_flight_occupancy[thief] -= d['thief_duration']
+            self.in_flight_occupancy[victim] += d['victim_duration']
+
+            if not self.in_flight:
+                self.in_flight_occupancy = defaultdict(lambda: 0)
+
             if (self.scheduler.task_state.get(key) != 'processing' or
                     self.scheduler.rprocessing[key] != victim):
-                self.scheduler.occupancy[thief] = sum(self.scheduler.processing[thief].values())
-                self.scheduler.occupancy[victim] = sum(self.scheduler.processing[victim].values())
+                old_thief = self.scheduler.occupancy[thief]
+                new_thief = sum(self.scheduler.processing[thief].values())
+                old_victim = self.scheduler.occupancy[victim]
+                new_victim = sum(self.scheduler.processing[victim].values())
+                self.scheduler.occupancy[thief] = new_thief
+                self.scheduler.occupancy[victim] = new_victim
+                self.scheduler.total_occupancy += new_thief - old_thief + new_victim - old_victim
                 return
 
             # One of the pair has left, punt and reschedule
@@ -216,10 +227,6 @@ def move_task_confirm(self, key=None, worker=None, state=None):
 
             # Victim had already started execution, reverse stealing
             if state in ('memory', 'executing', 'long-running', None):
-                self.scheduler.occupancy[thief] -= d['thief_duration']
-                self.scheduler.total_occupancy -= d['thief_duration']
-                self.scheduler.occupancy[victim] += d['victim_duration']
-                self.scheduler.total_occupancy += d['victim_duration']
                 self.log.append(('already-computing', key, victim, thief))
                 self.scheduler.check_idle_saturated(thief)
                 self.scheduler.check_idle_saturated(victim)
@@ -228,8 +235,13 @@ def move_task_confirm(self, key=None, worker=None, state=None):
             elif state in ('waiting', 'ready'):
                 self.remove_key_from_stealable(key)
                 self.scheduler.rprocessing[key] = thief
+                duration = self.scheduler.processing[victim][key]
+                self.scheduler.occupancy[victim] -= duration
+                self.scheduler.total_occupancy -= duration
                 del self.scheduler.processing[victim][key]
                 self.scheduler.processing[thief][key] = d['thief_duration']
+                self.scheduler.occupancy[thief] += d['thief_duration']
+                self.scheduler.total_occupancy += d['thief_duration']
                 self.put_key_in_stealable(key)
 
                 try:
@@ -247,13 +259,16 @@ def move_task_confirm(self, key=None, worker=None, state=None):
             raise
 
     def balance(self):
+        s = self.scheduler
+
+        def combined_occupancy(w):
+            return s.occupancy[w] + self.in_flight_occupancy[w]
+
         with log_errors():
             i = 0
-            s = self.scheduler
-            occupancy = s.occupancy
             idle = s.idle
             saturated = s.saturated
-            if not idle or len(idle) == len(self.scheduler.workers):
+            if not idle or len(idle) == len(s.workers):
                 return
 
             log = list()
@@ -263,15 +278,14 @@ def balance(self):
             acted = False
 
             if not s.saturated:
-                saturated = topk(10, s.workers, key=occupancy.get)
+                saturated = topk(10, s.workers, key=combined_occupancy)
                 saturated = [w for w in saturated
-                             if occupancy[w] > 0.2
+                             if combined_occupancy(w) > 0.2
                              and len(s.processing[w]) > s.ncores[w]]
             elif len(s.saturated) < 20:
-                saturated = sorted(saturated, key=occupancy.get, reverse=True)
-
+                saturated = sorted(saturated, key=combined_occupancy, reverse=True)
             if len(idle) < 20:
-                idle = sorted(idle, key=occupancy.get)
+                idle = sorted(idle, key=combined_occupancy)
 
             for level, cost_multiplier in enumerate(self.cost_multipliers):
                 if not idle:
@@ -299,14 +313,14 @@ def balance(self):
                             stealable.remove(key)
                             continue
 
-                        if (occupancy[idl] + cost_multiplier * duration
-                                <= occupancy[sat] - duration / 2):
+                        if (combined_occupancy(idl) + cost_multiplier * duration
+                                <= combined_occupancy(sat) - duration / 2):
                             self.move_task_request(key, sat, idl)
                             log.append((start, level, key, duration,
-                                        sat, occupancy[sat],
-                                        idl, occupancy[idl]))
-                            self.scheduler.check_idle_saturated(sat)
-                            self.scheduler.check_idle_saturated(idl)
+                                        sat, combined_occupancy(sat),
+                                        idl, combined_occupancy(idl)))
+                            s.check_idle_saturated(sat, occ=combined_occupancy(sat))
+                            s.check_idle_saturated(idl, occ=combined_occupancy(idl))
                             seen = True
 
                 if self.cost_multipliers[level] < 20:  # don't steal from public at cost
@@ -325,7 +339,7 @@ def balance(self):
                         except KeyError:
                             stealable.remove(key)
                             continue
-                        if occupancy[sat] < 0.2:
+                        if combined_occupancy(sat) < 0.2:
                             continue
                         if len(s.processing[sat]) <= s.ncores[sat]:
                             continue
@@ -334,14 +348,14 @@ def balance(self):
                         idl = idle[i % len(idle)]
                         duration = s.processing[sat][key]
 
-                        if (occupancy[idl] + cost_multiplier * duration
-                                <= occupancy[sat] - duration / 2):
+                        if (combined_occupancy(idl) + cost_multiplier * duration
+                                <= combined_occupancy(sat) - duration / 2):
                             self.move_task_request(key, sat, idl)
                             log.append((start, level, key, duration,
-                                        sat, occupancy[sat],
-                                        idl, occupancy[idl]))
-                            self.scheduler.check_idle_saturated(sat)
-                            self.scheduler.check_idle_saturated(idl)
+                                        sat, combined_occupancy(sat),
+                                        idl, combined_occupancy(idl)))
+                            s.check_idle_saturated(sat, occ=combined_occupancy(sat))
+                            s.check_idle_saturated(idl, occ=combined_occupancy(idl))
                             seen = True
 
                 if seen and not acted:
@@ -351,8 +365,8 @@ def balance(self):
                 self.log.append(log)
                 self.count += 1
             stop = time()
-            if self.scheduler.digests:
-                self.scheduler.digests['steal-duration'].add(stop - start)
+            if s.digests:
+                s.digests['steal-duration'].add(stop - start)
 
     def restart(self, scheduler):
         for stealable in self.stealable.values():


Test patch:
diff --git a/distributed/tests/test_scheduler.py b/distributed/tests/test_scheduler.py
--- a/distributed/tests/test_scheduler.py
+++ b/distributed/tests/test_scheduler.py
@@ -824,6 +824,21 @@ def test_learn_occupancy_2(c, s, a, b):
     assert 50 < s.total_occupancy < 200
 
 
+@gen_cluster(client=True)
+def test_occupancy_cleardown(c, s, a, b):
+    s.validate = False
+
+    # Inject excess values in s.occupancy
+    s.occupancy[a.address] = 2
+    s.total_occupancy += 2
+    futures = c.map(slowinc, range(100), delay=0.01)
+    yield wait(futures)
+
+    # Verify that occupancy values have been zeroed out
+    assert abs(s.total_occupancy) < 0.01
+    assert all(v == 0 for v in s.occupancy.values())
+
+
 @nodebug
 @gen_cluster(client=True, ncores=[('127.0.0.1', 1)] * 30)
 def test_balance_many_workers(c, s, *workers):
diff --git a/distributed/tests/test_steal.py b/distributed/tests/test_steal.py
--- a/distributed/tests/test_steal.py
+++ b/distributed/tests/test_steal.py
@@ -410,7 +410,7 @@ def assert_balanced(inp, expected, c, s, *workers):
 
         if result2 == expected2:
             return
-    raise Exception()
+    raise Exception('Expected: {}; got: {}'.format(str(expected2), str(result2)))
 
 
 @pytest.mark.parametrize('inp,expected', [