Repository: spotify/luigi
File: luigi/scheduler.py (Python)

Code before:
            self._reset_orphaned_batch_running_tasks(worker_id)

        locally_pending_tasks = 0
        running_tasks = []
        upstream_table = {}

        greedy_resources = collections.defaultdict(int)
        n_unique_pending = 0

        worker = self._state.get_worker(worker_id)
        if worker.is_trivial_worker(self._state):
            relevant_tasks = worker.get_pending_tasks(self._state)
            used_resources = collections.defaultdict(int)
            greedy_workers = dict()  # If there's no resources, then they can grab any task
        else:
            relevant_tasks = self._state.get_pending_tasks()
            used_resources = self._used_resources()
            activity_limit = time.time() - self._config.worker_disconnect_delay
            active_workers = self._state.get_active_workers(last_get_work_gt=activity_limit)
            greedy_workers = dict((worker.id, worker.info.get('workers', 1))
                                  for worker in active_workers)
        tasks = list(relevant_tasks)
        tasks.sort(key=self._rank, reverse=True)

        for task in tasks:
            in_workers = (assistant and getattr(task, 'runnable', bool(task.workers))) or worker_id in task.workers
            if task.status == RUNNING and in_workers:
                # Return a list of currently running tasks to the client,
                # makes it easier to troubleshoot
                other_worker = self._state.get_worker(task.worker_running)
                more_info = {'task_id': task.id, 'worker': str(other_worker)}
                if other_worker is not None:
                    more_info.update(other_worker.info)
                    running_tasks.append(more_info)

            if task.status == PENDING and in_workers:
                upstream_status = self._upstream_status(task.id, upstream_table)
                if upstream_status != UPSTREAM_DISABLED:
                    locally_pending_tasks += 1
                    if len(task.workers) == 1 and not assistant:
                        n_unique_pending += 1

            if (best_task and batched_params and task.family == best_task.family and
                    len(batched_tasks) < max_batch_size and task.is_batchable() and all(
                    task.params.get(name) == value for name, value in unbatched_params.items()) and
                    self._schedulable(task)):
                for name, params in batched_params.items():
                    params.append(task.params.get(name))
                batched_tasks.append(task)
            if best_task:
                continue

Review comment (refactor):
Did you think about if we were able to refactor out some of this logic to share between this method and `get_work`?


Code after:
        num_pending, num_unique_pending, num_pending_last_scheduled = 0, 0, 0
        running_tasks = []

        upstream_status_table = {}
        for task in worker.get_pending_tasks(self._state):
            if self._upstream_status(task.id, upstream_status_table) == UPSTREAM_DISABLED:
                continue
            if task.status == RUNNING:
                # Return a list of currently running tasks to the client,
                # makes it easier to troubleshoot
                other_worker = self._state.get_worker(task.worker_running)
                if other_worker is not None:
                    more_info = {'task_id': task.id, 'worker': str(other_worker)}
                    more_info.update(other_worker.info)
                    running_tasks.append(more_info)
            elif task.status == PENDING:
                num_pending += 1
                num_unique_pending += int(len(task.workers) == 1)
                num_pending_last_scheduled += int(task.workers.peek(last=True) == worker_id)
        return {
            'n_pending_tasks': num_pending,
            'n_unique_pending': num_unique_pending,
            'n_pending_last_scheduled': num_pending_last_scheduled,
            'worker_state': worker.state,
            'running_tasks': running_tasks,
        }

    @rpc_method(allow_null=False)
    def get_work(self, host=None, assistant=False, current_tasks=None, worker=None, **kwargs):
        # TODO: remove any expired nodes

        # Algo: iterate over all nodes, find the highest priority node no dependencies and available
        # resources.

        # Resource checking looks both at currently available resources and at which resources would
        # be available if all running tasks died and we rescheduled all workers greedily. We do both
        # checks in order to prevent a worker with many low-priority tasks from starving other
        # workers with higher priority tasks that share the same resources.

        # 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()

        assert worker is not None
        worker_id = worker
        worker = self._update_worker(
            worker_id,
            worker_reference={'host': host},
            get_work=True)