<filename>distributed/tests/test_worker.py<gh_stars>0
from operator import add
import os
import shutil
import sys
import traceback

import pytest
from toolz import pluck
from tornado import gen
from tornado.ioloop import TimeoutError

from distributed.center import Center
from distributed.core import rpc, dumps, loads
from distributed.sizeof import sizeof
from distributed.worker import Worker
from distributed.utils_test import loop, _test_cluster, inc, gen_cluster



def test_worker_ncores():
    from distributed.worker import _ncores
    w = Worker('127.0.0.1', 8019)
    try:
        assert w.executor._max_workers == _ncores
    finally:
        shutil.rmtree(w.local_dir)

def test_identity():
    w = Worker('127.0.0.1', 8019)
    ident = w.identity(None)
    assert ident['type'] == 'Worker'
    assert ident['center'] == ('127.0.0.1', 8019)


def test_worker(loop):
    @gen.coroutine
    def f(c, a, b):
        aa = rpc(ip=a.ip, port=a.port)
        bb = rpc(ip=b.ip, port=b.port)

        assert not a.active
        response, _ = yield aa.compute(key='x', function=add,
                                       args=[1, 2], needed=[],
                                       close=True)
        assert not a.active
        assert response == b'OK'
        assert a.data['x'] == 3
        assert c.who_has['x'] == set([(a.ip, a.port)])

        response, info = yield bb.compute(key='y', function=add,
                                          args=['x', 10], needed=['x'])
        assert response == b'OK'
        assert b.data['y'] == 13
        assert c.who_has['y'] == set([(b.ip, b.port)])
        assert info['nbytes'] == sizeof(b.data['y'])

        def bad_func():
            1 / 0

        response, content = yield bb.compute(key='z',
                function=bad_func, args=(), needed=(), close=True)
        assert not b.active
        assert response == b'error'
        assert isinstance(content['exception'], ZeroDivisionError)
        if sys.version_info[0] >= 3:
            assert any('1 / 0' in line
                      for line in pluck(3, traceback.extract_tb(content['traceback']))
                      if line)

        aa.close_streams()
        yield a._close()

        assert a.address not in c.ncores and b.address in c.ncores

        assert list(c.ncores.keys()) == [(b.ip, b.port)]

        assert isinstance(b.address_string, str)
        assert b.ip in b.address_string
        assert str(b.port) in b.address_string

        bb.close_streams()
        yield b._close()

    _test_cluster(f)


def test_compute_who_has(loop):
    @gen.coroutine
    def f():
        c = Center(ip='127.0.0.1')
        c.listen(0)
        x = Worker(c.ip, c.port, ip='127.0.0.1')
        y = Worker(c.ip, c.port, ip='127.0.0.1')
        z = Worker(c.ip, c.port, ip='127.0.0.1')
        x.data['a'] = 1
        y.data['a'] = 2
        yield [x._start(), y._start(), z._start()]

        zz = rpc(ip=z.ip, port=z.port)
        yield zz.compute(function=inc, args=('a',), needed=['a'],
                         who_has={'a': {x.address}}, key='b')
        assert z.data['b'] == 2

        yield zz.compute(function=inc, args=('a',), needed=['a'],
                         who_has={'a': {y.address}}, key='c')
        assert z.data['c'] == 3

        yield [x._close(), y._close(), z._close()]
        zz.close_streams()

    loop.run_sync(f, timeout=5)


def test_workers_update_center(loop):
    @gen.coroutine
    def f(c, a, b):
        aa = rpc(ip=a.ip, port=a.port)

        response, content = yield aa.update_data(data={'x': 1, 'y': 2})
        assert response == b'OK'
        assert content['nbytes'] == {'x': sizeof(1), 'y': sizeof(2)}

        assert a.data == {'x': 1, 'y': 2}
        assert c.who_has == {'x': {(a.ip, a.port)},
                             'y': {(a.ip, a.port)}}
        assert c.has_what[(a.ip, a.port)] == {'x', 'y'}

        yield aa.delete_data(keys=['x'], close=True)
        assert not c.who_has['x']
        assert all('x' not in s for s in c.has_what.values())

        aa.close_streams()

    _test_cluster(f)


