Updating code between runs while testing

[esahione]

Hi,

My current setup has Airflow running on Celery for task scheduling and I use Dask as the HPC layer, on a Mesos cluster.

One thing I noticed is that the Dask workers are using cached functions even though Airflow sees the right code.

The code is on a read-only NFS share and it is mounted across every worker node, and the package is installed with the editable option (-e).

The way I am currently trying to solve this issue is with a plugin that follows but I'm not sure why it isn't working as intended. Thoughts? Here's my current code.

class RestartIdleWorkerPlugin(SchedulerPlugin):
    def __init__(self, scheduler):
        self.scheduler = scheduler
        self.worker_to_task = defaultdict(set)
        self.task_to_worker_state = {}

    def transition(self, key, start, finish, *args, **kwargs):
        if finish not in ('forgotten', 'memory'):
            return
        if finish == 'memory':
            worker = kwargs['worker']
            worker_state = self.scheduler.workers[worker]
            self.task_to_worker_state[key] = worker_state
            self.worker_to_task[kwargs['worker']].add(key)
            return

        if finish == 'forgotten':
            worker = self.task_to_worker_state[key]
            del self.task_to_worker_state[key]
            if len(worker.processing) == 0 and len(worker.has_what) == 0:
                self._restart_worker(worker)

    def _restart_worker(self, worker, timeout=3):
        """ Restart worker.  Reset local state. """
        logger.info(f'Plugin is restarting worker {worker}')
        nanny_address = self.scheduler.get_worker_service_addr(worker, 'nanny')
        nanny = rpc(nanny_address, connection_args=self.scheduler.connection_args)
        nanny.restart(close=True, timeout=timeout * 0.8, executor_wait=True)
        nanny.close_rpc()

[esahione]

Hey guys,

I'm closing the thread. A few bugfixes and it works as intended. You may want to add a bit of a delay between no memory and rebooting the thread, it may still be transferring the data somewhere.

Also, you may want to catch exceptions properly.

Otherwise it works :p

Here's the code in case anyone is interested:
-- edit: didn't want anyone to use or see my shitty code so here's a decent version --

class KillWorkerWithEmptyTaskDelayedCheck(Thread):
    """
    Tries worker 5 times with the space delineated by delay. If at any point the worker has a
    task number greater than the task_limit, the thread terminates.
    Otherwise it goes on to restart the worker.
    """
    def __init__(self, scheduler, worker, kill_switch, delay=1.5, num_tries=25, task_limit=0):
        super().__init__()
        self.scheduler = scheduler
        self.worker = worker
        self.delay = delay
        self.num_tries = num_tries
        self.current_try = 0
        self.failures = 0
        self.task_limit = task_limit
        self.nanny_addr = self.scheduler.get_worker_service_addr(self.worker.address, 'nanny')
        self.kill_switch = kill_switch

    @property
    def successes(self):
        return self.current_try - self.failures

    def run(self):
        while not self.kill_switch.is_set():
            self.current_try += 1
            time.sleep(self.delay)
            if not self.worker.processing:
                self.failures += 1
            if self.successes > self.task_limit:
                self.current_try = 0
            if self.current_try >= self.num_tries:
                print(f'Plugin is restarting {str(self.worker)} with nanny @ {self.nanny_addr}')
                nanny = rpc(self.nanny_addr, connection_args=self.scheduler.connection_args)
                nanny.restart(close=True, timeout=3, executor_wait=True)
                nanny.close_rpc()
                self.current_try = 0


class RestartPlugin(SchedulerPlugin):
    def __init__(self, scheduler, num_hits=20):
        self.scheduler = scheduler
        self.current_hit_number = 0
        self.num_hits_before_cleanup = num_hits
        self.workers = {}

    def remove_worker(self, scheduler=None, worker=None):
        worker, killswitch, killthread = self.workers[worker]
        killswitch.set()
        killthread.join()
        del self.workers[worker]

    def add_worker(self, scheduler=None, worker=None, **kwargs):
        for ws in scheduler.workers.values():
            if ws.address == worker:
                kill_switch = threading.Event()
                killthread = KillWorkerWithEmptyTaskDelayedCheck(scheduler, ws, kill_switch)
                killthread.start()
                self.workers[worker] = (ws, kill_switch, killthread)