class DaskSchedulerWrapper: # TODO: allow user to customize dask scheduler startup def __init__(self, work_dir): self.logger = logging.getLogger(self.__class__.__name__) self.scheduler = Scheduler(port=const.DASK_SCHEDULER_PORT, dashboard=True, http_prefix="/dask") async def start(self, queue, **kwargs): self.logger.info("Starting DASK Scheduler ... ") self.scheduler = await self.scheduler await self.scheduler.finished() async def stop(self, queue, **kwargs): self.logger.info("Stoping DASK Scheduler ... ") try: for worker in self.scheduler.workers_list(workers=None): try: await self.scheduler.close_worker(worker=worker) except Exception as worker_close_ex: self.logger.error(f"Failed to close worker {worker}", exc_info=worker_close_ex) raise worker_close_ex except Exception as ex: self.logger.error("Failed to close workers", exc_info=ex) try: await self.scheduler.close() except Exception as ex: self.logger.error("Failed to close Dask Scheduler", exc_info=ex)