def test_delete_data_with_missing_worker(loop):
    @gen.coroutine
    def f(c, a, b):
        bad = ('127.0.0.1', 9001)  # this worker doesn't exist
        c.who_has['z'].add(bad)
        c.who_has['z'].add(a.address)
        c.has_what[bad].add('z')
        c.has_what[a.address].add('z')
        a.data['z'] = 5

        cc = rpc(ip=c.ip, port=c.port)

        yield cc.delete_data(keys=['z'])
        assert 'z' not in a.data
        assert not c.who_has['z']
        assert not c.has_what[bad]
        assert not c.has_what[a.address]

        cc.close_streams()

    _test_cluster(f)


def test_upload_file(loop):
    @gen.coroutine
    def f(c, a, b):
        assert not os.path.exists(os.path.join(a.local_dir, 'foobar.py'))
        assert not os.path.exists(os.path.join(b.local_dir, 'foobar.py'))
        assert a.local_dir != b.local_dir

        aa = rpc(ip=a.ip, port=a.port)
        bb = rpc(ip=b.ip, port=b.port)
        yield [aa.upload_file(filename='foobar.py', data=b'x = 123'),
               bb.upload_file(filename='foobar.py', data=b'x = 123')]

        assert os.path.exists(os.path.join(a.local_dir, 'foobar.py'))
        assert os.path.exists(os.path.join(b.local_dir, 'foobar.py'))

        def g():
            import foobar
            return foobar.x

        yield aa.compute(function=g, key='x')
        result = yield aa.get_data(keys=['x'])
        assert result == {'x': 123}

        yield a._close()
        yield b._close()
        aa.close_streams()
        bb.close_streams()
        assert not os.path.exists(os.path.join(a.local_dir, 'foobar.py'))

    _test_cluster(f)


def test_upload_egg(loop):
    @gen.coroutine
    def f(c, a, b):
        eggname = 'mytestegg-1.0.0-py3.4.egg'
        local_file = __file__.replace('test_worker.py', eggname)
        assert not os.path.exists(os.path.join(a.local_dir, eggname))
        assert not os.path.exists(os.path.join(b.local_dir, eggname))
        assert a.local_dir != b.local_dir

        aa = rpc(ip=a.ip, port=a.port)
        bb = rpc(ip=b.ip, port=b.port)
        with open(local_file, 'rb') as f:
            payload = f.read()
        yield [aa.upload_file(filename=eggname, data=payload),
               bb.upload_file(filename=eggname, data=payload)]

        assert os.path.exists(os.path.join(a.local_dir, eggname))
        assert os.path.exists(os.path.join(b.local_dir, eggname))

        def g(x):
            import testegg
            return testegg.inc(x)

        yield aa.compute(function=g, key='x', args=(10,))
        result = yield aa.get_data(keys=['x'])
        assert result == {'x': 10 + 1}

        yield a._close()
        yield b._close()
        aa.close_streams()
        bb.close_streams()
        assert not os.path.exists(os.path.join(a.local_dir, eggname))

    _test_cluster(f)


def test_broadcast(loop):
    @gen.coroutine
    def f(c, a, b):
        cc = rpc(ip=c.ip, port=c.port)
        results = yield cc.broadcast(msg={'op': 'ping'})
        assert results == {a.address: b'pong', b.address: b'pong'}

        cc.close_streams()

    _test_cluster(f)


def test_worker_with_port_zero(loop):
    @gen.coroutine
    def f():
        c = Center('127.0.0.1')
        c.listen(8007)
        w = Worker(c.ip, c.port, ip='127.0.0.1')
        yield w._start()
        assert isinstance(w.port, int)
        assert w.port > 1024

    loop.run_sync(f)

@pytest.mark.slow
def test_worker_waits_for_center_to_come_up(loop):
    @gen.coroutine
    def f():
        w = Worker('127.0.0.1', 8007, ip='127.0.0.1')
        yield w._start()

    try:
        loop.run_sync(f, timeout=4)
    except TimeoutError:
        pass


@gen_cluster()
def test_worker_task(s, a, b):
    aa = rpc(ip=a.ip, port=a.port)
    yield aa.compute(task=(inc, 1), key='x')

    assert a.data['x'] == 2


@gen_cluster()
def test_worker_task_data(s, a, b):
    aa = rpc(ip=a.ip, port=a.port)
    yield aa.compute(task=2, key='x')

    assert a.data['x'] == 2


@gen_cluster()
def test_worker_task_bytes(s, a, b):
    aa = rpc(ip=a.ip, port=a.port)

    yield aa.compute(task=dumps((inc, 1)), key='x', serialized=True)
    assert a.data['x'] == 2

    yield aa.compute(function=dumps(inc), args=dumps((10,)), key='y', serialized=True)
    assert a.data['y'] == 11

