diff --git a/dask_drmaa/core.py b/dask_drmaa/core.py index 18fe3d0..d746236 100644 --- a/dask_drmaa/core.py +++ b/dask_drmaa/core.py @@ -290,12 +290,15 @@ def stop_workers(self, worker_ids, sync=False): get_session().synchronize(worker_ids, dispose=True) @gen.coroutine - def scale_up(self, n, **kwargs): + def _scale_up(self, n, **kwargs): yield [self.start_workers(**kwargs) for _ in range(n - len(self.workers))] + def scale_up(self, n, **kwargs): + self.scheduler.loop.add_callback(self._scale_up, n, **kwargs) + @gen.coroutine - def scale_down(self, worker_ids): + def _scale_down(self, worker_ids): worker_ids = list(set(worker_ids)) for wid in worker_ids: @@ -312,6 +315,9 @@ def scale_down(self, worker_ids): logger.info("Scaling down workers %s", worker_ids) + def scale_down(self, worker_ids): + self.scheduler.loop.add_callback(self._scale_down, worker_ids) + def close(self): logger.info("Closing DRMAA cluster") self.stop_workers(self.workers, sync=True)