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
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
|
from __future__ import annotations
import asyncio
import warnings
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}
)
@gen_cluster(client=True, nthreads=[])
async def test_create_with_client(c, s):
await c.register_worker_plugin(MyPlugin(123))
async with Worker(s.address) as worker:
assert worker._my_plugin_status == "setup"
assert worker._my_plugin_data == 123
assert worker._my_plugin_status == "teardown"
@gen_cluster(client=True, nthreads=[])
async def test_remove_with_client(c, s):
await c.register_worker_plugin(MyPlugin(123), name="foo")
await c.register_worker_plugin(MyPlugin(546), name="bar")
async with Worker(s.address) as worker:
# remove the 'foo' plugin
await c.unregister_worker_plugin("foo")
assert worker._my_plugin_status == "teardown"
# check that on the scheduler registered worker plugins we only have 'bar'
assert len(s.worker_plugins) == 1
assert "bar" in s.worker_plugins
# check on the worker plugins that we only have 'bar'
assert len(worker.plugins) == 1
assert "bar" in worker.plugins
# let's remove 'bar' and we should have none worker plugins
await c.unregister_worker_plugin("bar")
assert worker._my_plugin_status == "teardown"
assert not s.worker_plugins
assert not worker.plugins
@gen_cluster(client=True, nthreads=[])
async def test_remove_with_client_raises(c, s):
await c.register_worker_plugin(MyPlugin(123), name="foo")
async with Worker(s.address):
with pytest.raises(ValueError, match="bar"):
await c.unregister_worker_plugin("bar")
@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": "released", "finish": "waiting"},
{"key": "task", "start": "waiting", "finish": "ready"},
{"key": "task", "start": "ready", "finish": "executing"},
{"key": "task", "start": "executing", "finish": "memory"},
{"key": "task", "start": "memory", "finish": "released"},
{"key": "task", "start": "released", "finish": "forgotten"},
]
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.state.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": "released", "finish": "waiting"},
{"key": "task", "start": "waiting", "finish": "ready"},
{"key": "task", "start": "ready", "finish": "executing"},
{"key": "task", "start": "executing", "finish": "error"},
{"key": "task", "start": "error", "finish": "released"},
{"key": "task", "start": "released", "finish": "forgotten"},
]
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": "released", "finish": "waiting"},
{"key": "task", "start": "waiting", "finish": "ready"},
{"key": "task", "start": "waiting", "finish": "constrained"},
{"key": "task", "start": "constrained", "finish": "executing"},
{"key": "task", "start": "executing", "finish": "memory"},
{"key": "task", "start": "memory", "finish": "released"},
{"key": "task", "start": "released", "finish": "forgotten"},
]
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.state.tasks, timeout=10)
@gen_cluster(nthreads=[("127.0.0.1", 1)], client=True)
async def test_dependent_tasks(c, s, w):
dsk = {"dep": 1, "task": (inc, "dep")}
expected_notifications = [
{"key": "dep", "start": "released", "finish": "waiting"},
{"key": "dep", "start": "waiting", "finish": "ready"},
{"key": "dep", "start": "ready", "finish": "executing"},
{"key": "dep", "start": "executing", "finish": "memory"},
{"key": "task", "start": "released", "finish": "waiting"},
{"key": "task", "start": "waiting", "finish": "ready"},
{"key": "task", "start": "ready", "finish": "executing"},
{"key": "task", "start": "executing", "finish": "memory"},
{"key": "dep", "start": "memory", "finish": "released"},
{"key": "task", "start": "memory", "finish": "released"},
{"key": "task", "start": "released", "finish": "forgotten"},
{"key": "dep", "start": "released", "finish": "forgotten"},
]
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.state.tasks, timeout=10)
@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())
@gen_cluster(nthreads=[("127.0.0.1", 1)], client=True)
async def test_default_name(c, s, w):
class MyCustomPlugin(WorkerPlugin):
pass
await c.register_worker_plugin(MyCustomPlugin())
assert len(w.plugins) == 1
assert next(iter(w.plugins)).startswith("MyCustomPlugin-")
@gen_cluster(client=True, nthreads=[("", 1)])
async def test_assert_no_warning_no_overload(c, s, a):
"""Assert we do not receive a deprecation warning if we do not overload any
methods
"""
class Dummy(WorkerPlugin):
pass
with warnings.catch_warnings(record=True) as record:
await c.register_worker_plugin(Dummy())
assert await c.submit(inc, 1, key="x") == 2
while "x" in a.state.tasks:
await asyncio.sleep(0.01)
assert not record
@gen_cluster(nthreads=[("127.0.0.1", 1)], client=True)
async def test_WorkerPlugin_overwrite(c, s, w):
class MyCustomPlugin(WorkerPlugin):
name = "custom"
def setup(self, worker):
self.worker = worker
self.worker.foo = 0
def transition(self, *args, **kwargs):
self.worker.foo = 123
def teardown(self, worker):
del self.worker.foo
await c.register_worker_plugin(MyCustomPlugin())
assert w.foo == 0
await c.submit(inc, 0)
assert w.foo == 123
while s.tasks or w.state.tasks:
await asyncio.sleep(0.01)
class MyCustomPlugin(WorkerPlugin):
name = "custom"
def setup(self, worker):
self.worker = worker
self.worker.bar = 0
def transition(self, *args, **kwargs):
self.worker.bar = 456
def teardown(self, worker):
del self.worker.bar
await c.register_worker_plugin(MyCustomPlugin())
assert not hasattr(w, "foo")
assert w.bar == 0
await c.submit(inc, 0)
assert w.bar == 456
|
