test_bcp.py 14.6 KB
Newer Older
1
2
3
#!/usr/bin/env python
# vim: set fileencoding=utf-8 :

Samuel GAIST's avatar
Samuel GAIST committed
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
###################################################################################
#                                                                                 #
# Copyright (c) 2019 Idiap Research Institute, http://www.idiap.ch/               #
# Contact: beat.support@idiap.ch                                                  #
#                                                                                 #
# Redistribution and use in source and binary forms, with or without              #
# modification, are permitted provided that the following conditions are met:     #
#                                                                                 #
# 1. Redistributions of source code must retain the above copyright notice, this  #
# list of conditions and the following disclaimer.                                #
#                                                                                 #
# 2. Redistributions in binary form must reproduce the above copyright notice,    #
# this list of conditions and the following disclaimer in the documentation       #
# and/or other materials provided with the distribution.                          #
#                                                                                 #
# 3. Neither the name of the copyright holder nor the names of its contributors   #
# may be used to endorse or promote products derived from this software without   #
# specific prior written permission.                                              #
#                                                                                 #
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND #
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED   #
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE          #
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE    #
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL      #
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR      #
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER      #
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,   #
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE   #
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.            #
#                                                                                 #
###################################################################################
35
36
37
38
39


# Tests for experiment execution

import multiprocessing
Samuel GAIST's avatar
Samuel GAIST committed
40
import os
41
import queue
Samuel GAIST's avatar
Samuel GAIST committed
42
import unittest
43

Samuel GAIST's avatar
Samuel GAIST committed
44
import simplejson as json
Samuel GAIST's avatar
Samuel GAIST committed
45
import zmq
46
47
from flaky import flaky

Samuel GAIST's avatar
Samuel GAIST committed
48
49
from ..bcp import broker
from ..bcp import worker
50
51
52
53
54
from ..bcpapi import BCP
from ..bcpapi.client import BeatComputationClient
from ..bcpapi.execution import ExecutionProcess
from ..bcpapi.processor import BeatComputationProcessor
from ..database import Database
55
from ..dock import Host
Samuel GAIST's avatar
Samuel GAIST committed
56
from ..utils import find_free_port
57
from . import VERBOSE_BCP_LOGGING
Samuel GAIST's avatar
Samuel GAIST committed
58
59
from . import prefix
from . import tmp_prefix
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79

# ----------------------------------------------------------


CONFIGURATION1 = {
    "queue": "queue",
    "inputs": {
        "in": {
            "set": "double",
            "protocol": "double",
            "database": "integers_db/1",
            "output": "a",
            "path": "ec/89/e5/6e161d2cb012ef6ac8acf59bf453a6328766f90dc9baba9eb14ea23c55.db",
            "endpoint": "a",
            "hash": "ec89e56e161d2cb012ef6ac8acf59bf453a6328766f90dc9baba9eb14ea23c55",
            "channel": "integers",
        }
    },
    "algorithm": "legacy/echo/1",
    "parameters": {},
80
    "environment": {"name": "Python for tests", "version": "1.3.0"},
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
    "outputs": {
        "out": {
            "path": "20/61/b6/2df3c3bedd5366f4a625c5d87ffbf5a26007c46c456e9abf21b46c6681",
            "endpoint": "out",
            "hash": "2061b62df3c3bedd5366f4a625c5d87ffbf5a26007c46c456e9abf21b46c6681",
            "channel": "integers",
        }
    },
    "nb_slots": 1,
    "channel": "integers",
}


# ----------------------------------------------------------


