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
40
41


# Tests for experiment execution

import os
import unittest
import multiprocessing
42
import queue
43
44
import zmq

Samuel GAIST's avatar
Samuel GAIST committed
45
46
import simplejson as json

47
48
from flaky import flaky

49
50
51
52
53
54
55
56
57
58
from ..bcpapi import BCP
from ..bcpapi.client import BeatComputationClient
from ..bcpapi.execution import ExecutionProcess
from ..bcpapi.processor import BeatComputationProcessor

from ..bcp import worker
from ..bcp import broker

from ..database import Database
from ..utils import find_free_port
59
from ..dock import Host
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82

from . import prefix, tmp_prefix
from . import VERBOSE_BCP_LOGGING

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


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": {},
83
    "environment": {"name": "Python for tests", "version": "1.3.0"},
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
    "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": {},
116
    "environment": {"name": "Python for tests", "version": "1.3.0"},
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
    "outputs": {
        "out": {
            "path": "40/61/b6/2df3c3bedd5366f4a625c5d87ffbf5a26007c46c456e9abf21b46c6681",
            "endpoint": "out",
            "hash": "4061b62df3c3bedd5366f4a625c5d87ffbf5a26007c46c456e9abf21b46c6681",
            "channel": "integers",
        }
    },
    "nb_slots": 1,
    "channel": "integers",
}


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


133
134
135
136
137
138
DEFAULT_MAX_ITERATION_COUNT = 30


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


139
140
141
142
143
144
145
146
147
148
149
150
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):
151
152
153
    def __init__(
        self, address, name, verbose, use_docker=False, docker_images_cache=None
    ):
154
155
156
157
158
        super(ZMQWorkerProcess, self).__init__()
        self.broker_address = address
        self.service_name = name
        self.verbose = verbose
        self.use_docker = use_docker
159
        self.docker_images_cache = None
160
161
162
163
164
165
166
167
168

    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,
169
            docker_images_cache=self.docker_images_cache,
170
171
172
173
174
175
        )


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


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

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
    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
217
218
        queue_messages = []
        while max_rounds > 0 and len(queue_messages) < 2:
219
            max_rounds -= 1
220
221
222
223
224
225
            try:
                message = self.queue.get(block=True, timeout=5)
            except queue.Empty:
                pass
            else:
                queue_messages.append(message)
226
227
228
229

        broker_p.terminate()
        broker_p.join()

230
231
232
        self.assertEqual(len(queue_messages), 2)
        self.assertEqual(queue_messages[0], "ready")
        self.assertEqual(queue_messages[1], "gone")
233
234


Samuel GAIST's avatar
Samuel GAIST committed
235
class TestBCP(ExecutionTestCase):
236
    use_docker = False
237
    docker_images_cache = None
238
239

    def setUp(self):
240
241
        super().setUp()

242
243
244
245
246
247
248
249
250
        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(
251
252
253
254
255
            broker_address,
            self.worker_name,
            VERBOSE_BCP_LOGGING,
            self.use_docker,
            self.docker_images_cache,
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
        )
        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
273
274
        iterations = 0
        while reply is None and iterations < self.MAX_ITERATION_COUNT:
275
276
277
278
            try:
                reply = self.client.recv()
            except KeyboardInterrupt:
                break
279
280
            else:
                iterations += 1
281

282
        self.assertTrue(iterations < self.MAX_ITERATION_COUNT)
283
284
285
286
287
288
289
290
291
292
293
294
295
296
        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 = []
297
298
        iterations = 0
        while len(messages) < 3 and iterations < self.MAX_ITERATION_COUNT:
299
300
301
302
303
304
305
            try:
                reply = self.client.recv()
            except KeyboardInterrupt:
                break
            else:
                if reply:
                    messages.append(reply)
306
                iterations += 1
307

308
        self.assertTrue(iterations < self.MAX_ITERATION_COUNT)
309
310
311
312
313
314
315
        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)

316
    @flaky(max_runs=3)
317
318
319
320
321
322
323
324
    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 = []
325
326
        iterations = 0
        while len(messages) < 3 and iterations < self.MAX_ITERATION_COUNT:
327
328
329
330
331
332
333
            try:
                reply = self.client.recv()
            except KeyboardInterrupt:
                break
            else:
                if reply:
                    messages.append(reply)
334
                iterations += 1
335

336
        self.assertTrue(iterations < self.MAX_ITERATION_COUNT)
337
338
339
340
341
342
343
344
345
346
347
348
        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

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

354

Samuel GAIST's avatar
Samuel GAIST committed
355
356
class TestExecutionProcess(ExecutionTestCase):
    REMOTE_ADDRESS = "ipc://execution_feed"
357
358

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

    def setup_process(self):
        self.prepare_databases(CONFIGURATION1)
        process = ExecutionProcess(
Samuel GAIST's avatar
Samuel GAIST committed
364
365
366
367
368
369
            self.REMOTE_ADDRESS,
            b"1",
            prefix,
            CONFIGURATION1,
            tmp_prefix,
            VERBOSE_BCP_LOGGING,
370
371
372
373
374
375
376
377
        )
        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
378
        socket.bind(self.REMOTE_ADDRESS)
379
380
381
382
383
384

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

        process = self.setup_process()

385
386
        done = False
        iterations = 0
387
        messages = []
388
389

        while True and iterations < self.MAX_ITERATION_COUNT:
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
            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
405
            iterations += 1
406
407
408
409
410

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

411
        self.assertTrue(iterations < self.MAX_ITERATION_COUNT)
412
413
414
415
416
417
        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
418
419
420
        processor = BeatComputationProcessor(
            poller, self.REMOTE_ADDRESS, VERBOSE_BCP_LOGGING
        )
421
422
423
424

        process = self.setup_process()

        done = False
425
        iterations = 0
426
427

        messages = []
428
429

        while True and iterations < self.MAX_ITERATION_COUNT:
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
            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
445
            iterations += 1
446
447
448
449

        process.terminate()
        process.join()

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