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
Samuel GAIST's avatar
Samuel GAIST committed
46

47
48
from flaky import flaky

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

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


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


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


131
132
133
134
135
136
DEFAULT_MAX_ITERATION_COUNT = 30


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


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

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


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


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

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

        broker_p.terminate()
        broker_p.join()

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


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

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

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

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

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

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

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

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

352

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

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

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

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

        process = self.setup_process()

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

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

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

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

        process = self.setup_process()

        done = False
423
        iterations = 0
424
425

        messages = []
426
427

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

        process.terminate()
        process.join()

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