CONFIGURATION2 = {
    "queue": "queue",
    "inputs": {
        "in": {
            "set": "double",
            "protocol": "double",
            "database": "integers_db/1",
            "output": "a",
            "path": "ec/89/e5/6e161d2cb012ef6ac8acf59bf453a6328766f90dc9baba9eb14ea23c55.db",
            "endpoint": "a",
            "hash": "ec89e56e161d2cb012ef6ac8acf59bf453a6328766f90dc9baba9eb14ea23c55",
            "channel": "integers",
        }
    },
    "algorithm": "legacy/echo/1",
    "parameters": {},
113
    "environment": {"name": "Python for tests", "version": "1.3.0"},
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
    "outputs": {
        "out": {
            "path": "40/61/b6/2df3c3bedd5366f4a625c5d87ffbf5a26007c46c456e9abf21b46c6681",
            "endpoint": "out",
            "hash": "4061b62df3c3bedd5366f4a625c5d87ffbf5a26007c46c456e9abf21b46c6681",
            "channel": "integers",
        }
    },
    "nb_slots": 1,
    "channel": "integers",
}


# ----------------------------------------------------------


130
131
132
133
134
135
DEFAULT_MAX_ITERATION_COUNT = 30


# ----------------------------------------------------------


136
137
138
139
140
141
142
143
144
145
146
147
class ZMQBrokerProcess(multiprocessing.Process):
    def __init__(self, port, verbose, callbacks=None):
        super(ZMQBrokerProcess, self).__init__()
        self.port = port
        self.verbose = verbose
        self.callbacks = callbacks

    def run(self):
        return broker.run(self.port, verbose=self.verbose, callbacks=self.callbacks)


class ZMQWorkerProcess(multiprocessing.Process):
148
149
150
    def __init__(
        self, address, name, verbose, use_docker=False, docker_images_cache=None
    ):
151
152
153
154
155
        super(ZMQWorkerProcess, self).__init__()
        self.broker_address = address
        self.service_name = name
        self.verbose = verbose
        self.use_docker = use_docker
156
        self.docker_images_cache = None
157
158
159
160
161
162
163
164
165

    def run(self):
        return worker.run(
            self.broker_address,
            service_name=self.service_name,
            verbose=self.verbose,
            prefix=prefix,
            cache=tmp_prefix,
            use_docker=self.use_docker,
166
            docker_images_cache=self.docker_images_cache,
167
168
169
170
171
172
        )


# ----------------------------------------------------------


Samuel GAIST's avatar
Samuel GAIST committed
173
class ExecutionTestCase(unittest.TestCase):
174
175
176
177
178
    def setUp(self):
        self.MAX_ITERATION_COUNT = int(
            os.environ.get("BPC_MAX_ITERATION_COUNT", DEFAULT_MAX_ITERATION_COUNT)
        )

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
    def prepare_databases(self, configuration):
        for _, input_cfg in configuration["inputs"].items():
            database = Database(prefix, input_cfg["database"])
            view = database.view(input_cfg["protocol"], input_cfg["set"])
            view.index(os.path.join(tmp_prefix, input_cfg["path"]))


class TestBroker(unittest.TestCase):
    def __on_ready(self, name):
        self.queue.put("ready")

    def __on_gone(self, name):
        self.queue.put("gone")

    def setUp(self):
        self.queue = multiprocessing.Queue()

    def test_callback(self):
        worker_name = b"test_worker"

        port = find_free_port()
        broker_address = "tcp://localhost:{}".format(port)

        broker_p = ZMQBrokerProcess(
            port, VERBOSE_BCP_LOGGING, (self.__on_ready, self.__on_gone)
        )
        broker_p.start()

        worker = ZMQWorkerProcess(broker_address, worker_name, VERBOSE_BCP_LOGGING)
        worker.start()
        worker.join(2)  # Give the worker enough time to announce itself
        worker.terminate()
        worker.join()

        max_rounds = 5
214
215
        queue_messages = []
        while max_rounds > 0 and len(queue_messages) < 2:
216
            max_rounds -= 1
217
218
219
220
221
222
            try:
                message = self.queue.get(block=True, timeout=5)
            except queue.Empty:
                pass
            else:
                queue_messages.append(message)
223
224
225
226

        broker_p.terminate()
        broker_p.join()

227
228
229
        self.assertEqual(len(queue_messages), 2)
        self.assertEqual(queue_messages[0], "ready")
        self.assertEqual(queue_messages[1], "gone")
