def balance(self):
    s = self.scheduler

    def combined_occupancy(ws):
        return (ws.occupancy + self.in_flight_occupancy[ws])

    def maybe_move_task(level, ts, sat, idl, duration, cost_multiplier):
        occ_idl = combined_occupancy(idl)
        occ_sat = combined_occupancy(sat)
        if ((occ_idl + (cost_multiplier * duration)) <= (occ_sat - (duration / 2))):
            self.move_task_request(ts, sat, idl)
            log.append((start, level, ts.key, duration, sat.address, occ_sat, idl.address, occ_idl))
            s.check_idle_saturated(sat, occ=occ_sat)
            s.check_idle_saturated(idl, occ=occ_idl)
    with log_errors():
        i = 0
        idle = s.idle.values()
        saturated = s.saturated
        if ((not idle) or (len(idle) == len(s.workers))):
            return
        log = []
        start = time()
        if (not s.saturated):
            saturated = topk(10, s.workers.values(), key=combined_occupancy)
            saturated = [ws for ws in saturated if ((combined_occupancy(ws) > 0.2) and (len(ws.processing) > ws.nthreads))]
        elif (len(s.saturated) < 20):
            saturated = sorted(saturated, key=combined_occupancy, reverse=True)
        if (len(idle) < 20):
            idle = sorted(idle, key=combined_occupancy)
        for (level, cost_multiplier) in enumerate(self.cost_multipliers):
            if (not idle):
                break
            for sat in list(saturated):
                stealable = self.stealable[sat.address][level]
                if ((not stealable) or (not idle)):
                    continue
                for ts in list(stealable):
                    if ((ts not in self.key_stealable) or (ts.processing_on is not sat)):
                        stealable.discard(ts)
                        continue
                    i += 1
                    if (not idle):
                        break
                    if _has_restrictions(ts):
                        thieves = [ws for ws in idle if _can_steal(ws, ts, sat)]
                    else:
                        thieves = idle
                    if (not thieves):
                        break
                    thief = thieves[(i % len(thieves))]
                    duration = sat.processing.get(ts)
                    if (duration is None):
                        stealable.discard(ts)
                        continue
                    maybe_move_task(level, ts, sat, thief, duration, cost_multiplier)
            if (self.cost_multipliers[level] < 20):
                stealable = self.stealable_all[level]
                for ts in list(stealable):
                    if (not idle):
                        break
                    if (ts not in self.key_stealable):
                        stealable.discard(ts)
                        continue
                    sat = ts.processing_on
                    if (sat is None):
                        stealable.discard(ts)
                        continue
                    if (combined_occupancy(sat) < 0.2):
                        continue
                    if (len(sat.processing) <= sat.nthreads):
                        continue
                    i += 1
                    if _has_restrictions(ts):
                        thieves = [ws for ws in idle if _can_steal(ws, ts, sat)]
                    else:
                        thieves = idle
                    if (not thieves):
                        continue
                    thief = thieves[(i % len(thieves))]
                    duration = sat.processing[ts]
                    maybe_move_task(level, ts, sat, thief, duration, cost_multiplier)
        if log:
            self.log(log)
            self.count += 1
        stop = time()
        if s.digests:
            s.digests['steal-duration'].add((stop - start))
def balance(self):
    s = self.scheduler

    def combined_occupancy(ws):
        return (ws.occupancy + self.in_flight_occupancy[ws])

    def maybe_move_task(level, ts, sat, idl, duration, cost_multiplier):
        occ_idl = combined_occupancy(idl)
        occ_sat = combined_occupancy(sat)
        if ((occ_idl + (cost_multiplier * duration)) <= (occ_sat - (duration / 2))):
            self.move_task_request(ts, sat, idl)
            log.append((start, level, ts.key, duration, sat.address, occ_sat, idl.address, occ_idl))
            s.check_idle_saturated(sat, occ=occ_sat)
            s.check_idle_saturated(idl, occ=occ_idl)
    with log_errors():
        i = 0
        idle = s.idle.values()
        saturated = s.saturated
        if ((not idle) or (len(idle) == len(s.workers))):
            return
        log = []
        start = time()
        if (not s.saturated):
            saturated = topk(10, s.workers.values(), key=combined_occupancy)
            saturated = [ws for ws in saturated if ((combined_occupancy(ws) > 0.2) and (len(ws.processing) > ws.nthreads))]
        elif (len(s.saturated) < 20):
            saturated = sorted(saturated, key=combined_occupancy, reverse=True)
        if (len(idle) < 20):
            idle = sorted(idle, key=combined_occupancy)
        for (level, cost_multiplier) in enumerate(self.cost_multipliers):
            if (not idle):
                break
            for sat in list(saturated):
                stealable = self.stealable[sat.address][level]
                if ((not stealable) or (not idle)):
                    continue
                for ts in list(stealable):
                    if ((ts not in self.key_stealable) or (ts.processing_on is not sat)):
                        stealable.discard(ts)
                        continue
                    i += 1
                    if (not idle):
                        break
                    if _has_restrictions(ts):
                        thieves = [ws for ws in idle if _can_steal(ws, ts, sat)]
                    else:
                        thieves = idle
                    if (not thieves):
                        break
                    thief = thieves[(i % len(thieves))]
                    duration = sat.processing.get(ts)
                    if (duration is None):
                        stealable.discard(ts)
                        continue
                    maybe_move_task(level, ts, sat, thief, duration, cost_multiplier)
            if (self.cost_multipliers[level] < 20):
                stealable = self.stealable_all[level]
                for ts in list(stealable):
                    if (not idle):
                        break
                    if (ts not in self.key_stealable):
                        stealable.discard(ts)
                        continue
                    sat = ts.processing_on
                    if (sat is None):
                        stealable.discard(ts)
                        continue
                    if (combined_occupancy(sat) < 0.2):
                        continue
                    if (len(sat.processing) <= sat.nthreads):
                        continue
                    i += 1
                    if _has_restrictions(ts):
                        thieves = [ws for ws in idle if _can_steal(ws, ts, sat)]
                    else:
                        thieves = idle
                    if (not thieves):
                        continue
                    thief = thieves[(i % len(thieves))]
                    duration = sat.processing[ts]
                    maybe_move_task(level, ts, sat, thief, duration, cost_multiplier)
        if log:
            self.log(log)
            self.count += 1
        stop = time()
        if s.digests:
            s.digests['steal-duration'].add((stop - start))