1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173 | import pytest
from distributed import Worker, WorkerPlugin
from distributed.utils_test import async_wait_for, gen_cluster, inc
class MyPlugin(WorkerPlugin):
name = "MyPlugin"
def __init__(self, data, expected_notifications=None):
self.data = data
self.expected_notifications = expected_notifications
def setup(self, worker):
assert isinstance(worker, Worker)
self.worker = worker
self.worker._my_plugin_status = "setup"
self.worker._my_plugin_data = self.data
self.observed_notifications = []
def teardown(self, worker):
self.worker._my_plugin_status = "teardown"
if self.expected_notifications is not None:
assert len(self.observed_notifications) == len(self.expected_notifications)
for expected, real in zip(
self.expected_notifications, self.observed_notifications
):
assert expected == real
def transition(self, key, start, finish, **kwargs):
self.observed_notifications.append(
{"key": key, "start": start, "finish": finish}
)
def release_key(self, key, state, cause, reason, report):
self.observed_notifications.append({"key": key, "state": state})
def release_dep(self, dep, state, report):
self.observed_notifications.append({"dep": dep, "state": state})
@gen_cluster(client=True, nthreads=[])
async def test_create_with_client(c, s):
await c.register_worker_plugin(MyPlugin(123))
worker = await Worker(s.address, loop=s.loop)
assert worker._my_plugin_status == "setup"
assert worker._my_plugin_data == 123
await worker.close()
assert worker._my_plugin_status == "teardown"
@gen_cluster(client=True, nthreads=[])
async def test_create_with_client_and_plugin_from_class(c, s):
await c.register_worker_plugin(MyPlugin, data=456)
worker = await Worker(s.address, loop=s.loop)
assert worker._my_plugin_status == "setup"
assert worker._my_plugin_data == 456
# Give the plugin a new name so that it registers
await c.register_worker_plugin(MyPlugin, name="new", data=789)
assert worker._my_plugin_data == 789
@gen_cluster(client=True, worker_kwargs={"plugins": [MyPlugin(5)]})
async def test_create_on_construction(c, s, a, b):
assert len(a.plugins) == len(b.plugins) == 1
assert a._my_plugin_status == "setup"
assert a._my_plugin_data == 5
@gen_cluster(nthreads=[("127.0.0.1", 1)], client=True)
async def test_normal_task_transitions_called(c, s, w):
expected_notifications = [
{"key": "task", "start": "waiting", "finish": "ready"},
{"key": "task", "start": "ready", "finish": "executing"},
{"key": "task", "start": "executing", "finish": "memory"},
{"key": "task", "state": "memory"},
]
plugin = MyPlugin(1, expected_notifications=expected_notifications)
await c.register_worker_plugin(plugin)
await c.submit(lambda x: x, 1, key="task")
await async_wait_for(lambda: not w.tasks, timeout=10)
@gen_cluster(nthreads=[("127.0.0.1", 1)], client=True)
async def test_failing_task_transitions_called(c, s, w):
def failing(x):
raise Exception()
expected_notifications = [
{"key": "task", "start": "waiting", "finish": "ready"},
{"key": "task", "start": "ready", "finish": "executing"},
{"key": "task", "start": "executing", "finish": "error"},
]
plugin = MyPlugin(1, expected_notifications=expected_notifications)
await c.register_worker_plugin(plugin)
with pytest.raises(Exception):
await c.submit(failing, 1, key="task")
@gen_cluster(
nthreads=[("127.0.0.1", 1)], client=True, worker_kwargs={"resources": {"X": 1}}
)
async def test_superseding_task_transitions_called(c, s, w):
expected_notifications = [
{"key": "task", "start": "waiting", "finish": "constrained"},
{"key": "task", "start": "constrained", "finish": "executing"},
{"key": "task", "start": "executing", "finish": "memory"},
{"key": "task", "state": "memory"},
]
plugin = MyPlugin(1, expected_notifications=expected_notifications)
await c.register_worker_plugin(plugin)
await c.submit(lambda x: x, 1, key="task", resources={"X": 1})
await async_wait_for(lambda: not w.tasks, timeout=10)
@gen_cluster(nthreads=[("127.0.0.1", 1)], client=True)
async def test_release_dep_called(c, s, w):
dsk = {"dep": 1, "task": (inc, "dep")}
expected_notifications = [
{"key": "dep", "start": "waiting", "finish": "ready"},
{"key": "dep", "start": "ready", "finish": "executing"},
{"key": "dep", "start": "executing", "finish": "memory"},
{"key": "task", "start": "waiting", "finish": "ready"},
{"key": "task", "start": "ready", "finish": "executing"},
{"key": "task", "start": "executing", "finish": "memory"},
{"key": "dep", "state": "memory"},
{"key": "task", "state": "memory"},
]
plugin = MyPlugin(1, expected_notifications=expected_notifications)
await c.register_worker_plugin(plugin)
await c.get(dsk, "task", sync=False)
await async_wait_for(lambda: not w.tasks, timeout=10)
@gen_cluster(nthreads=[("127.0.0.1", 1)], client=True)
async def test_registering_with_name_arg(c, s, w):
class FooWorkerPlugin:
def setup(self, worker):
if hasattr(worker, "foo"):
raise RuntimeError(f"Worker {worker.address} already has foo!")
worker.foo = True
responses = await c.register_worker_plugin(FooWorkerPlugin(), name="foo")
assert list(responses.values()) == [{"status": "OK"}]
async with Worker(s.address, loop=s.loop):
responses = await c.register_worker_plugin(FooWorkerPlugin(), name="foo")
assert list(responses.values()) == [{"status": "repeat"}] * 2
@gen_cluster(nthreads=[("127.0.0.1", 1)], client=True)
async def test_empty_plugin(c, s, w):
class EmptyPlugin:
pass
await c.register_worker_plugin(EmptyPlugin())
|