230
231


Samuel GAIST's avatar
Samuel GAIST committed
232
class TestBCP(ExecutionTestCase):
233
    use_docker = False
234
    docker_images_cache = None
235
236

    def setUp(self):
237
238
        super().setUp()

239
240
241
242
243
244
245
246
247
        self.worker_name = b"test_worker"

        port = find_free_port()
        broker_address = "tcp://localhost:{}".format(port)

        self.broker_p = ZMQBrokerProcess(port, VERBOSE_BCP_LOGGING)
        self.broker_p.start()

        self.worker = ZMQWorkerProcess(
248
249
250
251
252
            broker_address,
            self.worker_name,
            VERBOSE_BCP_LOGGING,
            self.use_docker,
            self.docker_images_cache,
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
        )
        self.worker.start()

        self.client = BeatComputationClient(broker_address, VERBOSE_BCP_LOGGING)

    def tearDown(self):
        self.worker.terminate()
        self.worker.join()
        self.broker_p.terminate()
        self.broker_p.join()
        self.client = None

    def test_cancel_unknown(self):
        request = [BCP.BCPE_CANCEL, b"1"]
        self.client.send(self.worker_name, request)

        reply = None
270
271
        iterations = 0
        while reply is None and iterations < self.MAX_ITERATION_COUNT:
272
273
274
275
            try:
                reply = self.client.recv()
            except KeyboardInterrupt:
                break
276
277
            else:
                iterations += 1
278

279
        self.assertTrue(iterations < self.MAX_ITERATION_COUNT)
280
281
282
283
284
285
286
287
288
289
290
291
292
293
        self.assertEqual(reply[1], BCP.BCPP_ERROR)
        self.assertEqual(reply[2], b"Unknown job: 1")

    def test_cancel(self):
        self.prepare_databases(CONFIGURATION1)
        job_id = b"1"

        request = [BCP.BCPE_EXECUTE, job_id, json.dumps(CONFIGURATION1).encode("utf-8")]
        self.client.send(self.worker_name, request)

        request = [BCP.BCPE_CANCEL, job_id]
        self.client.send(self.worker_name, request)

        messages = []
294
295
        iterations = 0
        while len(messages) < 3 and iterations < self.MAX_ITERATION_COUNT:
296
297
298
299
300
301
302
            try:
                reply = self.client.recv()
            except KeyboardInterrupt:
                break
            else:
                if reply:
                    messages.append(reply)
303
                iterations += 1
304

305
        self.assertTrue(iterations < self.MAX_ITERATION_COUNT)
306
307
308
309
310
311
312
        self.assertEqual(messages[0][1], BCP.BCPP_JOB_RECEIVED)
        self.assertEqual(messages[0][2], job_id)
        self.assertEqual(messages[1][1], BCP.BCPP_JOB_STARTED)
        self.assertEqual(messages[1][2], job_id)
        self.assertEqual(messages[2][1], BCP.BCPP_JOB_CANCELLED)
        self.assertEqual(messages[2][2], job_id)

313
    @flaky(max_runs=3)
314
315
316
317
318
319
320
321
    def test_execute(self):
        self.prepare_databases(CONFIGURATION1)
        job_id = b"1"

        request = [BCP.BCPE_EXECUTE, job_id, json.dumps(CONFIGURATION1).encode("utf-8")]
        self.client.send(self.worker_name, request)

        messages = []
322
323
        iterations = 0
        while len(messages) < 3 and iterations < self.MAX_ITERATION_COUNT:
324
325
326
327
328
329
330
            try:
                reply = self.client.recv()
            except KeyboardInterrupt:
                break
            else:
                if reply:
                    messages.append(reply)
331
                iterations += 1
332

333
        self.assertTrue(iterations < self.MAX_ITERATION_COUNT)
