async def dask_setup(worker: distributed.Worker) -> None: """This is a special function recognized by the dask worker when starting with flag --preload""" settings = Settings.create_from_envs() # set up logging logging.basicConfig(level=settings.LOG_LEVEL.value) logging.root.setLevel(level=settings.LOG_LEVEL.value) logger.setLevel(level=settings.LOG_LEVEL.value) logger.info("Setting up worker...") logger.info("Settings: %s", pformat(settings.dict())) print_banner() if threading.current_thread() is threading.main_thread(): loop = asyncio.get_event_loop() logger.info("We do have a running loop in the main thread: %s", f"{loop=}") if threading.current_thread() is threading.main_thread(): GracefulKiller(worker)
def test_worker_bad_args(c, s, a, b): class NoReprObj(object): """ This object cannot be properly represented as a string. """ def __str__(self): raise ValueError("I have no str representation.") def __repr__(self): raise ValueError("I have no repr representation.") x = c.submit(NoReprObj, workers=a.address) yield wait(x) assert not a.executing assert a.data def bad_func(*args, **kwargs): 1 / 0 class MockLoggingHandler(logging.Handler): """Mock logging handler to check for expected logs.""" def __init__(self, *args, **kwargs): self.reset() logging.Handler.__init__(self, *args, **kwargs) def emit(self, record): self.messages[record.levelname.lower()].append(record.getMessage()) def reset(self): self.messages = { 'debug': [], 'info': [], 'warning': [], 'error': [], 'critical': [], } hdlr = MockLoggingHandler() old_level = logger.level logger.setLevel(logging.DEBUG) logger.addHandler(hdlr) y = c.submit(bad_func, x, k=x, workers=b.address) yield wait(y) assert not b.executing assert y.status == 'error' # Make sure job died because of bad func and not because of bad # argument. with pytest.raises(ZeroDivisionError): yield y if sys.version_info[0] >= 3: tb = yield y._traceback() assert any('1 / 0' in line for line in pluck(3, traceback.extract_tb(tb)) if line) assert "Compute Failed" in hdlr.messages['warning'][0] logger.setLevel(old_level) # Now we check that both workers are still alive. xx = c.submit(add, 1, 2, workers=a.address) yy = c.submit(add, 3, 4, workers=b.address) results = yield c._gather([xx, yy]) assert tuple(results) == (3, 7)
def test_worker_bad_args(c, a, b): aa = rpc(ip=a.ip, port=a.port) bb = rpc(ip=b.ip, port=b.port) class NoReprObj(object): """ This object cannot be properly represented as a string. """ def __str__(self): raise ValueError("I have no str representation.") def __repr__(self): raise ValueError("I have no repr representation.") response = yield aa.compute(key='x', function=dumps(NoReprObj), args=dumps(()), who_has={}) assert not a.active assert response['status'] == 'OK' assert a.data['x'] assert isinstance(response['compute_start'], float) assert isinstance(response['compute_stop'], float) assert isinstance(response['thread'], Integral) def bad_func(*args, **kwargs): 1 / 0 class MockLoggingHandler(logging.Handler): """Mock logging handler to check for expected logs.""" def __init__(self, *args, **kwargs): self.reset() logging.Handler.__init__(self, *args, **kwargs) def emit(self, record): self.messages[record.levelname.lower()].append(record.getMessage()) def reset(self): self.messages = { 'debug': [], 'info': [], 'warning': [], 'error': [], 'critical': [], } hdlr = MockLoggingHandler() old_level = logger.level logger.setLevel(logging.DEBUG) logger.addHandler(hdlr) response = yield bb.compute(key='y', function=dumps(bad_func), args=dumps(['x']), kwargs=dumps({'k': 'x'}), who_has={'x': [a.address]}) assert not b.active assert response['status'] == 'error' # Make sure job died because of bad func and not because of bad # argument. assert isinstance(loads(response['exception']), ZeroDivisionError) if sys.version_info[0] >= 3: assert any('1 / 0' in line for line in pluck(3, traceback.extract_tb( loads(response['traceback']))) if line) assert hdlr.messages['warning'][0] == " Compute Failed\n" \ "Function: bad_func\n" \ "args: (< could not convert arg to str >)\n" \ "kwargs: {'k': < could not convert arg to str >}\n" assert re.match(r"^Send compute response to scheduler: y, " \ "\{.*'args': \(< could not convert arg to str >\), .*" \ "'kwargs': \{'k': < could not convert arg to str >\}.*\}", hdlr.messages['debug'][0]) or \ re.match("^Send compute response to scheduler: y, " \ "\{.*'kwargs': \{'k': < could not convert arg to str >\}, .*" \ "'args': \(< could not convert arg to str >\).*\}", hdlr.messages['debug'][0]) logger.setLevel(old_level) # Now we check that both workers are still alive. assert not a.active response = yield aa.compute(key='z', function=dumps(add), args=dumps([1, 2]), who_has={}, close=True) assert not a.active assert response['status'] == 'OK' assert a.data['z'] == 3 assert isinstance(response['compute_start'], float) assert isinstance(response['compute_stop'], float) assert isinstance(response['thread'], Integral) assert not b.active response = yield bb.compute(key='w', function=dumps(add), args=dumps([1, 2]), who_has={}, close=True) assert not b.active assert response['status'] == 'OK' assert b.data['w'] == 3 assert isinstance(response['compute_start'], float) assert isinstance(response['compute_stop'], float) assert isinstance(response['thread'], Integral) aa.close_streams() bb.close_streams()