class MalcolmServerConnect:
def __init__(self, socket):
self.socket = socket
self.request = None
self.response = None
self._loop = IOLoop()
@gen.coroutine
def message_coroutine(self):
conn = yield websocket_connect("ws://localhost:%s/ws" % self.socket)
msg = json.dumps(self.request)
conn.write_message(msg)
resp = yield conn.read_message()
self.response = json.loads(resp, object_pairs_hook=OrderedDict)
conn.close()
def send_subscribe(self, path, delta):
self.request = OrderedDict()
self.request['typeid'] = "malcolm:core/Subscribe:1.0"
self.request['id'] = "CANNED"
self.request['path'] = path
self.request['delta'] = delta
self._loop.run_sync(self.message_coroutine)
if delta:
for change in self.response['changes']:
if len(change) == 2 and isinstance(change[1], dict) and \
"timeStamp" in change[1]:
for endpoint in ("secondsPastEpoch", "nanoseconds"):
change[1]['timeStamp'][endpoint] = 0
def subscribe_request_response(self, path, delta=True):
self.send_subscribe(path, delta)
request_msg = json.dumps(self.request, indent=2)
response_msg = json.dumps(self.response, indent=2)
return request_msg, response_msg
def scheduler(cport, **kwargs):
q = Queue()
proc = Process(target=run_scheduler, args=(q, cport), kwargs=kwargs)
proc.daemon = True
proc.start()
sport = q.get()
s = rpc(ip='127.0.0.1', port=sport)
loop = IOLoop()
start = time()
while True:
ncores = loop.run_sync(s.ncores)
if ncores:
break
if time() - start > 5:
raise Exception("Timeout on cluster creation")
sleep(0.01)
try:
yield sport
finally:
loop = IOLoop()
with ignoring(socket.error, TimeoutError, StreamClosedError):
loop.run_sync(lambda: disconnect('127.0.0.1', sport), timeout=0.5)
proc.terminate()
def test_instance_pool(self):
instance1 = _Instance()
instance2 = _Instance()
pool = InstancePool([instance1, instance2])
pool.increment()
pool.increment()
self.assertEquals(instance1.value(), 1)
self.assertEquals(instance2.value(), 1)
pool.transaction(lambda i: i.increment())
pool.transaction(lambda i: i.increment())
self.assertEquals(instance1.value(), 2)
self.assertEquals(instance2.value(), 2)
@coroutine
def yield_tasks():
self.assertEquals((yield pool.await().increment()), 3)
self.assertEquals((yield pool.await().increment()), 3)
self.assertEquals(instance1.value(), 3)
self.assertEquals(instance2.value(), 3)
self.assertEquals((yield pool.await_transaction(lambda i: i.increment())), 4)
self.assertEquals((yield pool.await_transaction(lambda i: i.increment())), 4)
loop = IOLoop()
loop.make_current()
loop.run_sync(yield_tasks)
self.assertEquals(instance1.value(), 4)
self.assertEquals(instance2.value(), 4)
class SyncHTTPClientTest(unittest.TestCase):
def setUp(self):
self.server_ioloop = IOLoop()
event = threading.Event()
@gen.coroutine
def init_server():
sock, self.port = bind_unused_port()
app = Application([("/", HelloWorldHandler)])
self.server = HTTPServer(app)
self.server.add_socket(sock)
event.set()
def start():
self.server_ioloop.run_sync(init_server)
self.server_ioloop.start()
self.server_thread = threading.Thread(target=start)
self.server_thread.start()
event.wait()
self.http_client = HTTPClient()
def tearDown(self):
def stop_server():
self.server.stop()
# Delay the shutdown of the IOLoop by several iterations because
# the server may still have some cleanup work left when
# the client finishes with the response (this is noticeable
# with http/2, which leaves a Future with an unexamined
# StreamClosedError on the loop).
@gen.coroutine
def slow_stop():
# The number of iterations is difficult to predict. Typically,
# one is sufficient, although sometimes it needs more.
for i in range(5):
yield
self.server_ioloop.stop()
self.server_ioloop.add_callback(slow_stop)
self.server_ioloop.add_callback(stop_server)
self.server_thread.join()
self.http_client.close()
self.server_ioloop.close(all_fds=True)
def get_url(self, path):
return "http://127.0.0.1:%d%s" % (self.port, path)
def test_sync_client(self):
response = self.http_client.fetch(self.get_url("/"))
self.assertEqual(b"Hello world!", response.body)
def test_sync_client_error(self):
# Synchronous HTTPClient raises errors directly; no need for
# response.rethrow()
with self.assertRaises(HTTPError) as assertion:
self.http_client.fetch(self.get_url("/notfound"))
self.assertEqual(assertion.exception.code, 404)
def cluster(nworkers=2, nanny=False, worker_kwargs={}):
if nanny:
_run_worker = run_nanny
else:
_run_worker = run_worker
scheduler_q = Queue()
scheduler = Process(target=run_scheduler, args=(scheduler_q,))
scheduler.daemon = True
scheduler.start()
sport = scheduler_q.get()
workers = []
for i in range(nworkers):
q = Queue()
fn = '_test_worker-%s' % uuid.uuid1()
proc = Process(target=_run_worker, args=(q, sport),
kwargs=merge({'ncores': 1, 'local_dir': fn},
worker_kwargs))
workers.append({'proc': proc, 'queue': q, 'dir': fn})
for worker in workers:
worker['proc'].start()
for worker in workers:
worker['port'] = worker['queue'].get()
loop = IOLoop()
s = rpc(ip='127.0.0.1', port=sport)
start = time()
try:
while True:
ncores = loop.run_sync(s.ncores)
if len(ncores) == nworkers:
break
if time() - start > 5:
raise Exception("Timeout on cluster creation")
yield {'proc': scheduler, 'port': sport}, workers
finally:
logger.debug("Closing out test cluster")
with ignoring(socket.error, TimeoutError, StreamClosedError):
loop.run_sync(lambda: disconnect('127.0.0.1', sport), timeout=0.5)
scheduler.terminate()
scheduler.join(timeout=2)
for port in [w['port'] for w in workers]:
with ignoring(socket.error, TimeoutError, StreamClosedError):
loop.run_sync(lambda: disconnect('127.0.0.1', port),
timeout=0.5)
for proc in [w['proc'] for w in workers]:
with ignoring(Exception):
proc.terminate()
proc.join(timeout=2)
for q in [w['queue'] for w in workers]:
q.close()
for fn in glob('_test_worker-*'):
shutil.rmtree(fn)
loop.close(all_fds=True)
def atexit(self):
"""atexit callback"""
if self._atexit_ran:
return
self._atexit_ran = True
# run the cleanup step (in a new loop, because the interrupted one is unclean)
IOLoop.clear_current()
loop = IOLoop()
loop.make_current()
loop.run_sync(self.cleanup)
def _atexit(self):
if self._atexit_ran:
return
self._atexit_ran = True
self._stats_job.stop()
IOLoop.clear_current()
loop = IOLoop()
loop.make_current()
loop.run_sync(self._cleanup)
def run_nanny(port, center_port, **kwargs):
from distributed import Nanny
from tornado.ioloop import IOLoop, PeriodicCallback
import logging
IOLoop.clear_instance()
loop = IOLoop(); loop.make_current()
PeriodicCallback(lambda: None, 500).start()
logging.getLogger("tornado").setLevel(logging.CRITICAL)
worker = Nanny('127.0.0.1', port, port + 1000, '127.0.0.1', center_port, **kwargs)
loop.run_sync(worker._start)
loop.start()
def test_func():
IOLoop.clear_instance()
loop = IOLoop()
loop.make_current()
cor = gen.coroutine(func)
try:
loop.run_sync(cor, timeout=timeout)
finally:
loop.stop()
loop.close(all_fds=True)
def test_func():
IOLoop.clear_instance()
loop = IOLoop()
loop.make_current()
s, workers = loop.run_sync(lambda: start_cluster(ncores))
try:
loop.run_sync(lambda: cor(s, *workers), timeout=timeout)
finally:
loop.run_sync(lambda: end_cluster(s, workers))
loop.stop()
loop.close()
class HTTPClient(object):
"""A blocking HTTP client.
This interface is provided for convenience and testing; most applications
that are running an IOLoop will want to use `AsyncHTTPClient` instead.
Typical usage looks like this::
http_client = httpclient.HTTPClient()
try:
response = http_client.fetch("http://www.google.com/")
print(response.body)
except httpclient.HTTPError as e:
# HTTPError is raised for non-200 responses; the response
# can be found in e.response.
print("Error: " + str(e))
except Exception as e:
# Other errors are possible, such as IOError.
print("Error: " + str(e))
http_client.close()
"""
def __init__(self, async_client_class=None, **kwargs):
self._io_loop = IOLoop(make_current=False)
if async_client_class is None:
async_client_class = AsyncHTTPClient
# Create the client while our IOLoop is "current", without
# clobbering the thread's real current IOLoop (if any).
self._async_client = self._io_loop.run_sync(
gen.coroutine(lambda: async_client_class(**kwargs)))
self._closed = False
def __del__(self):
self.close()
def close(self):
"""Closes the HTTPClient, freeing any resources used."""
if not self._closed:
self._async_client.close()
self._io_loop.close()
self._closed = True
def fetch(self, request, **kwargs):
"""Executes a request, returning an `HTTPResponse`.
The request may be either a string URL or an `HTTPRequest` object.
If it is a string, we construct an `HTTPRequest` using any additional
kwargs: ``HTTPRequest(request, **kwargs)``
If an error occurs during the fetch, we raise an `HTTPError` unless
the ``raise_error`` keyword argument is set to False.
"""
response = self._io_loop.run_sync(functools.partial(
self._async_client.fetch, request, **kwargs))
return response
def run_worker(q, center_port, **kwargs):
from distributed import Worker
from tornado.ioloop import IOLoop, PeriodicCallback
import logging
with log_errors():
IOLoop.clear_instance()
loop = IOLoop(); loop.make_current()
PeriodicCallback(lambda: None, 500).start()
logging.getLogger("tornado").setLevel(logging.CRITICAL)
worker = Worker('127.0.0.1', center_port, ip='127.0.0.1', **kwargs)
loop.run_sync(lambda: worker._start(0))
q.put(worker.port)
loop.start()
class SyncHTTPClientTest(unittest.TestCase):
def setUp(self):
if IOLoop.configured_class().__name__ == 'TwistedIOLoop':
# TwistedIOLoop only supports the global reactor, so we can't have
# separate IOLoops for client and server threads.
raise unittest.SkipTest(
'Sync HTTPClient not compatible with TwistedIOLoop')
self.server_ioloop = IOLoop()
@gen.coroutine
def init_server():
sock, self.port = bind_unused_port()
app = Application([('/', HelloWorldHandler)])
self.server = HTTPServer(app)
self.server.add_socket(sock)
self.server_ioloop.run_sync(init_server)
self.server_thread = threading.Thread(target=self.server_ioloop.start)
self.server_thread.start()
self.http_client = HTTPClient()
def tearDown(self):
def stop_server():
self.server.stop()
# Delay the shutdown of the IOLoop by one iteration because
# the server may still have some cleanup work left when
# the client finishes with the response (this is noticeable
# with http/2, which leaves a Future with an unexamined
# StreamClosedError on the loop).
self.server_ioloop.add_callback(self.server_ioloop.stop)
self.server_ioloop.add_callback(stop_server)
self.server_thread.join()
self.http_client.close()
self.server_ioloop.close(all_fds=True)
def get_url(self, path):
return 'http://127.0.0.1:%d%s' % (self.port, path)
def test_sync_client(self):
response = self.http_client.fetch(self.get_url('/'))
self.assertEqual(b'Hello world!', response.body)
def test_sync_client_error(self):
# Synchronous HTTPClient raises errors directly; no need for
# response.rethrow()
with self.assertRaises(HTTPError) as assertion:
self.http_client.fetch(self.get_url('/notfound'))
self.assertEqual(assertion.exception.code, 404)
def run_worker(q, scheduler_port, **kwargs):
from distributed import Worker
from tornado.ioloop import IOLoop, PeriodicCallback
with log_errors():
IOLoop.clear_instance()
loop = IOLoop(); loop.make_current()
PeriodicCallback(lambda: None, 500).start()
worker = Worker('127.0.0.1', scheduler_port, ip='127.0.0.1',
loop=loop, validate=True, **kwargs)
loop.run_sync(lambda: worker._start(0))
q.put(worker.port)
try:
loop.start()
finally:
loop.close(all_fds=True)
def single_run_ws(url, delay=0, size=0, msgs=1):
"""Time a single websocket run"""
buf = hexlify(os.urandom(size // 2)).decode('ascii')
msg = json.dumps({'delay': delay, 'data': buf})
async def go():
ws = await websocket_connect(url.replace('http', 'ws') + '/ws')
for i in range(msgs):
ws.write_message(msg)
await ws.read_message()
asyncio.set_event_loop(asyncio.new_event_loop())
IOLoop.clear_current()
loop = IOLoop(make_current=True)
loop.run_sync(go)
def test_awaitable(self):
instance = _Instance()
instance.increment()
self.assertEquals(instance.value(), 1)
awaitable = AwaitableInstance(instance)
@coroutine
def yield_tasks():
self.assertEquals((yield awaitable.increment()), 2)
self.assertEquals((yield awaitable.increment()), 3)
self.assertEquals((yield awaitable.increment()), 4)
self.assertEquals((yield awaitable.value()), 4)
loop = IOLoop()
loop.make_current()
loop.run_sync(yield_tasks)
self.assertEquals(instance.value(), 4)
def _execute_cmd(self, host, command, params):
url = self.START_URL_TPL.format(host=self._wrap_host(host), port=self.minion_port)
data = self._update_query_parameters(
dst={'command': command},
src=params,
)
io_loop = IOLoop()
client = SimpleAsyncHTTPClient(io_loop)
logger.debug(
'ioloop for single async http request for host {} started: {}'.format(
host, command))
response = io_loop.run_sync(functools.partial(
client.fetch, url, method='POST',
headers=self.minion_headers,
body=urllib.urlencode(data, doseq=True),
request_timeout=MINIONS_CFG.get('request_timeout', 5.0),
allow_ipv6=True,
use_gzip=True))
logger.debug(
'ioloop for single async http request for host {} finished: {}'.format(
host, command))
try:
data = self._process_state(host, self._get_response(host, response), sync=True)
except HTTPError:
logger.exception('Execute cmd raised http error')
raise
return data
def run_nanny(q, center_port, **kwargs):
from distributed import Nanny
from tornado.ioloop import IOLoop, PeriodicCallback
import logging
with log_errors():
IOLoop.clear_instance()
loop = IOLoop(); loop.make_current()
PeriodicCallback(lambda: None, 500).start()
logging.getLogger("tornado").setLevel(logging.CRITICAL)
worker = Nanny('127.0.0.1', center_port, ip='127.0.0.1', loop=loop, **kwargs)
loop.run_sync(lambda: worker._start(0))
q.put(worker.port)
try:
loop.start()
finally:
loop.run_sync(worker._close)
loop.close(all_fds=True)
def websocket(self, id):
loop = IOLoop()
req = HTTPRequest(
url_path_join(self.base_url.replace('http', 'ws', 1), 'api/kernels', id, 'channels'),
headers=self.headers,
)
f = websocket_connect(req, io_loop=loop)
return loop.run_sync(lambda : f)
def test_yield_task(self):
queue = TaskQueue()
task_results = []
@coroutine
def yield_tasks():
task = lambda x: x
futures = []
futures.append(queue.yield_task(task, 1))
futures.append(queue.yield_task(task, 2))
futures.append(queue.yield_task(task, 3))
res = yield futures
task_results[:] = res
loop = IOLoop()
loop.make_current()
loop.run_sync(yield_tasks)
self.assertEquals(len(task_results), 3)
self.assertEquals(task_results, [1, 2, 3])
class ProxyTestServer(object):
def __init__(self,):
self.server_ioloop = IOLoop()
self.access_count = 0
@tornado.gen.coroutine
def init_server():
sock, self.port = bind_unused_port()
app = Application([('/(.*)', ProxyTestHandler, dict(server=self))])
self.server = HTTPServer(app)
self.server.add_socket(sock)
self.server_ioloop.run_sync(init_server)
self.server_thread = threading.Thread(target=self.server_ioloop.start)
def start(self):
self.server_thread.start()
def stop(self):
def stop_server():
self.server.stop()
@tornado.gen.coroutine
def slow_stop():
for i in range(5):
yield
self.server_ioloop.stop()
self.server_ioloop.add_callback(slow_stop)
self.server_ioloop.add_callback(stop_server)
self.server_thread.join()
self.server_ioloop.close(all_fds=True)
def get_access_count(self):
return self.access_count
def clear_access_count(self):
self.access_count = 0
def cluster(nworkers=2, nanny=False):
if nanny:
_run_worker = run_nanny
else:
_run_worker = run_worker
_port[0] += 1
cport = _port[0]
center = Process(target=run_center, args=(cport,))
workers = []
for i in range(nworkers):
_port[0] += 1
port = _port[0]
proc = Process(target=_run_worker, args=(port, cport),
kwargs={'ncores': 1, 'local_dir': '_test_worker-%d' % port})
workers.append({'port': port, 'proc': proc})
center.start()
for worker in workers:
worker['proc'].start()
sock = connect_sync('127.0.0.1', cport)
start = time()
try:
while True:
write_sync(sock, {'op': 'ncores'})
ncores = read_sync(sock)
if len(ncores) == nworkers:
break
if time() - start > 5:
raise Exception("Timeout on cluster creation")
yield {'proc': center, 'port': cport}, workers
finally:
loop = IOLoop()
logger.debug("Closing out test cluster")
for port in [cport] + [w['port'] for w in workers]:
with ignoring(socket.error, TimeoutError, StreamClosedError):
loop.run_sync(lambda: disconnect('127.0.0.1', port), timeout=10)
for proc in [center] + [w['proc'] for w in workers]:
with ignoring(Exception):
proc.terminate()
for fn in glob('_test_worker-*'):
shutil.rmtree(fn)
def run_scheduler(q, center_port=None, **kwargs):
from distributed import Scheduler
from tornado.ioloop import IOLoop, PeriodicCallback
import logging
IOLoop.clear_instance()
loop = IOLoop(); loop.make_current()
PeriodicCallback(lambda: None, 500).start()
logging.getLogger("tornado").setLevel(logging.CRITICAL)
center = ('127.0.0.1', center_port) if center_port else None
scheduler = Scheduler(center=center, **kwargs)
scheduler.listen(0)
if center_port:
loop.run_sync(scheduler.sync_center)
done = scheduler.start(0)
q.put(scheduler.port)
loop.start()
def test_yield_task_exception(self):
queue = TaskQueue()
task_results = []
@coroutine
def yield_tasks():
def task(x):
raise Exception('failure')
futures = []
futures.append(queue.yield_task(task, 1))
futures.append(queue.yield_task(task, 2))
futures.append(queue.yield_task(task, 3))
for f in futures:
try:
yield f
except Exception as e:
task_results.append(e)
loop = IOLoop()
loop.make_current()
loop.run_sync(yield_tasks)
self.assertEquals(len(task_results), 3)
for e in task_results:
self.assertEquals(e.message, 'failure')
def run_worker_fork(q, ip, scheduler_ip, scheduler_port, ncores, nanny_port,
worker_port, local_dir, services, name, memory_limit, reconnect,
resources, validate):
""" Function run by the Nanny when creating the worker """
from distributed import Worker # pragma: no cover
from tornado.ioloop import IOLoop # pragma: no cover
IOLoop.clear_instance() # pragma: no cover
loop = IOLoop() # pragma: no cover
loop.make_current() # pragma: no cover
worker = Worker(scheduler_ip, scheduler_port, ncores=ncores, ip=ip,
service_ports={'nanny': nanny_port}, local_dir=local_dir,
services=services, name=name, memory_limit=memory_limit,
reconnect=reconnect, validate=validate,
resources=resources, loop=loop) # pragma: no cover
@gen.coroutine # pragma: no cover
def run():
try: # pragma: no cover
yield worker._start(worker_port) # pragma: no cover
except Exception as e: # pragma: no cover
logger.exception(e) # pragma: no cover
q.put(e) # pragma: no cover
else:
assert worker.port # pragma: no cover
q.put({'port': worker.port, 'dir': worker.local_dir}) # pragma: no cover
while worker.status != 'closed':
yield gen.sleep(0.1)
logger.info("Worker closed")
try:
loop.run_sync(run)
finally:
loop.stop()
loop.close(all_fds=True)
def test_func():
IOLoop.clear_instance()
loop = IOLoop()
loop.make_current()
s, workers = loop.run_sync(lambda: start_cluster(ncores,
Worker=Worker))
args = [s] + workers
if executor:
e = Executor((s.ip, s.port), loop=loop, start=False)
loop.run_sync(e._start)
args = [e] + args
try:
loop.run_sync(lambda: cor(*args), timeout=timeout)
finally:
if executor:
loop.run_sync(e._shutdown)
loop.run_sync(lambda: end_cluster(s, workers))
loop.stop()
loop.close(all_fds=True)
def _terminate_cmd(self, host, uid):
url = self.TERMINATE_URL_TPL.format(host=host, port=self.minion_port)
data = {'cmd_uid': uid}
io_loop = IOLoop()
client = SimpleAsyncHTTPClient(io_loop)
response = io_loop.run_sync(functools.partial(
client.fetch, url, method='POST',
headers=self.minion_headers,
body=urllib.urlencode(data),
request_timeout=MINIONS_CFG.get('request_timeout', 5.0),
allow_ipv6=True,
use_gzip=True))
data = self._process_state(host, self._get_response(host, response), sync=True)
return data
class HTTPClient(object):
"""A blocking HTTP client.
This interface is provided for convenience and testing; most applications
that are running an IOLoop will want to use `AsyncHTTPClient` instead.
Typical usage looks like this::
http_client = httpclient.HTTPClient()
try:
response = http_client.fetch("http://www.google.com/")
print response.body
except httpclient.HTTPError as e:
print "Error:", e
http_client.close()
"""
def __init__(self, async_client_class=None, **kwargs):
self._io_loop = IOLoop()
if async_client_class is None:
async_client_class = AsyncHTTPClient
self._async_client = async_client_class(self._io_loop, **kwargs)
self._closed = False
def __del__(self):
self.close()
def close(self):
"""Closes the HTTPClient, freeing any resources used."""
if not self._closed:
self._async_client.close()
self._io_loop.close()
self._closed = True
def fetch(self, request, **kwargs):
"""Executes a request, returning an `HTTPResponse`.
The request may be either a string URL or an `HTTPRequest` object.
If it is a string, we construct an `HTTPRequest` using any additional
kwargs: ``HTTPRequest(request, **kwargs)``
If an error occurs during the fetch, we raise an `HTTPError`.
"""
response = self._io_loop.run_sync(functools.partial(
self._async_client.fetch, request, **kwargs))
response.rethrow()
return response
def _execute_cmd(self, host, command, params):
url = self.START_URL_TPL.format(host=host, port=self.minion_port)
data = {'command': command}
params_rename = {
'success_codes': 'success_code'
}
def check_param(key, val):
if not isinstance(val, (int, long, basestring)):
logger.warn('Failed parameter {0}, value {1}'.format(key, val))
raise ValueError('Only strings are accepted as command parameters')
for k, v in params.iteritems():
if k == 'command':
raise ValueError('Parameter "command" is not accepted as command parameter')
if isinstance(v, (list, tuple)):
for val in v:
check_param(k, val)
else:
check_param(k, v)
data[params_rename.get(k, k)] = v
io_loop = IOLoop()
client = SimpleAsyncHTTPClient(io_loop)
logger.debug(
'ioloop for single async http request for host {} started: {}'.format(
host, command))
response = io_loop.run_sync(functools.partial(
client.fetch, url, method='POST',
headers=self.minion_headers,
body=urllib.urlencode(data, doseq=True),
request_timeout=MINIONS_CFG.get('request_timeout', 5.0),
allow_ipv6=True,
use_gzip=True))
logger.debug(
'ioloop for single async http request for host {} finished: {}'.format(
host, command))
try:
data = self._process_state(host, self._get_response(host, response), sync=True)
except HTTPError:
logger.exception('Execute cmd raised http error')
raise
return data
class TestIOLoopRunSync(unittest.TestCase):
def setUp(self):
self.io_loop = IOLoop()
def tearDown(self):
self.io_loop.close()
def test_sync_result(self):
self.assertEqual(self.io_loop.run_sync(lambda: 42), 42)
def test_sync_exception(self):
with self.assertRaises(ZeroDivisionError):
self.io_loop.run_sync(lambda: 1 / 0)
def test_async_result(self):
@gen.coroutine
def f():
yield gen.Task(self.io_loop.add_callback)
raise gen.Return(42)
self.assertEqual(self.io_loop.run_sync(f), 42)
def test_async_exception(self):
@gen.coroutine
def f():
yield gen.Task(self.io_loop.add_callback)
1 / 0
with self.assertRaises(ZeroDivisionError):
self.io_loop.run_sync(f)
def test_current(self):
def f():
self.assertIs(IOLoop.current(), self.io_loop)
self.io_loop.run_sync(f)
def test_timeout(self):
@gen.coroutine
def f():
yield gen.Task(self.io_loop.add_timeout, self.io_loop.time() + 1)
self.assertRaises(TimeoutError, self.io_loop.run_sync, f, timeout=0.01)
class TestIOLoopRunSync(unittest.TestCase):
def setUp(self):
self.io_loop = IOLoop()
def tearDown(self):
self.io_loop.close()
def test_sync_result(self):
with self.assertRaises(gen.BadYieldError):
self.io_loop.run_sync(lambda: 42)
def test_sync_exception(self):
with self.assertRaises(ZeroDivisionError):
self.io_loop.run_sync(lambda: 1 / 0)
def test_async_result(self):
@gen.coroutine
def f():
yield gen.moment
raise gen.Return(42)
self.assertEqual(self.io_loop.run_sync(f), 42)
def test_async_exception(self):
@gen.coroutine
def f():
yield gen.moment
1 / 0
with self.assertRaises(ZeroDivisionError):
self.io_loop.run_sync(f)
def test_current(self):
def f():
self.assertIs(IOLoop.current(), self.io_loop)
self.io_loop.run_sync(f)
def test_timeout(self):
@gen.coroutine
def f():
yield gen.sleep(1)
self.assertRaises(TimeoutError, self.io_loop.run_sync, f, timeout=0.01)
@skipBefore35
def test_native_coroutine(self):
@gen.coroutine
def f1():
yield gen.moment
namespace = exec_test(globals(), locals(), """
async def f2():
await f1()
""")
self.io_loop.run_sync(namespace['f2'])
def test_open_sync(self):
loop = IOLoop()
cx = loop.run_sync(self.motor_client(io_loop=loop).open)
self.assertTrue(isinstance(cx, motor.MotorClient))
def _run(cls, worker_args, worker_kwargs, worker_start_args, silence_logs,
init_result_q, child_stop_q): # pragma: no cover
from distributed import Worker
try:
from dask.multiprocessing import initialize_worker_process
except ImportError: # old Dask version
pass
else:
initialize_worker_process()
if silence_logs:
logger.setLevel(silence_logs)
IOLoop.clear_instance()
loop = IOLoop()
loop.make_current()
worker = Worker(*worker_args, **worker_kwargs)
@gen.coroutine
def do_stop(timeout):
try:
yield worker._close(report=False, nanny=False)
finally:
loop.stop()
def watch_stop_q():
"""
Wait for an incoming stop message and then stop the
worker cleanly.
"""
while True:
try:
msg = child_stop_q.get(timeout=1000)
except Empty:
pass
else:
assert msg['op'] == 'stop'
loop.add_callback(do_stop, msg['timeout'])
break
t = threading.Thread(target=watch_stop_q,
name="Nanny stop queue watch")
t.daemon = True
t.start()
@gen.coroutine
def run():
"""
Try to start worker and inform parent of outcome.
"""
try:
yield worker._start(*worker_start_args)
except Exception as e:
logger.exception("Failed to start worker")
init_result_q.put(e)
else:
assert worker.address
init_result_q.put({
'address': worker.address,
'dir': worker.local_dir
})
yield worker.wait_until_closed()
logger.info("Worker closed")
try:
loop.run_sync(run)
except TimeoutError:
# Loop was stopped before wait_until_closed() returned, ignore
pass
except KeyboardInterrupt:
pass
class HTTPClient(object):
"""A blocking HTTP client.
This interface is provided to make it easier to share code between
synchronous and asynchronous applications. Applications that are
running an `.IOLoop` must use `AsyncHTTPClient` instead.
Typical usage looks like this::
http_client = httpclient.HTTPClient()
try:
response = http_client.fetch("http://www.google.com/")
print(response.body)
except httpclient.HTTPError as e:
# HTTPError is raised for non-200 responses; the response
# can be found in e.response.
print("Error: " + str(e))
except Exception as e:
# Other errors are possible, such as IOError.
print("Error: " + str(e))
http_client.close()
.. versionchanged:: 5.0
Due to limitations in `asyncio`, it is no longer possible to
use the synchronous ``HTTPClient`` while an `.IOLoop` is running.
Use `AsyncHTTPClient` instead.
"""
def __init__(self,
async_client_class: Type["AsyncHTTPClient"] = None,
**kwargs: Any) -> None:
# Initialize self._closed at the beginning of the constructor
# so that an exception raised here doesn't lead to confusing
# failures in __del__.
self._closed = True
self._io_loop = IOLoop(make_current=False)
if async_client_class is None:
async_client_class = AsyncHTTPClient
# Create the client while our IOLoop is "current", without
# clobbering the thread's real current IOLoop (if any).
async def make_client() -> "AsyncHTTPClient":
await gen.sleep(0)
assert async_client_class is not None
return async_client_class(**kwargs)
self._async_client = self._io_loop.run_sync(make_client)
self._closed = False
def __del__(self) -> None:
self.close()
def close(self) -> None:
"""Closes the HTTPClient, freeing any resources used."""
if not self._closed:
self._async_client.close()
self._io_loop.close()
self._closed = True
def fetch(self, request: Union["HTTPRequest", str],
**kwargs: Any) -> "HTTPResponse":
"""Executes a request, returning an `HTTPResponse`.
The request may be either a string URL or an `HTTPRequest` object.
If it is a string, we construct an `HTTPRequest` using any additional
kwargs: ``HTTPRequest(request, **kwargs)``
If an error occurs during the fetch, we raise an `HTTPError` unless
the ``raise_error`` keyword argument is set to False.
"""
response = self._io_loop.run_sync(
functools.partial(self._async_client.fetch, request, **kwargs))
return response
class TestSessionManager(TestCase):
def setUp(self):
self.sm = SessionManager(
kernel_manager=DummyMKM(),
contents_manager=ContentsManager(),
)
self.loop = IOLoop()
def tearDown(self):
self.loop.close(all_fds=True)
def create_sessions(self, *kwarg_list):
@gen.coroutine
def co_add():
sessions = []
for kwargs in kwarg_list:
session = yield self.sm.create_session(**kwargs)
sessions.append(session)
raise gen.Return(sessions)
return self.loop.run_sync(co_add)
def create_session(self, **kwargs):
return self.create_sessions(kwargs)[0]
def test_get_session(self):
sm = self.sm
session_id = self.create_session(path='/path/to/test.ipynb',
kernel_name='bar')['id']
model = sm.get_session(session_id=session_id)
expected = {
'id': session_id,
'notebook': {
'path': u'/path/to/test.ipynb'
},
'kernel': {
'id': u'A',
'name': 'bar'
}
}
self.assertEqual(model, expected)
def test_bad_get_session(self):
# Should raise error if a bad key is passed to the database.
sm = self.sm
session_id = self.create_session(path='/path/to/test.ipynb',
kernel_name='foo')['id']
self.assertRaises(TypeError, sm.get_session,
bad_id=session_id) # Bad keyword
def test_get_session_dead_kernel(self):
sm = self.sm
session = self.create_session(path='/path/to/1/test1.ipynb',
kernel_name='python')
# kill the kernel
sm.kernel_manager.shutdown_kernel(session['kernel']['id'])
with self.assertRaises(KeyError):
sm.get_session(session_id=session['id'])
# no sessions left
listed = sm.list_sessions()
self.assertEqual(listed, [])
def test_list_sessions(self):
sm = self.sm
sessions = self.create_sessions(
dict(path='/path/to/1/test1.ipynb', kernel_name='python'),
dict(path='/path/to/2/test2.ipynb', kernel_name='python'),
dict(path='/path/to/3/test3.ipynb', kernel_name='python'),
)
sessions = sm.list_sessions()
expected = [{
'id': sessions[0]['id'],
'notebook': {
'path': u'/path/to/1/test1.ipynb'
},
'kernel': {
'id': u'A',
'name': 'python'
}
}, {
'id': sessions[1]['id'],
'notebook': {
'path': u'/path/to/2/test2.ipynb'
},
'kernel': {
'id': u'B',
'name': 'python'
}
}, {
'id': sessions[2]['id'],
'notebook': {
'path': u'/path/to/3/test3.ipynb'
},
'kernel': {
'id': u'C',
'name': 'python'
}
}]
self.assertEqual(sessions, expected)
def test_list_sessions_dead_kernel(self):
sm = self.sm
sessions = self.create_sessions(
dict(path='/path/to/1/test1.ipynb', kernel_name='python'),
dict(path='/path/to/2/test2.ipynb', kernel_name='python'),
)
# kill one of the kernels
sm.kernel_manager.shutdown_kernel(sessions[0]['kernel']['id'])
listed = sm.list_sessions()
expected = [{
'id': sessions[1]['id'],
'notebook': {
'path': u'/path/to/2/test2.ipynb',
},
'kernel': {
'id': u'B',
'name': 'python',
}
}]
self.assertEqual(listed, expected)
def test_update_session(self):
sm = self.sm
session_id = self.create_session(path='/path/to/test.ipynb',
kernel_name='julia')['id']
sm.update_session(session_id, path='/path/to/new_name.ipynb')
model = sm.get_session(session_id=session_id)
expected = {
'id': session_id,
'notebook': {
'path': u'/path/to/new_name.ipynb'
},
'kernel': {
'id': u'A',
'name': 'julia'
}
}
self.assertEqual(model, expected)
def test_bad_update_session(self):
# try to update a session with a bad keyword ~ raise error
sm = self.sm
session_id = self.create_session(path='/path/to/test.ipynb',
kernel_name='ir')['id']
self.assertRaises(TypeError,
sm.update_session,
session_id=session_id,
bad_kw='test.ipynb') # Bad keyword
def test_delete_session(self):
sm = self.sm
sessions = self.create_sessions(
dict(path='/path/to/1/test1.ipynb', kernel_name='python'),
dict(path='/path/to/2/test2.ipynb', kernel_name='python'),
dict(path='/path/to/3/test3.ipynb', kernel_name='python'),
)
sm.delete_session(sessions[1]['id'])
new_sessions = sm.list_sessions()
expected = [{
'id': sessions[0]['id'],
'notebook': {
'path': u'/path/to/1/test1.ipynb'
},
'kernel': {
'id': u'A',
'name': 'python'
}
}, {
'id': sessions[2]['id'],
'notebook': {
'path': u'/path/to/3/test3.ipynb'
},
'kernel': {
'id': u'C',
'name': 'python'
}
}]
self.assertEqual(new_sessions, expected)
def test_bad_delete_session(self):
# try to delete a session that doesn't exist ~ raise error
sm = self.sm
self.create_session(path='/path/to/test.ipynb', kernel_name='python')
self.assertRaises(TypeError, sm.delete_session,
bad_kwarg='23424') # Bad keyword
self.assertRaises(web.HTTPError, sm.delete_session,
session_id='23424') # nonexistant
def async_run(func):
ioloop = IOLoop()
ioloop.make_current()
work = coroutine(func)
ioloop.run_sync(work, timeout=5)
def _run(
cls,
worker_kwargs,
worker_start_args,
silence_logs,
init_result_q,
child_stop_q,
uid,
env,
config,
Worker,
): # pragma: no cover
os.environ.update(env)
dask.config.set(config)
try:
from dask.multiprocessing import initialize_worker_process
except ImportError: # old Dask version
pass
else:
initialize_worker_process()
if silence_logs:
logger.setLevel(silence_logs)
IOLoop.clear_instance()
loop = IOLoop()
loop.make_current()
worker = Worker(**worker_kwargs)
async def do_stop(timeout=5, executor_wait=True):
try:
await worker.close(
report=False,
nanny=False,
executor_wait=executor_wait,
timeout=timeout,
)
finally:
loop.stop()
def watch_stop_q():
"""
Wait for an incoming stop message and then stop the
worker cleanly.
"""
while True:
try:
msg = child_stop_q.get(timeout=1000)
except Empty:
pass
else:
child_stop_q.close()
assert msg.pop("op") == "stop"
loop.add_callback(do_stop, **msg)
break
t = threading.Thread(target=watch_stop_q,
name="Nanny stop queue watch")
t.daemon = True
t.start()
async def run():
"""
Try to start worker and inform parent of outcome.
"""
try:
await worker
except Exception as e:
logger.exception("Failed to start worker")
init_result_q.put({"uid": uid, "exception": e})
init_result_q.close()
else:
try:
assert worker.address
except ValueError:
pass
else:
init_result_q.put({
"address": worker.address,
"dir": worker.local_directory,
"uid": uid,
})
init_result_q.close()
await worker.finished()
logger.info("Worker closed")
try:
loop.run_sync(run)
except (TimeoutError, gen.TimeoutError):
# Loop was stopped before wait_until_closed() returned, ignore
pass
except KeyboardInterrupt:
# At this point the loop is not running thus we have to run
# do_stop() explicitly.
loop.run_sync(do_stop)
class TestSessionManager(TestCase):
def setUp(self):
self.sm = SessionManager(
kernel_manager=DummyMKM(),
contents_manager=ContentsManager(),
)
self.loop = IOLoop()
self.addCleanup(partial(self.loop.close, all_fds=True))
def create_sessions(self, *kwarg_list):
@gen.coroutine
def co_add():
sessions = []
for kwargs in kwarg_list:
kwargs.setdefault('type', 'notebook')
session = yield self.sm.create_session(**kwargs)
sessions.append(session)
raise gen.Return(sessions)
return self.loop.run_sync(co_add)
def create_session(self, **kwargs):
return self.create_sessions(kwargs)[0]
def test_get_session(self):
sm = self.sm
session_id = self.create_session(path='/path/to/test.ipynb',
kernel_name='bar')['id']
model = self.loop.run_sync(
lambda: sm.get_session(session_id=session_id))
expected = {
'id': session_id,
'path': u'/path/to/test.ipynb',
'notebook': {
'path': u'/path/to/test.ipynb',
'name': None
},
'type': 'notebook',
'name': None,
'kernel': {
'id': 'A',
'name': 'bar',
'connections': 0,
'last_activity': dummy_date_s,
'execution_state': 'idle',
}
}
self.assertEqual(model, expected)
def test_bad_get_session(self):
# Should raise error if a bad key is passed to the database.
sm = self.sm
session_id = self.create_session(path='/path/to/test.ipynb',
kernel_name='foo')['id']
with self.assertRaises(TypeError):
self.loop.run_sync(
lambda: sm.get_session(bad_id=session_id)) # Bad keyword
def test_get_session_dead_kernel(self):
sm = self.sm
session = self.create_session(path='/path/to/1/test1.ipynb',
kernel_name='python')
# kill the kernel
sm.kernel_manager.shutdown_kernel(session['kernel']['id'])
with self.assertRaises(KeyError):
self.loop.run_sync(
lambda: sm.get_session(session_id=session['id']))
# no sessions left
listed = self.loop.run_sync(lambda: sm.list_sessions())
self.assertEqual(listed, [])
def test_list_sessions(self):
sm = self.sm
sessions = self.create_sessions(
dict(path='/path/to/1/test1.ipynb', kernel_name='python'),
dict(path='/path/to/2/test2.py', type='file',
kernel_name='python'),
dict(path='/path/to/3',
name='foo',
type='console',
kernel_name='python'),
)
sessions = self.loop.run_sync(lambda: sm.list_sessions())
expected = [{
'id': sessions[0]['id'],
'path': u'/path/to/1/test1.ipynb',
'type': 'notebook',
'notebook': {
'path': u'/path/to/1/test1.ipynb',
'name': None
},
'name': None,
'kernel': {
'id': 'A',
'name': 'python',
'connections': 0,
'last_activity': dummy_date_s,
'execution_state': 'idle',
}
}, {
'id': sessions[1]['id'],
'path': u'/path/to/2/test2.py',
'type': 'file',
'name': None,
'kernel': {
'id': 'B',
'name': 'python',
'connections': 0,
'last_activity': dummy_date_s,
'execution_state': 'idle',
}
}, {
'id': sessions[2]['id'],
'path': u'/path/to/3',
'type': 'console',
'name': 'foo',
'kernel': {
'id': 'C',
'name': 'python',
'connections': 0,
'last_activity': dummy_date_s,
'execution_state': 'idle',
}
}]
self.assertEqual(sessions, expected)
def test_list_sessions_dead_kernel(self):
sm = self.sm
sessions = self.create_sessions(
dict(path='/path/to/1/test1.ipynb', kernel_name='python'),
dict(path='/path/to/2/test2.ipynb', kernel_name='python'),
)
# kill one of the kernels
sm.kernel_manager.shutdown_kernel(sessions[0]['kernel']['id'])
listed = self.loop.run_sync(lambda: sm.list_sessions())
expected = [{
'id': sessions[1]['id'],
'path': u'/path/to/2/test2.ipynb',
'type': 'notebook',
'name': None,
'notebook': {
'path': u'/path/to/2/test2.ipynb',
'name': None
},
'kernel': {
'id': 'B',
'name': 'python',
'connections': 0,
'last_activity': dummy_date_s,
'execution_state': 'idle',
}
}]
self.assertEqual(listed, expected)
def test_update_session(self):
sm = self.sm
session_id = self.create_session(path='/path/to/test.ipynb',
kernel_name='julia')['id']
self.loop.run_sync(lambda: sm.update_session(
session_id, path='/path/to/new_name.ipynb'))
model = self.loop.run_sync(
lambda: sm.get_session(session_id=session_id))
expected = {
'id': session_id,
'path': u'/path/to/new_name.ipynb',
'type': 'notebook',
'name': None,
'notebook': {
'path': u'/path/to/new_name.ipynb',
'name': None
},
'kernel': {
'id': 'A',
'name': 'julia',
'connections': 0,
'last_activity': dummy_date_s,
'execution_state': 'idle',
}
}
self.assertEqual(model, expected)
def test_bad_update_session(self):
# try to update a session with a bad keyword ~ raise error
sm = self.sm
session_id = self.create_session(path='/path/to/test.ipynb',
kernel_name='ir')['id']
with self.assertRaises(TypeError):
self.loop.run_sync(lambda: sm.update_session(
session_id=session_id, bad_kw='test.ipynb')) # Bad keyword
def test_delete_session(self):
sm = self.sm
sessions = self.create_sessions(
dict(path='/path/to/1/test1.ipynb', kernel_name='python'),
dict(path='/path/to/2/test2.ipynb', kernel_name='python'),
dict(path='/path/to/3',
name='foo',
type='console',
kernel_name='python'),
)
self.loop.run_sync(lambda: sm.delete_session(sessions[1]['id']))
new_sessions = self.loop.run_sync(lambda: sm.list_sessions())
expected = [{
'id': sessions[0]['id'],
'path': u'/path/to/1/test1.ipynb',
'type': 'notebook',
'name': None,
'notebook': {
'path': u'/path/to/1/test1.ipynb',
'name': None
},
'kernel': {
'id': 'A',
'name': 'python',
'connections': 0,
'last_activity': dummy_date_s,
'execution_state': 'idle',
}
}, {
'id': sessions[2]['id'],
'type': 'console',
'path': u'/path/to/3',
'name': 'foo',
'kernel': {
'id': 'C',
'name': 'python',
'connections': 0,
'last_activity': dummy_date_s,
'execution_state': 'idle',
}
}]
self.assertEqual(new_sessions, expected)
def test_bad_delete_session(self):
# try to delete a session that doesn't exist ~ raise error
sm = self.sm
self.create_session(path='/path/to/test.ipynb', kernel_name='python')
with self.assertRaises(TypeError):
self.loop.run_sync(
lambda: sm.delete_session(bad_kwarg='23424')) # Bad keyword
with self.assertRaises(web.HTTPError):
self.loop.run_sync(
lambda: sm.delete_session(session_id='23424')) # nonexistent
def main(scheduler, host, worker_port, http_port, nanny_port, nthreads, nprocs,
nanny, name, memory_limit, pid_file, temp_filename, reconnect,
resources, bokeh, bokeh_port, local_directory):
if nanny:
port = nanny_port
else:
port = worker_port
if nprocs > 1 and worker_port != 0:
logger.error("Failed to launch worker. You cannot use the --port argument when nprocs > 1.")
exit(1)
if nprocs > 1 and name:
logger.error("Failed to launch worker. You cannot use the --name argument when nprocs > 1.")
exit(1)
if not nthreads:
nthreads = _ncores // nprocs
if pid_file:
with open(pid_file, 'w') as f:
f.write(str(os.getpid()))
def del_pid_file():
if os.path.exists(pid_file):
os.remove(pid_file)
atexit.register(del_pid_file)
services = {('http', http_port): HTTPWorker}
if bokeh:
try:
from distributed.bokeh.worker import BokehWorker
except ImportError:
pass
else:
services[('bokeh', bokeh_port)] = BokehWorker
if resources:
resources = resources.replace(',', ' ').split()
resources = dict(pair.split('=') for pair in resources)
resources = valmap(float, resources)
else:
resources = None
loop = IOLoop.current()
if nanny:
kwargs = {'worker_port': worker_port}
t = Nanny
else:
kwargs = {}
if nanny_port:
kwargs['service_ports'] = {'nanny': nanny_port}
t = Worker
nannies = [t(scheduler, ncores=nthreads,
services=services, name=name, loop=loop, resources=resources,
memory_limit=memory_limit, reconnect=reconnect,
local_dir=local_directory, **kwargs)
for i in range(nprocs)]
for n in nannies:
if host:
n.start((host, port))
else:
n.start(port)
if t is Nanny:
global_nannies.append(n)
if temp_filename:
@gen.coroutine
def f():
while nannies[0].status != 'running':
yield gen.sleep(0.01)
import json
msg = {'port': nannies[0].port,
'local_directory': nannies[0].local_dir}
with open(temp_filename, 'w') as f:
json.dump(msg, f)
loop.add_callback(f)
@gen.coroutine
def run():
while all(n.status != 'closed' for n in nannies):
yield gen.sleep(0.2)
try:
loop.run_sync(run)
except (KeyboardInterrupt, TimeoutError):
pass
finally:
logger.info("End worker")
loop.close()
# Clean exit: unregister all workers from scheduler
loop2 = IOLoop()
@gen.coroutine
def f():
scheduler = rpc(nannies[0].scheduler.address)
if nanny:
yield gen.with_timeout(timedelta(seconds=2),
All([scheduler.unregister(address=n.worker_address, close=True)
for n in nannies if n.process and n.worker_address]),
io_loop=loop2)
loop2.run_sync(f)
if nanny:
for n in nannies:
if isalive(n.process):
n.process.terminate()
if nanny:
start = time()
while (any(isalive(n.process) for n in nannies)
and time() < start + 1):
sleep(0.1)
for nanny in nannies:
nanny.stop()
class SyncHTTPClientTest(unittest.TestCase):
def setUp(self):
self.server_ioloop = IOLoop()
event = threading.Event()
@gen.coroutine
def init_server():
sock, self.port = bind_unused_port()
app = Application([("/", HelloWorldHandler)])
self.server = HTTPServer(app)
self.server.add_socket(sock)
event.set()
def start():
self.server_ioloop.run_sync(init_server)
self.server_ioloop.start()
self.server_thread = threading.Thread(target=start)
self.server_thread.start()
event.wait()
self.http_client = HTTPClient()
def tearDown(self):
def stop_server():
self.server.stop()
# Delay the shutdown of the IOLoop by several iterations because
# the server may still have some cleanup work left when
# the client finishes with the response (this is noticeable
# with http/2, which leaves a Future with an unexamined
# StreamClosedError on the loop).
@gen.coroutine
def slow_stop():
yield self.server.close_all_connections()
# The number of iterations is difficult to predict. Typically,
# one is sufficient, although sometimes it needs more.
for i in range(5):
yield
self.server_ioloop.stop()
self.server_ioloop.add_callback(slow_stop)
self.server_ioloop.add_callback(stop_server)
self.server_thread.join()
self.http_client.close()
self.server_ioloop.close(all_fds=True)
def get_url(self, path):
return "http://127.0.0.1:%d%s" % (self.port, path)
def test_sync_client(self):
response = self.http_client.fetch(self.get_url("/"))
self.assertEqual(b"Hello world!", response.body)
def test_sync_client_error(self):
# Synchronous HTTPClient raises errors directly; no need for
# response.rethrow()
with self.assertRaises(HTTPError) as assertion:
self.http_client.fetch(self.get_url("/notfound"))
self.assertEqual(assertion.exception.code, 404)
class TestFutureSocket(BaseZMQTestCase):
Context = future.Context
def setUp(self):
self.loop = IOLoop()
self.loop.make_current()
super(TestFutureSocket, self).setUp()
def tearDown(self):
super(TestFutureSocket, self).tearDown()
if self.loop:
self.loop.close(all_fds=True)
IOLoop.clear_current()
IOLoop.clear_instance()
def test_socket_class(self):
s = self.context.socket(zmq.PUSH)
assert isinstance(s, future.Socket)
s.close()
def test_instance_subclass_first(self):
actx = self.Context.instance()
ctx = zmq.Context.instance()
ctx.term()
actx.term()
assert type(ctx) is zmq.Context
assert type(actx) is self.Context
def test_instance_subclass_second(self):
ctx = zmq.Context.instance()
actx = self.Context.instance()
ctx.term()
actx.term()
assert type(ctx) is zmq.Context
assert type(actx) is self.Context
def test_recv_multipart(self):
async def test():
a, b = self.create_bound_pair(zmq.PUSH, zmq.PULL)
f = b.recv_multipart()
assert not f.done()
await a.send(b"hi")
recvd = await f
self.assertEqual(recvd, [b'hi'])
self.loop.run_sync(test)
def test_recv(self):
async def test():
a, b = self.create_bound_pair(zmq.PUSH, zmq.PULL)
f1 = b.recv()
f2 = b.recv()
assert not f1.done()
assert not f2.done()
await a.send_multipart([b"hi", b"there"])
recvd = await f2
assert f1.done()
self.assertEqual(f1.result(), b'hi')
self.assertEqual(recvd, b'there')
self.loop.run_sync(test)
def test_recv_cancel(self):
async def test():
a, b = self.create_bound_pair(zmq.PUSH, zmq.PULL)
f1 = b.recv()
f2 = b.recv_multipart()
assert f1.cancel()
assert f1.done()
assert not f2.done()
await a.send_multipart([b"hi", b"there"])
recvd = await f2
assert f1.cancelled()
assert f2.done()
self.assertEqual(recvd, [b'hi', b'there'])
self.loop.run_sync(test)
@pytest.mark.skipif(not hasattr(zmq, 'RCVTIMEO'),
reason="requires RCVTIMEO")
def test_recv_timeout(self):
async def test():
a, b = self.create_bound_pair(zmq.PUSH, zmq.PULL)
b.rcvtimeo = 100
f1 = b.recv()
b.rcvtimeo = 1000
f2 = b.recv_multipart()
with pytest.raises(zmq.Again):
await f1
await a.send_multipart([b"hi", b"there"])
recvd = await f2
assert f2.done()
self.assertEqual(recvd, [b'hi', b'there'])
self.loop.run_sync(test)
@pytest.mark.skipif(not hasattr(zmq, 'SNDTIMEO'),
reason="requires SNDTIMEO")
def test_send_timeout(self):
async def test():
s = self.socket(zmq.PUSH)
s.sndtimeo = 100
with pytest.raises(zmq.Again):
await s.send(b"not going anywhere")
self.loop.run_sync(test)
def test_send_noblock(self):
async def test():
s = self.socket(zmq.PUSH)
with pytest.raises(zmq.Again):
await s.send(b"not going anywhere", flags=zmq.NOBLOCK)
self.loop.run_sync(test)
def test_send_multipart_noblock(self):
async def test():
s = self.socket(zmq.PUSH)
with pytest.raises(zmq.Again):
await s.send_multipart([b"not going anywhere"],
flags=zmq.NOBLOCK)
self.loop.run_sync(test)
def test_recv_string(self):
async def test():
a, b = self.create_bound_pair(zmq.PUSH, zmq.PULL)
f = b.recv_string()
assert not f.done()
msg = u('πøøπ')
await a.send_string(msg)
recvd = await f
assert f.done()
self.assertEqual(f.result(), msg)
self.assertEqual(recvd, msg)
self.loop.run_sync(test)
def test_recv_json(self):
async def test():
a, b = self.create_bound_pair(zmq.PUSH, zmq.PULL)
f = b.recv_json()
assert not f.done()
obj = dict(a=5)
await a.send_json(obj)
recvd = await f
assert f.done()
self.assertEqual(f.result(), obj)
self.assertEqual(recvd, obj)
self.loop.run_sync(test)
def test_recv_json_cancelled(self):
async def test():
a, b = self.create_bound_pair(zmq.PUSH, zmq.PULL)
f = b.recv_json()
assert not f.done()
f.cancel()
# cycle eventloop to allow cancel events to fire
await gen.sleep(0)
obj = dict(a=5)
await a.send_json(obj)
with pytest.raises(future.CancelledError):
recvd = await f
assert f.done()
# give it a chance to incorrectly consume the event
events = await b.poll(timeout=5)
assert events
await gen.sleep(0)
# make sure cancelled recv didn't eat up event
recvd = await gen.with_timeout(timedelta(seconds=5), b.recv_json())
assert recvd == obj
self.loop.run_sync(test)
def test_recv_pyobj(self):
async def test():
a, b = self.create_bound_pair(zmq.PUSH, zmq.PULL)
f = b.recv_pyobj()
assert not f.done()
obj = dict(a=5)
await a.send_pyobj(obj)
recvd = await f
assert f.done()
self.assertEqual(f.result(), obj)
self.assertEqual(recvd, obj)
self.loop.run_sync(test)
def test_custom_serialize(self):
def serialize(msg):
frames = []
frames.extend(msg.get('identities', []))
content = json.dumps(msg['content']).encode('utf8')
frames.append(content)
return frames
def deserialize(frames):
identities = frames[:-1]
content = json.loads(frames[-1].decode('utf8'))
return {
'identities': identities,
'content': content,
}
async def test():
a, b = self.create_bound_pair(zmq.DEALER, zmq.ROUTER)
msg = {
'content': {
'a': 5,
'b': 'bee',
}
}
await a.send_serialized(msg, serialize)
recvd = await b.recv_serialized(deserialize)
assert recvd['content'] == msg['content']
assert recvd['identities']
# bounce back, tests identities
await b.send_serialized(recvd, serialize)
r2 = await a.recv_serialized(deserialize)
assert r2['content'] == msg['content']
assert not r2['identities']
self.loop.run_sync(test)
def test_custom_serialize_error(self):
async def test():
a, b = self.create_bound_pair(zmq.DEALER, zmq.ROUTER)
msg = {
'content': {
'a': 5,
'b': 'bee',
}
}
with pytest.raises(TypeError):
await a.send_serialized(json, json.dumps)
await a.send(b"not json")
with pytest.raises(TypeError):
recvd = await b.recv_serialized(json.loads)
self.loop.run_sync(test)
def test_poll(self):
async def test():
a, b = self.create_bound_pair(zmq.PUSH, zmq.PULL)
f = b.poll(timeout=0)
assert f.done()
self.assertEqual(f.result(), 0)
f = b.poll(timeout=1)
assert not f.done()
evt = await f
self.assertEqual(evt, 0)
f = b.poll(timeout=1000)
assert not f.done()
await a.send_multipart([b"hi", b"there"])
evt = await f
self.assertEqual(evt, zmq.POLLIN)
recvd = await b.recv_multipart()
self.assertEqual(recvd, [b'hi', b'there'])
self.loop.run_sync(test)
@pytest.mark.skipif(sys.platform.startswith('win'),
reason='Windows unsupported socket type')
def test_poll_base_socket(self):
async def test():
ctx = zmq.Context()
url = 'inproc://test'
a = ctx.socket(zmq.PUSH)
b = ctx.socket(zmq.PULL)
self.sockets.extend([a, b])
a.bind(url)
b.connect(url)
poller = future.Poller()
poller.register(b, zmq.POLLIN)
f = poller.poll(timeout=1000)
assert not f.done()
a.send_multipart([b'hi', b'there'])
evt = await f
self.assertEqual(evt, [(b, zmq.POLLIN)])
recvd = b.recv_multipart()
self.assertEqual(recvd, [b'hi', b'there'])
a.close()
b.close()
ctx.term()
self.loop.run_sync(test)
def test_close_all_fds(self):
s = self.socket(zmq.PUB)
s._get_loop()
self.loop.close(all_fds=True)
self.loop = None # avoid second close later
assert s.closed
@pytest.mark.skipif(
sys.platform.startswith('win'),
reason='Windows does not support polling on files',
)
def test_poll_raw(self):
async def test():
p = future.Poller()
# make a pipe
r, w = os.pipe()
r = os.fdopen(r, 'rb')
w = os.fdopen(w, 'wb')
# POLLOUT
p.register(r, zmq.POLLIN)
p.register(w, zmq.POLLOUT)
evts = await p.poll(timeout=1)
evts = dict(evts)
assert r.fileno() not in evts
assert w.fileno() in evts
assert evts[w.fileno()] == zmq.POLLOUT
# POLLIN
p.unregister(w)
w.write(b'x')
w.flush()
evts = await p.poll(timeout=1000)
evts = dict(evts)
assert r.fileno() in evts
assert evts[r.fileno()] == zmq.POLLIN
assert r.read(1) == b'x'
r.close()
w.close()
self.loop.run_sync(test)
def main(scheduler, host, worker_port, http_port, nanny_port, nthreads, nprocs,
nanny, name, memory_limit, pid_file, reconnect, resources, bokeh,
bokeh_port, local_directory, scheduler_file, interface, death_timeout,
preload, bokeh_prefix, tls_ca_file, tls_cert, tls_key):
sec = Security(
tls_ca_file=tls_ca_file,
tls_worker_cert=tls_cert,
tls_worker_key=tls_key,
)
if nanny:
port = nanny_port
else:
port = worker_port
if nprocs > 1 and worker_port != 0:
logger.error(
"Failed to launch worker. You cannot use the --port argument when nprocs > 1."
)
exit(1)
if nprocs > 1 and name:
logger.error(
"Failed to launch worker. You cannot use the --name argument when nprocs > 1."
)
exit(1)
if not nthreads:
nthreads = _ncores // nprocs
if pid_file:
with open(pid_file, 'w') as f:
f.write(str(os.getpid()))
def del_pid_file():
if os.path.exists(pid_file):
os.remove(pid_file)
atexit.register(del_pid_file)
services = {('http', http_port): HTTPWorker}
if bokeh:
try:
from distributed.bokeh.worker import BokehWorker
except ImportError:
pass
else:
if bokeh_prefix:
result = (BokehWorker, {'prefix': bokeh_prefix})
else:
result = BokehWorker
services[('bokeh', bokeh_port)] = result
if resources:
resources = resources.replace(',', ' ').split()
resources = dict(pair.split('=') for pair in resources)
resources = valmap(float, resources)
else:
resources = None
loop = IOLoop.current()
if nanny:
kwargs = {'worker_port': worker_port}
t = Nanny
else:
kwargs = {}
if nanny_port:
kwargs['service_ports'] = {'nanny': nanny_port}
t = Worker
if scheduler_file:
while not os.path.exists(scheduler_file):
sleep(0.01)
for i in range(10):
try:
with open(scheduler_file) as f:
cfg = json.load(f)
scheduler = cfg['address']
break
except (ValueError, KeyError): # race with scheduler on file
sleep(0.01)
if not scheduler:
raise ValueError("Need to provide scheduler address like\n"
"dask-worker SCHEDULER_ADDRESS:8786")
nannies = [
t(scheduler,
ncores=nthreads,
services=services,
name=name,
loop=loop,
resources=resources,
memory_limit=memory_limit,
reconnect=reconnect,
local_dir=local_directory,
death_timeout=death_timeout,
preload=preload,
security=sec,
**kwargs) for i in range(nprocs)
]
if interface:
if host:
raise ValueError("Can not specify both interface and host")
else:
host = get_ip_interface(interface)
if host or port:
addr = uri_from_host_port(host, port, 0)
else:
# Choose appropriate address for scheduler
addr = None
for n in nannies:
n.start(addr)
if t is Nanny:
global_nannies.append(n)
@gen.coroutine
def run():
while all(n.status != 'closed' for n in nannies):
yield gen.sleep(0.2)
try:
loop.run_sync(run)
except (KeyboardInterrupt, TimeoutError):
pass
finally:
logger.info("End worker")
loop.close()
# Clean exit: unregister all workers from scheduler
loop2 = IOLoop()
@gen.coroutine
def f():
if nanny:
w = nannies[0]
with w.rpc(w.scheduler.address) as scheduler:
yield gen.with_timeout(timeout=timedelta(seconds=2),
future=All([
scheduler.unregister(
address=n.worker_address,
close=True) for n in nannies
if n.process and n.worker_address
]),
io_loop=loop2)
loop2.run_sync(f)
loop2.close()
if nanny:
for n in nannies:
if isalive(n.process):
n.process.terminate()
if nanny:
start = time()
while (any(isalive(n.process) for n in nannies)
and time() < start + 1):
sleep(0.1)
for nanny in nannies:
nanny.stop()