334
335
336
337
338
339
340
341
342
343
344
345
        self.assertEqual(messages[0][1], BCP.BCPP_JOB_RECEIVED)
        self.assertEqual(messages[0][2], job_id)
        self.assertEqual(messages[1][1], BCP.BCPP_JOB_STARTED)
        self.assertEqual(messages[1][2], job_id)
        self.assertEqual(messages[2][1], BCP.BCPP_JOB_DONE)
        self.assertEqual(messages[2][2], job_id)
        self.assertEqual(len(messages[2]), 4)


class TestBCPDocker(TestBCP):
    use_docker = True

346
347
    @classmethod
    def setUpClass(cls):
348
        cls.docker_images_cache = os.path.join(tmp_prefix, "docker_images_cache.json")
349
350
        cls.host = Host(images_cache=cls.docker_images_cache, raise_on_errors=False)

351

Samuel GAIST's avatar
Samuel GAIST committed
352
353
class TestExecutionProcess(ExecutionTestCase):
    REMOTE_ADDRESS = "ipc://execution_feed"
354
355

    def tearDown(self):
Samuel GAIST's avatar
Samuel GAIST committed
356
        os.remove(self.REMOTE_ADDRESS.split("//")[1])
357
358
359
360

    def setup_process(self):
        self.prepare_databases(CONFIGURATION1)
        process = ExecutionProcess(
Samuel GAIST's avatar
Samuel GAIST committed
361
362
363
364
365
366
            self.REMOTE_ADDRESS,
            b"1",
            prefix,
            CONFIGURATION1,
            tmp_prefix,
            VERBOSE_BCP_LOGGING,
367
368
369
370
371
372
373
374
        )
        process.start()
        return process

    def test_execution_process(self):
        ctx = zmq.Context()
        socket = ctx.socket(zmq.ROUTER)
        socket.linger = 0
Samuel GAIST's avatar
Samuel GAIST committed
375
        socket.bind(self.REMOTE_ADDRESS)
376
377
378
379
380
381

        poller = zmq.Poller()
        poller.register(socket, zmq.POLLIN)

        process = self.setup_process()

382
383
        done = False
        iterations = 0
384
        messages = []
385
386

        while True and iterations < self.MAX_ITERATION_COUNT:
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
            try:
                items = poller.poll(1000)
            except KeyboardInterrupt:
                break

            if items:
                msg = socket.recv_multipart()
                messages.append(msg)

                result = msg[1]
                if result == BCP.BCPP_JOB_DONE:
                    done = True
                    break
                elif result in [BCP.BCPP_JOB_ERROR, BCP.BCPP_ERROR]:
                    break
402
            iterations += 1
403
404
405
406
407

        process.terminate()
        process.join()
        ctx.destroy()

408
        self.assertTrue(iterations < self.MAX_ITERATION_COUNT)
409
410
411
412
413
414
        self.assertTrue(done)
        self.assertEqual(process.queue.get(), "started")
        self.assertEqual(messages[0][1], BCP.BCPP_JOB_DONE)

    def test_processor(self):
        poller = zmq.Poller()
Samuel GAIST's avatar
Samuel GAIST committed
415
416
417
        processor = BeatComputationProcessor(
            poller, self.REMOTE_ADDRESS, VERBOSE_BCP_LOGGING
        )
418
419
420
421

        process = self.setup_process()

        done = False
422
        iterations = 0
423
424

        messages = []
425
426

        while True and iterations < self.MAX_ITERATION_COUNT:
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
            try:
                items = poller.poll(1000)
            except KeyboardInterrupt:
                break

            if items:
                msg = processor.process()
                messages.append(msg)

                result = msg[1]
                if result == BCP.BCPP_JOB_DONE:
                    done = True
                    break
                elif result in [BCP.BCPP_JOB_ERROR, BCP.BCPP_ERROR]:
                    break
442
            iterations += 1
443
444
445
446

        process.terminate()
        process.join()

447
        self.assertTrue(iterations < self.MAX_ITERATION_COUNT)
448
449
450
        self.assertTrue(done)
        self.assertEqual(process.queue.get(), "started")
        self.assertEqual(messages[0][1], BCP.BCPP_JOB_DONE)