<filename>distributed/tests/test_worker.py<gh_stars>0
from operator import add
import os
import shutil
import sys
import traceback

import pytest
from toolz import pluck
from tornado import gen
from tornado.ioloop import TimeoutError

from distributed.center import Center
from distributed.core import rpc, dumps, loads
from distributed.sizeof import sizeof
from distributed.worker import Worker
from distributed.utils_test import loop, _test_cluster, inc, gen_cluster



def test_worker_ncores():
    from distributed.worker import _ncores
    w = Worker('127.0.0.1', 8019)
    try:
        assert w.executor._max_workers == _ncores
    finally:
        shutil.rmtree(w.local_dir)

def test_identity():
    w = Worker('127.0.0.1', 8019)
    ident = w.identity(None)
    assert ident['type'] == 'Worker'
    assert ident['center'] == ('127.0.0.1', 8019)


def test_worker(loop):
    @gen.coroutine
    def f(c, a, b):
        aa = rpc(ip=a.ip, port=a.port)
        bb = rpc(ip=b.ip, port=b.port)

        assert not a.active
        response, _ = yield aa.compute(key='x', function=add,
                                       args=[1, 2], needed=[],
                                       close=True)
        assert not a.active
        assert response == b'OK'
        assert a.data['x'] == 3
        assert c.who_has['x'] == set([(a.ip, a.port)])

        response, info = yield bb.compute(key='y', function=add,
                                          args=['x', 10], needed=['x'])
        assert response == b'OK'
        assert b.data['y'] == 13
        assert c.who_has['y'] == set([(b.ip, b.port)])
        assert info['nbytes'] == sizeof(b.data['y'])

        def bad_func():
            1 / 0

        response, content = yield bb.compute(key='z',
                function=bad_func, args=(), needed=(), close=True)
        assert not b.active
        assert response == b'error'
        assert isinstance(content['exception'], ZeroDivisionError)
        if sys.version_info[0] >= 3:
            assert any('1 / 0' in line
                      for line in pluck(3, traceback.extract_tb(content['traceback']))
                      if line)

        aa.close_streams()
        yield a._close()

        assert a.address not in c.ncores and b.address in c.ncores

        assert list(c.ncores.keys()) == [(b.ip, b.port)]

        assert isinstance(b.address_string, str)
        assert b.ip in b.address_string
        assert str(b.port) in b.address_string

        bb.close_streams()
        yield b._close()

    _test_cluster(f)


def test_compute_who_has(loop):
    @gen.coroutine
    def f():
        c = Center(ip='127.0.0.1')
        c.listen(0)
        x = Worker(c.ip, c.port, ip='127.0.0.1')
        y = Worker(c.ip, c.port, ip='127.0.0.1')
        z = Worker(c.ip, c.port, ip='127.0.0.1')
        x.data['a'] = 1
        y.data['a'] = 2
        yield [x._start(), y._start(), z._start()]

        zz = rpc(ip=z.ip, port=z.port)
        yield zz.compute(function=inc, args=('a',), needed=['a'],
                         who_has={'a': {x.address}}, key='b')
        assert z.data['b'] == 2

        yield zz.compute(function=inc, args=('a',), needed=['a'],
                         who_has={'a': {y.address}}, key='c')
        assert z.data['c'] == 3

        yield [x._close(), y._close(), z._close()]
        zz.close_streams()

    loop.run_sync(f, timeout=5)


def test_workers_update_center(loop):
    @gen.coroutine
    def f(c, a, b):
        aa = rpc(ip=a.ip, port=a.port)

        response, content = yield aa.update_data(data={'x': 1, 'y': 2})
        assert response == b'OK'
        assert content['nbytes'] == {'x': sizeof(1), 'y': sizeof(2)}

        assert a.data == {'x': 1, 'y': 2}
        assert c.who_has == {'x': {(a.ip, a.port)},
                             'y': {(a.ip, a.port)}}
        assert c.has_what[(a.ip, a.port)] == {'x', 'y'}

        yield aa.delete_data(keys=['x'], close=True)
        assert not c.who_has['x']
        assert all('x' not in s for s in c.has_what.values())

        aa.close_streams()

    _test_cluster(f)


