test_bcp.py 14.5 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
49
50
51
52
53
54
55
56
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
57
from ..dock import Host
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

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": {},
    "environment": {"name": "Python 2.7", "version": "1.3.0"},
    "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": {},
    "environment": {"name": "Python 2.7", "version": "1.3.0"},
    "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
174
        )


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


class ExcecutionTestCase(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
233
234


class TestBCP(ExcecutionTestCase):
    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
314
315
316
317
318
319
320
321
        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)

    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
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376

class TestExcecutionProcess(ExcecutionTestCase):
    address = "ipc://execution_feed"

    def tearDown(self):
        os.remove(self.address.split("//")[1])

    def setup_process(self):
        self.prepare_databases(CONFIGURATION1)
        process = ExecutionProcess(
            self.address, b"1", prefix, CONFIGURATION1, tmp_prefix, VERBOSE_BCP_LOGGING
        )
        process.start()
        return process

    def test_execution_process(self):
        ctx = zmq.Context()
        socket = ctx.socket(zmq.ROUTER)
        socket.linger = 0
        socket.bind(self.address)

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

        process = self.setup_process()

377
378
        done = False
        iterations = 0
379
        messages = []
380
381

        while True and iterations < self.MAX_ITERATION_COUNT:
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
            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
397
            iterations += 1
398
399
400
401
402

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

403
        self.assertTrue(iterations < self.MAX_ITERATION_COUNT)
404
405
406
407
408
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()
        processor = BeatComputationProcessor(poller, self.address, VERBOSE_BCP_LOGGING)

        process = self.setup_process()

        done = False
415
        iterations = 0
416
417

        messages = []
418
419

        while True and iterations < self.MAX_ITERATION_COUNT:
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
            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
435
            iterations += 1
436
437
438
439

        process.terminate()
        process.join()

440
        self.assertTrue(iterations < self.MAX_ITERATION_COUNT)
441
442
443
        self.assertTrue(done)
        self.assertEqual(process.queue.get(), "started")
        self.assertEqual(messages[0][1], BCP.BCPP_JOB_DONE)