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
|
import pytest
from distributed import Scheduler, Worker, SchedulerPlugin
from distributed.utils_test import inc, gen_cluster, cleanup # noqa: F401
@gen_cluster(client=True)
async def test_simple(c, s, a, b):
class Counter(SchedulerPlugin):
def start(self, scheduler):
self.scheduler = scheduler
scheduler.add_plugin(self)
self.count = 0
def transition(self, key, start, finish, *args, **kwargs):
if start == "processing" and finish == "memory":
self.count += 1
counter = Counter()
counter.start(s)
assert counter in s.plugins
assert counter.count == 0
x = c.submit(inc, 1)
y = c.submit(inc, x)
z = c.submit(inc, y)
await z
assert counter.count == 3
s.remove_plugin(counter)
assert counter not in s.plugins
@gen_cluster(nthreads=[], client=False)
async def test_add_remove_worker(s):
events = []
class MyPlugin(SchedulerPlugin):
def add_worker(self, worker, scheduler):
assert scheduler is s
events.append(("add_worker", worker))
def remove_worker(self, worker, scheduler):
assert scheduler is s
events.append(("remove_worker", worker))
plugin = MyPlugin()
s.add_plugin(plugin)
assert events == []
a = Worker(s.address)
b = Worker(s.address)
await a
await b
await a.close()
await b.close()
assert events == [
("add_worker", a.address),
("add_worker", b.address),
("remove_worker", a.address),
("remove_worker", b.address),
]
events[:] = []
s.remove_plugin(plugin)
a = await Worker(s.address)
await a.close()
assert events == []
@gen_cluster(nthreads=[], client=False)
async def test_async_add_remove_worker(s):
events = []
class MyPlugin(SchedulerPlugin):
async def add_worker(self, worker, scheduler):
assert scheduler is s
events.append(("add_worker", worker))
async def remove_worker(self, worker, scheduler):
assert scheduler is s
events.append(("remove_worker", worker))
plugin = MyPlugin()
s.add_plugin(plugin)
assert events == []
async with Worker(s.address) as a:
async with Worker(s.address) as b:
pass
assert set(events) == {
("add_worker", a.address),
("add_worker", b.address),
("remove_worker", a.address),
("remove_worker", b.address),
}
events[:] = []
s.remove_plugin(plugin)
async with Worker(s.address):
pass
assert events == []
@pytest.mark.asyncio
async def test_lifecycle(cleanup):
class LifeCycle(SchedulerPlugin):
def __init__(self):
self.history = []
async def start(self, scheduler):
self.scheduler = scheduler
self.history.append("started")
async def close(self):
self.history.append("closed")
plugin = LifeCycle()
async with Scheduler(plugins=[plugin]) as s:
pass
assert plugin.history == ["started", "closed"]
assert plugin.scheduler is s
|