def test_delete_data_with_missing_worker(loop):
    @gen.coroutine
    def f(c, a, b):
        bad = ('127.0.0.1', 9001)  # this worker doesn't exist
        c.who_has['z'].add(bad)
        c.who_has['z'].add(a.address)
        c.has_what[bad].add('z')
        c.has_what[a.address].add('z')
        a.data['z'] = 5

        cc = rpc(ip=c.ip, port=c.port)

        yield cc.delete_data(keys=['z'])
        assert 'z' not in a.data
        assert not c.who_has['z']
        assert not c.has_what[bad]
        assert not c.has_what[a.address]

        cc.close_streams()

    _test_cluster(f)


def test_upload_file(loop):
    @gen.coroutine
    def f(c, a, b):
        assert not os.path.exists(os.path.join(a.local_dir, 'foobar.py'))
        assert not os.path.exists(os.path.join(b.local_dir, 'foobar.py'))
        assert a.local_dir != b.local_dir

        aa = rpc(ip=a.ip, port=a.port)
        bb = rpc(ip=b.ip, port=b.port)
        yield [aa.upload_file(filename='foobar.py', data=b'x = 123'),
               bb.upload_file(filename='foobar.py', data=b'x = 123')]

        assert os.path.exists(os.path.join(a.local_dir, 'foobar.py'))
        assert os.path.exists(os.path.join(b.local_dir, 'foobar.py'))

        def g():
            import foobar
            return foobar.x

        yield aa.compute(function=g, key='x')
        result = yield aa.get_data(keys=['x'])
        assert result == {'x': 123}

        yield a._close()
        yield b._close()
        aa.close_streams()
        bb.close_streams()
        assert not os.path.exists(os.path.join(a.local_dir, 'foobar.py'))

    _test_cluster(f)


def test_upload_egg(loop):
    @gen.coroutine
    def f(c, a, b):
        eggname = 'mytestegg-1.0.0-py3.4.egg'
        local_file = __file__.replace('test_worker.py', eggname)
        assert not os.path.exists(os.path.join(a.local_dir, eggname))
        assert not os.path.exists(os.path.join(b.local_dir, eggname))
        assert a.local_dir != b.local_dir

        aa = rpc(ip=a.ip, port=a.port)
        bb = rpc(ip=b.ip, port=b.port)
        with open(local_file, 'rb') as f:
            payload = f.read()
        yield [aa.upload_file(filename=eggname, data=payload),
               bb.upload_file(filename=eggname, data=payload)]

        assert os.path.exists(os.path.join(a.local_dir, eggname))
        assert os.path.exists(os.path.join(b.local_dir, eggname))

        def g(x):
            import testegg
            return testegg.inc(x)

        yield aa.compute(function=g, key='x', args=(10,))
        result = yield aa.get_data(keys=['x'])
        assert result == {'x': 10 + 1}

        yield a._close()
        yield b._close()
        aa.close_streams()
        bb.close_streams()
        assert not os.path.exists(os.path.join(a.local_dir, eggname))

    _test_cluster(f)


def test_broadcast(loop):
    @gen.coroutine
    def f(c, a, b):
        cc = rpc(ip=c.ip, port=c.port)
        results = yield cc.broadcast(msg={'op': 'ping'})
        assert results == {a.address: b'pong', b.address: b'pong'}

        cc.close_streams()

    _test_cluster(f)


def test_worker_with_port_zero(loop):
    @gen.coroutine
    def f():
        c = Center('127.0.0.1')
        c.listen(8007)
        w = Worker(c.ip, c.port, ip='127.0.0.1')
        yield w._start()
        assert isinstance(w.port, int)
        assert w.port > 1024

    loop.run_sync(f)

@pytest.mark.slow
def test_worker_waits_for_center_to_come_up(loop):
    @gen.coroutine
    def f():
        w = Worker('127.0.0.1', 8007, ip='127.0.0.1')
        yield w._start()

    try:
        loop.run_sync(f, timeout=4)
    except TimeoutError:
        pass


@gen_cluster()
def test_worker_task(s, a, b):
    aa = rpc(ip=a.ip, port=a.port)
    yield aa.compute(task=(inc, 1), key='x')

    assert a.data['x'] == 2


@gen_cluster()
def test_worker_task_data(s, a, b):
    aa = rpc(ip=a.ip, port=a.port)
    yield aa.compute(task=2, key='x')

    assert a.data['x'] == 2


@gen_cluster()
def test_worker_task_bytes(s, a, b):
    aa = rpc(ip=a.ip, port=a.port)

    yield aa.compute(task=dumps((inc, 1)), key='x', serialized=True)
    assert a.data['x'] == 2

    yield aa.compute(function=dumps(inc), args=dumps((10,)), key='y', serialized=True)
    assert a.data['y'] == 11

# this worker doesn't exist