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)
